diff --git a/deep_cluster/hardware/__init__.py b/deep_cluster/hardware/__init__.py
index 9dc1dac..dda35e4 100644
--- a/deep_cluster/hardware/__init__.py
+++ b/deep_cluster/hardware/__init__.py
@@ -6,8 +6,8 @@ from deep_cluster.hardware.device import (
     device_numa_node,
     get_actual_num_gpus,
     get_bf16_tflops_limit,
-    is_sm_major_10,
     is_sm_major_9,
+    is_sm_major_10,
     num_gpus,
     physical_device_count,
     physical_device_id,
diff --git a/deepinfer/h20_ops/deep_select_sm90.py b/deepinfer/h20_ops/deep_select_sm90.py
index 1a0257c..35447c7 100644
--- a/deepinfer/h20_ops/deep_select_sm90.py
+++ b/deepinfer/h20_ops/deep_select_sm90.py
@@ -58,11 +58,8 @@ def topk(
 
     values, indices = torch.topk(scores, k, dim=-1, sorted=True)
 
-    if valid_mask is None:
-        keep = None
-    else:
-        # `indices` are original column ids, so the mask still applies after the top-k pick.
-        keep = valid_mask.gather(1, indices)
+    # `indices` are original column ids, so the mask still applies after the top-k pick.
+    keep = None if valid_mask is None else valid_mask.gather(1, indices)
 
     out_idx = indices.to(indices_type)
     if output_idx_offset is not None:
@@ -80,9 +77,7 @@ def topk(
         if keep is not None:
             out_val = torch.where(keep, out_val, torch.full_like(out_val, value_oob_fill_value))
         if k < topk:
-            out_val = torch.cat(
-                [out_val, out_val.new_full((num_rows, topk - k), value_oob_fill_value)], dim=1
-            )
+            out_val = torch.cat([out_val, out_val.new_full((num_rows, topk - k), value_oob_fill_value)], dim=1)
     else:
         out_val = None
 
diff --git a/deepinfer/h20_ops/engram_local.py b/deepinfer/h20_ops/engram_local.py
index e999888..82778bb 100644
--- a/deepinfer/h20_ops/engram_local.py
+++ b/deepinfer/h20_ops/engram_local.py
@@ -171,12 +171,9 @@ class LocalEngramGather:
             rows, rows_sf = self._gather_layer(
                 self.weights[layer], self.sfs[layer] if self.sfs else None, flat, self.num_local_entries[layer]
             )
-            if rows_sf is not None:
-                payload = self._requant(rows, rows_sf, num_tokens)
-            else:
-                payload = rows.view(num_tokens, column)
+            payload = self._requant(rows, rows_sf, num_tokens) if rows_sf is not None else rows.view(num_tokens, column)
 
-            def hook(payload=payload):
+            def hook(payload: torch.Tensor = payload) -> torch.Tensor:
                 return payload
 
             hooks.append(hook)
diff --git a/deepinfer/h20_ops/quant.py b/deepinfer/h20_ops/quant.py
index 6a1e0c3..27e8604 100644
--- a/deepinfer/h20_ops/quant.py
+++ b/deepinfer/h20_ops/quant.py
@@ -60,7 +60,9 @@ def _requant(x: torch.Tensor, dst_block: int) -> tuple[torch.Tensor, torch.Tenso
 
 
 @torch.no_grad()
-def requant_dense_fp8(weight: torch.Tensor, sf: torch.Tensor, src_block: int = FP4_BLOCK, dst_block: int = 128):
+def requant_dense_fp8(
+    weight: torch.Tensor, sf: torch.Tensor, src_block: int = FP4_BLOCK, dst_block: int = 128
+) -> tuple[torch.Tensor, torch.Tensor]:
     """FP8 dense weight with per-``src_block`` scales → FP8 with per-``dst_block`` FP32 scales."""
     rows, cols = weight.shape
     x = weight.to(torch.float32) * _expand_blocks(sf.to(torch.float32), rows, cols, src_block)
diff --git a/deepinfer/h20_ops/sparse_mla_sm90.py b/deepinfer/h20_ops/sparse_mla_sm90.py
index c38a521..38ec025 100644
--- a/deepinfer/h20_ops/sparse_mla_sm90.py
+++ b/deepinfer/h20_ops/sparse_mla_sm90.py
@@ -18,9 +18,9 @@ the ``padded_H`` handling for head counts below 64, and the ``REPLICATE_H`` head
 
 from __future__ import annotations
 
-import torch
 import tilelang
 import tilelang.language as T
+import torch
 
 LOG2E = 1.44269504
 
@@ -33,9 +33,9 @@ LOG2E = 1.44269504
     },
 )
 def sparse_mla_fwd(
-    Q,
-    KV,
-    Indices,
+    Q: torch.Tensor,
+    KV: torch.Tensor,
+    Indices: torch.Tensor,
     heads: int,
     dim: int,
     topk: int,
@@ -45,12 +45,19 @@ def sparse_mla_fwd(
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
-):
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Sparse MLA forward: out[b, s, h, dim], lse[b, s, h].
 
     Q: [b, s, h, dim] bf16, KV: [b, s_kv, kv_group, dim] bf16, Indices: [b, s, kv_group, topk] int32.
-    An index is valid when ``0 <= idx <= q_position``; larger values (and the padding used to make
-    ``topk`` a multiple of ``block_I``) are masked out.
+    Negative indices (the padding that makes ``topk`` a multiple of ``block_I``) are always masked.
+
+    ``is_causal`` additionally drops indices above the query's row number. **The engine must pass
+    False**: `kv_index_cat` shifts the compressed-KV block by
+    `total_k_swa_cached + total_k_swa_cur`, so indices index the assembled
+    `[swa cache | current swa | compressed]` buffer and say nothing about the row number. Masking
+    on it drops nearly every index, `sumexp` stays 0, and the whole layer turns to NaN. Causality
+    is already guaranteed upstream -- the indexer only scores positions before each query, and the
+    sliding-window column block is a window by construction.
     """
     assert sm_scale is not None, "pass the explicit attention scale"
     scale = float(sm_scale * LOG2E)  # log2(e): the kernel folds the exponent into exp2
diff --git a/deepinfer/hugging_face/weights.py b/deepinfer/hugging_face/weights.py
index fb5a002..955a1ad 100644
--- a/deepinfer/hugging_face/weights.py
+++ b/deepinfer/hugging_face/weights.py
@@ -33,9 +33,9 @@ import torch
 
 from deepinfer import PLATFORM
 from deepinfer.hugging_face.names import hf_layer_module_prefix, hf_to_model, resolve_hf_name, scale_name
-from deepinfer.model_executor.layers.linear import get_fp8_casting_granularity
 from deepinfer.hugging_face.parameters import model_params, set_ue8m0_sf
 from deepinfer.logger import init_logger
+from deepinfer.model_executor.layers.linear import get_fp8_casting_granularity
 
 logger = init_logger(__name__)
 
diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py
index 4b939a2..6aa27d4 100644
--- a/deepinfer/model_executor/models/decoder_layer.py
+++ b/deepinfer/model_executor/models/decoder_layer.py
@@ -52,6 +52,7 @@ def _check_finite(tag: str, tensor: torch.Tensor) -> None:
         bad = int((~torch.isfinite(tensor)).sum())
         raise RuntimeError(f"[nan-debug] {tag}: {bad}/{tensor.numel()} non-finite values, shape={tuple(tensor.shape)}")
 
+
 logger = init_logger(__name__)
 
 
diff --git a/deepinfer/model_executor/models/v4_modules/index_attn.py b/deepinfer/model_executor/models/v4_modules/index_attn.py
index 550a0b1..a9ea4f1 100644
--- a/deepinfer/model_executor/models/v4_modules/index_attn.py
+++ b/deepinfer/model_executor/models/v4_modules/index_attn.py
@@ -124,13 +124,9 @@ class MQA(nn.Module):
         # which has no SM90 implementation (`DG_HOST_ASSERT(arch_major == 10)`), so Hopper runs
         # the dense path for every layer. Selection is identical, only the scoring is wider.
         self.is_candidate_topk_source = (
-            shared_topology is not None
-            and shared_topology.is_candidate_source_layer(self.layer_idx)
-            and not PLATFORM.is_sm90
-        )
-        self.use_candidate_topk = (
-            candidate_source is not None and not self.is_candidate_topk_source and not PLATFORM.is_sm90
+            shared_topology is not None and shared_topology.is_candidate_source_layer(self.layer_idx) and not PLATFORM.is_sm90
         )
+        self.use_candidate_topk = candidate_source is not None and not self.is_candidate_topk_source and not PLATFORM.is_sm90
         if self.use_candidate_topk:
             assert candidate_cfg is not None
             assert self.n_head <= 32 and self.head_dim == 128
@@ -518,9 +514,7 @@ class MQA(nn.Module):
         page_tokens = kv_cache.shape[1]
         page_bytes = page_tokens * (head_dim + 4)
         # One contiguous byte row per page, following the tensor's real block stride.
-        pages = torch.as_strided(
-            kv_cache, (kv_cache.shape[0], page_bytes), (kv_cache.stride(0), 1), kv_cache.storage_offset()
-        )
+        pages = torch.as_strided(kv_cache, (kv_cache.shape[0], page_bytes), (kv_cache.stride(0), 1), kv_cache.storage_offset())
 
         out = torch.full((num_rows, self.max_kv_len), float("-inf"), dtype=torch.float32, device=q.device)
         # One host read of the per-row context lengths; the engine forbids syncs during the
@@ -577,9 +571,7 @@ class MQA(nn.Module):
             f"decode_index_attn-part-{self.layer_idx}", enabled=self.layer_idx == second_layer, scale=len(same_layers)
         ):
             p_sum = (
-                self._decode_p_sum_sm90(
-                    q_fp8or4, q_sf, kv_cache, weights, klen, input_metadata.block_tables
-                )
+                self._decode_p_sum_sm90(q_fp8or4, q_sf, kv_cache, weights, klen, input_metadata.block_tables)
                 if PLATFORM.is_sm90
                 else self._compute_decode_p_sum(
                     (q_fp8or4, q_sf),
diff --git a/scripts/h20_ask.py b/scripts/h20_ask.py
index 4839dc4..faef608 100644
--- a/scripts/h20_ask.py
+++ b/scripts/h20_ask.py
@@ -76,9 +76,7 @@ def main(dp_size: int, dp_rank: int) -> None:
 
     if args.ids:
         ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
-        outputs = llm.generate(
-            prompts=None, prompt_token_ids=[ids], sampling_params=sampling_params[:1], use_tqdm=False
-        )
+        outputs = llm.generate(prompts=None, prompt_token_ids=[ids], sampling_params=sampling_params[:1], use_tqdm=False)
         shown = [f"<ids len={len(ids)}>"]
     else:
         outputs = llm.generate(prompts=args.prompt, sampling_params=sampling_params, use_tqdm=False)
diff --git a/scripts/h20_attn_check.py b/scripts/h20_attn_check.py
index c20e6f8..a951876 100644
--- a/scripts/h20_attn_check.py
+++ b/scripts/h20_attn_check.py
@@ -6,11 +6,9 @@ Shapes follow DeepSeek-V4.1-Flash: head_dim 512 with RoPE inside, 64 query heads
 """
 
 import torch
-
 from tilelang.profiler import do_bench
 
 from deepinfer.h20_ops.attn_sm90 import (
-    CACHE_ELEMENTS,
     CACHE_ELEMENTS_PER_SF,
     decode_mla_mqa_sparse,
     prefill_attn_varlen_sparse,
@@ -21,13 +19,13 @@ HEADS = 64
 DIM = 512
 
 
-def ref_attention(q, kv, indices, scale, sink=None):
+def ref_attention(q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, scale: float, sink: torch.Tensor | None = None) -> torch.Tensor:
     """q [n, h, d], kv [m, 1, d], indices [n, topk] -> [n, h, d] float32."""
     n, h, d = q.shape
     topk = indices.shape[1]
     valid = indices >= 0
     idx = indices.clamp_min(0)
-    k = kv[idx.reshape(-1), 0, :].view(n, topk, d).float()      # [n, topk, d]
+    k = kv[idx.reshape(-1), 0, :].view(n, topk, d).float()  # [n, topk, d]
     v = k
     scores = torch.einsum("nhd,nkd->nhk", q.float(), k) * scale  # [n, h, topk]
     scores = scores.masked_fill(~valid.unsqueeze(1), float("-inf"))
@@ -39,7 +37,7 @@ def ref_attention(q, kv, indices, scale, sink=None):
     return out, lse
 
 
-def make_indices(n, m, topk, gen):
+def make_indices(n: int, m: int, topk: int, gen: torch.Generator) -> torch.Tensor:
     """Causal top-k like the indexer produces: random subset of [0, position], -1 padded."""
     indices = torch.full((n, topk), -1, dtype=torch.int32, device=DEV)
     for t in range(n):
@@ -50,7 +48,7 @@ def make_indices(n, m, topk, gen):
     return indices
 
 
-def quantize_cache(kv, block_size=64):
+def quantize_cache(kv: torch.Tensor, block_size: int = 64) -> torch.Tensor:
     """Build a V41-format paged cache ([rows, 528] uint8) from bf16 kv [m, 1, d]."""
     m = kv.shape[0]
     values = kv[:, 0, :].float()
@@ -61,9 +59,7 @@ def quantize_cache(kv, block_size=64):
     exp = exp.clamp_min(-126)
     scale = torch.pow(2.0, exp)
     q = (groups / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
-    packed = torch.cat(
-        [q.reshape(m, DIM).view(torch.uint8), (exp + 127).to(torch.uint8)], dim=1
-    )  # [m, 528]
+    packed = torch.cat([q.reshape(m, DIM).view(torch.uint8), (exp + 127).to(torch.uint8)], dim=1)  # [m, 528]
     deq = (q.float() * scale.unsqueeze(-1)).reshape(m, DIM).to(torch.bfloat16)
     rows = packed.contiguous()
     n_blocks = (m + block_size - 1) // block_size
@@ -74,7 +70,7 @@ def quantize_cache(kv, block_size=64):
     return rows.view(n_blocks, block_size, 1, -1), deq
 
 
-def main():
+def main() -> None:
     torch.manual_seed(0)
     gen = torch.Generator(device="cpu").manual_seed(0)
     scale = DIM**-0.5
@@ -119,10 +115,7 @@ def main():
     kv_empty = torch.randn(64, 1, DIM, device=DEV, dtype=torch.bfloat16)
     idx_empty = torch.full((8, 128), -1, dtype=torch.int32, device=DEV)
     out_e, lse_e = prefill_attn_varlen_sparse(q_empty, kv_empty, idx_empty, scale)
-    print(
-        f"empty window: finite={bool(torch.isfinite(out_e).all())} "
-        f"zero={bool((out_e == 0).all())} lse_max={lse_e.max().item():.3e}"
-    )
+    print(f"empty window: finite={bool(torch.isfinite(out_e).all())} zero={bool((out_e == 0).all())} lse_max={lse_e.max().item():.3e}")
     assert torch.isfinite(out_e).all() and (out_e == 0).all()
 
     # ---------------- decode: paged fp8 cache + gather ----------------
@@ -132,15 +125,15 @@ def main():
     cache, kv_deq = quantize_cache(kv2)
 
     idx2 = torch.randint(0, m_dec, (n_dec, topk_dec), device=DEV, dtype=torch.int32)
-    dec, _ = decode_mla_mqa_sparse(
-        q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink
-    )
+    dec, _ = decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink)
     ref2, _ = ref_attention(q2, kv_deq.unsqueeze(1), idx2, scale, sink=sink)
     err2 = (dec.float() - ref2).abs().max().item()
     corr2 = torch.cosine_similarity(dec.float().flatten(), ref2.flatten(), dim=0).item()
     print(f"decode (n={n_dec}, cache={m_dec}, topk={topk_dec}): max_abs_err={err2:.4f} corr={corr2:.6f}")
 
-    ms2 = do_bench(lambda: decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink), warmup=5, rep=20)
+    ms2 = do_bench(
+        lambda: decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink), warmup=5, rep=20
+    )
     print(f"decode perf: batch={n_dec} topk={topk_dec} -> {ms2:.2f} ms/step")
 
     # ---------------- perf: full-size prefill ----------------
diff --git a/scripts/h20_engram_check.py b/scripts/h20_engram_check.py
index 503f1ef..e721e67 100644
--- a/scripts/h20_engram_check.py
+++ b/scripts/h20_engram_check.py
@@ -1,23 +1,21 @@
 """Validate LocalEngramGather against a torch reference, single-process and multi-rank.
 
-    torchrun --nproc_per_node=8 scripts/h20_engram_check.py
+torchrun --nproc_per_node=8 scripts/h20_engram_check.py
 """
 
-import os
-
 import torch
 import torch.distributed as dist
 
 from deepinfer.h20_ops.engram_local import LocalEngramGather
 
 HIDDEN = 256
-N_LOCAL = 4096          # rows per rank
-ENTRIES = 24            # 3 n-gram tables x 8 heads
+N_LOCAL = 4096  # rows per rank
+ENTRIES = 24  # 3 n-gram tables x 8 heads
 TOKENS = 128
 LAYERS = 2
 
 
-def quantize_table(rows, gen):
+def quantize_table(rows: torch.Tensor, gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]:
     """fp8 table + per-32 ue8m0 (stored as fp32) scales, like the checkpoint."""
     groups = rows.float().view(rows.shape[0], HIDDEN // 32, 32)
     amax = groups.abs().amax(-1).clamp_min(1e-30)
@@ -27,7 +25,7 @@ def quantize_table(rows, gen):
     return q, scale
 
 
-def main():
+def main() -> None:
     dist.init_process_group("nccl")
     rank, world = dist.get_rank(), dist.get_world_size()
     torch.cuda.set_device(rank)
@@ -63,7 +61,7 @@ def main():
         dist.all_gather_into_tensor(ref_rows[layer], tables[layer])
         dist.all_gather_into_tensor(ref_sf[layer], scales[layer])
 
-    def reference(layer):
+    def reference(layer: int) -> torch.Tensor:
         flat = idx[layer].reshape(-1).long()
         valid = flat >= 0
         safe = flat.clamp_min(0)
@@ -82,10 +80,7 @@ def main():
     ref_deq = ref.float().view(TOKENS, ENTRIES, HIDDEN)
     err = (got_vals - ref_deq).abs().max().item()
     rel = err / max(ref_deq.abs().max().item(), 1e-9)
-    print(
-        f"[rank {rank}] out={tuple(a.shape)} sf={None if sf is None else tuple(sf.shape)} "
-        f"max_abs={err:.4f} rel={rel:.4f}"
-    )
+    print(f"[rank {rank}] out={tuple(a.shape)} sf={None if sf is None else tuple(sf.shape)} max_abs={err:.4f} rel={rel:.4f}")
     print(f"[rank {rank}] invalid slot all-zero: {bool((got_vals[0, 0] == 0).all())}")
 
     # Non-owner traffic actually exercised?
diff --git a/scripts/h20_expert_weight_check.py b/scripts/h20_expert_weight_check.py
index 7077a27..3366c02 100644
--- a/scripts/h20_expert_weight_check.py
+++ b/scripts/h20_expert_weight_check.py
@@ -105,9 +105,7 @@ def main() -> None:
         assert ref_vals.shape == ours_vals.shape, f"{ref_vals.shape} vs {ours_vals.shape}"
         num = (ours_vals - ref_vals).norm().item()
         den = ref_vals.norm().item()
-        cos = torch.nn.functional.cosine_similarity(
-            ours_vals.reshape(-1), ref_vals.reshape(-1), dim=0
-        ).item()
+        cos = torch.nn.functional.cosine_similarity(ours_vals.reshape(-1), ref_vals.reshape(-1), dim=0).item()
         print(f"  ref scale: e8m0 {tuple(ref_sf.shape)}   ours: fp32 {tuple(ours_sf.shape)}")
         print(f"  |ref|={den:.4f}  |ours|={ours_vals.norm().item():.4f}  ratio={ours_vals.norm().item() / den:.5f}")
         print(f"  rel_err={num / den:.5f}  cos={cos:.7f}")
diff --git a/scripts/h20_real_requant_check.py b/scripts/h20_real_requant_check.py
index faf5fa5..8be9440 100644
--- a/scripts/h20_real_requant_check.py
+++ b/scripts/h20_real_requant_check.py
@@ -1,11 +1,15 @@
 """Validate the SM90 dense requant on a real checkpoint tensor (wq_a + shared expert)."""
+
 import json
+
 import torch
 from safetensors import safe_open
+
 from deepinfer.h20_ops.quant import requant_dense_fp8, ue8m0_to_fp32
 
 root = "/data/models/DeepSeek-V4.1-Flash-1017"
-idx = json.load(open(f"{root}/model.safetensors.index.json"))["weight_map"]
+with open(f"{root}/model.safetensors.index.json") as f:
+    idx = json.load(f)["weight_map"]
 
 targets = [
     "layers.0.attn.wq_a.weight",
diff --git a/scripts/h20_smoke_check.py b/scripts/h20_smoke_check.py
index 23af07e..208bbbe 100644
--- a/scripts/h20_smoke_check.py
+++ b/scripts/h20_smoke_check.py
@@ -1,6 +1,9 @@
 """SM90 sanity checks: DeepGEMM FP8 1D2D (fp32 scales) + the H20 requant helpers."""
-import torch, deep_gemm
-from deepinfer.h20_ops.quant import ue8m0_to_fp32, requant_dense_fp8, requant_expert_fp4, _unpack_fp4
+
+import deep_gemm
+import torch
+
+from deepinfer.h20_ops.quant import _unpack_fp4, requant_dense_fp8, requant_expert_fp4, ue8m0_to_fp32
 
 torch.manual_seed(0)
 dev = "cuda"
@@ -12,7 +15,8 @@ blk = 128
 x = torch.randn(M, K, device=dev, dtype=torch.bfloat16)
 w = torch.randn(N, K, device=dev, dtype=torch.bfloat16)
 
-def q_token(t, blk):
+
+def q_token(t: torch.Tensor, blk: int) -> tuple[torch.Tensor, torch.Tensor]:
     """A operand: per-token 1x128 blocks -> sf (rows, cols/128) fp32."""
     r, c = t.shape
     tf = t.float()
@@ -21,7 +25,8 @@ def q_token(t, blk):
     q = (tf.view(r, c // blk, blk) / s[:, :, None]).to(torch.float8_e4m3fn)
     return q.view(r, c), s.contiguous()
 
-def q_blocks(t, blk):
+
+def q_blocks(t: torch.Tensor, blk: int) -> tuple[torch.Tensor, torch.Tensor]:
     """B operand: 128x128 blocks -> sf (rows/128, cols/128) fp32."""
     r, c = t.shape
     tf = t.float()
@@ -30,26 +35,27 @@ def q_blocks(t, blk):
     q = (tf.view(r // blk, blk, c // blk, blk) / s[:, None, :, None]).to(torch.float8_e4m3fn)
     return q.view(r, c), s.contiguous()
 
+
 xq, xsf = q_token(x, blk)
 wq, wsf = q_blocks(w, blk)
 out = torch.empty(M, N, device=dev, dtype=torch.bfloat16)
 deep_gemm.fp8_fp4_gemm_nt((xq, xsf), (wq, wsf), out, None, disable_ue8m0_cast=True)
-print('recipe default (1,128,128) OK')
+print("recipe default (1,128,128) OK")
 ref = x @ w.t()
 err = (out.float() - ref).abs().max().item()
 rel = err / ref.abs().max().item()
 print(f"gemm 1D2D fp32-sf on SM90: max_abs_err={err:.4f} rel={rel:.5f}")
 
 # ---------------- 2. per-token activation quant via the runtime recipe ----------
-params = {"fmt": "e4m3", "num_per_channels": 128, "round_sf": True,
-          "use_packed_ue8m0": False, "use_tma_aligned_col_major_sf": True}
+params = {"fmt": "e4m3", "num_per_channels": 128, "round_sf": True, "use_packed_ue8m0": False, "use_tma_aligned_col_major_sf": True}
 import tile_kernels
+
 a8, asf = tile_kernels.quant.per_token_cast(x, **params)
 print("per_token_cast ->", a8.shape, a8.dtype, asf.shape, asf.dtype)
 
 # ---------------- 3. expert FP4 -> FP8 requant round-trip ----------------------
 n_out, k = 256, 512
-vals = torch.tensor([0., .5, 1., 1.5, 2., 3., 4., 6., 0., -.5, -1., -1.5, -2., -3., -4., -6.], device=dev)
+vals = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], device=dev)
 codes = torch.randint(0, 16, (n_out, k), device=dev)
 ref_x = vals[codes]
 packed = ((codes[..., 1::2] << 4) | codes[..., 0::2]).to(torch.uint8).view(torch.int8).contiguous()
@@ -61,14 +67,23 @@ x_true = ref_x * ue8m0_to_fp32(sf_e8m0).repeat_interleave(32, dim=1)
 
 w8, sf = requant_expert_fp4(packed, ue8m0_to_fp32(sf_e8m0), out_block=128)
 xr = w8.to(torch.float32) * sf.t().repeat_interleave(128, 0).repeat_interleave(128, 1)
-print("expert requant: fp8", tuple(w8.shape), "sf", tuple(sf.shape),
-      "max_abs_err", (xr - x_true).abs().max().item(),
-      "rel", ((xr - x_true).abs().max() / x_true.abs().max()).item())
+print(
+    "expert requant: fp8",
+    tuple(w8.shape),
+    "sf",
+    tuple(sf.shape),
+    "max_abs_err",
+    (xr - x_true).abs().max().item(),
+    "rel",
+    ((xr - x_true).abs().max() / x_true.abs().max()).item(),
+)
 
 # ---------------- 4. dense 32->128 requant ------------------------------------
-w_ref = (torch.randn(256, 512, device=dev, dtype=torch.bfloat16) / 8)
+w_ref = torch.randn(256, 512, device=dev, dtype=torch.bfloat16) / 8
 w8d, sfd = q_blocks(w_ref, 32)
-sf32 = ue8m0_to_fp32((torch.randint(118, 130, (256 // 32, 512 // 32), device=dev).to(torch.int32)).to(torch.uint8).view(torch.float8_e8m0fnu))
+sf32 = ue8m0_to_fp32(
+    (torch.randint(118, 130, (256 // 32, 512 // 32), device=dev).to(torch.int32)).to(torch.uint8).view(torch.float8_e8m0fnu)
+)
 src = w8d.float() * sf32.repeat_interleave(32, 0).repeat_interleave(32, 1)
 w8n, sfn = requant_dense_fp8(w8d, sf32, dst_block=128)
 rec = w8n.float() * sfn.repeat_interleave(128, 0).repeat_interleave(128, 1)
diff --git a/scripts/h20_topk_check.py b/scripts/h20_topk_check.py
index d4acd49..ca5f9d5 100644
--- a/scripts/h20_topk_check.py
+++ b/scripts/h20_topk_check.py
@@ -11,7 +11,14 @@ import torch
 from deepinfer.h20_ops.deep_select_sm90 import topk
 
 
-def ref_topk(scores, k, end=None, offset=None, fill=-1, val_fill=float("-inf")):
+def ref_topk(
+    scores: torch.Tensor,
+    k: int,
+    end: torch.Tensor | None = None,
+    offset: int | None = None,
+    fill: int = -1,
+    val_fill: float = float("-inf"),
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Reference: consider only [0, end) per row, take k largest, pad, then shift."""
     rows = scores.shape[0]
     idx_rows, val_rows = [], []
@@ -68,8 +75,14 @@ def main() -> None:
     view = storage[:, :k]
     offset = torch.arange(rows, dtype=torch.int32, device=dev) * 1000
     vals2, idxs2 = topk(
-        scores, k, end=end, indices_type=torch.int32, output_idx=view, output_idx_offset=offset,
-        idx_oob_fill_value=-1, return_value=False,
+        scores,
+        k,
+        end=end,
+        indices_type=torch.int32,
+        output_idx=view,
+        output_idx_offset=offset,
+        idx_oob_fill_value=-1,
+        return_value=False,
     )
     assert idxs2.data_ptr() == view.data_ptr()
     assert vals2 is None
diff --git a/scripts/ref_layer_dump.py b/scripts/ref_layer_dump.py
index c4d4cc9..76ba49a 100644
--- a/scripts/ref_layer_dump.py
+++ b/scripts/ref_layer_dump.py
@@ -20,10 +20,11 @@ import torch.distributed as dist
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
-from safetensors.torch import load_model  # noqa: E402
-from transformers import AutoTokenizer  # noqa: E402
+from collections.abc import Callable  # noqa: E402
 
 from model import ModelArgs, Transformer  # noqa: E402
+from safetensors.torch import load_model  # noqa: E402
+from transformers import AutoTokenizer  # noqa: E402
 
 
 def main() -> None:
@@ -59,27 +60,27 @@ def main() -> None:
 
     captured: dict[str, np.ndarray] = {}
 
-    def make_hook(layer_id: int):
-        def hook(_module, _inputs, output):
+    def make_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
+        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
             # Block.forward returns (x, pre_mix); x is [b, s, hc_mult, dim].
             x = output[0] if isinstance(output, tuple) else output
             captured[f"layer{layer_id}"] = x.mean(dim=2)[0, -1].float().cpu().numpy()
 
         return hook
 
-    def make_attn_hook(layer_id: int):
-        def hook(_module, _inputs, output):
+    def make_attn_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
+        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
             # The attention sub-layer output is already collapsed: [b, s, dim].
             x = output[0] if isinstance(output, tuple) else output
             captured[f"layer{layer_id}_attn"] = x[0, -1].float().cpu().numpy()
 
         return hook
 
-    def embed_hook(_module, _inputs, output):
+    def embed_hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
         captured["embed"] = output[0, -1].float().cpu().numpy()
 
-    def make_gate_hook(layer_id: int):
-        def hook(_module, _inputs, output):
+    def make_gate_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
+        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
             # Gate.forward returns (weights [n, topk], indices [n, topk]).
             weights, indices = output
             captured[f"layer{layer_id}_topk_idx"] = indices[-1].to(torch.int32).cpu().numpy()
@@ -87,14 +88,14 @@ def main() -> None:
 
         return hook
 
-    def make_shared_hook(layer_id: int):
-        def hook(_module, _inputs, output):
+    def make_shared_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
+        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
             captured[f"layer{layer_id}_shared_out"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()
 
         return hook
 
-    def make_moe_hook(layer_id: int):
-        def hook(_module, _inputs, output):
+    def make_moe_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
+        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
             # The MoE output already includes the shared expert, matching our `_combined`.
             captured[f"layer{layer_id}_combined"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()
 
diff --git a/scripts/ref_next_token.py b/scripts/ref_next_token.py
index 758b8e7..25ee053 100644
--- a/scripts/ref_next_token.py
+++ b/scripts/ref_next_token.py
@@ -19,9 +19,8 @@ import torch.distributed as dist
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
-from safetensors.torch import load_model  # noqa: E402
-
 from model import ModelArgs, Transformer  # noqa: E402
+from safetensors.torch import load_model  # noqa: E402
 
 
 def main() -> None: