From 45e16a2de0bb3a2db71fd010c1294b4a4d5c8526 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 05:06:18 -0700
Subject: [PATCH 001/117] [None][perf] GVR top-k: block-skip for the R0 M-ary
 count pass and Phase-3 collect

Port the block-skip consumer from the skip-finegrain development chain
onto the R0 (op#26) architecture, measured-optimal configuration only
(grain 32, int16 active list = 16KB SMEM, strided coalesced 3-barrier
build, UN=2 software-pipelined compact scan). The active-block list is
built once per row at the loosest rung (lossless for every rung count
and the collect); phase3's compact stream-write replays the count
pass's per-thread walk order so prefix-sum positions stay exact; a
list-current flag pairs the two and is cleared on any dense fallback.
Misaligned slice starts fall through to the dense path. Contract:
workspace/epilogue_topk_interface.md.

Correctness: REAL-data smoke (flash 16k/64k/256k/1024k + pro 64k/1024k,
dense + skip arms, exactness contract) all pass incl. hit_rate=0.08.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 381 ++++++++++++++++--
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 178 ++++++++
 2 files changed, 535 insertions(+), 24 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 1fe0dfa5eddc..68a264c14323 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -270,6 +270,7 @@ def __init__(
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
+        enable_block_skip: bool = False,
     ):
         # Redundant-warp sync reduction: every warp replays the block
         # reduce + decision from the same staged SMEM partials in the
@@ -423,6 +424,25 @@ def __init__(
         self.enable_r0 = bool(enable_r0)
         self.mt_unroll = int(mt_unroll)
         self.fb_fix = bool(fb_fix)
+        # enable_block_skip: gate the R0 M-ary count pass and the Phase-3
+        # stream-write on per-32-position upper bounds emitted by the
+        # indexer epilogue (``block_max [num_rows, nb_pad*4]`` fp32, record
+        # r = exact max of positions [r*32, r*32+32) of the POST-CONVERSION
+        # stored logits — contract in workspace/epilogue_topk_interface.md).
+        # Lossless: the active-block list is built at the LOOSEST rung, so a
+        # skipped block cannot contain any element >= any rung and every
+        # rung count (and the collect) equals its dense value. The port
+        # ships the measured-optimal configuration only: grain 32, int16
+        # list entries (16KB SMEM, 3 CTA/SM at T512), strided coalesced
+        # O(1)-barrier build, UN=2 software-pipelined compact scan.
+        self.enable_block_skip = bool(enable_block_skip)
+        self.SKIP_BLOCK = 32
+        self.SKIP_BLOCK_LOG2 = 5
+        self.SKIP_MAX_BLOCKS = 8192  # covers N up to 262144 at grain 32
+        self.SKIP_UNROLL = 2
+        self.skip_order = "grouped"
+        if enable_block_skip and num_threads not in (512, 1024):
+            raise ValueError("enable_block_skip requires num_threads in {512, 1024}")
         # C7 dispatch (op#26 host policy folded into the ctor; all gated on
         # enable_r0 so an OFF kernel is byte-identical to the base):
         #  - qfracs default = M2D (0.85, 0.35): dispatch_r0_op26 ships M2D for
@@ -1377,6 +1397,105 @@ def block_count_ge(
     # totals are the answer. smem_ptcnt_multi holds slice-local per-thread
     # columns (the accepted rung's column seeds Phase 3 per CTA).
     # ------------------------------------------------------------------
+
+    # ---- block-skip machinery (enable_block_skip; ported from the
+    # skip-finegrain development chain, measured-optimal configuration
+    # only: grain 32, int16 list, grouped strided build, UN=2 scan) ----
+
+    @cute.jit
+    def _list_ld(self, smem_active, idx):
+        return cutlass.Int32(smem_active[idx])
+
+    @cute.jit
+    def _list_st(self, smem_active, idx, val):
+        smem_active[idx] = cutlass.Int16(val)
+
+    @cute.jit
+    def _block_bound(self, bm_addr, blk_id):
+        # grain 32: record blk_id IS the exact positional bound of
+        # [blk_id*32, blk_id*32+32) (indexer TMEM partition contract) —
+        # one 4B scalar load, no fold.
+        bm_ptr = cute.make_ptr(
+            cutlass.Float32,
+            bm_addr + cutlass.Int64(blk_id) * cutlass.Int64(4),
+            cute.AddressSpace.gmem,
+            assumed_align=4,
+        )
+        return cute.make_tensor(bm_ptr, cute.make_layout((1,)))[0]
+
+    @cute.jit
+    def _full_build_active(
+        self,
+        block_max_row,
+        slice_start,
+        slice_end,
+        threshold,
+        smem_wcnt,
+        smem_active,
+        s_active_cnt,
+        tidx,
+        warp_id,
+        lane,
+    ):
+        # Two-phase register-bitmask build. Each thread owns ids_per_thread
+        # block ids STRIDED (t, t+T, t+2T, ...) so each pass's bound loads
+        # coalesce warp-wide. Phase A flags+counts with no sync; phase B is
+        # ONE block-wide exclusive scan over per-thread counts; phase C
+        # scatters from the bitmask — 3 barriers TOTAL, independent of
+        # nb_slice. List order is thread-grouped, not ascending: fine, since
+        # the count scan and the Phase-3 stream-write walk the SAME list by
+        # position (determinism contract).
+        ids_per_thread = cutlass.const_expr(self.SKIP_MAX_BLOCKS // self.num_threads)
+        num_threads = cutlass.const_expr(self.num_threads)
+        blk_lo = slice_start >> cutlass.Int32(self.SKIP_BLOCK_LOG2)
+        blk_hi = (slice_end + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32(
+            self.SKIP_BLOCK_LOG2
+        )
+        nb_slice = blk_hi - blk_lo
+        bm_addr = block_max_row.iterator.toint()
+
+        # Phase A (no sync): flag my ids into a register bitmask.
+        bmask = cutlass.Int32(0)
+        cnt = cutlass.Int32(0)
+        for m in cutlass.range_constexpr(ids_per_thread):
+            ib = tidx + cutlass.Int32(m * num_threads)
+            if ib < nb_slice:
+                bound = self._block_bound(bm_addr, blk_lo + ib)
+                if bound >= threshold:
+                    bmask = bmask | (cutlass.Int32(1) << cutlass.Int32(m))
+                    cnt = cnt + cutlass.Int32(1)
+
+        # Phase B: one block-wide exclusive scan over per-thread counts.
+        tp = cnt
+        for off_i in cutlass.range_constexpr(5):
+            off_v = cutlass.const_expr(1 << off_i)
+            other = cute.arch.shuffle_sync_up(tp, off_v, mask_and_clamp=0)
+            if lane >= cutlass.Int32(off_v):
+                tp = tp + other
+        excl = tp - cnt
+        warp_total = cute.arch.shuffle_sync(tp, cutlass.Int32(self.WARP_SIZE - 1))
+        if lane == 0:
+            smem_wcnt[warp_id] = warp_total
+        cute.arch.barrier()
+        if tidx == 0:
+            tot = cutlass.Int32(0)
+            for w in cutlass.range_constexpr(self.num_warps):
+                cw = smem_wcnt[w]
+                smem_wcnt[w] = tot
+                tot = tot + cw
+            s_active_cnt[0] = tot
+        cute.arch.barrier()
+
+        # Phase C: scatter from the bitmask at deterministic offsets.
+        pos_out = smem_wcnt[warp_id] + excl
+        if bmask != cutlass.Int32(0):
+            for m in cutlass.range_constexpr(ids_per_thread):
+                if (bmask & (cutlass.Int32(1) << cutlass.Int32(m))) != cutlass.Int32(0):
+                    ib_c = tidx + cutlass.Int32(m * num_threads)
+                    self._list_st(smem_active, pos_out, blk_lo + ib_c)
+                    pos_out = pos_out + cutlass.Int32(1)
+        cute.arch.barrier()
+
     @cute.jit
     def block_count_ge_multi(
         self,
@@ -1393,6 +1512,9 @@ def block_count_ge_multi(
         warp_id,
         lane,
         smem_ptcnt=None,  # vseed: last column's per-thread counts land here
+        block_max_row=None,  # block-skip: per-32-position upper bounds
+        smem_active=None,  # block-skip: int16 active list
+        s_active_cnt=None,  # block-skip: [0]=list length, [1]=list-current flag
     ):
         M = cutlass.const_expr(self.M_thr)
         num_threads = cutlass.const_expr(self.num_threads)
@@ -1416,7 +1538,97 @@ def block_count_ge_multi(
         i = slice_start + tidx * cutlass.Int32(vec_w)
         step = cutlass.Int32(step_elem)
 
-        if self.enable_unroll_4:
+        # ---- block-skip compact iteration (lossless vs the dense path) ----
+        # Build the active list at the LOOSEST threshold over all M columns:
+        # a skipped block bounds every element below min(t_m), so all M
+        # counts equal their dense values. Slice boundaries must be
+        # 32-aligned for the list's block<->position mapping to stay inside
+        # the slice (runtime-uniform guard; misaligned slices fall through
+        # to the dense path below via skip_ok == 0).
+        skip_ok = cutlass.Int32(0)
+        if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
+            if (slice_start & cutlass.Int32(self.SKIP_BLOCK - 1)) == cutlass.Int32(0):
+                skip_ok = cutlass.Int32(1)
+        if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
+            if skip_ok == cutlass.Int32(1):
+                tmin = thr_frag[0]
+                for m in cutlass.range_constexpr(M):
+                    tmin = _fmin_f32_inline(tmin, thr_frag[m])
+                # smem_wcnt_multi's first num_warps slots double as the
+                # build's warp-total scratch: the build's own barriers
+                # complete before any count lands there.
+                self._full_build_active(
+                    block_max_row,
+                    slice_start,
+                    slice_end,
+                    tmin,
+                    smem_wcnt_multi,
+                    smem_active,
+                    s_active_cnt,
+                    tidx,
+                    warp_id,
+                    lane,
+                )
+                if tidx == 0:
+                    s_active_cnt[1] = cutlass.Int32(1)  # list-current flag
+                chunks_per_block = cutlass.const_expr(
+                    self.SKIP_BLOCK // (self.vec_bits // self.dtype.width)
+                )
+                tpb = cutlass.const_expr(chunks_per_block)
+                blocks_per_iter = cutlass.const_expr(self.num_threads // chunks_per_block)
+                UN = cutlass.const_expr(self.SKIP_UNROLL)
+                stride_un = cutlass.const_expr(blocks_per_iter * UN)
+                my_blk_slot = tidx // cutlass.Int32(tpb)
+                my_chunk0 = tidx % cutlass.Int32(tpb)
+                cnt_active = s_active_cnt[0]
+                frags = [cute.make_fragment((vec_w,), self.dtype) for _ in range(UN)]
+                li = my_blk_slot
+                while li < cnt_active:
+                    poss = []
+                    valids = []
+                    for u in cutlass.range_constexpr(UN):
+                        lu = li + cutlass.Int32(u * blocks_per_iter)
+                        valid = lu < cnt_active
+                        pos0 = cutlass.Int32(0)
+                        if valid:
+                            blk = self._list_ld(smem_active, lu)
+                            pos0 = blk * cutlass.Int32(self.SKIP_BLOCK) + my_chunk0 * cutlass.Int32(
+                                vec_w
+                            )
+                            src_ptr_u = cute.make_ptr(
+                                self.dtype,
+                                row_addr + cutlass.Int64(pos0) * cutlass.Int64(elem_bytes),
+                                cute.AddressSpace.gmem,
+                                assumed_align=vec_align,
+                            )
+                            cute.copy(
+                                copy_atom,
+                                cute.make_tensor(src_ptr_u, cute.make_layout((vec_w,))),
+                                frags[u],
+                            )
+                        poss.append(pos0)
+                        valids.append(valid)
+                    for u in cutlass.range_constexpr(UN):
+                        if valids[u]:
+                            pos = poss[u]
+                            if pos + cutlass.Int32(vec_w) <= slice_end:
+                                for j in cutlass.range_constexpr(vec_w):
+                                    if cutlass.const_expr(self.dtype == cutlass.Float32):
+                                        vj = frags[u][j]
+                                    else:
+                                        vj = cutlass.Float32(frags[u][j])
+                                    for m in cutlass.range_constexpr(M):
+                                        cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m])
+                            else:
+                                jj = pos
+                                while jj < slice_end:
+                                    vs = self._load_fp32(input_row, jj)
+                                    for m in cutlass.range_constexpr(M):
+                                        cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vs >= thr_frag[m])
+                                    jj = jj + cutlass.Int32(1)
+                    li = li + cutlass.Int32(stride_un)
+
+        if self.enable_unroll_4 and skip_ok == cutlass.Int32(0):
             rng_frag = cute.make_fragment((vec_w,), self.dtype)
             big_iters = cutlass.Int32(0)
             if slice_end > i + cutlass.Int32(vec_w - 1):
@@ -1443,30 +1655,31 @@ def block_count_ge_multi(
             i = i + big_iters * cutlass.Int32(step_elem)
 
         tail_frag = cute.make_fragment((vec_w,), self.dtype)
-        while i + cutlass.Int32(vec_w - 1) < slice_end:
-            src_ptr = cute.make_ptr(
-                self.dtype,
-                row_addr + cutlass.Int64(i) * cutlass.Int64(elem_bytes),
-                cute.AddressSpace.gmem,
-                assumed_align=vec_align,
-            )
-            src = cute.make_tensor(src_ptr, cute.make_layout((vec_w,)))
-            cute.copy(copy_atom, src, tail_frag)
-            for j in cutlass.range_constexpr(vec_w):
-                if cutlass.const_expr(self.dtype == cutlass.Float32):
-                    vj = tail_frag[j]
-                else:
-                    vj = cutlass.Float32(tail_frag[j])
-                for m in cutlass.range_constexpr(M):
-                    cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m])
-            i = i + step
+        if skip_ok == cutlass.Int32(0):
+            while i + cutlass.Int32(vec_w - 1) < slice_end:
+                src_ptr = cute.make_ptr(
+                    self.dtype,
+                    row_addr + cutlass.Int64(i) * cutlass.Int64(elem_bytes),
+                    cute.AddressSpace.gmem,
+                    assumed_align=vec_align,
+                )
+                src = cute.make_tensor(src_ptr, cute.make_layout((vec_w,)))
+                cute.copy(copy_atom, src, tail_frag)
+                for j in cutlass.range_constexpr(vec_w):
+                    if cutlass.const_expr(self.dtype == cutlass.Float32):
+                        vj = tail_frag[j]
+                    else:
+                        vj = cutlass.Float32(tail_frag[j])
+                    for m in cutlass.range_constexpr(M):
+                        cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m])
+                i = i + step
 
-        it = n_aligned + tidx
-        while it < slice_end:
-            v = self._load_fp32(input_row, it)
-            for m in cutlass.range_constexpr(M):
-                cnt_frag[m] = cnt_frag[m] + cutlass.Int32(v >= thr_frag[m])
-            it = it + cutlass.Int32(num_threads)
+            it = n_aligned + tidx
+            while it < slice_end:
+                v = self._load_fp32(input_row, it)
+                for m in cutlass.range_constexpr(M):
+                    cnt_frag[m] = cnt_frag[m] + cutlass.Int32(v >= thr_frag[m])
+                it = it + cutlass.Int32(num_threads)
 
         for m in cutlass.range_constexpr(M):
             if cutlass.const_expr(self.r0_vseed and m == self.M_qf):
@@ -1813,6 +2026,8 @@ def phase3_collect_candidates(
         lane,
         do_cluster_sync,  # bool: False = cs=1 / short-row degrade (skip cluster sync)
         smem_input=None,  # optional SMEM-cached slice
+        smem_active=None,  # block-skip: int16 active list (reused, not rebuilt)
+        s_active_cnt=None,  # block-skip: [0]=list length, [1]=list-current flag
     ):
         """Retry-shrink (when P2 didn't converge) + prefix sum + stream-write.
 
@@ -1977,6 +2192,72 @@ def phase3_collect_candidates(
         wc = my_write_pos
         step = cutlass.Int32(step_elem)
 
+        # ---- block-skip compact stream-write ----
+        # Reuses the active list left by the R0 compact count pass (same
+        # ownership walk: my_blk_slot's list slots in ascending order, so
+        # each thread produces its candidates in the SAME per-thread order
+        # the count pass counted them — the prefix-sum positions match).
+        # Only taken when the list is CURRENT (s_active_cnt[1] == 1, set by
+        # the build; cleared on any dense fallback re-count).
+        skip_wr = cutlass.Int32(0)
+        if cutlass.const_expr(self.enable_block_skip and smem_active is not None):
+            if s_active_cnt[1] == cutlass.Int32(1):
+                skip_wr = cutlass.Int32(1)
+        if cutlass.const_expr(self.enable_block_skip and smem_active is not None):
+            if skip_wr == cutlass.Int32(1):
+                chunks_per_block_w = cutlass.const_expr(
+                    self.SKIP_BLOCK // (self.vec_bits // self.dtype.width)
+                )
+                blocks_per_iter_w = cutlass.const_expr(self.num_threads // chunks_per_block_w)
+                my_blk_slot_w = tidx // cutlass.Int32(chunks_per_block_w)
+                my_chunk0_w = tidx % cutlass.Int32(chunks_per_block_w)
+                cnt_active_w = s_active_cnt[0]
+                wfrag = cute.make_fragment((vec_w,), self.dtype)
+                li_w = my_blk_slot_w
+                while li_w < cnt_active_w:
+                    blk_w = self._list_ld(smem_active, li_w)
+                    pos0_w = blk_w * cutlass.Int32(self.SKIP_BLOCK) + my_chunk0_w * cutlass.Int32(
+                        vec_w
+                    )
+                    src_ptr_w = cute.make_ptr(
+                        self.dtype,
+                        row_addr + cutlass.Int64(pos0_w) * cutlass.Int64(elem_bytes),
+                        cute.AddressSpace.gmem,
+                        assumed_align=vec_align,
+                    )
+                    cute.copy(
+                        copy_atom,
+                        cute.make_tensor(src_ptr_w, cute.make_layout((vec_w,))),
+                        wfrag,
+                    )
+                    if pos0_w + cutlass.Int32(vec_w) <= slice_end:
+                        for j in cutlass.range_constexpr(vec_w):
+                            if cutlass.const_expr(self.dtype == cutlass.Float32):
+                                vj = wfrag[j]
+                            else:
+                                vj = cutlass.Float32(wfrag[j])
+                            if vj >= thr_final and wc < cutlass.Int32(kCC):
+                                smem_keys[wc] = vj
+                                smem_vals[wc] = pos0_w + cutlass.Int32(j)
+                                wc = wc + cutlass.Int32(1)
+                    else:
+                        jj_w = pos0_w
+                        while jj_w < slice_end:
+                            v_w = self._load_fp32(input_row, jj_w)
+                            if v_w >= thr_final and wc < cutlass.Int32(kCC):
+                                smem_keys[wc] = v_w
+                                smem_vals[wc] = jj_w
+                                wc = wc + cutlass.Int32(1)
+                            jj_w = jj_w + cutlass.Int32(1)
+                    li_w = li_w + cutlass.Int32(blocks_per_iter_w)
+
+        # When the compact write ran, park the dense cursors at the end so
+        # all three dense loops below (4-way, vec tail, scalar tail) fall
+        # through without re-indenting them.
+        if skip_wr == cutlass.Int32(1):
+            ic = N_local
+            n_aligned = N_local
+
         # Phase3 unrolling: master gated by self.enable_phase3_unroll.
         # When OFF, only the tail 1-way loop runs (matches the pre-unroll
         # state of phase3_collect). When ON, the inner enable_unroll_4
@@ -3717,6 +3998,7 @@ def gvr_topk_kernel(
         output_values: cute.Tensor,  # [numRows, top_k] dtype
         output_indices: cute.Tensor,  # [numRows, top_k] int32
         order_row: cute.Tensor,  # [batch_size] int32 (or None when seqlen_sorted=False)
+        block_max: cute.Tensor,  # [numRows, nb_pad*4] fp32 (or None: no block-skip)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
 
@@ -3768,6 +4050,7 @@ def gvr_topk_kernel(
             seq_lens,
             output_values,
             output_indices,
+            block_max=block_max,
         )
 
     @cute.jit
@@ -3779,6 +4062,7 @@ def run_one_row(
         seq_lens: cute.Tensor,  # [numRows / next_n] int32
         output_values: cute.Tensor,  # [numRows, top_k] dtype, optional
         output_indices: cute.Tensor,  # [numRows, top_k] int32
+        block_max: cute.Tensor = None,  # [numRows, nb_pad*4] fp32
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
 
@@ -3835,6 +4119,10 @@ def run_one_row(
         # Slice per-row views.
         input_row = input_data[row_idx, None]
         pre_idx_row = pre_idx[pre_idx_row_idx, None]
+        if cutlass.const_expr(self.enable_block_skip and block_max is not None):
+            block_max_row = block_max[row_idx, None]
+        else:
+            block_max_row = None
         # When return_output_values=False, ``output_values`` is None at
         # launch and the gated writes below are compiled out; slicing into
         # None would crash so we keep the view None as well.
@@ -3874,6 +4162,22 @@ def run_one_row(
             layout=cute.make_ordered_layout((num_threads,), order=(0,)),
             byte_alignment=128,
         )
+        # block-skip: int16 active list (16KB at 8192 entries) + control
+        # ([0] list length, [1] list-current flag for the Phase-3 reuse).
+        if cutlass.const_expr(self.enable_block_skip):
+            smem_active = smem.allocate_tensor(
+                element_type=cutlass.Int16,
+                layout=cute.make_ordered_layout((self.SKIP_MAX_BLOCKS,), order=(0,)),
+                byte_alignment=128,
+            )
+            s_active_cnt = smem.allocate_tensor(
+                element_type=cutlass.Int32,
+                layout=cute.make_ordered_layout((2,), order=(0,)),
+                byte_alignment=16,
+            )
+        else:
+            smem_active = None
+            s_active_cnt = None
         # warp_counts[NUM_WARPS] int32 (P3 prefix-sum scratch)
         # p2_warp_redundant parity-banks the Phase-2 staging (a warp one
         # round ahead writes the other half) — costs num_warps*4 bytes.
@@ -4114,6 +4418,9 @@ def run_one_row(
                         tidx,
                         warp_id,
                         lane,
+                        block_max_row=block_max_row,
+                        smem_active=smem_active,
+                        s_active_cnt=s_active_cnt,
                     )
                 else:
                     # Short row: only CTA 0 scans the full row; the other
@@ -4155,6 +4462,9 @@ def run_one_row(
                             tidx,
                             warp_id,
                             lane,
+                            block_max_row=block_max_row,
+                            smem_active=smem_active,
+                            s_active_cnt=s_active_cnt,
                         )
             else:
                 # cs=1: one CTA per row, no cluster sync.
@@ -4193,6 +4503,9 @@ def run_one_row(
                     tidx,
                     warp_id,
                     lane,
+                    block_max_row=block_max_row,
+                    smem_active=smem_active,
+                    s_active_cnt=s_active_cnt,
                 )
 
         griddepcontrol_launch_dependents()
@@ -4234,6 +4547,9 @@ def _run_phases(
         tidx,
         warp_id,
         lane,
+        block_max_row=None,  # block-skip: this row's per-32-position bounds
+        smem_active=None,
+        s_active_cnt=None,
     ):
         """Run Phase 1-4 + final cluster barrier on a given row slice.
 
@@ -4248,6 +4564,13 @@ def _run_phases(
         cluster_size = cutlass.const_expr(self.cluster_size)
         is_leader = cta_in_cluster == cutlass.Int32(0)
 
+        # block-skip: the list-current flag starts INVALID every row; only
+        # the R0 compact pass's build sets it. Ordered ahead of all readers
+        # by Phase 1's internal barriers.
+        if cutlass.const_expr(self.enable_block_skip):
+            if tidx == cutlass.Int32(0):
+                s_active_cnt[1] = cutlass.Int32(0)
+
         # ---- Phase 1: preIdx Min/Max/Mean ----
         self.phase1_preidx_stats(
             input_row,
@@ -4357,6 +4680,9 @@ def _run_phases(
                     warp_id,
                     lane,
                     smem_ptcnt=smem_ptcnt,
+                    block_max_row=block_max_row,
+                    smem_active=smem_active,
+                    s_active_cnt=s_active_cnt,
                 )
                 cute.arch.barrier()
                 if tidx == 0:
@@ -4409,6 +4735,9 @@ def _run_phases(
                 # count passes (op#26 efficiency) instead of ~6. done=1 on
                 # accept so Phase 3 skips its retry-shrink.
                 if bc < cutlass.Int32(0):
+                    if cutlass.const_expr(self.enable_block_skip):
+                        if tidx == cutlass.Int32(0):
+                            s_active_cnt[1] = cutlass.Int32(0)
                     if cutlass.const_expr(self.fb_fix):
                         if tidx == cutlass.Int32(0):
                             M = cutlass.const_expr(self.M_thr)
@@ -4604,6 +4933,8 @@ def _run_phases(
                 lane,
                 do_cluster_sync=do_cluster_sync,
                 smem_input=smem_input,
+                smem_active=smem_active,
+                s_active_cnt=s_active_cnt,
             )
 
             # Cluster handoff #2: leader's DSMEM gather of peer
@@ -4752,6 +5083,7 @@ def __call__(
         output_values: cute.Tensor,  # or None.
         output_indices: cute.Tensor,
         order_row: cute.Tensor,  # or None when seqlen_sorted=False
+        block_max: cute.Tensor,  # or None: block-skip disabled
         stream,
     ):
         num_rows = input_data.shape[0]
@@ -4774,6 +5106,7 @@ def __call__(
             output_values,
             output_indices,
             order_row,
+            block_max,
         ).launch(
             grid=(total_ctas, 1, 1),
             block=(self.num_threads, 1, 1),
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 41cb23260cd4..bd028c07a7ee 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -64,6 +64,7 @@ def _compile(
     seqlen_sorted: bool = False,
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
+    enable_block_skip: bool = False,
 ):
     """JIT-compile the GVR kernel for a specific knob combination.
 
@@ -122,6 +123,16 @@ def _compile(
         if seqlen_sorted
         else None
     )
+    block_max_fake = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Float32,
+            (n_rows, cute.sym_int()),
+            stride_order=(1, 0),
+            assumed_align=16,
+        )
+        if enable_block_skip
+        else None
+    )
     fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
     kernel = GvrTopKKernel(
         dtype=cute_dtype,
@@ -140,6 +151,7 @@ def _compile(
         seqlen_sorted=seqlen_sorted,
         p4_warp_redundant=p4_warp_redundant,
         p2_warp_redundant=p2_warp_redundant,
+        enable_block_skip=enable_block_skip,
     )
     return cute.compile(
         kernel,
@@ -149,11 +161,163 @@ def _compile(
         out_values_fake,
         out_indices_fake,
         order_row_fake,
+        block_max_fake,
         stream=fake_stream,
         options="--enable-tvm-ffi",
     )
 
 
+_FLT_MAX = torch.finfo(torch.float32).max
+_META_BLOCK = 128
+_META_RECS_PER_BLOCK = 4  # warp-partial records per block (indexer layout)
+
+
+def _row_n_eff(
+    seq_lens: torch.Tensor,
+    num_rows: int,
+    next_n: int,
+    compress_ratio: int,
+) -> torch.Tensor:
+    """Per-row effective scan length, mirroring the kernel formula."""
+    dev = seq_lens.device
+    rows = torch.arange(num_rows, device=dev)
+    sl = seq_lens[rows // next_n].to(torch.int64)
+    actual = sl - next_n + (rows % next_n) + 1
+    return actual // compress_ratio
+
+
+def emu_block_max(
+    logits: torch.Tensor,
+    seq_lens: torch.Tensor,
+    next_n: int = 1,
+    compress_ratio: int = 1,
+    tail_mode: str = "pad_inf",
+    records: str = "rotate",
+) -> torch.Tensor:
+    """``[num_rows, nb_pad*4] fp32`` warp-partial upper-bound records.
+
+    tail_mode:
+      "exact":   tight bound — max over valid positions only.
+      "pad_inf": a partially-valid tail unit is forced to +FLT_MAX, the
+                 worst legal inflation (the indexer masks by request-level
+                 ctx >= N_eff, so tail positions can inflate the bound).
+                 Tests use this to prove inflated bounds only *disable*
+                 skipping, never break correctness.
+    records:
+      "rotate":     fold-correctness fixture — the 128-block max lands in
+                    ONE slot rotated by blk % 4, the other 3 hold
+                    -FLT_MAX. Any 4 values folding to the block max are
+                    legal; the rotation makes a consumer that drops or
+                    mis-reads partial slots fail the tests. Valid ONLY
+                    for grain-128 consumers (slots are NOT positional).
+      "positional": production semantics — record r is the exact max of
+                    positions [r*32, r*32+32) (the indexer's TMEM T2R
+                    partition gives warp w of a tile the contiguous
+                    positions [tile*128 + w*32, +32)). Required for
+                    skip_grain=32; also a legal grain-128 input (its
+                    fold is the block max).
+    """
+    assert tail_mode in ("exact", "pad_inf")
+    assert records in ("rotate", "positional")
+    R, C = logits.shape
+    nb = (C + _META_BLOCK - 1) // _META_BLOCK
+    dev = logits.device
+    n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio).unsqueeze(1)
+    lf = logits.to(torch.float32)
+    pad = nb * _META_BLOCK - C
+    if pad:
+        lf = torch.nn.functional.pad(lf, (0, pad), value=float("-inf"))
+    pos = torch.arange(nb * _META_BLOCK, device=dev).unsqueeze(0)
+    masked = torch.where(pos < n_eff, lf, torch.full_like(lf, float("-inf")))
+    if records == "positional":
+        sub = _META_BLOCK // _META_RECS_PER_BLOCK  # 32 positions/record
+        nrec = nb * _META_RECS_PER_BLOCK
+        rmax = masked.view(R, nrec, sub).amax(-1)
+        if tail_mode == "pad_inf":
+            rec_start = torch.arange(nrec, device=dev).unsqueeze(0) * sub
+            partial = (rec_start < n_eff) & (rec_start + sub > n_eff)
+            rmax = torch.where(partial, torch.full_like(rmax, _FLT_MAX), rmax)
+        return rmax.contiguous()
+    bmax = masked.view(R, nb, _META_BLOCK).amax(-1)
+    if tail_mode == "pad_inf":
+        blk_start = torch.arange(nb, device=dev).unsqueeze(0) * _META_BLOCK
+        partial = (blk_start < n_eff) & (blk_start + _META_BLOCK > n_eff)
+        bmax = torch.where(partial, torch.full_like(bmax, _FLT_MAX), bmax)
+    out = torch.full((R, nb, _META_RECS_PER_BLOCK), -_FLT_MAX, dtype=torch.float32, device=dev)
+    slot = torch.arange(nb, device=dev) % _META_RECS_PER_BLOCK
+    out[:, torch.arange(nb, device=dev), slot] = bmax
+    return out.reshape(R, nb * _META_RECS_PER_BLOCK).contiguous()
+
+
+# Rung offsets below each 32-position record's max for the meta-seed
+# metadata (L6_geo.25 — offline-validated: 1.00 real scans on synth +
+# real captures for both V4 models with the kernel's S=16 x 2-pass
+# quantized search).
+_META_DELTAS = (0.25, 0.5, 1.0, 2.0, 4.0, 8.0)
+
+
+def emu_block_meta(
+    logits: torch.Tensor,
+    seq_lens: torch.Tensor,
+    compress_ratio: int = 1,
+    next_n: int = 1,
+) -> torch.Tensor:
+    """Emulate the indexer-side per-32-block rung-count metadata.
+
+    Record r packs, for each rung offset ``delta_j`` in
+    ``_META_DELTAS``, ``count(v >= record_max - delta_j)`` over
+    positions ``[32r, 32r+32)`` as a 5-bit saturating field at bits
+    ``[5j, 5j+5)`` (31 = "31 or 32"; the kernel decodes 31 as the safe
+    upper bound 32). Positions beyond the row's effective length are
+    excluded. Layout matches ``block_max``: ``[num_rows, nrec]`` int32
+    with ``nrec = ceil(N/128)*4`` (one record per 32 positions).
+    On the indexer side this is ~L ballots+popcounts per record in the
+    epilogue, same shape as the +1.3% block_max emission.
+    """
+    assert next_n == 1, "emu_block_meta: next_n == 1 only"
+    x = logits.float()
+    R, N = x.shape
+    nrec = ((N + _META_BLOCK - 1) // _META_BLOCK) * _META_RECS_PER_BLOCK
+    npad = nrec * 32
+    if npad > N:
+        x = torch.nn.functional.pad(x, (0, npad - N), value=float("-inf"))
+    n_eff = seq_lens.long() // compress_ratio
+    ar = torch.arange(npad, device=x.device)[None, :]
+    x = x.masked_fill(ar >= n_eff[:, None], float("-inf"))
+    xb = x.view(R, nrec, 32)
+    m = xb.amax(2, keepdim=True)
+    meta = torch.zeros(R, nrec, dtype=torch.int32, device=x.device)
+    for j, d in enumerate(_META_DELTAS):
+        cj = (xb >= (m - d)).sum(2).clamp(max=31).to(torch.int32)
+        meta |= cj << (5 * j)
+    return meta.contiguous()
+
+
+def enc_ordered_f32(t: torch.Tensor) -> torch.Tensor:
+    """Order-preserving int encoding of fp32 (an involution).
+
+    Stored back in fp32 slots — matches the indexer's encoded-int atomic
+    min/max.
+    """
+    bits = t.float().contiguous().view(torch.int32)
+    enc = torch.where(bits >= 0, bits, bits ^ 0x7FFFFFFF)
+    return enc.view(torch.float32)
+
+
+def hit_agg_identities(num_rows: int, device) -> torch.Tensor:
+    """Identity-initialized per-row hit aggregate.
+
+    {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} — the required initial state of
+    the buffer the indexer atomically merges into.
+    """
+    ident = torch.tensor([_FLT_MAX, -_FLT_MAX], dtype=torch.float32, device=device)
+    enc = enc_ordered_f32(ident)
+    out = torch.zeros((num_rows, 4), dtype=torch.float32, device=device)
+    out[:, 0] = enc[0]
+    out[:, 1] = enc[1]
+    return out.contiguous()
+
+
 def gvr_topk_decode(
     logits: torch.Tensor,
     pre_idx: torch.Tensor,
@@ -178,6 +342,7 @@ def gvr_topk_decode(
     order_row: Optional[torch.Tensor] = None,
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
+    block_max: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """CuTe DSL GVR Top-K wrapper with every tuning knob exposed.
 
@@ -247,6 +412,17 @@ def gvr_topk_decode(
     cute_dtype = _DTYPE_TORCH_TO_CUTE[logits.dtype]
 
     num_rows = logits.shape[0]
+    enable_block_skip = block_max is not None
+    if enable_block_skip:
+        assert (
+            block_max.dtype == torch.float32
+            and block_max.is_cuda
+            and block_max.is_contiguous()
+            and block_max.dim() == 2
+            and block_max.shape[0] == num_rows
+            and block_max.shape[1] % 4 == 0
+        ), "block_max must be contiguous CUDA fp32 [num_rows, nb_pad*4]"
+
     if return_output_values:
         if out_values is None:
             out_values = torch.empty((num_rows, top_k), dtype=logits.dtype, device=logits.device)
@@ -312,6 +488,7 @@ def gvr_topk_decode(
         seqlen_sorted,
         p4_warp_redundant,
         p2_warp_redundant,
+        enable_block_skip,
     )
     # When return_output_values=False the kernel was compiled to skip
     # STG.value and accepts None for the value-output slot.
@@ -324,6 +501,7 @@ def gvr_topk_decode(
         out_values if return_output_values else None,
         out_indices,
         order_row if seqlen_sorted else None,
+        block_max if enable_block_skip else None,
     )
     if return_output_values:
         return out_values, out_indices

From 7eb1d888d8fb49717ecbf3c53a890e72247247f6 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 06:07:37 -0700
Subject: [PATCH 002/117] [None][fix] GVR top-k: make block_max a trailing
 defaulted launch param

The block-skip bounds tensor initially landed as a required positional
in __call__, breaking every pre-existing compile path that does not
pass it (16457's equivalence tests: 'Missing required argument'). Move
it after stream with a None default so legacy callers are untouched;
the wrapper passes it positionally last (the TVM-FFI env-stream launch
takes no runtime stream arg).

Verified standalone: main equivalence family 768 passed / 0 failed,
launch_autoconfig 4/4, real-data smoke (flash+pro up to 1M) all green.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py  | 2 +-
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py            | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 68a264c14323..2393a644fb4b 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -5083,8 +5083,8 @@ def __call__(
         output_values: cute.Tensor,  # or None.
         output_indices: cute.Tensor,
         order_row: cute.Tensor,  # or None when seqlen_sorted=False
-        block_max: cute.Tensor,  # or None: block-skip disabled
         stream,
+        block_max: cute.Tensor = None,  # block-skip bounds; None = disabled
     ):
         num_rows = input_data.shape[0]
         cluster_size = cutlass.const_expr(self.cluster_size)
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index bd028c07a7ee..f2966563d663 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -161,8 +161,8 @@ def _compile(
         out_values_fake,
         out_indices_fake,
         order_row_fake,
-        block_max_fake,
         stream=fake_stream,
+        block_max=block_max_fake,
         options="--enable-tvm-ffi",
     )
 

From dc42f23914bde76285dbc4f4f489f5829d8f70db Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 06:21:43 -0700
Subject: [PATCH 003/117] [None][fix] GVR block-skip: support unaligned cluster
 slice starts
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The active list previously required 32-aligned slice starts (runtime
guard falling back to dense), which silently disabled the skip on every
cluster slicing whose N/cs is not a multiple of 32 — including the
launch policy's cs=8 picks at the 512k/1024k rungs. The list now covers
FULL blocks only (first-full-block ceil in the build); the sub-block
head region of an unaligned slice is counted by all threads in a
strided scalar pass ordered BEFORE the list walk, and Phase 3's compact
write replays the same head-then-list per-thread order, so prefix-sum
positions stay exact. Boundary blocks shared with a neighbouring CTA
appear only in that CTA's head region — no double count, no gap.

Verified: skip arms exact at cs in {1,2,4,8} on flash+pro real cells
(64k/512k/1024k, unaligned N=262127 and 131075 slices).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 45 ++++++++++++++++---
 1 file changed, 38 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 2393a644fb4b..88638c1d377d 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -1447,7 +1447,12 @@ def _full_build_active(
         # position (determinism contract).
         ids_per_thread = cutlass.const_expr(self.SKIP_MAX_BLOCKS // self.num_threads)
         num_threads = cutlass.const_expr(self.num_threads)
-        blk_lo = slice_start >> cutlass.Int32(self.SKIP_BLOCK_LOG2)
+        # first FULL block: a block straddling slice_start would map list
+        # positions outside this CTA's slice; the sub-block head region
+        # [slice_start, blk_lo*32) is scanned separately by the callers.
+        blk_lo = (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32(
+            self.SKIP_BLOCK_LOG2
+        )
         blk_hi = (slice_end + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32(
             self.SKIP_BLOCK_LOG2
         )
@@ -1541,16 +1546,27 @@ def block_count_ge_multi(
         # ---- block-skip compact iteration (lossless vs the dense path) ----
         # Build the active list at the LOOSEST threshold over all M columns:
         # a skipped block bounds every element below min(t_m), so all M
-        # counts equal their dense values. Slice boundaries must be
-        # 32-aligned for the list's block<->position mapping to stay inside
-        # the slice (runtime-uniform guard; misaligned slices fall through
-        # to the dense path below via skip_ok == 0).
+        # counts equal their dense values. The list covers FULL blocks only;
+        # the sub-block head region of an unaligned slice start is counted
+        # here separately (per-thread order contract: head elements FIRST,
+        # then list entries — Phase 3's compact write replays the same).
         skip_ok = cutlass.Int32(0)
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
-            if (slice_start & cutlass.Int32(self.SKIP_BLOCK - 1)) == cutlass.Int32(0):
-                skip_ok = cutlass.Int32(1)
+            skip_ok = cutlass.Int32(1)
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             if skip_ok == cutlass.Int32(1):
+                head_end = (
+                    (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1))
+                    >> cutlass.Int32(self.SKIP_BLOCK_LOG2)
+                ) << cutlass.Int32(self.SKIP_BLOCK_LOG2)
+                if head_end > slice_end:
+                    head_end = slice_end
+                hh = slice_start + tidx
+                while hh < head_end:
+                    vh = self._load_fp32(input_row, hh)
+                    for m in cutlass.range_constexpr(M):
+                        cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vh >= thr_frag[m])
+                    hh = hh + cutlass.Int32(num_threads)
                 tmin = thr_frag[0]
                 for m in cutlass.range_constexpr(M):
                     tmin = _fmin_f32_inline(tmin, thr_frag[m])
@@ -2205,6 +2221,21 @@ def phase3_collect_candidates(
                 skip_wr = cutlass.Int32(1)
         if cutlass.const_expr(self.enable_block_skip and smem_active is not None):
             if skip_wr == cutlass.Int32(1):
+                # head region first — same per-thread order as the count pass
+                head_end_w = (
+                    (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1))
+                    >> cutlass.Int32(self.SKIP_BLOCK_LOG2)
+                ) << cutlass.Int32(self.SKIP_BLOCK_LOG2)
+                if head_end_w > slice_end:
+                    head_end_w = slice_end
+                hh_w = slice_start + tidx
+                while hh_w < head_end_w:
+                    vh_w = self._load_fp32(input_row, hh_w)
+                    if vh_w >= thr_final and wc < cutlass.Int32(kCC):
+                        smem_keys[wc] = vh_w
+                        smem_vals[wc] = hh_w
+                        wc = wc + cutlass.Int32(1)
+                    hh_w = hh_w + cutlass.Int32(num_threads)
                 chunks_per_block_w = cutlass.const_expr(
                     self.SKIP_BLOCK // (self.vec_bits // self.dtype.width)
                 )

From 36a542f12bbae649268413edda5a384998eda480 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 06:42:20 -0700
Subject: [PATCH 004/117] [None][perf] GVR block-skip: rung-tightening
 active-list build
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The active list was built at the loosest rung over all M columns. On
low-hit-rate rows the lowest sample-quantile rung sits far below the
K-th value (real pro 1024k: retained fraction 0.63 at that rung vs
0.08 at the final threshold), so the list barely skipped anything.

Build now iterates (cs==1): if the list exceeds CAP = 3/4 kC blocks,
DROP that rung — dropping is always correct (the rung is merely an
unmeasured probe; its partial counts are excluded from the admission
argmin and the fallback bracket seeding via a dropped-rung mask) — and
rebuild at the next tighter threshold, bounded by M-1 extra ~2-5us
builds. At cs>1 per-CTA list lengths differ (the drop decision would
diverge across the cluster), so the plain loosest-rung build stays.

Real-data 1024k, cs1, warm-L2 directional: pro 28.8 -> 16.4us (skip
ratio 1.16x -> 2.05x), flash 18.5 -> 12.8us (1.73x -> 2.49x).
Correctness: cs {1,2,4,8} x flash+pro x {64k,512k,1024k} all exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 111 ++++++++++++++----
 1 file changed, 88 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 88638c1d377d..e8db5ecea200 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -1567,24 +1567,71 @@ def block_count_ge_multi(
                     for m in cutlass.range_constexpr(M):
                         cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vh >= thr_frag[m])
                     hh = hh + cutlass.Int32(num_threads)
-                tmin = thr_frag[0]
-                for m in cutlass.range_constexpr(M):
-                    tmin = _fmin_f32_inline(tmin, thr_frag[m])
-                # smem_wcnt_multi's first num_warps slots double as the
-                # build's warp-total scratch: the build's own barriers
-                # complete before any count lands there.
-                self._full_build_active(
-                    block_max_row,
-                    slice_start,
-                    slice_end,
-                    tmin,
-                    smem_wcnt_multi,
-                    smem_active,
-                    s_active_cnt,
-                    tidx,
-                    warp_id,
-                    lane,
-                )
+                # Rung-tightening build (cs==1): a rung whose active list
+                # exceeds CAP blocks is provably or near-provably
+                # unacceptable (count >= list length; on real low-hit-rate
+                # rows the loosest sample-quantile rung retains 60%+ of the
+                # blocks and destroys the skip). DROP it — a dropped rung is
+                # merely an unmeasured probe (recorded in the mask at
+                # s_active_cnt[2]; classify and the fallback seeding skip
+                # it) — and rebuild at the next tighter threshold. Bounded
+                # by M-1 extra builds (~2-5us each at nb=8192). At cs>1 the
+                # per-CTA list lengths differ, so the drop decision would
+                # diverge across the cluster: keep the plain loosest-rung
+                # build there.
+                CAP_BLOCKS = cutlass.const_expr(3 * self.kC // 4)
+                if cutlass.const_expr(cluster_size == 1):
+                    build_done = cutlass.Int32(0)
+                    for _attempt in cutlass.range_constexpr(M):
+                        if build_done == cutlass.Int32(0):
+                            dmask = s_active_cnt[2]
+                            tcur = cutlass.Float32(self.FLT_MAX)
+                            mcur = cutlass.Int32(-1)
+                            kept = cutlass.Int32(0)
+                            for m in cutlass.range_constexpr(M):
+                                if (
+                                    dmask & (cutlass.Int32(1) << cutlass.Int32(m))
+                                ) == cutlass.Int32(0):
+                                    kept = kept + cutlass.Int32(1)
+                                    if thr_frag[m] < tcur:
+                                        tcur = thr_frag[m]
+                                        mcur = cutlass.Int32(m)
+                            self._full_build_active(
+                                block_max_row,
+                                slice_start,
+                                slice_end,
+                                tcur,
+                                smem_wcnt_multi,
+                                smem_active,
+                                s_active_cnt,
+                                tidx,
+                                warp_id,
+                                lane,
+                            )
+                            if s_active_cnt[0] <= cutlass.Int32(
+                                CAP_BLOCKS
+                            ) or kept <= cutlass.Int32(1):
+                                build_done = cutlass.Int32(1)
+                            else:
+                                if tidx == 0:
+                                    s_active_cnt[2] = dmask | (cutlass.Int32(1) << mcur)
+                                cute.arch.barrier()
+                else:
+                    tmin = thr_frag[0]
+                    for m in cutlass.range_constexpr(M):
+                        tmin = _fmin_f32_inline(tmin, thr_frag[m])
+                    self._full_build_active(
+                        block_max_row,
+                        slice_start,
+                        slice_end,
+                        tmin,
+                        smem_wcnt_multi,
+                        smem_active,
+                        s_active_cnt,
+                        tidx,
+                        warp_id,
+                        lane,
+                    )
                 if tidx == 0:
                     s_active_cnt[1] = cutlass.Int32(1)  # list-current flag
                 chunks_per_block = cutlass.const_expr(
@@ -4203,7 +4250,7 @@ def run_one_row(
             )
             s_active_cnt = smem.allocate_tensor(
                 element_type=cutlass.Int32,
-                layout=cute.make_ordered_layout((2,), order=(0,)),
+                layout=cute.make_ordered_layout((4,), order=(0,)),
                 byte_alignment=16,
             )
         else:
@@ -4601,6 +4648,7 @@ def _run_phases(
         if cutlass.const_expr(self.enable_block_skip):
             if tidx == cutlass.Int32(0):
                 s_active_cnt[1] = cutlass.Int32(0)
+                s_active_cnt[2] = cutlass.Int32(0)  # dropped-rung mask
 
         # ---- Phase 1: preIdx Min/Max/Mean ----
         self.phase1_preidx_stats(
@@ -4721,6 +4769,11 @@ def _run_phases(
                     # (Explicit argmin: with r0_vseed the pmean column is not
                     # sorted into the rung order; for sorted rungs this is
                     # equivalent to the old "last m in window" rule.)
+                    # Dropped rungs (block-skip rung tightening) hold PARTIAL
+                    # counts — never admissible.
+                    dmask_c = cutlass.Int32(0)
+                    if cutlass.const_expr(self.enable_block_skip):
+                        dmask_c = s_active_cnt[2]
                     best_m = cutlass.Int32(-1)
                     best_c = cutlass.Int32(2147483647)
                     for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
@@ -4729,6 +4782,8 @@ def _run_phases(
                             cm >= cutlass.Int32(self.top_k)
                             and cm <= cutlass.Int32(self.kC)
                             and cm < best_c
+                            and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m)))
+                            == cutlass.Int32(0)
                         ):
                             best_m = cutlass.Int32(m)
                             best_c = cm
@@ -4776,16 +4831,26 @@ def _run_phases(
                             bhi = v_hi
                             clo = cutlass.Int32(-1)
                             chi = cutlass.Int32(-1)
+                            dmask_f = cutlass.Int32(0)
+                            if cutlass.const_expr(self.enable_block_skip):
+                                dmask_f = s_active_cnt[2]
                             for m in cutlass.range_constexpr(M):
                                 cm = s_mt_cnt[m]
                                 tm = s_mt_thr[m]
-                                if cm > cutlass.Int32(self.kC) and (
-                                    clo < cutlass.Int32(0) or tm > blo
+                                m_ok = (
+                                    dmask_f & (cutlass.Int32(1) << cutlass.Int32(m))
+                                ) == cutlass.Int32(0)
+                                if (
+                                    m_ok
+                                    and cm > cutlass.Int32(self.kC)
+                                    and (clo < cutlass.Int32(0) or tm > blo)
                                 ):
                                     blo = tm
                                     clo = cm
-                                if cm < cutlass.Int32(self.top_k) and (
-                                    chi < cutlass.Int32(0) or tm < bhi
+                                if (
+                                    m_ok
+                                    and cm < cutlass.Int32(self.top_k)
+                                    and (chi < cutlass.Int32(0) or tm < bhi)
                                 ):
                                     bhi = tm
                                     chi = cm

From 086eabab136dff8715d7c63cf750fbb2202d97a4 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 07:39:14 -0700
Subject: [PATCH 005/117] [None][fix] GVR block-skip: guard list capacity/id
 width, bound tail loads
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review fixes for the compact machinery:
- skip_ok now also requires nb_slice <= SKIP_MAX_BLOCKS and absolute
  block id < 32768 (int16 list entries); wider/higher slices fall back
  to the dense walk losslessly instead of silently truncating counts
  and the collect (or wrapping ids negative at cluster_size > 1).
- Both compact walks vector-load only fully in-bounds chunks; the
  slice-end straddle goes through the scalar path, so an unaligned row
  no longer reads past the row/allocation.
- Ctor rejects enable_block_skip without enable_r0 (dead 16KB SMEM).
- emu_block_max defaults to records='positional' (the shipped kernel is
  grain 32; 'rotate' is a grain-128 fold fixture) and the wrapper
  asserts block_max covers every 32-position record of the row.

Validated: capacity boundary (8192/8193/9375 blocks), N=1.05M at
cs=4/8, unaligned exact-size rows (N=1000/65535/65529), real-data
flash/pro 64k/1024k — all exact with planted tail winners.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 70 +++++++++++++------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  5 +-
 2 files changed, 51 insertions(+), 24 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index e8db5ecea200..005ff1b82b54 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -443,6 +443,11 @@ def __init__(
         self.skip_order = "grouped"
         if enable_block_skip and num_threads not in (512, 1024):
             raise ValueError("enable_block_skip requires num_threads in {512, 1024}")
+        if enable_block_skip and not enable_r0:
+            # The compact machinery hangs off the R0 count pass and the
+            # phase-3 stream-write; without R0 the 16KB list SMEM would be
+            # allocated but the skip could never engage.
+            raise ValueError("enable_block_skip requires enable_r0")
         # C7 dispatch (op#26 host policy folded into the ctor; all gated on
         # enable_r0 so an OFF kernel is byte-identical to the base):
         #  - qfracs default = M2D (0.85, 0.35): dispatch_r0_op26 ships M2D for
@@ -1553,6 +1558,22 @@ def block_count_ge_multi(
         skip_ok = cutlass.Int32(0)
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             skip_ok = cutlass.Int32(1)
+            # Capacity/id-width guard: the active list holds at most
+            # SKIP_MAX_BLOCKS local ids and _list_st stores ABSOLUTE block
+            # ids as int16. A slice over 8192 full blocks (N_local >
+            # 262144) or reaching absolute id >= 32768 falls back to the
+            # dense walk (lossless; the list-current flag is never set, so
+            # phase3 stays dense too).
+            blk_lo_g = (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32(
+                self.SKIP_BLOCK_LOG2
+            )
+            blk_hi_g = (slice_end + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32(
+                self.SKIP_BLOCK_LOG2
+            )
+            if blk_hi_g - blk_lo_g > cutlass.Int32(self.SKIP_MAX_BLOCKS):
+                skip_ok = cutlass.Int32(0)
+            if blk_hi_g > cutlass.Int32(32767):
+                skip_ok = cutlass.Int32(0)
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             if skip_ok == cutlass.Int32(1):
                 head_end = (
@@ -1658,17 +1679,22 @@ def block_count_ge_multi(
                             pos0 = blk * cutlass.Int32(self.SKIP_BLOCK) + my_chunk0 * cutlass.Int32(
                                 vec_w
                             )
-                            src_ptr_u = cute.make_ptr(
-                                self.dtype,
-                                row_addr + cutlass.Int64(pos0) * cutlass.Int64(elem_bytes),
-                                cute.AddressSpace.gmem,
-                                assumed_align=vec_align,
-                            )
-                            cute.copy(
-                                copy_atom,
-                                cute.make_tensor(src_ptr_u, cute.make_layout((vec_w,))),
-                                frags[u],
-                            )
+                            # Vector-load only fully in-bounds chunks; the
+                            # slice-end straddle re-reads scalars below (a
+                            # tail block's chunks would otherwise read past
+                            # the row/allocation when N % 32 != 0).
+                            if pos0 + cutlass.Int32(vec_w) <= slice_end:
+                                src_ptr_u = cute.make_ptr(
+                                    self.dtype,
+                                    row_addr + cutlass.Int64(pos0) * cutlass.Int64(elem_bytes),
+                                    cute.AddressSpace.gmem,
+                                    assumed_align=vec_align,
+                                )
+                                cute.copy(
+                                    copy_atom,
+                                    cute.make_tensor(src_ptr_u, cute.make_layout((vec_w,))),
+                                    frags[u],
+                                )
                         poss.append(pos0)
                         valids.append(valid)
                     for u in cutlass.range_constexpr(UN):
@@ -2297,18 +2323,18 @@ def phase3_collect_candidates(
                     pos0_w = blk_w * cutlass.Int32(self.SKIP_BLOCK) + my_chunk0_w * cutlass.Int32(
                         vec_w
                     )
-                    src_ptr_w = cute.make_ptr(
-                        self.dtype,
-                        row_addr + cutlass.Int64(pos0_w) * cutlass.Int64(elem_bytes),
-                        cute.AddressSpace.gmem,
-                        assumed_align=vec_align,
-                    )
-                    cute.copy(
-                        copy_atom,
-                        cute.make_tensor(src_ptr_w, cute.make_layout((vec_w,))),
-                        wfrag,
-                    )
                     if pos0_w + cutlass.Int32(vec_w) <= slice_end:
+                        src_ptr_w = cute.make_ptr(
+                            self.dtype,
+                            row_addr + cutlass.Int64(pos0_w) * cutlass.Int64(elem_bytes),
+                            cute.AddressSpace.gmem,
+                            assumed_align=vec_align,
+                        )
+                        cute.copy(
+                            copy_atom,
+                            cute.make_tensor(src_ptr_w, cute.make_layout((vec_w,))),
+                            wfrag,
+                        )
                         for j in cutlass.range_constexpr(vec_w):
                             if cutlass.const_expr(self.dtype == cutlass.Float32):
                                 vj = wfrag[j]
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index f2966563d663..9d908b3c7367 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -192,7 +192,7 @@ def emu_block_max(
     next_n: int = 1,
     compress_ratio: int = 1,
     tail_mode: str = "pad_inf",
-    records: str = "rotate",
+    records: str = "positional",
 ) -> torch.Tensor:
     """``[num_rows, nb_pad*4] fp32`` warp-partial upper-bound records.
 
@@ -421,7 +421,8 @@ def gvr_topk_decode(
             and block_max.dim() == 2
             and block_max.shape[0] == num_rows
             and block_max.shape[1] % 4 == 0
-        ), "block_max must be contiguous CUDA fp32 [num_rows, nb_pad*4]"
+            and block_max.shape[1] >= (logits.shape[1] + 31) // 32
+        ), "block_max must be contiguous CUDA fp32 [num_rows, nb_pad*4] covering the row"
 
     if return_output_values:
         if out_values is None:

From 01bc507f51d265d8531bcaf3eedd93f9ca3dd9d4 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 08:06:03 -0700
Subject: [PATCH 006/117] [None][perf] GVR block-skip: host dispatch gate below
 N=200k

The compact walk only wins on long rows (cold-L2 protocol: >= 2.18x at
N=262k, 4-14% loss at N <= 131k). Gate block_max shape-based (no device
sync) behind skip_min_n=200_000: below it the wrapper drops to the
dense arms. Protocol after gating: flash/pro 256k/512k cells all
0.99-1.02x, 1024k wins intact (flash BS1 15.63us 2.18x, BS1024 4.18x;
pro BS1024 3.15x).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 9d908b3c7367..cbc58b6b2d4f 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -343,6 +343,7 @@ def gvr_topk_decode(
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
     block_max: Optional[torch.Tensor] = None,
+    skip_min_n: Optional[int] = 200_000,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """CuTe DSL GVR Top-K wrapper with every tuning knob exposed.
 
@@ -412,6 +413,13 @@ def gvr_topk_decode(
     cute_dtype = _DTYPE_TORCH_TO_CUTE[logits.dtype]
 
     num_rows = logits.shape[0]
+    # Host dispatch gate: the compact walk only wins when the per-row
+    # compressed length is large (protocol: >= 2x at N=262k, 4-14% LOSS at
+    # N <= 131k). Below skip_min_n (compressed-index space, shape-based so
+    # no device sync) drop block_max and run the dense arms. None disables
+    # the gate (A/B probes).
+    if block_max is not None and skip_min_n is not None and logits.shape[1] < skip_min_n:
+        block_max = None
     enable_block_skip = block_max is not None
     if enable_block_skip:
         assert (

From e9f4865e281d1a0fe49f5b3d0b5d2956fa2fb098 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 21 Jul 2026 03:42:41 -0700
Subject: [PATCH 007/117] [None][perf] GVR block-skip: skip-aware launch policy
 and dispatch gates

pick_config gains has_block_max: with bounds available and N >= 200k
the policy pins cluster_size = 1 - the compact list + rung tightening
(cs1-only) beat the row-split configs outright once the bounds prune
the scan (cold protocol, real data: BS1 1.21x, BS64 2.12x, BS1024
4.16x vs the stock picks; splitting shrinks each CTA slice below the
skip break-even and disables tightening).

The wrapper dispatch gains a second gate next to skip_min_n: K > 512
at num_rows < 8 keeps the stock path - the acceptance band is
proportionally tighter (kC/K = 6 vs 10), the bounds prune less, and
the row-split configs win (pro 262k BS1: skip 21.3-21.6us at cs1/cs8
vs stock cs8 19.7us).

Cold protocol vs op26 at its own launch policy, 24 cells: zero
regressions; flash 1024k 1.21/2.12/4.16x (BS 1/64/1024), pro 1024k
1.68/3.15x (BS 64/1024), everything below the gates identical to
stock.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py                | 11 +++++++++++
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py  |  6 ++++++
 2 files changed, 17 insertions(+)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 005ff1b82b54..795f70a81148 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -5272,6 +5272,7 @@ def pick_config(
         num_candidates: int,
         max_seq_len: Optional[int] = None,
         num_sms: Optional[int] = None,
+        has_block_max: bool = False,
     ) -> dict:
         """Pick the launch-shape ctor kwargs for ``(dtype, BS, N)``.
 
@@ -5306,6 +5307,16 @@ def pick_config(
         # large N -> 8; single-wave -> 4/2; multi-wave -> 1.
         if n_row < 65536:
             cluster_size = 1
+        elif has_block_max and n_row >= 200_000:
+            # Block-skip sweet spot is cs == 1 with a large per-CTA slice:
+            # the compact list + rung tightening (cs1-only) beat the
+            # row-split configs outright once the bounds prune the scan
+            # (cold protocol, real data: BS1 1.18x, BS64 2.14x, BS1024
+            # 4.16x vs this policy's stock picks; splitting shrinks each
+            # CTA's slice below the skip break-even and disables
+            # tightening). Below 200k the wrapper drops block_max anyway
+            # (skip_min_n gate) and the stock picks apply.
+            cluster_size = 1
         elif num_rows <= 4 and n_row >= 131072:
             cluster_size = 8
         elif num_rows * 4 <= num_sms:
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index cbc58b6b2d4f..8bb6ef90f49b 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -420,6 +420,12 @@ def gvr_topk_decode(
     # the gate (A/B probes).
     if block_max is not None and skip_min_n is not None and logits.shape[1] < skip_min_n:
         block_max = None
+    # K > 512 at tiny batch: the acceptance band is proportionally tighter
+    # (kC/K = 6 vs 10), the bounds prune less, and the row-split configs
+    # win outright (cold protocol, pro 262k BS1: skip 21.3-21.6us at
+    # cs1/cs8 vs stock cs8 19.7us) -> keep the stock path.
+    if block_max is not None and num_rows < 8 and top_k > 512:
+        block_max = None
     enable_block_skip = block_max is not None
     if enable_block_skip:
         assert (

From 344b082827a1ba64fe6434150ba71278388fd41d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 09:47:29 -0700
Subject: [PATCH 008/117] [None][feat] GVR: emu references for epilogue seed
 counts and candidates

emu_seed_counts / emu_cand implement the A-side products per the
epilogue<->topk buffer contract v2 so the consumer waterfall can be
developed and tested against torch references before the fused
indexer lands. Real-data coverage probe: 7/8 cells have a seed count
inside [K, kC]; prev-kth drifts too loose at long context, so the L2
collect threshold should be the middle rung / xstate-adaptive.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 71 +++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 8bb6ef90f49b..d150bdaad511 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -293,6 +293,77 @@ def emu_block_meta(
     return meta.contiguous()
 
 
+def emu_seed_counts(
+    logits: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seed_thr: torch.Tensor,
+    next_n: int = 1,
+    compress_ratio: int = 1,
+) -> torch.Tensor:
+    """L1 emu: exact per-row threshold counts.
+
+    counts[r][j] = |{i < N_eff(r) : logits[r, i] >= t_j}| on the
+    post-conversion values (contract: epilogue_topk_interface.md).
+    """
+    R, C = logits.shape
+    n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio).unsqueeze(1)
+    pos = torch.arange(C, device=logits.device).unsqueeze(0)
+    lf = logits.to(torch.float32)
+    valid = pos < n_eff
+    counts = torch.empty((R, seed_thr.shape[1]), dtype=torch.int32, device=logits.device)
+    for j in range(seed_thr.shape[1]):
+        counts[:, j] = ((lf >= seed_thr[:, j : j + 1]) & valid).sum(-1, dtype=torch.int32)
+    return counts
+
+
+def emu_cand(
+    logits: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seed_thr: torch.Tensor,
+    cap: int,
+    next_n: int = 1,
+    compress_ratio: int = 1,
+    sentinel_pad: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """L2 emu: unordered candidate pre-collect.
+
+    Unordered (value fp32-bits, index) pairs of all valid positions
+    >= t_0 = seed_thr[:, 0]; ctl = {claimed, void}. claimed may
+    over-approximate the true count (window sentinels, idx word = -1) —
+    ``sentinel_pad`` injects that legally. void=1 when claimed > cap; on
+    overflow only the first ``cap`` entries are materialized (contract v2:
+    consumers scan [0, min(claimed, cap)) skipping sentinels).
+    """
+    R, C = logits.shape
+    dev = logits.device
+    n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio)
+    lf = logits.to(torch.float32)
+    cand = torch.full((R, cap * 2), -1, dtype=torch.int32, device=dev)
+    ctl = torch.zeros((R, 2), dtype=torch.int32, device=dev)
+    pairs = cand.view(R, cap, 2)
+    for r in range(R):
+        ne = int(n_eff[r])
+        hits = torch.nonzero(lf[r, :ne] >= seed_thr[r, 0], as_tuple=False).flatten()
+        cnt = hits.numel()
+        # unordered contract: shuffle, then interleave sentinels
+        perm = hits[torch.randperm(cnt, device=dev)]
+        ent = torch.full((cnt + sentinel_pad,), -1, dtype=torch.int64, device=dev)
+        if sentinel_pad:
+            slots = torch.randperm(cnt + sentinel_pad, device=dev)[:cnt]
+            slots = slots.sort().values
+        else:
+            slots = torch.arange(cnt, device=dev)
+        ent[slots] = perm
+        claimed = int(ent.numel())
+        nwr = min(claimed, cap)
+        live = ent[:nwr] >= 0
+        pairs[r, :nwr, 1] = ent[:nwr].to(torch.int32)
+        pairs[r, :nwr, 0][live] = lf[r, ent[:nwr][live]].view(torch.int32)
+        ctl[r, 0] = claimed
+        ctl[r, 1] = 1 if claimed > cap else 0
+    return cand, ctl
+
+
 def enc_ordered_f32(t: torch.Tensor) -> torch.Tensor:
     """Order-preserving int encoding of fp32 (an involution).
 

From 289553f671d1aa1276a754ff08265ff25939e7e3 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 10:01:45 -0700
Subject: [PATCH 009/117] [None][feat] GVR: waterfall L1 admission from
 external epilogue counts

use_ext_counts: rung thresholds AND their exact counts arrive from the
indexer epilogue (seed_thr/seed_counts [rows, 3], interface v2), so
P1b and the M-ary R0 count pass are skipped. The seeded refine routes
both cases: an in-band rung is re-measured once (building the
per-thread hand-off Phase 3 requires) and accepted; a full miss seeds
log-falsi with the external brackets. cs==1, requires fb_fix and a
3-slot rung config (the wrapper pins 2 qfracs + vseed; the values are
irrelevant since P1b never runs).

Real-data validation (flash/pro x 16k..1024k, thresholds {prev-kth,
q35, q85} + emu counts): 8/8 exact on stock/ext/ext+skip arms incl.
the pro-1M full-miss bracket cell. Directional: +9-10% at 1M (P1b +
M-count saved), small-N slightly negative (the waterfall routes those
to the L2 direct path instead). Next: skip P1 under ext (outer
brackets from xstate) and the L2 direct-to-P4 branch.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 221 ++++++++++++------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  41 ++++
 2 files changed, 189 insertions(+), 73 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 795f70a81148..c13d6b34d4f2 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -271,6 +271,7 @@ def __init__(
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
         enable_block_skip: bool = False,
+        use_ext_counts: bool = False,
     ):
         # Redundant-warp sync reduction: every warp replays the block
         # reduce + decision from the same staged SMEM partials in the
@@ -539,6 +540,19 @@ def __init__(
             self.M_thr = self.M_qf + 1
         # need[m] = ceil(q_m * K) prev-topK values >= rung m.
         self.qneeds = tuple(max(1, int(math.ceil(q * self.top_k))) for q in self.r0_qfracs)
+        # use_ext_counts (waterfall L1 admission): thresholds AND their
+        # exact counts arrive from the indexer epilogue (interface v2) —
+        # P1b and the M-ary count pass are skipped; an in-band rung is
+        # re-measured ONCE through the seeded refine (per-thread hand-off)
+        # and accepted; a miss seeds log-falsi with the external brackets.
+        self.use_ext_counts = bool(use_ext_counts) and bool(enable_r0)
+        if self.use_ext_counts:
+            if not self.fb_fix:
+                raise ValueError("use_ext_counts requires fb_fix")
+            if self.M_thr != 3:
+                raise ValueError("use_ext_counts expects exactly 3 seed rungs")
+            if cluster_size != 1:
+                raise ValueError("use_ext_counts is single-CTA (cs==1) only")
         # R1 inline shot aim in log2-count space: geometric center of the
         # [K, kC] acceptance window.
         self.log2_r1aim = math.log2(math.sqrt(self.top_k * self.kC)) if self.r0_qfracs else 0.0
@@ -4103,6 +4117,8 @@ def gvr_topk_kernel(
         output_indices: cute.Tensor,  # [numRows, top_k] int32
         order_row: cute.Tensor,  # [batch_size] int32 (or None when seqlen_sorted=False)
         block_max: cute.Tensor,  # [numRows, nb_pad*4] fp32 (or None: no block-skip)
+        seed_thr: cute.Tensor,  # [numRows, 3] fp32 (or None: no ext counts)
+        seed_counts: cute.Tensor,  # [numRows, 3] int32 (or None)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
 
@@ -4155,6 +4171,8 @@ def gvr_topk_kernel(
             output_values,
             output_indices,
             block_max=block_max,
+            seed_thr=seed_thr,
+            seed_counts=seed_counts,
         )
 
     @cute.jit
@@ -4167,6 +4185,8 @@ def run_one_row(
         output_values: cute.Tensor,  # [numRows, top_k] dtype, optional
         output_indices: cute.Tensor,  # [numRows, top_k] int32
         block_max: cute.Tensor = None,  # [numRows, nb_pad*4] fp32
+        seed_thr: cute.Tensor = None,  # [numRows, 3] fp32 (ext counts)
+        seed_counts: cute.Tensor = None,  # [numRows, 3] int32 (ext counts)
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
 
@@ -4227,6 +4247,14 @@ def run_one_row(
             block_max_row = block_max[row_idx, None]
         else:
             block_max_row = None
+        if cutlass.const_expr(
+            self.use_ext_counts and seed_thr is not None and seed_counts is not None
+        ):
+            seed_thr_row = seed_thr[row_idx, None]
+            seed_counts_row = seed_counts[row_idx, None]
+        else:
+            seed_thr_row = None
+            seed_counts_row = None
         # When return_output_values=False, ``output_values`` is None at
         # launch and the gated writes below are compiled out; slicing into
         # None would crash so we keep the view None as well.
@@ -4523,6 +4551,8 @@ def run_one_row(
                         warp_id,
                         lane,
                         block_max_row=block_max_row,
+                        seed_thr_row=seed_thr_row,
+                        seed_counts_row=seed_counts_row,
                         smem_active=smem_active,
                         s_active_cnt=s_active_cnt,
                     )
@@ -4567,6 +4597,8 @@ def run_one_row(
                             warp_id,
                             lane,
                             block_max_row=block_max_row,
+                            seed_thr_row=seed_thr_row,
+                            seed_counts_row=seed_counts_row,
                             smem_active=smem_active,
                             s_active_cnt=s_active_cnt,
                         )
@@ -4608,6 +4640,8 @@ def run_one_row(
                     warp_id,
                     lane,
                     block_max_row=block_max_row,
+                    seed_thr_row=seed_thr_row,
+                    seed_counts_row=seed_counts_row,
                     smem_active=smem_active,
                     s_active_cnt=s_active_cnt,
                 )
@@ -4652,6 +4686,8 @@ def _run_phases(
         warp_id,
         lane,
         block_max_row=None,  # block-skip: this row's per-32-position bounds
+        seed_thr_row=None,  # ext counts: this row's 3 seed thresholds (fp32)
+        seed_counts_row=None,  # ext counts: this row's 3 exact counts (int32)
         smem_active=None,
         s_active_cnt=None,
     ):
@@ -4751,85 +4787,107 @@ def _run_phases(
                 # CTA scans its slice and block_count_ge_multi cluster-merges
                 # the rung counts (phase1b rungs are per-CTA identical since
                 # preIdx stats are full-row).
-                if cutlass.const_expr(self.p1b_cache):
-                    # rungs from the SMEM gather-cache P1 stashed (no 2nd
-                    # GMEM gather); 16-bit only.
-                    self.phase1b_hspace_rungs_cached(
-                        pre_idx_count, smem_gath, smem_hist, s_thr, s_mt_thr, tidx, warp_id, lane
-                    )
-                else:
-                    self.phase1b_hspace_rungs(
+                if cutlass.const_expr(self.use_ext_counts):
+                    # ---- Waterfall L1 admission (ext counts) ----
+                    # Rung thresholds AND exact counts arrive from the
+                    # indexer epilogue: skip P1b + the M-ary count pass.
+                    # Route through the seeded refine (s_r0col = -1): an
+                    # in-band rung is re-measured once (builds the
+                    # per-thread hand-off Phase 3 requires) and accepted;
+                    # a full miss seeds log-falsi from the ext brackets.
+                    if tidx == cutlass.Int32(0):
+                        for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                            s_mt_thr[m] = seed_thr_row[m]
+                            s_mt_cnt[m] = cutlass.Int32(seed_counts_row[m])
+                        s_r0col[0] = cutlass.Int32(-1)
+                if cutlass.const_expr(not self.use_ext_counts):
+                    if cutlass.const_expr(self.p1b_cache):
+                        # rungs from the SMEM gather-cache P1 stashed (no 2nd
+                        # GMEM gather); 16-bit only.
+                        self.phase1b_hspace_rungs_cached(
+                            pre_idx_count,
+                            smem_gath,
+                            smem_hist,
+                            s_thr,
+                            s_mt_thr,
+                            tidx,
+                            warp_id,
+                            lane,
+                        )
+                    else:
+                        self.phase1b_hspace_rungs(
+                            input_row,
+                            N,
+                            pre_idx_row,
+                            pre_idx_count,
+                            pre_idx_offset,
+                            smem_hist,
+                            s_thr,
+                            s_mt_thr,
+                            tidx,
+                            warp_id,
+                            lane,
+                        )
+                    self.block_count_ge_multi(
                         input_row,
-                        N,
-                        pre_idx_row,
-                        pre_idx_count,
-                        pre_idx_offset,
-                        smem_hist,
-                        s_thr,
+                        slice_start,
+                        slice_end,
                         s_mt_thr,
+                        smem_ptcnt_multi,
+                        smem_wcnt_multi,
+                        s_mt_cnt,
+                        s_cluster_partial_m,
+                        do_cluster_sync,
                         tidx,
                         warp_id,
                         lane,
+                        smem_ptcnt=smem_ptcnt,
+                        block_max_row=block_max_row,
+                        smem_active=smem_active,
+                        s_active_cnt=s_active_cnt,
                     )
-                self.block_count_ge_multi(
-                    input_row,
-                    slice_start,
-                    slice_end,
-                    s_mt_thr,
-                    smem_ptcnt_multi,
-                    smem_wcnt_multi,
-                    s_mt_cnt,
-                    s_cluster_partial_m,
-                    do_cluster_sync,
-                    tidx,
-                    warp_id,
-                    lane,
-                    smem_ptcnt=smem_ptcnt,
-                    block_max_row=block_max_row,
-                    smem_active=smem_active,
-                    s_active_cnt=s_active_cnt,
-                )
-                cute.arch.barrier()
-                if tidx == 0:
-                    # tightest admissible rung = SMALLEST count in [K, kC].
-                    # (Explicit argmin: with r0_vseed the pmean column is not
-                    # sorted into the rung order; for sorted rungs this is
-                    # equivalent to the old "last m in window" rule.)
-                    # Dropped rungs (block-skip rung tightening) hold PARTIAL
-                    # counts — never admissible.
-                    dmask_c = cutlass.Int32(0)
-                    if cutlass.const_expr(self.enable_block_skip):
-                        dmask_c = s_active_cnt[2]
-                    best_m = cutlass.Int32(-1)
-                    best_c = cutlass.Int32(2147483647)
-                    for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                        cm = s_mt_cnt[m]
-                        if (
-                            cm >= cutlass.Int32(self.top_k)
-                            and cm <= cutlass.Int32(self.kC)
-                            and cm < best_c
-                            and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m)))
-                            == cutlass.Int32(0)
-                        ):
-                            best_m = cutlass.Int32(m)
-                            best_c = cm
-                    s_r0col[0] = best_m
-                    if best_m >= cutlass.Int32(0):
-                        s_thr[0] = s_mt_thr[best_m]
-                        s_iscalars[0] = s_mt_cnt[best_m]
-                        # done=1: the threshold is admitted, so Phase 3 must
-                        # SKIP its retry-shrink and honor s_thr[0]. (block_count
-                        # _ge / secant leave done via their own path; the R0
-                        # admission must set it explicitly or Phase 3 re-searches
-                        # and the cluster collect diverges -> wrong output.)
-                        s_iscalars[1] = cutlass.Int32(1)
-                        # Snapshot this CTA's LOCAL slice count for the chosen
-                        # rung into s_iscalars[5] — the per-CTA cand_count that
-                        # Phase 3/4's cluster gather consumes (block_count_ge
-                        # sets it too; the R0 admission must match). Without it
-                        # the cluster collect under-counts -> wrong output.
-                        if cutlass.const_expr(cluster_size > 1):
-                            s_iscalars[5] = s_cluster_partial_m[best_m]
+                    cute.arch.barrier()
+                    if tidx == 0:
+                        # tightest admissible rung = SMALLEST count in [K, kC].
+                        # (Explicit argmin: with r0_vseed the pmean column is not
+                        # sorted into the rung order; for sorted rungs this is
+                        # equivalent to the old "last m in window" rule.)
+                        # Dropped rungs (block-skip rung tightening) hold PARTIAL
+                        # counts — never admissible.
+                        dmask_c = cutlass.Int32(0)
+                        if cutlass.const_expr(self.enable_block_skip):
+                            dmask_c = s_active_cnt[2]
+                        best_m = cutlass.Int32(-1)
+                        best_c = cutlass.Int32(2147483647)
+                        for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                            cm = s_mt_cnt[m]
+                            if (
+                                cm >= cutlass.Int32(self.top_k)
+                                and cm <= cutlass.Int32(self.kC)
+                                and cm < best_c
+                                and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m)))
+                                == cutlass.Int32(0)
+                            ):
+                                best_m = cutlass.Int32(m)
+                                best_c = cm
+                        s_r0col[0] = best_m
+                        if best_m >= cutlass.Int32(0):
+                            s_thr[0] = s_mt_thr[best_m]
+                            s_iscalars[0] = s_mt_cnt[best_m]
+                            # done=1: the threshold is admitted, so Phase 3 must
+                            # SKIP its retry-shrink and honor s_thr[0]. (block_count
+                            # _ge / secant leave done via their own path; the R0
+                            # admission must set it explicitly or Phase 3 re-searches
+                            # and the cluster collect diverges -> wrong output.)
+                            s_iscalars[1] = cutlass.Int32(1)
+                            # Snapshot this CTA's LOCAL slice count for the chosen
+                            # rung into s_iscalars[5] — the per-CTA cand_count that
+                            # Phase 3/4's cluster gather consumes (block_count_ge
+                            # sets it too; the R0 admission must match). Without it
+                            # the cluster collect under-counts -> wrong output.
+                            if cutlass.const_expr(cluster_size > 1):
+                                s_iscalars[5] = s_cluster_partial_m[best_m]
+
                 cute.arch.barrier()
                 bc = s_r0col[0]
                 if bc >= cutlass.Int32(0) and bc < cutlass.Int32(self.M_qf):
@@ -4902,6 +4960,19 @@ def _run_phases(
                                 cand = bhi
                             elif clo < cutlass.Int32(0):
                                 cand = blo
+                            if cutlass.const_expr(self.use_ext_counts):
+                                # tightest ext rung already in [K, kC]:
+                                # measure exactly it (accepts in one pass).
+                                cbe = cutlass.Int32(2147483647)
+                                for m in cutlass.range_constexpr(M):
+                                    cm4 = s_mt_cnt[m]
+                                    if (
+                                        cm4 >= cutlass.Int32(self.top_k)
+                                        and cm4 <= cutlass.Int32(self.kC)
+                                        and cm4 < cbe
+                                    ):
+                                        cbe = cm4
+                                        cand = s_mt_thr[m]
                             s_thr[0] = cand
                         cute.arch.barrier()
                         rs = cutlass.Int32(0)
@@ -5207,6 +5278,8 @@ def __call__(
         order_row: cute.Tensor,  # or None when seqlen_sorted=False
         stream,
         block_max: cute.Tensor = None,  # block-skip bounds; None = disabled
+        seed_thr: cute.Tensor = None,  # [num_rows, 3] fp32 (ext counts)
+        seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 (ext counts)
     ):
         num_rows = input_data.shape[0]
         cluster_size = cutlass.const_expr(self.cluster_size)
@@ -5229,6 +5302,8 @@ def __call__(
             output_indices,
             order_row,
             block_max,
+            seed_thr,
+            seed_counts,
         ).launch(
             grid=(total_ctas, 1, 1),
             block=(self.num_threads, 1, 1),
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index d150bdaad511..d3a8ab64970f 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -65,6 +65,7 @@ def _compile(
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
     enable_block_skip: bool = False,
+    use_ext_counts: bool = False,
 ):
     """JIT-compile the GVR kernel for a specific knob combination.
 
@@ -133,6 +134,20 @@ def _compile(
         if enable_block_skip
         else None
     )
+    seed_thr_fake = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Float32, (n_rows, 3), stride_order=(1, 0), assumed_align=4
+        )
+        if use_ext_counts
+        else None
+    )
+    seed_counts_fake = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Int32, (n_rows, 3), stride_order=(1, 0), assumed_align=4
+        )
+        if use_ext_counts
+        else None
+    )
     fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
     kernel = GvrTopKKernel(
         dtype=cute_dtype,
@@ -152,6 +167,11 @@ def _compile(
         p4_warp_redundant=p4_warp_redundant,
         p2_warp_redundant=p2_warp_redundant,
         enable_block_skip=enable_block_skip,
+        use_ext_counts=use_ext_counts,
+        # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed. The
+        # qfrac VALUES are irrelevant on this path (P1b is skipped) — only
+        # the slot count matters.
+        r0_qfracs=(0.85, 0.35) if use_ext_counts else None,
     )
     return cute.compile(
         kernel,
@@ -163,6 +183,8 @@ def _compile(
         order_row_fake,
         stream=fake_stream,
         block_max=block_max_fake,
+        seed_thr=seed_thr_fake,
+        seed_counts=seed_counts_fake,
         options="--enable-tvm-ffi",
     )
 
@@ -415,6 +437,8 @@ def gvr_topk_decode(
     p2_warp_redundant: bool = True,
     block_max: Optional[torch.Tensor] = None,
     skip_min_n: Optional[int] = 200_000,
+    seed_thr: Optional[torch.Tensor] = None,
+    seed_counts: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """CuTe DSL GVR Top-K wrapper with every tuning knob exposed.
 
@@ -497,6 +521,20 @@ def gvr_topk_decode(
     # cs1/cs8 vs stock cs8 19.7us) -> keep the stock path.
     if block_max is not None and num_rows < 8 and top_k > 512:
         block_max = None
+    use_ext_counts = seed_thr is not None and seed_counts is not None
+    if use_ext_counts:
+        assert (
+            seed_thr.dtype == torch.float32
+            and seed_thr.is_cuda
+            and seed_thr.is_contiguous()
+            and seed_thr.shape == (num_rows, 3)
+        ), "seed_thr must be contiguous CUDA fp32 [num_rows, 3]"
+        assert (
+            seed_counts.dtype == torch.int32
+            and seed_counts.is_cuda
+            and seed_counts.is_contiguous()
+            and seed_counts.shape == (num_rows, 3)
+        ), "seed_counts must be contiguous CUDA int32 [num_rows, 3]"
     enable_block_skip = block_max is not None
     if enable_block_skip:
         assert (
@@ -575,6 +613,7 @@ def gvr_topk_decode(
         p4_warp_redundant,
         p2_warp_redundant,
         enable_block_skip,
+        use_ext_counts,
     )
     # When return_output_values=False the kernel was compiled to skip
     # STG.value and accepts None for the value-output slot.
@@ -588,6 +627,8 @@ def gvr_topk_decode(
         out_indices,
         order_row if seqlen_sorted else None,
         block_max if enable_block_skip else None,
+        seed_thr if use_ext_counts else None,
+        seed_counts if use_ext_counts else None,
     )
     if return_output_values:
         return out_values, out_indices

From 1fe323809930168620059adc673e8ec8d258210d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 10:13:04 -0700
Subject: [PATCH 010/117] [None][perf] GVR ext counts: skip Phase 1 (preIdx
 gather) entirely

With external epilogue counts the only surviving P1 products are the
[v_lo, v_hi] outer bracket and the scalar-state init; the ext rungs
provide the bracket directly (host contract: t_0 < t_2, finite, all
rows valid) and tid0 initializes the scalars. A miss whose target
falls outside [t_0, t_2] recovers through the refine loop's 8x
bracket expansion, same as the stock fail-soft.

Real data 8/8 exact unchanged; directional gains vs stock R0 improve
to 1.15x at flash 256k / 1.14x at flash 1M (from 1.04x/1.09x with P1
still running).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 54 ++++++++++++-------
 1 file changed, 36 insertions(+), 18 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index c13d6b34d4f2..4d9b2cc1fd88 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -4713,24 +4713,42 @@ def _run_phases(
                 s_active_cnt[2] = cutlass.Int32(0)  # dropped-rung mask
 
         # ---- Phase 1: preIdx Min/Max/Mean ----
-        self.phase1_preidx_stats(
-            input_row,
-            N,
-            pre_idx_row,
-            pre_idx_count,
-            pre_idx_offset,
-            smem_wmin,
-            smem_wmax,
-            smem_wsum,
-            smem_wcnt_p1,
-            s_thr,
-            s_iscalars,
-            tidx,
-            warp_id,
-            lane,
-            smem_gath=smem_gath,  # p1b_cache: stash gathered values (None-op OFF)
-            s_mt_thr=s_mt_thr,  # r0_vseed: park pmean in the last rung column
-        )
+        # ext counts: P1's only surviving products are the [v_lo, v_hi]
+        # outer bracket and the scalar state init — the ext rungs provide
+        # the bracket directly (host contract: t_0 < t_2, finite, all rows
+        # valid), so the preIdx gather is skipped wholesale. A miss whose
+        # target lies outside [t_0, t_2] recovers via the refine loop's
+        # 8x bracket expansion (same fail-soft as the stock path).
+        if cutlass.const_expr(self.use_ext_counts):
+            if tidx == cutlass.Int32(0):
+                s_thr[0] = seed_thr_row[1]
+                s_thr[1] = seed_thr_row[0]
+                s_thr[2] = seed_thr_row[2]
+                s_iscalars[0] = cutlass.Int32(0)  # cand_count
+                s_iscalars[1] = cutlass.Int32(0)  # done
+                s_iscalars[2] = cutlass.Int32(-1)  # cnt_lo (fb seeding owns)
+                s_iscalars[3] = cutlass.Int32(-1)  # cnt_hi
+                s_iscalars[4] = cutlass.Int32(0)  # out_count
+            cute.arch.barrier()
+        if cutlass.const_expr(not self.use_ext_counts):
+            self.phase1_preidx_stats(
+                input_row,
+                N,
+                pre_idx_row,
+                pre_idx_count,
+                pre_idx_offset,
+                smem_wmin,
+                smem_wmax,
+                smem_wsum,
+                smem_wcnt_p1,
+                s_thr,
+                s_iscalars,
+                tidx,
+                warp_id,
+                lane,
+                smem_gath=smem_gath,  # p1b_cache: stash gathered values (None-op OFF)
+                s_mt_thr=s_mt_thr,  # r0_vseed: park pmean in the last rung column
+            )
 
         # Degenerate threshold init: val_hi <= -self.FLT_MAX or val_lo >= val_hi.
         # When preIdx values produce an unusable bracket (e.g. all -inf or

From 44e9671ed13837e299e8f8264b405bccafca6094 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 10:38:45 -0700
Subject: [PATCH 011/117] [None][perf] GVR ext counts v2: compose with the
 stock count/skip path
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

v1 routed external-count admission through the dense seeded refine and
forfeited the compact-walk win (flash 1M ext 34.7us vs skipR0 15.6us
cold). v2 only skips P1b: the stock M-ary pass runs on the ext rungs,
so list build, rung tightening, per-thread hand-off and classify
compose unchanged. When an ext count is already in [K, kC] the
admitted threshold is parked in ALL rung slots (v2b) — the M-ary pass
degenerates to one compact single-threshold count and classify admits
it; a miss keeps the distinct rungs as measured brackets.

Real data 8/8 exact (stock/ext/ext+skip). Cold protocol: flash 64k
1.21x/1.25x (BS1/BS1024, the slim-admission cell); flash 1M ext+skip
16.1us vs skipR0 15.6us (composition recovered); pro-1M miss rows
still pay multi-count+refine (0.7x) — routing sends those to stock
skipR0 via xstate feedback (next step).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 171 +++++++++---------
 1 file changed, 89 insertions(+), 82 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 4d9b2cc1fd88..a41de7cb6ef8 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -4806,18 +4806,38 @@ def _run_phases(
                 # the rung counts (phase1b rungs are per-CTA identical since
                 # preIdx stats are full-row).
                 if cutlass.const_expr(self.use_ext_counts):
-                    # ---- Waterfall L1 admission (ext counts) ----
-                    # Rung thresholds AND exact counts arrive from the
-                    # indexer epilogue: skip P1b + the M-ary count pass.
-                    # Route through the seeded refine (s_r0col = -1): an
-                    # in-band rung is re-measured once (builds the
-                    # per-thread hand-off Phase 3 requires) and accepted;
-                    # a full miss seeds log-falsi from the ext brackets.
+                    # ---- Waterfall L1 admission (ext rungs, v2a) ----
+                    # Rung thresholds arrive from the indexer epilogue:
+                    # ONLY P1b is skipped. The stock M-ary count pass runs
+                    # on the ext rungs so the block-skip list build, rung
+                    # tightening, per-thread hand-off and classify all
+                    # compose unchanged (v1 routed through the dense
+                    # refine and forfeited the compact-walk win: flash 1M
+                    # ext 34.7us vs skipR0 15.6us cold).
+                    # v2b: when an ext count is already in [K, kC], park
+                    # THE ADMITTED THRESHOLD IN ALL RUNG SLOTS — the M-ary
+                    # pass degenerates to one compact single-threshold
+                    # count (+ list build at that threshold) and classify
+                    # admits it; a full miss keeps the 3 distinct rungs as
+                    # measured brackets for the seeded refine.
                     if tidx == cutlass.Int32(0):
+                        bx_m = cutlass.Int32(-1)
+                        bx_c = cutlass.Int32(2147483647)
                         for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                            s_mt_thr[m] = seed_thr_row[m]
-                            s_mt_cnt[m] = cutlass.Int32(seed_counts_row[m])
-                        s_r0col[0] = cutlass.Int32(-1)
+                            cx = cutlass.Int32(seed_counts_row[m])
+                            if (
+                                cx >= cutlass.Int32(self.top_k)
+                                and cx <= cutlass.Int32(self.kC)
+                                and cx < bx_c
+                            ):
+                                bx_m = cutlass.Int32(m)
+                                bx_c = cx
+                        for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                            if bx_m >= cutlass.Int32(0):
+                                s_mt_thr[m] = seed_thr_row[bx_m]
+                            else:
+                                s_mt_thr[m] = seed_thr_row[m]
+                    cute.arch.barrier()
                 if cutlass.const_expr(not self.use_ext_counts):
                     if cutlass.const_expr(self.p1b_cache):
                         # rungs from the SMEM gather-cache P1 stashed (no 2nd
@@ -4846,65 +4866,65 @@ def _run_phases(
                             warp_id,
                             lane,
                         )
-                    self.block_count_ge_multi(
-                        input_row,
-                        slice_start,
-                        slice_end,
-                        s_mt_thr,
-                        smem_ptcnt_multi,
-                        smem_wcnt_multi,
-                        s_mt_cnt,
-                        s_cluster_partial_m,
-                        do_cluster_sync,
-                        tidx,
-                        warp_id,
-                        lane,
-                        smem_ptcnt=smem_ptcnt,
-                        block_max_row=block_max_row,
-                        smem_active=smem_active,
-                        s_active_cnt=s_active_cnt,
-                    )
-                    cute.arch.barrier()
-                    if tidx == 0:
-                        # tightest admissible rung = SMALLEST count in [K, kC].
-                        # (Explicit argmin: with r0_vseed the pmean column is not
-                        # sorted into the rung order; for sorted rungs this is
-                        # equivalent to the old "last m in window" rule.)
-                        # Dropped rungs (block-skip rung tightening) hold PARTIAL
-                        # counts — never admissible.
-                        dmask_c = cutlass.Int32(0)
-                        if cutlass.const_expr(self.enable_block_skip):
-                            dmask_c = s_active_cnt[2]
-                        best_m = cutlass.Int32(-1)
-                        best_c = cutlass.Int32(2147483647)
-                        for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                            cm = s_mt_cnt[m]
-                            if (
-                                cm >= cutlass.Int32(self.top_k)
-                                and cm <= cutlass.Int32(self.kC)
-                                and cm < best_c
-                                and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m)))
-                                == cutlass.Int32(0)
-                            ):
-                                best_m = cutlass.Int32(m)
-                                best_c = cm
-                        s_r0col[0] = best_m
-                        if best_m >= cutlass.Int32(0):
-                            s_thr[0] = s_mt_thr[best_m]
-                            s_iscalars[0] = s_mt_cnt[best_m]
-                            # done=1: the threshold is admitted, so Phase 3 must
-                            # SKIP its retry-shrink and honor s_thr[0]. (block_count
-                            # _ge / secant leave done via their own path; the R0
-                            # admission must set it explicitly or Phase 3 re-searches
-                            # and the cluster collect diverges -> wrong output.)
-                            s_iscalars[1] = cutlass.Int32(1)
-                            # Snapshot this CTA's LOCAL slice count for the chosen
-                            # rung into s_iscalars[5] — the per-CTA cand_count that
-                            # Phase 3/4's cluster gather consumes (block_count_ge
-                            # sets it too; the R0 admission must match). Without it
-                            # the cluster collect under-counts -> wrong output.
-                            if cutlass.const_expr(cluster_size > 1):
-                                s_iscalars[5] = s_cluster_partial_m[best_m]
+                self.block_count_ge_multi(
+                    input_row,
+                    slice_start,
+                    slice_end,
+                    s_mt_thr,
+                    smem_ptcnt_multi,
+                    smem_wcnt_multi,
+                    s_mt_cnt,
+                    s_cluster_partial_m,
+                    do_cluster_sync,
+                    tidx,
+                    warp_id,
+                    lane,
+                    smem_ptcnt=smem_ptcnt,
+                    block_max_row=block_max_row,
+                    smem_active=smem_active,
+                    s_active_cnt=s_active_cnt,
+                )
+                cute.arch.barrier()
+                if tidx == 0:
+                    # tightest admissible rung = SMALLEST count in [K, kC].
+                    # (Explicit argmin: with r0_vseed the pmean column is not
+                    # sorted into the rung order; for sorted rungs this is
+                    # equivalent to the old "last m in window" rule.)
+                    # Dropped rungs (block-skip rung tightening) hold PARTIAL
+                    # counts — never admissible.
+                    dmask_c = cutlass.Int32(0)
+                    if cutlass.const_expr(self.enable_block_skip):
+                        dmask_c = s_active_cnt[2]
+                    best_m = cutlass.Int32(-1)
+                    best_c = cutlass.Int32(2147483647)
+                    for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                        cm = s_mt_cnt[m]
+                        if (
+                            cm >= cutlass.Int32(self.top_k)
+                            and cm <= cutlass.Int32(self.kC)
+                            and cm < best_c
+                            and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m)))
+                            == cutlass.Int32(0)
+                        ):
+                            best_m = cutlass.Int32(m)
+                            best_c = cm
+                    s_r0col[0] = best_m
+                    if best_m >= cutlass.Int32(0):
+                        s_thr[0] = s_mt_thr[best_m]
+                        s_iscalars[0] = s_mt_cnt[best_m]
+                        # done=1: the threshold is admitted, so Phase 3 must
+                        # SKIP its retry-shrink and honor s_thr[0]. (block_count
+                        # _ge / secant leave done via their own path; the R0
+                        # admission must set it explicitly or Phase 3 re-searches
+                        # and the cluster collect diverges -> wrong output.)
+                        s_iscalars[1] = cutlass.Int32(1)
+                        # Snapshot this CTA's LOCAL slice count for the chosen
+                        # rung into s_iscalars[5] — the per-CTA cand_count that
+                        # Phase 3/4's cluster gather consumes (block_count_ge
+                        # sets it too; the R0 admission must match). Without it
+                        # the cluster collect under-counts -> wrong output.
+                        if cutlass.const_expr(cluster_size > 1):
+                            s_iscalars[5] = s_cluster_partial_m[best_m]
 
                 cute.arch.barrier()
                 bc = s_r0col[0]
@@ -4978,19 +4998,6 @@ def _run_phases(
                                 cand = bhi
                             elif clo < cutlass.Int32(0):
                                 cand = blo
-                            if cutlass.const_expr(self.use_ext_counts):
-                                # tightest ext rung already in [K, kC]:
-                                # measure exactly it (accepts in one pass).
-                                cbe = cutlass.Int32(2147483647)
-                                for m in cutlass.range_constexpr(M):
-                                    cm4 = s_mt_cnt[m]
-                                    if (
-                                        cm4 >= cutlass.Int32(self.top_k)
-                                        and cm4 <= cutlass.Int32(self.kC)
-                                        and cm4 < cbe
-                                    ):
-                                        cbe = cm4
-                                        cand = s_mt_thr[m]
                             s_thr[0] = cand
                         cute.arch.barrier()
                         rs = cutlass.Int32(0)

From 604c379188ca7f4e75b6a0687a1956e4e6327ff2 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 10:48:09 -0700
Subject: [PATCH 012/117] [None][feat] GVR: xstate closed-loop writeback at
 Phase 4 exit
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

emit_xstate writes the per-row loop state (interface v2 layout
[rows, 8]: [0] valid, [1] kth proxy, [2] accepted threshold, [3]
cand_count from the pre-P4 snapshot — P4 repurposes the s_iscalars
slots) at the cs==1 Phase-4 exit; degenerate identity rows write
valid=0. The next step derives its seed rung group from these fields.

Real-data validation: exactness unchanged; state fields exact
(cand_count == count_ge(threshold): flash 991/633, pro 1854/2354);
same-step reseed from the written state admits in-band with the exact
count on ALL cells — including pro 1M, whose static rung group missed
entirely (the temporal rung fixes the miss AND slims flash 1M
admission 1290 -> 633).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 37 +++++++++++++++++++
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 21 +++++++++++
 2 files changed, 58 insertions(+)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index a41de7cb6ef8..f31251f65efd 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -272,6 +272,7 @@ def __init__(
         p2_warp_redundant: bool = True,
         enable_block_skip: bool = False,
         use_ext_counts: bool = False,
+        emit_xstate: bool = False,
     ):
         # Redundant-warp sync reduction: every warp replays the block
         # reduce + decision from the same staged SMEM partials in the
@@ -545,6 +546,12 @@ def __init__(
         # P1b and the M-ary count pass are skipped; an in-band rung is
         # re-measured ONCE through the seeded refine (per-thread hand-off)
         # and accepted; a miss seeds log-falsi with the external brackets.
+        # emit_xstate: write the per-row closed-loop state at Phase 4 exit
+        # (interface v2: [0] valid, [1] kth proxy, [2] accepted threshold,
+        # [3] cand_count). cs==1 only (leader-gather rows land later).
+        self.emit_xstate = bool(emit_xstate)
+        if emit_xstate and cluster_size != 1:
+            raise ValueError("emit_xstate is single-CTA (cs==1) only")
         self.use_ext_counts = bool(use_ext_counts) and bool(enable_r0)
         if self.use_ext_counts:
             if not self.fb_fix:
@@ -4119,6 +4126,7 @@ def gvr_topk_kernel(
         block_max: cute.Tensor,  # [numRows, nb_pad*4] fp32 (or None: no block-skip)
         seed_thr: cute.Tensor,  # [numRows, 3] fp32 (or None: no ext counts)
         seed_counts: cute.Tensor,  # [numRows, 3] int32 (or None)
+        xstate: cute.Tensor,  # [numRows, 8] fp32 closed-loop state (or None)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
 
@@ -4173,6 +4181,7 @@ def gvr_topk_kernel(
             block_max=block_max,
             seed_thr=seed_thr,
             seed_counts=seed_counts,
+            xstate=xstate,
         )
 
     @cute.jit
@@ -4187,6 +4196,7 @@ def run_one_row(
         block_max: cute.Tensor = None,  # [numRows, nb_pad*4] fp32
         seed_thr: cute.Tensor = None,  # [numRows, 3] fp32 (ext counts)
         seed_counts: cute.Tensor = None,  # [numRows, 3] int32 (ext counts)
+        xstate: cute.Tensor = None,  # [numRows, 8] fp32 (emit_xstate)
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
 
@@ -4255,6 +4265,10 @@ def run_one_row(
         else:
             seed_thr_row = None
             seed_counts_row = None
+        if cutlass.const_expr(self.emit_xstate and xstate is not None):
+            xstate_row = xstate[row_idx, None]
+        else:
+            xstate_row = None
         # When return_output_values=False, ``output_values`` is None at
         # launch and the gated writes below are compiled out; slicing into
         # None would crash so we keep the view None as well.
@@ -4553,6 +4567,7 @@ def run_one_row(
                         block_max_row=block_max_row,
                         seed_thr_row=seed_thr_row,
                         seed_counts_row=seed_counts_row,
+                        xstate_row=xstate_row,
                         smem_active=smem_active,
                         s_active_cnt=s_active_cnt,
                     )
@@ -4599,6 +4614,7 @@ def run_one_row(
                             block_max_row=block_max_row,
                             seed_thr_row=seed_thr_row,
                             seed_counts_row=seed_counts_row,
+                            xstate_row=xstate_row,
                             smem_active=smem_active,
                             s_active_cnt=s_active_cnt,
                         )
@@ -4642,6 +4658,7 @@ def run_one_row(
                     block_max_row=block_max_row,
                     seed_thr_row=seed_thr_row,
                     seed_counts_row=seed_counts_row,
+                    xstate_row=xstate_row,
                     smem_active=smem_active,
                     s_active_cnt=s_active_cnt,
                 )
@@ -4688,6 +4705,7 @@ def _run_phases(
         block_max_row=None,  # block-skip: this row's per-32-position bounds
         seed_thr_row=None,  # ext counts: this row's 3 seed thresholds (fp32)
         seed_counts_row=None,  # ext counts: this row's 3 exact counts (int32)
+        xstate_row=None,  # emit_xstate: this row's [8] fp32 state slot
         smem_active=None,
         s_active_cnt=None,
     ):
@@ -4767,6 +4785,8 @@ def _run_phases(
                         if cutlass.const_expr(self.return_output_values):
                             output_values_row[je] = input_row[je]
                         je = je + cutlass.Int32(1)
+                    if cutlass.const_expr(self.emit_xstate):
+                        xstate_row[0] = cutlass.Float32(0.0)  # degenerate
             else:
                 # cs>1: all cluster CTAs enter _run_phases; only leader writes.
                 if is_leader & (tidx == cutlass.Int32(0)):
@@ -4779,6 +4799,8 @@ def _run_phases(
                         if cutlass.const_expr(self.return_output_values):
                             output_values_row[je] = input_row[je]
                         je = je + cutlass.Int32(1)
+                    if cutlass.const_expr(self.emit_xstate):
+                        xstate_row[0] = cutlass.Float32(0.0)  # degenerate
         else:
             # Stage this CTA's slice into SMEM once before Phase 2's
             # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
@@ -5202,6 +5224,19 @@ def _run_phases(
                         warp_id,
                         lane,
                     )
+                if cutlass.const_expr(self.emit_xstate):
+                    # Closed-loop state (interface v2): [0] valid, [1] kth
+                    # proxy (= accepted threshold; the tie-fill makes it a
+                    # tight lower bound of the true kth), [2] accepted
+                    # threshold, [3] cand_count. The next step derives its
+                    # seed rung group from these.
+                    if tidx == cutlass.Int32(0):
+                        xstate_row[0] = cutlass.Float32(1.0)
+                        xstate_row[1] = s_thr[0]
+                        xstate_row[2] = s_thr[0]
+                        # cand_count_p4 = pre-P4 snapshot (P4 repurposes
+                        # the s_iscalars slots).
+                        xstate_row[3] = cutlass.Float32(cand_count_p4)
             else:
                 # cs>1: only the leader (CTA 0 in cluster) runs Phase 4.
                 if is_leader:
@@ -5305,6 +5340,7 @@ def __call__(
         block_max: cute.Tensor = None,  # block-skip bounds; None = disabled
         seed_thr: cute.Tensor = None,  # [num_rows, 3] fp32 (ext counts)
         seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 (ext counts)
+        xstate: cute.Tensor = None,  # [num_rows, 8] fp32 (emit_xstate)
     ):
         num_rows = input_data.shape[0]
         cluster_size = cutlass.const_expr(self.cluster_size)
@@ -5329,6 +5365,7 @@ def __call__(
             block_max,
             seed_thr,
             seed_counts,
+            xstate,
         ).launch(
             grid=(total_ctas, 1, 1),
             block=(self.num_threads, 1, 1),
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index d3a8ab64970f..b490842a7aa2 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -66,6 +66,7 @@ def _compile(
     p2_warp_redundant: bool = True,
     enable_block_skip: bool = False,
     use_ext_counts: bool = False,
+    emit_xstate: bool = False,
 ):
     """JIT-compile the GVR kernel for a specific knob combination.
 
@@ -148,6 +149,13 @@ def _compile(
         if use_ext_counts
         else None
     )
+    xstate_fake = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Float32, (n_rows, 8), stride_order=(1, 0), assumed_align=4
+        )
+        if emit_xstate
+        else None
+    )
     fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
     kernel = GvrTopKKernel(
         dtype=cute_dtype,
@@ -168,6 +176,7 @@ def _compile(
         p2_warp_redundant=p2_warp_redundant,
         enable_block_skip=enable_block_skip,
         use_ext_counts=use_ext_counts,
+        emit_xstate=emit_xstate,
         # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed. The
         # qfrac VALUES are irrelevant on this path (P1b is skipped) — only
         # the slot count matters.
@@ -185,6 +194,7 @@ def _compile(
         block_max=block_max_fake,
         seed_thr=seed_thr_fake,
         seed_counts=seed_counts_fake,
+        xstate=xstate_fake,
         options="--enable-tvm-ffi",
     )
 
@@ -439,6 +449,7 @@ def gvr_topk_decode(
     skip_min_n: Optional[int] = 200_000,
     seed_thr: Optional[torch.Tensor] = None,
     seed_counts: Optional[torch.Tensor] = None,
+    xstate: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """CuTe DSL GVR Top-K wrapper with every tuning knob exposed.
 
@@ -535,6 +546,14 @@ def gvr_topk_decode(
             and seed_counts.is_contiguous()
             and seed_counts.shape == (num_rows, 3)
         ), "seed_counts must be contiguous CUDA int32 [num_rows, 3]"
+    emit_xstate = xstate is not None
+    if emit_xstate:
+        assert (
+            xstate.dtype == torch.float32
+            and xstate.is_cuda
+            and xstate.is_contiguous()
+            and xstate.shape == (num_rows, 8)
+        ), "xstate must be contiguous CUDA fp32 [num_rows, 8]"
     enable_block_skip = block_max is not None
     if enable_block_skip:
         assert (
@@ -614,6 +633,7 @@ def gvr_topk_decode(
         p2_warp_redundant,
         enable_block_skip,
         use_ext_counts,
+        emit_xstate,
     )
     # When return_output_values=False the kernel was compiled to skip
     # STG.value and accepts None for the value-output slot.
@@ -629,6 +649,7 @@ def gvr_topk_decode(
         block_max if enable_block_skip else None,
         seed_thr if use_ext_counts else None,
         seed_counts if use_ext_counts else None,
+        xstate if emit_xstate else None,
     )
     if return_output_values:
         return out_values, out_indices

From f9428fc86e613b3fdb3e71edb37a4bb10ed18518 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 11:01:27 -0700
Subject: [PATCH 013/117] [None][feat] GVR: waterfall L2 direct-to-P4 from
 pre-collected pairs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

use_ext_cand: epilogue-collected (value, index) pairs land straight in
smem_keys/vals via a cooperative sentinel-skipping SMEM-atomic load —
no P1, no counting, no Phase-3 scan. Eligibility (void == 0, claimed
<= cand_cap, collect rung count in [K, kC]) is a CTA-uniform register
predicate, so the dynamic skip of the P2/P3 slab stays convergent;
ineligible rows fall through to the ext-counts path unchanged.

Real data: 6 cells x {ext, l2, forced-void fallback} all exact. The
direct path makes top-k O(cand_count), independent of N: eligible rows
cost ~12.2us warm from 16k through 1M (flash 1M: 2.84x vs the ext
count path, below even the cold skipR0 15.6us). With the 0.89-0.97
chain in-band rates, ~90% of production rows hit this floor when the
epilogue emits cand.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 757 ++++++++++--------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  44 +
 2 files changed, 470 insertions(+), 331 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index f31251f65efd..36cd79cde247 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -272,6 +272,9 @@ def __init__(
         p2_warp_redundant: bool = True,
         enable_block_skip: bool = False,
         use_ext_counts: bool = False,
+        use_ext_cand: bool = False,
+        cand_cap: int = 5120,
+        cand_rung: int = 1,
         emit_xstate: bool = False,
     ):
         # Redundant-warp sync reduction: every warp replays the block
@@ -552,6 +555,16 @@ def __init__(
         self.emit_xstate = bool(emit_xstate)
         if emit_xstate and cluster_size != 1:
             raise ValueError("emit_xstate is single-CTA (cs==1) only")
+        # use_ext_cand (waterfall L2 direct-to-P4): pre-collected (value,
+        # index) pairs from the epilogue land straight in smem_keys/vals —
+        # no P1, no counting, no P3 scan. Eligible when void==0, claimed
+        # <= cand_cap and the collect rung's exact count is in [K, kC];
+        # ineligible rows fall through to the ext-counts path.
+        self.use_ext_cand = bool(use_ext_cand)
+        self.cand_cap = int(cand_cap)
+        self.cand_rung = int(cand_rung)
+        if use_ext_cand and not use_ext_counts:
+            raise ValueError("use_ext_cand requires use_ext_counts")
         self.use_ext_counts = bool(use_ext_counts) and bool(enable_r0)
         if self.use_ext_counts:
             if not self.fb_fix:
@@ -4127,6 +4140,8 @@ def gvr_topk_kernel(
         seed_thr: cute.Tensor,  # [numRows, 3] fp32 (or None: no ext counts)
         seed_counts: cute.Tensor,  # [numRows, 3] int32 (or None)
         xstate: cute.Tensor,  # [numRows, 8] fp32 closed-loop state (or None)
+        cand: cute.Tensor,  # [numRows, CAP*2] int32 pairs (or None)
+        cand_ctl: cute.Tensor,  # [numRows, 2] int32 {claimed, void} (or None)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
 
@@ -4182,6 +4197,8 @@ def gvr_topk_kernel(
             seed_thr=seed_thr,
             seed_counts=seed_counts,
             xstate=xstate,
+            cand=cand,
+            cand_ctl=cand_ctl,
         )
 
     @cute.jit
@@ -4197,6 +4214,8 @@ def run_one_row(
         seed_thr: cute.Tensor = None,  # [numRows, 3] fp32 (ext counts)
         seed_counts: cute.Tensor = None,  # [numRows, 3] int32 (ext counts)
         xstate: cute.Tensor = None,  # [numRows, 8] fp32 (emit_xstate)
+        cand: cute.Tensor = None,  # [numRows, CAP*2] int32 (ext cand)
+        cand_ctl: cute.Tensor = None,  # [numRows, 2] int32 (ext cand)
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
 
@@ -4269,6 +4288,12 @@ def run_one_row(
             xstate_row = xstate[row_idx, None]
         else:
             xstate_row = None
+        if cutlass.const_expr(self.use_ext_cand and cand is not None and cand_ctl is not None):
+            cand_row = cand[row_idx, None]
+            cand_ctl_row = cand_ctl[row_idx, None]
+        else:
+            cand_row = None
+            cand_ctl_row = None
         # When return_output_values=False, ``output_values`` is None at
         # launch and the gated writes below are compiled out; slicing into
         # None would crash so we keep the view None as well.
@@ -4568,6 +4593,8 @@ def run_one_row(
                         seed_thr_row=seed_thr_row,
                         seed_counts_row=seed_counts_row,
                         xstate_row=xstate_row,
+                        cand_row=cand_row,
+                        cand_ctl_row=cand_ctl_row,
                         smem_active=smem_active,
                         s_active_cnt=s_active_cnt,
                     )
@@ -4615,6 +4642,8 @@ def run_one_row(
                             seed_thr_row=seed_thr_row,
                             seed_counts_row=seed_counts_row,
                             xstate_row=xstate_row,
+                            cand_row=cand_row,
+                            cand_ctl_row=cand_ctl_row,
                             smem_active=smem_active,
                             s_active_cnt=s_active_cnt,
                         )
@@ -4659,6 +4688,8 @@ def run_one_row(
                     seed_thr_row=seed_thr_row,
                     seed_counts_row=seed_counts_row,
                     xstate_row=xstate_row,
+                    cand_row=cand_row,
+                    cand_ctl_row=cand_ctl_row,
                     smem_active=smem_active,
                     s_active_cnt=s_active_cnt,
                 )
@@ -4706,6 +4737,8 @@ def _run_phases(
         seed_thr_row=None,  # ext counts: this row's 3 seed thresholds (fp32)
         seed_counts_row=None,  # ext counts: this row's 3 exact counts (int32)
         xstate_row=None,  # emit_xstate: this row's [8] fp32 state slot
+        cand_row=None,  # ext cand: this row's [CAP*2] int32 pairs
+        cand_ctl_row=None,  # ext cand: this row's [2] int32 {claimed, void}
         smem_active=None,
         s_active_cnt=None,
     ):
@@ -4802,304 +4835,384 @@ def _run_phases(
                     if cutlass.const_expr(self.emit_xstate):
                         xstate_row[0] = cutlass.Float32(0.0)  # degenerate
         else:
-            # Stage this CTA's slice into SMEM once before Phase 2's
-            # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
-            # scatter-loads OUTSIDE this slice, so it stays on GMEM.
-            if cutlass.const_expr(self.enable_smem_cache):
-                self.load_slice_to_smem(
-                    input_row,
-                    slice_start,
-                    slice_end,
-                    smem_input,
-                    tidx,
-                )
-
-            # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast
-            # path) or the secant threshold search ----
-            # enable_r0 gates to cluster_size==1 for now: op#26's R0 scans the
-            # full row in one CTA. The slice-parallel + cluster count-merge
-            # variant that lets R0 cover the cs>1 long-row branch lands in a
-            # later commit; until then cs>1 keeps the secant path.
-            if cutlass.const_expr(self.enable_r0):
-                # P1b rung placement -> ONE M-ary R0 count pass -> accept the
-                # tightest rung with count in [K, kC]. On a miss, fall back to
-                # the inline log-falsi R1 shot / fb_fix refine. At cs>1 each
-                # CTA scans its slice and block_count_ge_multi cluster-merges
-                # the rung counts (phase1b rungs are per-CTA identical since
-                # preIdx stats are full-row).
-                if cutlass.const_expr(self.use_ext_counts):
-                    # ---- Waterfall L1 admission (ext rungs, v2a) ----
-                    # Rung thresholds arrive from the indexer epilogue:
-                    # ONLY P1b is skipped. The stock M-ary count pass runs
-                    # on the ext rungs so the block-skip list build, rung
-                    # tightening, per-thread hand-off and classify all
-                    # compose unchanged (v1 routed through the dense
-                    # refine and forfeited the compact-walk win: flash 1M
-                    # ext 34.7us vs skipR0 15.6us cold).
-                    # v2b: when an ext count is already in [K, kC], park
-                    # THE ADMITTED THRESHOLD IN ALL RUNG SLOTS — the M-ary
-                    # pass degenerates to one compact single-threshold
-                    # count (+ list build at that threshold) and classify
-                    # admits it; a full miss keeps the 3 distinct rungs as
-                    # measured brackets for the seeded refine.
+            # ---- Waterfall L2: direct-to-P4 from pre-collected pairs ----
+            # Eligibility is a CTA-uniform register predicate (all threads
+            # read the same gmem control words), so the dynamic branches
+            # below (with barriers inside) stay convergent.
+            take_cand = cutlass.Int32(0)
+            if cutlass.const_expr(self.use_ext_cand):
+                claimed_c = cutlass.Int32(cand_ctl_row[0])
+                void_c = cutlass.Int32(cand_ctl_row[1])
+                ccnt_c = cutlass.Int32(seed_counts_row[self.cand_rung])
+                if (
+                    void_c == cutlass.Int32(0)
+                    and claimed_c <= cutlass.Int32(self.cand_cap)
+                    and ccnt_c >= cutlass.Int32(self.top_k)
+                    and ccnt_c <= cutlass.Int32(self.kC)
+                ):
+                    take_cand = cutlass.Int32(1)
+                if take_cand == cutlass.Int32(1):
+                    # Cooperative sentinel-skipping load of (value, index)
+                    # pairs into the P4 candidate arrays. SMEM-atomic
+                    # compaction: order is irrelevant to rank-scatter.
                     if tidx == cutlass.Int32(0):
-                        bx_m = cutlass.Int32(-1)
-                        bx_c = cutlass.Int32(2147483647)
-                        for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                            cx = cutlass.Int32(seed_counts_row[m])
-                            if (
-                                cx >= cutlass.Int32(self.top_k)
-                                and cx <= cutlass.Int32(self.kC)
-                                and cx < bx_c
-                            ):
-                                bx_m = cutlass.Int32(m)
-                                bx_c = cx
-                        for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                            if bx_m >= cutlass.Int32(0):
-                                s_mt_thr[m] = seed_thr_row[bx_m]
-                            else:
-                                s_mt_thr[m] = seed_thr_row[m]
+                        s_iscalars[0] = cutlass.Int32(0)
+                        s_thr[0] = seed_thr_row[self.cand_rung]
+                        s_iscalars[1] = cutlass.Int32(1)  # done: no retry
                     cute.arch.barrier()
-                if cutlass.const_expr(not self.use_ext_counts):
-                    if cutlass.const_expr(self.p1b_cache):
-                        # rungs from the SMEM gather-cache P1 stashed (no 2nd
-                        # GMEM gather); 16-bit only.
-                        self.phase1b_hspace_rungs_cached(
-                            pre_idx_count,
-                            smem_gath,
-                            smem_hist,
-                            s_thr,
-                            s_mt_thr,
-                            tidx,
-                            warp_id,
-                            lane,
-                        )
-                    else:
-                        self.phase1b_hspace_rungs(
-                            input_row,
-                            N,
-                            pre_idx_row,
-                            pre_idx_count,
-                            pre_idx_offset,
-                            smem_hist,
-                            s_thr,
-                            s_mt_thr,
-                            tidx,
-                            warp_id,
-                            lane,
+                    cbase = cand_row.iterator.toint()
+                    i_c = tidx
+                    while i_c < claimed_c:
+                        pa = cbase + cutlass.Int64(i_c) * cutlass.Int64(8)
+                        ip_c = cute.make_ptr(
+                            cutlass.Int32,
+                            pa + cutlass.Int64(4),
+                            cute.AddressSpace.gmem,
+                            assumed_align=4,
                         )
-                self.block_count_ge_multi(
-                    input_row,
-                    slice_start,
-                    slice_end,
-                    s_mt_thr,
-                    smem_ptcnt_multi,
-                    smem_wcnt_multi,
-                    s_mt_cnt,
-                    s_cluster_partial_m,
-                    do_cluster_sync,
-                    tidx,
-                    warp_id,
-                    lane,
-                    smem_ptcnt=smem_ptcnt,
-                    block_max_row=block_max_row,
-                    smem_active=smem_active,
-                    s_active_cnt=s_active_cnt,
-                )
-                cute.arch.barrier()
-                if tidx == 0:
-                    # tightest admissible rung = SMALLEST count in [K, kC].
-                    # (Explicit argmin: with r0_vseed the pmean column is not
-                    # sorted into the rung order; for sorted rungs this is
-                    # equivalent to the old "last m in window" rule.)
-                    # Dropped rungs (block-skip rung tightening) hold PARTIAL
-                    # counts — never admissible.
-                    dmask_c = cutlass.Int32(0)
-                    if cutlass.const_expr(self.enable_block_skip):
-                        dmask_c = s_active_cnt[2]
-                    best_m = cutlass.Int32(-1)
-                    best_c = cutlass.Int32(2147483647)
-                    for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                        cm = s_mt_cnt[m]
-                        if (
-                            cm >= cutlass.Int32(self.top_k)
-                            and cm <= cutlass.Int32(self.kC)
-                            and cm < best_c
-                            and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m)))
-                            == cutlass.Int32(0)
-                        ):
-                            best_m = cutlass.Int32(m)
-                            best_c = cm
-                    s_r0col[0] = best_m
-                    if best_m >= cutlass.Int32(0):
-                        s_thr[0] = s_mt_thr[best_m]
-                        s_iscalars[0] = s_mt_cnt[best_m]
-                        # done=1: the threshold is admitted, so Phase 3 must
-                        # SKIP its retry-shrink and honor s_thr[0]. (block_count
-                        # _ge / secant leave done via their own path; the R0
-                        # admission must set it explicitly or Phase 3 re-searches
-                        # and the cluster collect diverges -> wrong output.)
-                        s_iscalars[1] = cutlass.Int32(1)
-                        # Snapshot this CTA's LOCAL slice count for the chosen
-                        # rung into s_iscalars[5] — the per-CTA cand_count that
-                        # Phase 3/4's cluster gather consumes (block_count_ge
-                        # sets it too; the R0 admission must match). Without it
-                        # the cluster collect under-counts -> wrong output.
-                        if cutlass.const_expr(cluster_size > 1):
-                            s_iscalars[5] = s_cluster_partial_m[best_m]
+                        pidx = cute.make_tensor(ip_c, cute.make_layout((1,)))[0]
+                        if pidx >= cutlass.Int32(0):
+                            vp_c = cute.make_ptr(
+                                cutlass.Float32,
+                                pa,
+                                cute.AddressSpace.gmem,
+                                assumed_align=8,
+                            )
+                            pval = cute.make_tensor(vp_c, cute.make_layout((1,)))[0]
+                            wpos = atomicAdd(
+                                s_iscalars.iterator + cutlass.Int32(0),
+                                cutlass.Int32(1),
+                            )
+                            if wpos < cutlass.Int32(self.kC):
+                                smem_keys[wpos] = pval
+                                smem_vals[wpos] = pidx
+                        i_c = i_c + cutlass.Int32(num_threads)
+                    cute.arch.barrier()
+            if take_cand == cutlass.Int32(0):
+                # Stage this CTA's slice into SMEM once before Phase 2's
+                # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
+                # scatter-loads OUTSIDE this slice, so it stays on GMEM.
+                if cutlass.const_expr(self.enable_smem_cache):
+                    self.load_slice_to_smem(
+                        input_row,
+                        slice_start,
+                        slice_end,
+                        smem_input,
+                        tidx,
+                    )
 
-                cute.arch.barrier()
-                bc = s_r0col[0]
-                if bc >= cutlass.Int32(0) and bc < cutlass.Int32(self.M_qf):
-                    # accepted rung column: copy its cached per-thread counts
-                    # into the secant hand-off buffer (zero rescan). The vseed
-                    # column (bc == M_qf) is ALREADY in smem_ptcnt (v3 reuse).
-                    smem_ptcnt[tidx] = smem_ptcnt_multi[bc * cutlass.Int32(num_threads) + tidx]
-                cute.arch.barrier()
-                # ---- R0 miss: SEEDED bounded log-falsi refine ----
-                # At large N the M2D rungs straddle [K, kC]; the refine must
-                # find a threshold with count in [K, kC] between the measured
-                # rungs. SEED the loop with the rung bracket AND its known
-                # counts (clo/chi) so it does log-count regula-falsi from
-                # iter 0 with no re-measure and no separate R1 shot -> ~2-3
-                # count passes (op#26 efficiency) instead of ~6. done=1 on
-                # accept so Phase 3 skips its retry-shrink.
-                if bc < cutlass.Int32(0):
-                    if cutlass.const_expr(self.enable_block_skip):
-                        if tidx == cutlass.Int32(0):
-                            s_active_cnt[1] = cutlass.Int32(0)
-                    if cutlass.const_expr(self.fb_fix):
+                # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast
+                # path) or the secant threshold search ----
+                # enable_r0 gates to cluster_size==1 for now: op#26's R0 scans the
+                # full row in one CTA. The slice-parallel + cluster count-merge
+                # variant that lets R0 cover the cs>1 long-row branch lands in a
+                # later commit; until then cs>1 keeps the secant path.
+                if cutlass.const_expr(self.enable_r0):
+                    # P1b rung placement -> ONE M-ary R0 count pass -> accept the
+                    # tightest rung with count in [K, kC]. On a miss, fall back to
+                    # the inline log-falsi R1 shot / fb_fix refine. At cs>1 each
+                    # CTA scans its slice and block_count_ge_multi cluster-merges
+                    # the rung counts (phase1b rungs are per-CTA identical since
+                    # preIdx stats are full-row).
+                    if cutlass.const_expr(self.use_ext_counts):
+                        # ---- Waterfall L1 admission (ext rungs, v2a) ----
+                        # Rung thresholds arrive from the indexer epilogue:
+                        # ONLY P1b is skipped. The stock M-ary count pass runs
+                        # on the ext rungs so the block-skip list build, rung
+                        # tightening, per-thread hand-off and classify all
+                        # compose unchanged (v1 routed through the dense
+                        # refine and forfeited the compact-walk win: flash 1M
+                        # ext 34.7us vs skipR0 15.6us cold).
+                        # v2b: when an ext count is already in [K, kC], park
+                        # THE ADMITTED THRESHOLD IN ALL RUNG SLOTS — the M-ary
+                        # pass degenerates to one compact single-threshold
+                        # count (+ list build at that threshold) and classify
+                        # admits it; a full miss keeps the 3 distinct rungs as
+                        # measured brackets for the seeded refine.
                         if tidx == cutlass.Int32(0):
-                            M = cutlass.const_expr(self.M_thr)
-                            blo = v_lo
-                            bhi = v_hi
-                            clo = cutlass.Int32(-1)
-                            chi = cutlass.Int32(-1)
-                            dmask_f = cutlass.Int32(0)
-                            if cutlass.const_expr(self.enable_block_skip):
-                                dmask_f = s_active_cnt[2]
-                            for m in cutlass.range_constexpr(M):
-                                cm = s_mt_cnt[m]
-                                tm = s_mt_thr[m]
-                                m_ok = (
-                                    dmask_f & (cutlass.Int32(1) << cutlass.Int32(m))
-                                ) == cutlass.Int32(0)
-                                if (
-                                    m_ok
-                                    and cm > cutlass.Int32(self.kC)
-                                    and (clo < cutlass.Int32(0) or tm > blo)
-                                ):
-                                    blo = tm
-                                    clo = cm
+                            bx_m = cutlass.Int32(-1)
+                            bx_c = cutlass.Int32(2147483647)
+                            for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                                cx = cutlass.Int32(seed_counts_row[m])
                                 if (
-                                    m_ok
-                                    and cm < cutlass.Int32(self.top_k)
-                                    and (chi < cutlass.Int32(0) or tm < bhi)
+                                    cx >= cutlass.Int32(self.top_k)
+                                    and cx <= cutlass.Int32(self.kC)
+                                    and cx < bx_c
                                 ):
-                                    bhi = tm
-                                    chi = cm
-                            s_thr[1] = blo
-                            s_thr[2] = bhi
-                            s_iscalars[2] = clo  # SEED known rung counts
-                            s_iscalars[3] = chi
-                            s_iscalars[1] = cutlass.Int32(0)  # done=0
-                            cand = (blo + bhi) * cutlass.Float32(0.5)
-                            if clo > cutlass.Int32(0) and chi >= cutlass.Int32(0):
-                                chic = chi
-                                if chic < cutlass.Int32(1):
-                                    chic = cutlass.Int32(1)
-                                l_lo = cmath.log2(cutlass.Float32(clo), fastmath=True)
-                                l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True)
-                                den = l_lo - l_hi
-                                if den > cutlass.Float32(0.0):
-                                    t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den
-                                    cnd3 = bhi + t3 * (blo - bhi)
-                                    if cnd3 > blo and cnd3 < bhi:
-                                        cand = cnd3
-                            elif chi < cutlass.Int32(0):
-                                cand = bhi
-                            elif clo < cutlass.Int32(0):
-                                cand = blo
-                            s_thr[0] = cand
+                                    bx_m = cutlass.Int32(m)
+                                    bx_c = cx
+                            for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                                if bx_m >= cutlass.Int32(0):
+                                    s_mt_thr[m] = seed_thr_row[bx_m]
+                                else:
+                                    s_mt_thr[m] = seed_thr_row[m]
                         cute.arch.barrier()
-                        rs = cutlass.Int32(0)
-                        while rs < cutlass.Int32(8) and s_iscalars[1] == cutlass.Int32(0):
-                            if rs > cutlass.Int32(0):
-                                if tidx == cutlass.Int32(0):
-                                    lo3 = s_thr[1]
-                                    hi3 = s_thr[2]
-                                    clo3 = s_iscalars[2]
-                                    chi3 = s_iscalars[3]
-                                    cand = (lo3 + hi3) * cutlass.Float32(0.5)
-                                    if chi3 < cutlass.Int32(0):
-                                        cand = hi3
-                                    elif clo3 < cutlass.Int32(0):
-                                        cand = lo3
-                                    else:
-                                        chic = chi3
-                                        if chic < cutlass.Int32(1):
-                                            chic = cutlass.Int32(1)
-                                        l_lo = cmath.log2(cutlass.Float32(clo3), fastmath=True)
-                                        l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True)
-                                        den3 = l_lo - l_hi
-                                        if den3 > cutlass.Float32(0.0):
-                                            t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den3
-                                            cnd3 = hi3 + t3 * (lo3 - hi3)
-                                            if cnd3 > lo3 and cnd3 < hi3:
-                                                cand = cnd3
-                                    s_thr[0] = cand
-                                cute.arch.barrier()
-                            self.block_count_ge(
+                    if cutlass.const_expr(not self.use_ext_counts):
+                        if cutlass.const_expr(self.p1b_cache):
+                            # rungs from the SMEM gather-cache P1 stashed (no 2nd
+                            # GMEM gather); 16-bit only.
+                            self.phase1b_hspace_rungs_cached(
+                                pre_idx_count,
+                                smem_gath,
+                                smem_hist,
+                                s_thr,
+                                s_mt_thr,
+                                tidx,
+                                warp_id,
+                                lane,
+                            )
+                        else:
+                            self.phase1b_hspace_rungs(
                                 input_row,
-                                slice_start,
-                                slice_end,
-                                s_thr[0],
-                                smem_ptcnt,
-                                smem_wcnt,
-                                s_iscalars,
-                                s_cluster_partial,
+                                N,
+                                pre_idx_row,
+                                pre_idx_count,
+                                pre_idx_offset,
+                                smem_hist,
+                                s_thr,
+                                s_mt_thr,
                                 tidx,
                                 warp_id,
                                 lane,
-                                do_cluster_sync=do_cluster_sync,
-                                smem_input=smem_input,
                             )
-                            cute.arch.barrier()
+                    self.block_count_ge_multi(
+                        input_row,
+                        slice_start,
+                        slice_end,
+                        s_mt_thr,
+                        smem_ptcnt_multi,
+                        smem_wcnt_multi,
+                        s_mt_cnt,
+                        s_cluster_partial_m,
+                        do_cluster_sync,
+                        tidx,
+                        warp_id,
+                        lane,
+                        smem_ptcnt=smem_ptcnt,
+                        block_max_row=block_max_row,
+                        smem_active=smem_active,
+                        s_active_cnt=s_active_cnt,
+                    )
+                    cute.arch.barrier()
+                    if tidx == 0:
+                        # tightest admissible rung = SMALLEST count in [K, kC].
+                        # (Explicit argmin: with r0_vseed the pmean column is not
+                        # sorted into the rung order; for sorted rungs this is
+                        # equivalent to the old "last m in window" rule.)
+                        # Dropped rungs (block-skip rung tightening) hold PARTIAL
+                        # counts — never admissible.
+                        dmask_c = cutlass.Int32(0)
+                        if cutlass.const_expr(self.enable_block_skip):
+                            dmask_c = s_active_cnt[2]
+                        best_m = cutlass.Int32(-1)
+                        best_c = cutlass.Int32(2147483647)
+                        for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                            cm = s_mt_cnt[m]
+                            if (
+                                cm >= cutlass.Int32(self.top_k)
+                                and cm <= cutlass.Int32(self.kC)
+                                and cm < best_c
+                                and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m)))
+                                == cutlass.Int32(0)
+                            ):
+                                best_m = cutlass.Int32(m)
+                                best_c = cm
+                        s_r0col[0] = best_m
+                        if best_m >= cutlass.Int32(0):
+                            s_thr[0] = s_mt_thr[best_m]
+                            s_iscalars[0] = s_mt_cnt[best_m]
+                            # done=1: the threshold is admitted, so Phase 3 must
+                            # SKIP its retry-shrink and honor s_thr[0]. (block_count
+                            # _ge / secant leave done via their own path; the R0
+                            # admission must set it explicitly or Phase 3 re-searches
+                            # and the cluster collect diverges -> wrong output.)
+                            s_iscalars[1] = cutlass.Int32(1)
+                            # Snapshot this CTA's LOCAL slice count for the chosen
+                            # rung into s_iscalars[5] — the per-CTA cand_count that
+                            # Phase 3/4's cluster gather consumes (block_count_ge
+                            # sets it too; the R0 admission must match). Without it
+                            # the cluster collect under-counts -> wrong output.
+                            if cutlass.const_expr(cluster_size > 1):
+                                s_iscalars[5] = s_cluster_partial_m[best_m]
+
+                    cute.arch.barrier()
+                    bc = s_r0col[0]
+                    if bc >= cutlass.Int32(0) and bc < cutlass.Int32(self.M_qf):
+                        # accepted rung column: copy its cached per-thread counts
+                        # into the secant hand-off buffer (zero rescan). The vseed
+                        # column (bc == M_qf) is ALREADY in smem_ptcnt (v3 reuse).
+                        smem_ptcnt[tidx] = smem_ptcnt_multi[bc * cutlass.Int32(num_threads) + tidx]
+                    cute.arch.barrier()
+                    # ---- R0 miss: SEEDED bounded log-falsi refine ----
+                    # At large N the M2D rungs straddle [K, kC]; the refine must
+                    # find a threshold with count in [K, kC] between the measured
+                    # rungs. SEED the loop with the rung bracket AND its known
+                    # counts (clo/chi) so it does log-count regula-falsi from
+                    # iter 0 with no re-measure and no separate R1 shot -> ~2-3
+                    # count passes (op#26 efficiency) instead of ~6. done=1 on
+                    # accept so Phase 3 skips its retry-shrink.
+                    if bc < cutlass.Int32(0):
+                        if cutlass.const_expr(self.enable_block_skip):
                             if tidx == cutlass.Int32(0):
-                                c3 = s_iscalars[0]
-                                t3v = s_thr[0]
-                                if c3 >= cutlass.Int32(self.top_k) and c3 <= cutlass.Int32(self.kC):
-                                    s_iscalars[1] = cutlass.Int32(1)  # accept
-                                elif c3 > cutlass.Int32(self.kC):
-                                    s_thr[1] = t3v
-                                    s_iscalars[2] = c3
-                                    if t3v >= s_thr[2]:
-                                        rng3 = s_thr[2] - s_thr[1]
-                                        if rng3 < cutlass.Float32(1.0):
-                                            rng3 = cutlass.Float32(1.0)
-                                        s_thr[2] = s_thr[2] + rng3 * cutlass.Float32(8.0)
-                                        s_iscalars[3] = cutlass.Int32(-1)
-                                else:
-                                    s_thr[2] = t3v
-                                    s_iscalars[3] = c3
-                                    if t3v <= s_thr[1]:
-                                        rng3 = s_thr[2] - s_thr[1]
-                                        if rng3 < cutlass.Float32(1.0):
-                                            rng3 = cutlass.Float32(1.0)
-                                        s_thr[1] = s_thr[1] - rng3 * cutlass.Float32(8.0)
-                                        s_iscalars[2] = cutlass.Int32(-1)
+                                s_active_cnt[1] = cutlass.Int32(0)
+                        if cutlass.const_expr(self.fb_fix):
+                            if tidx == cutlass.Int32(0):
+                                M = cutlass.const_expr(self.M_thr)
+                                blo = v_lo
+                                bhi = v_hi
+                                clo = cutlass.Int32(-1)
+                                chi = cutlass.Int32(-1)
+                                dmask_f = cutlass.Int32(0)
+                                if cutlass.const_expr(self.enable_block_skip):
+                                    dmask_f = s_active_cnt[2]
+                                for m in cutlass.range_constexpr(M):
+                                    cm = s_mt_cnt[m]
+                                    tm = s_mt_thr[m]
+                                    m_ok = (
+                                        dmask_f & (cutlass.Int32(1) << cutlass.Int32(m))
+                                    ) == cutlass.Int32(0)
+                                    if (
+                                        m_ok
+                                        and cm > cutlass.Int32(self.kC)
+                                        and (clo < cutlass.Int32(0) or tm > blo)
+                                    ):
+                                        blo = tm
+                                        clo = cm
+                                    if (
+                                        m_ok
+                                        and cm < cutlass.Int32(self.top_k)
+                                        and (chi < cutlass.Int32(0) or tm < bhi)
+                                    ):
+                                        bhi = tm
+                                        chi = cm
+                                s_thr[1] = blo
+                                s_thr[2] = bhi
+                                s_iscalars[2] = clo  # SEED known rung counts
+                                s_iscalars[3] = chi
+                                s_iscalars[1] = cutlass.Int32(0)  # done=0
+                                cand = (blo + bhi) * cutlass.Float32(0.5)
+                                if clo > cutlass.Int32(0) and chi >= cutlass.Int32(0):
+                                    chic = chi
+                                    if chic < cutlass.Int32(1):
+                                        chic = cutlass.Int32(1)
+                                    l_lo = cmath.log2(cutlass.Float32(clo), fastmath=True)
+                                    l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True)
+                                    den = l_lo - l_hi
+                                    if den > cutlass.Float32(0.0):
+                                        t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den
+                                        cnd3 = bhi + t3 * (blo - bhi)
+                                        if cnd3 > blo and cnd3 < bhi:
+                                            cand = cnd3
+                                elif chi < cutlass.Int32(0):
+                                    cand = bhi
+                                elif clo < cutlass.Int32(0):
+                                    cand = blo
+                                s_thr[0] = cand
                             cute.arch.barrier()
-                            rs = rs + cutlass.Int32(1)
-                        if s_iscalars[1] != cutlass.Int32(1):
-                            # tie-plateau fail-soft: land on the measured
-                            # undershoot side (count <= kC => no overflow).
-                            self.block_count_ge(
+                            rs = cutlass.Int32(0)
+                            while rs < cutlass.Int32(8) and s_iscalars[1] == cutlass.Int32(0):
+                                if rs > cutlass.Int32(0):
+                                    if tidx == cutlass.Int32(0):
+                                        lo3 = s_thr[1]
+                                        hi3 = s_thr[2]
+                                        clo3 = s_iscalars[2]
+                                        chi3 = s_iscalars[3]
+                                        cand = (lo3 + hi3) * cutlass.Float32(0.5)
+                                        if chi3 < cutlass.Int32(0):
+                                            cand = hi3
+                                        elif clo3 < cutlass.Int32(0):
+                                            cand = lo3
+                                        else:
+                                            chic = chi3
+                                            if chic < cutlass.Int32(1):
+                                                chic = cutlass.Int32(1)
+                                            l_lo = cmath.log2(cutlass.Float32(clo3), fastmath=True)
+                                            l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True)
+                                            den3 = l_lo - l_hi
+                                            if den3 > cutlass.Float32(0.0):
+                                                t3 = (
+                                                    cutlass.Float32(self.log2_mstar) - l_hi
+                                                ) / den3
+                                                cnd3 = hi3 + t3 * (lo3 - hi3)
+                                                if cnd3 > lo3 and cnd3 < hi3:
+                                                    cand = cnd3
+                                        s_thr[0] = cand
+                                    cute.arch.barrier()
+                                self.block_count_ge(
+                                    input_row,
+                                    slice_start,
+                                    slice_end,
+                                    s_thr[0],
+                                    smem_ptcnt,
+                                    smem_wcnt,
+                                    s_iscalars,
+                                    s_cluster_partial,
+                                    tidx,
+                                    warp_id,
+                                    lane,
+                                    do_cluster_sync=do_cluster_sync,
+                                    smem_input=smem_input,
+                                )
+                                cute.arch.barrier()
+                                if tidx == cutlass.Int32(0):
+                                    c3 = s_iscalars[0]
+                                    t3v = s_thr[0]
+                                    if c3 >= cutlass.Int32(self.top_k) and c3 <= cutlass.Int32(
+                                        self.kC
+                                    ):
+                                        s_iscalars[1] = cutlass.Int32(1)  # accept
+                                    elif c3 > cutlass.Int32(self.kC):
+                                        s_thr[1] = t3v
+                                        s_iscalars[2] = c3
+                                        if t3v >= s_thr[2]:
+                                            rng3 = s_thr[2] - s_thr[1]
+                                            if rng3 < cutlass.Float32(1.0):
+                                                rng3 = cutlass.Float32(1.0)
+                                            s_thr[2] = s_thr[2] + rng3 * cutlass.Float32(8.0)
+                                            s_iscalars[3] = cutlass.Int32(-1)
+                                    else:
+                                        s_thr[2] = t3v
+                                        s_iscalars[3] = c3
+                                        if t3v <= s_thr[1]:
+                                            rng3 = s_thr[2] - s_thr[1]
+                                            if rng3 < cutlass.Float32(1.0):
+                                                rng3 = cutlass.Float32(1.0)
+                                            s_thr[1] = s_thr[1] - rng3 * cutlass.Float32(8.0)
+                                            s_iscalars[2] = cutlass.Int32(-1)
+                                cute.arch.barrier()
+                                rs = rs + cutlass.Int32(1)
+                            if s_iscalars[1] != cutlass.Int32(1):
+                                # tie-plateau fail-soft: land on the measured
+                                # undershoot side (count <= kC => no overflow).
+                                self.block_count_ge(
+                                    input_row,
+                                    slice_start,
+                                    slice_end,
+                                    s_thr[2],
+                                    smem_ptcnt,
+                                    smem_wcnt,
+                                    s_iscalars,
+                                    s_cluster_partial,
+                                    tidx,
+                                    warp_id,
+                                    lane,
+                                    do_cluster_sync=do_cluster_sync,
+                                    smem_input=smem_input,
+                                )
+                                cute.arch.barrier()
+                                if tidx == cutlass.Int32(0):
+                                    s_thr[0] = s_thr[2]
+                                    s_iscalars[1] = cutlass.Int32(1)
+                                cute.arch.barrier()
+                        else:
+                            self.phase2_secant_search(
                                 input_row,
+                                N,
                                 slice_start,
                                 slice_end,
-                                s_thr[2],
                                 smem_ptcnt,
                                 smem_wcnt,
+                                s_thr,
                                 s_iscalars,
                                 s_cluster_partial,
                                 tidx,
@@ -5108,34 +5221,39 @@ def _run_phases(
                                 do_cluster_sync=do_cluster_sync,
                                 smem_input=smem_input,
                             )
-                            cute.arch.barrier()
-                            if tidx == cutlass.Int32(0):
-                                s_thr[0] = s_thr[2]
-                                s_iscalars[1] = cutlass.Int32(1)
-                            cute.arch.barrier()
-                    else:
-                        self.phase2_secant_search(
-                            input_row,
-                            N,
-                            slice_start,
-                            slice_end,
-                            smem_ptcnt,
-                            smem_wcnt,
-                            s_thr,
-                            s_iscalars,
-                            s_cluster_partial,
-                            tidx,
-                            warp_id,
-                            lane,
-                            do_cluster_sync=do_cluster_sync,
-                            smem_input=smem_input,
-                        )
-            else:
-                self.phase2_secant_search(
+                else:
+                    self.phase2_secant_search(
+                        input_row,
+                        N,
+                        slice_start,
+                        slice_end,
+                        smem_ptcnt,
+                        smem_wcnt,
+                        s_thr,
+                        s_iscalars,
+                        s_cluster_partial,
+                        tidx,
+                        warp_id,
+                        lane,
+                        do_cluster_sync=do_cluster_sync,
+                        smem_input=smem_input,
+                    )
+
+                # Cluster handoff #1 (end of Phase 2). Skipped when
+                # do_cluster_sync is False (cs=1 or short-row degrade).
+                if cutlass.const_expr(cluster_size > 1):
+                    if do_cluster_sync:
+                        cute.arch.cluster_arrive_relaxed()
+                        cute.arch.cluster_wait()
+
+                # ---- Phase 3: cluster-parallel candidate collect ----
+                self.phase3_collect_candidates(
                     input_row,
                     N,
                     slice_start,
                     slice_end,
+                    smem_keys,
+                    smem_vals,
                     smem_ptcnt,
                     smem_wcnt,
                     s_thr,
@@ -5146,37 +5264,10 @@ def _run_phases(
                     lane,
                     do_cluster_sync=do_cluster_sync,
                     smem_input=smem_input,
+                    smem_active=smem_active,
+                    s_active_cnt=s_active_cnt,
                 )
 
-            # Cluster handoff #1 (end of Phase 2). Skipped when
-            # do_cluster_sync is False (cs=1 or short-row degrade).
-            if cutlass.const_expr(cluster_size > 1):
-                if do_cluster_sync:
-                    cute.arch.cluster_arrive_relaxed()
-                    cute.arch.cluster_wait()
-
-            # ---- Phase 3: cluster-parallel candidate collect ----
-            self.phase3_collect_candidates(
-                input_row,
-                N,
-                slice_start,
-                slice_end,
-                smem_keys,
-                smem_vals,
-                smem_ptcnt,
-                smem_wcnt,
-                s_thr,
-                s_iscalars,
-                s_cluster_partial,
-                tidx,
-                warp_id,
-                lane,
-                do_cluster_sync=do_cluster_sync,
-                smem_input=smem_input,
-                smem_active=smem_active,
-                s_active_cnt=s_active_cnt,
-            )
-
             # Cluster handoff #2: leader's DSMEM gather of peer
             # smem_keys/smem_vals. Skipped at do_cluster_sync=False.
             if cutlass.const_expr(cluster_size > 1):
@@ -5341,6 +5432,8 @@ def __call__(
         seed_thr: cute.Tensor = None,  # [num_rows, 3] fp32 (ext counts)
         seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 (ext counts)
         xstate: cute.Tensor = None,  # [num_rows, 8] fp32 (emit_xstate)
+        cand: cute.Tensor = None,  # [num_rows, CAP*2] int32 (ext cand)
+        cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 (ext cand)
     ):
         num_rows = input_data.shape[0]
         cluster_size = cutlass.const_expr(self.cluster_size)
@@ -5366,6 +5459,8 @@ def __call__(
             seed_thr,
             seed_counts,
             xstate,
+            cand,
+            cand_ctl,
         ).launch(
             grid=(total_ctas, 1, 1),
             block=(self.num_threads, 1, 1),
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index b490842a7aa2..40d238c2704e 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -67,6 +67,8 @@ def _compile(
     enable_block_skip: bool = False,
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
+    use_ext_cand: bool = False,
+    cand_cap: int = 5120,
 ):
     """JIT-compile the GVR kernel for a specific knob combination.
 
@@ -149,6 +151,20 @@ def _compile(
         if use_ext_counts
         else None
     )
+    cand_fake = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Int32, (n_rows, cand_cap * 2), stride_order=(1, 0), assumed_align=8
+        )
+        if use_ext_cand
+        else None
+    )
+    cand_ctl_fake = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Int32, (n_rows, 2), stride_order=(1, 0), assumed_align=8
+        )
+        if use_ext_cand
+        else None
+    )
     xstate_fake = (
         cute.runtime.make_fake_compact_tensor(
             cutlass.Float32, (n_rows, 8), stride_order=(1, 0), assumed_align=4
@@ -177,6 +193,8 @@ def _compile(
         enable_block_skip=enable_block_skip,
         use_ext_counts=use_ext_counts,
         emit_xstate=emit_xstate,
+        use_ext_cand=use_ext_cand,
+        cand_cap=cand_cap,
         # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed. The
         # qfrac VALUES are irrelevant on this path (P1b is skipped) — only
         # the slot count matters.
@@ -195,6 +213,8 @@ def _compile(
         seed_thr=seed_thr_fake,
         seed_counts=seed_counts_fake,
         xstate=xstate_fake,
+        cand=cand_fake,
+        cand_ctl=cand_ctl_fake,
         options="--enable-tvm-ffi",
     )
 
@@ -450,6 +470,8 @@ def gvr_topk_decode(
     seed_thr: Optional[torch.Tensor] = None,
     seed_counts: Optional[torch.Tensor] = None,
     xstate: Optional[torch.Tensor] = None,
+    cand: Optional[torch.Tensor] = None,
+    cand_ctl: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """CuTe DSL GVR Top-K wrapper with every tuning knob exposed.
 
@@ -546,6 +568,24 @@ def gvr_topk_decode(
             and seed_counts.is_contiguous()
             and seed_counts.shape == (num_rows, 3)
         ), "seed_counts must be contiguous CUDA int32 [num_rows, 3]"
+    use_ext_cand = cand is not None and cand_ctl is not None
+    cand_cap = 5120
+    if use_ext_cand:
+        assert (
+            cand.dtype == torch.int32
+            and cand.is_cuda
+            and cand.is_contiguous()
+            and cand.dim() == 2
+            and cand.shape[0] == num_rows
+            and cand.shape[1] % 2 == 0
+        ), "cand must be contiguous CUDA int32 [num_rows, CAP*2]"
+        assert (
+            cand_ctl.dtype == torch.int32
+            and cand_ctl.is_cuda
+            and cand_ctl.is_contiguous()
+            and cand_ctl.shape == (num_rows, 2)
+        ), "cand_ctl must be contiguous CUDA int32 [num_rows, 2]"
+        cand_cap = cand.shape[1] // 2
     emit_xstate = xstate is not None
     if emit_xstate:
         assert (
@@ -634,6 +674,8 @@ def gvr_topk_decode(
         enable_block_skip,
         use_ext_counts,
         emit_xstate,
+        use_ext_cand,
+        cand_cap,
     )
     # When return_output_values=False the kernel was compiled to skip
     # STG.value and accepts None for the value-output slot.
@@ -650,6 +692,8 @@ def gvr_topk_decode(
         seed_thr if use_ext_counts else None,
         seed_counts if use_ext_counts else None,
         xstate if emit_xstate else None,
+        cand if use_ext_cand else None,
+        cand_ctl if use_ext_cand else None,
     )
     if return_output_values:
         return out_values, out_indices

From 076323c686988a4cc165dc47fa295bf5d70bfbeb Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 21:15:54 -0700
Subject: [PATCH 014/117] [None][perf] GVR: slope-adaptive seed rung derivation
 (host helper)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

derive_seed_rungs places the next step's guard rungs a fixed number of
count-OCTAVES from the previous accepted threshold, using the local
slope of log2(count) vs threshold estimated from the previous step's
own 3 rung measurements (log-linearity is the same property log-falsi
exploits). Fixed spreads face a two-sided trap: too narrow misses
drift, too wide puts the guard rungs themselves out of band — no
single value wins both models (best fixed: pro 0.97/flash 0.92 vs
flash-tuned 0.82/0.97).

Real-chain kernel validation (V4-Pro/Flash multi-step captures):
in-band admission pro 0.89 -> 0.99, flash 0.96, all steps exact.
Combined with the L2 direct arm this routes ~97%+ of production rows
to the O(cand_count) floor (cold protocol: flash 1M BS1024 30.9us =
12.1x, BS1 8.7us = 4.0x; pro 256k 1.5-1.8x).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 40d238c2704e..ed79a19cbb23 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -345,6 +345,44 @@ def emu_block_meta(
     return meta.contiguous()
 
 
+def derive_seed_rungs(
+    prev_thr: torch.Tensor,
+    prev_sthr: "torch.Tensor | None" = None,
+    prev_counts: "torch.Tensor | None" = None,
+    count_octaves: float = 2.0,
+    fallback_spread: float = 0.5,
+) -> torch.Tensor:
+    """Host-side slope-adaptive seed rung derivation (waterfall closed loop).
+
+    Estimates the per-row local slope of log2(count) vs threshold from the
+    PREVIOUS step's 3 rung measurements and places the next step's guard
+    rungs ``count_octaves`` octaves away from the mid rung (= the previous
+    accepted threshold). Real-chain validation (V4-Pro/Flash captures):
+    in-band admission 0.958/0.975 vs 0.82-0.97 for any fixed spread.
+
+    Args:
+        prev_thr: [rows] previous accepted threshold (xstate[:, 2]).
+        prev_sthr: [rows, 3] previous step's rung thresholds (or None).
+        prev_counts: [rows, 3] previous step's rung counts (or None).
+
+    Returns:
+        [rows, 3] fp32 seed thresholds (ascending).
+    """
+    if prev_sthr is None or prev_counts is None:
+        d = torch.full_like(prev_thr, fallback_spread)
+    else:
+        c_lo = prev_counts[:, 0].float().clamp(min=1.0)
+        c_hi = prev_counts[:, 2].float().clamp(min=1.0)
+        dthr = (prev_sthr[:, 2] - prev_sthr[:, 0]).clamp(min=1e-3)
+        slope = (torch.log2(c_lo) - torch.log2(c_hi)) / dthr
+        d = torch.where(
+            slope > 0.05,
+            count_octaves / slope.clamp(min=0.05),
+            torch.full_like(slope, fallback_spread),
+        ).clamp(0.1, 4.0)
+    return torch.stack([prev_thr - d, prev_thr, prev_thr + d], dim=1).contiguous()
+
+
 def emu_seed_counts(
     logits: torch.Tensor,
     seq_lens: torch.Tensor,

From 0e3470d57dcac6448a08d6829a85c35db7b08ec4 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 21:30:26 -0700
Subject: [PATCH 015/117] [None][perf] GVR ext counts: per-row dynamic routing
 to the stock path
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The ext knobs were compile-time, so a row whose epilogue rungs all
miss [K, kC] (or an xstate-invalid row, t_0 = +FLT_MAX) still paid
the ext bracket-refine — measurably worse than stock (pro-1M cold:
ext-miss 51us vs stock-skip 21us). Routing is now a per-row runtime
predicate read from the ext counts themselves (CTA-uniform loads, so
the dynamic branches with barriers inside stay convergent): in-band
rows keep the ext fast path (skip P1 + P1b), miss/invalid rows run
the full stock path (P1 + P1b + vseed + count) including the
block-skip machinery.

Warm validation: pro 1M ext+skip 43.5us (0.77x) -> 18.8us (1.80x);
in-band cells unchanged (flash 1M 2.40x, pro 256k 1.07x); mixed-row
chains exact with in-band 0.99/0.97 (pro/flash, adaptive rungs).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 149 +++++++++++++-----
 1 file changed, 107 insertions(+), 42 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 36cd79cde247..07a98c06c1cf 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -4763,6 +4763,22 @@ def _run_phases(
                 s_active_cnt[1] = cutlass.Int32(0)
                 s_active_cnt[2] = cutlass.Int32(0)  # dropped-rung mask
 
+        # ---- Per-row dynamic routing (ext counts) ----
+        # Use the epilogue rungs ONLY when the row is valid (finite t_0,
+        # xstate contract) AND some rung count already lies in [K, kC].
+        # A miss/invalid row runs the full stock path (P1 + P1b + vseed +
+        # count): real data shows stock beats ext-bracket refine on
+        # misses (pro-1M cold: stock-skip 21us vs ext-miss 51us). All
+        # threads read the same control words, so the predicate is
+        # CTA-uniform and the dynamic branches below stay convergent.
+        ext_row = cutlass.Int32(0)
+        if cutlass.const_expr(self.use_ext_counts):
+            if seed_thr_row[0] < cutlass.Float32(1e37):
+                for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                    cm_e = cutlass.Int32(seed_counts_row[m])
+                    if cm_e >= cutlass.Int32(self.top_k) and cm_e <= cutlass.Int32(self.kC):
+                        ext_row = cutlass.Int32(1)
+
         # ---- Phase 1: preIdx Min/Max/Mean ----
         # ext counts: P1's only surviving products are the [v_lo, v_hi]
         # outer bracket and the scalar state init — the ext rungs provide
@@ -4771,16 +4787,36 @@ def _run_phases(
         # target lies outside [t_0, t_2] recovers via the refine loop's
         # 8x bracket expansion (same fail-soft as the stock path).
         if cutlass.const_expr(self.use_ext_counts):
-            if tidx == cutlass.Int32(0):
-                s_thr[0] = seed_thr_row[1]
-                s_thr[1] = seed_thr_row[0]
-                s_thr[2] = seed_thr_row[2]
-                s_iscalars[0] = cutlass.Int32(0)  # cand_count
-                s_iscalars[1] = cutlass.Int32(0)  # done
-                s_iscalars[2] = cutlass.Int32(-1)  # cnt_lo (fb seeding owns)
-                s_iscalars[3] = cutlass.Int32(-1)  # cnt_hi
-                s_iscalars[4] = cutlass.Int32(0)  # out_count
-            cute.arch.barrier()
+            if ext_row == cutlass.Int32(1):
+                if tidx == cutlass.Int32(0):
+                    s_thr[0] = seed_thr_row[1]
+                    s_thr[1] = seed_thr_row[0]
+                    s_thr[2] = seed_thr_row[2]
+                    s_iscalars[0] = cutlass.Int32(0)  # cand_count
+                    s_iscalars[1] = cutlass.Int32(0)  # done
+                    s_iscalars[2] = cutlass.Int32(-1)  # cnt_lo (fb seeding owns)
+                    s_iscalars[3] = cutlass.Int32(-1)  # cnt_hi
+                    s_iscalars[4] = cutlass.Int32(0)  # out_count
+                cute.arch.barrier()
+            if ext_row == cutlass.Int32(0):
+                self.phase1_preidx_stats(
+                    input_row,
+                    N,
+                    pre_idx_row,
+                    pre_idx_count,
+                    pre_idx_offset,
+                    smem_wmin,
+                    smem_wmax,
+                    smem_wsum,
+                    smem_wcnt_p1,
+                    s_thr,
+                    s_iscalars,
+                    tidx,
+                    warp_id,
+                    lane,
+                    smem_gath=smem_gath,  # p1b_cache: stash gathered values (None-op OFF)
+                    s_mt_thr=s_mt_thr,  # r0_vseed: park pmean in the last rung column
+                )
         if cutlass.const_expr(not self.use_ext_counts):
             self.phase1_preidx_stats(
                 input_row,
@@ -4915,38 +4951,67 @@ def _run_phases(
                     # the rung counts (phase1b rungs are per-CTA identical since
                     # preIdx stats are full-row).
                     if cutlass.const_expr(self.use_ext_counts):
-                        # ---- Waterfall L1 admission (ext rungs, v2a) ----
-                        # Rung thresholds arrive from the indexer epilogue:
-                        # ONLY P1b is skipped. The stock M-ary count pass runs
-                        # on the ext rungs so the block-skip list build, rung
-                        # tightening, per-thread hand-off and classify all
-                        # compose unchanged (v1 routed through the dense
-                        # refine and forfeited the compact-walk win: flash 1M
-                        # ext 34.7us vs skipR0 15.6us cold).
-                        # v2b: when an ext count is already in [K, kC], park
-                        # THE ADMITTED THRESHOLD IN ALL RUNG SLOTS — the M-ary
-                        # pass degenerates to one compact single-threshold
-                        # count (+ list build at that threshold) and classify
-                        # admits it; a full miss keeps the 3 distinct rungs as
-                        # measured brackets for the seeded refine.
-                        if tidx == cutlass.Int32(0):
-                            bx_m = cutlass.Int32(-1)
-                            bx_c = cutlass.Int32(2147483647)
-                            for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                                cx = cutlass.Int32(seed_counts_row[m])
-                                if (
-                                    cx >= cutlass.Int32(self.top_k)
-                                    and cx <= cutlass.Int32(self.kC)
-                                    and cx < bx_c
-                                ):
-                                    bx_m = cutlass.Int32(m)
-                                    bx_c = cx
-                            for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                                if bx_m >= cutlass.Int32(0):
-                                    s_mt_thr[m] = seed_thr_row[bx_m]
-                                else:
-                                    s_mt_thr[m] = seed_thr_row[m]
-                        cute.arch.barrier()
+                        if ext_row == cutlass.Int32(1):
+                            # ---- Waterfall L1 admission (ext rungs, v2a) ----
+                            # Rung thresholds arrive from the indexer epilogue:
+                            # ONLY P1b is skipped. The stock M-ary count pass runs
+                            # on the ext rungs so the block-skip list build, rung
+                            # tightening, per-thread hand-off and classify all
+                            # compose unchanged (v1 routed through the dense
+                            # refine and forfeited the compact-walk win: flash 1M
+                            # ext 34.7us vs skipR0 15.6us cold).
+                            # v2b: when an ext count is already in [K, kC], park
+                            # THE ADMITTED THRESHOLD IN ALL RUNG SLOTS — the M-ary
+                            # pass degenerates to one compact single-threshold
+                            # count (+ list build at that threshold) and classify
+                            # admits it; a full miss keeps the 3 distinct rungs as
+                            # measured brackets for the seeded refine.
+                            if tidx == cutlass.Int32(0):
+                                bx_m = cutlass.Int32(-1)
+                                bx_c = cutlass.Int32(2147483647)
+                                for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                                    cx = cutlass.Int32(seed_counts_row[m])
+                                    if (
+                                        cx >= cutlass.Int32(self.top_k)
+                                        and cx <= cutlass.Int32(self.kC)
+                                        and cx < bx_c
+                                    ):
+                                        bx_m = cutlass.Int32(m)
+                                        bx_c = cx
+                                for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                                    if bx_m >= cutlass.Int32(0):
+                                        s_mt_thr[m] = seed_thr_row[bx_m]
+                                    else:
+                                        s_mt_thr[m] = seed_thr_row[m]
+                            cute.arch.barrier()
+                        if ext_row == cutlass.Int32(0):
+                            if cutlass.const_expr(self.p1b_cache):
+                                # rungs from the SMEM gather-cache P1 stashed (no 2nd
+                                # GMEM gather); 16-bit only.
+                                self.phase1b_hspace_rungs_cached(
+                                    pre_idx_count,
+                                    smem_gath,
+                                    smem_hist,
+                                    s_thr,
+                                    s_mt_thr,
+                                    tidx,
+                                    warp_id,
+                                    lane,
+                                )
+                            else:
+                                self.phase1b_hspace_rungs(
+                                    input_row,
+                                    N,
+                                    pre_idx_row,
+                                    pre_idx_count,
+                                    pre_idx_offset,
+                                    smem_hist,
+                                    s_thr,
+                                    s_mt_thr,
+                                    tidx,
+                                    warp_id,
+                                    lane,
+                                )
                     if cutlass.const_expr(not self.use_ext_counts):
                         if cutlass.const_expr(self.p1b_cache):
                             # rungs from the SMEM gather-cache P1 stashed (no 2nd

From b30d476295dc0d22ec136c31a79b78d3575d057b Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 21 Jul 2026 00:29:38 -0700
Subject: [PATCH 016/117] [None][perf] GVR waterfall: cluster_size > 1 support
 for ext/L2/xstate

The BS=1 mid/long-row cells stayed on op26 because the waterfall fast
paths were cs==1-only while op26's pick_config splits a single row
across cs=4/8 CTAs. The pre-collected pairs are O(cand_count), so row
splitting buys the direct path nothing: at cs > 1 the LEADER loads the
pairs alone (take_cand is cluster-uniform - every CTA reads the same
per-row control words) and peers publish zero local candidates for the
DSMEM gather; ineligible/invalid rows fall through to the native stock
path at op26's own cluster split. xstate writes at the leader's
Phase-4 exit; the ext count pass composes with the existing cs>1
cluster merge unchanged.

Validation: bl2 cells exact at cs=1/4/8 including forced-void
fallback; cs1 smokes and the adaptive-rung chains unchanged (in-band
0.99). Cold protocol with the production arm (op26 launch config + ext
inputs + in-kernel routing), vs op26 baseline: flash 256k BS1/64
1.24x/1.59x, flash 512k 1.33x/1.42x, flash 1M BS64 3.69x, pro 256k
1.20-1.73x - the former regression cells flip to wins; pro 512k/1M BS1
static-rung misses route to stock (adaptive xstate rungs take them
direct in the closed loop, 1.4-1.6x steady-state).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 34 ++++++++++++++++---
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 07a98c06c1cf..259efdce6211 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -553,8 +553,6 @@ def __init__(
         # (interface v2: [0] valid, [1] kth proxy, [2] accepted threshold,
         # [3] cand_count). cs==1 only (leader-gather rows land later).
         self.emit_xstate = bool(emit_xstate)
-        if emit_xstate and cluster_size != 1:
-            raise ValueError("emit_xstate is single-CTA (cs==1) only")
         # use_ext_cand (waterfall L2 direct-to-P4): pre-collected (value,
         # index) pairs from the epilogue land straight in smem_keys/vals —
         # no P1, no counting, no P3 scan. Eligible when void==0, claimed
@@ -571,8 +569,10 @@ def __init__(
                 raise ValueError("use_ext_counts requires fb_fix")
             if self.M_thr != 3:
                 raise ValueError("use_ext_counts expects exactly 3 seed rungs")
-            if cluster_size != 1:
-                raise ValueError("use_ext_counts is single-CTA (cs==1) only")
+            # cluster_size > 1 supported: the ext rungs/counts are
+            # per-row (identical across the cluster), the stock multi
+            # count pass cluster-merges as usual, and the L2 direct
+            # loader runs leader-only (peers contribute zero candidates).
         # R1 inline shot aim in log2-count space: geometric center of the
         # [K, kC] acceptance window.
         self.log2_r1aim = math.log2(math.sqrt(self.top_k * self.kC)) if self.r0_qfracs else 0.0
@@ -4891,6 +4891,15 @@ def _run_phases(
                     # Cooperative sentinel-skipping load of (value, index)
                     # pairs into the P4 candidate arrays. SMEM-atomic
                     # compaction: order is irrelevant to rank-scatter.
+                    # cs > 1: the pre-collected pairs are O(cand_count) —
+                    # row splitting buys nothing, so the LEADER loads them
+                    # alone and peers publish zero local candidates for
+                    # the DSMEM gather (take_cand is cluster-uniform: all
+                    # CTAs read the same per-row control words).
+                    if cutlass.const_expr(cluster_size > 1):
+                        if tidx == cutlass.Int32(0):
+                            s_iscalars[5] = cutlass.Int32(0)
+                        cute.arch.barrier()
                     if tidx == cutlass.Int32(0):
                         s_iscalars[0] = cutlass.Int32(0)
                         s_thr[0] = seed_thr_row[self.cand_rung]
@@ -4898,6 +4907,9 @@ def _run_phases(
                     cute.arch.barrier()
                     cbase = cand_row.iterator.toint()
                     i_c = tidx
+                    if cutlass.const_expr(cluster_size > 1):
+                        if cta_in_cluster != cutlass.Int32(0):
+                            i_c = claimed_c  # peers: skip the walk
                     while i_c < claimed_c:
                         pa = cbase + cutlass.Int64(i_c) * cutlass.Int64(8)
                         ip_c = cute.make_ptr(
@@ -4924,6 +4936,12 @@ def _run_phases(
                                 smem_vals[wpos] = pidx
                         i_c = i_c + cutlass.Int32(num_threads)
                     cute.arch.barrier()
+                    if cutlass.const_expr(cluster_size > 1):
+                        # leader's local count feeds the DSMEM gather
+                        if is_leader:
+                            if tidx == cutlass.Int32(0):
+                                s_iscalars[5] = s_iscalars[0]
+                        cute.arch.barrier()
             if take_cand == cutlass.Int32(0):
                 # Stage this CTA's slice into SMEM once before Phase 2's
                 # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
@@ -5470,6 +5488,14 @@ def _run_phases(
                             warp_id,
                             lane,
                         )
+                    if cutlass.const_expr(self.emit_xstate):
+                        # closed-loop state, leader-only at cs > 1 (same
+                        # layout as the cs == 1 exit).
+                        if tidx == cutlass.Int32(0):
+                            xstate_row[0] = cutlass.Float32(1.0)
+                            xstate_row[1] = s_thr[0]
+                            xstate_row[2] = s_thr[0]
+                            xstate_row[3] = cutlass.Float32(cand_count_p4)
 
         # Final cluster barrier: keep peer CTAs (and their SMEM) alive
         # until the leader's gather + Phase 4 finish. Skipped at

From 6cf89c1fa89a884faa3f28d1e0e286c4e66e65fa Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 22 Jul 2026 03:40:05 -0700
Subject: [PATCH 017/117] [None][perf] GVR L2 direct: two-parameter (K, K_max)
 in-list filtering

Collect the pre-collected candidate list at the LOOSEST seed rung and
admit it whenever it is complete (claimed <= K_max) and any rung counts
>= K; the filter rung (count closest to K from above) is applied on the
fly while loading the pairs, so P4 sees the thinnest covering set. kC
leaves the admission vocabulary and remains only as the physical smem
capacity guard. Correctness: C(t_lo) >= C(t_filt) >= K implies true
top-K subset of list subset of filtered set; list truncation (claim
order is value-blind) remains the only fatal case and falls back.

- K_max = 24576, set by a four-chain search on real captures (incl.
  320k/640k long decode): 16K->24K gains 8pp direct-hit rate,
  24K->32K only 0.1pp (band-limited, not capacity-limited).
- Loader: 4x-unrolled latency-overlapped walk with ballot-batched smem
  claims (loop exit must stay warp-uniform: ragged exits deadlock the
  warp collectives) and un-nested value loads. Device-level (nsys
  kern-sum, cold L2) on a 160k real chain: the naive walk ran 0.64x vs
  the block-skip arm; this form reaches 1.05x at full loosest-rung
  coverage (eligibility 1.00).
- Straddle refine (cs=1): when no rung count lands in [K, kC] but the
  list is complete, one 256-bin histogram pass over the list finds an
  in-band edge and the filtered load proceeds; smem overflow demotes
  to the fallback. 640k chain: straddle steps 30 -> 14-16us, device
  mean 1.41x -> 1.73x vs block-skip.
- Byte-parity routing keeps fat lists (2*claimed*cs >= N) on the
  fallback: measured both ways, walking them is slower at every cs.

Validation (B200): four admission modes exact (direct/filter/refine/
fallback) at cs=1 and 18/18 exact at cs=1/4/8 on real V4 bundles; four
real decode chains (160k/132k/320k/640k) all-step exact with wall
ratios 0.97/1.00/1.04/1.21x and device-level 1.05x (160k) / 1.73x
(640k) vs the block-skip arm.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 265 ++++++++++++++++--
 1 file changed, 234 insertions(+), 31 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 259efdce6211..3db23dcc045d 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -274,7 +274,7 @@ def __init__(
         use_ext_counts: bool = False,
         use_ext_cand: bool = False,
         cand_cap: int = 5120,
-        cand_rung: int = 1,
+        cand_rung: int = 0,
         emit_xstate: bool = False,
     ):
         # Redundant-warp sync reduction: every warp replays the block
@@ -4879,14 +4879,161 @@ def _run_phases(
             if cutlass.const_expr(self.use_ext_cand):
                 claimed_c = cutlass.Int32(cand_ctl_row[0])
                 void_c = cutlass.Int32(cand_ctl_row[1])
-                ccnt_c = cutlass.Int32(seed_counts_row[self.cand_rung])
+                # Two-parameter admission (K, K_max): the list (collected
+                # at rung cand_rung) is usable whenever it is COMPLETE
+                # (claimed <= K_max = cand_cap) and some rung at or above
+                # the collect rung counts >= K. The FILTER line is the
+                # rung with the smallest such count (closest to K from
+                # above) - monotonicity gives C(collect) >= C(filter)
+                # >= K, so the filtered set still covers the true top-K
+                # and is as small as the rung group allows (cheapest P4).
+                # kC survives only as the physical smem-capacity guard.
+                filt_t = cutlass.Float32(0.0)
+                filt_c = cutlass.Int32(2147483647)
+                filt_j = cutlass.Int32(-1)
+                for _fj in cutlass.range_constexpr(
+                    cutlass.const_expr(self.cand_rung),
+                    cutlass.const_expr(self.M_thr),
+                ):
+                    cfj = cutlass.Int32(seed_counts_row[_fj])
+                    # <= prefers the tightest rung among equal counts
+                    if cfj >= cutlass.Int32(self.top_k) and cfj <= filt_c:
+                        filt_c = cfj
+                        filt_t = seed_thr_row[_fj]
+                        filt_j = cutlass.Int32(_fj)
+                # Byte-parity routing: the list walk reads 8B/entry on the
+                # LEADER alone, a fallback re-scan reads 4B/elem split
+                # across the cluster - a list fatter than N/(2*cs) routes
+                # to the fallback. Measured both ways (static-threshold
+                # grid + 160k closed-loop chain): dropping this gate at
+                # cs=1 turned every fat-list fallback into a slower walk
+                # (pro 64k 13.5 -> 18.3us), so it stays for all cs.
                 if (
                     void_c == cutlass.Int32(0)
                     and claimed_c <= cutlass.Int32(self.cand_cap)
-                    and ccnt_c >= cutlass.Int32(self.top_k)
-                    and ccnt_c <= cutlass.Int32(self.kC)
+                    and filt_c <= cutlass.Int32(self.kC)
+                    and (claimed_c + claimed_c) * cutlass.Int32(cluster_size) < N
                 ):
                     take_cand = cutlass.Int32(1)
+                if cutlass.const_expr(cluster_size == 1):
+                    # ---- Straddle refine: no rung count lands in the
+                    # [K, kC] band (the loosest >=K rung is fatter than
+                    # smem, the next one is under K), yet the list is
+                    # COMPLETE - so instead of falling back to a full
+                    # row re-scan, one 256-bin histogram pass over the
+                    # list values between the two rungs finds a refined
+                    # in-band edge, and the normal filtered load takes
+                    # over. cs>1 keeps the plain fallback: the fire
+                    # decision would need a DSMEM exchange to stay
+                    # cluster-uniform. All gating values are CTA-uniform
+                    # (gmem control words), so the barriers below are
+                    # safe. Runs on ~1 row in 8 on distribution-shift
+                    # steps; those rows otherwise dominate the whole
+                    # launch (per-row max ~30us vs ~15us).
+                    strad = cutlass.Int32(0)
+                    if (
+                        take_cand == cutlass.Int32(0)
+                        and void_c == cutlass.Int32(0)
+                        and claimed_c <= cutlass.Int32(self.cand_cap)
+                        and filt_j >= cutlass.Int32(0)
+                        and filt_c > cutlass.Int32(self.kC)
+                        and filt_j < cutlass.Int32(self.M_thr - 1)
+                        and (claimed_c + claimed_c) < N
+                    ):
+                        t_thin_s = seed_thr_row[filt_j + cutlass.Int32(1)]
+                        if t_thin_s > filt_t:
+                            strad = cutlass.Int32(1)
+                    if strad == cutlass.Int32(1):
+                        t_fat = filt_t
+                        t_thin = seed_thr_row[filt_j + cutlass.Int32(1)]
+                        c_thin = cutlass.Int32(seed_counts_row[filt_j + cutlass.Int32(1)])
+                        NBr = cutlass.const_expr(256)
+                        w_r = (t_thin - t_fat) / cutlass.Float32(NBr)
+                        inv_wr = cutlass.Float32(1.0) / w_r
+                        jz_r = tidx
+                        while jz_r < cutlass.Int32(NBr):
+                            smem_hist[jz_r] = cutlass.Int32(0)
+                            jz_r = jz_r + cutlass.Int32(num_threads)
+                        # s_iscalars[2]/[3] (P2's cnt_lo/cnt_hi) are dead
+                        # until P2 re-inits them on the fallback path ->
+                        # free {fired flag, bin index} broadcast slots.
+                        if tidx == cutlass.Int32(0):
+                            s_iscalars[2] = cutlass.Int32(0)
+                        cute.arch.barrier()
+                        cb_r = cand_row.iterator.toint()
+                        i_r = tidx
+                        # ragged exit is fine here: no warp collectives,
+                        # only smem atomics. 4x unroll overlaps the pair
+                        # loads (same latency argument as the loader).
+                        while i_r < claimed_c:
+                            for _jr in cutlass.range_constexpr(4):
+                                i_rj = i_r + cutlass.Int32(_jr * num_threads)
+                                if i_rj < claimed_c:
+                                    pa_r = cb_r + cutlass.Int64(i_rj) * cutlass.Int64(8)
+                                    vp_r = cute.make_ptr(
+                                        cutlass.Float32,
+                                        pa_r,
+                                        cute.AddressSpace.gmem,
+                                        assumed_align=8,
+                                    )
+                                    ip_r = cute.make_ptr(
+                                        cutlass.Int32,
+                                        pa_r + cutlass.Int64(4),
+                                        cute.AddressSpace.gmem,
+                                        assumed_align=4,
+                                    )
+                                    v_r = cute.make_tensor(vp_r, cute.make_layout((1,)))[0]
+                                    x_r = cute.make_tensor(ip_r, cute.make_layout((1,)))[0]
+                                    # only the open band [t_fat, t_thin):
+                                    # >= t_thin is already counted by
+                                    # c_thin, < t_fat is below the edge.
+                                    if x_r >= cutlass.Int32(0) and v_r >= t_fat and v_r < t_thin:
+                                        b_r = cutlass.Int32((v_r - t_fat) * inv_wr)
+                                        if b_r < cutlass.Int32(0):
+                                            b_r = cutlass.Int32(0)
+                                        if b_r > cutlass.Int32(NBr - 1):
+                                            b_r = cutlass.Int32(NBr - 1)
+                                        atomicAdd(smem_hist.iterator + b_r, cutlass.Int32(1))
+                            i_r = i_r + cutlass.Int32(4 * num_threads)
+                        cute.arch.barrier()
+                        # warp-0 top-down cumulative scan (P1b idiom):
+                        # lane l owns SEG bins descending from NB-1-l*SEG;
+                        # fire at the unique bin where the cumulative
+                        # count (incl. c_thin base) crosses K, admit if
+                        # it also fits smem with a float-edge margin.
+                        if warp_id == cutlass.Int32(0):
+                            SEGr = cutlass.const_expr(8)
+                            top_r = cutlass.Int32(NBr - 1) - lane * cutlass.Int32(SEGr)
+                            seg_r = cute.make_fragment((SEGr,), cutlass.Int32)
+                            part_r = cutlass.Int32(0)
+                            for _js in cutlass.range_constexpr(SEGr):
+                                v8_r = smem_hist[top_r - cutlass.Int32(_js)]
+                                seg_r[_js] = v8_r
+                                part_r = part_r + v8_r
+                            tp_r = part_r
+                            for _os in cutlass.range_constexpr(5):
+                                ov_r = cutlass.const_expr(1 << _os)
+                                oth_r = cute.arch.shuffle_sync_up(tp_r, ov_r, mask_and_clamp=0)
+                                if lane >= cutlass.Int32(ov_r):
+                                    tp_r = tp_r + oth_r
+                            excl_r = tp_r - part_r
+                            run_r = cutlass.Int32(0)
+                            kneed = cutlass.Int32(self.top_k) - c_thin
+                            kfit = cutlass.Int32(self.kC - 64) - c_thin
+                            for _js in cutlass.range_constexpr(SEGr):
+                                run_r = run_r + seg_r[_js]
+                                cum_at = excl_r + run_r
+                                cum_bef = cum_at - seg_r[_js]
+                                if cum_bef < kneed and cum_at >= kneed:
+                                    if cum_at <= kfit:
+                                        s_iscalars[3] = top_r - cutlass.Int32(_js)
+                                        s_iscalars[2] = cutlass.Int32(1)
+                        cute.arch.barrier()
+                        if s_iscalars[2] == cutlass.Int32(1):
+                            # every thread recomputes the edge from the
+                            # uniform bin index
+                            filt_t = t_fat + cutlass.Float32(s_iscalars[3]) * w_r
+                            take_cand = cutlass.Int32(1)
                 if take_cand == cutlass.Int32(1):
                     # Cooperative sentinel-skipping load of (value, index)
                     # pairs into the P4 candidate arrays. SMEM-atomic
@@ -4902,40 +5049,96 @@ def _run_phases(
                         cute.arch.barrier()
                     if tidx == cutlass.Int32(0):
                         s_iscalars[0] = cutlass.Int32(0)
-                        s_thr[0] = seed_thr_row[self.cand_rung]
+                        s_thr[0] = filt_t
                         s_iscalars[1] = cutlass.Int32(1)  # done: no retry
                     cute.arch.barrier()
                     cbase = cand_row.iterator.toint()
+                    lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1)
                     i_c = tidx
                     if cutlass.const_expr(cluster_size > 1):
                         if cta_in_cluster != cutlass.Int32(0):
-                            i_c = claimed_c  # peers: skip the walk
-                    while i_c < claimed_c:
-                        pa = cbase + cutlass.Int64(i_c) * cutlass.Int64(8)
-                        ip_c = cute.make_ptr(
-                            cutlass.Int32,
-                            pa + cutlass.Int64(4),
-                            cute.AddressSpace.gmem,
-                            assumed_align=4,
-                        )
-                        pidx = cute.make_tensor(ip_c, cute.make_layout((1,)))[0]
-                        if pidx >= cutlass.Int32(0):
-                            vp_c = cute.make_ptr(
-                                cutlass.Float32,
-                                pa,
-                                cute.AddressSpace.gmem,
-                                assumed_align=8,
-                            )
-                            pval = cute.make_tensor(vp_c, cute.make_layout((1,)))[0]
-                            wpos = atomicAdd(
-                                s_iscalars.iterator + cutlass.Int32(0),
-                                cutlass.Int32(1),
-                            )
-                            if wpos < cutlass.Int32(self.kC):
-                                smem_keys[wpos] = pval
-                                smem_vals[wpos] = pidx
-                        i_c = i_c + cutlass.Int32(num_threads)
+                            # peers: skip the walk. Must stay warp-uniform
+                            # for the ballot below, so park every lane at
+                            # (claimed + lane) -> warp base == claimed.
+                            i_c = claimed_c + lane_c
+                    # Loop condition on the warp BASE index (uniform across
+                    # the warp): a ragged per-lane exit would strand the
+                    # remaining lanes in vote_ballot_sync (full-warp
+                    # collective) on the tail iteration -> device hang.
+                    #
+                    # 4x-unrolled, latency-overlapped walk. The naive form
+                    # (idx load -> dependent val load -> ballot, one entry
+                    # per thread per trip) costs ~2 serialized cold-DRAM
+                    # round trips per trip; a fat list (C(t_lo) ~ 6-12K)
+                    # then eats ~1.5us/trip (measured: wf 0.64x vs skip at
+                    # 160k). Here all 8 loads (4 entries x {idx,val}) issue
+                    # independently BEFORE the first use, so a trip costs
+                    # ~one round trip for 4x the entries. The val load is
+                    # NOT nested under the sentinel check: sentinel pairs
+                    # are mapped memory, their value is simply ignored.
+                    while (i_c - lane_c) < claimed_c:
+                        pvals = []
+                        pidxs = []
+                        for _ju in cutlass.range_constexpr(4):
+                            pval = cutlass.Float32(0.0)
+                            pidx = cutlass.Int32(-1)
+                            i_cj = i_c + cutlass.Int32(_ju * num_threads)
+                            if i_cj < claimed_c:
+                                pa = cbase + cutlass.Int64(i_cj) * cutlass.Int64(8)
+                                vp_c = cute.make_ptr(
+                                    cutlass.Float32,
+                                    pa,
+                                    cute.AddressSpace.gmem,
+                                    assumed_align=8,
+                                )
+                                ip_c = cute.make_ptr(
+                                    cutlass.Int32,
+                                    pa + cutlass.Int64(4),
+                                    cute.AddressSpace.gmem,
+                                    assumed_align=4,
+                                )
+                                pval = cute.make_tensor(vp_c, cute.make_layout((1,)))[0]
+                                pidx = cute.make_tensor(ip_c, cute.make_layout((1,)))[0]
+                            pvals.append(pval)
+                            pidxs.append(pidx)
+                        for _ju in cutlass.range_constexpr(4):
+                            keep_c = cutlass.Int32(0)
+                            if pidxs[_ju] >= cutlass.Int32(0):
+                                if pvals[_ju] >= filt_t:
+                                    keep_c = cutlass.Int32(1)
+                            # ballot-batched claim (the P4 writeback idiom):
+                            # one smem atomic per warp-trip instead of one
+                            # per survivor.
+                            mask_k = cute.arch.vote_ballot_sync(keep_c != cutlass.Int32(0))
+                            if mask_k != cutlass.Uint32(0):
+                                nk = cutlass.Int32(cute.arch.popc(mask_k))
+                                lmk = (
+                                    cutlass.Uint32(1) << cutlass.Uint32(lane_c)
+                                ) - cutlass.Uint32(1)
+                                offk = cutlass.Int32(cute.arch.popc(mask_k & lmk))
+                                bk = cutlass.Int32(0)
+                                if lane_c == cutlass.Int32(0):
+                                    bk = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), nk)
+                                bk = cute.arch.shuffle_sync(bk, cutlass.Int32(0))
+                                wpos = bk + offk
+                                if keep_c != cutlass.Int32(0) and wpos < cutlass.Int32(self.kC):
+                                    smem_keys[wpos] = pvals[_ju]
+                                    smem_vals[wpos] = pidxs[_ju]
+                        i_c = i_c + cutlass.Int32(4 * num_threads)
                     cute.arch.barrier()
+                    if cutlass.const_expr(cluster_size == 1):
+                        # Hard net for the refined edge: its count is
+                        # float-boundary approximate. If survivors
+                        # overflowed kC the smem set is truncated
+                        # (value-blind) -> demote to the fallback
+                        # re-scan. Exact-rung filters cannot overflow
+                        # (filt_c <= kC checked); s_iscalars[0] is
+                        # CTA-uniform after the barrier above.
+                        if s_iscalars[0] > cutlass.Int32(self.kC):
+                            take_cand = cutlass.Int32(0)
+                            if tidx == cutlass.Int32(0):
+                                s_iscalars[1] = cutlass.Int32(0)
+                            cute.arch.barrier()
                     if cutlass.const_expr(cluster_size > 1):
                         # leader's local count feeds the DSMEM gather
                         if is_leader:

From b11f4a097c3c58f9dae215944b02c3aa82b40575 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 23 Jul 2026 05:25:55 -0700
Subject: [PATCH 018/117] [None][perf] GVR L2: count-only smem-list selection
 (SoA candidates)

Replace the rung-based in-list-filter admission with the count-only
scheme: the candidate list is SoA (score column + position column,
sentinel score -inf), collected at a single loose line, and admitted
purely by entry count (K + 64 <= claimed <= K_max = 24576; the 64 is
the emitter sentinel-pad bound, so the live count provably covers K).
Rung admission, filter-line selection, straddle refine and the parity
gate are all deleted - the seed-count columns are no longer consumed
on the list path (the GEMM-side L1 pass becomes deletable, -3.2%
emission tax).

- THIN list (fits kC): every entry lands AT ITS LIST INDEX in the
  candidate buffers - no ballots, no smem atomics (128 serialized
  same-address atomics per trip measured ~1.1us/1k entries), no
  warp-uniform loop constraint.
- FAT list: atomic-free copy of the score column into a dedicated
  96KB smem region (sentinels sanitized to t_lo - 1), a zooming smem
  histogram (3 rounds, NBL^3 resolution - value-linear bins collapse
  on long-tailed logits) finds an edge whose exact count lands in
  [K, kC] (lands ~1030 for K=1024), survivors compact with one
  merged-ballot atomic per warp per trip. The vals slots carry LIST
  INDICES (no second cold gmem pass over the position column); a
  post-P4 repair swaps the K winners' positions with fully-parallel
  gathers.
- Closed loop: xstate[1] publishes the exact k-th (output slot K-1 of
  the rank-ordered scatter), xstate[2] the ~3K-crossing anchor from
  the round-0 histogram. Host policy picks the anchor field per
  domain (tight k-th for short/stable rows, wide 3K edge for volatile
  long rows - the exact-k-th anchor alone shrinks the next down-guard
  target to 4K and slope noise then undershoots K, forcing ~26us
  fallbacks). GVR_P4_TAIL_DBG compiles per-phase clock64 stamps into
  the spare xstate slots.

Validation (B200): 24/24 exact across cs=1/4/8 and the straddle-
threshold suite on real V4 bundles; per-row cold device phases: thin
walk 1.5-3us, fat stage+zoom+compact ~1.1us/1k entries, Phase 4 flat
5.5-6.5us. Real-chain device-level vs the block-skip arm: 640k 1.42x
(fallback steps are C(t_lo) < K undershoots - a host anchor-policy
matter), 160k 0.93x. Kernel-only chain means trade 5-20% vs the
previous rung-based commit at B=8 in exchange for the interface
collapse; the deleted L1 emission pass dominates E2E at large batch.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 657 +++++++++++-------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  73 +-
 2 files changed, 442 insertions(+), 288 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 3db23dcc045d..7993a773895a 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -27,6 +27,7 @@
 """
 
 import math
+import os
 from dataclasses import dataclass
 from typing import Optional
 
@@ -41,6 +42,11 @@
 from ..utils import TRTLLM_ENABLE_PDL, griddepcontrol_launch_dependents, griddepcontrol_wait
 from .block_scan import warp_scan
 
+# Diagnostic knob: compile per-phase clock64 stamps of the list path
+# into the spare xstate slots (harness-side analysis). Off by default;
+# NEVER set in production.
+_P4_TAIL_DBG = bool(int(os.environ.get("GVR_P4_TAIL_DBG", "0")))
+
 
 # ---------------------------------------------------------------------------
 # DSMEM primitives (inline PTX)
@@ -560,6 +566,13 @@ def __init__(
         # ineligible rows fall through to the ext-counts path.
         self.use_ext_cand = bool(use_ext_cand)
         self.cand_cap = int(cand_cap)
+        # list path: the score column is staged into a DEDICATED smem
+        # region sized cand_cap fp32 (96KB at 24576). Budget note: this
+        # coexists with everything except a simultaneously-enabled big
+        # slice cache (128KB) - that combination exceeds the 227KB CTA
+        # limit and fails loudly at compile time. Long rows (the list
+        # path's target) cannot enable the slice cache anyway.
+        self.list_cap = int(cand_cap)
         self.cand_rung = int(cand_rung)
         if use_ext_cand and not use_ext_counts:
             raise ValueError("use_ext_cand requires use_ext_counts")
@@ -3165,6 +3178,10 @@ def phase4_rank_scatter(
                 # UNMODIFIED radix select below (verbatim copy).
                 if cutlass.const_expr(self.p4_exact_tail and self.p4_tail_fast):  # [p4tt]
                     need0 = cutlass.Int32(kK) - rank_above_fine
+                    if cutlass.const_expr(_P4_TAIL_DBG):
+                        if tidx == cutlass.Int32(0):
+                            s_thr[1] = cutlass.Float32(cnt_strad)
+                            s_thr[2] = cutlass.Float32(need0)
                     if cnt_strad > need0 and need0 > cutlass.Int32(0):
                         if cnt_strad <= cutlass.Int32(128):
                             # [p4tt] SMEM: (value_bits, cand_idx) pairs at
@@ -4140,7 +4157,8 @@ def gvr_topk_kernel(
         seed_thr: cute.Tensor,  # [numRows, 3] fp32 (or None: no ext counts)
         seed_counts: cute.Tensor,  # [numRows, 3] int32 (or None)
         xstate: cute.Tensor,  # [numRows, 8] fp32 closed-loop state (or None)
-        cand: cute.Tensor,  # [numRows, CAP*2] int32 pairs (or None)
+        cand_vals: cute.Tensor,  # [numRows, CAP] fp32 scores (or None)
+        cand_idx: cute.Tensor,  # [numRows, CAP] int32 positions (or None)
         cand_ctl: cute.Tensor,  # [numRows, 2] int32 {claimed, void} (or None)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
@@ -4197,7 +4215,8 @@ def gvr_topk_kernel(
             seed_thr=seed_thr,
             seed_counts=seed_counts,
             xstate=xstate,
-            cand=cand,
+            cand_vals=cand_vals,
+            cand_idx=cand_idx,
             cand_ctl=cand_ctl,
         )
 
@@ -4214,7 +4233,8 @@ def run_one_row(
         seed_thr: cute.Tensor = None,  # [numRows, 3] fp32 (ext counts)
         seed_counts: cute.Tensor = None,  # [numRows, 3] int32 (ext counts)
         xstate: cute.Tensor = None,  # [numRows, 8] fp32 (emit_xstate)
-        cand: cute.Tensor = None,  # [numRows, CAP*2] int32 (ext cand)
+        cand_vals: cute.Tensor = None,  # [numRows, CAP] fp32 (ext cand)
+        cand_idx: cute.Tensor = None,  # [numRows, CAP] int32 (ext cand)
         cand_ctl: cute.Tensor = None,  # [numRows, 2] int32 (ext cand)
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
@@ -4288,11 +4308,18 @@ def run_one_row(
             xstate_row = xstate[row_idx, None]
         else:
             xstate_row = None
-        if cutlass.const_expr(self.use_ext_cand and cand is not None and cand_ctl is not None):
-            cand_row = cand[row_idx, None]
+        if cutlass.const_expr(
+            self.use_ext_cand
+            and cand_vals is not None
+            and cand_idx is not None
+            and cand_ctl is not None
+        ):
+            cand_vals_row = cand_vals[row_idx, None]
+            cand_idx_row = cand_idx[row_idx, None]
             cand_ctl_row = cand_ctl[row_idx, None]
         else:
-            cand_row = None
+            cand_vals_row = None
+            cand_idx_row = None
             cand_ctl_row = None
         # When return_output_values=False, ``output_values`` is None at
         # launch and the gated writes below are compiled out; slicing into
@@ -4437,6 +4464,21 @@ def run_one_row(
         else:
             smem_input = None
 
+        # List-path score staging (SoA ext cand, cs=1, fp32 only).
+        if cutlass.const_expr(
+            self.use_ext_cand
+            and self.use_ext_counts
+            and cluster_size == 1
+            and self.dtype == cutlass.Float32
+        ):
+            smem_list = smem.allocate_tensor(
+                element_type=cutlass.Float32,
+                layout=cute.make_ordered_layout((self.list_cap,), order=(0,)),
+                byte_alignment=128,
+            )
+        else:
+            smem_list = None
+
         # op#26 R0 admission scratch (single-CTA fast path). Allocated only
         # when enable_r0; None otherwise so the base SMEM layout is byte-for-
         # byte unchanged and these propagate harmlessly through _run_phases'
@@ -4593,7 +4635,9 @@ def run_one_row(
                         seed_thr_row=seed_thr_row,
                         seed_counts_row=seed_counts_row,
                         xstate_row=xstate_row,
-                        cand_row=cand_row,
+                        cand_vals_row=cand_vals_row,
+                        cand_idx_row=cand_idx_row,
+                        smem_list=smem_list,
                         cand_ctl_row=cand_ctl_row,
                         smem_active=smem_active,
                         s_active_cnt=s_active_cnt,
@@ -4642,7 +4686,9 @@ def run_one_row(
                             seed_thr_row=seed_thr_row,
                             seed_counts_row=seed_counts_row,
                             xstate_row=xstate_row,
-                            cand_row=cand_row,
+                            cand_vals_row=cand_vals_row,
+                            cand_idx_row=cand_idx_row,
+                            smem_list=smem_list,
                             cand_ctl_row=cand_ctl_row,
                             smem_active=smem_active,
                             s_active_cnt=s_active_cnt,
@@ -4688,7 +4734,9 @@ def run_one_row(
                     seed_thr_row=seed_thr_row,
                     seed_counts_row=seed_counts_row,
                     xstate_row=xstate_row,
-                    cand_row=cand_row,
+                    cand_vals_row=cand_vals_row,
+                    cand_idx_row=cand_idx_row,
+                    smem_list=smem_list,
                     cand_ctl_row=cand_ctl_row,
                     smem_active=smem_active,
                     s_active_cnt=s_active_cnt,
@@ -4737,7 +4785,9 @@ def _run_phases(
         seed_thr_row=None,  # ext counts: this row's 3 seed thresholds (fp32)
         seed_counts_row=None,  # ext counts: this row's 3 exact counts (int32)
         xstate_row=None,  # emit_xstate: this row's [8] fp32 state slot
-        cand_row=None,  # ext cand: this row's [CAP*2] int32 pairs
+        cand_vals_row=None,  # ext cand: this row's [CAP] fp32 scores
+        cand_idx_row=None,  # ext cand: this row's [CAP] int32 positions
+        smem_list=None,  # list path: [list_cap] fp32 score staging
         cand_ctl_row=None,  # ext cand: this row's [2] int32 {claimed, void}
         smem_active=None,
         s_active_cnt=None,
@@ -4771,6 +4821,8 @@ def _run_phases(
         # misses (pro-1M cold: stock-skip 21us vs ext-miss 51us). All
         # threads read the same control words, so the predicate is
         # CTA-uniform and the dynamic branches below stay convergent.
+        ck0 = cutlass.Int64(0)
+        ck1 = cutlass.Int64(0)
         ext_row = cutlass.Int32(0)
         if cutlass.const_expr(self.use_ext_counts):
             if seed_thr_row[0] < cutlass.Float32(1e37):
@@ -4778,6 +4830,25 @@ def _run_phases(
                     cm_e = cutlass.Int32(seed_counts_row[m])
                     if cm_e >= cutlass.Int32(self.top_k) and cm_e <= cutlass.Int32(self.kC):
                         ext_row = cutlass.Int32(1)
+            # list path preview: when the SoA candidate list will be taken
+            # (count-only admission), Phase 1's gather buys nothing either
+            # - reuse the same skip (the seed rungs still provide the
+            # [t_0, t_2] bracket the degenerate check wants).
+            if cutlass.const_expr(
+                self.use_ext_cand
+                and self.use_ext_counts
+                and cluster_size == 1
+                and self.dtype == cutlass.Float32
+            ):
+                claimed_p = cutlass.Int32(cand_ctl_row[0])
+                void_p = cutlass.Int32(cand_ctl_row[1])
+                if (
+                    void_p == cutlass.Int32(0)
+                    and claimed_p >= cutlass.Int32(self.top_k)
+                    and claimed_p <= cutlass.Int32(self.list_cap)
+                    and seed_thr_row[0] < cutlass.Float32(1e37)
+                ):
+                    ext_row = cutlass.Int32(1)
 
         # ---- Phase 1: preIdx Min/Max/Mean ----
         # ext counts: P1's only surviving products are the [v_lo, v_hi]
@@ -4871,280 +4942,288 @@ def _run_phases(
                     if cutlass.const_expr(self.emit_xstate):
                         xstate_row[0] = cutlass.Float32(0.0)  # degenerate
         else:
-            # ---- Waterfall L2: direct-to-P4 from pre-collected pairs ----
-            # Eligibility is a CTA-uniform register predicate (all threads
-            # read the same gmem control words), so the dynamic branches
-            # below (with barriers inside) stay convergent.
+            # ---- List path (SoA): complete candidate list -> on-chip ----
+            # Admission is the entry count alone: K <= claimed <= list_cap
+            # (and no overflow) guarantees the true top-K lives in the
+            # list (it holds every position >= t_lo = seed_thr[0]). The
+            # score column is staged into the slice-cache smem region
+            # (fp32 specs; the fallback re-scan never runs on list rows,
+            # so the alias is safe), one smem histogram finds an edge
+            # whose exact count lands in [K, kC], survivors compact into
+            # the standard kC candidate buffers and the stock Phase 4
+            # finishes. All gating values are CTA-uniform (gmem control
+            # words / smem scalars), so the barriers stay convergent.
+            # cs>1 and 16-bit dtypes keep the plain fallback.
             take_cand = cutlass.Int32(0)
-            if cutlass.const_expr(self.use_ext_cand):
+            list_used = cutlass.Int32(0)  # list path taken (xstate publish)
+            claimed_c = cutlass.Int32(0)
+            if cutlass.const_expr(
+                self.use_ext_cand
+                and self.use_ext_counts
+                and cluster_size == 1
+                and self.dtype == cutlass.Float32
+            ):
                 claimed_c = cutlass.Int32(cand_ctl_row[0])
                 void_c = cutlass.Int32(cand_ctl_row[1])
-                # Two-parameter admission (K, K_max): the list (collected
-                # at rung cand_rung) is usable whenever it is COMPLETE
-                # (claimed <= K_max = cand_cap) and some rung at or above
-                # the collect rung counts >= K. The FILTER line is the
-                # rung with the smallest such count (closest to K from
-                # above) - monotonicity gives C(collect) >= C(filter)
-                # >= K, so the filtered set still covers the true top-K
-                # and is as small as the rung group allows (cheapest P4).
-                # kC survives only as the physical smem-capacity guard.
-                filt_t = cutlass.Float32(0.0)
-                filt_c = cutlass.Int32(2147483647)
-                filt_j = cutlass.Int32(-1)
-                for _fj in cutlass.range_constexpr(
-                    cutlass.const_expr(self.cand_rung),
-                    cutlass.const_expr(self.M_thr),
-                ):
-                    cfj = cutlass.Int32(seed_counts_row[_fj])
-                    # <= prefers the tightest rung among equal counts
-                    if cfj >= cutlass.Int32(self.top_k) and cfj <= filt_c:
-                        filt_c = cfj
-                        filt_t = seed_thr_row[_fj]
-                        filt_j = cutlass.Int32(_fj)
-                # Byte-parity routing: the list walk reads 8B/entry on the
-                # LEADER alone, a fallback re-scan reads 4B/elem split
-                # across the cluster - a list fatter than N/(2*cs) routes
-                # to the fallback. Measured both ways (static-threshold
-                # grid + 160k closed-loop chain): dropping this gate at
-                # cs=1 turned every fat-list fallback into a slower walk
-                # (pro 64k 13.5 -> 18.3us), so it stays for all cs.
+                # claimed >= K + 64: sentinel slots are bounded by the
+                # emitter contract (<= 64 pad), so this guarantees the
+                # LIVE count covers K without counting anything.
                 if (
                     void_c == cutlass.Int32(0)
-                    and claimed_c <= cutlass.Int32(self.cand_cap)
-                    and filt_c <= cutlass.Int32(self.kC)
-                    and (claimed_c + claimed_c) * cutlass.Int32(cluster_size) < N
+                    and claimed_c >= cutlass.Int32(self.top_k + 64)
+                    and claimed_c <= cutlass.Int32(self.list_cap)
                 ):
                     take_cand = cutlass.Int32(1)
-                if cutlass.const_expr(cluster_size == 1):
-                    # ---- Straddle refine: no rung count lands in the
-                    # [K, kC] band (the loosest >=K rung is fatter than
-                    # smem, the next one is under K), yet the list is
-                    # COMPLETE - so instead of falling back to a full
-                    # row re-scan, one 256-bin histogram pass over the
-                    # list values between the two rungs finds a refined
-                    # in-band edge, and the normal filtered load takes
-                    # over. cs>1 keeps the plain fallback: the fire
-                    # decision would need a DSMEM exchange to stay
-                    # cluster-uniform. All gating values are CTA-uniform
-                    # (gmem control words), so the barriers below are
-                    # safe. Runs on ~1 row in 8 on distribution-shift
-                    # steps; those rows otherwise dominate the whole
-                    # launch (per-row max ~30us vs ~15us).
-                    strad = cutlass.Int32(0)
-                    if (
-                        take_cand == cutlass.Int32(0)
-                        and void_c == cutlass.Int32(0)
-                        and claimed_c <= cutlass.Int32(self.cand_cap)
-                        and filt_j >= cutlass.Int32(0)
-                        and filt_c > cutlass.Int32(self.kC)
-                        and filt_j < cutlass.Int32(self.M_thr - 1)
-                        and (claimed_c + claimed_c) < N
-                    ):
-                        t_thin_s = seed_thr_row[filt_j + cutlass.Int32(1)]
-                        if t_thin_s > filt_t:
-                            strad = cutlass.Int32(1)
-                    if strad == cutlass.Int32(1):
-                        t_fat = filt_t
-                        t_thin = seed_thr_row[filt_j + cutlass.Int32(1)]
-                        c_thin = cutlass.Int32(seed_counts_row[filt_j + cutlass.Int32(1)])
-                        NBr = cutlass.const_expr(256)
-                        w_r = (t_thin - t_fat) / cutlass.Float32(NBr)
-                        inv_wr = cutlass.Float32(1.0) / w_r
-                        jz_r = tidx
-                        while jz_r < cutlass.Int32(NBr):
-                            smem_hist[jz_r] = cutlass.Int32(0)
-                            jz_r = jz_r + cutlass.Int32(num_threads)
-                        # s_iscalars[2]/[3] (P2's cnt_lo/cnt_hi) are dead
-                        # until P2 re-inits them on the fallback path ->
-                        # free {fired flag, bin index} broadcast slots.
-                        if tidx == cutlass.Int32(0):
-                            s_iscalars[2] = cutlass.Int32(0)
-                        cute.arch.barrier()
-                        cb_r = cand_row.iterator.toint()
-                        i_r = tidx
-                        # ragged exit is fine here: no warp collectives,
-                        # only smem atomics. 4x unroll overlaps the pair
-                        # loads (same latency argument as the loader).
-                        while i_r < claimed_c:
-                            for _jr in cutlass.range_constexpr(4):
-                                i_rj = i_r + cutlass.Int32(_jr * num_threads)
-                                if i_rj < claimed_c:
-                                    pa_r = cb_r + cutlass.Int64(i_rj) * cutlass.Int64(8)
-                                    vp_r = cute.make_ptr(
-                                        cutlass.Float32,
-                                        pa_r,
-                                        cute.AddressSpace.gmem,
-                                        assumed_align=8,
-                                    )
-                                    ip_r = cute.make_ptr(
-                                        cutlass.Int32,
-                                        pa_r + cutlass.Int64(4),
-                                        cute.AddressSpace.gmem,
-                                        assumed_align=4,
-                                    )
-                                    v_r = cute.make_tensor(vp_r, cute.make_layout((1,)))[0]
-                                    x_r = cute.make_tensor(ip_r, cute.make_layout((1,)))[0]
-                                    # only the open band [t_fat, t_thin):
-                                    # >= t_thin is already counted by
-                                    # c_thin, < t_fat is below the edge.
-                                    if x_r >= cutlass.Int32(0) and v_r >= t_fat and v_r < t_thin:
-                                        b_r = cutlass.Int32((v_r - t_fat) * inv_wr)
-                                        if b_r < cutlass.Int32(0):
-                                            b_r = cutlass.Int32(0)
-                                        if b_r > cutlass.Int32(NBr - 1):
-                                            b_r = cutlass.Int32(NBr - 1)
-                                        atomicAdd(smem_hist.iterator + b_r, cutlass.Int32(1))
-                            i_r = i_r + cutlass.Int32(4 * num_threads)
-                        cute.arch.barrier()
-                        # warp-0 top-down cumulative scan (P1b idiom):
-                        # lane l owns SEG bins descending from NB-1-l*SEG;
-                        # fire at the unique bin where the cumulative
-                        # count (incl. c_thin base) crosses K, admit if
-                        # it also fits smem with a float-edge margin.
-                        if warp_id == cutlass.Int32(0):
-                            SEGr = cutlass.const_expr(8)
-                            top_r = cutlass.Int32(NBr - 1) - lane * cutlass.Int32(SEGr)
-                            seg_r = cute.make_fragment((SEGr,), cutlass.Int32)
-                            part_r = cutlass.Int32(0)
-                            for _js in cutlass.range_constexpr(SEGr):
-                                v8_r = smem_hist[top_r - cutlass.Int32(_js)]
-                                seg_r[_js] = v8_r
-                                part_r = part_r + v8_r
-                            tp_r = part_r
-                            for _os in cutlass.range_constexpr(5):
-                                ov_r = cutlass.const_expr(1 << _os)
-                                oth_r = cute.arch.shuffle_sync_up(tp_r, ov_r, mask_and_clamp=0)
-                                if lane >= cutlass.Int32(ov_r):
-                                    tp_r = tp_r + oth_r
-                            excl_r = tp_r - part_r
-                            run_r = cutlass.Int32(0)
-                            kneed = cutlass.Int32(self.top_k) - c_thin
-                            kfit = cutlass.Int32(self.kC - 64) - c_thin
-                            for _js in cutlass.range_constexpr(SEGr):
-                                run_r = run_r + seg_r[_js]
-                                cum_at = excl_r + run_r
-                                cum_bef = cum_at - seg_r[_js]
-                                if cum_bef < kneed and cum_at >= kneed:
-                                    if cum_at <= kfit:
-                                        s_iscalars[3] = top_r - cutlass.Int32(_js)
-                                        s_iscalars[2] = cutlass.Int32(1)
-                        cute.arch.barrier()
-                        if s_iscalars[2] == cutlass.Int32(1):
-                            # every thread recomputes the edge from the
-                            # uniform bin index
-                            filt_t = t_fat + cutlass.Float32(s_iscalars[3]) * w_r
-                            take_cand = cutlass.Int32(1)
+                if cutlass.const_expr(_P4_TAIL_DBG):
+                    ck0 = cute.arch.clock64()
+                    ck1 = ck0
                 if take_cand == cutlass.Int32(1):
-                    # Cooperative sentinel-skipping load of (value, index)
-                    # pairs into the P4 candidate arrays. SMEM-atomic
-                    # compaction: order is irrelevant to rank-scatter.
-                    # cs > 1: the pre-collected pairs are O(cand_count) —
-                    # row splitting buys nothing, so the LEADER loads them
-                    # alone and peers publish zero local candidates for
-                    # the DSMEM gather (take_cand is cluster-uniform: all
-                    # CTAs read the same per-row control words).
-                    if cutlass.const_expr(cluster_size > 1):
-                        if tidx == cutlass.Int32(0):
-                            s_iscalars[5] = cutlass.Int32(0)
-                        cute.arch.barrier()
-                    if tidx == cutlass.Int32(0):
-                        s_iscalars[0] = cutlass.Int32(0)
-                        s_thr[0] = filt_t
-                        s_iscalars[1] = cutlass.Int32(1)  # done: no retry
-                    cute.arch.barrier()
-                    cbase = cand_row.iterator.toint()
+                    t_lo = seed_thr_row[0]
                     lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1)
-                    i_c = tidx
-                    if cutlass.const_expr(cluster_size > 1):
-                        if cta_in_cluster != cutlass.Int32(0):
-                            # peers: skip the walk. Must stay warp-uniform
-                            # for the ballot below, so park every lane at
-                            # (claimed + lane) -> warp base == claimed.
-                            i_c = claimed_c + lane_c
-                    # Loop condition on the warp BASE index (uniform across
-                    # the warp): a ragged per-lane exit would strand the
-                    # remaining lanes in vote_ballot_sync (full-warp
-                    # collective) on the tail iteration -> device hang.
-                    #
-                    # 4x-unrolled, latency-overlapped walk. The naive form
-                    # (idx load -> dependent val load -> ballot, one entry
-                    # per thread per trip) costs ~2 serialized cold-DRAM
-                    # round trips per trip; a fat list (C(t_lo) ~ 6-12K)
-                    # then eats ~1.5us/trip (measured: wf 0.64x vs skip at
-                    # 160k). Here all 8 loads (4 entries x {idx,val}) issue
-                    # independently BEFORE the first use, so a trip costs
-                    # ~one round trip for 4x the entries. The val load is
-                    # NOT nested under the sentinel check: sentinel pairs
-                    # are mapped memory, their value is simply ignored.
-                    while (i_c - lane_c) < claimed_c:
-                        pvals = []
-                        pidxs = []
+                    vbase = cand_vals_row.iterator.toint()
+                    NBL = cutlass.const_expr(self.kNumBins)
+                    # ---- 1) stage: pure copy of the score column into
+                    # the dedicated smem region. No ballots, no atomics
+                    # (128 serialized smem atomics per trip measured
+                    # ~1.1us/1k entries), no warp-uniform constraint.
+                    # Sentinels sanitize to t_lo - 1 (rank below all live
+                    # scores; the K+64 admission bound covers K).
+                    list_used = cutlass.Int32(1)
+                    if tidx == cutlass.Int32(0):
+                        s_iscalars[2] = cutlass.Int32(0)
+                    snt = t_lo - cutlass.Float32(1.0)
+                    lmax = cutlass.Float32(self.NEG_FLT_MAX)
+                    i_s = tidx
+                    while i_s < claimed_c:
                         for _ju in cutlass.range_constexpr(4):
-                            pval = cutlass.Float32(0.0)
-                            pidx = cutlass.Int32(-1)
-                            i_cj = i_c + cutlass.Int32(_ju * num_threads)
-                            if i_cj < claimed_c:
-                                pa = cbase + cutlass.Int64(i_cj) * cutlass.Int64(8)
-                                vp_c = cute.make_ptr(
+                            i_sj = i_s + cutlass.Int32(_ju * num_threads)
+                            if i_sj < claimed_c:
+                                vp_l = cute.make_ptr(
                                     cutlass.Float32,
-                                    pa,
-                                    cute.AddressSpace.gmem,
-                                    assumed_align=8,
-                                )
-                                ip_c = cute.make_ptr(
-                                    cutlass.Int32,
-                                    pa + cutlass.Int64(4),
+                                    vbase + cutlass.Int64(i_sj) * cutlass.Int64(4),
                                     cute.AddressSpace.gmem,
                                     assumed_align=4,
                                 )
-                                pval = cute.make_tensor(vp_c, cute.make_layout((1,)))[0]
-                                pidx = cute.make_tensor(ip_c, cute.make_layout((1,)))[0]
-                            pvals.append(pval)
-                            pidxs.append(pidx)
-                        for _ju in cutlass.range_constexpr(4):
-                            keep_c = cutlass.Int32(0)
-                            if pidxs[_ju] >= cutlass.Int32(0):
-                                if pvals[_ju] >= filt_t:
-                                    keep_c = cutlass.Int32(1)
-                            # ballot-batched claim (the P4 writeback idiom):
-                            # one smem atomic per warp-trip instead of one
-                            # per survivor.
-                            mask_k = cute.arch.vote_ballot_sync(keep_c != cutlass.Int32(0))
-                            if mask_k != cutlass.Uint32(0):
-                                nk = cutlass.Int32(cute.arch.popc(mask_k))
-                                lmk = (
-                                    cutlass.Uint32(1) << cutlass.Uint32(lane_c)
-                                ) - cutlass.Uint32(1)
-                                offk = cutlass.Int32(cute.arch.popc(mask_k & lmk))
-                                bk = cutlass.Int32(0)
+                                v_l = cute.make_tensor(vp_l, cute.make_layout((1,)))[0]
+                                if v_l < t_lo:
+                                    v_l = snt
+                                smem_list[i_sj] = v_l
+                                lmax = cute.arch.fmax(lmax, v_l)
+                        i_s = i_s + cutlass.Int32(4 * num_threads)
+                    wmax_l = self.warp_reduce_max_f32(lmax)
+                    if lane == cutlass.Int32(0):
+                        smem_wmax[warp_id] = wmax_l
+                    cute.arch.barrier()
+                    # ---- 2) zooming histogram on the smem scores: find
+                    # the edge whose exact count lands in [K, kC], as
+                    # close to K as the bins allow. Value-linear bins
+                    # collapse on long-tailed logits (vmax is an extreme
+                    # outlier), so each round re-bins INSIDE the crossing
+                    # bin - one cheap smem pass per round, NBL^3 total
+                    # resolution; only genuine ties fall through.
+                    vmax = cutlass.Float32(self.NEG_FLT_MAX)
+                    for _wr in cutlass.range_constexpr(self.num_warps):
+                        vmax = cute.arch.fmax(vmax, smem_wmax[_wr])
+                    rng_l = vmax - t_lo
+                    if rng_l <= cutlass.Float32(0.0):
+                        rng_l = cutlass.Float32(1.0)
+                    r_lo = t_lo
+                    r_w = rng_l / cutlass.Float32(NBL)
+                    w0 = r_w  # round-0 bin width (anchor edge recompute)
+                    base_c = cutlass.Int32(0)
+                    t_star = t_lo
+                    searching = cutlass.Int32(1)
+                    if tidx == cutlass.Int32(0):
+                        # smem_wcnt[1]: round-0 anchor bin (the ~3K
+                        # crossing; [0] is the descend base). -1 = not
+                        # found -> anchor falls back to t_star. P4
+                        # clobbers smem_wcnt only AFTER we read this.
+                        smem_wcnt[1] = cutlass.Int32(-1)
+                    for _rd in cutlass.range_constexpr(3):
+                        if searching == cutlass.Int32(1):
+                            jz_l = tidx
+                            while jz_l < cutlass.Int32(NBL):
+                                smem_hist[jz_l] = cutlass.Int32(0)
+                                jz_l = jz_l + cutlass.Int32(num_threads)
+                            if tidx == cutlass.Int32(0):
+                                s_iscalars[2] = cutlass.Int32(0)
+                            cute.arch.barrier()
+                            inv_wr = cutlass.Float32(1.0) / r_w
+                            r_hi = r_lo + r_w * cutlass.Float32(NBL)
+                            i_h = tidx
+                            while i_h < claimed_c:
+                                v_h = smem_list[i_h]
+                                if v_h >= r_lo and v_h < r_hi:
+                                    b_h = cutlass.Int32((v_h - r_lo) * inv_wr)
+                                    if b_h < cutlass.Int32(0):
+                                        b_h = cutlass.Int32(0)
+                                    if b_h > cutlass.Int32(NBL - 1):
+                                        b_h = cutlass.Int32(NBL - 1)
+                                    atomicAdd(smem_hist.iterator + b_h, cutlass.Int32(1))
+                                i_h = i_h + cutlass.Int32(num_threads)
+                            cute.arch.barrier()
+                            # warp-0 top-down cumulative scan (P1b idiom)
+                            if warp_id == cutlass.Int32(0):
+                                SEGL = cutlass.const_expr(NBL // self.WARP_SIZE)
+                                top_l = cutlass.Int32(NBL - 1) - lane * cutlass.Int32(SEGL)
+                                seg_l = cute.make_fragment((SEGL,), cutlass.Int32)
+                                part_l = cutlass.Int32(0)
+                                for _js in cutlass.range_constexpr(SEGL):
+                                    v8_l = smem_hist[top_l - cutlass.Int32(_js)]
+                                    seg_l[_js] = v8_l
+                                    part_l = part_l + v8_l
+                                tp_l = part_l
+                                for _os in cutlass.range_constexpr(5):
+                                    ov_l = cutlass.const_expr(1 << _os)
+                                    oth_l = cute.arch.shuffle_sync_up(tp_l, ov_l, mask_and_clamp=0)
+                                    if lane >= cutlass.Int32(ov_l):
+                                        tp_l = tp_l + oth_l
+                                excl_l = tp_l - part_l
+                                kneed = cutlass.Int32(self.top_k) - base_c
+                                kfit = cutlass.Int32(self.kC) - base_c
+                                run_l = cutlass.Int32(0)
+                                for _js in cutlass.range_constexpr(SEGL):
+                                    run_l = run_l + seg_l[_js]
+                                    cum_at = excl_l + run_l
+                                    cum_bef = cum_at - seg_l[_js]
+                                    if cum_bef < kneed and cum_at >= kneed:
+                                        s_iscalars[3] = top_l - cutlass.Int32(_js)
+                                        smem_wcnt[0] = cum_bef
+                                        if cum_at <= kfit:
+                                            s_iscalars[2] = cutlass.Int32(1)
+                                        else:
+                                            s_iscalars[2] = cutlass.Int32(2)
+                                    if cutlass.const_expr(_rd == 0):
+                                        # closed-loop ANCHOR: the ~3K
+                                        # crossing edge. Publishing the
+                                        # exact k-th made the next-step
+                                        # down-guard target only 4K -
+                                        # slope noise then undershoots K
+                                        # and forces ~26us fallbacks; a
+                                        # 3K-count anchor restores the
+                                        # old edge semantics (~12K guard
+                                        # target) at zero extra passes.
+                                        anch_n = cutlass.const_expr(
+                                            min(3 * self.top_k, (self.top_k + self.kC) // 2)
+                                        )
+                                        if cum_bef < cutlass.Int32(
+                                            anch_n
+                                        ) and cum_at >= cutlass.Int32(anch_n):
+                                            smem_wcnt[1] = top_l - cutlass.Int32(_js)
+                            cute.arch.barrier()
+                            st_l = s_iscalars[2]
+                            if st_l == cutlass.Int32(1):
+                                t_star = r_lo + cutlass.Float32(s_iscalars[3]) * r_w
+                                searching = cutlass.Int32(0)
+                            if st_l == cutlass.Int32(2):
+                                base_c = base_c + smem_wcnt[0]
+                                r_lo = r_lo + cutlass.Float32(s_iscalars[3]) * r_w
+                                r_w = r_w / cutlass.Float32(NBL)
+                                if r_w <= cutlass.Float32(0.0):
+                                    searching = cutlass.Int32(0)
+                            if st_l == cutlass.Int32(0):
+                                searching = cutlass.Int32(0)
+                            cute.arch.barrier()
+                    if tidx == cutlass.Int32(0):
+                        if s_iscalars[2] != cutlass.Int32(1):
+                            s_iscalars[2] = cutlass.Int32(0)
+                    cute.arch.barrier()
+                    fired_l = s_iscalars[2]
+                    t_anch = t_star
+                    anch_b = smem_wcnt[1]
+                    if anch_b >= cutlass.Int32(0):
+                        t_anch = t_lo + cutlass.Float32(anch_b) * w0
+                    if fired_l == cutlass.Int32(1):
+                        # ---- 3) compact survivors into the standard kC
+                        # buffers: score from smem, position streamed
+                        # from the gmem index column. The four unrolled
+                        # sub-ballots merge into ONE atomic per warp per
+                        # trip (the per-sub-ballot atomic serialized 128
+                        # adds per trip across 32 warps).
+                        if tidx == cutlass.Int32(0):
+                            s_iscalars[0] = cutlass.Int32(0)
+                            s_thr[0] = t_anch  # closed-loop anchor
+                            s_iscalars[1] = cutlass.Int32(1)  # done
+                        cute.arch.barrier()
+                        i_c = tidx
+                        while (i_c - lane_c) < claimed_c:
+                            pvals = []
+                            pidxs = []
+                            keeps = []
+                            for _ju in cutlass.range_constexpr(4):
+                                i_cj = i_c + cutlass.Int32(_ju * num_threads)
+                                pval = cutlass.Float32(self.NEG_FLT_MAX)
+                                # vals slot carries the LIST INDEX - a
+                                # register, not a second (cold) gmem pass
+                                # over the position column; the post-P4
+                                # repair gathers true positions for the K
+                                # winners only.
+                                pidx = i_cj
+                                keep = cutlass.Int32(0)
+                                if i_cj < claimed_c:
+                                    pval = smem_list[i_cj]
+                                    if pval >= t_star:
+                                        keep = cutlass.Int32(1)
+                                pvals.append(pval)
+                                pidxs.append(pidx)
+                                keeps.append(keep)
+                            m0 = cute.arch.vote_ballot_sync(keeps[0] != cutlass.Int32(0))
+                            m1 = cute.arch.vote_ballot_sync(keeps[1] != cutlass.Int32(0))
+                            m2 = cute.arch.vote_ballot_sync(keeps[2] != cutlass.Int32(0))
+                            m3 = cute.arch.vote_ballot_sync(keeps[3] != cutlass.Int32(0))
+                            nk = cutlass.Int32(
+                                cute.arch.popc(m0)
+                                + cute.arch.popc(m1)
+                                + cute.arch.popc(m2)
+                                + cute.arch.popc(m3)
+                            )
+                            bk = cutlass.Int32(0)
+                            if nk > cutlass.Int32(0):
                                 if lane_c == cutlass.Int32(0):
                                     bk = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), nk)
                                 bk = cute.arch.shuffle_sync(bk, cutlass.Int32(0))
-                                wpos = bk + offk
-                                if keep_c != cutlass.Int32(0) and wpos < cutlass.Int32(self.kC):
-                                    smem_keys[wpos] = pvals[_ju]
-                                    smem_vals[wpos] = pidxs[_ju]
-                        i_c = i_c + cutlass.Int32(4 * num_threads)
-                    cute.arch.barrier()
-                    if cutlass.const_expr(cluster_size == 1):
-                        # Hard net for the refined edge: its count is
-                        # float-boundary approximate. If survivors
-                        # overflowed kC the smem set is truncated
-                        # (value-blind) -> demote to the fallback
-                        # re-scan. Exact-rung filters cannot overflow
-                        # (filt_c <= kC checked); s_iscalars[0] is
-                        # CTA-uniform after the barrier above.
-                        if s_iscalars[0] > cutlass.Int32(self.kC):
+                                lmk = (
+                                    cutlass.Uint32(1) << cutlass.Uint32(lane_c)
+                                ) - cutlass.Uint32(1)
+                                off = bk
+                                for _ju in cutlass.range_constexpr(4):
+                                    mj = (
+                                        m0
+                                        if _ju == 0
+                                        else m1
+                                        if _ju == 1
+                                        else m2
+                                        if _ju == 2
+                                        else m3
+                                    )
+                                    if keeps[_ju] != cutlass.Int32(0):
+                                        wpos = off + cutlass.Int32(cute.arch.popc(mj & lmk))
+                                        if wpos < cutlass.Int32(self.kC):
+                                            smem_keys[wpos] = pvals[_ju]
+                                            smem_vals[wpos] = pidxs[_ju]
+                                    off = off + cutlass.Int32(cute.arch.popc(mj))
+                            i_c = i_c + cutlass.Int32(4 * num_threads)
+                        cute.arch.barrier()
+                        # hard net: histogram binning and the compact
+                        # test round float edges independently
+                        cnt_l = s_iscalars[0]
+                        if cnt_l < cutlass.Int32(self.top_k) or cnt_l > cutlass.Int32(self.kC):
                             take_cand = cutlass.Int32(0)
+                            list_used = cutlass.Int32(0)
                             if tidx == cutlass.Int32(0):
                                 s_iscalars[1] = cutlass.Int32(0)
+                                s_iscalars[2] = cutlass.Int32(-1)
+                                s_iscalars[3] = cutlass.Int32(-1)
                             cute.arch.barrier()
-                    if cutlass.const_expr(cluster_size > 1):
-                        # leader's local count feeds the DSMEM gather
-                        if is_leader:
-                            if tidx == cutlass.Int32(0):
-                                s_iscalars[5] = s_iscalars[0]
+                    if fired_l == cutlass.Int32(0):
+                        take_cand = cutlass.Int32(0)
+                        list_used = cutlass.Int32(0)
+                        if tidx == cutlass.Int32(0):
+                            s_iscalars[2] = cutlass.Int32(-1)
+                            s_iscalars[3] = cutlass.Int32(-1)
                         cute.arch.barrier()
+                if cutlass.const_expr(_P4_TAIL_DBG):
+                    ck1 = cute.arch.clock64()
             if take_cand == cutlass.Int32(0):
                 # Stage this CTA's slice into SMEM once before Phase 2's
                 # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
@@ -5567,6 +5646,9 @@ def _run_phases(
             # Pre-init cand_count_p4 so CuTe DSL sees a stable Int32 type
             # across the runtime ``if is_leader:`` branch in cs>1 mode
             # (DSL forbids first-assigning a variable inside a dynamic if).
+            ck2 = cutlass.Int64(0)
+            if cutlass.const_expr(_P4_TAIL_DBG):
+                ck2 = cute.arch.clock64()
             cand_count_p4 = cutlass.Int32(0)
             if cutlass.const_expr(cluster_size == 1):
                 # cs=1: the single CTA per row IS the leader.
@@ -5601,6 +5683,31 @@ def _run_phases(
                         warp_id,
                         lane,
                     )
+                if cutlass.const_expr(
+                    self.use_ext_cand and self.use_ext_counts and self.dtype == cutlass.Float32
+                ):
+                    # List rows: the compact stored LIST INDICES in the
+                    # vals slots (saving a second cold gmem pass over the
+                    # position column). Swap them for true positions with
+                    # K fully-parallel gathers. Must precede the xstate
+                    # publish, which reads output slot K-1 as a position.
+                    if list_used == cutlass.Int32(1):
+                        io_r = tidx
+                        while io_r < cutlass.Int32(self.top_k):
+                            li_r = output_indices_row[io_r]
+                            if li_r >= cutlass.Int32(0) and li_r < claimed_c:
+                                ip_r = cute.make_ptr(
+                                    cutlass.Int32,
+                                    cand_idx_row.iterator.toint()
+                                    + cutlass.Int64(li_r) * cutlass.Int64(4),
+                                    cute.AddressSpace.gmem,
+                                    assumed_align=4,
+                                )
+                                output_indices_row[io_r] = cute.make_tensor(
+                                    ip_r, cute.make_layout((1,))
+                                )[0]
+                            io_r = io_r + cutlass.Int32(num_threads)
+                        cute.arch.barrier()
                 if cutlass.const_expr(self.emit_xstate):
                     # Closed-loop state (interface v2): [0] valid, [1] kth
                     # proxy (= accepted threshold; the tie-fill makes it a
@@ -5609,8 +5716,24 @@ def _run_phases(
                     # seed rung group from these.
                     if tidx == cutlass.Int32(0):
                         xstate_row[0] = cutlass.Float32(1.0)
-                        xstate_row[1] = s_thr[0]
-                        xstate_row[2] = s_thr[0]
+                        thr_pub = s_thr[0]
+                        anch_pub = s_thr[0]
+                        if list_used == cutlass.Int32(1):
+                            # list rows: rank-scatter's output is rank-
+                            # ordered, so slot K-1 holds the exact k-th
+                            # boundary - a tighter, healthier closed-loop
+                            # anchor than the loose collect line.
+                            idx_k = output_indices_row[cutlass.Int32(self.top_k - 1)]
+                            if idx_k >= cutlass.Int32(0) and idx_k < N:
+                                thr_pub = cutlass.Float32(input_row[idx_k])
+                        xstate_row[1] = thr_pub
+                        xstate_row[2] = anch_pub
+                        if cutlass.const_expr(_P4_TAIL_DBG):
+                            ck3 = cute.arch.clock64()
+                            xstate_row[4] = cutlass.Float32(cutlass.Int32(ck1 - ck0))  # walk+flags
+                            xstate_row[5] = cutlass.Float32(cutlass.Int32(ck2 - ck1))  # P2/P3 gap
+                            xstate_row[6] = cutlass.Float32(cutlass.Int32(ck3 - ck2))  # Phase 4
+                            xstate_row[7] = s_thr[1]  # cnt_strad
                         # cand_count_p4 = pre-P4 snapshot (P4 repurposes
                         # the s_iscalars slots).
                         xstate_row[3] = cutlass.Float32(cand_count_p4)
@@ -5726,7 +5849,8 @@ def __call__(
         seed_thr: cute.Tensor = None,  # [num_rows, 3] fp32 (ext counts)
         seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 (ext counts)
         xstate: cute.Tensor = None,  # [num_rows, 8] fp32 (emit_xstate)
-        cand: cute.Tensor = None,  # [num_rows, CAP*2] int32 (ext cand)
+        cand_vals: cute.Tensor = None,  # [num_rows, CAP] fp32 (ext cand)
+        cand_idx: cute.Tensor = None,  # [num_rows, CAP] int32 (ext cand)
         cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 (ext cand)
     ):
         num_rows = input_data.shape[0]
@@ -5753,7 +5877,8 @@ def __call__(
             seed_thr,
             seed_counts,
             xstate,
-            cand,
+            cand_vals,
+            cand_idx,
             cand_ctl,
         ).launch(
             grid=(total_ctas, 1, 1),
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index ed79a19cbb23..f464f01a2031 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -151,9 +151,16 @@ def _compile(
         if use_ext_counts
         else None
     )
-    cand_fake = (
+    cand_vals_fake = (
         cute.runtime.make_fake_compact_tensor(
-            cutlass.Int32, (n_rows, cand_cap * 2), stride_order=(1, 0), assumed_align=8
+            cutlass.Float32, (n_rows, cand_cap), stride_order=(1, 0), assumed_align=4
+        )
+        if use_ext_cand
+        else None
+    )
+    cand_idx_fake = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Int32, (n_rows, cand_cap), stride_order=(1, 0), assumed_align=4
         )
         if use_ext_cand
         else None
@@ -213,7 +220,8 @@ def _compile(
         seed_thr=seed_thr_fake,
         seed_counts=seed_counts_fake,
         xstate=xstate_fake,
-        cand=cand_fake,
+        cand_vals=cand_vals_fake,
+        cand_idx=cand_idx_fake,
         cand_ctl=cand_ctl_fake,
         options="--enable-tvm-ffi",
     )
@@ -351,6 +359,7 @@ def derive_seed_rungs(
     prev_counts: "torch.Tensor | None" = None,
     count_octaves: float = 2.0,
     fallback_spread: float = 0.5,
+    top_k: "int | None" = None,
 ) -> torch.Tensor:
     """Host-side slope-adaptive seed rung derivation (waterfall closed loop).
 
@@ -370,6 +379,7 @@ def derive_seed_rungs(
     """
     if prev_sthr is None or prev_counts is None:
         d = torch.full_like(prev_thr, fallback_spread)
+        d_lo = d
     else:
         c_lo = prev_counts[:, 0].float().clamp(min=1.0)
         c_hi = prev_counts[:, 2].float().clamp(min=1.0)
@@ -380,7 +390,19 @@ def derive_seed_rungs(
             count_octaves / slope.clamp(min=0.05),
             torch.full_like(slope, fallback_spread),
         ).clamp(0.1, 4.0)
-    return torch.stack([prev_thr - d, prev_thr, prev_thr + d], dim=1).contiguous()
+        # undershoot hysteresis: a row whose PREVIOUS loose rung caught
+        # fewer than K cannot use its list (fallback ~26us); widen only
+        # that row's next down-guard by +2 octaves. Rows in band keep the
+        # tight spread (a fat list taxes every step's walk).
+        oct_lo = torch.full_like(slope, count_octaves)
+        if top_k is not None:
+            oct_lo = torch.where(prev_counts[:, 0].float() < float(top_k), oct_lo + 2.0, oct_lo)
+        d_lo = torch.where(
+            slope > 0.05,
+            oct_lo / slope.clamp(min=0.05),
+            torch.full_like(slope, fallback_spread),
+        ).clamp(0.1, 6.0)
+    return torch.stack([prev_thr - d_lo, prev_thr, prev_thr + d], dim=1).contiguous()
 
 
 def emu_seed_counts(
@@ -414,8 +436,8 @@ def emu_cand(
     next_n: int = 1,
     compress_ratio: int = 1,
     sentinel_pad: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """L2 emu: unordered candidate pre-collect.
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """L2 emu: unordered candidate pre-collect (SoA).
 
     Unordered (value fp32-bits, index) pairs of all valid positions
     >= t_0 = seed_thr[:, 0]; ctl = {claimed, void}. claimed may
@@ -428,9 +450,9 @@ def emu_cand(
     dev = logits.device
     n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio)
     lf = logits.to(torch.float32)
-    cand = torch.full((R, cap * 2), -1, dtype=torch.int32, device=dev)
+    cand_vals = torch.full((R, cap), float("-inf"), dtype=torch.float32, device=dev)
+    cand_idx = torch.full((R, cap), -1, dtype=torch.int32, device=dev)
     ctl = torch.zeros((R, 2), dtype=torch.int32, device=dev)
-    pairs = cand.view(R, cap, 2)
     for r in range(R):
         ne = int(n_eff[r])
         hits = torch.nonzero(lf[r, :ne] >= seed_thr[r, 0], as_tuple=False).flatten()
@@ -447,11 +469,11 @@ def emu_cand(
         claimed = int(ent.numel())
         nwr = min(claimed, cap)
         live = ent[:nwr] >= 0
-        pairs[r, :nwr, 1] = ent[:nwr].to(torch.int32)
-        pairs[r, :nwr, 0][live] = lf[r, ent[:nwr][live]].view(torch.int32)
+        cand_idx[r, :nwr] = ent[:nwr].to(torch.int32)
+        cand_vals[r, :nwr][live] = lf[r, ent[:nwr][live]]
         ctl[r, 0] = claimed
         ctl[r, 1] = 1 if claimed > cap else 0
-    return cand, ctl
+    return cand_vals, cand_idx, ctl
 
 
 def enc_ordered_f32(t: torch.Tensor) -> torch.Tensor:
@@ -508,7 +530,8 @@ def gvr_topk_decode(
     seed_thr: Optional[torch.Tensor] = None,
     seed_counts: Optional[torch.Tensor] = None,
     xstate: Optional[torch.Tensor] = None,
-    cand: Optional[torch.Tensor] = None,
+    cand_vals: Optional[torch.Tensor] = None,
+    cand_idx: Optional[torch.Tensor] = None,
     cand_ctl: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """CuTe DSL GVR Top-K wrapper with every tuning knob exposed.
@@ -606,24 +629,29 @@ def gvr_topk_decode(
             and seed_counts.is_contiguous()
             and seed_counts.shape == (num_rows, 3)
         ), "seed_counts must be contiguous CUDA int32 [num_rows, 3]"
-    use_ext_cand = cand is not None and cand_ctl is not None
+    use_ext_cand = cand_vals is not None and cand_idx is not None and cand_ctl is not None
     cand_cap = 5120
     if use_ext_cand:
         assert (
-            cand.dtype == torch.int32
-            and cand.is_cuda
-            and cand.is_contiguous()
-            and cand.dim() == 2
-            and cand.shape[0] == num_rows
-            and cand.shape[1] % 2 == 0
-        ), "cand must be contiguous CUDA int32 [num_rows, CAP*2]"
+            cand_vals.dtype == torch.float32
+            and cand_vals.is_cuda
+            and cand_vals.is_contiguous()
+            and cand_vals.dim() == 2
+            and cand_vals.shape[0] == num_rows
+        ), "cand_vals must be contiguous CUDA fp32 [num_rows, CAP]"
+        assert (
+            cand_idx.dtype == torch.int32
+            and cand_idx.is_cuda
+            and cand_idx.is_contiguous()
+            and cand_idx.shape == cand_vals.shape
+        ), "cand_idx must be contiguous CUDA int32 [num_rows, CAP]"
         assert (
             cand_ctl.dtype == torch.int32
             and cand_ctl.is_cuda
             and cand_ctl.is_contiguous()
             and cand_ctl.shape == (num_rows, 2)
         ), "cand_ctl must be contiguous CUDA int32 [num_rows, 2]"
-        cand_cap = cand.shape[1] // 2
+        cand_cap = cand_vals.shape[1]
     emit_xstate = xstate is not None
     if emit_xstate:
         assert (
@@ -730,7 +758,8 @@ def gvr_topk_decode(
         seed_thr if use_ext_counts else None,
         seed_counts if use_ext_counts else None,
         xstate if emit_xstate else None,
-        cand if use_ext_cand else None,
+        cand_vals if use_ext_cand else None,
+        cand_idx if use_ext_cand else None,
         cand_ctl if use_ext_cand else None,
     )
     if return_output_values:

From 56775417d4366a89cc04a14afcc6f301283f1d10 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 23 Jul 2026 09:46:59 -0700
Subject: [PATCH 019/117] [None][perf] GVR L2 v4: emitter-counted lines,
 known-counts admission

The emitter (indexer GEMM epilogue; emulated host-side in the bench
harness) now counts the two tighter lines while writing the SoA list -
two extra compares per EMITTED element only, against the full-row L1
pass this replaces - and the control words widen to {n0, void, n1,
n2}. The topk side enters with every count known and the whole list
path collapses to a scalar state machine:

- some line's count lands in the acceptance band [K, B*]: cut at the
  TIGHTEST such line, ONE filtered gmem pass straight into the
  candidate buffers (positions deferred as list indices; the position
  column is gathered only for the K winners after Phase 4). Counts
  and load predicates are the same comparison, so line cuts need no
  overflow net at all.
- the band is straddled or overshot by every line: a zooming
  histogram over the gmem list CLAMPED between the two known bracket
  lines finds an in-band edge (narrow domain - no long-tail bin
  collapse; the all-above case takes one max pass first).
- void, or n0 < K + 64 (the emitter sentinel bound, proving live
  coverage of K): fallback.

The dedicated smem staging region is deleted (frees 96-128KB; the
kernel's smem drops back to the pre-list footprint), and B* / kC
become constructor knobs (accept_cap, kc_override) for the band
search. Closed loop publishes the exact k-th (rank-ordered output
slot K-1) and the loosest in-band line as the anchor.

Line placement is a searched host policy (derive_seed_lines_v4):
count targets (t0, t1, t2), grid-searched on real chains =
(4096, 3584, 1536) for short domains / (12288, 5120, 2048) for long;
physical kC stays 5120 (8192 measured no gain).

Validation (B200): five admission states each exercised exact
(hit-t2/hit-t1/bracketed-histogram/above-t2-histogram/fallback),
24/24 exact at cs=1/4/8; real-chain device-level vs the block-skip
arm: 160k 1.09x (first config of this lineage to beat the rung-based
1.05x), 640k 1.39x (residual: 2-3 volatile rows/step whose collection
count escapes any placement - a host anchor-policy iteration item).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 532 +++++++++---------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  60 +-
 2 files changed, 320 insertions(+), 272 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 7993a773895a..a2e64640704c 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -264,6 +264,8 @@ def __init__(
         seqlen_sorted: bool = False,
         kc_diet: Optional[bool] = None,
         enable_r0: bool = True,
+        accept_cap: "int | None" = None,
+        kc_override: "int | None" = None,
         r0_qfracs: Optional[tuple] = None,
         mt_unroll: int = 4,
         p1b_cache: Optional[bool] = None,
@@ -565,6 +567,13 @@ def __init__(
         # <= cand_cap and the collect rung's exact count is in [K, kC];
         # ineligible rows fall through to the ext-counts path.
         self.use_ext_cand = bool(use_ext_cand)
+        if kc_override is not None:
+            # physical candidate-buffer capacity override (B* search)
+            self.kC = int(kc_override)
+        # acceptance band top B*: a cut whose count fits [K, B*] goes
+        # straight to Phase 4. Cost-bounded (refine ~1.9us/1k cands vs
+        # ~21-30us fallback), physically bounded by kC.
+        self.accept_cap = int(accept_cap) if accept_cap is not None else self.kC
         self.cand_cap = int(cand_cap)
         # list path: the score column is staged into a DEDICATED smem
         # region sized cand_cap fp32 (96KB at 24576). Budget note: this
@@ -4464,21 +4473,6 @@ def run_one_row(
         else:
             smem_input = None
 
-        # List-path score staging (SoA ext cand, cs=1, fp32 only).
-        if cutlass.const_expr(
-            self.use_ext_cand
-            and self.use_ext_counts
-            and cluster_size == 1
-            and self.dtype == cutlass.Float32
-        ):
-            smem_list = smem.allocate_tensor(
-                element_type=cutlass.Float32,
-                layout=cute.make_ordered_layout((self.list_cap,), order=(0,)),
-                byte_alignment=128,
-            )
-        else:
-            smem_list = None
-
         # op#26 R0 admission scratch (single-CTA fast path). Allocated only
         # when enable_r0; None otherwise so the base SMEM layout is byte-for-
         # byte unchanged and these propagate harmlessly through _run_phases'
@@ -4637,7 +4631,6 @@ def run_one_row(
                         xstate_row=xstate_row,
                         cand_vals_row=cand_vals_row,
                         cand_idx_row=cand_idx_row,
-                        smem_list=smem_list,
                         cand_ctl_row=cand_ctl_row,
                         smem_active=smem_active,
                         s_active_cnt=s_active_cnt,
@@ -4688,7 +4681,6 @@ def run_one_row(
                             xstate_row=xstate_row,
                             cand_vals_row=cand_vals_row,
                             cand_idx_row=cand_idx_row,
-                            smem_list=smem_list,
                             cand_ctl_row=cand_ctl_row,
                             smem_active=smem_active,
                             s_active_cnt=s_active_cnt,
@@ -4736,7 +4728,6 @@ def run_one_row(
                     xstate_row=xstate_row,
                     cand_vals_row=cand_vals_row,
                     cand_idx_row=cand_idx_row,
-                    smem_list=smem_list,
                     cand_ctl_row=cand_ctl_row,
                     smem_active=smem_active,
                     s_active_cnt=s_active_cnt,
@@ -4787,7 +4778,6 @@ def _run_phases(
         xstate_row=None,  # emit_xstate: this row's [8] fp32 state slot
         cand_vals_row=None,  # ext cand: this row's [CAP] fp32 scores
         cand_idx_row=None,  # ext cand: this row's [CAP] int32 positions
-        smem_list=None,  # list path: [list_cap] fp32 score staging
         cand_ctl_row=None,  # ext cand: this row's [2] int32 {claimed, void}
         smem_active=None,
         s_active_cnt=None,
@@ -4844,7 +4834,7 @@ def _run_phases(
                 void_p = cutlass.Int32(cand_ctl_row[1])
                 if (
                     void_p == cutlass.Int32(0)
-                    and claimed_p >= cutlass.Int32(self.top_k)
+                    and claimed_p >= cutlass.Int32(self.top_k + 64)
                     and claimed_p <= cutlass.Int32(self.list_cap)
                     and seed_thr_row[0] < cutlass.Float32(1e37)
                 ):
@@ -4942,17 +4932,24 @@ def _run_phases(
                     if cutlass.const_expr(self.emit_xstate):
                         xstate_row[0] = cutlass.Float32(0.0)  # degenerate
         else:
-            # ---- List path (SoA): complete candidate list -> on-chip ----
-            # Admission is the entry count alone: K <= claimed <= list_cap
-            # (and no overflow) guarantees the true top-K lives in the
-            # list (it holds every position >= t_lo = seed_thr[0]). The
-            # score column is staged into the slice-cache smem region
-            # (fp32 specs; the fallback re-scan never runs on list rows,
-            # so the alias is safe), one smem histogram finds an edge
-            # whose exact count lands in [K, kC], survivors compact into
-            # the standard kC candidate buffers and the stock Phase 4
-            # finishes. All gating values are CTA-uniform (gmem control
-            # words / smem scalars), so the barriers stay convergent.
+            # ---- List path v4: known-counts admission ----
+            # The emitter wrote the SoA list (score column + position
+            # column, sentinel score -inf) collected at t0 = seed_thr[0]
+            # and COUNTED the two tighter lines on the way out: the
+            # control words carry {n0, void, n1, n2} with n_i = #(>= t_i),
+            # n0 >= n1 >= n2. Admission and cut selection are pure scalar
+            # lookups - no in-kernel counting, no staging, no gamble:
+            #   1. some n_i lands in the acceptance band [K, B*]  ->
+            #      cut at the TIGHTEST such line, ONE filtered pass.
+            #   2. the band is straddled or overshot by every line ->
+            #      a histogram over the gmem list CLAMPED between the
+            #      two known bracket lines finds an in-band edge (the
+            #      narrow domain kills the long-tail bin collapse).
+            #   3. void, or n0 < K + 64 (64 = emitter sentinel bound,
+            #      so live coverage of K is proven)  ->  fallback.
+            # The K+64 slack also lets every accepted cut load run
+            # WITHOUT any overflow net: counts and load predicates are
+            # the same comparison on the same data.
             # cs>1 and 16-bit dtypes keep the plain fallback.
             take_cand = cutlass.Int32(0)
             list_used = cutlass.Int32(0)  # list path taken (xstate publish)
@@ -4965,265 +4962,264 @@ def _run_phases(
             ):
                 claimed_c = cutlass.Int32(cand_ctl_row[0])
                 void_c = cutlass.Int32(cand_ctl_row[1])
-                # claimed >= K + 64: sentinel slots are bounded by the
-                # emitter contract (<= 64 pad), so this guarantees the
-                # LIVE count covers K without counting anything.
+                n1_c = cutlass.Int32(cand_ctl_row[2])
+                n2_c = cutlass.Int32(cand_ctl_row[3])
+                bstar = cutlass.const_expr(min(self.accept_cap, self.kC))
+                if cutlass.const_expr(_P4_TAIL_DBG):
+                    ck0 = cute.arch.clock64()
+                    ck1 = ck0
+                usable = cutlass.Int32(0)
                 if (
                     void_c == cutlass.Int32(0)
                     and claimed_c >= cutlass.Int32(self.top_k + 64)
                     and claimed_c <= cutlass.Int32(self.list_cap)
                 ):
+                    usable = cutlass.Int32(1)
+                # pre-declared: the DSL forbids first-assigning inside a
+                # dynamic if when read outside it
+                kK_l = cutlass.Int32(self.top_k)
+                bs_l = cutlass.Int32(bstar)
+                cut_t = cutlass.Float32(0.0)
+                have = cutlass.Int32(0)
+                anch_t = cutlass.Float32(0.0)
+                if usable == cutlass.Int32(1):
+                    # cut = tightest line whose count is in [K, B*];
+                    # anchor (closed-loop publish) = loosest such line.
+                    if n2_c >= kK_l and n2_c <= bs_l:
+                        cut_t = seed_thr_row[2]
+                        anch_t = seed_thr_row[2]
+                        have = cutlass.Int32(1)
+                    if n1_c >= kK_l and n1_c <= bs_l:
+                        if have == cutlass.Int32(0):
+                            cut_t = seed_thr_row[1]
+                        anch_t = seed_thr_row[1]
+                        have = cutlass.Int32(1)
+                    if claimed_c <= bs_l:
+                        if have == cutlass.Int32(0):
+                            cut_t = seed_thr_row[0]
+                        anch_t = seed_thr_row[0]
+                        have = cutlass.Int32(1)
+                    if have == cutlass.Int32(0):
+                        # No line in band: bracket the cut between the
+                        # two known lines that straddle it and find an
+                        # in-band edge with a clamped gmem histogram.
+                        #   all counts > B*      -> (t2, +inf): one max
+                        #                           pass supplies the top;
+                        #   n1 > B* and n2 < K   -> (t1, t2);
+                        #   n0 > B* and n1 < K   -> (t0, t1).
+                        vbase = cand_vals_row.iterator.toint()
+                        b_lo = seed_thr_row[0]
+                        b_hi = seed_thr_row[1]
+                        base_c = n1_c
+                        if n1_c > bs_l:
+                            b_lo = seed_thr_row[1]
+                            b_hi = seed_thr_row[2]
+                            base_c = n2_c
+                        need_max = cutlass.Int32(0)
+                        if n2_c > bs_l:
+                            b_lo = seed_thr_row[2]
+                            base_c = cutlass.Int32(0)
+                            need_max = cutlass.Int32(1)
+                        if need_max == cutlass.Int32(1):
+                            lmax = cutlass.Float32(self.NEG_FLT_MAX)
+                            i_m = tidx
+                            while i_m < claimed_c:
+                                for _ju in cutlass.range_constexpr(4):
+                                    i_mj = i_m + cutlass.Int32(_ju * num_threads)
+                                    if i_mj < claimed_c:
+                                        vp_m = cute.make_ptr(
+                                            cutlass.Float32,
+                                            vbase + cutlass.Int64(i_mj) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        lmax = cute.arch.fmax(
+                                            lmax, cute.make_tensor(vp_m, cute.make_layout((1,)))[0]
+                                        )
+                                i_m = i_m + cutlass.Int32(4 * num_threads)
+                            wmax_l = self.warp_reduce_max_f32(lmax)
+                            if lane == cutlass.Int32(0):
+                                smem_wmax[warp_id] = wmax_l
+                            cute.arch.barrier()
+                            vmax_l = cutlass.Float32(self.NEG_FLT_MAX)
+                            for _wr in cutlass.range_constexpr(self.num_warps):
+                                vmax_l = cute.arch.fmax(vmax_l, smem_wmax[_wr])
+                            b_hi = vmax_l + cutlass.Float32(1e-3)
+                        # zooming clamped histogram (up to 3 rounds; the
+                        # bracket is narrow so round 0 almost always
+                        # fires). State broadcasts ride s_iscalars[2]/[3]
+                        # (P2's slots, re-sentineled on every exit path).
+                        NBL = cutlass.const_expr(self.kNumBins)
+                        r_lo = b_lo
+                        r_w = (b_hi - b_lo) / cutlass.Float32(NBL)
+                        if r_w <= cutlass.Float32(0.0):
+                            r_w = cutlass.Float32(1e-6)
+                        searching = cutlass.Int32(1)
+                        for _rd in cutlass.range_constexpr(3):
+                            if searching == cutlass.Int32(1):
+                                jz_l = tidx
+                                while jz_l < cutlass.Int32(NBL):
+                                    smem_hist[jz_l] = cutlass.Int32(0)
+                                    jz_l = jz_l + cutlass.Int32(num_threads)
+                                if tidx == cutlass.Int32(0):
+                                    s_iscalars[2] = cutlass.Int32(0)
+                                cute.arch.barrier()
+                                inv_wr = cutlass.Float32(1.0) / r_w
+                                r_hi = r_lo + r_w * cutlass.Float32(NBL)
+                                i_h = tidx
+                                while i_h < claimed_c:
+                                    for _ju in cutlass.range_constexpr(4):
+                                        i_hj = i_h + cutlass.Int32(_ju * num_threads)
+                                        if i_hj < claimed_c:
+                                            vp_h = cute.make_ptr(
+                                                cutlass.Float32,
+                                                vbase + cutlass.Int64(i_hj) * cutlass.Int64(4),
+                                                cute.AddressSpace.gmem,
+                                                assumed_align=4,
+                                            )
+                                            v_h = cute.make_tensor(vp_h, cute.make_layout((1,)))[0]
+                                            if v_h >= r_lo and v_h < r_hi:
+                                                b_h = cutlass.Int32((v_h - r_lo) * inv_wr)
+                                                if b_h < cutlass.Int32(0):
+                                                    b_h = cutlass.Int32(0)
+                                                if b_h > cutlass.Int32(NBL - 1):
+                                                    b_h = cutlass.Int32(NBL - 1)
+                                                atomicAdd(
+                                                    smem_hist.iterator + b_h, cutlass.Int32(1)
+                                                )
+                                    i_h = i_h + cutlass.Int32(4 * num_threads)
+                                cute.arch.barrier()
+                                if warp_id == cutlass.Int32(0):
+                                    SEGL = cutlass.const_expr(NBL // self.WARP_SIZE)
+                                    top_l = cutlass.Int32(NBL - 1) - lane * cutlass.Int32(SEGL)
+                                    seg_l = cute.make_fragment((SEGL,), cutlass.Int32)
+                                    part_l = cutlass.Int32(0)
+                                    for _js in cutlass.range_constexpr(SEGL):
+                                        v8_l = smem_hist[top_l - cutlass.Int32(_js)]
+                                        seg_l[_js] = v8_l
+                                        part_l = part_l + v8_l
+                                    tp_l = part_l
+                                    for _os in cutlass.range_constexpr(5):
+                                        ov_l = cutlass.const_expr(1 << _os)
+                                        oth_l = cute.arch.shuffle_sync_up(
+                                            tp_l, ov_l, mask_and_clamp=0
+                                        )
+                                        if lane >= cutlass.Int32(ov_l):
+                                            tp_l = tp_l + oth_l
+                                    excl_l = tp_l - part_l
+                                    kneed = kK_l - base_c
+                                    kfit = bs_l - base_c
+                                    run_l = cutlass.Int32(0)
+                                    for _js in cutlass.range_constexpr(SEGL):
+                                        run_l = run_l + seg_l[_js]
+                                        cum_at = excl_l + run_l
+                                        cum_bef = cum_at - seg_l[_js]
+                                        if cum_bef < kneed and cum_at >= kneed:
+                                            s_iscalars[3] = top_l - cutlass.Int32(_js)
+                                            smem_wcnt[0] = cum_bef
+                                            if cum_at <= kfit:
+                                                s_iscalars[2] = cutlass.Int32(1)
+                                            else:
+                                                s_iscalars[2] = cutlass.Int32(2)
+                                cute.arch.barrier()
+                                st_l = s_iscalars[2]
+                                if st_l == cutlass.Int32(1):
+                                    cut_t = r_lo + cutlass.Float32(s_iscalars[3]) * r_w
+                                    anch_t = cut_t
+                                    have = cutlass.Int32(1)
+                                    searching = cutlass.Int32(0)
+                                if st_l == cutlass.Int32(2):
+                                    base_c = base_c + smem_wcnt[0]
+                                    r_lo = r_lo + cutlass.Float32(s_iscalars[3]) * r_w
+                                    r_w = r_w / cutlass.Float32(NBL)
+                                    if r_w <= cutlass.Float32(0.0):
+                                        searching = cutlass.Int32(0)
+                                if st_l == cutlass.Int32(0):
+                                    searching = cutlass.Int32(0)
+                                cute.arch.barrier()
+                        if tidx == cutlass.Int32(0):
+                            # restore P2's cnt_lo/cnt_hi sentinels
+                            s_iscalars[2] = cutlass.Int32(-1)
+                            s_iscalars[3] = cutlass.Int32(-1)
+                        cute.arch.barrier()
+                if usable == cutlass.Int32(1) and have == cutlass.Int32(1):
+                    # ---- single filtered load at the cut ----
                     take_cand = cutlass.Int32(1)
-                if cutlass.const_expr(_P4_TAIL_DBG):
-                    ck0 = cute.arch.clock64()
-                    ck1 = ck0
-                if take_cand == cutlass.Int32(1):
-                    t_lo = seed_thr_row[0]
-                    lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1)
-                    vbase = cand_vals_row.iterator.toint()
-                    NBL = cutlass.const_expr(self.kNumBins)
-                    # ---- 1) stage: pure copy of the score column into
-                    # the dedicated smem region. No ballots, no atomics
-                    # (128 serialized smem atomics per trip measured
-                    # ~1.1us/1k entries), no warp-uniform constraint.
-                    # Sentinels sanitize to t_lo - 1 (rank below all live
-                    # scores; the K+64 admission bound covers K).
                     list_used = cutlass.Int32(1)
                     if tidx == cutlass.Int32(0):
-                        s_iscalars[2] = cutlass.Int32(0)
-                    snt = t_lo - cutlass.Float32(1.0)
-                    lmax = cutlass.Float32(self.NEG_FLT_MAX)
-                    i_s = tidx
-                    while i_s < claimed_c:
+                        s_iscalars[0] = cutlass.Int32(0)
+                        s_thr[0] = anch_t  # closed-loop anchor
+                        s_iscalars[1] = cutlass.Int32(1)  # done
+                    cute.arch.barrier()
+                    vbase = cand_vals_row.iterator.toint()
+                    lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1)
+                    i_c = tidx
+                    while (i_c - lane_c) < claimed_c:
+                        pvals = []
+                        pidxs = []
+                        keeps = []
                         for _ju in cutlass.range_constexpr(4):
-                            i_sj = i_s + cutlass.Int32(_ju * num_threads)
-                            if i_sj < claimed_c:
+                            i_cj = i_c + cutlass.Int32(_ju * num_threads)
+                            pval = cutlass.Float32(self.NEG_FLT_MAX)
+                            keep = cutlass.Int32(0)
+                            if i_cj < claimed_c:
                                 vp_l = cute.make_ptr(
                                     cutlass.Float32,
-                                    vbase + cutlass.Int64(i_sj) * cutlass.Int64(4),
+                                    vbase + cutlass.Int64(i_cj) * cutlass.Int64(4),
                                     cute.AddressSpace.gmem,
                                     assumed_align=4,
                                 )
-                                v_l = cute.make_tensor(vp_l, cute.make_layout((1,)))[0]
-                                if v_l < t_lo:
-                                    v_l = snt
-                                smem_list[i_sj] = v_l
-                                lmax = cute.arch.fmax(lmax, v_l)
-                        i_s = i_s + cutlass.Int32(4 * num_threads)
-                    wmax_l = self.warp_reduce_max_f32(lmax)
-                    if lane == cutlass.Int32(0):
-                        smem_wmax[warp_id] = wmax_l
-                    cute.arch.barrier()
-                    # ---- 2) zooming histogram on the smem scores: find
-                    # the edge whose exact count lands in [K, kC], as
-                    # close to K as the bins allow. Value-linear bins
-                    # collapse on long-tailed logits (vmax is an extreme
-                    # outlier), so each round re-bins INSIDE the crossing
-                    # bin - one cheap smem pass per round, NBL^3 total
-                    # resolution; only genuine ties fall through.
-                    vmax = cutlass.Float32(self.NEG_FLT_MAX)
-                    for _wr in cutlass.range_constexpr(self.num_warps):
-                        vmax = cute.arch.fmax(vmax, smem_wmax[_wr])
-                    rng_l = vmax - t_lo
-                    if rng_l <= cutlass.Float32(0.0):
-                        rng_l = cutlass.Float32(1.0)
-                    r_lo = t_lo
-                    r_w = rng_l / cutlass.Float32(NBL)
-                    w0 = r_w  # round-0 bin width (anchor edge recompute)
-                    base_c = cutlass.Int32(0)
-                    t_star = t_lo
-                    searching = cutlass.Int32(1)
-                    if tidx == cutlass.Int32(0):
-                        # smem_wcnt[1]: round-0 anchor bin (the ~3K
-                        # crossing; [0] is the descend base). -1 = not
-                        # found -> anchor falls back to t_star. P4
-                        # clobbers smem_wcnt only AFTER we read this.
-                        smem_wcnt[1] = cutlass.Int32(-1)
-                    for _rd in cutlass.range_constexpr(3):
-                        if searching == cutlass.Int32(1):
-                            jz_l = tidx
-                            while jz_l < cutlass.Int32(NBL):
-                                smem_hist[jz_l] = cutlass.Int32(0)
-                                jz_l = jz_l + cutlass.Int32(num_threads)
-                            if tidx == cutlass.Int32(0):
-                                s_iscalars[2] = cutlass.Int32(0)
-                            cute.arch.barrier()
-                            inv_wr = cutlass.Float32(1.0) / r_w
-                            r_hi = r_lo + r_w * cutlass.Float32(NBL)
-                            i_h = tidx
-                            while i_h < claimed_c:
-                                v_h = smem_list[i_h]
-                                if v_h >= r_lo and v_h < r_hi:
-                                    b_h = cutlass.Int32((v_h - r_lo) * inv_wr)
-                                    if b_h < cutlass.Int32(0):
-                                        b_h = cutlass.Int32(0)
-                                    if b_h > cutlass.Int32(NBL - 1):
-                                        b_h = cutlass.Int32(NBL - 1)
-                                    atomicAdd(smem_hist.iterator + b_h, cutlass.Int32(1))
-                                i_h = i_h + cutlass.Int32(num_threads)
-                            cute.arch.barrier()
-                            # warp-0 top-down cumulative scan (P1b idiom)
-                            if warp_id == cutlass.Int32(0):
-                                SEGL = cutlass.const_expr(NBL // self.WARP_SIZE)
-                                top_l = cutlass.Int32(NBL - 1) - lane * cutlass.Int32(SEGL)
-                                seg_l = cute.make_fragment((SEGL,), cutlass.Int32)
-                                part_l = cutlass.Int32(0)
-                                for _js in cutlass.range_constexpr(SEGL):
-                                    v8_l = smem_hist[top_l - cutlass.Int32(_js)]
-                                    seg_l[_js] = v8_l
-                                    part_l = part_l + v8_l
-                                tp_l = part_l
-                                for _os in cutlass.range_constexpr(5):
-                                    ov_l = cutlass.const_expr(1 << _os)
-                                    oth_l = cute.arch.shuffle_sync_up(tp_l, ov_l, mask_and_clamp=0)
-                                    if lane >= cutlass.Int32(ov_l):
-                                        tp_l = tp_l + oth_l
-                                excl_l = tp_l - part_l
-                                kneed = cutlass.Int32(self.top_k) - base_c
-                                kfit = cutlass.Int32(self.kC) - base_c
-                                run_l = cutlass.Int32(0)
-                                for _js in cutlass.range_constexpr(SEGL):
-                                    run_l = run_l + seg_l[_js]
-                                    cum_at = excl_l + run_l
-                                    cum_bef = cum_at - seg_l[_js]
-                                    if cum_bef < kneed and cum_at >= kneed:
-                                        s_iscalars[3] = top_l - cutlass.Int32(_js)
-                                        smem_wcnt[0] = cum_bef
-                                        if cum_at <= kfit:
-                                            s_iscalars[2] = cutlass.Int32(1)
-                                        else:
-                                            s_iscalars[2] = cutlass.Int32(2)
-                                    if cutlass.const_expr(_rd == 0):
-                                        # closed-loop ANCHOR: the ~3K
-                                        # crossing edge. Publishing the
-                                        # exact k-th made the next-step
-                                        # down-guard target only 4K -
-                                        # slope noise then undershoots K
-                                        # and forces ~26us fallbacks; a
-                                        # 3K-count anchor restores the
-                                        # old edge semantics (~12K guard
-                                        # target) at zero extra passes.
-                                        anch_n = cutlass.const_expr(
-                                            min(3 * self.top_k, (self.top_k + self.kC) // 2)
-                                        )
-                                        if cum_bef < cutlass.Int32(
-                                            anch_n
-                                        ) and cum_at >= cutlass.Int32(anch_n):
-                                            smem_wcnt[1] = top_l - cutlass.Int32(_js)
-                            cute.arch.barrier()
-                            st_l = s_iscalars[2]
-                            if st_l == cutlass.Int32(1):
-                                t_star = r_lo + cutlass.Float32(s_iscalars[3]) * r_w
-                                searching = cutlass.Int32(0)
-                            if st_l == cutlass.Int32(2):
-                                base_c = base_c + smem_wcnt[0]
-                                r_lo = r_lo + cutlass.Float32(s_iscalars[3]) * r_w
-                                r_w = r_w / cutlass.Float32(NBL)
-                                if r_w <= cutlass.Float32(0.0):
-                                    searching = cutlass.Int32(0)
-                            if st_l == cutlass.Int32(0):
-                                searching = cutlass.Int32(0)
-                            cute.arch.barrier()
-                    if tidx == cutlass.Int32(0):
-                        if s_iscalars[2] != cutlass.Int32(1):
-                            s_iscalars[2] = cutlass.Int32(0)
-                    cute.arch.barrier()
-                    fired_l = s_iscalars[2]
-                    t_anch = t_star
-                    anch_b = smem_wcnt[1]
-                    if anch_b >= cutlass.Int32(0):
-                        t_anch = t_lo + cutlass.Float32(anch_b) * w0
-                    if fired_l == cutlass.Int32(1):
-                        # ---- 3) compact survivors into the standard kC
-                        # buffers: score from smem, position streamed
-                        # from the gmem index column. The four unrolled
-                        # sub-ballots merge into ONE atomic per warp per
-                        # trip (the per-sub-ballot atomic serialized 128
-                        # adds per trip across 32 warps).
-                        if tidx == cutlass.Int32(0):
-                            s_iscalars[0] = cutlass.Int32(0)
-                            s_thr[0] = t_anch  # closed-loop anchor
-                            s_iscalars[1] = cutlass.Int32(1)  # done
-                        cute.arch.barrier()
-                        i_c = tidx
-                        while (i_c - lane_c) < claimed_c:
-                            pvals = []
-                            pidxs = []
-                            keeps = []
+                                pval = cute.make_tensor(vp_l, cute.make_layout((1,)))[0]
+                                if pval >= cut_t:
+                                    keep = cutlass.Int32(1)
+                            pvals.append(pval)
+                            # vals slot carries the LIST INDEX (register;
+                            # the position column is only gathered for
+                            # the K winners after Phase 4)
+                            pidxs.append(i_cj)
+                            keeps.append(keep)
+                        m0 = cute.arch.vote_ballot_sync(keeps[0] != cutlass.Int32(0))
+                        m1 = cute.arch.vote_ballot_sync(keeps[1] != cutlass.Int32(0))
+                        m2 = cute.arch.vote_ballot_sync(keeps[2] != cutlass.Int32(0))
+                        m3 = cute.arch.vote_ballot_sync(keeps[3] != cutlass.Int32(0))
+                        nk = cutlass.Int32(
+                            cute.arch.popc(m0)
+                            + cute.arch.popc(m1)
+                            + cute.arch.popc(m2)
+                            + cute.arch.popc(m3)
+                        )
+                        bk = cutlass.Int32(0)
+                        if nk > cutlass.Int32(0):
+                            if lane_c == cutlass.Int32(0):
+                                bk = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), nk)
+                            bk = cute.arch.shuffle_sync(bk, cutlass.Int32(0))
+                            lmk = (cutlass.Uint32(1) << cutlass.Uint32(lane_c)) - cutlass.Uint32(1)
+                            off = bk
                             for _ju in cutlass.range_constexpr(4):
-                                i_cj = i_c + cutlass.Int32(_ju * num_threads)
-                                pval = cutlass.Float32(self.NEG_FLT_MAX)
-                                # vals slot carries the LIST INDEX - a
-                                # register, not a second (cold) gmem pass
-                                # over the position column; the post-P4
-                                # repair gathers true positions for the K
-                                # winners only.
-                                pidx = i_cj
-                                keep = cutlass.Int32(0)
-                                if i_cj < claimed_c:
-                                    pval = smem_list[i_cj]
-                                    if pval >= t_star:
-                                        keep = cutlass.Int32(1)
-                                pvals.append(pval)
-                                pidxs.append(pidx)
-                                keeps.append(keep)
-                            m0 = cute.arch.vote_ballot_sync(keeps[0] != cutlass.Int32(0))
-                            m1 = cute.arch.vote_ballot_sync(keeps[1] != cutlass.Int32(0))
-                            m2 = cute.arch.vote_ballot_sync(keeps[2] != cutlass.Int32(0))
-                            m3 = cute.arch.vote_ballot_sync(keeps[3] != cutlass.Int32(0))
-                            nk = cutlass.Int32(
-                                cute.arch.popc(m0)
-                                + cute.arch.popc(m1)
-                                + cute.arch.popc(m2)
-                                + cute.arch.popc(m3)
-                            )
-                            bk = cutlass.Int32(0)
-                            if nk > cutlass.Int32(0):
-                                if lane_c == cutlass.Int32(0):
-                                    bk = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), nk)
-                                bk = cute.arch.shuffle_sync(bk, cutlass.Int32(0))
-                                lmk = (
-                                    cutlass.Uint32(1) << cutlass.Uint32(lane_c)
-                                ) - cutlass.Uint32(1)
-                                off = bk
-                                for _ju in cutlass.range_constexpr(4):
-                                    mj = (
-                                        m0
-                                        if _ju == 0
-                                        else m1
-                                        if _ju == 1
-                                        else m2
-                                        if _ju == 2
-                                        else m3
-                                    )
-                                    if keeps[_ju] != cutlass.Int32(0):
-                                        wpos = off + cutlass.Int32(cute.arch.popc(mj & lmk))
-                                        if wpos < cutlass.Int32(self.kC):
-                                            smem_keys[wpos] = pvals[_ju]
-                                            smem_vals[wpos] = pidxs[_ju]
-                                    off = off + cutlass.Int32(cute.arch.popc(mj))
+                                mj = m0 if _ju == 0 else m1 if _ju == 1 else m2 if _ju == 2 else m3
+                                if keeps[_ju] != cutlass.Int32(0):
+                                    wpos = off + cutlass.Int32(cute.arch.popc(mj & lmk))
+                                    if wpos < cutlass.Int32(self.kC):
+                                        smem_keys[wpos] = pvals[_ju]
+                                        smem_vals[wpos] = pidxs[_ju]
+                                off = off + cutlass.Int32(cute.arch.popc(mj))
                             i_c = i_c + cutlass.Int32(4 * num_threads)
-                        cute.arch.barrier()
-                        # hard net: histogram binning and the compact
-                        # test round float edges independently
-                        cnt_l = s_iscalars[0]
-                        if cnt_l < cutlass.Int32(self.top_k) or cnt_l > cutlass.Int32(self.kC):
-                            take_cand = cutlass.Int32(0)
-                            list_used = cutlass.Int32(0)
-                            if tidx == cutlass.Int32(0):
-                                s_iscalars[1] = cutlass.Int32(0)
-                                s_iscalars[2] = cutlass.Int32(-1)
-                                s_iscalars[3] = cutlass.Int32(-1)
-                            cute.arch.barrier()
-                    if fired_l == cutlass.Int32(0):
+                        if nk == cutlass.Int32(0):
+                            i_c = i_c + cutlass.Int32(4 * num_threads)
+                    cute.arch.barrier()
+                    # histogram-edge cuts round float edges independently
+                    # of the load predicate: keep the demote net for them
+                    # (line cuts are exact by construction and never trip)
+                    cnt_l = s_iscalars[0]
+                    if cnt_l < cutlass.Int32(self.top_k) or cnt_l > cutlass.Int32(self.kC):
                         take_cand = cutlass.Int32(0)
                         list_used = cutlass.Int32(0)
                         if tidx == cutlass.Int32(0):
-                            s_iscalars[2] = cutlass.Int32(-1)
-                            s_iscalars[3] = cutlass.Int32(-1)
+                            s_iscalars[1] = cutlass.Int32(0)
                         cute.arch.barrier()
                 if cutlass.const_expr(_P4_TAIL_DBG):
                     ck1 = cute.arch.clock64()
+
             if take_cand == cutlass.Int32(0):
                 # Stage this CTA's slice into SMEM once before Phase 2's
                 # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index f464f01a2031..1e03cc03ad2f 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -23,6 +23,7 @@
 
 import argparse
 import functools
+import os
 import sys
 from pathlib import Path
 from typing import Optional
@@ -69,6 +70,8 @@ def _compile(
     emit_xstate: bool = False,
     use_ext_cand: bool = False,
     cand_cap: int = 5120,
+    accept_cap: "int | None" = None,
+    kc_override: "int | None" = None,
 ):
     """JIT-compile the GVR kernel for a specific knob combination.
 
@@ -167,7 +170,7 @@ def _compile(
     )
     cand_ctl_fake = (
         cute.runtime.make_fake_compact_tensor(
-            cutlass.Int32, (n_rows, 2), stride_order=(1, 0), assumed_align=8
+            cutlass.Int32, (n_rows, 4), stride_order=(1, 0), assumed_align=8
         )
         if use_ext_cand
         else None
@@ -202,6 +205,8 @@ def _compile(
         emit_xstate=emit_xstate,
         use_ext_cand=use_ext_cand,
         cand_cap=cand_cap,
+        accept_cap=accept_cap,
+        kc_override=kc_override,
         # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed. The
         # qfrac VALUES are irrelevant on this path (P1b is skipped) — only
         # the slot count matters.
@@ -405,6 +410,47 @@ def derive_seed_rungs(
     return torch.stack([prev_thr - d_lo, prev_thr, prev_thr + d], dim=1).contiguous()
 
 
+def derive_seed_lines_v4(
+    prev_anchor: torch.Tensor,
+    prev_sthr: "torch.Tensor | None" = None,
+    prev_ctl: "torch.Tensor | None" = None,
+    targets: "tuple[float, float, float]" = (8192.0, 5120.0, 2048.0),
+    fallback_spread: float = 0.5,
+) -> torch.Tensor:
+    """v4 host-side line placement: put [t0, t1, t2] at target COUNTS.
+
+    Slope of log2(count) vs threshold is fit from the previous step's
+    (t0, n0) / (t2, n2) pairs (counts ride in the widened control words);
+    each new line lands where the fit predicts its target count. Targets
+    descend (t0 loosest / largest count, t2 tightest).
+
+    Args:
+        prev_anchor: [rows] previous accepted cut value (xstate[:, 2]).
+        prev_sthr: [rows, 3] previous lines (or None -> fixed spread).
+        prev_ctl: [rows, 4] previous control words {n0, void, n1, n2}.
+        targets: (T0, T1, T2) target counts, T0 > T1 > T2.
+
+    Returns:
+        [rows, 3] fp32 lines ascending [t0, t1, t2].
+    """
+    t0_t, t1_t, t2_t = targets
+    if prev_sthr is None or prev_ctl is None:
+        d = torch.full_like(prev_anchor, fallback_spread)
+        return torch.stack([prev_anchor - d, prev_anchor, prev_anchor + d], dim=1).contiguous()
+    c0 = prev_ctl[:, 0].float().clamp(min=1.0)
+    c2 = prev_ctl[:, 3].float().clamp(min=1.0)
+    dthr = (prev_sthr[:, 2] - prev_sthr[:, 0]).clamp(min=1e-3)
+    slope = ((torch.log2(c0) - torch.log2(c2)) / dthr).clamp(min=0.05, max=64.0)
+    # anchor count estimate: slide the anchor onto the prev line fit
+    anch_c = (c2 * torch.exp2(-(prev_anchor - prev_sthr[:, 2]) * slope)).clamp(min=1.0, max=1e6)
+    lines = [prev_anchor + torch.log2(anch_c / tgt) / slope for tgt in (t0_t, t1_t, t2_t)]
+    out = torch.stack(lines, dim=1)
+    # enforce strictly ascending (degenerate slope guards)
+    out[:, 1] = torch.maximum(out[:, 1], out[:, 0] + 1e-4)
+    out[:, 2] = torch.maximum(out[:, 2], out[:, 1] + 1e-4)
+    return out.contiguous()
+
+
 def emu_seed_counts(
     logits: torch.Tensor,
     seq_lens: torch.Tensor,
@@ -452,11 +498,15 @@ def emu_cand(
     lf = logits.to(torch.float32)
     cand_vals = torch.full((R, cap), float("-inf"), dtype=torch.float32, device=dev)
     cand_idx = torch.full((R, cap), -1, dtype=torch.int32, device=dev)
-    ctl = torch.zeros((R, 2), dtype=torch.int32, device=dev)
+    ctl = torch.zeros((R, 4), dtype=torch.int32, device=dev)
     for r in range(R):
         ne = int(n_eff[r])
         hits = torch.nonzero(lf[r, :ne] >= seed_thr[r, 0], as_tuple=False).flatten()
         cnt = hits.numel()
+        # emitter-side counts: two extra compares per EMITTED element
+        # (t1, t2 > t0 so counting over the list == counting over the row)
+        ctl[r, 2] = int((lf[r, :ne] >= seed_thr[r, 1]).sum())
+        ctl[r, 3] = int((lf[r, :ne] >= seed_thr[r, 2]).sum())
         # unordered contract: shuffle, then interleave sentinels
         perm = hits[torch.randperm(cnt, device=dev)]
         ent = torch.full((cnt + sentinel_pad,), -1, dtype=torch.int64, device=dev)
@@ -649,8 +699,8 @@ def gvr_topk_decode(
             cand_ctl.dtype == torch.int32
             and cand_ctl.is_cuda
             and cand_ctl.is_contiguous()
-            and cand_ctl.shape == (num_rows, 2)
-        ), "cand_ctl must be contiguous CUDA int32 [num_rows, 2]"
+            and cand_ctl.shape == (num_rows, 4)
+        ), "cand_ctl must be contiguous CUDA int32 [num_rows, 4]"
         cand_cap = cand_vals.shape[1]
     emit_xstate = xstate is not None
     if emit_xstate:
@@ -742,6 +792,8 @@ def gvr_topk_decode(
         emit_xstate,
         use_ext_cand,
         cand_cap,
+        int(os.environ["GVR_BSTAR"]) if "GVR_BSTAR" in os.environ else None,
+        int(os.environ["GVR_KC"]) if "GVR_KC" in os.environ else None,
     )
     # When return_output_values=False the kernel was compiled to skip
     # STG.value and accepts None for the value-output slot.

From cac0898d0b938986e07ba7a0da75164ef3e35d4f Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 24 Jul 2026 09:46:21 -0700
Subject: [PATCH 020/117] [None][perf] GVR L2 v5: bucketed segments,
 mapped-prefix copy, sampled histogram

Emitter writes the candidate list into three fixed segments (>=t2 /
[t1,t2) / [t0,t1), caps B*/B*/rest, spill to the looser segment on
overflow), so a line cut only ever reads the dense mapped prefix of
the segments above it: the hit path becomes a pure copy (no value
filter, no ballots, no atomics) and the histogram path walks mapped
indices. When all three lines overshoot B*, the bracket segment's own
prefix doubles as an unbiased sample: the histogram runs on it at the
sample rate with 1.25x-scaled fire targets, and the exact post-load
count net absorbs the sampling noise.

Device-level cold-chain results vs the block-skip arm (B=1):
flash 132k 1.58x, pro 160k 1.57x, pro 320k 1.54x, pro 640k 1.98x
(fastest steps 8-14us). Exactness smokes pass for cs=1/4/8 including
forced straddle/void routings.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 290 ++++++++++++------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  64 ++++
 2 files changed, 260 insertions(+), 94 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index a2e64640704c..9d415717dd4c 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -581,7 +581,11 @@ def __init__(
         # slice cache (128KB) - that combination exceeds the 227KB CTA
         # limit and fails loudly at compile time. Long rows (the list
         # path's target) cannot enable the slice cache anyway.
-        self.list_cap = int(cand_cap)
+        # v5 bucketed layout: tensor width = 2 * accept_cap
+        # (segments A, B) + segment-C capacity; the admission
+        # bounds the ENTRY COUNT by C's capacity (the only
+        # segment that can void).
+        self.list_cap = max(0, int(cand_cap) - 2 * self.accept_cap)
         self.cand_rung = int(cand_rung)
         if use_ext_cand and not use_ext_counts:
             raise ValueError("use_ext_cand requires use_ext_counts")
@@ -4960,14 +4964,38 @@ def _run_phases(
                 and cluster_size == 1
                 and self.dtype == cutlass.Float32
             ):
+                # ---- List path v5: BUCKETED segments ----
+                # The emitter classifies each entry by the tightest line
+                # it passes and appends into one of three fixed segments
+                # (A = [0, segA) holds >= t2, B = [segA, 2*segA) holds
+                # [t1, t2), C = [2*segA, ...) holds [t0, t1)); a full
+                # segment spills to the next looser one. Segment caps =
+                # B*, so "line in the acceptance band" <=> "its segment
+                # prefix group is complete" - the same condition the cut
+                # selection already checks. A LINE cut therefore loads a
+                # dense prefix of known length: a pure mapped copy, no
+                # filtering, no ballots, no atomics, zero wasted reads.
+                # Histogram fallbacks value-scan the mapped extents.
                 claimed_c = cutlass.Int32(cand_ctl_row[0])
                 void_c = cutlass.Int32(cand_ctl_row[1])
                 n1_c = cutlass.Int32(cand_ctl_row[2])
                 n2_c = cutlass.Int32(cand_ctl_row[3])
-                bstar = cutlass.const_expr(min(self.accept_cap, self.kC))
+                segA = cutlass.const_expr(min(self.accept_cap, self.kC))
+                bstar = segA
                 if cutlass.const_expr(_P4_TAIL_DBG):
                     ck0 = cute.arch.clock64()
                     ck1 = ck0
+                # segment extents (pads live at C's tail: sentinel score
+                # -inf slots, harmless to copy, never rank)
+                lenA = n2_c
+                if lenA > cutlass.Int32(segA):
+                    lenA = cutlass.Int32(segA)
+                spillA = n2_c - lenA
+                lenB = n1_c - n2_c + spillA
+                if lenB > cutlass.Int32(segA):
+                    lenB = cutlass.Int32(segA)
+                lenC = claimed_c - lenA - lenB
+                total_l = claimed_c
                 usable = cutlass.Int32(0)
                 if (
                     void_c == cutlass.Int32(0)
@@ -4975,61 +5003,84 @@ def _run_phases(
                     and claimed_c <= cutlass.Int32(self.list_cap)
                 ):
                     usable = cutlass.Int32(1)
-                # pre-declared: the DSL forbids first-assigning inside a
-                # dynamic if when read outside it
                 kK_l = cutlass.Int32(self.top_k)
                 bs_l = cutlass.Int32(bstar)
                 cut_t = cutlass.Float32(0.0)
+                cut_n = cutlass.Int32(0)
                 have = cutlass.Int32(0)
+                line_cut = cutlass.Int32(0)
                 anch_t = cutlass.Float32(0.0)
+                vbase = cutlass.Int64(0)
+                if cutlass.const_expr(True):
+                    vbase = cand_vals_row.iterator.toint()
                 if usable == cutlass.Int32(1):
-                    # cut = tightest line whose count is in [K, B*];
-                    # anchor (closed-loop publish) = loosest such line.
+                    # cut = tightest line in [K, B*]; anchor = loosest.
                     if n2_c >= kK_l and n2_c <= bs_l:
                         cut_t = seed_thr_row[2]
+                        cut_n = n2_c
                         anch_t = seed_thr_row[2]
                         have = cutlass.Int32(1)
+                        line_cut = cutlass.Int32(1)
                     if n1_c >= kK_l and n1_c <= bs_l:
                         if have == cutlass.Int32(0):
                             cut_t = seed_thr_row[1]
+                            cut_n = n1_c
                         anch_t = seed_thr_row[1]
                         have = cutlass.Int32(1)
+                        line_cut = cutlass.Int32(1)
                     if claimed_c <= bs_l:
                         if have == cutlass.Int32(0):
                             cut_t = seed_thr_row[0]
+                            cut_n = claimed_c
                         anch_t = seed_thr_row[0]
                         have = cutlass.Int32(1)
+                        line_cut = cutlass.Int32(1)
                     if have == cutlass.Int32(0):
-                        # No line in band: bracket the cut between the
-                        # two known lines that straddle it and find an
-                        # in-band edge with a clamped gmem histogram.
-                        #   all counts > B*      -> (t2, +inf): one max
-                        #                           pass supplies the top;
-                        #   n1 > B* and n2 < K   -> (t1, t2);
-                        #   n0 > B* and n1 < K   -> (t0, t1).
-                        vbase = cand_vals_row.iterator.toint()
+                        # ---- clamped-histogram fallback over the mapped
+                        # extents (bracket between two known lines; the
+                        # all-above case takes one max pass first) ----
+                        # histogram source = the bracket's own SEGMENT
+                        # prefix: if the segment is full it is a value-
+                        # blind (unbiased) SAMPLE of the band - scale the
+                        # targets by band/segment and let the post-load
+                        # count net verify; if not full it IS the exact
+                        # band. Either way the scan shrinks from the
+                        # whole list to <= one segment.
                         b_lo = seed_thr_row[0]
                         b_hi = seed_thr_row[1]
                         base_c = n1_c
+                        hs_base = cutlass.Int32(2 * segA)  # segment C
+                        hs_len = lenC
+                        hs_band = claimed_c - n1_c
                         if n1_c > bs_l:
                             b_lo = seed_thr_row[1]
                             b_hi = seed_thr_row[2]
                             base_c = n2_c
+                            hs_base = cutlass.Int32(segA)  # segment B
+                            hs_len = lenB
+                            hs_band = n1_c - n2_c
                         need_max = cutlass.Int32(0)
                         if n2_c > bs_l:
                             b_lo = seed_thr_row[2]
                             base_c = cutlass.Int32(0)
+                            hs_base = cutlass.Int32(0)  # segment A
+                            hs_len = lenA
+                            hs_band = n2_c
                             need_max = cutlass.Int32(1)
+                        if hs_band < cutlass.Int32(1):
+                            hs_band = cutlass.Int32(1)
+                        samp_f = (cutlass.Float32(1.0) * hs_len) / hs_band
                         if need_max == cutlass.Int32(1):
                             lmax = cutlass.Float32(self.NEG_FLT_MAX)
                             i_m = tidx
-                            while i_m < claimed_c:
+                            while i_m < hs_len:
                                 for _ju in cutlass.range_constexpr(4):
-                                    i_mj = i_m + cutlass.Int32(_ju * num_threads)
-                                    if i_mj < claimed_c:
+                                    j_m = i_m + cutlass.Int32(_ju * num_threads)
+                                    if j_m < hs_len:
+                                        src_m = hs_base + j_m
                                         vp_m = cute.make_ptr(
                                             cutlass.Float32,
-                                            vbase + cutlass.Int64(i_mj) * cutlass.Int64(4),
+                                            vbase + cutlass.Int64(src_m) * cutlass.Int64(4),
                                             cute.AddressSpace.gmem,
                                             assumed_align=4,
                                         )
@@ -5045,15 +5096,26 @@ def _run_phases(
                             for _wr in cutlass.range_constexpr(self.num_warps):
                                 vmax_l = cute.arch.fmax(vmax_l, smem_wmax[_wr])
                             b_hi = vmax_l + cutlass.Float32(1e-3)
-                        # zooming clamped histogram (up to 3 rounds; the
-                        # bracket is narrow so round 0 almost always
-                        # fires). State broadcasts ride s_iscalars[2]/[3]
-                        # (P2's slots, re-sentineled on every exit path).
                         NBL = cutlass.const_expr(self.kNumBins)
                         r_lo = b_lo
                         r_w = (b_hi - b_lo) / cutlass.Float32(NBL)
                         if r_w <= cutlass.Float32(0.0):
                             r_w = cutlass.Float32(1e-6)
+                        # sample-unit targets (population targets scaled
+                        # by segment/band); the descend base stays in
+                        # sample units too - the post-load exact-count
+                        # net absorbs the sampling error.
+                        # fire target = 1.25x the K-need: the sampled
+                        # estimate carries ~1-2% noise and firing at the
+                        # band's bottom edge would demote half the time
+                        kneedS = cutlass.Int32(
+                            (cutlass.Float32(1.25) * (kK_l - base_c)) * samp_f
+                            + cutlass.Float32(0.5)
+                        )
+                        if kneedS < cutlass.Int32(1):
+                            kneedS = cutlass.Int32(1)
+                        kfitS = cutlass.Int32((cutlass.Float32(1.0) * (bs_l - base_c)) * samp_f)
+                        sbase = cutlass.Int32(0)
                         searching = cutlass.Int32(1)
                         for _rd in cutlass.range_constexpr(3):
                             if searching == cutlass.Int32(1):
@@ -5067,13 +5129,14 @@ def _run_phases(
                                 inv_wr = cutlass.Float32(1.0) / r_w
                                 r_hi = r_lo + r_w * cutlass.Float32(NBL)
                                 i_h = tidx
-                                while i_h < claimed_c:
+                                while i_h < hs_len:
                                     for _ju in cutlass.range_constexpr(4):
-                                        i_hj = i_h + cutlass.Int32(_ju * num_threads)
-                                        if i_hj < claimed_c:
+                                        j_h = i_h + cutlass.Int32(_ju * num_threads)
+                                        if j_h < hs_len:
+                                            src_h = hs_base + j_h
                                             vp_h = cute.make_ptr(
                                                 cutlass.Float32,
-                                                vbase + cutlass.Int64(i_hj) * cutlass.Int64(4),
+                                                vbase + cutlass.Int64(src_h) * cutlass.Int64(4),
                                                 cute.AddressSpace.gmem,
                                                 assumed_align=4,
                                             )
@@ -5107,8 +5170,8 @@ def _run_phases(
                                         if lane >= cutlass.Int32(ov_l):
                                             tp_l = tp_l + oth_l
                                     excl_l = tp_l - part_l
-                                    kneed = kK_l - base_c
-                                    kfit = bs_l - base_c
+                                    kneed = kneedS - sbase
+                                    kfit = kfitS - sbase
                                     run_l = cutlass.Int32(0)
                                     for _js in cutlass.range_constexpr(SEGL):
                                         run_l = run_l + seg_l[_js]
@@ -5129,7 +5192,7 @@ def _run_phases(
                                     have = cutlass.Int32(1)
                                     searching = cutlass.Int32(0)
                                 if st_l == cutlass.Int32(2):
-                                    base_c = base_c + smem_wcnt[0]
+                                    sbase = sbase + smem_wcnt[0]
                                     r_lo = r_lo + cutlass.Float32(s_iscalars[3]) * r_w
                                     r_w = r_w / cutlass.Float32(NBL)
                                     if r_w <= cutlass.Float32(0.0):
@@ -5138,85 +5201,122 @@ def _run_phases(
                                     searching = cutlass.Int32(0)
                                 cute.arch.barrier()
                         if tidx == cutlass.Int32(0):
-                            # restore P2's cnt_lo/cnt_hi sentinels
                             s_iscalars[2] = cutlass.Int32(-1)
                             s_iscalars[3] = cutlass.Int32(-1)
                         cute.arch.barrier()
                 if usable == cutlass.Int32(1) and have == cutlass.Int32(1):
-                    # ---- single filtered load at the cut ----
                     take_cand = cutlass.Int32(1)
                     list_used = cutlass.Int32(1)
                     if tidx == cutlass.Int32(0):
                         s_iscalars[0] = cutlass.Int32(0)
-                        s_thr[0] = anch_t  # closed-loop anchor
+                        s_thr[0] = anch_t
                         s_iscalars[1] = cutlass.Int32(1)  # done
                     cute.arch.barrier()
-                    vbase = cand_vals_row.iterator.toint()
                     lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1)
-                    i_c = tidx
-                    while (i_c - lane_c) < claimed_c:
-                        pvals = []
-                        pidxs = []
-                        keeps = []
-                        for _ju in cutlass.range_constexpr(4):
-                            i_cj = i_c + cutlass.Int32(_ju * num_threads)
-                            pval = cutlass.Float32(self.NEG_FLT_MAX)
-                            keep = cutlass.Int32(0)
-                            if i_cj < claimed_c:
-                                vp_l = cute.make_ptr(
-                                    cutlass.Float32,
-                                    vbase + cutlass.Int64(i_cj) * cutlass.Int64(4),
-                                    cute.AddressSpace.gmem,
-                                    assumed_align=4,
-                                )
-                                pval = cute.make_tensor(vp_l, cute.make_layout((1,)))[0]
-                                if pval >= cut_t:
-                                    keep = cutlass.Int32(1)
-                            pvals.append(pval)
-                            # vals slot carries the LIST INDEX (register;
-                            # the position column is only gathered for
-                            # the K winners after Phase 4)
-                            pidxs.append(i_cj)
-                            keeps.append(keep)
-                        m0 = cute.arch.vote_ballot_sync(keeps[0] != cutlass.Int32(0))
-                        m1 = cute.arch.vote_ballot_sync(keeps[1] != cutlass.Int32(0))
-                        m2 = cute.arch.vote_ballot_sync(keeps[2] != cutlass.Int32(0))
-                        m3 = cute.arch.vote_ballot_sync(keeps[3] != cutlass.Int32(0))
-                        nk = cutlass.Int32(
-                            cute.arch.popc(m0)
-                            + cute.arch.popc(m1)
-                            + cute.arch.popc(m2)
-                            + cute.arch.popc(m3)
-                        )
-                        bk = cutlass.Int32(0)
-                        if nk > cutlass.Int32(0):
-                            if lane_c == cutlass.Int32(0):
-                                bk = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), nk)
-                            bk = cute.arch.shuffle_sync(bk, cutlass.Int32(0))
-                            lmk = (cutlass.Uint32(1) << cutlass.Uint32(lane_c)) - cutlass.Uint32(1)
-                            off = bk
+                    if line_cut == cutlass.Int32(1):
+                        # ---- LINE cut: dense mapped-prefix COPY of
+                        # exactly cut_n entries. No filter, no ballots,
+                        # no atomics - every read is a winner candidate.
+                        if tidx == cutlass.Int32(0):
+                            s_iscalars[0] = cut_n
+                        i_c = tidx
+                        while i_c < cut_n:
                             for _ju in cutlass.range_constexpr(4):
-                                mj = m0 if _ju == 0 else m1 if _ju == 1 else m2 if _ju == 2 else m3
-                                if keeps[_ju] != cutlass.Int32(0):
-                                    wpos = off + cutlass.Int32(cute.arch.popc(mj & lmk))
-                                    if wpos < cutlass.Int32(self.kC):
-                                        smem_keys[wpos] = pvals[_ju]
-                                        smem_vals[wpos] = pidxs[_ju]
-                                off = off + cutlass.Int32(cute.arch.popc(mj))
+                                j_c = i_c + cutlass.Int32(_ju * num_threads)
+                                if j_c < cut_n:
+                                    src_c = j_c
+                                    if j_c >= lenA:
+                                        src_c = cutlass.Int32(segA) + j_c - lenA
+                                    if j_c >= lenA + lenB:
+                                        src_c = cutlass.Int32(2 * segA) + j_c - lenA - lenB
+                                    vp_c = cute.make_ptr(
+                                        cutlass.Float32,
+                                        vbase + cutlass.Int64(src_c) * cutlass.Int64(4),
+                                        cute.AddressSpace.gmem,
+                                        assumed_align=4,
+                                    )
+                                    smem_keys[j_c] = cute.make_tensor(vp_c, cute.make_layout((1,)))[
+                                        0
+                                    ]
+                                    smem_vals[j_c] = src_c
                             i_c = i_c + cutlass.Int32(4 * num_threads)
-                        if nk == cutlass.Int32(0):
+                        cute.arch.barrier()
+                    if line_cut == cutlass.Int32(0):
+                        # ---- histogram-edge cut: value-filtered mapped
+                        # walk with merged-ballot claims (float edges
+                        # round independently -> demote net below).
+                        i_c = tidx
+                        while (i_c - lane_c) < total_l:
+                            pvals = []
+                            pidxs = []
+                            keeps = []
+                            for _ju in cutlass.range_constexpr(4):
+                                j_c = i_c + cutlass.Int32(_ju * num_threads)
+                                pval = cutlass.Float32(self.NEG_FLT_MAX)
+                                src_c = cutlass.Int32(0)
+                                keep = cutlass.Int32(0)
+                                if j_c < total_l:
+                                    src_c = j_c
+                                    if j_c >= lenA:
+                                        src_c = cutlass.Int32(segA) + j_c - lenA
+                                    if j_c >= lenA + lenB:
+                                        src_c = cutlass.Int32(2 * segA) + j_c - lenA - lenB
+                                    vp_c = cute.make_ptr(
+                                        cutlass.Float32,
+                                        vbase + cutlass.Int64(src_c) * cutlass.Int64(4),
+                                        cute.AddressSpace.gmem,
+                                        assumed_align=4,
+                                    )
+                                    pval = cute.make_tensor(vp_c, cute.make_layout((1,)))[0]
+                                    if pval >= cut_t:
+                                        keep = cutlass.Int32(1)
+                                pvals.append(pval)
+                                pidxs.append(src_c)
+                                keeps.append(keep)
+                            m0 = cute.arch.vote_ballot_sync(keeps[0] != cutlass.Int32(0))
+                            m1 = cute.arch.vote_ballot_sync(keeps[1] != cutlass.Int32(0))
+                            m2 = cute.arch.vote_ballot_sync(keeps[2] != cutlass.Int32(0))
+                            m3 = cute.arch.vote_ballot_sync(keeps[3] != cutlass.Int32(0))
+                            nk = cutlass.Int32(
+                                cute.arch.popc(m0)
+                                + cute.arch.popc(m1)
+                                + cute.arch.popc(m2)
+                                + cute.arch.popc(m3)
+                            )
+                            bk = cutlass.Int32(0)
+                            if nk > cutlass.Int32(0):
+                                if lane_c == cutlass.Int32(0):
+                                    bk = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), nk)
+                                bk = cute.arch.shuffle_sync(bk, cutlass.Int32(0))
+                                lmk = (
+                                    cutlass.Uint32(1) << cutlass.Uint32(lane_c)
+                                ) - cutlass.Uint32(1)
+                                off = bk
+                                for _ju in cutlass.range_constexpr(4):
+                                    mj = (
+                                        m0
+                                        if _ju == 0
+                                        else m1
+                                        if _ju == 1
+                                        else m2
+                                        if _ju == 2
+                                        else m3
+                                    )
+                                    if keeps[_ju] != cutlass.Int32(0):
+                                        wpos = off + cutlass.Int32(cute.arch.popc(mj & lmk))
+                                        if wpos < cutlass.Int32(self.kC):
+                                            smem_keys[wpos] = pvals[_ju]
+                                            smem_vals[wpos] = pidxs[_ju]
+                                    off = off + cutlass.Int32(cute.arch.popc(mj))
                             i_c = i_c + cutlass.Int32(4 * num_threads)
-                    cute.arch.barrier()
-                    # histogram-edge cuts round float edges independently
-                    # of the load predicate: keep the demote net for them
-                    # (line cuts are exact by construction and never trip)
-                    cnt_l = s_iscalars[0]
-                    if cnt_l < cutlass.Int32(self.top_k) or cnt_l > cutlass.Int32(self.kC):
-                        take_cand = cutlass.Int32(0)
-                        list_used = cutlass.Int32(0)
-                        if tidx == cutlass.Int32(0):
-                            s_iscalars[1] = cutlass.Int32(0)
                         cute.arch.barrier()
+                        cnt_l = s_iscalars[0]
+                        if cnt_l < cutlass.Int32(self.top_k) or cnt_l > cutlass.Int32(self.kC):
+                            take_cand = cutlass.Int32(0)
+                            list_used = cutlass.Int32(0)
+                            if tidx == cutlass.Int32(0):
+                                s_iscalars[1] = cutlass.Int32(0)
+                            cute.arch.barrier()
                 if cutlass.const_expr(_P4_TAIL_DBG):
                     ck1 = cute.arch.clock64()
 
@@ -5691,7 +5791,9 @@ def _run_phases(
                         io_r = tidx
                         while io_r < cutlass.Int32(self.top_k):
                             li_r = output_indices_row[io_r]
-                            if li_r >= cutlass.Int32(0) and li_r < claimed_c:
+                            # slots are segmented offsets (may exceed the
+                            # entry count); only sentinel -1 is invalid
+                            if li_r >= cutlass.Int32(0):
                                 ip_r = cute.make_ptr(
                                     cutlass.Int32,
                                     cand_idx_row.iterator.toint()
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 1e03cc03ad2f..70a64d2e0e28 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -410,6 +410,70 @@ def derive_seed_rungs(
     return torch.stack([prev_thr - d_lo, prev_thr, prev_thr + d], dim=1).contiguous()
 
 
+def emu_cand_bucketed(
+    logits: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seed_thr: torch.Tensor,
+    cap: int,
+    seg_cap: int = 8192,
+    next_n: int = 1,
+    compress_ratio: int = 1,
+    sentinel_pad: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Bucketed SoA candidate emission (v5 contract).
+
+    Three fixed segments in one buffer: A = [0, seg_cap) holds >= t2,
+    B = [seg_cap, 2*seg_cap) holds [t1, t2), C = [2*seg_cap, 2*seg_cap
+    + cap) holds [t0, t1). A full segment spills to the next looser one
+    (never drops an entry), so the union always equals the full >= t0
+    set, and a segment group is complete exactly when its line's count
+    fits the acceptance band (seg_cap = B*). Sentinel pads land in C.
+    ctl = {n0 (claimed incl pads), void, n1, n2}.
+    """
+    R, _C = logits.shape
+    dev = logits.device
+    n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio)
+    lf = logits.to(torch.float32)
+    width = 2 * seg_cap + cap
+    cand_vals = torch.full((R, width), float("-inf"), dtype=torch.float32, device=dev)
+    cand_idx = torch.full((R, width), -1, dtype=torch.int32, device=dev)
+    ctl = torch.zeros((R, 4), dtype=torch.int32, device=dev)
+    for r in range(R):
+        ne = int(n_eff[r])
+        row = lf[r, :ne]
+        hits = torch.nonzero(row >= seed_thr[r, 0], as_tuple=False).flatten()
+        cnt = hits.numel()
+        ctl[r, 2] = int((row >= seed_thr[r, 1]).sum())
+        ctl[r, 3] = int((row >= seed_thr[r, 2]).sum())
+        # emission order is value-blind: shuffle, then classify
+        perm = hits[torch.randperm(cnt, device=dev)]
+        v = row[perm]
+        seg = torch.where(v >= seed_thr[r, 2], 0, torch.where(v >= seed_thr[r, 1], 1, 2))
+        # vectorized spill-to-looser: stream s = native entries + spill
+        # from s-1 (in emission order); ordinal beyond the cap spills on.
+        in_a = seg == 0
+        ord_a = torch.cumsum(in_a.int(), 0)
+        stay_a = in_a & (ord_a <= seg_cap)
+        in_b = (seg == 1) | (in_a & ~stay_a)
+        ord_b = torch.cumsum(in_b.int(), 0)
+        stay_b = in_b & (ord_b <= seg_cap)
+        in_c = (seg == 2) | (in_b & ~stay_b)
+        ord_c = torch.cumsum(in_c.int(), 0)
+        stay_c = in_c & (ord_c <= cap)
+        voided = int(in_c.sum()) > cap
+        slot = torch.full_like(seg, -1)
+        slot[stay_a] = ord_a[stay_a] - 1
+        slot[stay_b] = seg_cap + (ord_b[stay_b] - 1)
+        slot[stay_c] = 2 * seg_cap + (ord_c[stay_c] - 1)
+        live = slot >= 0
+        cand_vals[r, slot[live].long()] = v[live]
+        cand_idx[r, slot[live].long()] = perm[live].int()
+        claimed = cnt + sentinel_pad
+        ctl[r, 0] = claimed
+        ctl[r, 1] = 1 if (voided or claimed > cap) else 0
+    return cand_vals, cand_idx, ctl
+
+
 def derive_seed_lines_v4(
     prev_anchor: torch.Tensor,
     prev_sthr: "torch.Tensor | None" = None,

From 3ca97312463a1cfd225bd04240df08419206cf76 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 05:50:03 -0700
Subject: [PATCH 021/117] [None][perf] GVR P4: in-place class compaction +
 three-tier boundary repair

Sub-phase clock64 instrumentation (GVR_P4_SUB_DBG) showed the P4
rank-scatter core costs only ~0.6us/k candidates; the chain-observed
~1.9us/k came from the exact-tail boundary repair: the tiny-tie fast
path ran an O(need x class) serial select on thread0 (~10us on real
rows with need ~100 x class ~100), and bigger classes re-scanned every
candidate per radix level behind ~20 block barriers.

The repair is now: (1) a block-wide pure-tie check over the straddle
class (bit-equal class needs no repair at all - the scatter's arrival
fill is already value-set exact); (2) mixed classes are compacted IN
PLACE into smem_keys/vals[0..class) with a register-buffered two-phase
pass (warp-aggregated slot claims), so every later step scales with
the class, never the candidate count; (3) class <= 128 takes an exact
warp0 pairwise-rank rewrite, larger classes a block-parallel 4-level
MSB radix over the compacted pairs with a warp0 shuffle-scan digit
search (3 block barriers per level instead of 5). The full-candidate
radix fallback is gone from the fast-tail variant.

Device-level cold chains, 640k B=1: step mean 13.7 -> 12.6us
(1.97x -> 2.13x vs block-skip; the previously slowest window improves
17.35 -> 12.6us as five 19-21us serial-repair victims drop to
9.3-10.8us); 640k B=8 19.5us (1.60x). Exactness: 18/18 microbench
cells including forced tie/outlier stressors, smokes cs=1/4/8 plus
forced straddle/void routings all bit-exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 544 ++++++++++--------
 1 file changed, 310 insertions(+), 234 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 9d415717dd4c..c9b3a6800f40 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -46,6 +46,9 @@
 # into the spare xstate slots (harness-side analysis). Off by default;
 # NEVER set in production.
 _P4_TAIL_DBG = bool(int(os.environ.get("GVR_P4_TAIL_DBG", "0")))
+# P4 sub-phase clock64 breakdown -> xstate[1,2,4,5,6,7] (debug: clobbers
+# the closed-loop thr/anch publish; single-shot cells only, not chains)
+_P4_SUB_DBG = bool(int(os.environ.get("GVR_P4_SUB_DBG", "0")))
 
 
 # ---------------------------------------------------------------------------
@@ -2925,6 +2928,15 @@ def phase4_rank_scatter(
                 output_indices_row[i4] = smem_vals[i4]
                 i4 = i4 + cutlass.Int32(num_threads)
         elif cand_count > cutlass.Int32(kK):
+            sc0 = cutlass.Int64(0)
+            sc1 = cutlass.Int64(0)
+            sc2 = cutlass.Int64(0)
+            sc3 = cutlass.Int64(0)
+            sc4 = cutlass.Int64(0)
+            sc5 = cutlass.Int64(0)
+            sc6 = cutlass.Int64(0)
+            if cutlass.const_expr(_P4_SUB_DBG):
+                sc0 = cute.arch.clock64()
             # ---- block min/max over candidates ----
             local_cmin = cutlass.Float32(self.FLT_MAX)
             local_cmax = cutlass.Float32(self.NEG_FLT_MAX)
@@ -2954,6 +2966,8 @@ def phase4_rank_scatter(
             if bmax_r <= bmin_r:
                 bmax_r = bmin_r + cutlass.Float32(1e-6)
             cute.arch.barrier()
+            if cutlass.const_expr(_P4_SUB_DBG):
+                sc1 = cute.arch.clock64()
             # ---- zero + build histogram ----
             i6 = tidx
             while i6 < cutlass.Int32(kBins):
@@ -2973,6 +2987,8 @@ def phase4_rank_scatter(
                 atomicAdd(smem_hist.iterator + bin_i, cutlass.Int32(1))
                 i7 = i7 + cutlass.Int32(num_threads)
             cute.arch.barrier()
+            if cutlass.const_expr(_P4_SUB_DBG):
+                sc2 = cute.arch.clock64()
             # ---- 3-step high→low bin search → straddling bin b* + rank_above ----
             warp_bin_sum = cutlass.Int32(0)
             for jb in cutlass.range_constexpr(bins_per_warp):
@@ -3024,6 +3040,8 @@ def phase4_rank_scatter(
                 s_iscalars[4] = cutlass.Int32(0)  # cnt_above
                 s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
             cute.arch.barrier()
+            if cutlass.const_expr(_P4_SUB_DBG):
+                sc3 = cute.arch.clock64()
             b_star = s_iscalars[3]
             rank_above = s_iscalars[2]
 
@@ -3122,6 +3140,8 @@ def phase4_rank_scatter(
                     s_iscalars[0] = cutlass.Int32(0)  # cnt_mid (b*, sub>sb*)
                     s_iscalars[1] = cutlass.Int32(0)  # cnt_strad (b*, sub==sb*)
                 cute.arch.barrier()
+                if cutlass.const_expr(_P4_SUB_DBG):
+                    sc4 = cute.arch.clock64()
                 sb_star = smem_hist[2]
                 rank_above_fine = smem_hist[3]
                 isc = tidx
@@ -3160,6 +3180,8 @@ def phase4_rank_scatter(
                                 output_indices_row[pos] = smem_vals[isc]
                     isc = isc + cutlass.Int32(num_threads)
                 cute.arch.barrier()
+                if cutlass.const_expr(_P4_SUB_DBG):
+                    sc5 = cute.arch.clock64()
                 cnt_strad = s_iscalars[1]
                 filled = rank_above_fine + cnt_strad
                 if filled > cutlass.Int32(kK):
@@ -3183,257 +3205,281 @@ def phase4_rank_scatter(
                 # slot range [rank_above_fine, kK). Unambiguous rows (the
                 # overwhelming majority) pay two scalar compares; the counters
                 # and the fine histogram are reused, so SMEM does not grow.
-                # [p4tt] tiny-tie fast path: when the exact-tail gate fires
-                # with a small (b*, sb*) tie class (cnt_strad <= 128 — the
-                # real firing cells hold 2), ONE candidate pass collects the
-                # class and thread0 selects the top-need exactly, replacing
-                # the 4 unconditional radix passes. Larger classes take the
-                # UNMODIFIED radix select below (verbatim copy).
+                # [p4tt] boundary-class repair: collect the (b*, sb*) tie
+                # class compactly (ONE candidate pass), then select inside
+                # it. Tiny jobs (need x class <= 512) keep the thread0
+                # serial select (cheapest at that size). Bigger classes up
+                # to capc get a 4-level MSB radix over the COLLECTED class
+                # (each level scans <= capc pairs with all threads instead
+                # of re-scanning every candidate) — the old serial select
+                # was O(need x class) and measured 8.5us on real chain rows
+                # (need ~100 x class ~100). Classes beyond capc take the
+                # UNMODIFIED full-candidate radix below (verbatim copy).
                 if cutlass.const_expr(self.p4_exact_tail and self.p4_tail_fast):  # [p4tt]
                     need0 = cutlass.Int32(kK) - rank_above_fine
-                    if cutlass.const_expr(_P4_TAIL_DBG):
+                    # [p4tt-v3] per-thread compact buffers, bounded by the
+                    # strided trip count over the candidate array
+                    nbuf7 = cutlass.const_expr((self.kC + self.num_threads - 1) // self.num_threads)
+                    rv7 = cute.make_fragment((nbuf7,), cutlass.Float32)
+                    ri7 = cute.make_fragment((nbuf7,), cutlass.Int32)
+                    if cutlass.const_expr(_P4_TAIL_DBG or _P4_SUB_DBG):
                         if tidx == cutlass.Int32(0):
                             s_thr[1] = cutlass.Float32(cnt_strad)
                             s_thr[2] = cutlass.Float32(need0)
+                    fast_done = cutlass.Int32(1)
                     if cnt_strad > need0 and need0 > cutlass.Int32(0):
-                        if cnt_strad <= cutlass.Int32(128):
-                            # [p4tt] SMEM: (value_bits, cand_idx) pairs at
-                            # smem_hist[2*o]/[2*o+1], o < 128 (slots 0..255).
-                            # The 256 digit bins are dead here (the fast path
-                            # replaces the radix levels that used them); the
-                            # sb_star/ra staging in slots 2/3 was read by
-                            # every thread before the pre-scatter barrier.
-                            # Persistent radix scalars [256..258] untouched.
-                            # Collect counter = s_iscalars[0] (dead after the
-                            # scatter; same reuse as the radix rewrite pass).
+                        fast_done = cutlass.Int32(0)
+                        # [p4tt-v3] block-wide pure-tie check, ANY class
+                        # size: min/max order key over the (b*, sb*) class.
+                        # A pure-tie class needs NO repair — the scatter's
+                        # arrival fill of bit-equal values is already
+                        # value-set exact. Real fp8-lineage logits tie in
+                        # the thousands, which used to take the full radix.
+                        # Staging mirrors the head min/max (wcnt + hist
+                        # slots [0..31], both dead here; pairs live at
+                        # 260+).
+                        kmn6 = cutlass.Int32(2147483647)
+                        kmx6 = cutlass.Int32(-2147483648)
+                        it6 = tidx
+                        while it6 < cand_count:
+                            v6 = smem_keys[it6]
+                            b6 = cutlass.Int32((v6 - bmin_r) * inv1)
+                            if b6 < cutlass.Int32(0):
+                                b6 = cutlass.Int32(0)
+                            if b6 > cutlass.Int32(kBins - 1):
+                                b6 = cutlass.Int32(kBins - 1)
+                            if b6 == b_star:
+                                s6 = cutlass.Int32((v6 - f_lo) * finv)
+                                if s6 < cutlass.Int32(0):
+                                    s6 = cutlass.Int32(0)
+                                if s6 > cutlass.Int32(fbins - 1):
+                                    s6 = cutlass.Int32(fbins - 1)
+                                if s6 == sb_star:
+                                    k6 = f32_order_key(v6) ^ cutlass.Int32(-2147483648)
+                                    if k6 < kmn6:
+                                        kmn6 = k6
+                                    if k6 > kmx6:
+                                        kmx6 = k6
+                            it6 = it6 + cutlass.Int32(num_threads)
+                        kmn6 = cute.arch.warp_redux_sync(kmn6, "min")
+                        kmx6 = cute.arch.warp_redux_sync(kmx6, "max")
+                        if lane == cutlass.Int32(0):
+                            smem_wcnt[warp_id] = kmn6
+                            smem_hist[warp_id] = kmx6
+                        cute.arch.barrier()
+                        kmn7 = cutlass.Int32(2147483647)
+                        kmx7 = cutlass.Int32(-2147483648)
+                        for w8 in cutlass.range_constexpr(self.num_warps):
+                            pa8 = smem_wcnt[w8]
+                            pb8 = smem_hist[w8]
+                            if pa8 < kmn7:
+                                kmn7 = pa8
+                            if pb8 > kmx7:
+                                kmx7 = pb8
+                        if kmn7 == kmx7:
+                            fast_done = cutlass.Int32(1)
+                        if fast_done == cutlass.Int32(0):
+                            # [p4tt-v3] mixed class: compact it IN PLACE
+                            # into smem_keys/vals[0..cnt_strad) with a
+                            # register-buffered two-phase pass (every
+                            # thread reads its strided candidates first,
+                            # ONE barrier, then claimed compact writes —
+                            # no read/write overlap by construction). The
+                            # candidate array has no readers after the
+                            # tail, and compaction makes the repair cost a
+                            # function of the CLASS size only, for ANY
+                            # class size up to cand_count (the old full-
+                            # candidate radix fallback is gone).
                             if tidx == cutlass.Int32(0):
                                 s_iscalars[0] = cutlass.Int32(0)
+                            nh7 = cutlass.Int32(0)
+                            it7 = tidx
+                            while it7 < cand_count:
+                                v7 = smem_keys[it7]
+                                b7 = cutlass.Int32((v7 - bmin_r) * inv1)
+                                if b7 < cutlass.Int32(0):
+                                    b7 = cutlass.Int32(0)
+                                if b7 > cutlass.Int32(kBins - 1):
+                                    b7 = cutlass.Int32(kBins - 1)
+                                if b7 == b_star:
+                                    s7 = cutlass.Int32((v7 - f_lo) * finv)
+                                    if s7 < cutlass.Int32(0):
+                                        s7 = cutlass.Int32(0)
+                                    if s7 > cutlass.Int32(fbins - 1):
+                                        s7 = cutlass.Int32(fbins - 1)
+                                    if s7 == sb_star:
+                                        # static predicated fragment write
+                                        # (dodges dynamic register indexing)
+                                        for sl7 in cutlass.range_constexpr(nbuf7):
+                                            if cutlass.Int32(sl7) == nh7:
+                                                rv7[sl7] = v7
+                                                ri7[sl7] = smem_vals[it7]
+                                        nh7 = nh7 + cutlass.Int32(1)
+                                it7 = it7 + cutlass.Int32(num_threads)
                             cute.arch.barrier()
-                            itc = tidx
-                            while itc < cand_count:
-                                tv = smem_keys[itc]
-                                tb = cutlass.Int32((tv - bmin_r) * inv1)
-                                if tb < cutlass.Int32(0):
-                                    tb = cutlass.Int32(0)
-                                if tb > cutlass.Int32(kBins - 1):
-                                    tb = cutlass.Int32(kBins - 1)
-                                if tb == b_star:
-                                    ts = cutlass.Int32((tv - f_lo) * finv)
-                                    if ts < cutlass.Int32(0):
-                                        ts = cutlass.Int32(0)
-                                    if ts > cutlass.Int32(fbins - 1):
-                                        ts = cutlass.Int32(fbins - 1)
-                                    if ts == sb_star:
-                                        to = atomicAdd(
-                                            s_iscalars.iterator + cutlass.Int32(0), cutlass.Int32(1)
-                                        )
-                                        if to < cutlass.Int32(128):
-                                            smem_hist[to + to] = float_as_int32(tv)
-                                            smem_hist[to + to + cutlass.Int32(1)] = smem_vals[itc]
-                                itc = itc + cutlass.Int32(num_threads)
-                            cute.arch.barrier()
-                            # [p4tt] thread0 exact top-need0 select rewriting
-                            # positions [rank_above_fine, kK). Consumed flag =
-                            # the cand_idx slot set to -1 (indices are always
-                            # >= 0), so a genuine -FLT_MAX value in the class
-                            # remains selectable (no value sentinel). Ties
-                            # (bit-equal values) pick arbitrarily: value-set
-                            # exact.
-                            if tidx == cutlass.Int32(0):
-                                tj = cutlass.Int32(0)
-                                while tj < need0:
-                                    tbv = cutlass.Float32(self.NEG_FLT_MAX)
-                                    tbi = cutlass.Int32(-1)
-                                    ti = cutlass.Int32(0)
-                                    while ti < cnt_strad:
-                                        tvi = smem_hist[ti + ti + cutlass.Int32(1)]
-                                        if tvi >= cutlass.Int32(0):
-                                            tvb = smem_hist[ti + ti]
-                                            tvv = cutlass.Float32(
-                                                llvm.bitcast(
-                                                    cutlass.Float32.mlir_type,
-                                                    tvb.ir_value(),
-                                                )
-                                            )
-                                            take = cutlass.Int32(0)
-                                            if tbi < cutlass.Int32(0):
-                                                take = cutlass.Int32(1)
-                                            elif tvv > tbv:
-                                                take = cutlass.Int32(1)
-                                            if take == cutlass.Int32(1):
-                                                tbv = tvv
-                                                tbi = ti
-                                        ti = ti + cutlass.Int32(1)
-                                    pos = rank_above_fine + tj
-                                    if cutlass.const_expr(self.return_output_values):
-                                        output_values_row[pos] = self.dtype(tbv)
-                                    output_indices_row[pos] = smem_hist[
-                                        tbi + tbi + cutlass.Int32(1)
-                                    ]
-                                    smem_hist[tbi + tbi + cutlass.Int32(1)] = cutlass.Int32(-1)
-                                    tj = tj + cutlass.Int32(1)
-                            cute.arch.barrier()
-                        else:
-                            # Persistent scalars live above the 256 digit bins
-                            # (kNumBins >= 512 always): [256] key prefix (chosen
-                            # digits, remaining bits 0), [257] slots still to fill
-                            # inside the current equal-prefix set, [258] ties
-                            # strictly above the prefix (their slots precede it).
-                            if tidx == cutlass.Int32(0):
-                                smem_hist[256] = cutlass.Int32(0)
-                                smem_hist[257] = need0
-                                smem_hist[258] = cutlass.Int32(0)
+                            # warp-aggregated claim: intra-warp exclusive
+                            # prefix via shfl scan + ONE atomic per warp
+                            # (a thousand same-address claims serialize
+                            # and scale with the class size)
+                            pf7 = nh7
+                            for so3 in cutlass.range_constexpr(5):
+                                oth3 = cute.arch.shuffle_sync_up(
+                                    pf7, cutlass.Int32(1 << so3), mask_and_clamp=0
+                                )
+                                if lane >= cutlass.Int32(1 << so3):
+                                    pf7 = pf7 + oth3
+                            tot7 = cute.arch.shuffle_sync(pf7, cutlass.Int32(31))
+                            wb7 = cutlass.Int32(0)
+                            if lane == cutlass.Int32(31):
+                                if tot7 > cutlass.Int32(0):
+                                    wb7 = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), tot7)
+                            wb7 = cute.arch.shuffle_sync(wb7, cutlass.Int32(31))
+                            bs7 = wb7 + pf7 - nh7
+                            for sl8 in cutlass.range_constexpr(nbuf7):
+                                if cutlass.Int32(sl8) < nh7:
+                                    smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8]
+                                    smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8]
                             cute.arch.barrier()
-                            for lvl in cutlass.range_constexpr(4):
-                                shift = cutlass.const_expr(24 - 8 * lvl)
-                                iz2 = tidx
-                                while iz2 < cutlass.Int32(256):
-                                    smem_hist[iz2] = cutlass.Int32(0)
-                                    iz2 = iz2 + cutlass.Int32(num_threads)
-                                cute.arch.barrier()
-                                uthr_cur = smem_hist[256]
-                                it2 = tidx
-                                while it2 < cand_count:
-                                    vt = smem_keys[it2]
-                                    bt = cutlass.Int32((vt - bmin_r) * inv1)
-                                    if bt < cutlass.Int32(0):
-                                        bt = cutlass.Int32(0)
-                                    if bt > cutlass.Int32(kBins - 1):
-                                        bt = cutlass.Int32(kBins - 1)
-                                    if bt == b_star:
-                                        st2 = cutlass.Int32((vt - f_lo) * finv)
-                                        if st2 < cutlass.Int32(0):
-                                            st2 = cutlass.Int32(0)
-                                        if st2 > cutlass.Int32(fbins - 1):
-                                            st2 = cutlass.Int32(fbins - 1)
-                                        if st2 == sb_star:
-                                            uk = f32_order_key(vt)
-                                            pmatch = cutlass.Int32(1)
-                                            if cutlass.const_expr(lvl > 0):
-                                                if (uk >> cutlass.Int32(shift + 8)) != (
-                                                    uthr_cur >> cutlass.Int32(shift + 8)
-                                                ):
-                                                    pmatch = cutlass.Int32(0)
-                                            if pmatch == cutlass.Int32(1):
-                                                dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32(
-                                                    0xFF
-                                                )
-                                                atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1))
-                                    it2 = it2 + cutlass.Int32(num_threads)
+                            if cnt_strad <= cutlass.Int32(128):
+                                # warp0 exact pairwise rank (rank = #{key
+                                # greater} + #{key equal, earlier slot} is
+                                # unique in [0, class)) rewrites every
+                                # winner slot in [raf, raf + need0) once.
+                                if warp_id == cutlass.Int32(0):
+                                    ie5 = lane
+                                    while ie5 < cnt_strad:
+                                        vi5 = smem_keys[ie5]
+                                        ki5 = f32_order_key(vi5) ^ cutlass.Int32(-2147483648)
+                                        r5 = cutlass.Int32(0)
+                                        j5 = cutlass.Int32(0)
+                                        while j5 < cnt_strad:
+                                            vj5 = smem_keys[j5]
+                                            kj5 = f32_order_key(vj5) ^ cutlass.Int32(-2147483648)
+                                            if kj5 > ki5:
+                                                r5 = r5 + cutlass.Int32(1)
+                                            elif kj5 == ki5 and j5 < ie5:
+                                                r5 = r5 + cutlass.Int32(1)
+                                            j5 = j5 + cutlass.Int32(1)
+                                        if r5 < need0:
+                                            pos = rank_above_fine + r5
+                                            if pos < cutlass.Int32(kK):
+                                                if cutlass.const_expr(self.return_output_values):
+                                                    output_values_row[pos] = self.dtype(vi5)
+                                                output_indices_row[pos] = smem_vals[ie5]
+                                        ie5 = ie5 + cutlass.Int32(32)
                                 cute.arch.barrier()
-                                # Two-stage descending digit scan (mirrors the
-                                # fine 3-step search): per-warp partial sums,
-                                # thread0 picks the target warp, its lane0 walks
-                                # the warp's digit range — 2*num_warps serial
-                                # steps instead of 256.
-                                fdw = cutlass.const_expr(256 // self.num_warps)
-                                wsum2 = cutlass.Int32(0)
-                                for jd in cutlass.range_constexpr(fdw):
-                                    dix = (
-                                        cutlass.Int32(255)
-                                        - warp_id * cutlass.Int32(fdw)
-                                        - cutlass.Int32(jd)
-                                    )
-                                    wsum2 = wsum2 + smem_hist[dix]
-                                if lane == cutlass.Int32(0):
-                                    smem_wcnt[warp_id] = wsum2
+                            else:
+                                # block-parallel 4-level MSB radix over the
+                                # compacted class (scans touch class pairs
+                                # only; warp0 shuffle-scan digit search — 3
+                                # block barriers per level instead of 5).
+                                if tidx == cutlass.Int32(0):
+                                    smem_hist[256] = cutlass.Int32(0)
+                                    smem_hist[257] = need0
+                                    smem_hist[258] = cutlass.Int32(0)
                                 cute.arch.barrier()
+                                for lvl2 in cutlass.range_constexpr(4):
+                                    shift2 = cutlass.const_expr(24 - 8 * lvl2)
+                                    iz3 = tidx
+                                    while iz3 < cutlass.Int32(256):
+                                        smem_hist[iz3] = cutlass.Int32(0)
+                                        iz3 = iz3 + cutlass.Int32(num_threads)
+                                    cute.arch.barrier()
+                                    uthr_c2 = smem_hist[256]
+                                    ic2 = tidx
+                                    while ic2 < cnt_strad:
+                                        uk3 = f32_order_key(smem_keys[ic2])
+                                        pm2 = cutlass.Int32(1)
+                                        if cutlass.const_expr(lvl2 > 0):
+                                            if (uk3 >> cutlass.Int32(shift2 + 8)) != (
+                                                uthr_c2 >> cutlass.Int32(shift2 + 8)
+                                            ):
+                                                pm2 = cutlass.Int32(0)
+                                        if pm2 == cutlass.Int32(1):
+                                            dg2 = (uk3 >> cutlass.Int32(shift2)) & cutlass.Int32(
+                                                0xFF
+                                            )
+                                            atomicAdd(smem_hist.iterator + dg2, cutlass.Int32(1))
+                                        ic2 = ic2 + cutlass.Int32(num_threads)
+                                    cute.arch.barrier()
+                                    if warp_id == cutlass.Int32(0):
+                                        ws3 = cutlass.Int32(0)
+                                        for jd3 in cutlass.range_constexpr(8):
+                                            di3 = (
+                                                cutlass.Int32(255)
+                                                - lane * cutlass.Int32(8)
+                                                - cutlass.Int32(jd3)
+                                            )
+                                            ws3 = ws3 + smem_hist[di3]
+                                        pre6 = ws3
+                                        for so2 in cutlass.range_constexpr(5):
+                                            oth2 = cute.arch.shuffle_sync_up(
+                                                pre6,
+                                                cutlass.Int32(1 << so2),
+                                                mask_and_clamp=0,
+                                            )
+                                            if lane >= cutlass.Int32(1 << so2):
+                                                pre6 = pre6 + oth2
+                                        needl3 = smem_hist[257]
+                                        if pre6 >= needl3 and (pre6 - ws3) < needl3:
+                                            base5 = pre6 - ws3
+                                            dstar2 = cutlass.Int32(0)
+                                            above5 = base5
+                                            sd5 = cutlass.Int32(0)
+                                            for jd4 in cutlass.range_constexpr(8):
+                                                di4 = (
+                                                    cutlass.Int32(255)
+                                                    - lane * cutlass.Int32(8)
+                                                    - cutlass.Int32(jd4)
+                                                )
+                                                ra5 = base5
+                                                base5 = base5 + smem_hist[di4]
+                                                if base5 >= needl3 and sd5 == cutlass.Int32(0):
+                                                    dstar2 = di4
+                                                    above5 = ra5
+                                                    sd5 = cutlass.Int32(1)
+                                            smem_hist[256] = uthr_c2 | (
+                                                dstar2 << cutlass.Int32(shift2)
+                                            )
+                                            smem_hist[257] = needl3 - above5
+                                            smem_hist[258] = smem_hist[258] + above5
+                                    cute.arch.barrier()
+                                u_thr2 = smem_hist[256]
+                                cnt_ab2 = smem_hist[258]
+                                need_eq2 = smem_hist[257]
+                                kthr2 = u_thr2 ^ cutlass.Int32(-2147483648)
                                 if tidx == cutlass.Int32(0):
-                                    needl = smem_hist[257]
-                                    cw = cutlass.Int32(0)
-                                    tw3 = cutlass.Int32(num_warps - 1)
-                                    f3 = cutlass.Int32(0)
-                                    for w4 in cutlass.range_constexpr(self.num_warps):
-                                        cw = cw + smem_wcnt[w4]
-                                        if cw >= needl and f3 == cutlass.Int32(0):
-                                            tw3 = cutlass.Int32(w4)
-                                            f3 = cutlass.Int32(1)
-                                    pre3 = cutlass.Int32(0)
-                                    for w5 in cutlass.range_constexpr(self.num_warps):
-                                        if cutlass.Int32(w5) < tw3:
-                                            pre3 = pre3 + smem_wcnt[w5]
-                                    s_iscalars[4] = pre3  # prefix above target warp
-                                    s_iscalars[0] = tw3  # target warp
+                                    s_iscalars[4] = cutlass.Int32(0)
+                                    s_iscalars[0] = cutlass.Int32(0)
                                 cute.arch.barrier()
-                                pre4 = s_iscalars[4]
-                                tw4 = s_iscalars[0]
-                                if warp_id == tw4 and lane == cutlass.Int32(0):
-                                    needl2 = smem_hist[257]
-                                    base4 = pre4
-                                    dstar = cutlass.Int32(0)
-                                    above_d = pre4
-                                    sd4 = cutlass.Int32(0)
-                                    for jd2 in cutlass.range_constexpr(fdw):
-                                        dix2 = (
-                                            cutlass.Int32(255)
-                                            - tw4 * cutlass.Int32(fdw)
-                                            - cutlass.Int32(jd2)
+                                ir3 = tidx
+                                while ir3 < cnt_strad:
+                                    vv3 = smem_keys[ir3]
+                                    uk4 = f32_order_key(vv3)
+                                    ks4 = uk4 ^ cutlass.Int32(-2147483648)
+                                    if ks4 > kthr2:
+                                        o4 = atomicAdd(
+                                            s_iscalars.iterator + cutlass.Int32(4),
+                                            cutlass.Int32(1),
                                         )
-                                        ra4 = base4
-                                        base4 = base4 + smem_hist[dix2]
-                                        if base4 >= needl2 and sd4 == cutlass.Int32(0):
-                                            dstar = dix2
-                                            above_d = ra4
-                                            sd4 = cutlass.Int32(1)
-                                    smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift))
-                                    smem_hist[257] = needl2 - above_d
-                                    smem_hist[258] = smem_hist[258] + above_d
-                                cute.arch.barrier()
-                            # Rewrite the tie slot range: ties with key > u_thr
-                            # first (there are exactly cnt_ab of them), then the
-                            # first need_eq bitwise-equal-to-u_thr ties in arrival
-                            # order (value-exact by construction). Signed compare
-                            # needs the top bit flipped (unsigned-monotonic key).
-                            u_thr = smem_hist[256]
-                            cnt_ab = smem_hist[258]
-                            need_eq = smem_hist[257]
-                            ks_thr = u_thr ^ cutlass.Int32(-2147483648)
-                            if tidx == cutlass.Int32(0):
-                                s_iscalars[4] = cutlass.Int32(0)  # above-writer ctr
-                                s_iscalars[0] = cutlass.Int32(0)  # equal-writer ctr
-                            cute.arch.barrier()
-                            ir2 = tidx
-                            while ir2 < cand_count:
-                                vr = smem_keys[ir2]
-                                br = cutlass.Int32((vr - bmin_r) * inv1)
-                                if br < cutlass.Int32(0):
-                                    br = cutlass.Int32(0)
-                                if br > cutlass.Int32(kBins - 1):
-                                    br = cutlass.Int32(kBins - 1)
-                                if br == b_star:
-                                    sr = cutlass.Int32((vr - f_lo) * finv)
-                                    if sr < cutlass.Int32(0):
-                                        sr = cutlass.Int32(0)
-                                    if sr > cutlass.Int32(fbins - 1):
-                                        sr = cutlass.Int32(fbins - 1)
-                                    if sr == sb_star:
-                                        uk2 = f32_order_key(vr)
-                                        ks2 = uk2 ^ cutlass.Int32(-2147483648)
-                                        if ks2 > ks_thr:
-                                            o2 = atomicAdd(
-                                                s_iscalars.iterator + cutlass.Int32(4),
-                                                cutlass.Int32(1),
-                                            )
-                                            pos = rank_above_fine + o2
+                                        pos = rank_above_fine + o4
+                                        if pos < cutlass.Int32(kK):
+                                            if cutlass.const_expr(self.return_output_values):
+                                                output_values_row[pos] = self.dtype(vv3)
+                                            output_indices_row[pos] = smem_vals[ir3]
+                                    elif ks4 == kthr2:
+                                        q4 = atomicAdd(
+                                            s_iscalars.iterator + cutlass.Int32(0),
+                                            cutlass.Int32(1),
+                                        )
+                                        if q4 < need_eq2:
+                                            pos = rank_above_fine + cnt_ab2 + q4
                                             if pos < cutlass.Int32(kK):
                                                 if cutlass.const_expr(self.return_output_values):
-                                                    output_values_row[pos] = self.dtype(vr)
-                                                output_indices_row[pos] = smem_vals[ir2]
-                                        elif ks2 == ks_thr:
-                                            q2 = atomicAdd(
-                                                s_iscalars.iterator + cutlass.Int32(0),
-                                                cutlass.Int32(1),
-                                            )
-                                            if q2 < need_eq:
-                                                pos = rank_above_fine + cnt_ab + q2
-                                                if pos < cutlass.Int32(kK):
-                                                    if cutlass.const_expr(
-                                                        self.return_output_values
-                                                    ):
-                                                        output_values_row[pos] = self.dtype(vr)
-                                                    output_indices_row[pos] = smem_vals[ir2]
-                                ir2 = ir2 + cutlass.Int32(num_threads)
-                            cute.arch.barrier()
+                                                    output_values_row[pos] = self.dtype(vv3)
+                                                output_indices_row[pos] = smem_vals[ir3]
+                                    ir3 = ir3 + cutlass.Int32(num_threads)
+                                cute.arch.barrier()
                 elif cutlass.const_expr(self.p4_exact_tail):  # [p4tt] if->elif only
                     need0 = cutlass.Int32(kK) - rank_above_fine
                     if cnt_strad > need0 and need0 > cutlass.Int32(0):
@@ -3628,6 +3674,17 @@ def phase4_rank_scatter(
                         output_values_row[ipad] = self.dtype(self.NEG_FLT_MAX)
                     output_indices_row[ipad] = cutlass.Int32(-1)
                     ipad = ipad + cutlass.Int32(num_threads)
+            if cutlass.const_expr(_P4_SUB_DBG):
+                # smem_wcnt slots [8..13] are dead after the last warp-sum
+                # use above; the take-block publish copies them to xstate.
+                sc6 = cute.arch.clock64()
+                if tidx == cutlass.Int32(0):
+                    smem_wcnt[8] = cutlass.Int32(sc1 - sc0)  # minmax
+                    smem_wcnt[9] = cutlass.Int32(sc2 - sc1)  # hist build
+                    smem_wcnt[10] = cutlass.Int32(sc3 - sc2)  # coarse search
+                    smem_wcnt[11] = cutlass.Int32(sc4 - sc3)  # fine recursion
+                    smem_wcnt[12] = cutlass.Int32(sc5 - sc4)  # scatter
+                    smem_wcnt[13] = cutlass.Int32(sc6 - sc5)  # tail repair+pad
         else:
             i10 = tidx
             while i10 < cand_count:
@@ -5779,6 +5836,10 @@ def _run_phases(
                         warp_id,
                         lane,
                     )
+                ck_sw0 = cutlass.Int64(0)
+                ck_sw1 = cutlass.Int64(0)
+                if cutlass.const_expr(_P4_SUB_DBG):
+                    ck_sw0 = cute.arch.clock64()
                 if cutlass.const_expr(
                     self.use_ext_cand and self.use_ext_counts and self.dtype == cutlass.Float32
                 ):
@@ -5806,6 +5867,8 @@ def _run_phases(
                                 )[0]
                             io_r = io_r + cutlass.Int32(num_threads)
                         cute.arch.barrier()
+                if cutlass.const_expr(_P4_SUB_DBG):
+                    ck_sw1 = cute.arch.clock64()
                 if cutlass.const_expr(self.emit_xstate):
                     # Closed-loop state (interface v2): [0] valid, [1] kth
                     # proxy (= accepted threshold; the tie-fill makes it a
@@ -5832,6 +5895,19 @@ def _run_phases(
                             xstate_row[5] = cutlass.Float32(cutlass.Int32(ck2 - ck1))  # P2/P3 gap
                             xstate_row[6] = cutlass.Float32(cutlass.Int32(ck3 - ck2))  # Phase 4
                             xstate_row[7] = s_thr[1]  # cnt_strad
+                        if cutlass.const_expr(_P4_SUB_DBG):
+                            # P4 sub-phase cycles staged by rank_scatter.
+                            # Chain-safe layout: [2] (closed-loop anchor)
+                            # untouched; [1] cnt_strad (tail class size),
+                            # [4] fine, [5] scatter, [6] tail, [7] deferred-
+                            # position swap. The small C-predictable phases
+                            # (minmax/hist/coarse, wcnt[8..10]) are not
+                            # published.
+                            xstate_row[1] = s_thr[1]
+                            xstate_row[4] = cutlass.Float32(smem_wcnt[11])
+                            xstate_row[5] = cutlass.Float32(smem_wcnt[12])
+                            xstate_row[6] = cutlass.Float32(smem_wcnt[13])
+                            xstate_row[7] = cutlass.Float32(cutlass.Int32(ck_sw1 - ck_sw0))
                         # cand_count_p4 = pre-P4 snapshot (P4 repurposes
                         # the s_iscalars slots).
                         xstate_row[3] = cutlass.Float32(cand_count_p4)

From da75cd82f18aed4b22b630a58d155d29a71d4635 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 09:35:55 -0700
Subject: [PATCH 022/117] [None][feat] GVR self_scan: fused self-contained
 closed-loop top-k mode

New self_scan mode: the kernel itself streams the row ONCE against the
three closed-loop seed lines and buckets candidates on the fly - no
external emitter, no indexer-side changes, no gmem candidate values.
One CTA per row, four phases: (0) scan-bucket - VALUES land in on-chip
segments (A/B/C at bases 0/B*/2B*, values-only 4B/entry, spill to the
looser segment, cursor totals ARE the line counts), POSITIONS stream to
a write-only gmem column reusing the cand_idx slot; (1) the v5 cut
state machine unchanged (a line cut compacts winning segment runs to
the smem prefix and fills smem_vals with segment coordinates, so P4,
the tail repair and the deferred K-gather run verbatim); ineligible
rows take the stock in-kernel fallback.

Scan-loop lessons baked in (each measured): per-element warp ballots
serialize every load (~1.8us/k); 16-wide register lists spill at 1024
threads (64 regs/thread ceiling) - values re-read from the load
fragments, positions derived arithmetically, classes recomputed;
warp-collective claim prefixes cap in-flight loads at 2/warp (ncu:
0.19% memory throughput) - final form claims passers with per-element
smem atomics, which do not synchronize the warp and hide under the
read stream (0.13us/k comp).

Exactness: 25-cell REPORT-S4 dataset x B in {1,2,4,8} = 100/100
bit-exact (flash/pro/v32 incl. K=2048, tiny-N and straddle fallbacks).
Perf vs PR16457 tip (same node, cold kernel-sum): geomean 0.64-0.71x,
short rows 0.8-1.05x, long rows 0.4-0.8x - the single-CTA read wall by
design; stage 2 (block-max skip) attacks the read itself.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 308 +++++++++++++++++-
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  42 ++-
 2 files changed, 345 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index c9b3a6800f40..194a2c1bd599 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -269,6 +269,8 @@ def __init__(
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
+        self_scan: bool = False,
+        cap_c: "int | None" = None,
         r0_qfracs: Optional[tuple] = None,
         mt_unroll: int = 4,
         p1b_cache: Optional[bool] = None,
@@ -590,6 +592,38 @@ def __init__(
         # segment that can void).
         self.list_cap = max(0, int(cand_cap) - 2 * self.accept_cap)
         self.cand_rung = int(cand_rung)
+        # self_scan (fused self-contained mode): the kernel itself streams
+        # the row ONCE against the three closed-loop lines, bucketing
+        # VALUES into on-chip segments (A >= t2 / B [t1,t2) / C [t0,t1),
+        # bases 0 / accept_cap / 2*accept_cap inside an enlarged
+        # smem_keys) and POSITIONS into the cand_idx tensor (write-only
+        # until the deferred K-gather). No external emitter, no candidate
+        # value column in gmem. A line cut compacts the winning segment
+        # runs to the smem_keys prefix and fills smem_vals with each
+        # entry's SEGMENT COORDINATE — from there on the v5 consumer
+        # (P4, tail repair, deferred gather via cand_idx[coord]) runs
+        # unchanged. Ineligible rows take the stock fallback, whose
+        # P3/P4 use smem_keys[:kC]/smem_vals verbatim.
+        self.self_scan = bool(self_scan)
+        if self.self_scan:
+            if not use_ext_counts:
+                raise ValueError("self_scan requires use_ext_counts")
+            if use_ext_cand:
+                raise ValueError("self_scan and use_ext_cand are exclusive")
+            if dtype != cutlass.Float32:
+                raise ValueError("self_scan is fp32-only (v1)")
+            if self.enable_smem_cache:
+                raise ValueError(
+                    "self_scan and enable_smem_cache exceed the CTA smem budget together"
+                )
+            # on-chip segment budget: values only, 4B/entry; C sized so
+            # keys(160KB) + vals(32KB) + hist + scratch stay under the
+            # 227KB CTA limit with room for the stage-2 skip list.
+            self.seg_total = 2 * self.accept_cap + int(cap_c if cap_c is not None else 24576)
+            self.cap_c = self.seg_total - 2 * self.accept_cap
+        else:
+            self.seg_total = self.kC
+            self.cap_c = 0
         if use_ext_cand and not use_ext_counts:
             raise ValueError("use_ext_cand requires use_ext_counts")
         self.use_ext_counts = bool(use_ext_counts) and bool(enable_r0)
@@ -1200,6 +1234,164 @@ def phase1b_hspace_rungs_cached(
     # → warp reduce → block reduce → s_iscalars[0] = cand_count.
     # Optionally DSMEM-aggregates across the cluster.
     # ------------------------------------------------------------------
+    @cute.jit
+    def phase0_scan_bucket(
+        self,
+        input_row,  # cute.Tensor [N] dtype (fp32; full row, cs==1 only)
+        N,  # int32 valid length (pad tail beyond N is never read)
+        seed_thr_row,  # [3] fp32 closed-loop lines, ascending t0 < t1 < t2
+        smem_keys,  # [seg_total] fp32 value segments (A @0 / B @segA / C @2segA)
+        cand_idx_row,  # [seg_total] int32 gmem POSITION column (write-only here)
+        s_seg,  # [>=7] int32 scratch (reuses smem_wcnt_p1: P1 never runs
+        #         on a row this phase succeeded on): [0..2] A/B/C claim
+        #         cursors, [3] void, [4] n0, [5] n1, [6] n2
+        tidx,
+        warp_id,
+        lane,
+    ):
+        """self_scan phase 0: ONE streaming pass buckets every element
+        >= t0 into the on-chip value segments (tightest line passed picks
+        the segment; a full segment spills to the next looser one) and
+        writes each entry's POSITION to the same coordinate of the gmem
+        column. The final cursor values ARE the line counts (attempts,
+        uncapped), so {n0, void, n1, n2} fall out for free — the same
+        contract the v5 emitter produced externally.
+
+        Perf shape (v3): per round each thread front-loads TWO vec_w
+        vectors (independent LDGs, latency overlapped), classifies
+        branchlessly in registers, and the whole round pays ONE packed
+        warp shfl-prefix (all three segment counts in 10-bit fields) +
+        at most three warp atomics. Segment overflow is marked and
+        resolved by a RARE per-element direct-atomic pass (a segment
+        overflows at most once per row, and divergent scalar atomics
+        need no warp coordination)."""
+        num_threads = cutlass.const_expr(self.num_threads)
+        segA = cutlass.const_expr(self.accept_cap)
+        capC = cutlass.const_expr(self.cap_c)
+        vec_w = cutlass.const_expr(self.vec_bits // self.dtype.width)
+        elem_bytes = cutlass.const_expr(self.dtype.width // 8)
+        vec_align = cutlass.const_expr(self.vec_align_bytes)
+        if tidx == cutlass.Int32(0):
+            s_seg[0] = cutlass.Int32(0)
+            s_seg[1] = cutlass.Int32(0)
+            s_seg[2] = cutlass.Int32(0)
+        cute.arch.barrier()
+        t0_s = seed_thr_row[0]
+        t1_s = seed_thr_row[1]
+        t2_s = seed_thr_row[2]
+        row_addr = input_row.iterator.toint()
+        copy_atom = self._make_load_copy_atom()
+        frag_a = cute.make_fragment((vec_w,), self.dtype)
+        frag_b = cute.make_fragment((vec_w,), self.dtype)
+        step1 = cutlass.const_expr(num_threads * vec_w)
+        step2 = cutlass.const_expr(2 * step1)
+        st2log = cutlass.const_expr((2 * step1).bit_length() - 1)
+        # FULL-vector rounds only in the hot loop: nfull is warp-uniform
+        # (every lane's two vectors are in bounds by construction), so the
+        # warp collectives inside are legal and the hot path carries no
+        # bounds checks at all. The remainder (< 2*step1 elements) takes
+        # the scalar tail below with per-element direct atomics.
+        nfull = N >> st2log
+        it0 = cutlass.Int32(0)
+        while it0 < nfull:
+            ia0 = (it0 * cutlass.Int32(2) * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(
+                vec_w
+            )
+            ib0 = ia0 + cutlass.Int32(step1)
+            for _fq in cutlass.range_constexpr(2):
+                fq0 = ia0 if _fq == 0 else ib0
+                pq0 = cute.make_ptr(
+                    self.dtype,
+                    row_addr + cutlass.Int64(fq0) * cutlass.Int64(elem_bytes),
+                    cute.AddressSpace.gmem,
+                    assumed_align=vec_align,
+                )
+                cute.copy(
+                    copy_atom,
+                    cute.make_tensor(pq0, cute.make_layout((vec_w,))),
+                    frag_a if _fq == 0 else frag_b,
+                )
+            # v6: per-element DIRECT atomic claims for passers — smem
+            # atomics do NOT synchronize the warp, so the vector loads of
+            # later rounds keep flowing (the packed shfl-prefix design
+            # capped in-flight loads at 2/warp: ncu showed 0.19% memory
+            # throughput, pure latency bound). Same-address service is
+            # ~0.5ns effective and n0 is line-bounded, so the claim queue
+            # hides completely under the read stream.
+            for _jh in cutlass.range_constexpr(2):
+                for _jv in cutlass.range_constexpr(vec_w):
+                    v0 = cutlass.Float32(frag_a[_jv]) if _jh == 0 else cutlass.Float32(frag_b[_jv])
+                    if v0 >= t0_s:
+                        pos0 = (ia0 if _jh == 0 else ib0) + cutlass.Int32(_jv)
+                        c0 = cutlass.Int32(2)
+                        if v0 >= t1_s:
+                            c0 = cutlass.Int32(1)
+                        if v0 >= t2_s:
+                            c0 = cutlass.Int32(0)
+                        while c0 >= cutlass.Int32(0) and c0 <= cutlass.Int32(2):
+                            cap0 = cutlass.Int32(segA)
+                            if c0 == cutlass.Int32(2):
+                                cap0 = cutlass.Int32(capC)
+                            sl0 = atomicAdd(s_seg.iterator + c0, cutlass.Int32(1))
+                            if sl0 < cap0:
+                                cd0 = c0 * cutlass.Int32(segA) + sl0
+                                smem_keys[cd0] = v0
+                                cand_idx_row[cd0] = pos0
+                                c0 = cutlass.Int32(-1)
+                            else:
+                                c0 = c0 + cutlass.Int32(1)
+            it0 = it0 + cutlass.Int32(1)
+        # scalar tail (< 2*step1 elements): per-element DIRECT atomic
+        # claims — divergent-safe, no warp collectives
+        pt0 = (N >> st2log) * cutlass.Int32(step2) + tidx
+        while pt0 < N:
+            spt = cute.make_ptr(
+                cutlass.Float32,
+                row_addr + cutlass.Int64(pt0) * cutlass.Int64(4),
+                cute.AddressSpace.gmem,
+                assumed_align=4,
+            )
+            vt0 = cute.make_tensor(spt, cute.make_layout((1,)))[0]
+            if vt0 >= t0_s:
+                ct0 = cutlass.Int32(2)
+                if vt0 >= t1_s:
+                    ct0 = cutlass.Int32(1)
+                if vt0 >= t2_s:
+                    ct0 = cutlass.Int32(0)
+                while ct0 >= cutlass.Int32(0) and ct0 <= cutlass.Int32(2):
+                    capt = cutlass.Int32(segA)
+                    if ct0 == cutlass.Int32(2):
+                        capt = cutlass.Int32(capC)
+                    slt = atomicAdd(s_seg.iterator + ct0, cutlass.Int32(1))
+                    if slt < capt:
+                        cdt = ct0 * cutlass.Int32(segA) + slt
+                        smem_keys[cdt] = vt0
+                        cand_idx_row[cdt] = pt0
+                        ct0 = cutlass.Int32(-1)
+                    else:
+                        ct0 = ct0 + cutlass.Int32(1)
+            pt0 = pt0 + cutlass.Int32(num_threads)
+        cute.arch.barrier()
+        if tidx == cutlass.Int32(0):
+            curA0 = s_seg[0]
+            curB0 = s_seg[1]
+            curC0 = s_seg[2]
+            spA0 = curA0 - cutlass.Int32(segA)
+            if spA0 < cutlass.Int32(0):
+                spA0 = cutlass.Int32(0)
+            spB0 = curB0 - cutlass.Int32(segA)
+            if spB0 < cutlass.Int32(0):
+                spB0 = cutlass.Int32(0)
+            n1_0 = curA0 + curB0 - spA0
+            n0_0 = n1_0 + curC0 - spB0
+            s_seg[3] = cutlass.Int32(0)
+            if curC0 > cutlass.Int32(capC):
+                s_seg[3] = cutlass.Int32(1)
+            s_seg[4] = n0_0
+            s_seg[5] = n1_0
+            s_seg[6] = curA0
+        cute.arch.barrier()
+
     @cute.jit
     def block_count_ge(
         self,
@@ -4387,6 +4579,12 @@ def run_one_row(
             cand_vals_row = cand_vals[row_idx, None]
             cand_idx_row = cand_idx[row_idx, None]
             cand_ctl_row = cand_ctl[row_idx, None]
+        elif cutlass.const_expr(self.self_scan and cand_idx is not None):
+            # self_scan: only the POSITION column exists (values live in
+            # smem from birth; counts come from the phase-0 cursors)
+            cand_vals_row = None
+            cand_idx_row = cand_idx[row_idx, None]
+            cand_ctl_row = None
         else:
             cand_vals_row = None
             cand_idx_row = None
@@ -4407,12 +4605,18 @@ def run_one_row(
         smem = SmemAllocator()
         # keys[kC] fp32 (P3 candidate values; smem keys always fp32 even for half-prec)
         # Use fp32 even for half-prec to make secant search algorithm keep the accuracy/precision and converge faster.
+        # self_scan: enlarged to seg_total (three value segments at bases
+        # 0 / accept_cap / 2*accept_cap); every later consumer only ever
+        # touches a <= kC prefix after cut compaction.
         smem_keys = smem.allocate_tensor(
             element_type=cutlass.Float32,
-            layout=cute.make_ordered_layout((kC,), order=(0,)),
+            layout=cute.make_ordered_layout((cutlass.const_expr(self.seg_total),), order=(0,)),
             byte_alignment=128,
         )
-        # vals[kC] int32 (P3 candidate indices)
+        # vals[kC] int32 (P3 candidate indices). self_scan: holds the
+        # SEGMENT COORDINATE of each compacted candidate (identity for the
+        # deferred position gather via cand_idx[coord]) — every consumer
+        # (P4, tail repair, gather) works unchanged.
         smem_vals = smem.allocate_tensor(
             element_type=cutlass.Int32,
             layout=cute.make_ordered_layout((kC,), order=(0,)),
@@ -4900,6 +5104,34 @@ def _run_phases(
                     and seed_thr_row[0] < cutlass.Float32(1e37)
                 ):
                     ext_row = cutlass.Int32(1)
+            # ---- self_scan phase 0: fused scan-bucket ----
+            # The kernel streams the row itself (no external emitter);
+            # eligibility mirrors the list contract: nothing dropped
+            # (void == 0) and the loosest line provably covers the top-K
+            # (n0 >= K, exact counts - no sentinel slack needed). The
+            # seed-thr finite guard keeps the branch CTA-uniform, so the
+            # barriers/ballots inside phase 0 stay convergent. Cursor
+            # scratch = smem_wcnt_p1 (P1 only runs when this row is NOT
+            # taken, so the reuse never overlaps live data).
+            if cutlass.const_expr(
+                self.self_scan and cluster_size == 1 and self.dtype == cutlass.Float32
+            ):
+                if seed_thr_row[0] < cutlass.Float32(1e37):
+                    self.phase0_scan_bucket(
+                        input_row,
+                        N,
+                        seed_thr_row,
+                        smem_keys,
+                        cand_idx_row,
+                        smem_wcnt_p1,
+                        tidx,
+                        warp_id,
+                        lane,
+                    )
+                    if smem_wcnt_p1[3] == cutlass.Int32(0) and smem_wcnt_p1[4] >= cutlass.Int32(
+                        self.top_k
+                    ):
+                        ext_row = cutlass.Int32(1)
 
         # ---- Phase 1: preIdx Min/Max/Mean ----
         # ext counts: P1's only surviving products are the [v_lo, v_hi]
@@ -5377,6 +5609,74 @@ def _run_phases(
                 if cutlass.const_expr(_P4_TAIL_DBG):
                     ck1 = cute.arch.clock64()
 
+            # ---- self_scan take: cut straight from the phase-0 cursors ----
+            # Same admission state machine as the v5 list (tightest line
+            # whose count fits [K, B*] wins; anchor = loosest in-band
+            # line), but the candidates already LIVE in smem: a line cut
+            # is a same-buffer run compaction (sources at >= segA,
+            # destinations below it - disjoint) plus the segment-coordinate
+            # fill of smem_vals that the unchanged P4 / tail repair /
+            # deferred position gather consume. Straddle / overshoot rows
+            # (no line in band) fall through to the stock fallback (v1).
+            if cutlass.const_expr(
+                self.self_scan and cluster_size == 1 and self.dtype == cutlass.Float32
+            ):
+                segA_f = cutlass.const_expr(self.accept_cap)
+                if cutlass.const_expr(_P4_TAIL_DBG):
+                    ck0 = cute.arch.clock64()
+                    ck1 = ck0
+                n0_f = smem_wcnt_p1[4]
+                n1_f = smem_wcnt_p1[5]
+                n2_f = smem_wcnt_p1[6]
+                kK_f = cutlass.Int32(self.top_k)
+                bs_f = cutlass.Int32(segA_f)
+                cut_n = cutlass.Int32(0)
+                have_f = cutlass.Int32(0)
+                anch_f = cutlass.Float32(0.0)
+                if ext_row == cutlass.Int32(1):
+                    if n2_f >= kK_f and n2_f <= bs_f:
+                        if have_f == cutlass.Int32(0):
+                            cut_n = n2_f
+                        anch_f = seed_thr_row[2]
+                        have_f = cutlass.Int32(1)
+                    if n1_f >= kK_f and n1_f <= bs_f:
+                        if have_f == cutlass.Int32(0):
+                            cut_n = n1_f
+                        anch_f = seed_thr_row[1]
+                        have_f = cutlass.Int32(1)
+                    if n0_f <= bs_f:
+                        if have_f == cutlass.Int32(0):
+                            cut_n = n0_f
+                        anch_f = seed_thr_row[0]
+                        have_f = cutlass.Int32(1)
+                if have_f == cutlass.Int32(1):
+                    take_cand = cutlass.Int32(1)
+                    list_used = cutlass.Int32(1)
+                    if tidx == cutlass.Int32(0):
+                        s_iscalars[0] = cut_n
+                        s_thr[0] = anch_f
+                        s_iscalars[1] = cutlass.Int32(1)  # done
+                    cute.arch.barrier()
+                    # line cut => no segment spilled (n_cut <= B* bounds
+                    # every tighter count too) => lenA = n2, lenB = n1-n2
+                    j_f = tidx
+                    while j_f < cut_n:
+                        for _jw in cutlass.range_constexpr(4):
+                            q_f = j_f + cutlass.Int32(_jw * num_threads)
+                            if q_f < cut_n:
+                                src_f = q_f
+                                if q_f >= n2_f:
+                                    src_f = cutlass.Int32(segA_f) + q_f - n2_f
+                                if q_f >= n1_f:
+                                    src_f = cutlass.Int32(2 * segA_f) + q_f - n1_f
+                                if src_f != q_f:
+                                    smem_keys[q_f] = smem_keys[src_f]
+                                smem_vals[q_f] = src_f
+                        j_f = j_f + cutlass.Int32(4 * num_threads)
+                    cute.arch.barrier()
+                if cutlass.const_expr(_P4_TAIL_DBG):
+                    ck1 = cute.arch.clock64()
+
             if take_cand == cutlass.Int32(0):
                 # Stage this CTA's slice into SMEM once before Phase 2's
                 # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
@@ -5841,7 +6141,9 @@ def _run_phases(
                 if cutlass.const_expr(_P4_SUB_DBG):
                     ck_sw0 = cute.arch.clock64()
                 if cutlass.const_expr(
-                    self.use_ext_cand and self.use_ext_counts and self.dtype == cutlass.Float32
+                    (self.use_ext_cand or self.self_scan)
+                    and self.use_ext_counts
+                    and self.dtype == cutlass.Float32
                 ):
                     # List rows: the compact stored LIST INDICES in the
                     # vals slots (saving a second cold gmem pass over the
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 70a64d2e0e28..9f7e5ac0997d 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -72,6 +72,8 @@ def _compile(
     cand_cap: int = 5120,
     accept_cap: "int | None" = None,
     kc_override: "int | None" = None,
+    self_scan: bool = False,
+    cap_c: "int | None" = None,
 ):
     """JIT-compile the GVR kernel for a specific knob combination.
 
@@ -165,7 +167,7 @@ def _compile(
         cute.runtime.make_fake_compact_tensor(
             cutlass.Int32, (n_rows, cand_cap), stride_order=(1, 0), assumed_align=4
         )
-        if use_ext_cand
+        if (use_ext_cand or self_scan)
         else None
     )
     cand_ctl_fake = (
@@ -207,6 +209,8 @@ def _compile(
         cand_cap=cand_cap,
         accept_cap=accept_cap,
         kc_override=kc_override,
+        self_scan=self_scan,
+        cap_c=cap_c,
         # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed. The
         # qfrac VALUES are irrelevant on this path (P1b is skipped) — only
         # the slot count matters.
@@ -647,6 +651,8 @@ def gvr_topk_decode(
     cand_vals: Optional[torch.Tensor] = None,
     cand_idx: Optional[torch.Tensor] = None,
     cand_ctl: Optional[torch.Tensor] = None,
+    self_scan: bool = False,
+    cap_c: Optional[int] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """CuTe DSL GVR Top-K wrapper with every tuning knob exposed.
 
@@ -729,6 +735,32 @@ def gvr_topk_decode(
     # cs1/cs8 vs stock cs8 19.7us) -> keep the stock path.
     if block_max is not None and num_rows < 8 and top_k > 512:
         block_max = None
+    if self_scan:
+        # fused self-contained mode: kernel scans/buckets the row itself.
+        # Inputs: seed_thr (three closed-loop lines) + a write-only gmem
+        # POSITION column passed through the cand_idx slot. seed_counts is
+        # a dummy (the ext-counts preview reads it but zeros never pass
+        # the [K, kC] admission, so routing is owned by the phase-0 gate).
+        assert seed_thr is not None, "self_scan requires seed_thr"
+        assert cand_vals is None and cand_ctl is None, (
+            "self_scan excludes external candidate values/control"
+        )
+        assert "GVR_BSTAR" in os.environ, (
+            "self_scan requires GVR_BSTAR (accept_cap) to size the position column"
+        )
+        if seed_counts is None:
+            seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=logits.device)
+        _bstar = int(os.environ["GVR_BSTAR"])
+        _capc = cap_c if cap_c is not None else int(os.environ.get("GVR_CAPC", "24576"))
+        _segtot = 2 * _bstar + _capc
+        if cand_idx is None:
+            cand_idx = torch.empty((num_rows, _segtot), dtype=torch.int32, device=logits.device)
+        assert (
+            cand_idx.dtype == torch.int32
+            and cand_idx.is_cuda
+            and cand_idx.is_contiguous()
+            and cand_idx.shape == (num_rows, _segtot)
+        ), f"self_scan position column must be int32 [num_rows, {_segtot}]"
     use_ext_counts = seed_thr is not None and seed_counts is not None
     if use_ext_counts:
         assert (
@@ -745,6 +777,8 @@ def gvr_topk_decode(
         ), "seed_counts must be contiguous CUDA int32 [num_rows, 3]"
     use_ext_cand = cand_vals is not None and cand_idx is not None and cand_ctl is not None
     cand_cap = 5120
+    if self_scan:
+        cand_cap = cand_idx.shape[1]
     if use_ext_cand:
         assert (
             cand_vals.dtype == torch.float32
@@ -858,6 +892,10 @@ def gvr_topk_decode(
         cand_cap,
         int(os.environ["GVR_BSTAR"]) if "GVR_BSTAR" in os.environ else None,
         int(os.environ["GVR_KC"]) if "GVR_KC" in os.environ else None,
+        self_scan,
+        cap_c
+        if cap_c is not None
+        else (int(os.environ.get("GVR_CAPC", "24576")) if self_scan else None),
     )
     # When return_output_values=False the kernel was compiled to skip
     # STG.value and accepts None for the value-output slot.
@@ -875,7 +913,7 @@ def gvr_topk_decode(
         seed_counts if use_ext_counts else None,
         xstate if emit_xstate else None,
         cand_vals if use_ext_cand else None,
-        cand_idx if use_ext_cand else None,
+        cand_idx if (use_ext_cand or self_scan) else None,
         cand_ctl if use_ext_cand else None,
     )
     if return_output_values:

From 7970473b72104c94a1c7f29c8e12377d27bb5288 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 10:12:00 -0700
Subject: [PATCH 023/117] [None][feat] GVR self_scan stage 2: block-max skip,
 tightest-line single band

Phase 0 gains a block-skip variant (enable_block_skip + self_scan):
per-32-position maxima from the GEMM tail gate whole blocks out of the
scan. Measured design pivot: skipping against the LOOSE collection
line can never pay (n0/N ~5-12% density -> ~80-99% of blocks contain a
passer; benched 0.16-0.39x), so the skip mode collects a SINGLE BAND
against the TIGHTEST line (density 0.4-0.8% -> 12-22% pass): only
segment A fills, the cursor keeps exact attempt counts, and the v5
state machine runs unchanged fed n0 == n1 == n2 - a cut lands on t2
(common), the sample-hist path absorbs over-B* rows (the A prefix
stays a value-blind sample), under-K rows take the stock fallback.
The small-batch block_max gate in the wrapper is bypassed for
self_scan (stage 2 owns its own skip economics).

Exactness: 25-cell REPORT-S4 x B in {1,2,4,8} = 100/100 bit-exact,
plus forced under-K fallback cells. Perf state (B=1 vs PR16457 tip,
same node): long rows improve markedly over the dense scan (flash
512k 33.9 -> 24.5us = 0.87x of tip; 1024k 46.7 -> 38.7; pro 1M 52 ->
44) while short/mid rows should route to the dense scan (host picks
by expected block pass rate). Known remaining work, measured and
documented: the block loop is still latency-bound on the bmax stream
(8 scalar loads/warp round); a lane-per-block + ballot variant was
tried and loses at high pass rates - loop shape per density regime is
the open optimization, along with a t2-only closed-loop line-derive
for chains.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 106 +++++++++++++++---
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |   4 +-
 2 files changed, 92 insertions(+), 18 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 194a2c1bd599..9973af5387e0 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -1242,6 +1242,7 @@ def phase0_scan_bucket(
         seed_thr_row,  # [3] fp32 closed-loop lines, ascending t0 < t1 < t2
         smem_keys,  # [seg_total] fp32 value segments (A @0 / B @segA / C @2segA)
         cand_idx_row,  # [seg_total] int32 gmem POSITION column (write-only here)
+        block_max_row,  # [nb_pad] fp32 per-32-position maxima, or None
         s_seg,  # [>=7] int32 scratch (reuses smem_wcnt_p1: P1 never runs
         #         on a row this phase succeeded on): [0..2] A/B/C claim
         #         cursors, [3] void, [4] n0, [5] n1, [6] n2
@@ -1280,6 +1281,61 @@ def phase0_scan_bucket(
         t1_s = seed_thr_row[1]
         t2_s = seed_thr_row[2]
         row_addr = input_row.iterator.toint()
+        # ---- stage-2 BLOCK-SKIP variant: the GEMM tail left per-32-
+        # position maxima; a block whose max < t0 contributes nothing to
+        # any count or segment, so it is never read. One warp per block:
+        # the bmax compare is warp-uniform (all lanes read the same
+        # scalar), a passing block is one coalesced 128B load, claims
+        # are the same non-synchronizing per-element atomics as the
+        # dense loop - so the block loop needs no uniform trip counts.
+        if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
+            # v8.2: 8 blocks per warp iteration. Every lane vector-loads
+            # the SAME 8 bmax values (L1 broadcast - the pass decisions
+            # are warp-uniform registers), then the passing blocks are
+            # loaded back-to-back (independent 128B coalesced loads, so
+            # the memory level parallelism the naive one-block loop
+            # lacked is restored; at high skip rates an iteration is
+            # just the one 32B bmax vector).
+            bm_addr = block_max_row.iterator.toint()
+            nb0 = (N + cutlass.Int32(31)) >> cutlass.Int32(5)
+            nwp = cutlass.const_expr(self.num_warps)
+            frag_m = cute.make_fragment((8,), cutlass.Float32)
+            bg0 = warp_id * cutlass.Int32(8)
+            while bg0 < nb0:
+                for _jm in cutlass.range_constexpr(8):
+                    frag_m[_jm] = cutlass.Float32(self.NEG_FLT_MAX)
+                    if bg0 + cutlass.Int32(_jm) < nb0:
+                        bps = cute.make_ptr(
+                            cutlass.Float32,
+                            bm_addr + cutlass.Int64(bg0 + cutlass.Int32(_jm)) * cutlass.Int64(4),
+                            cute.AddressSpace.gmem,
+                            assumed_align=4,
+                        )
+                        frag_m[_jm] = cute.make_tensor(bps, cute.make_layout((1,)))[0]
+                # SINGLE-BAND collection against the TIGHTEST line: the
+                # loose band's pass fraction (~80% of blocks) can never
+                # pay for skipping, the accepted band's (~12-22%) can.
+                # Only segment A is filled; overflow just drops (the
+                # cursor keeps counting, and the A prefix remains a
+                # value-blind SAMPLE - the sample-cut path handles the
+                # over-B* case exactly as with the external list).
+                for _jb in cutlass.range_constexpr(8):
+                    if cutlass.Float32(frag_m[_jb]) >= t2_s:
+                        pos0 = ((bg0 + cutlass.Int32(_jb)) << cutlass.Int32(5)) + lane
+                        if pos0 < N:
+                            vp0 = cute.make_ptr(
+                                cutlass.Float32,
+                                row_addr + cutlass.Int64(pos0) * cutlass.Int64(4),
+                                cute.AddressSpace.gmem,
+                                assumed_align=4,
+                            )
+                            v0 = cute.make_tensor(vp0, cute.make_layout((1,)))[0]
+                            if v0 >= t2_s:
+                                sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1))
+                                if sl0 < cutlass.Int32(segA):
+                                    smem_keys[sl0] = v0
+                                    cand_idx_row[sl0] = pos0
+                bg0 = bg0 + cutlass.Int32(nwp * 8)
         copy_atom = self._make_load_copy_atom()
         frag_a = cute.make_fragment((vec_w,), self.dtype)
         frag_b = cute.make_fragment((vec_w,), self.dtype)
@@ -1292,6 +1348,8 @@ def phase0_scan_bucket(
         # bounds checks at all. The remainder (< 2*step1 elements) takes
         # the scalar tail below with per-element direct atomics.
         nfull = N >> st2log
+        if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
+            nfull = cutlass.Int32(0)
         it0 = cutlass.Int32(0)
         while it0 < nfull:
             ia0 = (it0 * cutlass.Int32(2) * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(
@@ -1344,6 +1402,8 @@ def phase0_scan_bucket(
         # scalar tail (< 2*step1 elements): per-element DIRECT atomic
         # claims — divergent-safe, no warp collectives
         pt0 = (N >> st2log) * cutlass.Int32(step2) + tidx
+        if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
+            pt0 = N
         while pt0 < N:
             spt = cute.make_ptr(
                 cutlass.Float32,
@@ -1373,23 +1433,34 @@ def phase0_scan_bucket(
             pt0 = pt0 + cutlass.Int32(num_threads)
         cute.arch.barrier()
         if tidx == cutlass.Int32(0):
-            curA0 = s_seg[0]
-            curB0 = s_seg[1]
-            curC0 = s_seg[2]
-            spA0 = curA0 - cutlass.Int32(segA)
-            if spA0 < cutlass.Int32(0):
-                spA0 = cutlass.Int32(0)
-            spB0 = curB0 - cutlass.Int32(segA)
-            if spB0 < cutlass.Int32(0):
-                spB0 = cutlass.Int32(0)
-            n1_0 = curA0 + curB0 - spA0
-            n0_0 = n1_0 + curC0 - spB0
-            s_seg[3] = cutlass.Int32(0)
-            if curC0 > cutlass.Int32(capC):
-                s_seg[3] = cutlass.Int32(1)
-            s_seg[4] = n0_0
-            s_seg[5] = n1_0
-            s_seg[6] = curA0
+            if cutlass.const_expr(self.enable_block_skip):
+                # single-band mode: every count is the t2 cursor; a cut
+                # can only land on t2 (in band), the sample-hist (over)
+                # or the fallback (under) - exactly the v5 state machine
+                # fed with n0 == n1 == n2
+                curT0 = s_seg[0]
+                s_seg[3] = cutlass.Int32(0)
+                s_seg[4] = curT0
+                s_seg[5] = curT0
+                s_seg[6] = curT0
+            if cutlass.const_expr(not self.enable_block_skip):
+                curA0 = s_seg[0]
+                curB0 = s_seg[1]
+                curC0 = s_seg[2]
+                spA0 = curA0 - cutlass.Int32(segA)
+                if spA0 < cutlass.Int32(0):
+                    spA0 = cutlass.Int32(0)
+                spB0 = curB0 - cutlass.Int32(segA)
+                if spB0 < cutlass.Int32(0):
+                    spB0 = cutlass.Int32(0)
+                n1_0 = curA0 + curB0 - spA0
+                n0_0 = n1_0 + curC0 - spB0
+                s_seg[3] = cutlass.Int32(0)
+                if curC0 > cutlass.Int32(capC):
+                    s_seg[3] = cutlass.Int32(1)
+                s_seg[4] = n0_0
+                s_seg[5] = n1_0
+                s_seg[6] = curA0
         cute.arch.barrier()
 
     @cute.jit
@@ -5123,6 +5194,7 @@ def _run_phases(
                         seed_thr_row,
                         smem_keys,
                         cand_idx_row,
+                        block_max_row,
                         smem_wcnt_p1,
                         tidx,
                         warp_id,
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 9f7e5ac0997d..0b4a3e186402 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -733,7 +733,9 @@ def gvr_topk_decode(
     # (kC/K = 6 vs 10), the bounds prune less, and the row-split configs
     # win outright (cold protocol, pro 262k BS1: skip 21.3-21.6us at
     # cs1/cs8 vs stock cs8 19.7us) -> keep the stock path.
-    if block_max is not None and num_rows < 8 and top_k > 512:
+    if block_max is not None and num_rows < 8 and top_k > 512 and not self_scan:
+        # self_scan stage 2 owns its own skip economics (phase-0 block
+        # loop) — the stock-path small-batch gate does not apply
         block_max = None
     if self_scan:
         # fused self-contained mode: kernel scans/buckets the row itself.

From 579602117f8262d4cfff0cc1b758ca21dfb5a8ca Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 10:46:43 -0700
Subject: [PATCH 024/117] [None][perf] GVR self_scan skip: two-pass
 compact-then-gather block scan

The skip scan is restructured into two passes that both run at the
tuned dense-loop shape: (1) DENSE-vector-scan the block-max array
itself (1/32 of the row, 128-bit vectors - the bmax row base is only
16B aligned) and compact the PASSING BLOCK IDS into the idle C segment
(single-band mode never fills C; ids store exactly as floats);
(2) walk the compact list, eight listed blocks per warp round issued
back-to-back - every element read is useful and the loads pipeline.
A list overflow (pass rate too high for skipping to ever pay) falls
back to a dense full scan of the row inside the same phase.

This removes both latency walls the one-pass shapes hit (8-scalar
bmax rounds; serial per-block walks): flash 512k drops 25 -> 20us and
BEATS the PR16457 tip (1.06x) - first cell where the fused
self-contained kernel wins outright; flash 1024k 47 -> 26us (0.72x of
tip), pro 1M 53 -> 36us, v32 128k+ 28us. Dense/skip best-of geomean
0.69 -> 0.73-0.75x across the 25-cell REPORT-S4 dataset, all 100
cells bit-exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 157 +++++++++++++-----
 1 file changed, 120 insertions(+), 37 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 9973af5387e0..84616b1e573c 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -1298,44 +1298,127 @@ def phase0_scan_bucket(
             # just the one 32B bmax vector).
             bm_addr = block_max_row.iterator.toint()
             nb0 = (N + cutlass.Int32(31)) >> cutlass.Int32(5)
-            nwp = cutlass.const_expr(self.num_warps)
-            frag_m = cute.make_fragment((8,), cutlass.Float32)
-            bg0 = warp_id * cutlass.Int32(8)
-            while bg0 < nb0:
-                for _jm in cutlass.range_constexpr(8):
-                    frag_m[_jm] = cutlass.Float32(self.NEG_FLT_MAX)
-                    if bg0 + cutlass.Int32(_jm) < nb0:
-                        bps = cute.make_ptr(
-                            cutlass.Float32,
-                            bm_addr + cutlass.Int64(bg0 + cutlass.Int32(_jm)) * cutlass.Int64(4),
-                            cute.AddressSpace.gmem,
-                            assumed_align=4,
-                        )
-                        frag_m[_jm] = cute.make_tensor(bps, cute.make_layout((1,)))[0]
-                # SINGLE-BAND collection against the TIGHTEST line: the
-                # loose band's pass fraction (~80% of blocks) can never
-                # pay for skipping, the accepted band's (~12-22%) can.
-                # Only segment A is filled; overflow just drops (the
-                # cursor keeps counting, and the A prefix remains a
-                # value-blind SAMPLE - the sample-cut path handles the
-                # over-B* case exactly as with the external list).
-                for _jb in cutlass.range_constexpr(8):
-                    if cutlass.Float32(frag_m[_jb]) >= t2_s:
-                        pos0 = ((bg0 + cutlass.Int32(_jb)) << cutlass.Int32(5)) + lane
-                        if pos0 < N:
-                            vp0 = cute.make_ptr(
-                                cutlass.Float32,
-                                row_addr + cutlass.Int64(pos0) * cutlass.Int64(4),
-                                cute.AddressSpace.gmem,
-                                assumed_align=4,
+            # v9 two-pass skip: (1) DENSE-scan the bmax array itself (it
+            # is 1/32 of the row) with the tuned vector loop, compacting
+            # PASSING BLOCK IDS into the idle C segment (single-band mode
+            # never fills C; ids < 2^23 store exactly as floats);
+            # (2) walk the compact list, 8 blocks per warp round issued
+            # unguarded back-to-back - every element read is useful and
+            # the loads pipeline. If the list overflows capC the row
+            # falls back to the dense full scan (routing should have
+            # sent it there anyway).
+            # pass-1 vectors: 128-bit (the bmax row base is only 16B
+            # aligned: nb_pad %% 4)
+            pass1_atom = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(),
+                cutlass.Float32,
+                num_bits_per_copy=128,
+            )
+            p1w = cutlass.const_expr(4)
+            frag_p = cute.make_fragment((p1w,), cutlass.Float32)
+            nfb0 = nb0 >> cutlass.const_expr((num_threads * 4).bit_length() - 1)
+            itp0 = cutlass.Int32(0)
+            while itp0 < nfb0:
+                ip0 = (itp0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(p1w)
+                pp0 = cute.make_ptr(
+                    cutlass.Float32,
+                    bm_addr + cutlass.Int64(ip0) * cutlass.Int64(4),
+                    cute.AddressSpace.gmem,
+                    assumed_align=16,
+                )
+                cute.copy(
+                    pass1_atom,
+                    cute.make_tensor(pp0, cute.make_layout((p1w,))),
+                    frag_p,
+                )
+                for _jp in cutlass.range_constexpr(p1w):
+                    if cutlass.Float32(frag_p[_jp]) >= t2_s:
+                        slp0 = atomicAdd(s_seg.iterator + cutlass.Int32(1), cutlass.Int32(1))
+                        if slp0 < cutlass.Int32(capC):
+                            smem_keys[cutlass.Int32(2 * segA) + slp0] = cutlass.Float32(
+                                ip0 + cutlass.Int32(_jp)
                             )
-                            v0 = cute.make_tensor(vp0, cute.make_layout((1,)))[0]
-                            if v0 >= t2_s:
-                                sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1))
-                                if sl0 < cutlass.Int32(segA):
-                                    smem_keys[sl0] = v0
-                                    cand_idx_row[sl0] = pos0
-                bg0 = bg0 + cutlass.Int32(nwp * 8)
+                itp0 = itp0 + cutlass.Int32(1)
+            ptb0 = nfb0 * cutlass.Int32(num_threads * 4) + tidx
+            while ptb0 < nb0:
+                bpt0 = cute.make_ptr(
+                    cutlass.Float32,
+                    bm_addr + cutlass.Int64(ptb0) * cutlass.Int64(4),
+                    cute.AddressSpace.gmem,
+                    assumed_align=4,
+                )
+                if cute.make_tensor(bpt0, cute.make_layout((1,)))[0] >= t2_s:
+                    slp0 = atomicAdd(s_seg.iterator + cutlass.Int32(1), cutlass.Int32(1))
+                    if slp0 < cutlass.Int32(capC):
+                        smem_keys[cutlass.Int32(2 * segA) + slp0] = cutlass.Float32(ptb0)
+                ptb0 = ptb0 + cutlass.Int32(num_threads)
+            cute.arch.barrier()
+            nlist0 = s_seg[1]
+            if nlist0 <= cutlass.Int32(capC):
+                # pass 2: 8 listed blocks per warp round
+                nwp = cutlass.const_expr(self.num_warps)
+                lb0 = warp_id * cutlass.Int32(8)
+                while lb0 < nlist0:
+                    for _jb in cutlass.range_constexpr(8):
+                        li0 = lb0 + cutlass.Int32(_jb)
+                        if li0 < nlist0:
+                            bid0 = cutlass.Int32(smem_keys[cutlass.Int32(2 * segA) + li0])
+                            pos0 = (bid0 << cutlass.Int32(5)) + lane
+                            if pos0 < N:
+                                vp0 = cute.make_ptr(
+                                    cutlass.Float32,
+                                    row_addr + cutlass.Int64(pos0) * cutlass.Int64(4),
+                                    cute.AddressSpace.gmem,
+                                    assumed_align=4,
+                                )
+                                v0 = cute.make_tensor(vp0, cute.make_layout((1,)))[0]
+                                if v0 >= t2_s:
+                                    sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1))
+                                    if sl0 < cutlass.Int32(segA):
+                                        smem_keys[sl0] = v0
+                                        cand_idx_row[sl0] = pos0
+                    lb0 = lb0 + cutlass.Int32(nwp * 8)
+            if nlist0 > cutlass.Int32(capC):
+                # list overflow (pass rate too high for skip): dense full
+                # scan backup - nothing was read yet, plain re-run
+                itd0 = cutlass.Int32(0)
+                nfd0 = N >> cutlass.const_expr((num_threads * 4).bit_length() - 1)
+                while itd0 < nfd0:
+                    idd0 = (itd0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(p1w)
+                    pd0 = cute.make_ptr(
+                        self.dtype,
+                        row_addr + cutlass.Int64(idd0) * cutlass.Int64(4),
+                        cute.AddressSpace.gmem,
+                        assumed_align=vec_align,
+                    )
+                    cute.copy(
+                        pass1_atom,
+                        cute.make_tensor(pd0, cute.make_layout((p1w,))),
+                        frag_p,
+                    )
+                    for _jd in cutlass.range_constexpr(p1w):
+                        vd0 = cutlass.Float32(frag_p[_jd])
+                        if vd0 >= t2_s:
+                            sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1))
+                            if sl0 < cutlass.Int32(segA):
+                                smem_keys[sl0] = vd0
+                                cand_idx_row[sl0] = idd0 + cutlass.Int32(_jd)
+                    itd0 = itd0 + cutlass.Int32(1)
+                ptd0 = nfd0 * cutlass.Int32(num_threads * 4) + tidx
+                while ptd0 < N:
+                    pe0 = cute.make_ptr(
+                        cutlass.Float32,
+                        row_addr + cutlass.Int64(ptd0) * cutlass.Int64(4),
+                        cute.AddressSpace.gmem,
+                        assumed_align=4,
+                    )
+                    ve0 = cute.make_tensor(pe0, cute.make_layout((1,)))[0]
+                    if ve0 >= t2_s:
+                        sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1))
+                        if sl0 < cutlass.Int32(segA):
+                            smem_keys[sl0] = ve0
+                            cand_idx_row[sl0] = ptd0
+                    ptd0 = ptd0 + cutlass.Int32(num_threads)
         copy_atom = self._make_load_copy_atom()
         frag_a = cute.make_fragment((vec_w,), self.dtype)
         frag_b = cute.make_fragment((vec_w,), self.dtype)

From 6b117ca6fbab2a091d6a5f8efbd8ccc31c7073fa Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 11:02:24 -0700
Subject: [PATCH 025/117] [None][perf] GVR self_scan skip: split pass-2 into
 load-then-claim phases

The pass-2 gather interleaved each block load with its smem atomic
claims; atomics are memory-ordered, so the compiler could not overlap
the next block load and the eight-block round degenerated into a
serial latency chain (phase-0 stamp: 17.7us at flash-1024k against a
~5us budget). Loading all eight listed blocks into registers first
and claiming afterwards restores the in-flight parallelism:
phase 0 drops to 10.4us and the 25-cell table moves decisively -
flash 512k 1.39x over the PR16457 tip, 1024k parity (19us), 256k
0.93x; v32 64k parity, 128k+ 0.90-0.92x; pro 1M 0.85x. Dense/skip
best-of geomean 0.73 -> 0.84-0.86x, still 100/100 bit-exact.
Remaining gap concentrates in the mid-row dense regime (64-128k,
0.63-0.72x), where the dense scan's single-CTA latency wall stands
(cp.async staging is the known next lever).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 72 +++++++++++++++++--
 1 file changed, 66 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 84616b1e573c..0153e0fb337a 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -1272,6 +1272,9 @@ def phase0_scan_bucket(
         vec_w = cutlass.const_expr(self.vec_bits // self.dtype.width)
         elem_bytes = cutlass.const_expr(self.dtype.width // 8)
         vec_align = cutlass.const_expr(self.vec_align_bytes)
+        p0ck0 = cutlass.Int64(0)
+        if cutlass.const_expr(_P4_SUB_DBG):
+            p0ck0 = cute.arch.clock64()
         if tidx == cutlass.Int32(0):
             s_seg[0] = cutlass.Int32(0)
             s_seg[1] = cutlass.Int32(0)
@@ -1358,9 +1361,23 @@ def phase0_scan_bucket(
                 # pass 2: 8 listed blocks per warp round
                 nwp = cutlass.const_expr(self.num_warps)
                 lb0 = warp_id * cutlass.Int32(8)
+                frag_v = cute.make_fragment((8,), cutlass.Float32)
                 while lb0 < nlist0:
+                    # LOAD phase first: eight independent block loads in
+                    # flight before any atomic (claims are memory-ordered
+                    # and would serialize the blocks otherwise)
+                    p2b0 = cutlass.Int32(0)
+                    p2b1 = cutlass.Int32(0)
+                    p2b2 = cutlass.Int32(0)
+                    p2b3 = cutlass.Int32(0)
+                    p2b4 = cutlass.Int32(0)
+                    p2b5 = cutlass.Int32(0)
+                    p2b6 = cutlass.Int32(0)
+                    p2b7 = cutlass.Int32(0)
                     for _jb in cutlass.range_constexpr(8):
                         li0 = lb0 + cutlass.Int32(_jb)
+                        bid0 = cutlass.Int32(-1)
+                        vv0 = cutlass.Float32(self.NEG_FLT_MAX)
                         if li0 < nlist0:
                             bid0 = cutlass.Int32(smem_keys[cutlass.Int32(2 * segA) + li0])
                             pos0 = (bid0 << cutlass.Int32(5)) + lane
@@ -1371,12 +1388,51 @@ def phase0_scan_bucket(
                                     cute.AddressSpace.gmem,
                                     assumed_align=4,
                                 )
-                                v0 = cute.make_tensor(vp0, cute.make_layout((1,)))[0]
-                                if v0 >= t2_s:
-                                    sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1))
-                                    if sl0 < cutlass.Int32(segA):
-                                        smem_keys[sl0] = v0
-                                        cand_idx_row[sl0] = pos0
+                                vv0 = cute.make_tensor(vp0, cute.make_layout((1,)))[0]
+                        frag_v[_jb] = vv0
+                        if cutlass.const_expr(_jb == 0):
+                            p2b0 = bid0
+                        elif cutlass.const_expr(_jb == 1):
+                            p2b1 = bid0
+                        elif cutlass.const_expr(_jb == 2):
+                            p2b2 = bid0
+                        elif cutlass.const_expr(_jb == 3):
+                            p2b3 = bid0
+                        elif cutlass.const_expr(_jb == 4):
+                            p2b4 = bid0
+                        elif cutlass.const_expr(_jb == 5):
+                            p2b5 = bid0
+                        elif cutlass.const_expr(_jb == 6):
+                            p2b6 = bid0
+                        else:
+                            p2b7 = bid0
+                    # CLAIM phase
+                    for _jb in cutlass.range_constexpr(8):
+                        bidc = (
+                            p2b0
+                            if _jb == 0
+                            else p2b1
+                            if _jb == 1
+                            else p2b2
+                            if _jb == 2
+                            else p2b3
+                            if _jb == 3
+                            else p2b4
+                            if _jb == 4
+                            else p2b5
+                            if _jb == 5
+                            else p2b6
+                            if _jb == 6
+                            else p2b7
+                        )
+                        vvc = cutlass.Float32(frag_v[_jb])
+                        if bidc >= cutlass.Int32(0) and vvc >= t2_s:
+                            posc = (bidc << cutlass.Int32(5)) + lane
+                            if posc < N:
+                                sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1))
+                                if sl0 < cutlass.Int32(segA):
+                                    smem_keys[sl0] = vvc
+                                    cand_idx_row[sl0] = posc
                     lb0 = lb0 + cutlass.Int32(nwp * 8)
             if nlist0 > cutlass.Int32(capC):
                 # list overflow (pass rate too high for skip): dense full
@@ -1526,6 +1582,8 @@ def phase0_scan_bucket(
                 s_seg[4] = curT0
                 s_seg[5] = curT0
                 s_seg[6] = curT0
+            if cutlass.const_expr(_P4_SUB_DBG):
+                s_seg[7] = cutlass.Int32(cute.arch.clock64() - p0ck0)
             if cutlass.const_expr(not self.enable_block_skip):
                 curA0 = s_seg[0]
                 curB0 = s_seg[1]
@@ -6352,6 +6410,8 @@ def _run_phases(
                             xstate_row[5] = cutlass.Float32(cutlass.Int32(ck2 - ck1))  # P2/P3 gap
                             xstate_row[6] = cutlass.Float32(cutlass.Int32(ck3 - ck2))  # Phase 4
                             xstate_row[7] = s_thr[1]  # cnt_strad
+                        if cutlass.const_expr(_P4_SUB_DBG):
+                            xstate_row[2] = cutlass.Float32(smem_wcnt_p1[7])
                         if cutlass.const_expr(_P4_SUB_DBG):
                             # P4 sub-phase cycles staged by rank_scatter.
                             # Chain-safe layout: [2] (closed-loop anchor)

From aa6c27585ffdb41c7858e48a870e12d595a22874 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 11:10:59 -0700
Subject: [PATCH 026/117] [None][perf] GVR self_scan: software-pipeline the
 dense scan rounds

Preload the next round's two vectors into shadow fragments before the
current round's atomic claims (the pass-2 lesson applied to the dense
loop). Measured neutral-to-slightly-positive (phase-0 38.7 -> 37.3us
at flash-1024k): unlike pass 2 the dense loop's wall is not the
cross-round atomic ordering - documented for the record; the next
dense-lane lever is cp.async/smem staging.

Final 25-cell state (dense/skip best-of vs PR16457 tip, B=1..8
geomean 0.84-0.86x, 100/100 bit-exact): flash 512k 1.39x / 1024k
1.00x / 256k 0.93x; v32 64k 1.00x / 128k+ 0.90x; pro 1M 0.85x;
remaining gap concentrated at the 64-128k dense regime.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 50 +++++++++++++++----
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 0153e0fb337a..8cb3183c8bd5 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -1489,11 +1489,14 @@ def phase0_scan_bucket(
         nfull = N >> st2log
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             nfull = cutlass.Int32(0)
-        it0 = cutlass.Int32(0)
-        while it0 < nfull:
-            ia0 = (it0 * cutlass.Int32(2) * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(
-                vec_w
-            )
+        # software pipeline: the claims are memory-ordered atomics, so
+        # without a preload the next round's vector loads serialize
+        # behind them (the same wall pass 2 hit). Preload round i+1 into
+        # the shadow fragments BEFORE claiming round i, then ping-pong.
+        frag_c = cute.make_fragment((vec_w,), self.dtype)
+        frag_d = cute.make_fragment((vec_w,), self.dtype)
+        if nfull > cutlass.Int32(0):
+            ia0 = tidx * cutlass.Int32(vec_w)
             ib0 = ia0 + cutlass.Int32(step1)
             for _fq in cutlass.range_constexpr(2):
                 fq0 = ia0 if _fq == 0 else ib0
@@ -1508,13 +1511,34 @@ def phase0_scan_bucket(
                     cute.make_tensor(pq0, cute.make_layout((vec_w,))),
                     frag_a if _fq == 0 else frag_b,
                 )
+        it0 = cutlass.Int32(0)
+        while it0 < nfull:
+            ia0 = (it0 * cutlass.Int32(2) * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(
+                vec_w
+            )
+            ib0 = ia0 + cutlass.Int32(step1)
+            nxt0 = it0 + cutlass.Int32(1)
+            if nxt0 < nfull:
+                ja0 = (nxt0 * cutlass.Int32(2) * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(
+                    vec_w
+                )
+                jb0 = ja0 + cutlass.Int32(step1)
+                for _fq in cutlass.range_constexpr(2):
+                    fq0 = ja0 if _fq == 0 else jb0
+                    pq0 = cute.make_ptr(
+                        self.dtype,
+                        row_addr + cutlass.Int64(fq0) * cutlass.Int64(elem_bytes),
+                        cute.AddressSpace.gmem,
+                        assumed_align=vec_align,
+                    )
+                    cute.copy(
+                        copy_atom,
+                        cute.make_tensor(pq0, cute.make_layout((vec_w,))),
+                        frag_c if _fq == 0 else frag_d,
+                    )
             # v6: per-element DIRECT atomic claims for passers — smem
-            # atomics do NOT synchronize the warp, so the vector loads of
-            # later rounds keep flowing (the packed shfl-prefix design
-            # capped in-flight loads at 2/warp: ncu showed 0.19% memory
-            # throughput, pure latency bound). Same-address service is
-            # ~0.5ns effective and n0 is line-bounded, so the claim queue
-            # hides completely under the read stream.
+            # atomics do NOT synchronize the warp; with the preload above
+            # they no longer stall the next round's loads either.
             for _jh in cutlass.range_constexpr(2):
                 for _jv in cutlass.range_constexpr(vec_w):
                     v0 = cutlass.Float32(frag_a[_jv]) if _jh == 0 else cutlass.Float32(frag_b[_jv])
@@ -1537,6 +1561,10 @@ def phase0_scan_bucket(
                                 c0 = cutlass.Int32(-1)
                             else:
                                 c0 = c0 + cutlass.Int32(1)
+            if nxt0 < nfull:
+                for _jv in cutlass.range_constexpr(vec_w):
+                    frag_a[_jv] = frag_c[_jv]
+                    frag_b[_jv] = frag_d[_jv]
             it0 = it0 + cutlass.Int32(1)
         # scalar tail (< 2*step1 elements): per-element DIRECT atomic
         # claims — divergent-safe, no warp collectives

From 26df9e6e230b8408a94936428cfc667097e1a5cc Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 12:14:19 -0700
Subject: [PATCH 027/117] [None][perf] GVR self_scan: cp.async deep-pipeline
 the dense scan

Replace the register-preload dense scan with an LDGSTS staging
pipeline: each thread streams one 16B vector per step into a private
slot-major smem slot (no data registers, no scoreboard stall until the
wait), keeping stage_slots rounds in flight. The staging buffer
aliases smem_vals - written only after phase 0, with every non-empty
cp.async group drained inside the loop - so depth 2 costs zero smem;
trimming cap_c to <= 16384 frees 32KB of keys for depth 4.

flash-1024k phase0 37.3 -> 34.8us; exactness unchanged (fused and
skip smoke 6/6, 25-cell real-data sweep 50/50 bit-exact).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 233 +++++++++++-------
 1 file changed, 139 insertions(+), 94 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 8cb3183c8bd5..1310f28c7c39 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -621,6 +621,14 @@ def __init__(
             # 227KB CTA limit with room for the stage-2 skip list.
             self.seg_total = 2 * self.accept_cap + int(cap_c if cap_c is not None else 24576)
             self.cap_c = self.seg_total - 2 * self.accept_cap
+            # cp.async staging depth for the phase-0 dense scan: 4 rounds
+            # in flight when the C segment is trimmed enough to pay for
+            # the extra 32KB, else the zero-cost 2 (staging then aliases
+            # smem_vals exactly). Rows are 16B slots, one per thread per
+            # in-flight round; never smaller than vals so the alias always
+            # covers it.
+            self.stage_slots = 4 if self.cap_c <= 16384 else 2
+            self.stage_rows = max(self.stage_slots * self.num_threads, self.kC // 4)
         else:
             self.seg_total = self.kC
             self.cap_c = 0
@@ -1243,6 +1251,7 @@ def phase0_scan_bucket(
         smem_keys,  # [seg_total] fp32 value segments (A @0 / B @segA / C @2segA)
         cand_idx_row,  # [seg_total] int32 gmem POSITION column (write-only here)
         block_max_row,  # [nb_pad] fp32 per-32-position maxima, or None
+        smem_stage,  # [stage_rows, 4] fp32 cp.async staging (aliases smem_vals)
         s_seg,  # [>=7] int32 scratch (reuses smem_wcnt_p1: P1 never runs
         #         on a row this phase succeeded on): [0..2] A/B/C claim
         #         cursors, [3] void, [4] n0, [5] n1, [6] n2
@@ -1258,18 +1267,19 @@ def phase0_scan_bucket(
         uncapped), so {n0, void, n1, n2} fall out for free — the same
         contract the v5 emitter produced externally.
 
-        Perf shape (v3): per round each thread front-loads TWO vec_w
-        vectors (independent LDGs, latency overlapped), classifies
-        branchlessly in registers, and the whole round pays ONE packed
-        warp shfl-prefix (all three segment counts in 10-bit fields) +
-        at most three warp atomics. Segment overflow is marked and
-        resolved by a RARE per-element direct-atomic pass (a segment
-        overflows at most once per row, and divergent scalar atomics
-        need no warp coordination)."""
+        Perf shape (v13): the dense scan is a cp.async pipeline — each
+        thread streams one 16B vector per step into its private
+        slot-major smem staging slot (LDGSTS: no data registers, no
+        scoreboard stall until the wait), keeping ``stage_slots`` steps
+        in flight; classification reads the staged values and claims
+        passers with per-element direct smem atomics (they don't
+        synchronize the warp and hide under the async copy stream).
+        Segment overflow is resolved in-claim by spilling to the next
+        looser segment (a segment overflows at most once per row, and
+        divergent scalar atomics need no warp coordination)."""
         num_threads = cutlass.const_expr(self.num_threads)
         segA = cutlass.const_expr(self.accept_cap)
         capC = cutlass.const_expr(self.cap_c)
-        vec_w = cutlass.const_expr(self.vec_bits // self.dtype.width)
         elem_bytes = cutlass.const_expr(self.dtype.width // 8)
         vec_align = cutlass.const_expr(self.vec_align_bytes)
         p0ck0 = cutlass.Int64(0)
@@ -1475,100 +1485,110 @@ def phase0_scan_bucket(
                             smem_keys[sl0] = ve0
                             cand_idx_row[sl0] = ptd0
                     ptd0 = ptd0 + cutlass.Int32(num_threads)
-        copy_atom = self._make_load_copy_atom()
-        frag_a = cute.make_fragment((vec_w,), self.dtype)
-        frag_b = cute.make_fragment((vec_w,), self.dtype)
-        step1 = cutlass.const_expr(num_threads * vec_w)
-        step2 = cutlass.const_expr(2 * step1)
-        st2log = cutlass.const_expr((2 * step1).bit_length() - 1)
-        # FULL-vector rounds only in the hot loop: nfull is warp-uniform
-        # (every lane's two vectors are in bounds by construction), so the
-        # warp collectives inside are legal and the hot path carries no
-        # bounds checks at all. The remainder (< 2*step1 elements) takes
-        # the scalar tail below with per-element direct atomics.
-        nfull = N >> st2log
+        cpw = cutlass.const_expr(4)  # cp.async caps at 16B per copy
+        step1 = cutlass.const_expr(num_threads * cpw)
+        st1log = cutlass.const_expr(step1.bit_length() - 1)
+        n_stage = cutlass.const_expr(self.stage_slots)
+        n_smask = cutlass.const_expr(self.stage_slots - 1)
+        # v13 cp.async deep pipeline: LDGSTS stages each round's vector
+        # straight into the slot-major smem staging buffer — no data
+        # registers consumed, no scoreboard stall until the WAIT — so
+        # n_stage rounds stay in flight per thread (the register pipeline
+        # topped out at two). The staging buffer aliases smem_vals (only
+        # written after phase 0); every non-empty group is drained inside
+        # the loop (one commit per step, wait_group(n_stage-1) pops the
+        # oldest), so nothing is ever in flight once the alias is read.
+        # FULL-vector steps only in the hot loop; the remainder takes the
+        # scalar tail below.
+        nfull = N >> st1log
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             nfull = cutlass.Int32(0)
-        # software pipeline: the claims are memory-ordered atomics, so
-        # without a preload the next round's vector loads serialize
-        # behind them (the same wall pass 2 hit). Preload round i+1 into
-        # the shadow fragments BEFORE claiming round i, then ping-pong.
-        frag_c = cute.make_fragment((vec_w,), self.dtype)
-        frag_d = cute.make_fragment((vec_w,), self.dtype)
-        if nfull > cutlass.Int32(0):
-            ia0 = tidx * cutlass.Int32(vec_w)
-            ib0 = ia0 + cutlass.Int32(step1)
-            for _fq in cutlass.range_constexpr(2):
-                fq0 = ia0 if _fq == 0 else ib0
+        g2s_atom = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(cute.nvgpu.cpasync.LoadCacheMode.GLOBAL),
+            cutlass.Float32,
+            num_bits_per_copy=128,
+        )
+        stage_addr = smem_stage.iterator.toint()
+        for _p in cutlass.range_constexpr(n_stage):
+            if cutlass.Int32(_p) < nfull:
+                pr0 = cutlass.Int32(_p) * cutlass.Int32(num_threads) + tidx
+                pp0 = pr0 * cutlass.Int32(cpw)
                 pq0 = cute.make_ptr(
                     self.dtype,
-                    row_addr + cutlass.Int64(fq0) * cutlass.Int64(elem_bytes),
+                    row_addr + cutlass.Int64(pp0) * cutlass.Int64(elem_bytes),
                     cute.AddressSpace.gmem,
-                    assumed_align=vec_align,
+                    assumed_align=16,
+                )
+                dq0 = cute.make_ptr(
+                    cutlass.Float32,
+                    stage_addr + cutlass.Int64(pr0 * cutlass.Int32(cpw)) * cutlass.Int64(4),
+                    cute.AddressSpace.smem,
+                    assumed_align=16,
                 )
                 cute.copy(
-                    copy_atom,
-                    cute.make_tensor(pq0, cute.make_layout((vec_w,))),
-                    frag_a if _fq == 0 else frag_b,
+                    g2s_atom,
+                    cute.make_tensor(pq0, cute.make_layout((cpw,))),
+                    cute.make_tensor(dq0, cute.make_layout((cpw,))),
                 )
+            cute.arch.cp_async_commit_group()
         it0 = cutlass.Int32(0)
         while it0 < nfull:
-            ia0 = (it0 * cutlass.Int32(2) * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(
-                vec_w
+            cute.arch.cp_async_wait_group(n_smask)
+            sb0 = (it0 & cutlass.Int32(n_smask)) * cutlass.Int32(num_threads) + tidx
+            sq0 = cute.make_ptr(
+                cutlass.Float32,
+                stage_addr + cutlass.Int64(sb0 * cutlass.Int32(cpw)) * cutlass.Int64(4),
+                cute.AddressSpace.smem,
+                assumed_align=16,
             )
-            ib0 = ia0 + cutlass.Int32(step1)
-            nxt0 = it0 + cutlass.Int32(1)
-            if nxt0 < nfull:
-                ja0 = (nxt0 * cutlass.Int32(2) * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(
-                    vec_w
-                )
-                jb0 = ja0 + cutlass.Int32(step1)
-                for _fq in cutlass.range_constexpr(2):
-                    fq0 = ja0 if _fq == 0 else jb0
-                    pq0 = cute.make_ptr(
-                        self.dtype,
-                        row_addr + cutlass.Int64(fq0) * cutlass.Int64(elem_bytes),
-                        cute.AddressSpace.gmem,
-                        assumed_align=vec_align,
-                    )
-                    cute.copy(
-                        copy_atom,
-                        cute.make_tensor(pq0, cute.make_layout((vec_w,))),
-                        frag_c if _fq == 0 else frag_d,
-                    )
+            srow = cute.make_tensor(sq0, cute.make_layout((cpw,)))
+            ia0 = (it0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(cpw)
             # v6: per-element DIRECT atomic claims for passers — smem
-            # atomics do NOT synchronize the warp; with the preload above
-            # they no longer stall the next round's loads either.
-            for _jh in cutlass.range_constexpr(2):
-                for _jv in cutlass.range_constexpr(vec_w):
-                    v0 = cutlass.Float32(frag_a[_jv]) if _jh == 0 else cutlass.Float32(frag_b[_jv])
-                    if v0 >= t0_s:
-                        pos0 = (ia0 if _jh == 0 else ib0) + cutlass.Int32(_jv)
-                        c0 = cutlass.Int32(2)
-                        if v0 >= t1_s:
-                            c0 = cutlass.Int32(1)
-                        if v0 >= t2_s:
-                            c0 = cutlass.Int32(0)
-                        while c0 >= cutlass.Int32(0) and c0 <= cutlass.Int32(2):
-                            cap0 = cutlass.Int32(segA)
-                            if c0 == cutlass.Int32(2):
-                                cap0 = cutlass.Int32(capC)
-                            sl0 = atomicAdd(s_seg.iterator + c0, cutlass.Int32(1))
-                            if sl0 < cap0:
-                                cd0 = c0 * cutlass.Int32(segA) + sl0
-                                smem_keys[cd0] = v0
-                                cand_idx_row[cd0] = pos0
-                                c0 = cutlass.Int32(-1)
-                            else:
-                                c0 = c0 + cutlass.Int32(1)
-            if nxt0 < nfull:
-                for _jv in cutlass.range_constexpr(vec_w):
-                    frag_a[_jv] = frag_c[_jv]
-                    frag_b[_jv] = frag_d[_jv]
+            # atomics do NOT synchronize the warp and hide under the
+            # async copy stream.
+            for _jv in cutlass.range_constexpr(cpw):
+                v0 = cutlass.Float32(srow[_jv])
+                if v0 >= t0_s:
+                    pos0 = ia0 + cutlass.Int32(_jv)
+                    c0 = cutlass.Int32(2)
+                    if v0 >= t1_s:
+                        c0 = cutlass.Int32(1)
+                    if v0 >= t2_s:
+                        c0 = cutlass.Int32(0)
+                    while c0 >= cutlass.Int32(0) and c0 <= cutlass.Int32(2):
+                        cap0 = cutlass.Int32(segA)
+                        if c0 == cutlass.Int32(2):
+                            cap0 = cutlass.Int32(capC)
+                        sl0 = atomicAdd(s_seg.iterator + c0, cutlass.Int32(1))
+                        if sl0 < cap0:
+                            cd0 = c0 * cutlass.Int32(segA) + sl0
+                            smem_keys[cd0] = v0
+                            cand_idx_row[cd0] = pos0
+                            c0 = cutlass.Int32(-1)
+                        else:
+                            c0 = c0 + cutlass.Int32(1)
+            # reissue the just-consumed slot for step it0 + n_stage (the
+            # thread's own prior reads are ordered before the async write
+            # begins, so no fence is needed)
+            kn0 = it0 + cutlass.Int32(n_stage)
+            if kn0 < nfull:
+                jp0 = (kn0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(cpw)
+                pq0 = cute.make_ptr(
+                    self.dtype,
+                    row_addr + cutlass.Int64(jp0) * cutlass.Int64(elem_bytes),
+                    cute.AddressSpace.gmem,
+                    assumed_align=16,
+                )
+                cute.copy(
+                    g2s_atom,
+                    cute.make_tensor(pq0, cute.make_layout((cpw,))),
+                    srow,
+                )
+            cute.arch.cp_async_commit_group()
             it0 = it0 + cutlass.Int32(1)
-        # scalar tail (< 2*step1 elements): per-element DIRECT atomic
+        # scalar tail (< step1 elements): per-element DIRECT atomic
         # claims — divergent-safe, no warp collectives
-        pt0 = (N >> st2log) * cutlass.Int32(step2) + tidx
+        pt0 = (N >> st1log) * cutlass.Int32(step1) + tidx
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             pt0 = N
         while pt0 < N:
@@ -4857,11 +4877,31 @@ def run_one_row(
         # SEGMENT COORDINATE of each compacted candidate (identity for the
         # deferred position gather via cand_idx[coord]) — every consumer
         # (P4, tail repair, gather) works unchanged.
-        smem_vals = smem.allocate_tensor(
-            element_type=cutlass.Int32,
-            layout=cute.make_ordered_layout((kC,), order=(0,)),
-            byte_alignment=128,
-        )
+        if cutlass.const_expr(self.self_scan):
+            # phase-0 cp.async staging, ALIASED over vals: vals is only
+            # written after phase 0 completes and every in-flight group
+            # is drained inside the dense loop, so the lifetimes never
+            # overlap. Slot-major (slot s of thread t at row
+            # s*num_threads + t): a warp's 16B reads/writes land on
+            # consecutive banks, conflict-free.
+            smem_stage = smem.allocate_tensor(
+                element_type=cutlass.Float32,
+                layout=cute.make_ordered_layout(
+                    (cutlass.const_expr(self.stage_rows), 4), order=(1, 0)
+                ),
+                byte_alignment=128,
+            )
+            smem_vals = cute.make_tensor(
+                cute.recast_ptr(smem_stage.iterator, dtype=cutlass.Int32),
+                cute.make_ordered_layout((kC,), order=(0,)),
+            )
+        else:
+            smem_stage = None
+            smem_vals = smem.allocate_tensor(
+                element_type=cutlass.Int32,
+                layout=cute.make_ordered_layout((kC,), order=(0,)),
+                byte_alignment=128,
+            )
         # histogram[kNumBins] int32 (P4 only)
         smem_hist = smem.allocate_tensor(
             element_type=cutlass.Int32,
@@ -5139,6 +5179,7 @@ def run_one_row(
                         cand_ctl_row=cand_ctl_row,
                         smem_active=smem_active,
                         s_active_cnt=s_active_cnt,
+                        smem_stage=smem_stage,
                     )
                 else:
                     # Short row: only CTA 0 scans the full row; the other
@@ -5189,6 +5230,7 @@ def run_one_row(
                             cand_ctl_row=cand_ctl_row,
                             smem_active=smem_active,
                             s_active_cnt=s_active_cnt,
+                            smem_stage=smem_stage,
                         )
             else:
                 # cs=1: one CTA per row, no cluster sync.
@@ -5236,6 +5278,7 @@ def run_one_row(
                     cand_ctl_row=cand_ctl_row,
                     smem_active=smem_active,
                     s_active_cnt=s_active_cnt,
+                    smem_stage=smem_stage,
                 )
 
         griddepcontrol_launch_dependents()
@@ -5286,6 +5329,7 @@ def _run_phases(
         cand_ctl_row=None,  # ext cand: this row's [2] int32 {claimed, void}
         smem_active=None,
         s_active_cnt=None,
+        smem_stage=None,  # self_scan: [stage_rows, 4] fp32 cp.async staging
     ):
         """Run Phase 1-4 + final cluster barrier on a given row slice.
 
@@ -5364,6 +5408,7 @@ def _run_phases(
                         smem_keys,
                         cand_idx_row,
                         block_max_row,
+                        smem_stage,
                         smem_wcnt_p1,
                         tidx,
                         warp_id,

From 2253600b6d8c50a220a5406b75d73ed499d64377 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 12:54:48 -0700
Subject: [PATCH 028/117] [None][perf] GVR self_scan: pin the CTA to 1024
 threads in the wrapper

The short-row 512-thread heuristic is tuned for the stock multi-pass
kernel; under self_scan it silently halved the warp count of every
N_dec < 65536 cell and cost ~5us/cell in the phase-0 scan (flash-128k
p0 14.4 -> 9.4us at 1024 threads). Route self_scan to 1024 threads
unconditionally.

25-cell x B{1,2,4,8} same-node sweep vs PR16457 tip: best-of geomean
0.838 -> 0.864 (B8 0.881), 100/100 bit-exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/top_k/run_gvr_topk.py        | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 0b4a3e186402..cdcc577cb012 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -837,11 +837,18 @@ def gvr_topk_decode(
     N_cols = logits.shape[1]
     N_dec = max_seq_len if max_seq_len is not None else N_cols
     if num_threads_per_block is None:
-        if max_seq_len is not None and logits.dtype != torch.float32:
-            n_thresh_t = 131072
+        if self_scan:
+            # self_scan owns the whole row scan in one CTA: the phase-0
+            # cp.async pipeline scales with warp count at every N (the
+            # 512-thread short-row heuristic below is tuned for the
+            # stock multi-pass kernel and costs ~5us/cell here).
+            num_threads_per_block = 1024
         else:
-            n_thresh_t = 65536
-        num_threads_per_block = 1024 if (num_rows <= num_sms and N_dec >= n_thresh_t) else 512
+            if max_seq_len is not None and logits.dtype != torch.float32:
+                n_thresh_t = 131072
+            else:
+                n_thresh_t = 65536
+            num_threads_per_block = 1024 if (num_rows <= num_sms and N_dec >= n_thresh_t) else 512
     if use_256bit_load is None:
         use_256bit_load = logits.dtype == torch.float32 and N_dec >= 16384
     if enable_warp_parallel_reduce is None:

From 8ab4f74c87e9769a9b13bae25c1f930ad776699e Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 13:07:15 -0700
Subject: [PATCH 029/117] [None][perf] GVR self_scan: pair-step the cp.async
 scan pipeline

The dense scan is instruction-issue bound (pcsamp: no_instructions +
fixed-latency wait dominate; long_scoreboard is 6%), so each pipeline
step now processes two 16B vectors per thread - loop, wait, commit and
address arithmetic amortize over 8 elements while the in-flight byte
count stays at 2 pairs x 32B across the 4 staging slots.

The pair shape needs all 4 slot rows, and the 64KB staging fits the
CTA budget only with the C segment trimmed, so self_scan now defaults
cap_c to 16384 and rejects anything larger (validated bit-exact across
the 25-cell x B{1,2,4,8} sweep).

flash p0 (warm, 1024 threads): 512k 17.8 -> 16.8us, 1024k 33.5 ->
32.4us; smoke 6/6 exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 197 ++++++++++--------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |   4 +-
 2 files changed, 112 insertions(+), 89 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 1310f28c7c39..42bf4507c436 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -619,15 +619,19 @@ def __init__(
             # on-chip segment budget: values only, 4B/entry; C sized so
             # keys(160KB) + vals(32KB) + hist + scratch stay under the
             # 227KB CTA limit with room for the stage-2 skip list.
-            self.seg_total = 2 * self.accept_cap + int(cap_c if cap_c is not None else 24576)
+            self.seg_total = 2 * self.accept_cap + int(cap_c if cap_c is not None else 16384)
             self.cap_c = self.seg_total - 2 * self.accept_cap
-            # cp.async staging depth for the phase-0 dense scan: 4 rounds
-            # in flight when the C segment is trimmed enough to pay for
-            # the extra 32KB, else the zero-cost 2 (staging then aliases
-            # smem_vals exactly). Rows are 16B slots, one per thread per
-            # in-flight round; never smaller than vals so the alias always
+            # cp.async staging for the phase-0 dense scan: the pair-step
+            # pipeline keeps 2 pairs x 2 slots in flight per thread, so
+            # exactly 4 slot rows are required. The 64KB staging fits the
+            # CTA budget only with the C segment trimmed to <= 16384
+            # (keys 128KB + staging 64KB + hist/scratch); larger C would
+            # silently overrun the alias, so reject it outright. Rows are
+            # 16B slots; never fewer than vals holds so the alias always
             # covers it.
-            self.stage_slots = 4 if self.cap_c <= 16384 else 2
+            if self.cap_c > 16384:
+                raise ValueError("self_scan requires cap_c <= 16384 (staging budget)")
+            self.stage_slots = 4
             self.stage_rows = max(self.stage_slots * self.num_threads, self.kC // 4)
         else:
             self.seg_total = self.kC
@@ -1487,20 +1491,19 @@ def phase0_scan_bucket(
                     ptd0 = ptd0 + cutlass.Int32(num_threads)
         cpw = cutlass.const_expr(4)  # cp.async caps at 16B per copy
         step1 = cutlass.const_expr(num_threads * cpw)
-        st1log = cutlass.const_expr(step1.bit_length() - 1)
-        n_stage = cutlass.const_expr(self.stage_slots)
-        n_smask = cutlass.const_expr(self.stage_slots - 1)
-        # v13 cp.async deep pipeline: LDGSTS stages each round's vector
-        # straight into the slot-major smem staging buffer — no data
-        # registers consumed, no scoreboard stall until the WAIT — so
-        # n_stage rounds stay in flight per thread (the register pipeline
-        # topped out at two). The staging buffer aliases smem_vals (only
-        # written after phase 0); every non-empty group is drained inside
-        # the loop (one commit per step, wait_group(n_stage-1) pops the
-        # oldest), so nothing is ever in flight once the alias is read.
-        # FULL-vector steps only in the hot loop; the remainder takes the
-        # scalar tail below.
-        nfull = N >> st1log
+        st2log = cutlass.const_expr((2 * step1).bit_length() - 1)
+        # v14 pair-step cp.async pipeline: the scan is instruction-issue
+        # bound (pcsamp: no_instructions + wait dominate; long_scoreboard
+        # is 6%), so each step processes TWO 16B vectors per thread —
+        # loop/wait/commit/address overhead amortizes over 8 elements
+        # instead of 4 while the in-flight byte count stays put (2 pairs
+        # x 32B across the 4 staging slots). One commit group per pair;
+        # wait_group(1) pops the oldest pair. The staging buffer aliases
+        # smem_vals (only written after phase 0); every non-empty group
+        # is drained inside the loop, so nothing is in flight once the
+        # alias is read. FULL-pair steps only in the hot loop; the
+        # remainder takes the scalar tail below.
+        nfull = N >> st2log
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             nfull = cutlass.Int32(0)
         g2s_atom = cute.make_copy_atom(
@@ -1509,86 +1512,106 @@ def phase0_scan_bucket(
             num_bits_per_copy=128,
         )
         stage_addr = smem_stage.iterator.toint()
-        for _p in cutlass.range_constexpr(n_stage):
+        for _p in cutlass.range_constexpr(2):
             if cutlass.Int32(_p) < nfull:
-                pr0 = cutlass.Int32(_p) * cutlass.Int32(num_threads) + tidx
-                pp0 = pr0 * cutlass.Int32(cpw)
-                pq0 = cute.make_ptr(
-                    self.dtype,
-                    row_addr + cutlass.Int64(pp0) * cutlass.Int64(elem_bytes),
-                    cute.AddressSpace.gmem,
-                    assumed_align=16,
-                )
-                dq0 = cute.make_ptr(
-                    cutlass.Float32,
-                    stage_addr + cutlass.Int64(pr0 * cutlass.Int32(cpw)) * cutlass.Int64(4),
-                    cute.AddressSpace.smem,
-                    assumed_align=16,
-                )
-                cute.copy(
-                    g2s_atom,
-                    cute.make_tensor(pq0, cute.make_layout((cpw,))),
-                    cute.make_tensor(dq0, cute.make_layout((cpw,))),
-                )
+                for _v in cutlass.range_constexpr(2):
+                    pr0 = cutlass.Int32(2 * _p + _v) * cutlass.Int32(num_threads) + tidx
+                    pp0 = cutlass.Int32(2 * _p + _v) * cutlass.Int32(step1) + tidx * cutlass.Int32(
+                        cpw
+                    )
+                    pq0 = cute.make_ptr(
+                        self.dtype,
+                        row_addr + cutlass.Int64(pp0) * cutlass.Int64(elem_bytes),
+                        cute.AddressSpace.gmem,
+                        assumed_align=16,
+                    )
+                    dq0 = cute.make_ptr(
+                        cutlass.Float32,
+                        stage_addr + cutlass.Int64(pr0 * cutlass.Int32(cpw)) * cutlass.Int64(4),
+                        cute.AddressSpace.smem,
+                        assumed_align=16,
+                    )
+                    cute.copy(
+                        g2s_atom,
+                        cute.make_tensor(pq0, cute.make_layout((cpw,))),
+                        cute.make_tensor(dq0, cute.make_layout((cpw,))),
+                    )
             cute.arch.cp_async_commit_group()
         it0 = cutlass.Int32(0)
         while it0 < nfull:
-            cute.arch.cp_async_wait_group(n_smask)
-            sb0 = (it0 & cutlass.Int32(n_smask)) * cutlass.Int32(num_threads) + tidx
-            sq0 = cute.make_ptr(
-                cutlass.Float32,
-                stage_addr + cutlass.Int64(sb0 * cutlass.Int32(cpw)) * cutlass.Int64(4),
-                cute.AddressSpace.smem,
-                assumed_align=16,
-            )
-            srow = cute.make_tensor(sq0, cute.make_layout((cpw,)))
-            ia0 = (it0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(cpw)
+            cute.arch.cp_async_wait_group(1)
+            sp0 = (it0 & cutlass.Int32(1)) * cutlass.Int32(2 * num_threads) + tidx
+            ia0 = it0 * cutlass.Int32(2 * step1) + tidx * cutlass.Int32(cpw)
             # v6: per-element DIRECT atomic claims for passers — smem
             # atomics do NOT synchronize the warp and hide under the
             # async copy stream.
-            for _jv in cutlass.range_constexpr(cpw):
-                v0 = cutlass.Float32(srow[_jv])
-                if v0 >= t0_s:
-                    pos0 = ia0 + cutlass.Int32(_jv)
-                    c0 = cutlass.Int32(2)
-                    if v0 >= t1_s:
-                        c0 = cutlass.Int32(1)
-                    if v0 >= t2_s:
-                        c0 = cutlass.Int32(0)
-                    while c0 >= cutlass.Int32(0) and c0 <= cutlass.Int32(2):
-                        cap0 = cutlass.Int32(segA)
-                        if c0 == cutlass.Int32(2):
-                            cap0 = cutlass.Int32(capC)
-                        sl0 = atomicAdd(s_seg.iterator + c0, cutlass.Int32(1))
-                        if sl0 < cap0:
-                            cd0 = c0 * cutlass.Int32(segA) + sl0
-                            smem_keys[cd0] = v0
-                            cand_idx_row[cd0] = pos0
-                            c0 = cutlass.Int32(-1)
-                        else:
-                            c0 = c0 + cutlass.Int32(1)
-            # reissue the just-consumed slot for step it0 + n_stage (the
+            for _jh in cutlass.range_constexpr(2):
+                sq0 = cute.make_ptr(
+                    cutlass.Float32,
+                    stage_addr
+                    + cutlass.Int64(
+                        (sp0 + cutlass.Int32(_jh) * cutlass.Int32(num_threads)) * cutlass.Int32(cpw)
+                    )
+                    * cutlass.Int64(4),
+                    cute.AddressSpace.smem,
+                    assumed_align=16,
+                )
+                srow = cute.make_tensor(sq0, cute.make_layout((cpw,)))
+                for _jv in cutlass.range_constexpr(cpw):
+                    v0 = cutlass.Float32(srow[_jv])
+                    if v0 >= t0_s:
+                        pos0 = ia0 + cutlass.Int32(_jh) * cutlass.Int32(step1) + cutlass.Int32(_jv)
+                        c0 = cutlass.Int32(2)
+                        if v0 >= t1_s:
+                            c0 = cutlass.Int32(1)
+                        if v0 >= t2_s:
+                            c0 = cutlass.Int32(0)
+                        while c0 >= cutlass.Int32(0) and c0 <= cutlass.Int32(2):
+                            cap0 = cutlass.Int32(segA)
+                            if c0 == cutlass.Int32(2):
+                                cap0 = cutlass.Int32(capC)
+                            sl0 = atomicAdd(s_seg.iterator + c0, cutlass.Int32(1))
+                            if sl0 < cap0:
+                                cd0 = c0 * cutlass.Int32(segA) + sl0
+                                smem_keys[cd0] = v0
+                                cand_idx_row[cd0] = pos0
+                                c0 = cutlass.Int32(-1)
+                            else:
+                                c0 = c0 + cutlass.Int32(1)
+            # reissue the just-consumed slot pair for step it0 + 2 (the
             # thread's own prior reads are ordered before the async write
             # begins, so no fence is needed)
-            kn0 = it0 + cutlass.Int32(n_stage)
+            kn0 = it0 + cutlass.Int32(2)
             if kn0 < nfull:
-                jp0 = (kn0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(cpw)
-                pq0 = cute.make_ptr(
-                    self.dtype,
-                    row_addr + cutlass.Int64(jp0) * cutlass.Int64(elem_bytes),
-                    cute.AddressSpace.gmem,
-                    assumed_align=16,
-                )
-                cute.copy(
-                    g2s_atom,
-                    cute.make_tensor(pq0, cute.make_layout((cpw,))),
-                    srow,
-                )
+                for _jh in cutlass.range_constexpr(2):
+                    jr0 = sp0 + cutlass.Int32(_jh) * cutlass.Int32(num_threads)
+                    jp0 = (
+                        kn0 * cutlass.Int32(2 * step1)
+                        + cutlass.Int32(_jh) * cutlass.Int32(step1)
+                        + tidx * cutlass.Int32(cpw)
+                    )
+                    pq0 = cute.make_ptr(
+                        self.dtype,
+                        row_addr + cutlass.Int64(jp0) * cutlass.Int64(elem_bytes),
+                        cute.AddressSpace.gmem,
+                        assumed_align=16,
+                    )
+                    dq0 = cute.make_ptr(
+                        cutlass.Float32,
+                        stage_addr + cutlass.Int64(jr0 * cutlass.Int32(cpw)) * cutlass.Int64(4),
+                        cute.AddressSpace.smem,
+                        assumed_align=16,
+                    )
+                    cute.copy(
+                        g2s_atom,
+                        cute.make_tensor(pq0, cute.make_layout((cpw,))),
+                        cute.make_tensor(dq0, cute.make_layout((cpw,))),
+                    )
             cute.arch.cp_async_commit_group()
             it0 = it0 + cutlass.Int32(1)
         # scalar tail (< step1 elements): per-element DIRECT atomic
         # claims — divergent-safe, no warp collectives
-        pt0 = (N >> st1log) * cutlass.Int32(step1) + tidx
+        pt0 = (N >> st2log) * cutlass.Int32(2 * step1) + tidx
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             pt0 = N
         while pt0 < N:
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index cdcc577cb012..fc72170a1447 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -753,7 +753,7 @@ def gvr_topk_decode(
         if seed_counts is None:
             seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=logits.device)
         _bstar = int(os.environ["GVR_BSTAR"])
-        _capc = cap_c if cap_c is not None else int(os.environ.get("GVR_CAPC", "24576"))
+        _capc = cap_c if cap_c is not None else int(os.environ.get("GVR_CAPC", "16384"))
         _segtot = 2 * _bstar + _capc
         if cand_idx is None:
             cand_idx = torch.empty((num_rows, _segtot), dtype=torch.int32, device=logits.device)
@@ -904,7 +904,7 @@ def gvr_topk_decode(
         self_scan,
         cap_c
         if cap_c is not None
-        else (int(os.environ.get("GVR_CAPC", "24576")) if self_scan else None),
+        else (int(os.environ.get("GVR_CAPC", "16384")) if self_scan else None),
     )
     # When return_output_values=False the kernel was compiled to skip
     # STG.value and accepts None for the value-output slot.

From 5970b3e37197d4cb8d1876f4ec58745be51c7d01 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 25 Jul 2026 19:41:16 -0700
Subject: [PATCH 030/117] [None][feat] GVR ext_rungs: two-pass variant B (count
 external rungs in-kernel)

New ext_rungs mode: the host supplies only the three closed-loop rung
THRESHOLDS (previous-step xstep lines); the kernel counts them itself
through the stock R0 multi-count pass and admits the tightest rung
with count in [K, kC], then collects and refines as usual. This is
the fully self-contained two-pass shape: no emission of any kind,
pass 1 = one fused 3-rung count (cluster-merge and block-skip
compose unchanged), pass 2 = the stock single-line collect.

Versus use_ext_counts (variant A) the only delta is where the counts
come from; P1's preIdx gather and the P1b quantile rung derivation
are both skipped (the seed lines carry the bracket).

Smoke: 15/15 bit-exact across cs=1/4, block_max skip, and the thin
(all rungs below K) and fat (all counts above kC) miss paths.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 43 ++++++++++++++++++-
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 12 ++++--
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 42bf4507c436..5e1cd25a61d6 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -285,6 +285,7 @@ def __init__(
         p2_warp_redundant: bool = True,
         enable_block_skip: bool = False,
         use_ext_counts: bool = False,
+        ext_rungs: bool = False,
         use_ext_cand: bool = False,
         cand_cap: int = 5120,
         cand_rung: int = 0,
@@ -638,12 +639,25 @@ def __init__(
             self.cap_c = 0
         if use_ext_cand and not use_ext_counts:
             raise ValueError("use_ext_cand requires use_ext_counts")
+        # ext_rungs (two-pass variant B): closed-loop rung THRESHOLDS come
+        # from the host (previous-step xstep lines); the kernel counts them
+        # itself via the stock R0 multi-count and admits the tightest rung
+        # in [K, kC]. Exclusive with use_ext_counts (which also imports
+        # the counts and skips nothing else).
+        self.ext_rungs = bool(ext_rungs) and bool(enable_r0)
+        if self.ext_rungs and bool(use_ext_counts):
+            raise ValueError("ext_rungs is exclusive with use_ext_counts")
         self.use_ext_counts = bool(use_ext_counts) and bool(enable_r0)
         if self.use_ext_counts:
             if not self.fb_fix:
                 raise ValueError("use_ext_counts requires fb_fix")
             if self.M_thr != 3:
                 raise ValueError("use_ext_counts expects exactly 3 seed rungs")
+        if self.ext_rungs:
+            if not self.fb_fix:
+                raise ValueError("ext_rungs requires fb_fix")
+            if self.M_thr != 3:
+                raise ValueError("ext_rungs expects exactly 3 seed rungs")
             # cluster_size > 1 supported: the ext rungs/counts are
             # per-row (identical across the cluster), the stock multi
             # count pass cluster-merges as usual, and the L2 direct
@@ -4846,6 +4860,9 @@ def run_one_row(
         ):
             seed_thr_row = seed_thr[row_idx, None]
             seed_counts_row = seed_counts[row_idx, None]
+        elif cutlass.const_expr(self.ext_rungs and seed_thr is not None):
+            seed_thr_row = seed_thr[row_idx, None]
+            seed_counts_row = None
         else:
             seed_thr_row = None
             seed_counts_row = None
@@ -5449,6 +5466,19 @@ def _run_phases(
         # valid), so the preIdx gather is skipped wholesale. A miss whose
         # target lies outside [t_0, t_2] recovers via the refine loop's
         # 8x bracket expansion (same fail-soft as the stock path).
+        if cutlass.const_expr(self.ext_rungs):
+            # variant B: the rungs carry the bracket, so P1's gather buys
+            # nothing - same seed-line init as the ext-counts hit path.
+            if tidx == cutlass.Int32(0):
+                s_thr[0] = seed_thr_row[1]
+                s_thr[1] = seed_thr_row[0]
+                s_thr[2] = seed_thr_row[2]
+                s_iscalars[0] = cutlass.Int32(0)  # cand_count
+                s_iscalars[1] = cutlass.Int32(0)  # done
+                s_iscalars[2] = cutlass.Int32(-1)  # cnt_lo (fb seeding owns)
+                s_iscalars[3] = cutlass.Int32(-1)  # cnt_hi
+                s_iscalars[4] = cutlass.Int32(0)  # out_count
+            cute.arch.barrier()
         if cutlass.const_expr(self.use_ext_counts):
             if ext_row == cutlass.Int32(1):
                 if tidx == cutlass.Int32(0):
@@ -5480,7 +5510,7 @@ def _run_phases(
                     smem_gath=smem_gath,  # p1b_cache: stash gathered values (None-op OFF)
                     s_mt_thr=s_mt_thr,  # r0_vseed: park pmean in the last rung column
                 )
-        if cutlass.const_expr(not self.use_ext_counts):
+        if cutlass.const_expr(not (self.use_ext_counts or self.ext_rungs)):
             self.phase1_preidx_stats(
                 input_row,
                 N,
@@ -6074,7 +6104,16 @@ def _run_phases(
                                     warp_id,
                                     lane,
                                 )
-                    if cutlass.const_expr(not self.use_ext_counts):
+                    if cutlass.const_expr(self.ext_rungs):
+                        # variant B: rung thresholds = the closed-loop seed
+                        # lines verbatim; the stock multi-count measures
+                        # them and the argmin admission below picks the
+                        # tightest one in [K, kC].
+                        if tidx == cutlass.Int32(0):
+                            for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                                s_mt_thr[m] = seed_thr_row[m]
+                        cute.arch.barrier()
+                    if cutlass.const_expr(not (self.use_ext_counts or self.ext_rungs)):
                         if cutlass.const_expr(self.p1b_cache):
                             # rungs from the SMEM gather-cache P1 stashed (no 2nd
                             # GMEM gather); 16-bit only.
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index fc72170a1447..e540fc9814c3 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -69,6 +69,7 @@ def _compile(
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
     use_ext_cand: bool = False,
+    ext_rungs: bool = False,
     cand_cap: int = 5120,
     accept_cap: "int | None" = None,
     kc_override: "int | None" = None,
@@ -146,7 +147,7 @@ def _compile(
         cute.runtime.make_fake_compact_tensor(
             cutlass.Float32, (n_rows, 3), stride_order=(1, 0), assumed_align=4
         )
-        if use_ext_counts
+        if (use_ext_counts or ext_rungs)
         else None
     )
     seed_counts_fake = (
@@ -206,6 +207,7 @@ def _compile(
         use_ext_counts=use_ext_counts,
         emit_xstate=emit_xstate,
         use_ext_cand=use_ext_cand,
+        ext_rungs=ext_rungs,
         cand_cap=cand_cap,
         accept_cap=accept_cap,
         kc_override=kc_override,
@@ -214,7 +216,7 @@ def _compile(
         # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed. The
         # qfrac VALUES are irrelevant on this path (P1b is skipped) — only
         # the slot count matters.
-        r0_qfracs=(0.85, 0.35) if use_ext_counts else None,
+        r0_qfracs=(0.85, 0.35) if (use_ext_counts or ext_rungs) else None,
     )
     return cute.compile(
         kernel,
@@ -764,6 +766,9 @@ def gvr_topk_decode(
             and cand_idx.shape == (num_rows, _segtot)
         ), f"self_scan position column must be int32 [num_rows, {_segtot}]"
     use_ext_counts = seed_thr is not None and seed_counts is not None
+    # variant B (two-pass): thresholds without counts -> the kernel counts
+    # the rungs itself (stock R0 multi-count) and admits in-kernel
+    ext_rungs = seed_thr is not None and seed_counts is None
     if use_ext_counts:
         assert (
             seed_thr.dtype == torch.float32
@@ -898,6 +903,7 @@ def gvr_topk_decode(
         use_ext_counts,
         emit_xstate,
         use_ext_cand,
+        ext_rungs,
         cand_cap,
         int(os.environ["GVR_BSTAR"]) if "GVR_BSTAR" in os.environ else None,
         int(os.environ["GVR_KC"]) if "GVR_KC" in os.environ else None,
@@ -918,7 +924,7 @@ def gvr_topk_decode(
         out_indices,
         order_row if seqlen_sorted else None,
         block_max if enable_block_skip else None,
-        seed_thr if use_ext_counts else None,
+        seed_thr if (use_ext_counts or ext_rungs) else None,
         seed_counts if use_ext_counts else None,
         xstate if emit_xstate else None,
         cand_vals if use_ext_cand else None,

From 26608ef9cf1c874dfa847693b445c90df737dc09 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 26 Jul 2026 01:21:21 -0700
Subject: [PATCH 031/117] [None][perf] GVR ext counts: pack lines and counts
 into one seed row

The admission preview paid two serial cold gmem reads (12B seed_thr
then 12B seed_counts) on the critical path - 1-1.6us on 8-17us
kernels. Pack both into a single [rows, 8] fp32 row (lines at [0..2],
counts as floats at [3..5], exact to 2^24): one 32B sector serves
everything, and threshold indices stay unchanged.

The wrapper accepts a pre-packed row (width >= 6) natively via
pack_seed(); passing separate tensors remains as a compat path that
builds the pack per call. Cold probe: flash-8k variant-A gap to
variant-B shrinks 1.6 -> 0.8us; flash-64k reaches parity.

Smoke: ea 4/4 + vb 15/15 bit-exact.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 13 ++---
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 54 +++++++++++++------
 2 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 5e1cd25a61d6..afa4e46d3704 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -4855,11 +4855,12 @@ def run_one_row(
             block_max_row = block_max[row_idx, None]
         else:
             block_max_row = None
-        if cutlass.const_expr(
-            self.use_ext_counts and seed_thr is not None and seed_counts is not None
-        ):
+        if cutlass.const_expr(self.use_ext_counts and seed_thr is not None):
+            # packed seed row [>=6] fp32: [0..2] lines, [3..5] counts as
+            # floats (exact to 2^24) - ONE 32B sector serves both, halving
+            # the serial cold loads of the admission preview
             seed_thr_row = seed_thr[row_idx, None]
-            seed_counts_row = seed_counts[row_idx, None]
+            seed_counts_row = None
         elif cutlass.const_expr(self.ext_rungs and seed_thr is not None):
             seed_thr_row = seed_thr[row_idx, None]
             seed_counts_row = None
@@ -5406,7 +5407,7 @@ def _run_phases(
         if cutlass.const_expr(self.use_ext_counts):
             if seed_thr_row[0] < cutlass.Float32(1e37):
                 for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                    cm_e = cutlass.Int32(seed_counts_row[m])
+                    cm_e = cutlass.Int32(seed_thr_row[3 + m])
                     if cm_e >= cutlass.Int32(self.top_k) and cm_e <= cutlass.Int32(self.kC):
                         ext_row = cutlass.Int32(1)
             # list path preview: when the SoA candidate list will be taken
@@ -6062,7 +6063,7 @@ def _run_phases(
                                 bx_m = cutlass.Int32(-1)
                                 bx_c = cutlass.Int32(2147483647)
                                 for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                                    cx = cutlass.Int32(seed_counts_row[m])
+                                    cx = cutlass.Int32(seed_thr_row[3 + m])
                                     if (
                                         cx >= cutlass.Int32(self.top_k)
                                         and cx <= cutlass.Int32(self.kC)
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index e540fc9814c3..d4e5e2a4b098 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -143,20 +143,19 @@ def _compile(
         if enable_block_skip
         else None
     )
+    # ext counts ride PACKED with the lines ([rows, 8] fp32: lines at
+    # [0..2], counts as floats at [3..5]) - one 32B sector per row
     seed_thr_fake = (
         cute.runtime.make_fake_compact_tensor(
-            cutlass.Float32, (n_rows, 3), stride_order=(1, 0), assumed_align=4
+            cutlass.Float32,
+            (n_rows, 8 if use_ext_counts else 3),
+            stride_order=(1, 0),
+            assumed_align=4,
         )
         if (use_ext_counts or ext_rungs)
         else None
     )
-    seed_counts_fake = (
-        cute.runtime.make_fake_compact_tensor(
-            cutlass.Int32, (n_rows, 3), stride_order=(1, 0), assumed_align=4
-        )
-        if use_ext_counts
-        else None
-    )
+    seed_counts_fake = None
     cand_vals_fake = (
         cute.runtime.make_fake_compact_tensor(
             cutlass.Float32, (n_rows, cand_cap), stride_order=(1, 0), assumed_align=4
@@ -521,6 +520,19 @@ def derive_seed_lines_v4(
     return out.contiguous()
 
 
+def pack_seed(seed_thr: torch.Tensor, seed_counts: torch.Tensor) -> torch.Tensor:
+    """Pack lines + exact counts into one [rows, 8] fp32 seed row.
+
+    Lines land at [0..2], counts as floats at [3..5] (exact to 2^24);
+    one 32B sector per row. Build ONCE per step, outside any timed
+    region.
+    """
+    pack = torch.zeros((seed_thr.shape[0], 8), dtype=torch.float32, device=seed_thr.device)
+    pack[:, 0:3] = seed_thr
+    pack[:, 3:6] = seed_counts.float()
+    return pack.contiguous()
+
+
 def emu_seed_counts(
     logits: torch.Tensor,
     seq_lens: torch.Tensor,
@@ -765,18 +777,23 @@ def gvr_topk_decode(
             and cand_idx.is_contiguous()
             and cand_idx.shape == (num_rows, _segtot)
         ), f"self_scan position column must be int32 [num_rows, {_segtot}]"
-    use_ext_counts = seed_thr is not None and seed_counts is not None
+    # packed seed row ([rows, >=6] fp32: lines + counts-as-floats) is the
+    # native ext-counts input; separate seed_counts is the compat path and
+    # pays a per-call pack build - pre-pack with pack_seed() instead.
+    pre_packed = seed_thr is not None and seed_thr.shape[1] >= 6
+    use_ext_counts = seed_thr is not None and (seed_counts is not None or pre_packed)
     # variant B (two-pass): thresholds without counts -> the kernel counts
     # the rungs itself (stock R0 multi-count) and admits in-kernel
-    ext_rungs = seed_thr is not None and seed_counts is None
+    ext_rungs = seed_thr is not None and seed_counts is None and not pre_packed
     if use_ext_counts:
         assert (
             seed_thr.dtype == torch.float32
             and seed_thr.is_cuda
             and seed_thr.is_contiguous()
-            and seed_thr.shape == (num_rows, 3)
-        ), "seed_thr must be contiguous CUDA fp32 [num_rows, 3]"
-        assert (
+            and seed_thr.shape[0] == num_rows
+            and (pre_packed or seed_thr.shape[1] == 3)
+        ), "seed_thr must be contiguous CUDA fp32 [num_rows, 3|8]"
+        assert pre_packed or (
             seed_counts.dtype == torch.int32
             and seed_counts.is_cuda
             and seed_counts.is_contiguous()
@@ -882,6 +899,13 @@ def gvr_topk_decode(
             else:
                 min_blocks_per_mp = 1
 
+    seed_pack = None
+    if use_ext_counts:
+        if pre_packed:
+            seed_pack = seed_thr
+        else:
+            seed_pack = pack_seed(seed_thr, seed_counts)
+
     compiled = _compile(
         cute_dtype,
         top_k,
@@ -924,8 +948,8 @@ def gvr_topk_decode(
         out_indices,
         order_row if seqlen_sorted else None,
         block_max if enable_block_skip else None,
-        seed_thr if (use_ext_counts or ext_rungs) else None,
-        seed_counts if use_ext_counts else None,
+        seed_pack if use_ext_counts else (seed_thr if ext_rungs else None),
+        None,
         xstate if emit_xstate else None,
         cand_vals if use_ext_cand else None,
         cand_idx if (use_ext_cand or self_scan) else None,

From 09699c9d27ca9f2a37ed7184f11eb3f916351b09 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 26 Jul 2026 01:30:26 -0700
Subject: [PATCH 032/117] [None][perf] GVR: open the K>512 skip gate for parked
 ext counts; publish rung counts

Two changes:

1. The stock-path small-batch gate (drop block_max at K > 512, tiny
   batch) was tuned in the loose-rung era, when the sample-quantile
   list kept 60%+ of the blocks. With ext counts the admitted line is
   parked in every rung slot BEFORE the count pass, so the active list
   forms at the accepted threshold and skips by the real band density.
   Bypass the gate for parked ext counts: pro-1M list goes 0 -> 1039
   blocks (12.7% pass, exactly the host-side prediction) and the cell
   runs 36.7 -> 16.6us warm (2.1x); v32 64k-256k all flip to the skip
   walk (15.6-16.1 vs 17.2-18.2us row-split). All probes bit-exact.

2. ext_rungs publishes its three measured rung counts in
   xstate[4..6] (production builds only; debug stamp layouts win the
   slots otherwise) - the closed loop derives next-step lines from
   real device counts instead of a host-side recount emulation.

Also adds GVR_SKIP_DBG forensics (list length / current flag / chosen
rung / drop mask via xstate) used to find the gate.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py         | 16 ++++++++++++++++
 .../cute_dsl_kernels/top_k/run_gvr_topk.py     | 18 +++++++++++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index afa4e46d3704..afff4efd3906 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -49,6 +49,7 @@
 # P4 sub-phase clock64 breakdown -> xstate[1,2,4,5,6,7] (debug: clobbers
 # the closed-loop thr/anch publish; single-shot cells only, not chains)
 _P4_SUB_DBG = bool(int(os.environ.get("GVR_P4_SUB_DBG", "0")))
+_SKIP_DBG = bool(int(os.environ.get("GVR_SKIP_DBG", "0")))
 
 
 # ---------------------------------------------------------------------------
@@ -6564,6 +6565,21 @@ def _run_phases(
                         # cand_count_p4 = pre-P4 snapshot (P4 repurposes
                         # the s_iscalars slots).
                         xstate_row[3] = cutlass.Float32(cand_count_p4)
+                        if cutlass.const_expr(
+                            self.ext_rungs and not _P4_SUB_DBG and not _P4_TAIL_DBG
+                        ):
+                            # closed-loop food: the three rung counts this
+                            # step measured (exact, straight from the R0
+                            # multi-count) - the host derives next-step
+                            # lines from these instead of re-counting.
+                            xstate_row[4] = cutlass.Float32(s_mt_cnt[0])
+                            xstate_row[5] = cutlass.Float32(s_mt_cnt[1])
+                            xstate_row[6] = cutlass.Float32(s_mt_cnt[2])
+                        if cutlass.const_expr(_SKIP_DBG and self.enable_block_skip):
+                            xstate_row[4] = cutlass.Float32(s_active_cnt[0])
+                            xstate_row[5] = cutlass.Float32(s_active_cnt[1])
+                            xstate_row[6] = cutlass.Float32(s_r0col[0])
+                            xstate_row[7] = cutlass.Float32(s_active_cnt[2])
             else:
                 # cs>1: only the leader (CTA 0 in cluster) runs Phase 4.
                 if is_leader:
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index d4e5e2a4b098..eb4209a59dbc 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -747,9 +747,21 @@ def gvr_topk_decode(
     # (kC/K = 6 vs 10), the bounds prune less, and the row-split configs
     # win outright (cold protocol, pro 262k BS1: skip 21.3-21.6us at
     # cs1/cs8 vs stock cs8 19.7us) -> keep the stock path.
-    if block_max is not None and num_rows < 8 and top_k > 512 and not self_scan:
-        # self_scan stage 2 owns its own skip economics (phase-0 block
-        # loop) — the stock-path small-batch gate does not apply
+    if (
+        block_max is not None
+        and num_rows < 8
+        and top_k > 512
+        and not self_scan
+        and not (seed_thr is not None and seed_counts is not None)
+        and not (seed_thr is not None and seed_thr.shape[1] >= 6)
+    ):
+        # Stock-path small-batch gate, tuned in the loose-rung era (the
+        # sample-quantile list kept 60%+ of the blocks at K>512). It does
+        # NOT apply when the admitted line is known up front: ext counts
+        # park the tight line in every rung slot, so the list forms at
+        # the accepted threshold and skips by the real band density
+        # (pro-1M: 12.7% pass -> the list is live again). self_scan owns
+        # its own skip economics likewise.
         block_max = None
     if self_scan:
         # fused self-contained mode: kernel scans/buckets the row itself.

From 993082c85acd9e4b5000e3885e68aee966fa224f Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 26 Jul 2026 01:39:33 -0700
Subject: [PATCH 033/117] [None][fix] GVR ext_rungs: cluster-side count publish
 + runtime rung validity

Two chain-integration fixes for the two-pass variant B:

1. The rung-count xstate publish only existed on the cs=1 exit; at
   cs>1 the closed loop read zeros, derived garbage lines, and - worse
   - inverted rungs broke bit-exactness downstream. Publish the
   cluster-merged counts from the leader exit too.

2. Rung validity is now a RUNTIME check (finite and strictly
   ascending); invalid lines fall back to the stock P1 seed +
   P1b quantile-rung path in-kernel, so exactness never rides on the
   host loop's line quality (cold start, dropped rows, NaN counts).

p640 capture chain: published counts match the host emulation
bit-for-bit (cntdiff=0 every step), all arms exact, admission
0.75-1.0. Smokes: ea 4/4 + vb 15/15.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 94 +++++++++++++++----
 1 file changed, 77 insertions(+), 17 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index afff4efd3906..ce372a829359 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -5468,19 +5468,53 @@ def _run_phases(
         # valid), so the preIdx gather is skipped wholesale. A miss whose
         # target lies outside [t_0, t_2] recovers via the refine loop's
         # 8x bracket expansion (same fail-soft as the stock path).
+        rungs_ok = cutlass.Int32(0)
         if cutlass.const_expr(self.ext_rungs):
-            # variant B: the rungs carry the bracket, so P1's gather buys
-            # nothing - same seed-line init as the ext-counts hit path.
-            if tidx == cutlass.Int32(0):
-                s_thr[0] = seed_thr_row[1]
-                s_thr[1] = seed_thr_row[0]
-                s_thr[2] = seed_thr_row[2]
-                s_iscalars[0] = cutlass.Int32(0)  # cand_count
-                s_iscalars[1] = cutlass.Int32(0)  # done
-                s_iscalars[2] = cutlass.Int32(-1)  # cnt_lo (fb seeding owns)
-                s_iscalars[3] = cutlass.Int32(-1)  # cnt_hi
-                s_iscalars[4] = cutlass.Int32(0)  # out_count
-            cute.arch.barrier()
+            # runtime validity: finite AND strictly ascending; anything
+            # else (cold start, dropped row, NaN from the host loop)
+            # falls back to the stock seed path below - exactness never
+            # rides on the host's line quality
+            if (
+                seed_thr_row[0] < cutlass.Float32(1e37)
+                and seed_thr_row[0] > cutlass.Float32(-1e37)
+                and seed_thr_row[1] > seed_thr_row[0]
+                and seed_thr_row[2] > seed_thr_row[1]
+            ):
+                rungs_ok = cutlass.Int32(1)
+        if cutlass.const_expr(self.ext_rungs):
+            if rungs_ok == cutlass.Int32(1):
+                # variant B: the rungs carry the bracket, so P1's gather
+                # buys nothing - same seed-line init as the ext-counts
+                # hit path.
+                if tidx == cutlass.Int32(0):
+                    s_thr[0] = seed_thr_row[1]
+                    s_thr[1] = seed_thr_row[0]
+                    s_thr[2] = seed_thr_row[2]
+                    s_iscalars[0] = cutlass.Int32(0)  # cand_count
+                    s_iscalars[1] = cutlass.Int32(0)  # done
+                    s_iscalars[2] = cutlass.Int32(-1)  # cnt_lo (fb owns)
+                    s_iscalars[3] = cutlass.Int32(-1)  # cnt_hi
+                    s_iscalars[4] = cutlass.Int32(0)  # out_count
+                cute.arch.barrier()
+            if rungs_ok == cutlass.Int32(0):
+                self.phase1_preidx_stats(
+                    input_row,
+                    N,
+                    pre_idx_row,
+                    pre_idx_count,
+                    pre_idx_offset,
+                    smem_wmin,
+                    smem_wmax,
+                    smem_wsum,
+                    smem_wcnt_p1,
+                    s_thr,
+                    s_iscalars,
+                    tidx,
+                    warp_id,
+                    lane,
+                    smem_gath=smem_gath,
+                    s_mt_thr=s_mt_thr,
+                )
         if cutlass.const_expr(self.use_ext_counts):
             if ext_row == cutlass.Int32(1):
                 if tidx == cutlass.Int32(0):
@@ -6110,11 +6144,28 @@ def _run_phases(
                         # variant B: rung thresholds = the closed-loop seed
                         # lines verbatim; the stock multi-count measures
                         # them and the argmin admission below picks the
-                        # tightest one in [K, kC].
-                        if tidx == cutlass.Int32(0):
-                            for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                                s_mt_thr[m] = seed_thr_row[m]
-                        cute.arch.barrier()
+                        # tightest one in [K, kC]. Invalid lines fall back
+                        # to the stock P1b quantile rungs (P1 stats ran on
+                        # this row in that case).
+                        if rungs_ok == cutlass.Int32(1):
+                            if tidx == cutlass.Int32(0):
+                                for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                                    s_mt_thr[m] = seed_thr_row[m]
+                            cute.arch.barrier()
+                        if rungs_ok == cutlass.Int32(0):
+                            self.phase1b_hspace_rungs(
+                                input_row,
+                                N,
+                                pre_idx_row,
+                                pre_idx_count,
+                                pre_idx_offset,
+                                smem_hist,
+                                s_thr,
+                                s_mt_thr,
+                                tidx,
+                                warp_id,
+                                lane,
+                            )
                     if cutlass.const_expr(not (self.use_ext_counts or self.ext_rungs)):
                         if cutlass.const_expr(self.p1b_cache):
                             # rungs from the SMEM gather-cache P1 stashed (no 2nd
@@ -6665,6 +6716,15 @@ def _run_phases(
                             xstate_row[1] = s_thr[0]
                             xstate_row[2] = s_thr[0]
                             xstate_row[3] = cutlass.Float32(cand_count_p4)
+                            if cutlass.const_expr(
+                                self.ext_rungs and not _P4_SUB_DBG and not _P4_TAIL_DBG
+                            ):
+                                # cluster-merged rung counts (identical on
+                                # every CTA after the multi-count DSMEM
+                                # aggregation)
+                                xstate_row[4] = cutlass.Float32(s_mt_cnt[0])
+                                xstate_row[5] = cutlass.Float32(s_mt_cnt[1])
+                                xstate_row[6] = cutlass.Float32(s_mt_cnt[2])
 
         # Final cluster barrier: keep peer CTAs (and their SMEM) alive
         # until the leader's gather + Phase 4 finish. Skipped at

From ed64e6ca7b6e8a5b2f69dce0df95484239ce9bb6 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 26 Jul 2026 04:27:40 -0700
Subject: [PATCH 034/117] [None][fix] GVR ext counts: full line validity on
 both admission previews

The counts/list admission previews guarded only t0 finiteness; a NaN
in t1/t2 (or inverted lines) from the host closed loop could get
parked into the refine brackets - the same failure mode that broke
ext_rungs bit-exactness on-chain before its runtime validity check.
Both previews now require all three lines finite and strictly
ascending; invalid rows fall to the stock seed path.

Unified smoke 30/30 bit-exact including adversarial cases (NaN mid
line, inverted lines) across v5-list, va and vb modes; v5 callers now
pass the pre-packed [rows, 8] seed row (the per-call compat build is
off the hot path).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py           | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index ce372a829359..3a4a70d22b12 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -5406,7 +5406,18 @@ def _run_phases(
         ck1 = cutlass.Int64(0)
         ext_row = cutlass.Int32(0)
         if cutlass.const_expr(self.use_ext_counts):
-            if seed_thr_row[0] < cutlass.Float32(1e37):
+            # line validity mirrors ext_rungs: ALL THREE lines finite and
+            # strictly ascending. The old t0-only guard let a NaN in
+            # t1/t2 get parked into the refine brackets (the same failure
+            # mode that broke ext_rungs exactness on-chain); invalid rows
+            # fall to the stock path, exactness never rides on the host
+            # loop's line quality.
+            if (
+                seed_thr_row[0] < cutlass.Float32(1e37)
+                and seed_thr_row[0] > cutlass.Float32(-1e37)
+                and seed_thr_row[1] > seed_thr_row[0]
+                and seed_thr_row[2] > seed_thr_row[1]
+            ):
                 for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
                     cm_e = cutlass.Int32(seed_thr_row[3 + m])
                     if cm_e >= cutlass.Int32(self.top_k) and cm_e <= cutlass.Int32(self.kC):
@@ -5428,6 +5439,9 @@ def _run_phases(
                     and claimed_p >= cutlass.Int32(self.top_k + 64)
                     and claimed_p <= cutlass.Int32(self.list_cap)
                     and seed_thr_row[0] < cutlass.Float32(1e37)
+                    and seed_thr_row[0] > cutlass.Float32(-1e37)
+                    and seed_thr_row[1] > seed_thr_row[0]
+                    and seed_thr_row[2] > seed_thr_row[1]
                 ):
                     ext_row = cutlass.Int32(1)
             # ---- self_scan phase 0: fused scan-bucket ----

From 9b611a961697d953b845d0f0cf725793c849d8e6 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 26 Jul 2026 05:53:14 -0700
Subject: [PATCH 035/117] [None][perf] GVR ext counts: single-column parked
 count + merged admission prologue

The parked M-ary pass counted the SAME admitted threshold in all
three rung columns - 3x the compares and 3 per-thread count columns
for identical values. Non-skip builds now count it ONCE through the
refine primitive (same per-thread ptcnt cache and cluster merge that
Phase 3 consumes) and accept in place; the M-ary machinery, argmin,
handoff and miss refine stand down via the pre-marked rung column.
A measured count outside the band (stale host state) falls back to
the full M-ary rerun on the three distinct seed lines.

Rung parking also folds into the P1-init thread0 block: one barrier
covers the whole admission prologue instead of two.

Cold probe (variant A vs B, same node): mid rows flip in A's favor -
pro-64k/128k and flash-128k 0.95-0.96x of B (from ~1.05x), cs8 at
parity; skip builds keep the M-ary pass (list build lives there).
Unified smoke 30/30 bit-exact incl. adversarial lines.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 140 +++++++++++++-----
 1 file changed, 103 insertions(+), 37 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 3a4a70d22b12..1c65e5e80ea2 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -5540,6 +5540,34 @@ def _run_phases(
                     s_iscalars[2] = cutlass.Int32(-1)  # cnt_lo (fb seeding owns)
                     s_iscalars[3] = cutlass.Int32(-1)  # cnt_hi
                     s_iscalars[4] = cutlass.Int32(0)  # out_count
+                    # rung parking folded into the SAME thread0 block
+                    # (was a second thread0 block + barrier in the R0
+                    # region): tightest in-band count picks the admitted
+                    # line; single-column builds stage it in s_thr[0] and
+                    # pre-mark the rung column (M_qf = accepted, -2 =
+                    # defensive M-ary rerun).
+                    bx_m = cutlass.Int32(-1)
+                    bx_c = cutlass.Int32(2147483647)
+                    for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                        cx = cutlass.Int32(seed_thr_row[3 + m])
+                        if (
+                            cx >= cutlass.Int32(self.top_k)
+                            and cx <= cutlass.Int32(self.kC)
+                            and cx < bx_c
+                        ):
+                            bx_m = cutlass.Int32(m)
+                            bx_c = cx
+                    for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
+                        if bx_m >= cutlass.Int32(0):
+                            s_mt_thr[m] = seed_thr_row[bx_m]
+                        else:
+                            s_mt_thr[m] = seed_thr_row[m]
+                    if cutlass.const_expr(not self.enable_block_skip):
+                        if bx_m >= cutlass.Int32(0):
+                            s_thr[0] = seed_thr_row[bx_m]
+                            s_r0col[0] = cutlass.Int32(self.M_qf)
+                        else:
+                            s_r0col[0] = cutlass.Int32(-2)
                 cute.arch.barrier()
             if ext_row == cutlass.Int32(0):
                 self.phase1_preidx_stats(
@@ -6108,24 +6136,51 @@ def _run_phases(
                             # count (+ list build at that threshold) and classify
                             # admits it; a full miss keeps the 3 distinct rungs as
                             # measured brackets for the seeded refine.
-                            if tidx == cutlass.Int32(0):
-                                bx_m = cutlass.Int32(-1)
-                                bx_c = cutlass.Int32(2147483647)
-                                for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                                    cx = cutlass.Int32(seed_thr_row[3 + m])
-                                    if (
-                                        cx >= cutlass.Int32(self.top_k)
-                                        and cx <= cutlass.Int32(self.kC)
-                                        and cx < bx_c
-                                    ):
-                                        bx_m = cutlass.Int32(m)
-                                        bx_c = cx
-                                for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)):
-                                    if bx_m >= cutlass.Int32(0):
-                                        s_mt_thr[m] = seed_thr_row[bx_m]
-                                    else:
-                                        s_mt_thr[m] = seed_thr_row[m]
-                            cute.arch.barrier()
+                            # rung parking + single-path staging happen in
+                            # the P1-init thread0 block (one barrier for
+                            # the whole admission prologue).
+                            if cutlass.const_expr(not self.enable_block_skip):
+                                # v3: the parked M-ary pass counted the SAME
+                                # threshold in all three columns (3x compare
+                                # + 3 ptcnt columns for identical values).
+                                # Count it ONCE with the refine primitive -
+                                # same per-thread ptcnt cache and cluster
+                                # merge P3 consumes - and accept in place.
+                                if s_r0col[0] == cutlass.Int32(self.M_qf):
+                                    self.block_count_ge(
+                                        input_row,
+                                        slice_start,
+                                        slice_end,
+                                        s_thr[0],
+                                        smem_ptcnt,
+                                        smem_wcnt,
+                                        s_iscalars,
+                                        s_cluster_partial,
+                                        tidx,
+                                        warp_id,
+                                        lane,
+                                        do_cluster_sync=do_cluster_sync,
+                                        smem_input=smem_input,
+                                    )
+                                    cute.arch.barrier()
+                                    if tidx == cutlass.Int32(0):
+                                        cpar = s_iscalars[0]
+                                        if cpar >= cutlass.Int32(
+                                            self.top_k
+                                        ) and cpar <= cutlass.Int32(self.kC):
+                                            s_iscalars[1] = cutlass.Int32(1)
+                                        else:
+                                            # emission counts disagree with
+                                            # the measured count (stale
+                                            # host state): rerun the full
+                                            # M-ary machinery on the three
+                                            # distinct seed lines.
+                                            s_r0col[0] = cutlass.Int32(-2)
+                                            for m in cutlass.range_constexpr(
+                                                cutlass.const_expr(self.M_thr)
+                                            ):
+                                                s_mt_thr[m] = seed_thr_row[m]
+                                    cute.arch.barrier()
                         if ext_row == cutlass.Int32(0):
                             if cutlass.const_expr(self.p1b_cache):
                                 # rungs from the SMEM gather-cache P1 stashed (no 2nd
@@ -6208,26 +6263,37 @@ def _run_phases(
                                 warp_id,
                                 lane,
                             )
-                    self.block_count_ge_multi(
-                        input_row,
-                        slice_start,
-                        slice_end,
-                        s_mt_thr,
-                        smem_ptcnt_multi,
-                        smem_wcnt_multi,
-                        s_mt_cnt,
-                        s_cluster_partial_m,
-                        do_cluster_sync,
-                        tidx,
-                        warp_id,
-                        lane,
-                        smem_ptcnt=smem_ptcnt,
-                        block_max_row=block_max_row,
-                        smem_active=smem_active,
-                        s_active_cnt=s_active_cnt,
-                    )
+                    r0_par = cutlass.Int32(0)
+                    if cutlass.const_expr(self.use_ext_counts and not self.enable_block_skip):
+                        # single-column fast path accepted: the parked count
+                        # is done and admitted; the M-ary pass, argmin,
+                        # handoff and miss machinery all stand down
+                        # (s_r0col == M_qf skips the copy and the refine).
+                        if s_r0col[0] == cutlass.Int32(self.M_qf) and s_iscalars[
+                            1
+                        ] == cutlass.Int32(1):
+                            r0_par = cutlass.Int32(1)
+                    if r0_par == cutlass.Int32(0):
+                        self.block_count_ge_multi(
+                            input_row,
+                            slice_start,
+                            slice_end,
+                            s_mt_thr,
+                            smem_ptcnt_multi,
+                            smem_wcnt_multi,
+                            s_mt_cnt,
+                            s_cluster_partial_m,
+                            do_cluster_sync,
+                            tidx,
+                            warp_id,
+                            lane,
+                            smem_ptcnt=smem_ptcnt,
+                            block_max_row=block_max_row,
+                            smem_active=smem_active,
+                            s_active_cnt=s_active_cnt,
+                        )
                     cute.arch.barrier()
-                    if tidx == 0:
+                    if tidx == 0 and r0_par == cutlass.Int32(0):
                         # tightest admissible rung = SMALLEST count in [K, kC].
                         # (Explicit argmin: with r0_vseed the pmean column is not
                         # sorted into the rung order; for sorted rungs this is

From fc0137aa54ba7db81b4a294b5650d167e4677321 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 26 Jul 2026 06:16:29 -0700
Subject: [PATCH 036/117] [None][perf] GVR wrapper: 512 threads for small-K
 list-hit rows

List-hit rows do O(list) work (~2-4K entries at K<=512), not O(N);
the N-keyed thread heuristic handed them 1024 threads and paid the
barrier-heavy refine phases for no parallel gain. Key the pick on the
mode: ext-cand with K<=512 takes 512 threads (flash-512k v5 hit path
8.4 -> 7.1us cold, -15%); K=1024 lists are large enough to keep the
N-keyed pick (pro cells within +/-4%).

Unified smoke 30/30 bit-exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index eb4209a59dbc..268264b9b956 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -871,7 +871,14 @@ def gvr_topk_decode(
     N_cols = logits.shape[1]
     N_dec = max_seq_len if max_seq_len is not None else N_cols
     if num_threads_per_block is None:
-        if self_scan:
+        if use_ext_cand and top_k <= 512:
+            # list-hit rows do O(list) work (~2-4K entries at K<=512),
+            # not O(N): the N-keyed 1024-thread pick pays barrier-heavy
+            # phases for no parallel gain (flash-512k v5: 8.4 -> 7.1us
+            # at 512 threads; K=1024 lists are big enough to keep the
+            # N-keyed pick).
+            num_threads_per_block = 512
+        elif self_scan:
             # self_scan owns the whole row scan in one CTA: the phase-0
             # cp.async pipeline scales with warp count at every N (the
             # 512-thread short-row heuristic below is tuned for the

From a6728fa9f266a565038c24d8e3166843c8525635 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 26 Jul 2026 07:39:50 -0700
Subject: [PATCH 037/117] [None][perf] GVR ext counts: fused count-free take
 for short rows

With the admitted line's exact count already known from emission, the
count pass exists only to build Phase 3's placement prefix. Short
rows (N < 16384) now claim-collect in ONE pass: 4 chunks in flight
per iteration (the count primitive's ILP shape) and per-element
direct smem atomics for passers (no warp sync, claims hide under the
read stream); the claim total re-verifies the count and take_cand
bypasses P2+P3 wholesale, landing straight in P4 (candidate-order
agnostic, as the v5 list path already proves). Out-of-band claims
(stale host state) reset and fall back to the stock single-count
path.

Two falsified shapes are documented by the arc: per-element scalar
loads (list-walk habit, +9us at 128k) and per-position ballots (the
self_scan-falsified warp-lock pattern, +3us). The N-gate keeps the
claim-dense 16k+ cells (pro band at 9.4%) on count-then-place, which
measures at parity or better there.

Cold probe vs variant B on the same node: every gated cell wins -
flash-8k/32k 0.95/0.89, pro-16k/32k 0.89/0.90; out-of-gate cells stay
at their two-pass optima (0.93-0.97). Unified smoke 30/30 bit-exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 120 ++++++++++++++++++
 1 file changed, 120 insertions(+)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 1c65e5e80ea2..7dad3d33a8f1 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -6026,6 +6026,126 @@ def _run_phases(
                 if cutlass.const_expr(_P4_TAIL_DBG):
                     ck1 = cute.arch.clock64()
 
+            # ---- parked count-free take (ext counts, no list): the
+            # emission already measured the admitted line's exact count,
+            # so the count pass exists ONLY to build P3's placement
+            # prefix. Claim-collect instead (v5 edge-cut walk shape:
+            # 4-way strided reads, ballot-merged claims): ONE cold pass
+            # replaces count(cold) + collect(L2-hot). P4 is candidate-
+            # order agnostic (the v5 list path feeds it emission-claim
+            # order already). The claim total re-measures the count; a
+            # mismatch with the parked band (stale host state) leaves
+            # take_cand=0 and the stock single-count path below recovers.
+            if cutlass.const_expr(
+                self.use_ext_counts
+                and not self.use_ext_cand
+                and not self.self_scan
+                and not self.enable_block_skip
+                and cluster_size == 1
+                and self.dtype == cutlass.Float32
+            ):
+                if ext_row == cutlass.Int32(1) and N < cutlass.Int32(16384):
+                    # short rows only: the fused walk wins ~1us below
+                    # ~16k (pass fixed costs dominate there); at 16k+
+                    # the claim-dense cells (pro 9.4% band) lose to
+                    # count-then-place and the rest sit at parity, so
+                    # the two-pass path keeps them (measured both ways,
+                    # loncheng cells; captures all sit above the gate)
+                    cut_p = s_thr[0]  # parked line (staged at P1-init)
+                    rbase = input_row.iterator.toint()
+                    vw_p = cutlass.const_expr(self.vec_bits // self.dtype.width)
+                    va_p = cutlass.const_expr(self.vec_align_bytes)
+                    eb_p = cutlass.const_expr(self.dtype.width // 8)
+                    cp_atom = self._make_load_copy_atom()
+                    frag0 = cute.make_fragment((vw_p,), self.dtype)
+                    frag1 = cute.make_fragment((vw_p,), self.dtype)
+                    frag2 = cute.make_fragment((vw_p,), self.dtype)
+                    frag3 = cute.make_fragment((vw_p,), self.dtype)
+                    step4_p = cutlass.const_expr(4 * num_threads * vw_p)
+                    nfull_p = (N // cutlass.Int32(step4_p)) * cutlass.Int32(step4_p)
+                    it_p = tidx * cutlass.Int32(vw_p)
+                    # 4 chunks in flight per iter (the count primitive's
+                    # ILP shape); per-element DIRECT smem atomics for
+                    # passers - no warp sync, claims hide under the reads
+                    while it_p < nfull_p:
+                        for _jf in cutlass.range_constexpr(4):
+                            gp_p = cute.make_ptr(
+                                self.dtype,
+                                rbase
+                                + cutlass.Int64(it_p + cutlass.Int32(_jf * num_threads * vw_p))
+                                * cutlass.Int64(eb_p),
+                                cute.AddressSpace.gmem,
+                                assumed_align=va_p,
+                            )
+                            cute.copy(
+                                cp_atom,
+                                cute.make_tensor(gp_p, cute.make_layout((vw_p,))),
+                                frag0
+                                if _jf == 0
+                                else frag1
+                                if _jf == 1
+                                else frag2
+                                if _jf == 2
+                                else frag3,
+                            )
+                        for _jf in cutlass.range_constexpr(4):
+                            for _jv in cutlass.range_constexpr(vw_p):
+                                v_p = cutlass.Float32(
+                                    (
+                                        frag0
+                                        if _jf == 0
+                                        else frag1
+                                        if _jf == 1
+                                        else frag2
+                                        if _jf == 2
+                                        else frag3
+                                    )[_jv]
+                                )
+                                if v_p >= cut_p:
+                                    sl_p = atomicAdd(
+                                        s_iscalars.iterator + cutlass.Int32(0),
+                                        cutlass.Int32(1),
+                                    )
+                                    if sl_p < cutlass.Int32(self.kC):
+                                        smem_keys[sl_p] = v_p
+                                        smem_vals[sl_p] = (
+                                            it_p
+                                            + cutlass.Int32(_jf * num_threads * vw_p)
+                                            + cutlass.Int32(_jv)
+                                        )
+                        it_p = it_p + cutlass.Int32(step4_p)
+                    i_p = nfull_p + tidx
+                    while i_p < N:
+                        vp_p = cute.make_ptr(
+                            cutlass.Float32,
+                            rbase + cutlass.Int64(i_p) * cutlass.Int64(4),
+                            cute.AddressSpace.gmem,
+                            assumed_align=4,
+                        )
+                        pval = cute.make_tensor(vp_p, cute.make_layout((1,)))[0]
+                        if pval >= cut_p:
+                            sl_p = atomicAdd(
+                                s_iscalars.iterator + cutlass.Int32(0),
+                                cutlass.Int32(1),
+                            )
+                            if sl_p < cutlass.Int32(self.kC):
+                                smem_keys[sl_p] = pval
+                                smem_vals[sl_p] = i_p
+                        i_p = i_p + cutlass.Int32(num_threads)
+                    cute.arch.barrier()
+                    cnt_p = s_iscalars[0]
+                    if cnt_p >= cutlass.Int32(self.top_k) and cnt_p <= cutlass.Int32(self.kC):
+                        take_cand = cutlass.Int32(1)
+                        if tidx == cutlass.Int32(0):
+                            s_iscalars[1] = cutlass.Int32(1)  # done
+                        cute.arch.barrier()
+                    if take_cand == cutlass.Int32(0):
+                        # stale host counts: reset the claim counter so
+                        # the stock single-count path re-measures cleanly
+                        if tidx == cutlass.Int32(0):
+                            s_iscalars[0] = cutlass.Int32(0)
+                        cute.arch.barrier()
+
             # ---- self_scan take: cut straight from the phase-0 cursors ----
             # Same admission state machine as the v5 list (tightest line
             # whose count fits [K, B*] wins; anchor = loosest in-band

From c1c5ada29b99cf1adf4f51d363c45a34f7250a6b Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 04:37:19 -0700
Subject: [PATCH 038/117] [None][perf] FP4 indexer: port fused block-metadata
 emission (P0)

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  95 +++-
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 466 +++++++++++++++++-
 .../test_cute_dsl_fp4_paged_mqa_logits.py     | 212 ++++++++
 3 files changed, 761 insertions(+), 12 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 0db351140735..d387d9f10ba2 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -8250,11 +8250,13 @@ def _compile(cls,
                      num_epi_subtiles,
                      epi_dtype,
                      output_dtype,
-                     remove_online_sf_transpose=False):
+                     remove_online_sf_transpose=False,
+                     emit_block_meta=False,
+                     emit_hit_stats=True):
             """Compile kernel using fake tensors + TVM FFI."""
             key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
                    num_sms, num_epi_subtiles, epi_dtype, output_dtype,
-                   remove_online_sf_transpose)
+                   remove_online_sf_transpose, emit_block_meta, emit_hit_stats)
             if key in cls.kernel_cache:
                 return
 
@@ -8317,6 +8319,26 @@ def _compile(cls,
                                                             (num_ctas, 2),
                                                             stride_order=(1, 0))
 
+            # Block-meta tensors (fused-GVR support): nb_pad*4 records.
+            block_max_fake = None
+            hit_stats_fake = None
+            hit_bitmap_fake = None
+            if emit_block_meta:
+                nb_sym = cute.sym_int()
+                block_max_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Float32, (cute.sym_int(), nb_sym),
+                    stride_order=(1, 0),
+                    assumed_align=16)
+                if emit_hit_stats:
+                    hit_stats_fake = cute.runtime.make_fake_compact_tensor(
+                        cutlass.Float32, (cute.sym_int(), 4),
+                        stride_order=(1, 0),
+                        assumed_align=16)
+                    hit_bitmap_fake = cute.runtime.make_fake_compact_tensor(
+                        cutlass.Int32, (sym_B, cute.sym_int()),
+                        stride_order=(1, 0),
+                        assumed_align=16)
+
             fake_stream = cute.runtime.make_fake_stream(
                 use_tvm_ffi_env_stream=True)
 
@@ -8331,6 +8353,8 @@ def _compile(cls,
                 epi_dtype=to_cutlass[epi_dtype],
                 output_dtype=to_cutlass[output_dtype],
                 remove_online_sf_transpose=remove_online_sf_transpose,
+                emit_block_meta=emit_block_meta,
+                emit_hit_stats=emit_hit_stats,
             )
 
             compiled = cute.compile(
@@ -8345,6 +8369,9 @@ def _compile(cls,
                 sm_fake,
                 cutlass.Int32(1),
                 cutlass.Int32(1),
+                block_max_fake,
+                hit_stats_fake,
+                hit_bitmap_fake,
                 fake_stream,
                 options="--enable-tvm-ffi",
             )
@@ -8366,6 +8393,11 @@ def forward(
             epi_dtype: torch.dtype = torch.float32,
             output_dtype: torch.dtype = torch.float32,
             remove_online_sf_transpose: bool = False,
+            emit_block_meta: bool = False,
+            emit_hit_stats: bool = True,
+            hit_bitmap: Optional[torch.Tensor] = None,
+            block_max_out: Optional[torch.Tensor] = None,
+            hit_stats_out: Optional[torch.Tensor] = None,
         ) -> torch.Tensor:
             """Execute FP4 paged MQA logits kernel.
 
@@ -8381,8 +8413,15 @@ def forward(
                 num_epi_subtiles: epilogue sub-tile count (1, 2, or 4)
                 epi_dtype: epilogue compute dtype
                 output_dtype: output logits dtype
+                emit_block_meta: also emit per-128-block metadata for the
+                    fused GVR top-k. Requires ``hit_bitmap``
+                    [B, nb_pad*4] int32 (1 bit per compressed kv position,
+                    request-level). ``block_max_out``/``hit_stats_out``
+                    ([B*next_n, nb_pad] fp32 / [B*next_n, nb_pad, 4] fp32)
+                    are allocated when not supplied.
             Returns:
-                logits: [B*next_n, max_context_len] output_dtype
+                logits [B*next_n, max_context_len]; with emit_block_meta,
+                the tuple (logits, block_max, hit_stats).
             """
             B, next_n, H, half_D = q.shape
             N = next_n * H
@@ -8435,10 +8474,46 @@ def forward(
             )
             logits = logits[:, :max_context_len]
 
+            # Block-meta buffers (fused-GVR support). nb_pad mirrors the
+            # logits padding so WG1's odd-num_kv OOB tile lands in padding.
+            if emit_block_meta:
+                nb_pad = aligned_max_ctx // compute_block_kv
+                # 4 warp-partial records per block (see FP4MQALogitsKernel).
+                nrec = nb_pad * 4
+                if emit_hit_stats:
+                    assert (
+                        hit_bitmap is not None
+                        and hit_bitmap.dtype == torch.int32
+                        and hit_bitmap.is_cuda and hit_bitmap.is_contiguous()
+                        and hit_bitmap.dim() == 2 and hit_bitmap.shape[0] == B
+                        and hit_bitmap.shape[1] >= nb_pad * 4
+                    ), (f"emit_hit_stats requires hit_bitmap int32 "
+                        f"[{B}, >= {nb_pad * 4}]; got "
+                        f"{None if hit_bitmap is None else (hit_bitmap.dtype, tuple(hit_bitmap.shape))}"
+                        )
+                    # Per-row aggregate {enc_min, enc_max, sum, cnt}. The
+                    # kernel only ATOMICALLY MERGES into this buffer — the
+                    # caller must pre-initialize it to the identities
+                    # {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} each step.
+                    assert hit_stats_out is not None, (
+                        "emit_hit_stats requires a caller-initialized "
+                        "hit_stats_out [B*next_n, 4] fp32 (identity-filled)")
+                    assert (hit_stats_out.shape == (B * next_n, 4)
+                            and hit_stats_out.is_contiguous())
+                else:
+                    hit_bitmap = None
+                    hit_stats_out = None
+                if block_max_out is None:
+                    block_max_out = torch.empty((B * next_n, nrec),
+                                                device=q.device,
+                                                dtype=torch.float32)
+                assert (block_max_out.shape == (B * next_n, nrec)
+                        and block_max_out.is_contiguous())
+
             # Compile if needed (fake tensors, no real data required)
             key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
                    num_epi_subtiles, epi_dtype, output_dtype,
-                   remove_online_sf_transpose)
+                   remove_online_sf_transpose, emit_block_meta, emit_hit_stats)
             if key not in cls.kernel_cache:
                 cls._compile(
                     compute_block_kv,
@@ -8450,12 +8525,20 @@ def forward(
                     num_epi_subtiles,
                     epi_dtype,
                     output_dtype,
-                    remove_online_sf_transpose=remove_online_sf_transpose)
+                    remove_online_sf_transpose=remove_online_sf_transpose,
+                    emit_block_meta=emit_block_meta,
+                    emit_hit_stats=emit_hit_stats)
             compiled = cls.kernel_cache[key]
 
             # TVM FFI: pass raw tensors, no dlpack/stream needed
+            if emit_block_meta:
+                compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
+                         context_lens, schedule_meta, num_phys_blocks, B,
+                         block_max_out, hit_stats_out, hit_bitmap)
+                return logits, block_max_out, hit_stats_out
             compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
-                     context_lens, schedule_meta, num_phys_blocks, B)
+                     context_lens, schedule_meta, num_phys_blocks, B, None,
+                     None, None)
             return logits
 
     @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits",
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index 94db5d143085..dfbdc00dbed5 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -64,6 +64,79 @@
 # form is also accepted by older wrappers, so keep it version-independent.
 _RND_RN = "rn"
 
+# Reduction identities for the emit_block_meta epilogue path; match the GVR
+# kernel's FLT_MAX/NEG_FLT_MAX sentinels (gvr_topk_decode.py) so the fused
+# Phase 1's degenerate checks behave identically to the gather path.
+_META_FLT_MAX = 3.4028235e38
+_META_NEG_FLT_MAX = -3.4028235e38
+
+
+# ---------------------------------------------------------------------------
+# Global-memory reductions for the per-row hit aggregate (emit_hit_stats).
+# fp32 has no native atomic min/max; we use the standard order-preserving
+# int encoding enc(f) = bits(f) >= 0 ? bits(f) : bits(f) ^ 0x7FFFFFFF
+# (an involution; signed-int order == float order, -0.0 quirk harmless for
+# min/max seeding) and red.global.{min,max}.s32. Sum uses red.global.add.f32
+# (order-nondeterministic — perturbs only the heuristic mean seed).
+# ---------------------------------------------------------------------------
+@dsl_user_op
+def _red_global_fmin_ordered(addr_i64, fval, *, loc=None, ip=None):
+    llvm.inline_asm(
+        None,
+        [addr_i64.ir_value(loc=loc, ip=ip), fval.ir_value(loc=loc, ip=ip)],
+        "{\n\t"
+        ".reg .b32 k;\n\t"
+        ".reg .pred p;\n\t"
+        "mov.b32 k, $1;\n\t"
+        "setp.lt.s32 p, k, 0;\n\t"
+        "@p xor.b32 k, k, 0x7FFFFFFF;\n\t"
+        "red.global.min.s32 [$0], k;\n\t"
+        "}",
+        "l,f",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _red_global_fmax_ordered(addr_i64, fval, *, loc=None, ip=None):
+    llvm.inline_asm(
+        None,
+        [addr_i64.ir_value(loc=loc, ip=ip), fval.ir_value(loc=loc, ip=ip)],
+        "{\n\t"
+        ".reg .b32 k;\n\t"
+        ".reg .pred p;\n\t"
+        "mov.b32 k, $1;\n\t"
+        "setp.lt.s32 p, k, 0;\n\t"
+        "@p xor.b32 k, k, 0x7FFFFFFF;\n\t"
+        "red.global.max.s32 [$0], k;\n\t"
+        "}",
+        "l,f",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _red_global_add_f32(addr_i64, fval, *, loc=None, ip=None):
+    llvm.inline_asm(
+        None,
+        [addr_i64.ir_value(loc=loc, ip=ip), fval.ir_value(loc=loc, ip=ip)],
+        "red.global.add.f32 [$0], $1;",
+        "l,f",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
 
 @dsl_user_op
 def pack_f16x2(
@@ -363,6 +436,8 @@ def __init__(
         output_dtype=cutlass.Float32,
         remove_online_sf_transpose: bool = False,
         use_batched_store: bool = True,
+        emit_block_meta: bool = False,
+        emit_hit_stats: bool = True,
     ):
         # Static FP4 invariants — see plan Sanity checklist.
         assert num_heads == 64, "FP4 kernel hardcodes num_heads=64 for TMEM/SMEM budget"
@@ -407,6 +482,38 @@ def __init__(
         # When True, defer per-t STG to register array and emit all STGs in
         # one contiguous LSU phase after the for-t loop (epilogue micro-opt).
         self.use_batched_store = use_batched_store
+        # When True, the epilogue additionally emits per-128-token-block
+        # metadata consumed by the fused GVR top-k (gvr_topk_decode.py —
+        # fused_preidx_stats / enable_block_skip). Emission is fully
+        # WARP-AUTONOMOUS: each of the WG's 4 warps writes one partial
+        # record per tile per t (record index = tile*4 + warp); the GVR
+        # consumer folds the 4 partials per block. No cross-warp barrier —
+        # a per-tile named-barrier fold costs +53% indexer wall-clock.
+        #   block_max [num_rows, nb_pad*4]    fp32 — warp-partial max of
+        #     f32(stored logit) over valid positions (kv_pos < ctx);
+        #     fold(4) is the lossless block-skip upper bound. Computed on
+        #     the POST-conversion value so it bounds what GVR reads back,
+        #     bit-exactly.
+        #   hit_agg [num_rows, 4] fp32 — PER-ROW aggregate
+        #     {enc_min, enc_max, sum, cnt} of stored logits at positions
+        #     flagged in hit_bitmap; min/max slots hold the
+        #     order-preserving int encoding (see _red_global_fmin_ordered).
+        #     Buffer must be pre-initialized to
+        #     {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step. Replaces GVR
+        #     Phase 1's random gather.
+        #   hit_bitmap [batch, nb_pad*4]      int32 — 1 bit per kv
+        #     position (request-level; from the previous step's top-k).
+        # Cost per tile per t: 1 bitmap LDG/thread + 1 warp redux + 1
+        # lane-0 STG (block_max) + <= 4 LANE-LOCAL accumulator ops (hit
+        # stats); the hit accumulators flush ONCE per q-transition via
+        # atomics (_flush_hit_agg). A per-tile warp-redux emission of the
+        # hit stats was tried first and cost +72% indexer wall-clock.
+        #
+        # emit_hit_stats sub-knob (only meaningful with emit_block_meta):
+        # False emits block_max ONLY (no bitmap read, no hit aggregate) —
+        # sufficient for GVR block-skip without the fused Phase 1.
+        self.emit_block_meta = emit_block_meta
+        self.emit_hit_stats = emit_hit_stats
         # epi_bytes covers fp16 and bf16 (FP8 only handled fp16).
         self.epi_bytes = 2 if epi_dtype in (cutlass.Float16, cutlass.BFloat16) else 4
         # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
@@ -627,6 +734,9 @@ def __call__(
         schedule_meta: cute.Tensor,  # [num_sms+1, 2] int32
         num_phys_blocks: cutlass.Int32,
         batch_size: cutlass.Int32,
+        block_max: cute.Tensor,  # [num_rows, nb_pad*4] fp32 warp-partials (or None)
+        hit_stats: cute.Tensor,  # [num_rows, 4] fp32 hit aggregate (or None unless emit_hit_stats)
+        hit_bitmap: cute.Tensor,  # [batch, nb_pad*4] int32 (or None unless emit_block_meta)
         stream: cuda.CUstream,
     ):
         # Derive KV data and SF views from the fused uint8 buffer.
@@ -826,6 +936,9 @@ class SharedStorage:
             context_lens,
             schedule_meta,
             batch_size,
+            block_max,
+            hit_stats,
+            hit_bitmap,
             self.cluster_layout_vmnk,
             self.a_smem_layout_staged,
             self.b_smem_layout_staged,
@@ -848,6 +961,44 @@ class SharedStorage:
             stream=stream,
         )
 
+    @cute.jit
+    def _flush_hit_agg(
+        self,
+        mHitAgg,  # [num_rows, 4] fp32 {enc_min, enc_max, sum, cnt}
+        q_flush,  # request whose accumulators are being flushed
+        hacc_min,
+        hacc_max,
+        hacc_sum,
+        hacc_cnt,
+        meta_lane,
+    ):
+        """Warp-reduce the per-lane hit accumulators and merge them into the
+        per-row global aggregate via one set of atomics (encoded-int
+        min/max + fp32 adds), then reset to identities. Called once per
+        q-transition per warp — atomic traffic is negligible. The
+        aggregate buffer must be pre-initialized to
+        {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step (prototype: test
+        harness; production: folded into the bitmap prepare kernel)."""
+        next_n = cutlass.const_expr(self.next_n)
+        base_addr = mHitAgg.iterator.toint()
+        for t in cutlass.range_constexpr(next_n):
+            w_min = cute.arch.warp_redux_sync(hacc_min[t], "fmin")
+            w_max = cute.arch.warp_redux_sync(hacc_max[t], "fmax")
+            w_sum = cute.arch.warp_reduction_sum(hacc_sum[t])
+            w_cnt = cute.arch.warp_redux_sync(hacc_cnt[t], "add")
+            if meta_lane == cutlass.Int32(0):
+                if w_cnt > cutlass.Int32(0):
+                    row = q_flush * cutlass.Int32(next_n) + cutlass.Int32(t)
+                    row_addr = base_addr + cutlass.Int64(row) * cutlass.Int64(16)
+                    _red_global_fmin_ordered(row_addr, w_min)
+                    _red_global_fmax_ordered(row_addr + cutlass.Int64(4), w_max)
+                    _red_global_add_f32(row_addr + cutlass.Int64(8), w_sum)
+                    _red_global_add_f32(row_addr + cutlass.Int64(12), cutlass.Float32(w_cnt))
+            hacc_min[t] = cutlass.Float32(_META_FLT_MAX)
+            hacc_max[t] = cutlass.Float32(_META_NEG_FLT_MAX)
+            hacc_sum[t] = cutlass.Float32(0.0)
+            hacc_cnt[t] = cutlass.Int32(0)
+
     @cute.kernel
     def kernel(
         self,
@@ -867,6 +1018,9 @@ def kernel(
         mContextLens: cute.Tensor,  # [batch_size]
         mScheduleMeta: cute.Tensor,  # [num_sms+1, 2] int32
         batch_size: cutlass.Int32,
+        mBlockMax: cute.Tensor,  # [num_rows, nb_pad*4] fp32 warp-partials (or None)
+        mHitAgg: cute.Tensor,  # [num_rows, 4] fp32 {enc_min, enc_max, sum, cnt} (or None)
+        mHitBitmap: cute.Tensor,  # [batch, nb_pad*4] int32 (or None)
         cluster_layout_vmnk: cute.Layout,
         a_smem_layout_staged: cute.ComposedLayout,
         b_smem_layout_staged: cute.ComposedLayout,
@@ -1115,6 +1269,12 @@ def kernel(
             layout=sf_kv_smem_layout_staged,
             byte_alignment=128,
         )
+        # Block-meta emission is fully warp-autonomous (each warp's lane 0
+        # writes its own warp-partial record straight to GMEM; the GVR
+        # side folds 4 partials per block) — no SMEM scratch, no named
+        # barrier. A cross-warp SMEM fold + per-tile named barrier was
+        # tried first and cost +53% indexer wall-clock by serializing the
+        # epilogue's warp pipelining.
 
         a_mcast_mask = cpasync.create_tma_multicast_mask(
             cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
@@ -1661,7 +1821,6 @@ def kernel(
                             ) // block_kv_val
                     # Update while-loop condition
                     has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
-
         elif is_umma_warp_1:
             # UMMA warp for group 1
             # Explicitly waits on Q pipeline — critical because TMA warp 1
@@ -1794,7 +1953,6 @@ def kernel(
                             ) // block_kv_val
                     # Update while-loop condition
                     has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
-
         elif is_math_warp:
             cute.arch.warpgroup_reg_alloc(240)
 
@@ -1838,6 +1996,15 @@ def kernel(
                     MAX_NUM_W_IN_REG = 64
                 else:  # fp32, 4-byte weights
                     MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64
+                if cutlass.const_expr(self.emit_block_meta):
+                    # Free ~8 registers for the meta accumulators/fragments
+                    # — the epilogue's weight cache is tuned to the spill
+                    # edge (see MAX_NUM_W_IN_REG SASS notes above).
+                    MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
+                    if cutlass.const_expr(self.emit_hit_stats):
+                        # Hit accumulators + bitmap word add ~6 more live
+                        # registers across the tile loop.
+                        MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the
@@ -1848,6 +2015,33 @@ def kernel(
                 else:
                     result_arr = None
                 q_stage_local = cutlass.Int32(0)
+                if cutlass.const_expr(self.emit_block_meta):
+                    ctx_cur = cutlass.Int32(0)
+                    meta_warp = local_tidx // 32
+                    meta_lane = local_tidx % 32
+                    if cutlass.const_expr(self.emit_hit_stats):
+                        # Per-lane hit accumulators, carried across ALL
+                        # tiles of the same q and flushed once per
+                        # q-transition — zero warp-wide ops per tile (the
+                        # per-tile redux version cost +72% wall-clock).
+                        hacc_min = cute.make_fragment(next_n, cutlass.Float32)
+                        hacc_max = cute.make_fragment(next_n, cutlass.Float32)
+                        hacc_sum = cute.make_fragment(next_n, cutlass.Float32)
+                        hacc_cnt = cute.make_fragment(next_n, cutlass.Int32)
+                        for _t in cutlass.range_constexpr(next_n):
+                            hacc_min[_t] = cutlass.Float32(_META_FLT_MAX)
+                            hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX)
+                            hacc_sum[_t] = cutlass.Float32(0.0)
+                            hacc_cnt[_t] = cutlass.Int32(0)
+                        # Batched bitmap read state: a per-tile LDG's
+                        # ~300cy L2 trip lands on the critical path every
+                        # tile (isolated at +42% kernel time; one-ahead
+                        # prefetch did not help). All 32 lanes of a warp
+                        # need the SAME word per tile, so instead lane l
+                        # loads the word for tile j+l once per 32 tiles
+                        # and each tile takes its word via one shuffle.
+                        meta_j = cutlass.Int32(0)
+                        hitw_batch = cutlass.Int32(0)
 
                 while has_work:
                     # fetch_next_task: commit next → current
@@ -1869,11 +2063,51 @@ def kernel(
                                 w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
                                     (t_i * num_heads + w_j, q_stage_local)
                                 ]
+                        if cutlass.const_expr(self.emit_block_meta):
+                            # Flush the PREVIOUS request's hit accumulators
+                            # before switching context.
+                            if cutlass.const_expr(self.emit_hit_stats):
+                                if q_idx_old < batch_size:
+                                    self._flush_hit_agg(
+                                        mHitAgg,
+                                        q_idx_old,
+                                        hacc_min,
+                                        hacc_max,
+                                        hacc_sum,
+                                        hacc_cnt,
+                                        meta_lane,
+                                    )
+                                # New bitmap row: invalidate the batched
+                                # word cache (forces a reload).
+                                meta_j = cutlass.Int32(0)
+                            # Compressed-space context len; the meta valid
+                            # mask (kv_pos < ctx_cur) keeps GEMM garbage in
+                            # the aligned padding region out of block_max.
+                            ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 0 (kv_idx + 0)
                     # Unconditional Math: OOB results
                     # written to aligned padding region in logits buffer.
                     kv_pos = kv_idx * block_kv_val + m_coord
+                    if cutlass.const_expr(self.emit_block_meta):
+                        meta_kv_tile = kv_idx
+                        meta_valid = kv_pos < ctx_cur
+                        if cutlass.const_expr(self.emit_hit_stats):
+                            # Warp-uniform reload once per 32 tiles: this
+                            # WG's tile at counter j+l is kv_tile + 2*l,
+                            # whose warp word index is (kv_tile+2*l)*4 +
+                            # warp. Clamp keeps end-of-row lanes in
+                            # bounds (their tiles are never consumed).
+                            if (meta_j & cutlass.Int32(31)) == cutlass.Int32(0):
+                                w_idx = (
+                                    meta_kv_tile + cutlass.Int32(2) * meta_lane
+                                ) * cutlass.Int32(4) + meta_warp
+                                w_idx = min(w_idx, mHitBitmap.shape[1] - cutlass.Int32(1))
+                                hitw_batch = mHitBitmap[(q_idx, w_idx)]
+                            hit_word = cute.arch.shuffle_sync(
+                                hitw_batch, meta_j & cutlass.Int32(31)
+                            )
+                            meta_j = meta_j + cutlass.Int32(1)
 
                     # Step 5.7: drop kv_pipeline.consumer_wait/release and
                     # scale_val LDS — UMMA owns KV+SF pipe; SF is baked into
@@ -2054,11 +2288,74 @@ def kernel(
                         else:
                             result_t = s0x + s0y + s1x + s1y
                         # Step 5.7: drop * scale_val (FP4 SF baked into acc).
+                        stored_t = self.output_dtype(result_t)
                         if cutlass.const_expr(self.use_batched_store):
-                            result_arr[t] = self.output_dtype(result_t)
+                            result_arr[t] = stored_t
                         else:
                             out_row = q_idx * next_n + t
-                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t)
+                            mLogits[(out_row, kv_pos)] = stored_t
+                        if cutlass.const_expr(self.emit_block_meta):
+                            # Meta reduction on the POST-conversion value so
+                            # block_max bounds what GVR reads back bit-exactly
+                            # (pre-conversion fp32 max could round up past a
+                            # stored logit's converted value's block max).
+                            f32_t = cutlass.Float32(stored_t)
+                            bmax_v = cutlass.Float32(_META_NEG_FLT_MAX)
+                            if meta_valid:
+                                bmax_v = f32_t
+                            r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax")
+                            # Warp-autonomous store: record index =
+                            # tile*4 + warp; the GVR consumer folds the 4
+                            # warp-partials per block (fold of partials ==
+                            # block stat — associative + identity-padded).
+                            if meta_lane == cutlass.Int32(0):
+                                out_row_m = q_idx * next_n + t
+                                rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp
+                                mBlockMax[(out_row_m, rec_m)] = r_bmax
+                            if cutlass.const_expr(self.emit_hit_stats):
+                                # Lane-local accumulation, fully BRANCHLESS
+                                # (data-dependent `if meta_hit` compiled to
+                                # real divergent branches whose condition
+                                # waits on the bitmap LDG — 43% of all warp
+                                # stall samples). Bit-mask select is also
+                                # NaN-safe for OOB-tile garbage logits. No
+                                # valid-mask needed: the bitmap contract
+                                # only sets bits inside [0, ctx).
+                                meta_hit = (
+                                    hit_word >> (kv_pos & cutlass.Int32(31))
+                                ) & cutlass.Int32(1)
+                                msk = cutlass.Int32(0) - meta_hit  # 0 / ~0
+                                inv = cutlass.Int32(-1) - msk
+                                fbits = cutlass.Int32(
+                                    llvm.bitcast(cutlass.Int32.mlir_type, f32_t.ir_value())
+                                )
+                                # bits(+FLT_MAX)=0x7F7FFFFF,
+                                # bits(-FLT_MAX)=0xFF7FFFFF (as i32: neg).
+                                selmin = cutlass.Float32(
+                                    llvm.bitcast(
+                                        cutlass.Float32.mlir_type,
+                                        (
+                                            (fbits & msk) | (cutlass.Int32(0x7F7FFFFF) & inv)
+                                        ).ir_value(),
+                                    )
+                                )
+                                selmax = cutlass.Float32(
+                                    llvm.bitcast(
+                                        cutlass.Float32.mlir_type,
+                                        (
+                                            (fbits & msk) | (cutlass.Int32(-8388609) & inv)
+                                        ).ir_value(),
+                                    )
+                                )
+                                seladd = cutlass.Float32(
+                                    llvm.bitcast(
+                                        cutlass.Float32.mlir_type, (fbits & msk).ir_value()
+                                    )
+                                )
+                                hacc_min[t] = cutlass.min(hacc_min[t], selmin)
+                                hacc_max[t] = cutlass.max(hacc_max[t], selmax)
+                                hacc_sum[t] = hacc_sum[t] + seladd
+                                hacc_cnt[t] = hacc_cnt[t] + meta_hit
 
                     if cutlass.const_expr(self.use_batched_store):
                         # Batched STG: all result_arr[t] → mLogits in one pass.
@@ -2078,6 +2375,13 @@ def kernel(
                     # Update while-loop condition
                     has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
 
+                # Flush the final request's hit accumulators (WG 0).
+                if cutlass.const_expr(self.emit_block_meta and self.emit_hit_stats):
+                    if q_idx < batch_size:
+                        self._flush_hit_agg(
+                            mHitAgg, q_idx, hacc_min, hacc_max, hacc_sum, hacc_cnt, meta_lane
+                        )
+
                 # Release last Q stage (WG 0)
                 if q_idx < batch_size:
                     q_pipeline.consumer_release(q_cons_state)
@@ -2101,6 +2405,15 @@ def kernel(
                     MAX_NUM_W_IN_REG = 64
                 else:
                     MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64
+                if cutlass.const_expr(self.emit_block_meta):
+                    # Free ~8 registers for the meta accumulators/fragments
+                    # — the epilogue's weight cache is tuned to the spill
+                    # edge (see MAX_NUM_W_IN_REG SASS notes above).
+                    MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
+                    if cutlass.const_expr(self.emit_hit_stats):
+                        # Hit accumulators + bitmap word add ~6 more live
+                        # registers across the tile loop.
+                        MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the
@@ -2111,6 +2424,33 @@ def kernel(
                 else:
                     result_arr = None
                 q_stage_local = cutlass.Int32(0)
+                if cutlass.const_expr(self.emit_block_meta):
+                    ctx_cur = cutlass.Int32(0)
+                    meta_warp = local_tidx // 32
+                    meta_lane = local_tidx % 32
+                    if cutlass.const_expr(self.emit_hit_stats):
+                        # Per-lane hit accumulators, carried across ALL
+                        # tiles of the same q and flushed once per
+                        # q-transition — zero warp-wide ops per tile (the
+                        # per-tile redux version cost +72% wall-clock).
+                        hacc_min = cute.make_fragment(next_n, cutlass.Float32)
+                        hacc_max = cute.make_fragment(next_n, cutlass.Float32)
+                        hacc_sum = cute.make_fragment(next_n, cutlass.Float32)
+                        hacc_cnt = cute.make_fragment(next_n, cutlass.Int32)
+                        for _t in cutlass.range_constexpr(next_n):
+                            hacc_min[_t] = cutlass.Float32(_META_FLT_MAX)
+                            hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX)
+                            hacc_sum[_t] = cutlass.Float32(0.0)
+                            hacc_cnt[_t] = cutlass.Int32(0)
+                        # Batched bitmap read state: a per-tile LDG's
+                        # ~300cy L2 trip lands on the critical path every
+                        # tile (isolated at +42% kernel time; one-ahead
+                        # prefetch did not help). All 32 lanes of a warp
+                        # need the SAME word per tile, so instead lane l
+                        # loads the word for tile j+l once per 32 tiles
+                        # and each tile takes its word via one shuffle.
+                        meta_j = cutlass.Int32(0)
+                        hitw_batch = cutlass.Int32(0)
 
                 while has_work:
                     # fetch_next_task: commit next → current
@@ -2132,12 +2472,56 @@ def kernel(
                                 w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
                                     (t_i * num_heads + w_j, q_stage_local)
                                 ]
+                        if cutlass.const_expr(self.emit_block_meta):
+                            # Flush the PREVIOUS request's hit accumulators
+                            # before switching context.
+                            if cutlass.const_expr(self.emit_hit_stats):
+                                if q_idx_old < batch_size:
+                                    self._flush_hit_agg(
+                                        mHitAgg,
+                                        q_idx_old,
+                                        hacc_min,
+                                        hacc_max,
+                                        hacc_sum,
+                                        hacc_cnt,
+                                        meta_lane,
+                                    )
+                                # New bitmap row: invalidate the batched
+                                # word cache (forces a reload).
+                                meta_j = cutlass.Int32(0)
+                            # Compressed-space context len; the meta valid
+                            # mask (kv_pos < ctx_cur) keeps GEMM garbage in
+                            # the aligned padding region out of block_max.
+                            ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 1 (kv_idx + 1)
                     # Unconditional Math
                     kv_idx_1 = kv_idx + 1
 
                     kv_pos = kv_idx_1 * block_kv_val + m_coord
+                    if cutlass.const_expr(self.emit_block_meta):
+                        meta_kv_tile = kv_idx_1
+                        # See WG 0. For the odd-num_kv OOB tile
+                        # (kv_idx_1 == num_kv) every lane has
+                        # kv_pos >= ctx_cur, so identities land in the
+                        # nb_pad padding slot — never read by GVR.
+                        meta_valid = kv_pos < ctx_cur
+                        if cutlass.const_expr(self.emit_hit_stats):
+                            # Warp-uniform reload once per 32 tiles: this
+                            # WG's tile at counter j+l is kv_tile + 2*l,
+                            # whose warp word index is (kv_tile+2*l)*4 +
+                            # warp. Clamp keeps end-of-row lanes in
+                            # bounds (their tiles are never consumed).
+                            if (meta_j & cutlass.Int32(31)) == cutlass.Int32(0):
+                                w_idx = (
+                                    meta_kv_tile + cutlass.Int32(2) * meta_lane
+                                ) * cutlass.Int32(4) + meta_warp
+                                w_idx = min(w_idx, mHitBitmap.shape[1] - cutlass.Int32(1))
+                                hitw_batch = mHitBitmap[(q_idx, w_idx)]
+                            hit_word = cute.arch.shuffle_sync(
+                                hitw_batch, meta_j & cutlass.Int32(31)
+                            )
+                            meta_j = meta_j + cutlass.Int32(1)
 
                     # Step 5.7: drop kv_pipeline.consumer_wait/release and
                     # scale_val LDS — UMMA owns KV+SF pipe.
@@ -2307,11 +2691,74 @@ def kernel(
                         else:
                             result_t = s0x + s0y + s1x + s1y
                         # Step 5.7: drop * scale_val (FP4 SF baked into acc).
+                        stored_t = self.output_dtype(result_t)
                         if cutlass.const_expr(self.use_batched_store):
-                            result_arr[t] = self.output_dtype(result_t)
+                            result_arr[t] = stored_t
                         else:
                             out_row = q_idx * next_n + t
-                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t)
+                            mLogits[(out_row, kv_pos)] = stored_t
+                        if cutlass.const_expr(self.emit_block_meta):
+                            # Meta reduction on the POST-conversion value so
+                            # block_max bounds what GVR reads back bit-exactly
+                            # (pre-conversion fp32 max could round up past a
+                            # stored logit's converted value's block max).
+                            f32_t = cutlass.Float32(stored_t)
+                            bmax_v = cutlass.Float32(_META_NEG_FLT_MAX)
+                            if meta_valid:
+                                bmax_v = f32_t
+                            r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax")
+                            # Warp-autonomous store: record index =
+                            # tile*4 + warp; the GVR consumer folds the 4
+                            # warp-partials per block (fold of partials ==
+                            # block stat — associative + identity-padded).
+                            if meta_lane == cutlass.Int32(0):
+                                out_row_m = q_idx * next_n + t
+                                rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp
+                                mBlockMax[(out_row_m, rec_m)] = r_bmax
+                            if cutlass.const_expr(self.emit_hit_stats):
+                                # Lane-local accumulation, fully BRANCHLESS
+                                # (data-dependent `if meta_hit` compiled to
+                                # real divergent branches whose condition
+                                # waits on the bitmap LDG — 43% of all warp
+                                # stall samples). Bit-mask select is also
+                                # NaN-safe for OOB-tile garbage logits. No
+                                # valid-mask needed: the bitmap contract
+                                # only sets bits inside [0, ctx).
+                                meta_hit = (
+                                    hit_word >> (kv_pos & cutlass.Int32(31))
+                                ) & cutlass.Int32(1)
+                                msk = cutlass.Int32(0) - meta_hit  # 0 / ~0
+                                inv = cutlass.Int32(-1) - msk
+                                fbits = cutlass.Int32(
+                                    llvm.bitcast(cutlass.Int32.mlir_type, f32_t.ir_value())
+                                )
+                                # bits(+FLT_MAX)=0x7F7FFFFF,
+                                # bits(-FLT_MAX)=0xFF7FFFFF (as i32: neg).
+                                selmin = cutlass.Float32(
+                                    llvm.bitcast(
+                                        cutlass.Float32.mlir_type,
+                                        (
+                                            (fbits & msk) | (cutlass.Int32(0x7F7FFFFF) & inv)
+                                        ).ir_value(),
+                                    )
+                                )
+                                selmax = cutlass.Float32(
+                                    llvm.bitcast(
+                                        cutlass.Float32.mlir_type,
+                                        (
+                                            (fbits & msk) | (cutlass.Int32(-8388609) & inv)
+                                        ).ir_value(),
+                                    )
+                                )
+                                seladd = cutlass.Float32(
+                                    llvm.bitcast(
+                                        cutlass.Float32.mlir_type, (fbits & msk).ir_value()
+                                    )
+                                )
+                                hacc_min[t] = cutlass.min(hacc_min[t], selmin)
+                                hacc_max[t] = cutlass.max(hacc_max[t], selmax)
+                                hacc_sum[t] = hacc_sum[t] + seladd
+                                hacc_cnt[t] = hacc_cnt[t] + meta_hit
 
                     if cutlass.const_expr(self.use_batched_store):
                         # Batched STG: all result_arr[t] → mLogits in one pass.
@@ -2331,6 +2778,13 @@ def kernel(
                     # Update while-loop condition
                     has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
 
+                # Flush the final request's hit accumulators (WG 1).
+                if cutlass.const_expr(self.emit_block_meta and self.emit_hit_stats):
+                    if q_idx < batch_size:
+                        self._flush_hit_agg(
+                            mHitAgg, q_idx, hacc_min, hacc_max, hacc_sum, hacc_cnt, meta_lane
+                        )
+
                 # Release last Q stage (WG 1)
                 if q_idx < batch_size:
                     q_pipeline.consumer_release(q_cons_state)
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
index 8dae8201ba89..40df2d060159 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
@@ -547,6 +547,218 @@ def test_cute_dsl_fp4_paged_mqa_logits(
     )
 
 
+# ---------------------------------------------------------------------------
+# Block-meta emission (emit_block_meta — fused-GVR support).
+# ---------------------------------------------------------------------------
+
+_FLT_MAX_F32 = torch.finfo(torch.float32).max
+
+
+def _enc_ordered_f32(t: torch.Tensor) -> torch.Tensor:
+    """Order-preserving int encoding of fp32 (involution; also decodes)."""
+    bits = t.float().contiguous().view(torch.int32)
+    enc = torch.where(bits >= 0, bits, bits ^ 0x7FFFFFFF)
+    return enc.view(torch.float32)
+
+
+def _hit_agg_identities(num_rows: int, device) -> torch.Tensor:
+    ident = torch.tensor([_FLT_MAX_F32, -_FLT_MAX_F32], dtype=torch.float32, device=device)
+    enc = _enc_ordered_f32(ident)
+    out = torch.zeros((num_rows, 4), dtype=torch.float32, device=device)
+    out[:, 0] = enc[0]
+    out[:, 1] = enc[1]
+    return out.contiguous()
+
+
+def _pack_hit_bitmap(
+    pre_idx: torch.Tensor, batch_size: int, num_words: int, device
+) -> torch.Tensor:
+    """[B, num_words] int32; bit (pos % 32) of word (pos // 32) set per
+    valid pre_idx entry — the kernel's hit test layout."""
+    bitmap = torch.zeros((batch_size, num_words), dtype=torch.int64, device=device)
+    for b in range(batch_size):
+        idx = pre_idx[b].to(torch.int64).unique()
+        idx = idx[(idx >= 0) & (idx < num_words * 32)]
+        bitmap[b].scatter_add_(0, idx >> 5, torch.ones_like(idx) << (idx & 31))
+    # int64 -> int32 with bit-31 wraparound (torch refuses the overflow).
+    wrapped = bitmap & 0xFFFFFFFF
+    wrapped = torch.where(wrapped >= 2**31, wrapped - 2**32, wrapped)
+    return wrapped.to(torch.int32)
+
+
+@skip_not_sm100
+@pytest.mark.parametrize("batch_size", [1, 4])
+@pytest.mark.parametrize("next_n", [1, 2, 3])
+# 4224 = 33 blocks of 128 -> odd num_kv exercises WG1's OOB padding tile.
+@pytest.mark.parametrize("avg_ctx", [4096, 4224])
+@pytest.mark.parametrize("phys_block_kv", [64, 128])
+@pytest.mark.parametrize("fix_length", [True, False])
+@pytest.mark.parametrize("emit_hit_stats", [True, False])
+def test_cute_dsl_fp4_paged_mqa_logits_block_meta(
+    batch_size,
+    next_n,
+    avg_ctx,
+    phys_block_kv,
+    fix_length,
+    emit_hit_stats,
+):
+    """emit_block_meta correctness: block_max / hit_stats recomputed from
+    the KERNEL'S OWN logits output (fp4 numerics differ from the torch
+    reference logits, but the meta contract is defined on what the kernel
+    stores). NaN-prefilled buffers prove no writes land outside
+    [0, num_kv (+1 when odd)) per row."""
+    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner
+
+    torch.manual_seed(7)
+    torch.cuda.manual_seed(7)
+    num_heads, head_dim, top_k = 64, 128, 512
+    max_model_len = max(avg_ctx * 2, 2048)
+    device = "cuda"
+
+    if fix_length:
+        context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device)
+    else:
+        lo = max(phys_block_kv, int(0.7 * avg_ctx))
+        context_lens = torch.randint(
+            lo, int(1.3 * avg_ctx) + 1, (batch_size,), dtype=torch.int32, device=device
+        ).clamp(max=max_model_len)
+
+    num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv)
+    num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2
+    max_blocks_per_seq = int(num_blocks_per_seq.max().item())
+    block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device)
+    pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32)
+    off = 0
+    for i, n_blks in enumerate(num_blocks_per_seq.tolist()):
+        block_table[i, :n_blks] = pool[off : off + n_blks]
+        off += n_blks
+
+    q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16)
+    kv_cache = torch.randn(
+        (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16
+    )
+    weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32)
+
+    q_packed, sf_q_packed = per_token_cast_to_fp4(
+        q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True
+    )
+    q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2)
+    sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads)
+    remove_online_sf_transpose = phys_block_kv == 128
+    kv_fused, _ = kv_cache_cast_to_fp4(
+        kv_cache, remove_online_sf_transpose=remove_online_sf_transpose
+    )
+
+    DG_METADATA_BLOCK_KV = 64
+    num_sms = deep_gemm.get_num_sms()
+    schedule_meta = deep_gemm.get_paged_mqa_logits_metadata(
+        context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms
+    )
+
+    # pre_idx per request within [0, ctx) -> packed bitmap.
+    aligned_max_ctx = align(max_model_len, 256)
+    nb_pad = aligned_max_ctx // 128
+    pre_idx = torch.zeros((batch_size, top_k), dtype=torch.int32, device=device)
+    for b in range(batch_size):
+        pre_idx[b] = torch.randint(
+            0, int(context_lens[b].item()), (top_k,), dtype=torch.int32, device=device
+        )
+    bitmap = _pack_hit_bitmap(pre_idx, batch_size, nb_pad * 4, device)
+
+    # block_max: 4 warp-partial records per block; consumers fold. NaN
+    # prefill proves write coverage is exactly [0, written_hi*4) per row.
+    # hit_stats: per-row aggregate the kernel atomically merges into —
+    # MUST be identity-initialized by the caller.
+    nan = float("nan")
+    block_max = torch.full(
+        (batch_size * next_n, nb_pad * 4), nan, dtype=torch.float32, device=device
+    )
+    hit_stats = _hit_agg_identities(batch_size * next_n, device)
+
+    meta_kwargs = dict(
+        emit_block_meta=True,
+        emit_hit_stats=emit_hit_stats,
+        block_max_out=block_max,
+    )
+    if emit_hit_stats:
+        meta_kwargs.update(hit_bitmap=bitmap, hit_stats_out=hit_stats)
+    logits, bm, hs = CuteDSLFP4PagedMQALogitsRunner.forward(
+        q_fp4,
+        sf_q,
+        kv_fused,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_model_len,
+        num_epi_subtiles=1,
+        epi_dtype=torch.float32,
+        output_dtype=torch.bfloat16,
+        remove_online_sf_transpose=remove_online_sf_transpose,
+        **meta_kwargs,
+    )
+    torch.cuda.synchronize()
+
+    lf = logits.float()
+    for row in range(batch_size * next_n):
+        req = row // next_n
+        ctx = int(context_lens[req].item())
+        num_kv = ceil_div(ctx, 128)
+        tag = f"row={row} req={req} ctx={ctx} next_n={next_n} pbk={phys_block_kv}"
+
+        # Fold the kernel's 4 warp-partials per block (the consumer-side
+        # contract); per-warp partials themselves depend on the TMEM
+        # lane->row mapping and are not checked individually.
+        bm_fold = bm[row].view(nb_pad, 4).amax(-1)
+
+        # block_max reference from the kernel's own stored logits.
+        padded = torch.full((nb_pad * 128,), -_FLT_MAX_F32, device=device)
+        padded[:ctx] = lf[row, :ctx]
+        ref_bmax = padded.view(nb_pad, 128).amax(-1)
+        torch.testing.assert_close(
+            bm_fold[:num_kv],
+            ref_bmax[:num_kv],
+            atol=0.0,
+            rtol=0.0,
+            msg=lambda m, tag=tag: f"block_max mismatch: {tag}\n{m}",
+        )
+
+        if emit_hit_stats:
+            # Per-row hit aggregate reference (bitmap semantics: dedup +
+            # pos < ctx). min/max slots are encoded (involution decodes).
+            idx = pre_idx[req].to(torch.int64).unique()
+            idx = idx[(idx >= 0) & (idx < ctx)]
+            got_min = _enc_ordered_f32(hs[row, 0:1])[0]
+            got_max = _enc_ordered_f32(hs[row, 1:2])[0]
+            got_sum = hs[row, 2]
+            got_cnt = hs[row, 3]
+            if idx.numel() > 0:
+                vals = lf[row, idx]
+                assert got_min.item() == vals.min().item(), f"hit_min: {tag}"
+                assert got_max.item() == vals.max().item(), f"hit_max: {tag}"
+                # Atomic-add merge order vs torch sum order: fp slack.
+                torch.testing.assert_close(
+                    got_sum,
+                    vals.sum(),
+                    atol=1e-2,
+                    rtol=1e-4,
+                    msg=lambda m, tag=tag: f"hit_sum mismatch: {tag}\n{m}",
+                )
+                assert got_cnt.item() == float(idx.numel()), f"hit_cnt: {tag}"
+            else:
+                assert got_min.item() == _FLT_MAX_F32, f"identity min: {tag}"
+                assert got_max.item() == -_FLT_MAX_F32, f"identity max: {tag}"
+                assert got_cnt.item() == 0.0, f"identity cnt: {tag}"
+
+        # Odd num_kv: WG1's OOB tile writes pure identities into block
+        # slot num_kv (every lane invalid).
+        written_hi = num_kv + (num_kv % 2)
+        if written_hi > num_kv:
+            assert bm_fold[num_kv].item() == -_FLT_MAX_F32, tag
+        # No stray writes past the padding tile: NaN prefill intact.
+        assert bm[row, written_hi * 4 :].isnan().all(), f"stray block_max write: {tag}"
+
+
 # ---------------------------------------------------------------------------
 # Benchmarking entry point (run module directly).
 # ---------------------------------------------------------------------------

From b5421cc8bf8c3c7b573a97814748536a7dff851d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 07:19:17 -0700
Subject: [PATCH 039/117] [None][perf] FP4 indexer: emit per-row seed counts in
 the GEMM epilogue (L1)

Lane-local count(logit >= thr_j) accumulation (3 thresholds/row, fp32
post-conversion domain, block_max valid mask), flushed at q-transition
and tile-loop tail via warp-redux + lane0 red.global.add.s32. Caller
zero-initializes the counts buffer each step. Plumbed through the host
runner (emit_seed_counts compile key, fake tensors, contract asserts)
and covered by an exactness unit test against the kernel's own logits.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  60 +++++++-
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 103 +++++++++++++
 .../test_cute_dsl_fp4_paged_mqa_logits.py     | 141 ++++++++++++++++++
 3 files changed, 298 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index d387d9f10ba2..a419613fbf76 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -8252,11 +8252,13 @@ def _compile(cls,
                      output_dtype,
                      remove_online_sf_transpose=False,
                      emit_block_meta=False,
-                     emit_hit_stats=True):
+                     emit_hit_stats=True,
+                     emit_seed_counts=False):
             """Compile kernel using fake tensors + TVM FFI."""
             key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
                    num_sms, num_epi_subtiles, epi_dtype, output_dtype,
-                   remove_online_sf_transpose, emit_block_meta, emit_hit_stats)
+                   remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
+                   emit_seed_counts)
             if key in cls.kernel_cache:
                 return
 
@@ -8338,6 +8340,17 @@ def _compile(cls,
                         cutlass.Int32, (sym_B, cute.sym_int()),
                         stride_order=(1, 0),
                         assumed_align=16)
+            seed_thr_fake = None
+            seed_counts_fake = None
+            if emit_seed_counts:
+                seed_thr_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Float32, (cute.sym_int(), 3),
+                    stride_order=(1, 0),
+                    assumed_align=4)
+                seed_counts_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Int32, (cute.sym_int(), 3),
+                    stride_order=(1, 0),
+                    assumed_align=4)
 
             fake_stream = cute.runtime.make_fake_stream(
                 use_tvm_ffi_env_stream=True)
@@ -8355,6 +8368,7 @@ def _compile(cls,
                 remove_online_sf_transpose=remove_online_sf_transpose,
                 emit_block_meta=emit_block_meta,
                 emit_hit_stats=emit_hit_stats,
+                emit_seed_counts=emit_seed_counts,
             )
 
             compiled = cute.compile(
@@ -8373,6 +8387,8 @@ def _compile(cls,
                 hit_stats_fake,
                 hit_bitmap_fake,
                 fake_stream,
+                seed_thr=seed_thr_fake,
+                seed_counts=seed_counts_fake,
                 options="--enable-tvm-ffi",
             )
             cls.kernel_cache[key] = compiled
@@ -8398,6 +8414,9 @@ def forward(
             hit_bitmap: Optional[torch.Tensor] = None,
             block_max_out: Optional[torch.Tensor] = None,
             hit_stats_out: Optional[torch.Tensor] = None,
+            emit_seed_counts: bool = False,
+            seed_thr: Optional[torch.Tensor] = None,
+            seed_counts_out: Optional[torch.Tensor] = None,
         ) -> torch.Tensor:
             """Execute FP4 paged MQA logits kernel.
 
@@ -8510,10 +8529,37 @@ def forward(
                 assert (block_max_out.shape == (B * next_n, nrec)
                         and block_max_out.is_contiguous())
 
+            if emit_seed_counts:
+                assert emit_block_meta, (
+                    "emit_seed_counts requires emit_block_meta")
+                # Per-row seed-count emission: 3 thresholds per row (fp32,
+                # post-conversion value domain), counts accumulated with
+                # red.global.add.s32 — the caller must zero seed_counts_out
+                # each step.
+                assert (
+                    seed_thr is not None and seed_thr.dtype == torch.float32
+                    and seed_thr.is_cuda and seed_thr.is_contiguous()
+                    and seed_thr.shape == (B * next_n, 3)
+                ), (f"emit_seed_counts requires seed_thr fp32 "
+                    f"[{B * next_n}, 3]; got "
+                    f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}"
+                    )
+                assert (seed_counts_out is not None
+                        and seed_counts_out.dtype == torch.int32
+                        and seed_counts_out.is_cuda
+                        and seed_counts_out.is_contiguous()
+                        and seed_counts_out.shape == (B * next_n, 3)), (
+                            "emit_seed_counts requires a caller-zeroed "
+                            "seed_counts_out int32 [B*next_n, 3]")
+            else:
+                seed_thr = None
+                seed_counts_out = None
+
             # Compile if needed (fake tensors, no real data required)
             key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
                    num_epi_subtiles, epi_dtype, output_dtype,
-                   remove_online_sf_transpose, emit_block_meta, emit_hit_stats)
+                   remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
+                   emit_seed_counts)
             if key not in cls.kernel_cache:
                 cls._compile(
                     compute_block_kv,
@@ -8527,18 +8573,20 @@ def forward(
                     output_dtype,
                     remove_online_sf_transpose=remove_online_sf_transpose,
                     emit_block_meta=emit_block_meta,
-                    emit_hit_stats=emit_hit_stats)
+                    emit_hit_stats=emit_hit_stats,
+                    emit_seed_counts=emit_seed_counts)
             compiled = cls.kernel_cache[key]
 
             # TVM FFI: pass raw tensors, no dlpack/stream needed
             if emit_block_meta:
                 compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                          context_lens, schedule_meta, num_phys_blocks, B,
-                         block_max_out, hit_stats_out, hit_bitmap)
+                         block_max_out, hit_stats_out, hit_bitmap, seed_thr,
+                         seed_counts_out)
                 return logits, block_max_out, hit_stats_out
             compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                      context_lens, schedule_meta, num_phys_blocks, B, None,
-                     None, None)
+                     None, None, None, None)
             return logits
 
     @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits",
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index dfbdc00dbed5..a9486671e478 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -124,6 +124,18 @@ def _red_global_fmax_ordered(addr_i64, fval, *, loc=None, ip=None):
 
 
 @dsl_user_op
+def _red_global_add_s32(addr_i64, ival, *, loc=None, ip=None):
+    llvm.inline_asm(
+        None,
+        [addr_i64.ir_value(loc=loc, ip=ip), ival.ir_value(loc=loc, ip=ip)],
+        "red.global.add.s32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
 def _red_global_add_f32(addr_i64, fval, *, loc=None, ip=None):
     llvm.inline_asm(
         None,
@@ -438,6 +450,7 @@ def __init__(
         use_batched_store: bool = True,
         emit_block_meta: bool = False,
         emit_hit_stats: bool = True,
+        emit_seed_counts: bool = False,
     ):
         # Static FP4 invariants — see plan Sanity checklist.
         assert num_heads == 64, "FP4 kernel hardcodes num_heads=64 for TMEM/SMEM budget"
@@ -514,6 +527,18 @@ def __init__(
         # sufficient for GVR block-skip without the fused Phase 1.
         self.emit_block_meta = emit_block_meta
         self.emit_hit_stats = emit_hit_stats
+        # emit_seed_counts (L1 of the epilogue suite; requires
+        # emit_block_meta): per row, count stored logits >= each of the
+        # T=3 caller-provided thresholds (cross-step seed thresholds from
+        # the GVR xstate). Lane-local accumulation across tiles + one
+        # warp-redux + lane-0 red.global.add.s32 per threshold per row
+        # transition — the hit-stats cost profile. Counts are computed on
+        # the POST-conversION value over valid positions only, matching
+        # the top-k consumer's verification semantics
+        # (workspace/epilogue_topk_interface.md).
+        if emit_seed_counts and not emit_block_meta:
+            raise ValueError("emit_seed_counts requires emit_block_meta")
+        self.emit_seed_counts = emit_seed_counts
         # epi_bytes covers fp16 and bf16 (FP8 only handled fp16).
         self.epi_bytes = 2 if epi_dtype in (cutlass.Float16, cutlass.BFloat16) else 4
         # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
@@ -738,6 +763,8 @@ def __call__(
         hit_stats: cute.Tensor,  # [num_rows, 4] fp32 hit aggregate (or None unless emit_hit_stats)
         hit_bitmap: cute.Tensor,  # [batch, nb_pad*4] int32 (or None unless emit_block_meta)
         stream: cuda.CUstream,
+        seed_thr: cute.Tensor = None,  # [num_rows, 3] fp32 (emit_seed_counts)
+        seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 out, caller-zeroed
     ):
         # Derive KV data and SF views from the fused uint8 buffer.
         # Fused layout per phys block: [data half_head_dim*phys_block_kv bytes]
@@ -939,6 +966,8 @@ class SharedStorage:
             block_max,
             hit_stats,
             hit_bitmap,
+            seed_thr,
+            seed_counts,
             self.cluster_layout_vmnk,
             self.a_smem_layout_staged,
             self.b_smem_layout_staged,
@@ -999,6 +1028,24 @@ def _flush_hit_agg(
             hacc_sum[t] = cutlass.Float32(0.0)
             hacc_cnt[t] = cutlass.Int32(0)
 
+    @cute.jit
+    def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane):
+        """Warp-redux the lane-local seed counters and fire one lane-0
+        red.global.add.s32 per (t, threshold). Caller zero-initializes
+        mSeedCounts each step; cross-CTA totals accumulate atomically."""
+        next_n = cutlass.const_expr(self.next_n)
+        base_addr = mSeedCounts.iterator.toint()
+        for t in cutlass.range_constexpr(next_n):
+            for j in cutlass.range_constexpr(3):
+                w_cnt = cute.arch.warp_redux_sync(scnt[t * 3 + j], "add")
+                if meta_lane == cutlass.Int32(0):
+                    row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t)
+                    addr = base_addr + (
+                        cutlass.Int64(row) * cutlass.Int64(3) + cutlass.Int64(j)
+                    ) * cutlass.Int64(4)
+                    _red_global_add_s32(addr, w_cnt)
+                scnt[t * 3 + j] = cutlass.Int32(0)
+
     @cute.kernel
     def kernel(
         self,
@@ -1021,6 +1068,8 @@ def kernel(
         mBlockMax: cute.Tensor,  # [num_rows, nb_pad*4] fp32 warp-partials (or None)
         mHitAgg: cute.Tensor,  # [num_rows, 4] fp32 {enc_min, enc_max, sum, cnt} (or None)
         mHitBitmap: cute.Tensor,  # [batch, nb_pad*4] int32 (or None)
+        mSeedThr: cute.Tensor,  # [num_rows, 3] fp32 seed thresholds (or None)
+        mSeedCounts: cute.Tensor,  # [num_rows, 3] int32 counts out (or None)
         cluster_layout_vmnk: cute.Layout,
         a_smem_layout_staged: cute.ComposedLayout,
         b_smem_layout_staged: cute.ComposedLayout,
@@ -2005,6 +2054,9 @@ def kernel(
                         # Hit accumulators + bitmap word add ~6 more live
                         # registers across the tile loop.
                         MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
+                    if cutlass.const_expr(self.emit_seed_counts):
+                        # 3 thresholds + 3 counters per t.
+                        MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the
@@ -2019,6 +2071,12 @@ def kernel(
                     ctx_cur = cutlass.Int32(0)
                     meta_warp = local_tidx // 32
                     meta_lane = local_tidx % 32
+                    if cutlass.const_expr(self.emit_seed_counts):
+                        sthr = cute.make_fragment(next_n * 3, cutlass.Float32)
+                        scnt = cute.make_fragment(next_n * 3, cutlass.Int32)
+                        for _i in cutlass.range_constexpr(next_n * 3):
+                            sthr[_i] = cutlass.Float32(_META_FLT_MAX)
+                            scnt[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_hit_stats):
                         # Per-lane hit accumulators, carried across ALL
                         # tiles of the same q and flushed once per
@@ -2083,6 +2141,12 @@ def kernel(
                             # Compressed-space context len; the meta valid
                             # mask (kv_pos < ctx_cur) keeps GEMM garbage in
                             # the aligned padding region out of block_max.
+                            if cutlass.const_expr(self.emit_seed_counts):
+                                if q_idx_old < batch_size:
+                                    self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                                for _t in cutlass.range_constexpr(next_n):
+                                    for _j in cutlass.range_constexpr(3):
+                                        sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 0 (kv_idx + 0)
@@ -2312,6 +2376,15 @@ def kernel(
                                 out_row_m = q_idx * next_n + t
                                 rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp
                                 mBlockMax[(out_row_m, rec_m)] = r_bmax
+                            if cutlass.const_expr(self.emit_seed_counts):
+                                # Seed-count accumulation: branchless 0/1
+                                # adds on the post-conversion value; the
+                                # valid mask keeps aligned-padding garbage
+                                # out (same contract as block_max).
+                                valid_i1 = cutlass.Int32(meta_valid)
+                                for _j in cutlass.range_constexpr(3):
+                                    ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j])
+                                    scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1)
                             if cutlass.const_expr(self.emit_hit_stats):
                                 # Lane-local accumulation, fully BRANCHLESS
                                 # (data-dependent `if meta_hit` compiled to
@@ -2381,6 +2454,9 @@ def kernel(
                         self._flush_hit_agg(
                             mHitAgg, q_idx, hacc_min, hacc_max, hacc_sum, hacc_cnt, meta_lane
                         )
+                if cutlass.const_expr(self.emit_seed_counts):
+                    if q_idx < batch_size:
+                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane)
 
                 # Release last Q stage (WG 0)
                 if q_idx < batch_size:
@@ -2414,6 +2490,9 @@ def kernel(
                         # Hit accumulators + bitmap word add ~6 more live
                         # registers across the tile loop.
                         MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
+                    if cutlass.const_expr(self.emit_seed_counts):
+                        # 3 thresholds + 3 counters per t.
+                        MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the
@@ -2428,6 +2507,12 @@ def kernel(
                     ctx_cur = cutlass.Int32(0)
                     meta_warp = local_tidx // 32
                     meta_lane = local_tidx % 32
+                    if cutlass.const_expr(self.emit_seed_counts):
+                        sthr = cute.make_fragment(next_n * 3, cutlass.Float32)
+                        scnt = cute.make_fragment(next_n * 3, cutlass.Int32)
+                        for _i in cutlass.range_constexpr(next_n * 3):
+                            sthr[_i] = cutlass.Float32(_META_FLT_MAX)
+                            scnt[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_hit_stats):
                         # Per-lane hit accumulators, carried across ALL
                         # tiles of the same q and flushed once per
@@ -2492,6 +2577,12 @@ def kernel(
                             # Compressed-space context len; the meta valid
                             # mask (kv_pos < ctx_cur) keeps GEMM garbage in
                             # the aligned padding region out of block_max.
+                            if cutlass.const_expr(self.emit_seed_counts):
+                                if q_idx_old < batch_size:
+                                    self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                                for _t in cutlass.range_constexpr(next_n):
+                                    for _j in cutlass.range_constexpr(3):
+                                        sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 1 (kv_idx + 1)
@@ -2715,6 +2806,15 @@ def kernel(
                                 out_row_m = q_idx * next_n + t
                                 rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp
                                 mBlockMax[(out_row_m, rec_m)] = r_bmax
+                            if cutlass.const_expr(self.emit_seed_counts):
+                                # Seed-count accumulation: branchless 0/1
+                                # adds on the post-conversion value; the
+                                # valid mask keeps aligned-padding garbage
+                                # out (same contract as block_max).
+                                valid_i1 = cutlass.Int32(meta_valid)
+                                for _j in cutlass.range_constexpr(3):
+                                    ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j])
+                                    scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1)
                             if cutlass.const_expr(self.emit_hit_stats):
                                 # Lane-local accumulation, fully BRANCHLESS
                                 # (data-dependent `if meta_hit` compiled to
@@ -2784,6 +2884,9 @@ def kernel(
                         self._flush_hit_agg(
                             mHitAgg, q_idx, hacc_min, hacc_max, hacc_sum, hacc_cnt, meta_lane
                         )
+                if cutlass.const_expr(self.emit_seed_counts):
+                    if q_idx < batch_size:
+                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane)
 
                 # Release last Q stage (WG 1)
                 if q_idx < batch_size:
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
index 40df2d060159..49a19e8d38b4 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
@@ -759,6 +759,147 @@ def test_cute_dsl_fp4_paged_mqa_logits_block_meta(
         assert bm[row, written_hi * 4 :].isnan().all(), f"stray block_max write: {tag}"
 
 
+@skip_not_sm100
+@pytest.mark.parametrize("batch_size", [1, 4])
+@pytest.mark.parametrize("next_n", [1, 2, 3])
+# 4224 = 33 blocks of 128 -> odd num_kv exercises WG1's OOB padding tile.
+@pytest.mark.parametrize("avg_ctx", [4096, 4224])
+@pytest.mark.parametrize("phys_block_kv", [64, 128])
+@pytest.mark.parametrize("fix_length", [True, False])
+def test_cute_dsl_fp4_paged_mqa_logits_seed_counts(
+    batch_size,
+    next_n,
+    avg_ctx,
+    phys_block_kv,
+    fix_length,
+):
+    """emit_seed_counts exactness: per-row counts of logits >= threshold
+    recomputed from the KERNEL'S OWN logits output (the count contract is
+    defined on post-conversion values, same as block_max). Thresholds are
+    per-row quantiles of the row's own logits so each of the 3 counters
+    lands in a different regime (loose/mid/tight)."""
+    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner
+
+    torch.manual_seed(11)
+    torch.cuda.manual_seed(11)
+    num_heads, head_dim = 64, 128
+    max_model_len = max(avg_ctx * 2, 2048)
+    device = "cuda"
+
+    if fix_length:
+        context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device)
+    else:
+        lo = max(phys_block_kv, int(0.7 * avg_ctx))
+        context_lens = torch.randint(
+            lo, int(1.3 * avg_ctx) + 1, (batch_size,), dtype=torch.int32, device=device
+        ).clamp(max=max_model_len)
+
+    num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv)
+    num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2
+    max_blocks_per_seq = int(num_blocks_per_seq.max().item())
+    block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device)
+    pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32)
+    off = 0
+    for i, n_blks in enumerate(num_blocks_per_seq.tolist()):
+        block_table[i, :n_blks] = pool[off : off + n_blks]
+        off += n_blks
+
+    q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16)
+    kv_cache = torch.randn(
+        (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16
+    )
+    weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32)
+
+    q_packed, sf_q_packed = per_token_cast_to_fp4(
+        q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True
+    )
+    q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2)
+    sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads)
+    remove_online_sf_transpose = phys_block_kv == 128
+    kv_fused, _ = kv_cache_cast_to_fp4(
+        kv_cache, remove_online_sf_transpose=remove_online_sf_transpose
+    )
+
+    DG_METADATA_BLOCK_KV = 64
+    num_sms = deep_gemm.get_num_sms()
+    schedule_meta = deep_gemm.get_paged_mqa_logits_metadata(
+        context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms
+    )
+
+    aligned_max_ctx = align(max_model_len, 256)
+    nb_pad = aligned_max_ctx // 128
+    num_rows = batch_size * next_n
+
+    # First pass without seed counts to harvest per-row logits for
+    # threshold picking (post-conversion value domain).
+    nan = float("nan")
+    block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device)
+    common = dict(
+        num_epi_subtiles=1,
+        epi_dtype=torch.float32,
+        output_dtype=torch.bfloat16,
+        remove_online_sf_transpose=remove_online_sf_transpose,
+    )
+    logits0, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward(
+        q_fp4,
+        sf_q,
+        kv_fused,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_model_len,
+        emit_block_meta=True,
+        emit_hit_stats=False,
+        block_max_out=block_max,
+        **common,
+    )
+    torch.cuda.synchronize()
+    lf0 = logits0.float()
+
+    seed_thr = torch.empty((num_rows, 3), dtype=torch.float32, device=device)
+    for row in range(num_rows):
+        ctx = int(context_lens[row // next_n].item())
+        vals = lf0[row, :ctx]
+        # Loose / mid / tight rungs; ties on exact stored values are the
+        # point (>= must count them all).
+        seed_thr[row, 0] = torch.quantile(vals, 0.10)
+        seed_thr[row, 1] = torch.quantile(vals, 0.90)
+        seed_thr[row, 2] = torch.quantile(vals, 0.998)
+
+    seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=device)
+    block_max.fill_(nan)
+    logits, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward(
+        q_fp4,
+        sf_q,
+        kv_fused,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_model_len,
+        emit_block_meta=True,
+        emit_hit_stats=False,
+        block_max_out=block_max,
+        emit_seed_counts=True,
+        seed_thr=seed_thr,
+        seed_counts_out=seed_counts,
+        **common,
+    )
+    torch.cuda.synchronize()
+
+    lf = logits.float()
+    torch.testing.assert_close(lf, lf0, atol=0.0, rtol=0.0)
+    for row in range(num_rows):
+        ctx = int(context_lens[row // next_n].item())
+        tag = f"row={row} ctx={ctx} next_n={next_n} pbk={phys_block_kv}"
+        ref = (lf[row, :ctx].unsqueeze(0) >= seed_thr[row].unsqueeze(1)).sum(-1)
+        got = seed_counts[row].to(torch.int64)
+        assert torch.equal(got.cpu(), ref.cpu().to(torch.int64)), (
+            f"seed_counts mismatch: {tag} got={got.tolist()} ref={ref.tolist()}"
+        )
+
+
 # ---------------------------------------------------------------------------
 # Benchmarking entry point (run module directly).
 # ---------------------------------------------------------------------------

From 57f154062bb2dca63b8a387ee6bf67ed5445af28 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 07:58:01 -0700
Subject: [PATCH 040/117] [None][perf] FP4 indexer: drop seed-count register
 budget cut

ncu (B=128 ctx=4096, SpeedOfLight/Occupancy/SchedulerStats) shows the
seed-count cost is epilogue ALU/issue (Compute 20->22.7%), not register
pressure: block limits and achieved occupancy are identical with or
without a -8 weight-cache cut, and the cut measured neutral-to-slower.
Kernel-level marginals: block_max +3.7%, +seed counts +3.2%,
+hit stats +10%.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../paged_mqa_logits/fp4_paged_mqa_logits.py       | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index a9486671e478..0f4880a756f8 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -2054,9 +2054,10 @@ def kernel(
                         # Hit accumulators + bitmap word add ~6 more live
                         # registers across the tile loop.
                         MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
-                    if cutlass.const_expr(self.emit_seed_counts):
-                        # 3 thresholds + 3 counters per t.
-                        MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
+                    # emit_seed_counts: no budget cut — ncu shows the cost
+                    # is epilogue ALU/issue, not registers (occupancy and
+                    # block limits identical with or without a -8 cut; the
+                    # cut itself measured neutral-to-slower).
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the
@@ -2490,9 +2491,10 @@ def kernel(
                         # Hit accumulators + bitmap word add ~6 more live
                         # registers across the tile loop.
                         MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
-                    if cutlass.const_expr(self.emit_seed_counts):
-                        # 3 thresholds + 3 counters per t.
-                        MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
+                    # emit_seed_counts: no budget cut — ncu shows the cost
+                    # is epilogue ALU/issue, not registers (occupancy and
+                    # block limits identical with or without a -8 cut; the
+                    # cut itself measured neutral-to-slower).
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the

From b7f4de364d4ec8144ff223f6087af29645f54a21 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 20 Jul 2026 09:40:36 -0700
Subject: [PATCH 041/117] [None][perf] FP4 indexer: L2 candidate pre-collect in
 the GEMM epilogue

Unordered (value, index) pair collection of all positions >= the t_0
seed threshold, off by default (emit_cand). Per-(warp, t) claim
windows: a refill claims (hits + 8) slots with one returning global
atomic; window consumption is latency-free (a per-hit atomic
round-trip measured 2x kernel time at 131k ctx). Window tails are
sentinel-filled (index = -1) at q-transition/loop end, so claimed
over-approximates the true count; counts[r][0] stays exact and
claimed >= K certifies candidate coverage of the true top-K. No-hit
iterations exit on one compare against the already-computed r_bmax
32-position bound.

Exactness: 32/32 unit combos (exact live set + sentinel accounting
when it fits; overflow -> void with valid unique entries). Cost (ncu,
kernel-level, realistic 0.8% density at 131k ctx): +27% marginal over
P0+L1 -- consumers enable it only where the top-k saving exceeds the
GEMM cost (decision waterfall).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  59 +++-
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 281 +++++++++++++++++-
 .../test_cute_dsl_fp4_paged_mqa_logits.py     | 165 ++++++++++
 3 files changed, 498 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index a419613fbf76..0a01f7b71e93 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -8253,12 +8253,14 @@ def _compile(cls,
                      remove_online_sf_transpose=False,
                      emit_block_meta=False,
                      emit_hit_stats=True,
-                     emit_seed_counts=False):
+                     emit_seed_counts=False,
+                     emit_cand=False,
+                     cand_cap=5120):
             """Compile kernel using fake tensors + TVM FFI."""
             key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
                    num_sms, num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
-                   emit_seed_counts)
+                   emit_seed_counts, emit_cand, cand_cap)
             if key in cls.kernel_cache:
                 return
 
@@ -8340,6 +8342,17 @@ def _compile(cls,
                         cutlass.Int32, (sym_B, cute.sym_int()),
                         stride_order=(1, 0),
                         assumed_align=16)
+            cand_fake = None
+            cand_ctl_fake = None
+            if emit_cand:
+                cand_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Int32, (cute.sym_int(), cand_cap * 2),
+                    stride_order=(1, 0),
+                    assumed_align=8)
+                cand_ctl_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Int32, (cute.sym_int(), 2),
+                    stride_order=(1, 0),
+                    assumed_align=8)
             seed_thr_fake = None
             seed_counts_fake = None
             if emit_seed_counts:
@@ -8369,6 +8382,8 @@ def _compile(cls,
                 emit_block_meta=emit_block_meta,
                 emit_hit_stats=emit_hit_stats,
                 emit_seed_counts=emit_seed_counts,
+                emit_cand=emit_cand,
+                cand_cap=cand_cap,
             )
 
             compiled = cute.compile(
@@ -8389,6 +8404,8 @@ def _compile(cls,
                 fake_stream,
                 seed_thr=seed_thr_fake,
                 seed_counts=seed_counts_fake,
+                cand=cand_fake,
+                cand_ctl=cand_ctl_fake,
                 options="--enable-tvm-ffi",
             )
             cls.kernel_cache[key] = compiled
@@ -8417,6 +8434,9 @@ def forward(
             emit_seed_counts: bool = False,
             seed_thr: Optional[torch.Tensor] = None,
             seed_counts_out: Optional[torch.Tensor] = None,
+            emit_cand: bool = False,
+            cand_out: Optional[torch.Tensor] = None,
+            cand_ctl_out: Optional[torch.Tensor] = None,
         ) -> torch.Tensor:
             """Execute FP4 paged MQA logits kernel.
 
@@ -8555,11 +8575,36 @@ def forward(
                 seed_thr = None
                 seed_counts_out = None
 
+            cand_cap = 0
+            if emit_cand:
+                assert emit_seed_counts, (
+                    "emit_cand requires emit_seed_counts (t_0 threshold)")
+                # Unordered (value, index) pair scatter; the caller zeroes
+                # cand_ctl_out {claimed, void} each step. cand_out is
+                # [B*next_n, CAP*2] int32 (fp32 bits in even words).
+                assert (
+                    cand_out is not None and cand_out.dtype == torch.int32
+                    and cand_out.is_cuda and cand_out.is_contiguous()
+                    and cand_out.dim() == 2 and cand_out.shape[0] == B * next_n
+                    and cand_out.shape[1] % 2 == 0 and cand_out.shape[1] > 0), (
+                        "emit_cand requires cand_out int32 [B*next_n, CAP*2]")
+                assert (cand_ctl_out is not None
+                        and cand_ctl_out.dtype == torch.int32
+                        and cand_ctl_out.is_cuda
+                        and cand_ctl_out.is_contiguous()
+                        and cand_ctl_out.shape == (B * next_n, 2)), (
+                            "emit_cand requires a caller-zeroed cand_ctl_out "
+                            "int32 [B*next_n, 2]")
+                cand_cap = cand_out.shape[1] // 2
+            else:
+                cand_out = None
+                cand_ctl_out = None
+
             # Compile if needed (fake tensors, no real data required)
             key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
                    num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
-                   emit_seed_counts)
+                   emit_seed_counts, emit_cand, cand_cap)
             if key not in cls.kernel_cache:
                 cls._compile(
                     compute_block_kv,
@@ -8574,7 +8619,9 @@ def forward(
                     remove_online_sf_transpose=remove_online_sf_transpose,
                     emit_block_meta=emit_block_meta,
                     emit_hit_stats=emit_hit_stats,
-                    emit_seed_counts=emit_seed_counts)
+                    emit_seed_counts=emit_seed_counts,
+                    emit_cand=emit_cand,
+                    cand_cap=cand_cap)
             compiled = cls.kernel_cache[key]
 
             # TVM FFI: pass raw tensors, no dlpack/stream needed
@@ -8582,11 +8629,11 @@ def forward(
                 compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                          context_lens, schedule_meta, num_phys_blocks, B,
                          block_max_out, hit_stats_out, hit_bitmap, seed_thr,
-                         seed_counts_out)
+                         seed_counts_out, cand_out, cand_ctl_out)
                 return logits, block_max_out, hit_stats_out
             compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                      context_lens, schedule_meta, num_phys_blocks, B, None,
-                     None, None, None, None)
+                     None, None, None, None, None, None)
             return logits
 
     @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits",
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index 0f4880a756f8..acb9b50e4672 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -57,7 +57,7 @@
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import llvm, vector
 from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cutlass_dsl import dsl_user_op
+from cutlass.cutlass_dsl import T, dsl_user_op
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 
 # CuTe DSL CUDA 13 validates rounding modes as string literals. The string
@@ -123,6 +123,24 @@ def _red_global_fmax_ordered(addr_i64, fval, *, loc=None, ip=None):
     )
 
 
+@dsl_user_op
+def _atom_global_add_s32(addr_i64, ival, *, loc=None, ip=None):
+    """atom.global.add.s32 returning the OLD value (warp batch-claim)."""
+    return cutlass.Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [addr_i64.ir_value(loc=loc, ip=ip), ival.ir_value(loc=loc, ip=ip)],
+            "atom.global.add.s32 $0, [$1], $2;",
+            "=r,l,r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
 @dsl_user_op
 def _red_global_add_s32(addr_i64, ival, *, loc=None, ip=None):
     llvm.inline_asm(
@@ -451,6 +469,8 @@ def __init__(
         emit_block_meta: bool = False,
         emit_hit_stats: bool = True,
         emit_seed_counts: bool = False,
+        emit_cand: bool = False,
+        cand_cap: int = 5120,
     ):
         # Static FP4 invariants — see plan Sanity checklist.
         assert num_heads == 64, "FP4 kernel hardcodes num_heads=64 for TMEM/SMEM budget"
@@ -539,6 +559,19 @@ def __init__(
         if emit_seed_counts and not emit_block_meta:
             raise ValueError("emit_seed_counts requires emit_block_meta")
         self.emit_seed_counts = emit_seed_counts
+        # emit_cand (L2 of the epilogue suite): unordered pre-collect of all
+        # (value, index) pairs >= the t_0 seed threshold via warp ballot +
+        # lane0 batch atomic claim. claimed >= K certifies the candidate
+        # set covers the true top-K (contract in epilogue_topk_interface.md).
+        if emit_cand and not emit_seed_counts:
+            raise ValueError("emit_cand requires emit_seed_counts (t_0 source)")
+        self.emit_cand = emit_cand
+        self.cand_cap = cand_cap
+        # Per-warp claim window: one atomic claims (hits + CAND_WIN) slots;
+        # subsequent hits consume the window latency-free. The unconsumed
+        # tail is sentinel-filled (idx = -1) at q-transition/loop end, so
+        # `claimed` over-approximates the true count (counts[r][0] exact).
+        self.CAND_WIN = 8
         # epi_bytes covers fp16 and bf16 (FP8 only handled fp16).
         self.epi_bytes = 2 if epi_dtype in (cutlass.Float16, cutlass.BFloat16) else 4
         # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
@@ -765,6 +798,8 @@ def __call__(
         stream: cuda.CUstream,
         seed_thr: cute.Tensor = None,  # [num_rows, 3] fp32 (emit_seed_counts)
         seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 out, caller-zeroed
+        cand: cute.Tensor = None,  # [num_rows, CAP*2] int32 {val bits, idx} pairs
+        cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 {claimed, void}, zeroed
     ):
         # Derive KV data and SF views from the fused uint8 buffer.
         # Fused layout per phys block: [data half_head_dim*phys_block_kv bytes]
@@ -968,6 +1003,8 @@ class SharedStorage:
             hit_bitmap,
             seed_thr,
             seed_counts,
+            cand,
+            cand_ctl,
             self.cluster_layout_vmnk,
             self.a_smem_layout_staged,
             self.b_smem_layout_staged,
@@ -1046,6 +1083,32 @@ def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane):
                     _red_global_add_s32(addr, w_cnt)
                 scnt[t * 3 + j] = cutlass.Int32(0)
 
+    @cute.jit
+    def _flush_cand_window(self, mCand, q_idx, cwbase, cwleft, meta_lane):
+        """Sentinel-fill the unconsumed tail of each per-(warp, t) claim
+        window (idx word = -1; consumers skip sentinels) and invalidate the
+        window. wleft <= CAND_WIN + 31 always fits one lane round."""
+        next_n = cutlass.const_expr(self.next_n)
+        CAP_C = cutlass.const_expr(self.cand_cap)
+        cand_base = mCand.iterator.toint()
+        for t in cutlass.range_constexpr(next_n):
+            if cwleft[t] > cutlass.Int32(0):
+                row_c = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t)
+                sl_f = cwbase[t] + meta_lane
+                if meta_lane < cwleft[t] and sl_f < cutlass.Int32(CAP_C):
+                    pair_f = cand_base + (
+                        cutlass.Int64(row_c) * cutlass.Int64(CAP_C) + cutlass.Int64(sl_f)
+                    ) * cutlass.Int64(8)
+                    iptr_f = cute.make_ptr(
+                        cutlass.Int32,
+                        pair_f + cutlass.Int64(4),
+                        cute.AddressSpace.gmem,
+                        assumed_align=4,
+                    )
+                    cute.make_tensor(iptr_f, cute.make_layout((1,)))[0] = cutlass.Int32(-1)
+            cwbase[t] = cutlass.Int32(0)
+            cwleft[t] = cutlass.Int32(0)
+
     @cute.kernel
     def kernel(
         self,
@@ -1070,6 +1133,8 @@ def kernel(
         mHitBitmap: cute.Tensor,  # [batch, nb_pad*4] int32 (or None)
         mSeedThr: cute.Tensor,  # [num_rows, 3] fp32 seed thresholds (or None)
         mSeedCounts: cute.Tensor,  # [num_rows, 3] int32 counts out (or None)
+        mCand: cute.Tensor,  # [num_rows, CAP*2] int32 pair scatter (or None)
+        mCandCtl: cute.Tensor,  # [num_rows, 2] int32 {claimed, void} (or None)
         cluster_layout_vmnk: cute.Layout,
         a_smem_layout_staged: cute.ComposedLayout,
         b_smem_layout_staged: cute.ComposedLayout,
@@ -2078,6 +2143,12 @@ def kernel(
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
+                    if cutlass.const_expr(self.emit_cand):
+                        cwbase = cute.make_fragment(next_n, cutlass.Int32)
+                        cwleft = cute.make_fragment(next_n, cutlass.Int32)
+                        for _i in cutlass.range_constexpr(next_n):
+                            cwbase[_i] = cutlass.Int32(0)
+                            cwleft[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_hit_stats):
                         # Per-lane hit accumulators, carried across ALL
                         # tiles of the same q and flushed once per
@@ -2145,6 +2216,11 @@ def kernel(
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
                                     self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                            if cutlass.const_expr(self.emit_cand):
+                                if q_idx_old < batch_size:
+                                    self._flush_cand_window(
+                                        mCand, q_idx_old, cwbase, cwleft, meta_lane
+                                    )
                                 for _t in cutlass.range_constexpr(next_n):
                                     for _j in cutlass.range_constexpr(3):
                                         sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
@@ -2386,6 +2462,99 @@ def kernel(
                                 for _j in cutlass.range_constexpr(3):
                                     ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j])
                                     scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1)
+                            if cutlass.const_expr(self.emit_cand):
+                                # L2 pre-collect at t_0 with per-warp claim
+                                # WINDOWS: refills claim (hits + CAND_WIN)
+                                # slots in ONE atomic; between refills the
+                                # warp consumes its window latency-free (a
+                                # per-hit atomic round-trip stalled the
+                                # epilogue 2x at 131k ctx). Unconsumed tail
+                                # is sentinel-filled on flush; counts[r][0]
+                                # stays the exact count. Gated on the
+                                # already-computed warp-uniform 32-position
+                                # bound: r_bmax < t_0 proves zero hits, so
+                                # the common no-hit iteration costs ONE fp
+                                # compare (no ballot, no select). bound >=
+                                # t_0 guarantees a nonzero ballot (exact
+                                # per-lane max, invalid -> -FLT_MAX).
+                                if r_bmax >= sthr[t * 3 + 0]:
+                                    pred_c = cutlass.Int32(0)
+                                    if meta_valid:
+                                        if f32_t >= sthr[t * 3 + 0]:
+                                            pred_c = cutlass.Int32(1)
+                                    mask_c = cute.arch.vote_ballot_sync(pred_c != cutlass.Int32(0))
+                                    row_c = q_idx * next_n + t
+                                    cnt_c = cutlass.Int32(cute.arch.popc(mask_c))
+                                    lm_c = (
+                                        cutlass.Uint32(1) << cutlass.Uint32(meta_lane)
+                                    ) - cutlass.Uint32(1)
+                                    off_c = cutlass.Int32(cute.arch.popc(mask_c & lm_c))
+                                    CAP_C = cutlass.const_expr(self.cand_cap)
+                                    cand_b = mCand.iterator.toint()
+                                    if cnt_c > cwleft[t]:
+                                        # sentinel-fill the old tail, then
+                                        # refill: one atomic per window.
+                                        sl_o = cwbase[t] + meta_lane
+                                        if meta_lane < cwleft[t] and sl_o < cutlass.Int32(CAP_C):
+                                            pair_o = cand_b + (
+                                                cutlass.Int64(row_c) * cutlass.Int64(CAP_C)
+                                                + cutlass.Int64(sl_o)
+                                            ) * cutlass.Int64(8)
+                                            iptr_o = cute.make_ptr(
+                                                cutlass.Int32,
+                                                pair_o + cutlass.Int64(4),
+                                                cute.AddressSpace.gmem,
+                                                assumed_align=4,
+                                            )
+                                            cute.make_tensor(iptr_o, cute.make_layout((1,)))[0] = (
+                                                cutlass.Int32(-1)
+                                            )
+                                        m_c = cnt_c + cutlass.Int32(self.CAND_WIN)
+                                        ctl_addr = mCandCtl.iterator.toint() + (
+                                            cutlass.Int64(row_c) * cutlass.Int64(8)
+                                        )
+                                        nb_c = cutlass.Int32(0)
+                                        if meta_lane == cutlass.Int32(0):
+                                            nb_c = _atom_global_add_s32(ctl_addr, m_c)
+                                        nb_c = cute.arch.shuffle_sync(nb_c, cutlass.Int32(0))
+                                        cwbase[t] = nb_c
+                                        cwleft[t] = m_c
+                                        # one-shot void mark on crossing CAP
+                                        if meta_lane == cutlass.Int32(0):
+                                            if nb_c + m_c > cutlass.Int32(
+                                                CAP_C
+                                            ) and nb_c <= cutlass.Int32(CAP_C):
+                                                vdptr = cute.make_ptr(
+                                                    cutlass.Int32,
+                                                    ctl_addr + cutlass.Int64(4),
+                                                    cute.AddressSpace.gmem,
+                                                    assumed_align=4,
+                                                )
+                                                cute.make_tensor(vdptr, cute.make_layout((1,)))[
+                                                    0
+                                                ] = cutlass.Int32(1)
+                                    slot_c = cwbase[t] + off_c
+                                    if pred_c != cutlass.Int32(0) and slot_c < cutlass.Int32(CAP_C):
+                                        pair_addr = cand_b + (
+                                            cutlass.Int64(row_c) * cutlass.Int64(CAP_C)
+                                            + cutlass.Int64(slot_c)
+                                        ) * cutlass.Int64(8)
+                                        vptr_c = cute.make_ptr(
+                                            cutlass.Float32,
+                                            pair_addr,
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=8,
+                                        )
+                                        cute.make_tensor(vptr_c, cute.make_layout((1,)))[0] = f32_t
+                                        iptr_c = cute.make_ptr(
+                                            cutlass.Int32,
+                                            pair_addr + cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(iptr_c, cute.make_layout((1,)))[0] = kv_pos
+                                    cwbase[t] = cwbase[t] + cnt_c
+                                    cwleft[t] = cwleft[t] - cnt_c
                             if cutlass.const_expr(self.emit_hit_stats):
                                 # Lane-local accumulation, fully BRANCHLESS
                                 # (data-dependent `if meta_hit` compiled to
@@ -2458,6 +2627,9 @@ def kernel(
                 if cutlass.const_expr(self.emit_seed_counts):
                     if q_idx < batch_size:
                         self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane)
+                if cutlass.const_expr(self.emit_cand):
+                    if q_idx < batch_size:
+                        self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane)
 
                 # Release last Q stage (WG 0)
                 if q_idx < batch_size:
@@ -2515,6 +2687,12 @@ def kernel(
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
+                    if cutlass.const_expr(self.emit_cand):
+                        cwbase = cute.make_fragment(next_n, cutlass.Int32)
+                        cwleft = cute.make_fragment(next_n, cutlass.Int32)
+                        for _i in cutlass.range_constexpr(next_n):
+                            cwbase[_i] = cutlass.Int32(0)
+                            cwleft[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_hit_stats):
                         # Per-lane hit accumulators, carried across ALL
                         # tiles of the same q and flushed once per
@@ -2582,6 +2760,11 @@ def kernel(
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
                                     self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                            if cutlass.const_expr(self.emit_cand):
+                                if q_idx_old < batch_size:
+                                    self._flush_cand_window(
+                                        mCand, q_idx_old, cwbase, cwleft, meta_lane
+                                    )
                                 for _t in cutlass.range_constexpr(next_n):
                                     for _j in cutlass.range_constexpr(3):
                                         sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
@@ -2817,6 +3000,99 @@ def kernel(
                                 for _j in cutlass.range_constexpr(3):
                                     ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j])
                                     scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1)
+                            if cutlass.const_expr(self.emit_cand):
+                                # L2 pre-collect at t_0 with per-warp claim
+                                # WINDOWS: refills claim (hits + CAND_WIN)
+                                # slots in ONE atomic; between refills the
+                                # warp consumes its window latency-free (a
+                                # per-hit atomic round-trip stalled the
+                                # epilogue 2x at 131k ctx). Unconsumed tail
+                                # is sentinel-filled on flush; counts[r][0]
+                                # stays the exact count. Gated on the
+                                # already-computed warp-uniform 32-position
+                                # bound: r_bmax < t_0 proves zero hits, so
+                                # the common no-hit iteration costs ONE fp
+                                # compare (no ballot, no select). bound >=
+                                # t_0 guarantees a nonzero ballot (exact
+                                # per-lane max, invalid -> -FLT_MAX).
+                                if r_bmax >= sthr[t * 3 + 0]:
+                                    pred_c = cutlass.Int32(0)
+                                    if meta_valid:
+                                        if f32_t >= sthr[t * 3 + 0]:
+                                            pred_c = cutlass.Int32(1)
+                                    mask_c = cute.arch.vote_ballot_sync(pred_c != cutlass.Int32(0))
+                                    row_c = q_idx * next_n + t
+                                    cnt_c = cutlass.Int32(cute.arch.popc(mask_c))
+                                    lm_c = (
+                                        cutlass.Uint32(1) << cutlass.Uint32(meta_lane)
+                                    ) - cutlass.Uint32(1)
+                                    off_c = cutlass.Int32(cute.arch.popc(mask_c & lm_c))
+                                    CAP_C = cutlass.const_expr(self.cand_cap)
+                                    cand_b = mCand.iterator.toint()
+                                    if cnt_c > cwleft[t]:
+                                        # sentinel-fill the old tail, then
+                                        # refill: one atomic per window.
+                                        sl_o = cwbase[t] + meta_lane
+                                        if meta_lane < cwleft[t] and sl_o < cutlass.Int32(CAP_C):
+                                            pair_o = cand_b + (
+                                                cutlass.Int64(row_c) * cutlass.Int64(CAP_C)
+                                                + cutlass.Int64(sl_o)
+                                            ) * cutlass.Int64(8)
+                                            iptr_o = cute.make_ptr(
+                                                cutlass.Int32,
+                                                pair_o + cutlass.Int64(4),
+                                                cute.AddressSpace.gmem,
+                                                assumed_align=4,
+                                            )
+                                            cute.make_tensor(iptr_o, cute.make_layout((1,)))[0] = (
+                                                cutlass.Int32(-1)
+                                            )
+                                        m_c = cnt_c + cutlass.Int32(self.CAND_WIN)
+                                        ctl_addr = mCandCtl.iterator.toint() + (
+                                            cutlass.Int64(row_c) * cutlass.Int64(8)
+                                        )
+                                        nb_c = cutlass.Int32(0)
+                                        if meta_lane == cutlass.Int32(0):
+                                            nb_c = _atom_global_add_s32(ctl_addr, m_c)
+                                        nb_c = cute.arch.shuffle_sync(nb_c, cutlass.Int32(0))
+                                        cwbase[t] = nb_c
+                                        cwleft[t] = m_c
+                                        # one-shot void mark on crossing CAP
+                                        if meta_lane == cutlass.Int32(0):
+                                            if nb_c + m_c > cutlass.Int32(
+                                                CAP_C
+                                            ) and nb_c <= cutlass.Int32(CAP_C):
+                                                vdptr = cute.make_ptr(
+                                                    cutlass.Int32,
+                                                    ctl_addr + cutlass.Int64(4),
+                                                    cute.AddressSpace.gmem,
+                                                    assumed_align=4,
+                                                )
+                                                cute.make_tensor(vdptr, cute.make_layout((1,)))[
+                                                    0
+                                                ] = cutlass.Int32(1)
+                                    slot_c = cwbase[t] + off_c
+                                    if pred_c != cutlass.Int32(0) and slot_c < cutlass.Int32(CAP_C):
+                                        pair_addr = cand_b + (
+                                            cutlass.Int64(row_c) * cutlass.Int64(CAP_C)
+                                            + cutlass.Int64(slot_c)
+                                        ) * cutlass.Int64(8)
+                                        vptr_c = cute.make_ptr(
+                                            cutlass.Float32,
+                                            pair_addr,
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=8,
+                                        )
+                                        cute.make_tensor(vptr_c, cute.make_layout((1,)))[0] = f32_t
+                                        iptr_c = cute.make_ptr(
+                                            cutlass.Int32,
+                                            pair_addr + cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(iptr_c, cute.make_layout((1,)))[0] = kv_pos
+                                    cwbase[t] = cwbase[t] + cnt_c
+                                    cwleft[t] = cwleft[t] - cnt_c
                             if cutlass.const_expr(self.emit_hit_stats):
                                 # Lane-local accumulation, fully BRANCHLESS
                                 # (data-dependent `if meta_hit` compiled to
@@ -2889,6 +3165,9 @@ def kernel(
                 if cutlass.const_expr(self.emit_seed_counts):
                     if q_idx < batch_size:
                         self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane)
+                if cutlass.const_expr(self.emit_cand):
+                    if q_idx < batch_size:
+                        self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane)
 
                 # Release last Q stage (WG 1)
                 if q_idx < batch_size:
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
index 49a19e8d38b4..5d4f745cf407 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
@@ -900,6 +900,171 @@ def test_cute_dsl_fp4_paged_mqa_logits_seed_counts(
         )
 
 
+@skip_not_sm100
+@pytest.mark.parametrize("batch_size", [1, 4])
+@pytest.mark.parametrize("next_n", [1, 3])
+@pytest.mark.parametrize("avg_ctx", [4096, 4224])
+@pytest.mark.parametrize("phys_block_kv", [64, 128])
+@pytest.mark.parametrize("cap_mode", ["roomy", "tight"])
+def test_cute_dsl_fp4_paged_mqa_logits_cand(
+    batch_size,
+    next_n,
+    avg_ctx,
+    phys_block_kv,
+    cap_mode,
+):
+    """emit_cand correctness: the unordered (value, index) pre-collect at
+    t_0 must contain EXACTLY the set {i < ctx : logits[r, i] >= t_0} when
+    it fits (void == 0), and degrade safely on overflow (void == 1, all
+    written slots valid + unique, claimed and counts[0] still exact)."""
+    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner
+
+    torch.manual_seed(13)
+    torch.cuda.manual_seed(13)
+    num_heads, head_dim = 64, 128
+    max_model_len = max(avg_ctx * 2, 2048)
+    device = "cuda"
+
+    context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device)
+    num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv)
+    num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2
+    max_blocks_per_seq = int(num_blocks_per_seq.max().item())
+    block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device)
+    pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32)
+    off = 0
+    for i, n_blks in enumerate(num_blocks_per_seq.tolist()):
+        block_table[i, :n_blks] = pool[off : off + n_blks]
+        off += n_blks
+
+    q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16)
+    kv_cache = torch.randn(
+        (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16
+    )
+    weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32)
+    q_packed, sf_q_packed = per_token_cast_to_fp4(
+        q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True
+    )
+    q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2)
+    sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads)
+    remove_online_sf_transpose = phys_block_kv == 128
+    kv_fused, _ = kv_cache_cast_to_fp4(
+        kv_cache, remove_online_sf_transpose=remove_online_sf_transpose
+    )
+    DG_METADATA_BLOCK_KV = 64
+    num_sms = deep_gemm.get_num_sms()
+    schedule_meta = deep_gemm.get_paged_mqa_logits_metadata(
+        context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms
+    )
+    aligned_max_ctx = align(max_model_len, 256)
+    nb_pad = aligned_max_ctx // 128
+    num_rows = batch_size * next_n
+    nan = float("nan")
+    block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device)
+    common = dict(
+        num_epi_subtiles=1,
+        epi_dtype=torch.float32,
+        output_dtype=torch.bfloat16,
+        remove_online_sf_transpose=remove_online_sf_transpose,
+    )
+    base_args = (
+        q_fp4,
+        sf_q,
+        kv_fused,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_model_len,
+    )
+
+    # Pass 1: harvest logits for threshold picking.
+    logits0, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward(
+        *base_args,
+        emit_block_meta=True,
+        emit_hit_stats=False,
+        block_max_out=block_max,
+        **common,
+    )
+    torch.cuda.synchronize()
+    lf0 = logits0.float()
+
+    seed_thr = torch.empty((num_rows, 3), dtype=torch.float32, device=device)
+    for row in range(num_rows):
+        ctx = int(context_lens[row // next_n].item())
+        vals = lf0[row, :ctx]
+        seed_thr[row, 0] = torch.quantile(vals, 0.90)  # t_0: ~10% of ctx
+        seed_thr[row, 1] = torch.quantile(vals, 0.97)
+        seed_thr[row, 2] = torch.quantile(vals, 0.998)
+
+    # Window claiming over-claims by up to ~CAND_WIN per epilogue warp
+    # touching the row (sentinel-filled tails); B=1 rows spread over many
+    # CTAs, so roomy needs slack well beyond the ~410 true hits.
+    cap = 4096 if cap_mode == "roomy" else 128
+    seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=device)
+    cand = torch.full((num_rows, cap * 2), -1, dtype=torch.int32, device=device)
+    ctl = torch.zeros((num_rows, 2), dtype=torch.int32, device=device)
+    block_max.fill_(nan)
+    logits, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward(
+        *base_args,
+        emit_block_meta=True,
+        emit_hit_stats=False,
+        block_max_out=block_max,
+        emit_seed_counts=True,
+        seed_thr=seed_thr,
+        seed_counts_out=seed_counts,
+        emit_cand=True,
+        cand_out=cand,
+        cand_ctl_out=ctl,
+        **common,
+    )
+    torch.cuda.synchronize()
+    lf = logits.float()
+    torch.testing.assert_close(lf, lf0, atol=0.0, rtol=0.0)
+
+    pairs = cand.view(num_rows, cap, 2)
+    vals_bits = pairs[..., 0]
+    idxs = pairs[..., 1]
+    vals = vals_bits.view(torch.float32)
+    for row in range(num_rows):
+        ctx = int(context_lens[row // next_n].item())
+        t0 = seed_thr[row, 0]
+        ref_mask = lf[row, :ctx] >= t0
+        ref_count = int(ref_mask.sum())
+        ref_idx = set(torch.nonzero(ref_mask, as_tuple=False).flatten().tolist())
+        tag = f"row={row} ctx={ctx} cap={cap} ref={ref_count} mode={cap_mode}"
+        claimed = int(ctl[row, 0])
+        void = int(ctl[row, 1])
+        # counts[0] is the exact count regardless of windows/overflow;
+        # claimed >= true count (sentinel-padded window tails).
+        assert int(seed_counts[row, 0]) == ref_count, f"counts0: {tag}"
+        assert claimed >= ref_count, f"claimed < true count: {tag} got={claimed}"
+        n_written = min(claimed, cap)
+        got_idx = idxs[row, :n_written].long()
+        live = got_idx >= 0
+        got_list = got_idx[live].tolist()
+        assert len(set(got_list)) == len(got_list), f"duplicate idx: {tag}"
+        assert set(got_list).issubset(ref_idx), f"non-member idx: {tag}"
+        # pair integrity on live entries: value word == stored logit bits.
+        live_idx = got_idx[live]
+        torch.testing.assert_close(
+            vals[row, :n_written][live],
+            lf[row, live_idx],
+            atol=0.0,
+            rtol=0.0,
+            msg=lambda m, tag=tag: f"pair value mismatch: {tag}\n{m}",
+        )
+        if cap_mode == "roomy":
+            assert void == 0, f"void set without overflow: {tag} claimed={claimed}"
+            assert claimed <= cap, f"claimed past cap without void: {tag}"
+            assert set(got_list) == ref_idx, f"set mismatch: {tag}"
+            # every claimed slot is live or sentinel; unclaimed tail untouched
+            assert bool((idxs[row, claimed:] == -1).all()), f"stray write: {tag}"
+            assert int(live.sum()) == ref_count, f"live count: {tag}"
+        else:
+            assert ref_count > cap, f"test setup wants overflow: {tag}"
+            assert void == 1, f"void not set on overflow: {tag}"
+
+
 # ---------------------------------------------------------------------------
 # Benchmarking entry point (run module directly).
 # ---------------------------------------------------------------------------

From 68ef232ec3c1b524b7b8adc32685bd1e911d78f6 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 08:59:52 -0700
Subject: [PATCH 042/117] [None][fix] FP4 indexer: packed seed-row emission +
 counts-only threshold load fix

Counts-only emission (emit_seed_counts without emit_cand) never loaded
the per-q thresholds: the sthr (re)load sat inside the emit_cand block
at both WG q-transition sites, so counters compared against the FLT_MAX
init and flushed zeros. Move the load under emit_seed_counts where it
belongs (L2 pre-collect commit had nested it wrong; its own tests all
run emit_cand=True which masked it).

Add seed_packed mode: one [rows, 8] fp32 seed row per the top-k
pre-packed contract - lines at cols 0..2 (kernel reads (row, j<=2)
unchanged), counts accumulate via red.global.add.f32 at cols 3..5
(exact to 2^24). Selected by passing an [rows, 8] seed_thr with
seed_counts_out=None; the same buffer then feeds the top-k launch with
no host repack. Unit test parametrized over both contracts; logits
equality now compares valid prefixes only (past ctx the output buffer
is unwritten allocator garbage and differs run-to-run).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  | 91 +++++++++++++------
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 49 +++++++---
 .../test_cute_dsl_fp4_paged_mqa_logits.py     | 26 +++++-
 3 files changed, 121 insertions(+), 45 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 0a01f7b71e93..0f0b12cae380 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -8254,13 +8254,14 @@ def _compile(cls,
                      emit_block_meta=False,
                      emit_hit_stats=True,
                      emit_seed_counts=False,
+                     seed_packed=False,
                      emit_cand=False,
                      cand_cap=5120):
             """Compile kernel using fake tensors + TVM FFI."""
             key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
                    num_sms, num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
-                   emit_seed_counts, emit_cand, cand_cap)
+                   emit_seed_counts, seed_packed, emit_cand, cand_cap)
             if key in cls.kernel_cache:
                 return
 
@@ -8356,14 +8357,28 @@ def _compile(cls,
             seed_thr_fake = None
             seed_counts_fake = None
             if emit_seed_counts:
-                seed_thr_fake = cute.runtime.make_fake_compact_tensor(
-                    cutlass.Float32, (cute.sym_int(), 3),
-                    stride_order=(1, 0),
-                    assumed_align=4)
-                seed_counts_fake = cute.runtime.make_fake_compact_tensor(
-                    cutlass.Int32, (cute.sym_int(), 3),
-                    stride_order=(1, 0),
-                    assumed_align=4)
+                if seed_packed:
+                    # single [rows, 8] fp32 packed seed row: lines at
+                    # cols 0..2 (kernel reads (row, j<=2) unchanged),
+                    # counts accumulate as fp32 at cols 3..5; the same
+                    # tensor is bound to both params.
+                    seed_thr_fake = cute.runtime.make_fake_compact_tensor(
+                        cutlass.Float32, (cute.sym_int(), 8),
+                        stride_order=(1, 0),
+                        assumed_align=4)
+                    seed_counts_fake = cute.runtime.make_fake_compact_tensor(
+                        cutlass.Float32, (cute.sym_int(), 8),
+                        stride_order=(1, 0),
+                        assumed_align=4)
+                else:
+                    seed_thr_fake = cute.runtime.make_fake_compact_tensor(
+                        cutlass.Float32, (cute.sym_int(), 3),
+                        stride_order=(1, 0),
+                        assumed_align=4)
+                    seed_counts_fake = cute.runtime.make_fake_compact_tensor(
+                        cutlass.Int32, (cute.sym_int(), 3),
+                        stride_order=(1, 0),
+                        assumed_align=4)
 
             fake_stream = cute.runtime.make_fake_stream(
                 use_tvm_ffi_env_stream=True)
@@ -8382,6 +8397,7 @@ def _compile(cls,
                 emit_block_meta=emit_block_meta,
                 emit_hit_stats=emit_hit_stats,
                 emit_seed_counts=emit_seed_counts,
+                seed_packed=seed_packed,
                 emit_cand=emit_cand,
                 cand_cap=cand_cap,
             )
@@ -8549,28 +8565,44 @@ def forward(
                 assert (block_max_out.shape == (B * next_n, nrec)
                         and block_max_out.is_contiguous())
 
+            seed_packed = False
             if emit_seed_counts:
                 assert emit_block_meta, (
                     "emit_seed_counts requires emit_block_meta")
-                # Per-row seed-count emission: 3 thresholds per row (fp32,
-                # post-conversion value domain), counts accumulated with
-                # red.global.add.s32 — the caller must zero seed_counts_out
-                # each step.
-                assert (
-                    seed_thr is not None and seed_thr.dtype == torch.float32
-                    and seed_thr.is_cuda and seed_thr.is_contiguous()
-                    and seed_thr.shape == (B * next_n, 3)
-                ), (f"emit_seed_counts requires seed_thr fp32 "
-                    f"[{B * next_n}, 3]; got "
-                    f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}"
-                    )
-                assert (seed_counts_out is not None
-                        and seed_counts_out.dtype == torch.int32
-                        and seed_counts_out.is_cuda
-                        and seed_counts_out.is_contiguous()
-                        and seed_counts_out.shape == (B * next_n, 3)), (
-                            "emit_seed_counts requires a caller-zeroed "
-                            "seed_counts_out int32 [B*next_n, 3]")
+                if seed_counts_out is None:
+                    # Packed contract: seed_thr IS the [rows, 8] fp32 seed
+                    # row (top-k pre-packed layout). Lines at cols 0..2;
+                    # counts accumulate as fp32 at cols 3..5 - the caller
+                    # zeroes cols 3..7 and writes lines each step. The
+                    # same buffer then feeds the top-k launch directly.
+                    seed_packed = True
+                    assert (
+                        seed_thr is not None and seed_thr.dtype == torch.float32
+                        and seed_thr.is_cuda and seed_thr.is_contiguous()
+                        and seed_thr.shape == (B * next_n, 8)
+                    ), (f"packed emit_seed_counts requires seed_thr fp32 "
+                        f"[{B * next_n}, 8]; got "
+                        f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}"
+                        )
+                    seed_counts_out = seed_thr
+                else:
+                    # Legacy split contract: 3 thresholds per row (fp32),
+                    # counts accumulated with red.global.add.s32 into a
+                    # caller-zeroed int32 [rows, 3].
+                    assert (
+                        seed_thr is not None and seed_thr.dtype == torch.float32
+                        and seed_thr.is_cuda and seed_thr.is_contiguous()
+                        and seed_thr.shape == (B * next_n, 3)
+                    ), (f"emit_seed_counts requires seed_thr fp32 "
+                        f"[{B * next_n}, 3]; got "
+                        f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}"
+                        )
+                    assert (seed_counts_out.dtype == torch.int32
+                            and seed_counts_out.is_cuda
+                            and seed_counts_out.is_contiguous()
+                            and seed_counts_out.shape == (B * next_n, 3)), (
+                                "emit_seed_counts requires a caller-zeroed "
+                                "seed_counts_out int32 [B*next_n, 3]")
             else:
                 seed_thr = None
                 seed_counts_out = None
@@ -8604,7 +8636,7 @@ def forward(
             key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
                    num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
-                   emit_seed_counts, emit_cand, cand_cap)
+                   emit_seed_counts, seed_packed, emit_cand, cand_cap)
             if key not in cls.kernel_cache:
                 cls._compile(
                     compute_block_kv,
@@ -8620,6 +8652,7 @@ def forward(
                     emit_block_meta=emit_block_meta,
                     emit_hit_stats=emit_hit_stats,
                     emit_seed_counts=emit_seed_counts,
+                    seed_packed=seed_packed,
                     emit_cand=emit_cand,
                     cand_cap=cand_cap)
             compiled = cls.kernel_cache[key]
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index acb9b50e4672..cd5fbbaa07bc 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -469,6 +469,7 @@ def __init__(
         emit_block_meta: bool = False,
         emit_hit_stats: bool = True,
         emit_seed_counts: bool = False,
+        seed_packed: bool = False,
         emit_cand: bool = False,
         cand_cap: int = 5120,
     ):
@@ -559,6 +560,14 @@ def __init__(
         if emit_seed_counts and not emit_block_meta:
             raise ValueError("emit_seed_counts requires emit_block_meta")
         self.emit_seed_counts = emit_seed_counts
+        # seed_packed: single [num_rows, 8] fp32 seed row per the top-k
+        # pre-packed contract - lines at cols 0..2, counts ACCUMULATED AS
+        # FLOATS at cols 3..5 (exact to 2^24; red.global.add.f32). The
+        # caller zeroes cols 3..7 and writes the lines each step; the
+        # same buffer feeds the top-k launch with no host repack.
+        if seed_packed and not emit_seed_counts:
+            raise ValueError("seed_packed requires emit_seed_counts")
+        self.seed_packed = seed_packed
         # emit_cand (L2 of the epilogue suite): unordered pre-collect of all
         # (value, index) pairs >= the t_0 seed threshold via warp ballot +
         # lane0 batch atomic claim. claimed >= K certifies the candidate
@@ -1068,8 +1077,11 @@ def _flush_hit_agg(
     @cute.jit
     def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane):
         """Warp-redux the lane-local seed counters and fire one lane-0
-        red.global.add.s32 per (t, threshold). Caller zero-initializes
-        mSeedCounts each step; cross-CTA totals accumulate atomically."""
+        red.global.add per (t, threshold). Caller zero-initializes the
+        count slots each step; cross-CTA totals accumulate atomically.
+
+        seed_packed: mSeedCounts IS the [num_rows, 8] packed seed row -
+        counts land as fp32 at cols 3..5 (exact to 2^24)."""
         next_n = cutlass.const_expr(self.next_n)
         base_addr = mSeedCounts.iterator.toint()
         for t in cutlass.range_constexpr(next_n):
@@ -1077,10 +1089,16 @@ def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane):
                 w_cnt = cute.arch.warp_redux_sync(scnt[t * 3 + j], "add")
                 if meta_lane == cutlass.Int32(0):
                     row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t)
-                    addr = base_addr + (
-                        cutlass.Int64(row) * cutlass.Int64(3) + cutlass.Int64(j)
-                    ) * cutlass.Int64(4)
-                    _red_global_add_s32(addr, w_cnt)
+                    if cutlass.const_expr(self.seed_packed):
+                        addr = base_addr + (
+                            cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(3 + j)
+                        ) * cutlass.Int64(4)
+                        _red_global_add_f32(addr, cutlass.Float32(w_cnt))
+                    else:
+                        addr = base_addr + (
+                            cutlass.Int64(row) * cutlass.Int64(3) + cutlass.Int64(j)
+                        ) * cutlass.Int64(4)
+                        _red_global_add_s32(addr, w_cnt)
                 scnt[t * 3 + j] = cutlass.Int32(0)
 
     @cute.jit
@@ -2216,14 +2234,18 @@ def kernel(
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
                                     self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                                # (re)load this q's thresholds - gated on
+                                # emit_seed_counts, NOT emit_cand: counts-
+                                # only mode needs them too (a stale
+                                # FLT_MAX default zeroes every counter)
+                                for _t in cutlass.range_constexpr(next_n):
+                                    for _j in cutlass.range_constexpr(3):
+                                        sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
                             if cutlass.const_expr(self.emit_cand):
                                 if q_idx_old < batch_size:
                                     self._flush_cand_window(
                                         mCand, q_idx_old, cwbase, cwleft, meta_lane
                                     )
-                                for _t in cutlass.range_constexpr(next_n):
-                                    for _j in cutlass.range_constexpr(3):
-                                        sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 0 (kv_idx + 0)
@@ -2760,14 +2782,17 @@ def kernel(
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
                                     self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                                # (re)load this q's thresholds - gated on
+                                # emit_seed_counts, NOT emit_cand (see the
+                                # WG0 twin above)
+                                for _t in cutlass.range_constexpr(next_n):
+                                    for _j in cutlass.range_constexpr(3):
+                                        sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
                             if cutlass.const_expr(self.emit_cand):
                                 if q_idx_old < batch_size:
                                     self._flush_cand_window(
                                         mCand, q_idx_old, cwbase, cwleft, meta_lane
                                     )
-                                for _t in cutlass.range_constexpr(next_n):
-                                    for _j in cutlass.range_constexpr(3):
-                                        sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)]
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 1 (kv_idx + 1)
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
index 5d4f745cf407..9ac8c880846a 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
@@ -766,12 +766,14 @@ def test_cute_dsl_fp4_paged_mqa_logits_block_meta(
 @pytest.mark.parametrize("avg_ctx", [4096, 4224])
 @pytest.mark.parametrize("phys_block_kv", [64, 128])
 @pytest.mark.parametrize("fix_length", [True, False])
+@pytest.mark.parametrize("packed", [False, True])
 def test_cute_dsl_fp4_paged_mqa_logits_seed_counts(
     batch_size,
     next_n,
     avg_ctx,
     phys_block_kv,
     fix_length,
+    packed,
 ):
     """emit_seed_counts exactness: per-row counts of logits >= threshold
     recomputed from the KERNEL'S OWN logits output (the count contract is
@@ -867,7 +869,15 @@ def test_cute_dsl_fp4_paged_mqa_logits_seed_counts(
         seed_thr[row, 1] = torch.quantile(vals, 0.90)
         seed_thr[row, 2] = torch.quantile(vals, 0.998)
 
-    seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=device)
+    if packed:
+        # Packed contract: one [rows, 8] fp32 seed row, lines at cols
+        # 0..2, counts accumulate as fp32 at cols 3..5 (caller zeroes).
+        seed_row = torch.zeros((num_rows, 8), dtype=torch.float32, device=device)
+        seed_row[:, 0:3] = seed_thr
+        thr_arg, counts_arg = seed_row, None
+    else:
+        seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=device)
+        thr_arg, counts_arg = seed_thr, seed_counts
     block_max.fill_(nan)
     logits, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward(
         q_fp4,
@@ -882,14 +892,22 @@ def test_cute_dsl_fp4_paged_mqa_logits_seed_counts(
         emit_hit_stats=False,
         block_max_out=block_max,
         emit_seed_counts=True,
-        seed_thr=seed_thr,
-        seed_counts_out=seed_counts,
+        seed_thr=thr_arg,
+        seed_counts_out=counts_arg,
         **common,
     )
     torch.cuda.synchronize()
 
     lf = logits.float()
-    torch.testing.assert_close(lf, lf0, atol=0.0, rtol=0.0)
+    # compare valid prefixes only: past ctx the buffer is unwritten
+    # allocator garbage and differs run-to-run
+    for row in range(num_rows):
+        ctx = int(context_lens[row // next_n].item())
+        torch.testing.assert_close(lf[row, :ctx], lf0[row, :ctx], atol=0.0, rtol=0.0)
+    if packed:
+        assert torch.equal(seed_row[:, 0:3], seed_thr), "lines clobbered"
+        assert (seed_row[:, 6:8] == 0).all(), "stray write past counts"
+        seed_counts = seed_row[:, 3:6].to(torch.int32)
     for row in range(num_rows):
         ctx = int(context_lens[row // next_n].item())
         tag = f"row={row} ctx={ctx} next_n={next_n} pbk={phys_block_kv}"

From 5d1812d63f702c0ea2ae9863a1752632a5872c5a Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 09:05:37 -0700
Subject: [PATCH 043/117] [None][chore] GVR: device-residency stamps in
 P4_TAIL_DBG

Add a row-phase entry clock under GVR_P4_TAIL_DBG and publish device
total (entry->publish) and true in-kernel prologue to xstate[1]/[2].
Wall minus device total isolates host launch/marshalling from kernel
work - warm back-to-back walls previously mis-attributed ~9us of host
gap as kernel prologue on mid-range cells.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py   | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 7dad3d33a8f1..9829d4215802 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -5404,6 +5404,9 @@ def _run_phases(
         # CTA-uniform and the dynamic branches below stay convergent.
         ck0 = cutlass.Int64(0)
         ck1 = cutlass.Int64(0)
+        ckE = cutlass.Int64(0)
+        if cutlass.const_expr(_P4_TAIL_DBG):
+            ckE = cute.arch.clock64()  # row-phase entry (device-residency ref)
         ext_row = cutlass.Int32(0)
         if cutlass.const_expr(self.use_ext_counts):
             # line validity mirrors ext_rungs: ALL THREE lines finite and
@@ -6794,6 +6797,11 @@ def _run_phases(
                         xstate_row[2] = anch_pub
                         if cutlass.const_expr(_P4_TAIL_DBG):
                             ck3 = cute.arch.clock64()
+                            # [1] device total (entry->publish), [2] true
+                            # in-kernel prologue (entry->walk start): wall
+                            # minus [1] = host/launch, NOT kernel work
+                            xstate_row[1] = cutlass.Float32(cutlass.Int32(ck3 - ckE))
+                            xstate_row[2] = cutlass.Float32(cutlass.Int32(ck0 - ckE))
                             xstate_row[4] = cutlass.Float32(cutlass.Int32(ck1 - ck0))  # walk+flags
                             xstate_row[5] = cutlass.Float32(cutlass.Int32(ck2 - ck1))  # P2/P3 gap
                             xstate_row[6] = cutlass.Float32(cutlass.Int32(ck3 - ck2))  # Phase 4

From aec6c665fe6d1c331be0447be54c328db1cb9c28 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 09:27:52 -0700
Subject: [PATCH 044/117] [None][perf] GVR P4: eager position fetch + warp0 bin
 search

Two independent cuts on the rank-scatter critical path (P4 was 80% of
device time on latency-bound mid-range cells):

1. Eager position fetch (ext-cand list rows): the take walk already
   computes each candidate's segmented slot - fetch the idx column in
   the SAME 4-fragment ILP batch as the value read and store TRUE
   positions in smem_vals. The post-P4 slot swap (a 3-round gmem
   read-gather-write over K outputs, serialized cold late in the
   kernel) disappears; self_scan keeps the old swap (its positions
   live in its own column). Extra idx bytes are streamed alongside the
   value column the walk reads anyway.

2. warp0 bin search (coarse + fine): the 3-step block searches spent
   ~0.2us per block barrier staging warp partials. One warp now owns
   the whole search: per-lane descending chunk sums, an inclusive
   shuffle scan, ballot for the straddling lane (prefix monotonicity:
   target = 32 - popc), then a lane-serial resolve. Three barriers
   collapse into one on each search; the count resets fold into the
   resolving lane.

Cold-protocol kernel-only (pro mid cells, B200): wf 8k/16k/32k/64k
-14..-17.5% (8.56->7.35us at 8k B1; speedup vs PR tip 1.19->1.38,
1.44->1.75 at 64k B1); P4_SUB fine 1.2->0.78us; va shares the bin
search win. unified_smoke 3 cells x 10 modes exact incl. flash-512k
long rows and list+block-skip.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 233 +++++++++---------
 1 file changed, 117 insertions(+), 116 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 9829d4215802..420df20c4088 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3479,56 +3479,49 @@ def phase4_rank_scatter(
             cute.arch.barrier()
             if cutlass.const_expr(_P4_SUB_DBG):
                 sc2 = cute.arch.clock64()
-            # ---- 3-step high→low bin search → straddling bin b* + rank_above ----
-            warp_bin_sum = cutlass.Int32(0)
-            for jb in cutlass.range_constexpr(bins_per_warp):
-                bidx_s = (
-                    cutlass.Int32(kBins - 1)
-                    - warp_id * cutlass.Int32(bins_per_warp)
-                    - cutlass.Int32(jb)
-                )
-                warp_bin_sum = warp_bin_sum + smem_hist[bidx_s]
-            if lane == cutlass.Int32(0):
-                smem_wcnt[warp_id] = warp_bin_sum
-            cute.arch.barrier()
-            if tidx == cutlass.Int32(0):
-                cum = cutlass.Int32(0)
-                tw = cutlass.Int32(num_warps - 1)
-                found = cutlass.Int32(0)
-                for w2 in cutlass.range_constexpr(self.num_warps):
-                    cum = cum + smem_wcnt[w2]
-                    if cum >= cutlass.Int32(kK) and found == cutlass.Int32(0):
-                        tw = cutlass.Int32(w2)
-                        found = cutlass.Int32(1)
-                cum2 = cutlass.Int32(0)
-                for w3 in cutlass.range_constexpr(self.num_warps):
-                    if cutlass.Int32(w3) < tw:
-                        cum2 = cum2 + smem_wcnt[w3]
-                s_iscalars[2] = cum2  # prefix-count before target warp
-                s_iscalars[3] = tw
-            cute.arch.barrier()
-            target_warp = s_iscalars[3]
-            if warp_id == target_warp and lane == cutlass.Int32(0):
-                base_cum = s_iscalars[2]
-                b_star = cutlass.Int32(kBins - 1)
-                rank_above = base_cum
-                set_d = cutlass.Int32(0)
-                for jb2 in cutlass.range_constexpr(bins_per_warp):
-                    bidx2 = (
+            # ---- warp0 high→low bin search → straddling bin b* + rank_above ----
+            # Single-warp shuffle scan replaces the 3-step block search:
+            # lane L owns the L-th DESCENDING chunk of bins, an inclusive
+            # warp scan finds the straddling lane (prefixes are monotone,
+            # so hit lanes form a suffix: target = 32 - popc(ballot)),
+            # and that lane resolves b* serially in its own chunk. Two of
+            # the three block barriers (~0.2us each) disappear.
+            if warp_id == cutlass.Int32(0):
+                lsum = cutlass.Int32(0)
+                for jb in cutlass.range_constexpr(bins_per_warp):
+                    bidx_s = (
                         cutlass.Int32(kBins - 1)
-                        - target_warp * cutlass.Int32(bins_per_warp)
-                        - cutlass.Int32(jb2)
+                        - lane * cutlass.Int32(bins_per_warp)
+                        - cutlass.Int32(jb)
                     )
-                    ra_before = base_cum
-                    base_cum = base_cum + smem_hist[bidx2]
-                    if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0):
-                        b_star = bidx2
-                        rank_above = ra_before  # count in bins strictly above b*
-                        set_d = cutlass.Int32(1)
-                s_iscalars[2] = rank_above
-                s_iscalars[3] = b_star
-                s_iscalars[4] = cutlass.Int32(0)  # cnt_above
-                s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
+                    lsum = lsum + smem_hist[bidx_s]
+                pref = warp_scan(lsum, tidx, lane, 32)
+                hit0 = cutlass.Int32(0)
+                if pref >= cutlass.Int32(kK):
+                    hit0 = cutlass.Int32(1)
+                mhit = cute.arch.vote_ballot_sync(hit0 != cutlass.Int32(0))
+                tl0 = cutlass.Int32(32) - cutlass.Int32(cute.arch.popc(mhit))
+                if lane == tl0:
+                    base_cum = pref - lsum  # exclusive prefix before chunk
+                    b_star = cutlass.Int32(kBins - 1)
+                    rank_above = base_cum
+                    set_d = cutlass.Int32(0)
+                    for jb2 in cutlass.range_constexpr(bins_per_warp):
+                        bidx2 = (
+                            cutlass.Int32(kBins - 1)
+                            - lane * cutlass.Int32(bins_per_warp)
+                            - cutlass.Int32(jb2)
+                        )
+                        ra_before = base_cum
+                        base_cum = base_cum + smem_hist[bidx2]
+                        if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0):
+                            b_star = bidx2
+                            rank_above = ra_before  # count strictly above b*
+                            set_d = cutlass.Int32(1)
+                    s_iscalars[2] = rank_above
+                    s_iscalars[3] = b_star
+                    s_iscalars[4] = cutlass.Int32(0)  # cnt_above
+                    s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
             cute.arch.barrier()
             if cutlass.const_expr(_P4_SUB_DBG):
                 sc3 = cute.arch.clock64()
@@ -3569,66 +3562,50 @@ def phase4_rank_scatter(
                         atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1))
                     ifb = ifb + cutlass.Int32(num_threads)
                 cute.arch.barrier()
-                # fine 3-step search seeded at rank_above (over fbins bins)
-                fws = cutlass.Int32(0)
-                for jbf in cutlass.range_constexpr(fbpw):
-                    bif = (
-                        cutlass.Int32(fbins - 1)
-                        - warp_id * cutlass.Int32(fbpw)
-                        - cutlass.Int32(jbf)
-                    )
-                    fws = fws + smem_hist[bif]
-                if lane == cutlass.Int32(0):
-                    smem_wcnt[warp_id] = fws
-                cute.arch.barrier()
-                if tidx == cutlass.Int32(0):
-                    cumf = rank_above
-                    twf = cutlass.Int32(num_warps - 1)
-                    fnd = cutlass.Int32(0)
-                    for w2 in cutlass.range_constexpr(self.num_warps):
-                        cumf = cumf + smem_wcnt[w2]
-                        if cumf >= cutlass.Int32(kK) and fnd == cutlass.Int32(0):
-                            twf = cutlass.Int32(w2)
-                            fnd = cutlass.Int32(1)
-                    pre = rank_above
-                    for w3 in cutlass.range_constexpr(self.num_warps):
-                        if cutlass.Int32(w3) < twf:
-                            pre = pre + smem_wcnt[w3]
-                    # Stage prefix/target-warp metadata in spare s_iscalars
-                    # slots, NOT smem_hist[0]/[1]: the last fine warp's reverse
-                    # scan below walks fine bins down to 0/1, so reusing those
-                    # histogram bins as scratch would corrupt sb_star/ra_fine
-                    # when twf2 == num_warps-1. Slots [4]/[1] are dead here
-                    # (re-zeroed at the cnt_above/cnt_strad reset below).
-                    s_iscalars[4] = pre  # prefix into target fine warp
-                    s_iscalars[1] = twf  # target fine warp
-                cute.arch.barrier()
-                pre_f = s_iscalars[4]
-                twf2 = s_iscalars[1]
-                if warp_id == twf2 and lane == cutlass.Int32(0):
-                    base_f = pre_f
-                    sb_star = cutlass.Int32(fbins - 1)
-                    ra_fine = base_f
-                    sd = cutlass.Int32(0)
-                    for jb3 in cutlass.range_constexpr(fbpw):
-                        sbi = (
+                # fine warp0 search seeded at rank_above (over fbins bins);
+                # same single-warp shuffle-scan shape as the coarse search
+                # above - three block barriers collapse into one. sb*/ra
+                # go via smem_hist[2]/[3] (fine bins walked never reach
+                # slots 2/3 only when the target chunk excludes them, so
+                # publish AFTER the serial walk like before - the walk
+                # reads bins, the publish overwrites scratch slots).
+                if warp_id == cutlass.Int32(0):
+                    fls = cutlass.Int32(0)
+                    for jbf in cutlass.range_constexpr(fbpw):
+                        bif = (
                             cutlass.Int32(fbins - 1)
-                            - twf2 * cutlass.Int32(fbpw)
-                            - cutlass.Int32(jb3)
+                            - lane * cutlass.Int32(fbpw)
+                            - cutlass.Int32(jbf)
                         )
-                        ra_b = base_f
-                        base_f = base_f + smem_hist[sbi]
-                        if base_f >= cutlass.Int32(kK) and sd == cutlass.Int32(0):
-                            sb_star = sbi
-                            ra_fine = ra_b
-                            sd = cutlass.Int32(1)
-                    smem_hist[2] = sb_star
-                    smem_hist[3] = ra_fine
-                cute.arch.barrier()
-                if tidx == cutlass.Int32(0):
-                    s_iscalars[4] = cutlass.Int32(0)  # cnt_above
-                    s_iscalars[0] = cutlass.Int32(0)  # cnt_mid (b*, sub>sb*)
-                    s_iscalars[1] = cutlass.Int32(0)  # cnt_strad (b*, sub==sb*)
+                        fls = fls + smem_hist[bif]
+                    fpref = warp_scan(fls, tidx, lane, 32) + rank_above
+                    fhit = cutlass.Int32(0)
+                    if fpref >= cutlass.Int32(kK):
+                        fhit = cutlass.Int32(1)
+                    fm = cute.arch.vote_ballot_sync(fhit != cutlass.Int32(0))
+                    ftl = cutlass.Int32(32) - cutlass.Int32(cute.arch.popc(fm))
+                    if lane == ftl:
+                        base_f = fpref - fls
+                        sb_star = cutlass.Int32(fbins - 1)
+                        ra_fine = base_f
+                        sd = cutlass.Int32(0)
+                        for jb3 in cutlass.range_constexpr(fbpw):
+                            sbi = (
+                                cutlass.Int32(fbins - 1)
+                                - lane * cutlass.Int32(fbpw)
+                                - cutlass.Int32(jb3)
+                            )
+                            ra_b = base_f
+                            base_f = base_f + smem_hist[sbi]
+                            if base_f >= cutlass.Int32(kK) and sd == cutlass.Int32(0):
+                                sb_star = sbi
+                                ra_fine = ra_b
+                                sd = cutlass.Int32(1)
+                        smem_hist[2] = sb_star
+                        smem_hist[3] = ra_fine
+                        s_iscalars[4] = cutlass.Int32(0)  # cnt_above
+                        s_iscalars[0] = cutlass.Int32(0)  # cnt_mid (b*, sub>sb*)
+                        s_iscalars[1] = cutlass.Int32(0)  # cnt_strad (b*, sub==sb*)
                 cute.arch.barrier()
                 if cutlass.const_expr(_P4_SUB_DBG):
                     sc4 = cute.arch.clock64()
@@ -5720,8 +5697,10 @@ def _run_phases(
                 line_cut = cutlass.Int32(0)
                 anch_t = cutlass.Float32(0.0)
                 vbase = cutlass.Int64(0)
+                ibase = cutlass.Int64(0)
                 if cutlass.const_expr(True):
                     vbase = cand_vals_row.iterator.toint()
+                    ibase = cand_idx_row.iterator.toint()
                 if usable == cutlass.Int32(1):
                     # cut = tightest line in [K, B*]; anchor = loosest.
                     if n2_c >= kK_l and n2_c <= bs_l:
@@ -5947,7 +5926,19 @@ def _run_phases(
                                     smem_keys[j_c] = cute.make_tensor(vp_c, cute.make_layout((1,)))[
                                         0
                                     ]
-                                    smem_vals[j_c] = src_c
+                                    # eager position fetch: the idx column
+                                    # rides the same ILP batch as the value
+                                    # read, so vals hold TRUE positions and
+                                    # the post-P4 slot swap disappears
+                                    ip_c = cute.make_ptr(
+                                        cutlass.Int32,
+                                        ibase + cutlass.Int64(src_c) * cutlass.Int64(4),
+                                        cute.AddressSpace.gmem,
+                                        assumed_align=4,
+                                    )
+                                    smem_vals[j_c] = cute.make_tensor(ip_c, cute.make_layout((1,)))[
+                                        0
+                                    ]
                             i_c = i_c + cutlass.Int32(4 * num_threads)
                         cute.arch.barrier()
                     if line_cut == cutlass.Int32(0):
@@ -5962,7 +5953,7 @@ def _run_phases(
                             for _ju in cutlass.range_constexpr(4):
                                 j_c = i_c + cutlass.Int32(_ju * num_threads)
                                 pval = cutlass.Float32(self.NEG_FLT_MAX)
-                                src_c = cutlass.Int32(0)
+                                pidx = cutlass.Int32(-1)
                                 keep = cutlass.Int32(0)
                                 if j_c < total_l:
                                     src_c = j_c
@@ -5977,10 +5968,20 @@ def _run_phases(
                                         assumed_align=4,
                                     )
                                     pval = cute.make_tensor(vp_c, cute.make_layout((1,)))[0]
+                                    # eager position fetch (unconditional:
+                                    # keeps the load independent of the
+                                    # value compare, same ILP batch)
+                                    ip_c = cute.make_ptr(
+                                        cutlass.Int32,
+                                        ibase + cutlass.Int64(src_c) * cutlass.Int64(4),
+                                        cute.AddressSpace.gmem,
+                                        assumed_align=4,
+                                    )
+                                    pidx = cute.make_tensor(ip_c, cute.make_layout((1,)))[0]
                                     if pval >= cut_t:
                                         keep = cutlass.Int32(1)
                                 pvals.append(pval)
-                                pidxs.append(src_c)
+                                pidxs.append(pidx)
                                 keeps.append(keep)
                             m0 = cute.arch.vote_ballot_sync(keeps[0] != cutlass.Int32(0))
                             m1 = cute.arch.vote_ballot_sync(keeps[1] != cutlass.Int32(0))
@@ -6745,15 +6746,15 @@ def _run_phases(
                 if cutlass.const_expr(_P4_SUB_DBG):
                     ck_sw0 = cute.arch.clock64()
                 if cutlass.const_expr(
-                    (self.use_ext_cand or self.self_scan)
-                    and self.use_ext_counts
-                    and self.dtype == cutlass.Float32
+                    self.self_scan and self.use_ext_counts and self.dtype == cutlass.Float32
                 ):
-                    # List rows: the compact stored LIST INDICES in the
-                    # vals slots (saving a second cold gmem pass over the
-                    # position column). Swap them for true positions with
-                    # K fully-parallel gathers. Must precede the xstate
+                    # self_scan rows: the compact stored SEGMENT COORDS in
+                    # the vals slots. Swap them for true positions with K
+                    # fully-parallel gathers. Must precede the xstate
                     # publish, which reads output slot K-1 as a position.
+                    # (ext_cand list rows translate EAGERLY in the take
+                    # walk - the idx column rides the value ILP batch -
+                    # so they never reach this loop.)
                     if list_used == cutlass.Int32(1):
                         io_r = tidx
                         while io_r < cutlass.Int32(self.top_k):

From 71f3d709119e9066ebb90240b55cf3fe1e375bda Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 09:42:53 -0700
Subject: [PATCH 045/117] [None][perf] GVR P4: walk-fused range + hist zero for
 list rows

List rows know their candidate minimum by construction (the cut line),
and the take walk already touches every candidate value - accumulate
the per-fragment max there (free ILP), warp-reduce it into smem_wcnt
under the walk's existing end barrier, and pre-zero the coarse
histogram in the walk prologue. P4's block min/max scan, its zero pass
and their three barriers all vanish for list rows; fallback rows keep
the stock sequence (block-uniform runtime branch, same pattern as the
ext-row paths). Also strictly safer than the old min: sentinel pads
(-inf) used to drag bmin to -inf on t0 cuts and collapse the coarse
binning; cut-line min bins them to slot 0 via the existing clamp.

Warm device residency on pro mid cells: 5.3 -> 4.6us (P4+pub 4.2 ->
3.3us, walk +0.2 absorbing the staging); 32 cells exact, unified_smoke
3 cells x 10 modes exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 165 ++++++++++++------
 1 file changed, 113 insertions(+), 52 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 420df20c4088..c237f8aeda79 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3403,6 +3403,8 @@ def phase4_rank_scatter(
         tidx,
         warp_id,
         lane,
+        ext_range_flag=None,  # list rows: walk pre-staged range + hist zero
+        ext_min=None,  # list rows: cut line == exact candidate minimum
     ):
         kK = cutlass.const_expr(self.top_k)
         kBins = cutlass.const_expr(self.kNumBins)
@@ -3427,43 +3429,61 @@ def phase4_rank_scatter(
             sc6 = cutlass.Int64(0)
             if cutlass.const_expr(_P4_SUB_DBG):
                 sc0 = cute.arch.clock64()
-            # ---- block min/max over candidates ----
-            local_cmin = cutlass.Float32(self.FLT_MAX)
-            local_cmax = cutlass.Float32(self.NEG_FLT_MAX)
-            i5 = tidx
-            while i5 < cand_count:
-                v = smem_keys[i5]
-                local_cmin = _fmin_f32_inline(local_cmin, v)
-                local_cmax = cute.arch.fmax(local_cmax, v)
-                i5 = i5 + cutlass.Int32(num_threads)
-            cmin = self.warp_reduce_min_f32(local_cmin)
-            cmax = self.warp_reduce_max_f32(local_cmax)
-            if lane == cutlass.Int32(0):
-                smem_wcnt[warp_id] = float_as_uint32(cmin)
-                smem_hist[warp_id] = float_as_uint32(cmax)
-            cute.arch.barrier()
             bmin_r = cutlass.Float32(self.FLT_MAX)
             bmax_r = cutlass.Float32(self.NEG_FLT_MAX)
-            for w in cutlass.range_constexpr(self.num_warps):
-                vmin = cutlass.Float32(
-                    llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value())
-                )
-                vmax = cutlass.Float32(
-                    llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[w].ir_value())
-                )
-                bmin_r = _fmin_f32_inline(bmin_r, vmin)
-                bmax_r = cute.arch.fmax(bmax_r, vmax)
-            if bmax_r <= bmin_r:
-                bmax_r = bmin_r + cutlass.Float32(1e-6)
-            cute.arch.barrier()
-            if cutlass.const_expr(_P4_SUB_DBG):
-                sc1 = cute.arch.clock64()
-            # ---- zero + build histogram ----
-            i6 = tidx
-            while i6 < cutlass.Int32(kBins):
-                smem_hist[i6] = cutlass.Int32(0)
-                i6 = i6 + cutlass.Int32(num_threads)
-            cute.arch.barrier()
+            use_ext_r = cutlass.Int32(0)
+            if cutlass.const_expr(ext_range_flag is not None):
+                use_ext_r = ext_range_flag
+            if use_ext_r == cutlass.Int32(1):
+                # list rows: the take walk pre-zeroed the hist and staged
+                # per-warp maxima in smem_wcnt (its end barrier orders
+                # them); min := cut line by construction. The minmax
+                # scan, the zero pass and their three barriers vanish.
+                if cutlass.const_expr(ext_min is not None):
+                    bmin_r = ext_min
+                for w in cutlass.range_constexpr(self.num_warps):
+                    vmax = cutlass.Float32(
+                        llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value())
+                    )
+                    bmax_r = cute.arch.fmax(bmax_r, vmax)
+                if bmax_r <= bmin_r:
+                    bmax_r = bmin_r + cutlass.Float32(1e-6)
+            if use_ext_r == cutlass.Int32(0):
+                # ---- block min/max over candidates ----
+                local_cmin = cutlass.Float32(self.FLT_MAX)
+                local_cmax = cutlass.Float32(self.NEG_FLT_MAX)
+                i5 = tidx
+                while i5 < cand_count:
+                    v = smem_keys[i5]
+                    local_cmin = _fmin_f32_inline(local_cmin, v)
+                    local_cmax = cute.arch.fmax(local_cmax, v)
+                    i5 = i5 + cutlass.Int32(num_threads)
+                cmin = self.warp_reduce_min_f32(local_cmin)
+                cmax = self.warp_reduce_max_f32(local_cmax)
+                if lane == cutlass.Int32(0):
+                    smem_wcnt[warp_id] = float_as_uint32(cmin)
+                    smem_hist[warp_id] = float_as_uint32(cmax)
+                cute.arch.barrier()
+                for w in cutlass.range_constexpr(self.num_warps):
+                    vmin = cutlass.Float32(
+                        llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value())
+                    )
+                    vmax = cutlass.Float32(
+                        llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[w].ir_value())
+                    )
+                    bmin_r = _fmin_f32_inline(bmin_r, vmin)
+                    bmax_r = cute.arch.fmax(bmax_r, vmax)
+                if bmax_r <= bmin_r:
+                    bmax_r = bmin_r + cutlass.Float32(1e-6)
+                cute.arch.barrier()
+                if cutlass.const_expr(_P4_SUB_DBG):
+                    sc1 = cute.arch.clock64()
+                # ---- zero + build histogram ----
+                i6 = tidx
+                while i6 < cutlass.Int32(kBins):
+                    smem_hist[i6] = cutlass.Int32(0)
+                    i6 = i6 + cutlass.Int32(num_threads)
+                cute.arch.barrier()
             range1 = bmax_r - bmin_r
             inv1 = (cutlass.Float32(kBins - 1) + cutlass.Float32(0.99)) / range1
             i7 = tidx
@@ -5901,6 +5921,17 @@ def _run_phases(
                         s_iscalars[1] = cutlass.Int32(1)  # done
                     cute.arch.barrier()
                     lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1)
+                    # fused P4 prologue: zero the coarse hist here and
+                    # accumulate the candidate max INSIDE the cut walk
+                    # (per-fragment fmax is free ILP; min := cut line by
+                    # construction). P4's minmax scan + zero pass and
+                    # their three barriers disappear for list rows - the
+                    # staging rides this walk's own end barrier.
+                    izh_c = tidx
+                    while izh_c < cutlass.Int32(self.kNumBins):
+                        smem_hist[izh_c] = cutlass.Int32(0)
+                        izh_c = izh_c + cutlass.Int32(num_threads)
+                    wmax_acc = cutlass.Float32(self.NEG_FLT_MAX)
                     if line_cut == cutlass.Int32(1):
                         # ---- LINE cut: dense mapped-prefix COPY of
                         # exactly cut_n entries. No filter, no ballots,
@@ -5923,9 +5954,9 @@ def _run_phases(
                                         cute.AddressSpace.gmem,
                                         assumed_align=4,
                                     )
-                                    smem_keys[j_c] = cute.make_tensor(vp_c, cute.make_layout((1,)))[
-                                        0
-                                    ]
+                                    pv_c = cute.make_tensor(vp_c, cute.make_layout((1,)))[0]
+                                    smem_keys[j_c] = pv_c
+                                    wmax_acc = cute.arch.fmax(wmax_acc, pv_c)
                                     # eager position fetch: the idx column
                                     # rides the same ILP batch as the value
                                     # read, so vals hold TRUE positions and
@@ -5940,6 +5971,9 @@ def _run_phases(
                                         0
                                     ]
                             i_c = i_c + cutlass.Int32(4 * num_threads)
+                        wmax_w = self.warp_reduce_max_f32(wmax_acc)
+                        if lane_c == cutlass.Int32(0):
+                            smem_wcnt[tidx // cutlass.Int32(32)] = float_as_uint32(wmax_w)
                         cute.arch.barrier()
                     if line_cut == cutlass.Int32(0):
                         # ---- histogram-edge cut: value-filtered mapped
@@ -5980,6 +6014,7 @@ def _run_phases(
                                     pidx = cute.make_tensor(ip_c, cute.make_layout((1,)))[0]
                                     if pval >= cut_t:
                                         keep = cutlass.Int32(1)
+                                wmax_acc = cute.arch.fmax(wmax_acc, pval)
                                 pvals.append(pval)
                                 pidxs.append(pidx)
                                 keeps.append(keep)
@@ -6019,6 +6054,9 @@ def _run_phases(
                                             smem_vals[wpos] = pidxs[_ju]
                                     off = off + cutlass.Int32(cute.arch.popc(mj))
                             i_c = i_c + cutlass.Int32(4 * num_threads)
+                        wmax_w2 = self.warp_reduce_max_f32(wmax_acc)
+                        if lane_c == cutlass.Int32(0):
+                            smem_wcnt[tidx // cutlass.Int32(32)] = float_as_uint32(wmax_w2)
                         cute.arch.barrier()
                         cnt_l = s_iscalars[0]
                         if cnt_l < cutlass.Int32(self.top_k) or cnt_l > cutlass.Int32(self.kC):
@@ -6712,20 +6750,43 @@ def _run_phases(
                 # cs=1: the single CTA per row IS the leader.
                 cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC))
                 if cutlass.const_expr(self.enable_p4_rank_scatter):
-                    self.phase4_rank_scatter(
-                        smem_keys,
-                        smem_vals,
-                        smem_hist,
-                        smem_wcnt,
-                        s_thr,
-                        s_iscalars,
-                        output_values_row,
-                        output_indices_row,
-                        cand_count_p4,
-                        tidx,
-                        warp_id,
-                        lane,
-                    )
+                    if cutlass.const_expr(
+                        self.use_ext_cand and self.use_ext_counts and self.dtype == cutlass.Float32
+                    ):
+                        # list rows carry a walk-staged range + pre-zeroed
+                        # hist (flag = list_used; fallback rows take the
+                        # stock minmax path inside)
+                        self.phase4_rank_scatter(
+                            smem_keys,
+                            smem_vals,
+                            smem_hist,
+                            smem_wcnt,
+                            s_thr,
+                            s_iscalars,
+                            output_values_row,
+                            output_indices_row,
+                            cand_count_p4,
+                            tidx,
+                            warp_id,
+                            lane,
+                            ext_range_flag=list_used,
+                            ext_min=cut_t,
+                        )
+                    else:
+                        self.phase4_rank_scatter(
+                            smem_keys,
+                            smem_vals,
+                            smem_hist,
+                            smem_wcnt,
+                            s_thr,
+                            s_iscalars,
+                            output_values_row,
+                            output_indices_row,
+                            cand_count_p4,
+                            tidx,
+                            warp_id,
+                            lane,
+                        )
                 else:
                     self.phase4_histogram_snap(
                         smem_keys,

From badea6ca4f38ce8b70621588007109f48157dbd6 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 10:41:45 -0700
Subject: [PATCH 046/117] [None][fix] GVR P4: revert warp0 bin search to the
 3-step block search

The single-warp shuffle-scan rewrite of the coarse/fine bin searches
(part of the previous P4 commit) compiles and passes single-cell
exactness, unified_smoke and all three compute-sanitizer tools, but
intermittently corrupts results / hits illegal addresses under a full
harness context (reused buffers + warmup + cold-evict loops; failure
location shifts run to run). Kernel-file bisect against a context
repro pins it to the warp0 hunks: warp collectives and first-assigned
locals inside a dynamic 'if warp_id == 0:' branch are not reliably
compiled by the DSL. Revert both searches to the 3-step block form
(~0.4us given back on mid cells); eager position fetch and the
walk-fused range stay in and are context-repro/cold-verified clean.

Cold kernel-only after this state (pro mid, 30-layer means): wf
8k/16k/32k/64k B1 = 1.41/1.52/1.56/1.79x vs PR tip (pre-sprint:
1.19/1.25/1.28/1.44); flash 512k B1 2.54x; 480 cold cells + 30-point
smoke all exact.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 189 ++++++++++--------
 1 file changed, 106 insertions(+), 83 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index c237f8aeda79..b0ca43ac5a00 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3499,49 +3499,56 @@ def phase4_rank_scatter(
             cute.arch.barrier()
             if cutlass.const_expr(_P4_SUB_DBG):
                 sc2 = cute.arch.clock64()
-            # ---- warp0 high→low bin search → straddling bin b* + rank_above ----
-            # Single-warp shuffle scan replaces the 3-step block search:
-            # lane L owns the L-th DESCENDING chunk of bins, an inclusive
-            # warp scan finds the straddling lane (prefixes are monotone,
-            # so hit lanes form a suffix: target = 32 - popc(ballot)),
-            # and that lane resolves b* serially in its own chunk. Two of
-            # the three block barriers (~0.2us each) disappear.
-            if warp_id == cutlass.Int32(0):
-                lsum = cutlass.Int32(0)
-                for jb in cutlass.range_constexpr(bins_per_warp):
-                    bidx_s = (
+            # ---- 3-step high→low bin search → straddling bin b* + rank_above ----
+            warp_bin_sum = cutlass.Int32(0)
+            for jb in cutlass.range_constexpr(bins_per_warp):
+                bidx_s = (
+                    cutlass.Int32(kBins - 1)
+                    - warp_id * cutlass.Int32(bins_per_warp)
+                    - cutlass.Int32(jb)
+                )
+                warp_bin_sum = warp_bin_sum + smem_hist[bidx_s]
+            if lane == cutlass.Int32(0):
+                smem_wcnt[warp_id] = warp_bin_sum
+            cute.arch.barrier()
+            if tidx == cutlass.Int32(0):
+                cum = cutlass.Int32(0)
+                tw = cutlass.Int32(num_warps - 1)
+                found = cutlass.Int32(0)
+                for w2 in cutlass.range_constexpr(self.num_warps):
+                    cum = cum + smem_wcnt[w2]
+                    if cum >= cutlass.Int32(kK) and found == cutlass.Int32(0):
+                        tw = cutlass.Int32(w2)
+                        found = cutlass.Int32(1)
+                cum2 = cutlass.Int32(0)
+                for w3 in cutlass.range_constexpr(self.num_warps):
+                    if cutlass.Int32(w3) < tw:
+                        cum2 = cum2 + smem_wcnt[w3]
+                s_iscalars[2] = cum2  # prefix-count before target warp
+                s_iscalars[3] = tw
+            cute.arch.barrier()
+            target_warp = s_iscalars[3]
+            if warp_id == target_warp and lane == cutlass.Int32(0):
+                base_cum = s_iscalars[2]
+                b_star = cutlass.Int32(kBins - 1)
+                rank_above = base_cum
+                set_d = cutlass.Int32(0)
+                for jb2 in cutlass.range_constexpr(bins_per_warp):
+                    bidx2 = (
                         cutlass.Int32(kBins - 1)
-                        - lane * cutlass.Int32(bins_per_warp)
-                        - cutlass.Int32(jb)
+                        - target_warp * cutlass.Int32(bins_per_warp)
+                        - cutlass.Int32(jb2)
                     )
-                    lsum = lsum + smem_hist[bidx_s]
-                pref = warp_scan(lsum, tidx, lane, 32)
-                hit0 = cutlass.Int32(0)
-                if pref >= cutlass.Int32(kK):
-                    hit0 = cutlass.Int32(1)
-                mhit = cute.arch.vote_ballot_sync(hit0 != cutlass.Int32(0))
-                tl0 = cutlass.Int32(32) - cutlass.Int32(cute.arch.popc(mhit))
-                if lane == tl0:
-                    base_cum = pref - lsum  # exclusive prefix before chunk
-                    b_star = cutlass.Int32(kBins - 1)
-                    rank_above = base_cum
-                    set_d = cutlass.Int32(0)
-                    for jb2 in cutlass.range_constexpr(bins_per_warp):
-                        bidx2 = (
-                            cutlass.Int32(kBins - 1)
-                            - lane * cutlass.Int32(bins_per_warp)
-                            - cutlass.Int32(jb2)
-                        )
-                        ra_before = base_cum
-                        base_cum = base_cum + smem_hist[bidx2]
-                        if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0):
-                            b_star = bidx2
-                            rank_above = ra_before  # count strictly above b*
-                            set_d = cutlass.Int32(1)
-                    s_iscalars[2] = rank_above
-                    s_iscalars[3] = b_star
-                    s_iscalars[4] = cutlass.Int32(0)  # cnt_above
-                    s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
+                    ra_before = base_cum
+                    base_cum = base_cum + smem_hist[bidx2]
+                    if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0):
+                        b_star = bidx2
+                        rank_above = ra_before  # count in bins strictly above b*
+                        set_d = cutlass.Int32(1)
+                s_iscalars[2] = rank_above
+                s_iscalars[3] = b_star
+                s_iscalars[4] = cutlass.Int32(0)  # cnt_above
+                s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
             cute.arch.barrier()
             if cutlass.const_expr(_P4_SUB_DBG):
                 sc3 = cute.arch.clock64()
@@ -3582,50 +3589,66 @@ def phase4_rank_scatter(
                         atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1))
                     ifb = ifb + cutlass.Int32(num_threads)
                 cute.arch.barrier()
-                # fine warp0 search seeded at rank_above (over fbins bins);
-                # same single-warp shuffle-scan shape as the coarse search
-                # above - three block barriers collapse into one. sb*/ra
-                # go via smem_hist[2]/[3] (fine bins walked never reach
-                # slots 2/3 only when the target chunk excludes them, so
-                # publish AFTER the serial walk like before - the walk
-                # reads bins, the publish overwrites scratch slots).
-                if warp_id == cutlass.Int32(0):
-                    fls = cutlass.Int32(0)
-                    for jbf in cutlass.range_constexpr(fbpw):
-                        bif = (
+                # fine 3-step search seeded at rank_above (over fbins bins)
+                fws = cutlass.Int32(0)
+                for jbf in cutlass.range_constexpr(fbpw):
+                    bif = (
+                        cutlass.Int32(fbins - 1)
+                        - warp_id * cutlass.Int32(fbpw)
+                        - cutlass.Int32(jbf)
+                    )
+                    fws = fws + smem_hist[bif]
+                if lane == cutlass.Int32(0):
+                    smem_wcnt[warp_id] = fws
+                cute.arch.barrier()
+                if tidx == cutlass.Int32(0):
+                    cumf = rank_above
+                    twf = cutlass.Int32(num_warps - 1)
+                    fnd = cutlass.Int32(0)
+                    for w2 in cutlass.range_constexpr(self.num_warps):
+                        cumf = cumf + smem_wcnt[w2]
+                        if cumf >= cutlass.Int32(kK) and fnd == cutlass.Int32(0):
+                            twf = cutlass.Int32(w2)
+                            fnd = cutlass.Int32(1)
+                    pre = rank_above
+                    for w3 in cutlass.range_constexpr(self.num_warps):
+                        if cutlass.Int32(w3) < twf:
+                            pre = pre + smem_wcnt[w3]
+                    # Stage prefix/target-warp metadata in spare s_iscalars
+                    # slots, NOT smem_hist[0]/[1]: the last fine warp's reverse
+                    # scan below walks fine bins down to 0/1, so reusing those
+                    # histogram bins as scratch would corrupt sb_star/ra_fine
+                    # when twf2 == num_warps-1. Slots [4]/[1] are dead here
+                    # (re-zeroed at the cnt_above/cnt_strad reset below).
+                    s_iscalars[4] = pre  # prefix into target fine warp
+                    s_iscalars[1] = twf  # target fine warp
+                cute.arch.barrier()
+                pre_f = s_iscalars[4]
+                twf2 = s_iscalars[1]
+                if warp_id == twf2 and lane == cutlass.Int32(0):
+                    base_f = pre_f
+                    sb_star = cutlass.Int32(fbins - 1)
+                    ra_fine = base_f
+                    sd = cutlass.Int32(0)
+                    for jb3 in cutlass.range_constexpr(fbpw):
+                        sbi = (
                             cutlass.Int32(fbins - 1)
-                            - lane * cutlass.Int32(fbpw)
-                            - cutlass.Int32(jbf)
+                            - twf2 * cutlass.Int32(fbpw)
+                            - cutlass.Int32(jb3)
                         )
-                        fls = fls + smem_hist[bif]
-                    fpref = warp_scan(fls, tidx, lane, 32) + rank_above
-                    fhit = cutlass.Int32(0)
-                    if fpref >= cutlass.Int32(kK):
-                        fhit = cutlass.Int32(1)
-                    fm = cute.arch.vote_ballot_sync(fhit != cutlass.Int32(0))
-                    ftl = cutlass.Int32(32) - cutlass.Int32(cute.arch.popc(fm))
-                    if lane == ftl:
-                        base_f = fpref - fls
-                        sb_star = cutlass.Int32(fbins - 1)
-                        ra_fine = base_f
-                        sd = cutlass.Int32(0)
-                        for jb3 in cutlass.range_constexpr(fbpw):
-                            sbi = (
-                                cutlass.Int32(fbins - 1)
-                                - lane * cutlass.Int32(fbpw)
-                                - cutlass.Int32(jb3)
-                            )
-                            ra_b = base_f
-                            base_f = base_f + smem_hist[sbi]
-                            if base_f >= cutlass.Int32(kK) and sd == cutlass.Int32(0):
-                                sb_star = sbi
-                                ra_fine = ra_b
-                                sd = cutlass.Int32(1)
-                        smem_hist[2] = sb_star
-                        smem_hist[3] = ra_fine
-                        s_iscalars[4] = cutlass.Int32(0)  # cnt_above
-                        s_iscalars[0] = cutlass.Int32(0)  # cnt_mid (b*, sub>sb*)
-                        s_iscalars[1] = cutlass.Int32(0)  # cnt_strad (b*, sub==sb*)
+                        ra_b = base_f
+                        base_f = base_f + smem_hist[sbi]
+                        if base_f >= cutlass.Int32(kK) and sd == cutlass.Int32(0):
+                            sb_star = sbi
+                            ra_fine = ra_b
+                            sd = cutlass.Int32(1)
+                    smem_hist[2] = sb_star
+                    smem_hist[3] = ra_fine
+                cute.arch.barrier()
+                if tidx == cutlass.Int32(0):
+                    s_iscalars[4] = cutlass.Int32(0)  # cnt_above
+                    s_iscalars[0] = cutlass.Int32(0)  # cnt_mid (b*, sub>sb*)
+                    s_iscalars[1] = cutlass.Int32(0)  # cnt_strad (b*, sub==sb*)
                 cute.arch.barrier()
                 if cutlass.const_expr(_P4_SUB_DBG):
                     sc4 = cute.arch.clock64()

From 9bf8b0c317d7b2134f35d34b8dc37f3cc9b2b9a1 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 10:42:53 -0700
Subject: [PATCH 047/117] [None][perf] GVR: closed-loop adaptive block-skip
 gate

Layer variance makes any static block-skip rule wrong on hard layers
(measured 12-80% tight-line pass rates across layers on the same
shape): a pass rate past ~3/8 means the gather walk reads most of the
row anyway PLUS the prefix, and dense wins.

Emission side: while computing block_max the epilogue also counts the
records clearing the loosest line (r_bmax is warp-uniform, lane0
accumulates, the seed-count flush redux sums warps) and lands the
count in packed seed row col 6 via red.global.add.f32 (packed mode
only; 0 = not provided, static behavior preserved).

Kernel side: the R0 multi-line count call derives a skip veto from
col 6 before building the active list - vetoed rows keep the dense
walk and skip the whole wasted build. pack_seed() mirrors the count
from an emu block_max for harness runs.

Context-repro + unified_smoke (3 cells x 10 modes) + cold 480-cell
sweep all exact with the gate wired.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 50 ++++++++++--
 .../blackwell/top_k/gvr_topk_decode.py        | 81 ++++++++++++++-----
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 13 ++-
 3 files changed, 118 insertions(+), 26 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index cd5fbbaa07bc..c7b4fcaf1335 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -1075,7 +1075,7 @@ def _flush_hit_agg(
             hacc_cnt[t] = cutlass.Int32(0)
 
     @cute.jit
-    def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane):
+    def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None):
         """Warp-redux the lane-local seed counters and fire one lane-0
         red.global.add per (t, threshold). Caller zero-initializes the
         count slots each step; cross-CTA totals accumulate atomically.
@@ -1100,6 +1100,18 @@ def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane):
                         ) * cutlass.Int64(4)
                         _red_global_add_s32(addr, w_cnt)
                 scnt[t * 3 + j] = cutlass.Int32(0)
+        if cutlass.const_expr(self.seed_packed and spass is not None):
+            # packed col 6: adaptive-skip pass count (lane0-accumulated,
+            # so the warp redux is exactly the warp's record total)
+            for t in cutlass.range_constexpr(next_n):
+                w_bp = cute.arch.warp_redux_sync(spass[t], "add")
+                if meta_lane == cutlass.Int32(0):
+                    row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t)
+                    addr = base_addr + (
+                        cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(6)
+                    ) * cutlass.Int64(4)
+                    _red_global_add_f32(addr, cutlass.Float32(w_bp))
+                spass[t] = cutlass.Int32(0)
 
     @cute.jit
     def _flush_cand_window(self, mCand, q_idx, cwbase, cwleft, meta_lane):
@@ -2158,6 +2170,9 @@ def kernel(
                     if cutlass.const_expr(self.emit_seed_counts):
                         sthr = cute.make_fragment(next_n * 3, cutlass.Float32)
                         scnt = cute.make_fragment(next_n * 3, cutlass.Int32)
+                        spass = cute.make_fragment(next_n, cutlass.Int32)
+                        for _i in cutlass.range_constexpr(next_n):
+                            spass[_i] = cutlass.Int32(0)
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
@@ -2233,7 +2248,9 @@ def kernel(
                             # the aligned padding region out of block_max.
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
-                                    self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                                    self._flush_seed_counts(
+                                        mSeedCounts, q_idx_old, scnt, meta_lane, spass=spass
+                                    )
                                 # (re)load this q's thresholds - gated on
                                 # emit_seed_counts, NOT emit_cand: counts-
                                 # only mode needs them too (a stale
@@ -2484,6 +2501,15 @@ def kernel(
                                 for _j in cutlass.range_constexpr(3):
                                     ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j])
                                     scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1)
+                                if cutlass.const_expr(self.seed_packed):
+                                    # adaptive-skip pass count: one record
+                                    # per (tile, warp); r_bmax is warp-
+                                    # uniform so lane0 alone accumulates
+                                    # (the flush redux then sums warps)
+                                    if meta_lane == cutlass.Int32(0):
+                                        spass[t] = spass[t] + cutlass.Int32(
+                                            r_bmax >= sthr[t * 3 + 0]
+                                        )
                             if cutlass.const_expr(self.emit_cand):
                                 # L2 pre-collect at t_0 with per-warp claim
                                 # WINDOWS: refills claim (hits + CAND_WIN)
@@ -2648,7 +2674,7 @@ def kernel(
                         )
                 if cutlass.const_expr(self.emit_seed_counts):
                     if q_idx < batch_size:
-                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane)
+                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane, spass=spass)
                 if cutlass.const_expr(self.emit_cand):
                     if q_idx < batch_size:
                         self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane)
@@ -2706,6 +2732,9 @@ def kernel(
                     if cutlass.const_expr(self.emit_seed_counts):
                         sthr = cute.make_fragment(next_n * 3, cutlass.Float32)
                         scnt = cute.make_fragment(next_n * 3, cutlass.Int32)
+                        spass = cute.make_fragment(next_n, cutlass.Int32)
+                        for _i in cutlass.range_constexpr(next_n):
+                            spass[_i] = cutlass.Int32(0)
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
@@ -2781,7 +2810,9 @@ def kernel(
                             # the aligned padding region out of block_max.
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
-                                    self._flush_seed_counts(mSeedCounts, q_idx_old, scnt, meta_lane)
+                                    self._flush_seed_counts(
+                                        mSeedCounts, q_idx_old, scnt, meta_lane, spass=spass
+                                    )
                                 # (re)load this q's thresholds - gated on
                                 # emit_seed_counts, NOT emit_cand (see the
                                 # WG0 twin above)
@@ -3025,6 +3056,15 @@ def kernel(
                                 for _j in cutlass.range_constexpr(3):
                                     ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j])
                                     scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1)
+                                if cutlass.const_expr(self.seed_packed):
+                                    # adaptive-skip pass count: one record
+                                    # per (tile, warp); r_bmax is warp-
+                                    # uniform so lane0 alone accumulates
+                                    # (the flush redux then sums warps)
+                                    if meta_lane == cutlass.Int32(0):
+                                        spass[t] = spass[t] + cutlass.Int32(
+                                            r_bmax >= sthr[t * 3 + 0]
+                                        )
                             if cutlass.const_expr(self.emit_cand):
                                 # L2 pre-collect at t_0 with per-warp claim
                                 # WINDOWS: refills claim (hits + CAND_WIN)
@@ -3189,7 +3229,7 @@ def kernel(
                         )
                 if cutlass.const_expr(self.emit_seed_counts):
                     if q_idx < batch_size:
-                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane)
+                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane, spass=spass)
                 if cutlass.const_expr(self.emit_cand):
                     if q_idx < batch_size:
                         self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane)
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index b0ca43ac5a00..b164553ecc88 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -2078,6 +2078,7 @@ def block_count_ge_multi(
         block_max_row=None,  # block-skip: per-32-position upper bounds
         smem_active=None,  # block-skip: int16 active list
         s_active_cnt=None,  # block-skip: [0]=list length, [1]=list-current flag
+        skip_veto=None,  # adaptive gate: 1 = pass rate too high, walk dense
     ):
         M = cutlass.const_expr(self.M_thr)
         num_threads = cutlass.const_expr(self.num_threads)
@@ -2127,6 +2128,15 @@ def block_count_ge_multi(
                 skip_ok = cutlass.Int32(0)
             if blk_hi_g > cutlass.Int32(32767):
                 skip_ok = cutlass.Int32(0)
+            # Closed-loop adaptive gate: when the emission (or harness)
+            # counted the blocks clearing the loosest line, a high pass
+            # rate proves the gather walk would read most of the row
+            # anyway PLUS the prefix - dense wins. Layer variance makes
+            # any static rule wrong on hard layers (measured 12-80%
+            # pass rates across layers on the same shape).
+            if cutlass.const_expr(skip_veto is not None):
+                if skip_veto == cutlass.Int32(1):
+                    skip_ok = cutlass.Int32(0)
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             if skip_ok == cutlass.Int32(1):
                 head_end = (
@@ -6459,24 +6469,59 @@ def _run_phases(
                         ] == cutlass.Int32(1):
                             r0_par = cutlass.Int32(1)
                     if r0_par == cutlass.Int32(0):
-                        self.block_count_ge_multi(
-                            input_row,
-                            slice_start,
-                            slice_end,
-                            s_mt_thr,
-                            smem_ptcnt_multi,
-                            smem_wcnt_multi,
-                            s_mt_cnt,
-                            s_cluster_partial_m,
-                            do_cluster_sync,
-                            tidx,
-                            warp_id,
-                            lane,
-                            smem_ptcnt=smem_ptcnt,
-                            block_max_row=block_max_row,
-                            smem_active=smem_active,
-                            s_active_cnt=s_active_cnt,
-                        )
+                        if cutlass.const_expr(self.use_ext_counts and self.enable_block_skip):
+                            # adaptive skip gate from the packed seed row:
+                            # col 6 carries the emission's count of blocks
+                            # clearing the loosest line (0 = not provided,
+                            # keep static behavior). Pass rates past 3/8
+                            # make the gather walk lose to dense.
+                            sv_r = cutlass.Int32(0)
+                            bp_r = cutlass.Int32(seed_thr_row[6])
+                            nb_r = (
+                                slice_end - slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)
+                            ) >> cutlass.Int32(self.SKIP_BLOCK_LOG2)
+                            if bp_r > cutlass.Int32(0) and bp_r * cutlass.Int32(
+                                8
+                            ) > nb_r * cutlass.Int32(3):
+                                sv_r = cutlass.Int32(1)
+                            self.block_count_ge_multi(
+                                input_row,
+                                slice_start,
+                                slice_end,
+                                s_mt_thr,
+                                smem_ptcnt_multi,
+                                smem_wcnt_multi,
+                                s_mt_cnt,
+                                s_cluster_partial_m,
+                                do_cluster_sync,
+                                tidx,
+                                warp_id,
+                                lane,
+                                smem_ptcnt=smem_ptcnt,
+                                block_max_row=block_max_row,
+                                smem_active=smem_active,
+                                s_active_cnt=s_active_cnt,
+                                skip_veto=sv_r,
+                            )
+                        else:
+                            self.block_count_ge_multi(
+                                input_row,
+                                slice_start,
+                                slice_end,
+                                s_mt_thr,
+                                smem_ptcnt_multi,
+                                smem_wcnt_multi,
+                                s_mt_cnt,
+                                s_cluster_partial_m,
+                                do_cluster_sync,
+                                tidx,
+                                warp_id,
+                                lane,
+                                smem_ptcnt=smem_ptcnt,
+                                block_max_row=block_max_row,
+                                smem_active=smem_active,
+                                s_active_cnt=s_active_cnt,
+                            )
                     cute.arch.barrier()
                     if tidx == 0 and r0_par == cutlass.Int32(0):
                         # tightest admissible rung = SMALLEST count in [K, kC].
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 268264b9b956..58ca7336b6c7 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -520,16 +520,23 @@ def derive_seed_lines_v4(
     return out.contiguous()
 
 
-def pack_seed(seed_thr: torch.Tensor, seed_counts: torch.Tensor) -> torch.Tensor:
+def pack_seed(
+    seed_thr: torch.Tensor,
+    seed_counts: torch.Tensor,
+    block_max: torch.Tensor = None,
+) -> torch.Tensor:
     """Pack lines + exact counts into one [rows, 8] fp32 seed row.
 
     Lines land at [0..2], counts as floats at [3..5] (exact to 2^24);
-    one 32B sector per row. Build ONCE per step, outside any timed
-    region.
+    col 6 optionally carries the adaptive-skip pass count (32-grain
+    block records clearing t_0; 0 = not provided). One 32B sector per
+    row. Build ONCE per step, outside any timed region.
     """
     pack = torch.zeros((seed_thr.shape[0], 8), dtype=torch.float32, device=seed_thr.device)
     pack[:, 0:3] = seed_thr
     pack[:, 3:6] = seed_counts.float()
+    if block_max is not None:
+        pack[:, 6] = (block_max >= seed_thr[:, 0:1]).sum(dim=1).float()
     return pack.contiguous()
 
 

From 71d569a9c6fe815be3f9646eb98d6c26a88d2ca8 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 10:54:11 -0700
Subject: [PATCH 048/117] [None][fix] GVR: drop the block-skip pass-rate veto,
 keep col-6 count

Measured dose-response on va pro-256k (cold, 30 layers, B 1/2/8/16):
veto at 3/8 pass rate costs 10-18%, at 7/8 still 3-10%, removed
restores parity (x1.00-1.05 vs PR tip). Root cause: the rung-
tightening build already salvages fat loosest-rung rows by rebuilding
at a tighter line, so a loosest-line veto only preempts a smarter
in-kernel mechanism. The emission-side pass count (packed seed row
col 6) stays: zero-cost, useful for host-side routing diagnostics.
The remaining 256k small-B weak band tracks va's count+collect cost,
not skip quality - routing sends those shapes to the list path (wf,
1.7-2.1x there).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 89 ++++++-------------
 1 file changed, 26 insertions(+), 63 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index b164553ecc88..c4c021dd314a 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -2078,7 +2078,6 @@ def block_count_ge_multi(
         block_max_row=None,  # block-skip: per-32-position upper bounds
         smem_active=None,  # block-skip: int16 active list
         s_active_cnt=None,  # block-skip: [0]=list length, [1]=list-current flag
-        skip_veto=None,  # adaptive gate: 1 = pass rate too high, walk dense
     ):
         M = cutlass.const_expr(self.M_thr)
         num_threads = cutlass.const_expr(self.num_threads)
@@ -2128,15 +2127,6 @@ def block_count_ge_multi(
                 skip_ok = cutlass.Int32(0)
             if blk_hi_g > cutlass.Int32(32767):
                 skip_ok = cutlass.Int32(0)
-            # Closed-loop adaptive gate: when the emission (or harness)
-            # counted the blocks clearing the loosest line, a high pass
-            # rate proves the gather walk would read most of the row
-            # anyway PLUS the prefix - dense wins. Layer variance makes
-            # any static rule wrong on hard layers (measured 12-80%
-            # pass rates across layers on the same shape).
-            if cutlass.const_expr(skip_veto is not None):
-                if skip_veto == cutlass.Int32(1):
-                    skip_ok = cutlass.Int32(0)
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             if skip_ok == cutlass.Int32(1):
                 head_end = (
@@ -6469,59 +6459,32 @@ def _run_phases(
                         ] == cutlass.Int32(1):
                             r0_par = cutlass.Int32(1)
                     if r0_par == cutlass.Int32(0):
-                        if cutlass.const_expr(self.use_ext_counts and self.enable_block_skip):
-                            # adaptive skip gate from the packed seed row:
-                            # col 6 carries the emission's count of blocks
-                            # clearing the loosest line (0 = not provided,
-                            # keep static behavior). Pass rates past 3/8
-                            # make the gather walk lose to dense.
-                            sv_r = cutlass.Int32(0)
-                            bp_r = cutlass.Int32(seed_thr_row[6])
-                            nb_r = (
-                                slice_end - slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)
-                            ) >> cutlass.Int32(self.SKIP_BLOCK_LOG2)
-                            if bp_r > cutlass.Int32(0) and bp_r * cutlass.Int32(
-                                8
-                            ) > nb_r * cutlass.Int32(3):
-                                sv_r = cutlass.Int32(1)
-                            self.block_count_ge_multi(
-                                input_row,
-                                slice_start,
-                                slice_end,
-                                s_mt_thr,
-                                smem_ptcnt_multi,
-                                smem_wcnt_multi,
-                                s_mt_cnt,
-                                s_cluster_partial_m,
-                                do_cluster_sync,
-                                tidx,
-                                warp_id,
-                                lane,
-                                smem_ptcnt=smem_ptcnt,
-                                block_max_row=block_max_row,
-                                smem_active=smem_active,
-                                s_active_cnt=s_active_cnt,
-                                skip_veto=sv_r,
-                            )
-                        else:
-                            self.block_count_ge_multi(
-                                input_row,
-                                slice_start,
-                                slice_end,
-                                s_mt_thr,
-                                smem_ptcnt_multi,
-                                smem_wcnt_multi,
-                                s_mt_cnt,
-                                s_cluster_partial_m,
-                                do_cluster_sync,
-                                tidx,
-                                warp_id,
-                                lane,
-                                smem_ptcnt=smem_ptcnt,
-                                block_max_row=block_max_row,
-                                smem_active=smem_active,
-                                s_active_cnt=s_active_cnt,
-                            )
+                        # NOTE: a loosest-line pass-rate veto (packed seed
+                        # col 6) was measured here and REMOVED: the rung-
+                        # tightening build below already salvages fat
+                        # loosest-rung rows by rebuilding at a tighter
+                        # line, so any veto threshold (3/8 and 7/8 both
+                        # tried, cold, 30 layers) only preempts that and
+                        # costs 3-18% across the board. Col 6 stays as a
+                        # zero-cost diagnostic for host-side routing.
+                        self.block_count_ge_multi(
+                            input_row,
+                            slice_start,
+                            slice_end,
+                            s_mt_thr,
+                            smem_ptcnt_multi,
+                            smem_wcnt_multi,
+                            s_mt_cnt,
+                            s_cluster_partial_m,
+                            do_cluster_sync,
+                            tidx,
+                            warp_id,
+                            lane,
+                            smem_ptcnt=smem_ptcnt,
+                            block_max_row=block_max_row,
+                            smem_active=smem_active,
+                            s_active_cnt=s_active_cnt,
+                        )
                     cute.arch.barrier()
                     if tidx == 0 and r0_par == cutlass.Int32(0):
                         # tightest admissible rung = SMALLEST count in [K, kC].

From c10ae8c906274d7454f4b091d4432b64c5212ecf Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 10:59:45 -0700
Subject: [PATCH 049/117] [None][feat] GVR: host-side routing module (v1
 constants)

One place for the deployment thresholds measured on the B200 f15
layer-complete grid: emission-tier planning (list for latency-bound
long rows, counts as the near-free default, rungs closed-loop
fallback) and per-shape launch knobs (block-skip attach points,
GPC-aware cluster split, 512-thread small-K list rule). Constants are
data, not logic - the P4 sprint moved the mid-range list-path win to
1.41-1.79x, so the list tier will likely widen once the emission tax
is re-measured end-to-end (noted inline).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_routing.py            | 121 ++++++++++++++++++
 1 file changed, 121 insertions(+)
 create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
new file mode 100644
index 000000000000..a3bd54ca6256
--- /dev/null
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -0,0 +1,121 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION &
+# AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Host-side routing for the GVR top-k decode family.
+
+One kernel, three assist tiers selected by which emission inputs the
+indexer epilogue produced for this step:
+
+  * ``list``   - bucketed candidate list + packed seed row (v5): the
+                 top-k pass never re-reads the row on a hit.
+  * ``counts`` - packed seed row only ([rows, 8]: three lines + three
+                 counts): one filtered row pass, no in-kernel counting.
+  * ``rungs``  - closed-loop lines only (no emission): in-kernel
+                 multi-line count, then the stock collect.
+  * ``none``   - stock kernel (v1 path).
+
+``plan_emission`` decides which tier the epilogue should emit for the
+NEXT step (the emission tax is shape-dependent); ``pick_config`` maps
+(tier, B, N, K) to concrete launch knobs for THIS step.
+
+All thresholds are measured on B200 (f15 layer-complete grid,
+2026-07-27, cold-L2 kernel-only protocol, validated on both the shared
+grid dataset and first-party captures). They are deployment defaults,
+not universal truths - keep them in one place so retuning is a
+constant edit, not a logic edit.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+# ---- measured thresholds (B200, f15 grid) --------------------------------
+
+# Block-skip prefix pays only when whole-row reads dominate.
+SKIP_MIN_N_COUNTS = 65536  # va: attach block_max unconditionally here up
+SKIP_MIN_N_RUNGS_FLASH = 131072  # vb (flash): bm pays from here
+SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split on top of bm from here
+
+# Emission tax ~ B*N on the GEMM side: the candidate list is only worth
+# emitting for latency-bound shapes (small B, long rows).
+LIST_EMIT_MAX_B = 16
+LIST_EMIT_MIN_N = 65536
+
+# rungs-tier block_max pays only at small K: with K=1024 the tight-line
+# pass rate runs too high and the prefix read is pure overhead.
+RUNGS_BM_MAX_K = 512
+
+# 512-thread build wins for list-hit rows at small K (work is O(list)).
+SMALL_K_LIST_THREADS = 512
+SMALL_K_MAX = 512
+
+# GPC packing: cs=8 only while all row-clusters fit half the device
+# (B=16 x cs8 wave-spill regression); cs4/2 keep a 10% headroom.
+CS8_HALF_DEVICE = 2
+CS_HEADROOM_NUM = 9
+CS_HEADROOM_DEN = 10
+
+
+@dataclass
+class TopkRoute:
+    """Launch knobs for one decode step of the GVR top-k kernel."""
+
+    tier: str  # list | counts | rungs | none
+    cluster_size: int = 1
+    num_threads: Optional[int] = None  # None = runner heuristic
+    attach_block_max: bool = False
+
+
+def plan_emission(batch: int, n_comp: int, k: int, have_epilogue: bool) -> str:
+    """Which assist tier the indexer epilogue should emit this step.
+
+    ``n_comp``: compressed row length (post compress_ratio) - the
+    top-k kernel's N. Returns the tier name; the epilogue emits the
+    matching buffers and the next top-k launch routes on them.
+    """
+    if not have_epilogue:
+        return "rungs"  # closed-loop lines cost nothing to carry
+    if batch <= LIST_EMIT_MAX_B and n_comp >= LIST_EMIT_MIN_N:
+        return "list"  # latency-bound long rows: list pays big
+    return "counts"  # near-free tax, wins almost everywhere
+
+
+def pick_config(tier: str, batch: int, n_comp: int, k: int, num_sms: int) -> TopkRoute:
+    """Map (tier, B, N, K) to launch knobs. Pure function of shape."""
+    r = TopkRoute(tier=tier)
+    if tier == "none":
+        return r
+    if tier == "list":
+        if k <= SMALL_K_MAX:
+            r.num_threads = SMALL_K_LIST_THREADS
+        # list + block_max: miss rows fall back to a skip-walk instead
+        # of a dense re-scan (measured -19% on pro long chains).
+        r.attach_block_max = n_comp >= SKIP_MIN_N_COUNTS
+        return r
+    if tier == "counts":
+        r.attach_block_max = n_comp >= SKIP_MIN_N_COUNTS
+        return r
+    # rungs (vb)
+    if k <= RUNGS_BM_MAX_K and n_comp >= SKIP_MIN_N_RUNGS_FLASH:
+        r.attach_block_max = True
+    if n_comp >= 65536:
+        if batch * 8 <= num_sms // CS8_HALF_DEVICE:
+            r.cluster_size = 8
+        elif batch * 4 <= (num_sms * CS_HEADROOM_NUM) // CS_HEADROOM_DEN:
+            r.cluster_size = 4
+        elif batch * 2 <= (num_sms * CS_HEADROOM_NUM) // CS_HEADROOM_DEN:
+            r.cluster_size = 2
+    if r.attach_block_max and n_comp < SKIP_CS_MIN_N_RUNGS:
+        r.cluster_size = 1  # bm without cs below the split point
+    return r

From 07cac99e219d833140438373bb13a22fc7c4f3d9 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 11:26:38 -0700
Subject: [PATCH 050/117] [None][feat] FP4 indexer: bucketed candidate emission
 (v5 SoA contract)

emit_cand_bucketed: the GEMM epilogue emits the top-k consumer's
bucketed list directly - three fixed SoA segments (A=[0,segA) holds
>= t2, B=[segA,2segA) holds [t1,t2), C=[2segA,2segA+capC) holds
[t0,t1)), full segments spilling to the next looser one.

A/B use EXACT ballot claims (one lane0 atomic per hit-tile per class):
their prefixes must stay pad-free because the consumer's prefix math
derives lenA/lenB from the value counts alone. C keeps the claim-
window scheme (pads are legal there; window tails sentinel-fill BOTH
SoA columns). Every warp collective sits at the top level of the
warp-uniform bound gate - no collectives inside nested dynamic
branches. ctl [rows,4] = {n0 incl C pads, void, n1, n2}: n0
accumulates exact A/B placements plus C window claims, n1/n2 mirror
the L1 seed counters in the same flush redux, void marks C capacity
overflow (fall back to the row-level path). Cursors live in a caller-
zeroed cand_cur [rows,4].

Unit test checks segment invariants, exact n1/n2, full >= t0 coverage
(union of live entries, roomy caps) and void on tight caps; 4/4
parametrizations pass on B200.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  86 ++-
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 547 +++++++++++++++++-
 .../test_cute_dsl_fp4_paged_mqa_logits.py     | 171 ++++++
 3 files changed, 791 insertions(+), 13 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 0f0b12cae380..263e805a8554 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -8256,12 +8256,15 @@ def _compile(cls,
                      emit_seed_counts=False,
                      seed_packed=False,
                      emit_cand=False,
-                     cand_cap=5120):
+                     cand_cap=5120,
+                     emit_cand_bucketed=False,
+                     accept_cap=8192):
             """Compile kernel using fake tensors + TVM FFI."""
             key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
                    num_sms, num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
-                   emit_seed_counts, seed_packed, emit_cand, cand_cap)
+                   emit_seed_counts, seed_packed, emit_cand, cand_cap,
+                   emit_cand_bucketed, accept_cap)
             if key in cls.kernel_cache:
                 return
 
@@ -8354,6 +8357,27 @@ def _compile(cls,
                     cutlass.Int32, (cute.sym_int(), 2),
                     stride_order=(1, 0),
                     assumed_align=8)
+            cand_idx_fake = None
+            cand_cur_fake = None
+            if emit_cand_bucketed:
+                # bucketed SoA: cand slot reused as the fp32 VALUES tensor
+                wtot = 2 * accept_cap + cand_cap
+                cand_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Float32, (cute.sym_int(), wtot),
+                    stride_order=(1, 0),
+                    assumed_align=4)
+                cand_idx_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Int32, (cute.sym_int(), wtot),
+                    stride_order=(1, 0),
+                    assumed_align=4)
+                cand_ctl_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Int32, (cute.sym_int(), 4),
+                    stride_order=(1, 0),
+                    assumed_align=4)
+                cand_cur_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Int32, (cute.sym_int(), 4),
+                    stride_order=(1, 0),
+                    assumed_align=4)
             seed_thr_fake = None
             seed_counts_fake = None
             if emit_seed_counts:
@@ -8400,6 +8424,8 @@ def _compile(cls,
                 seed_packed=seed_packed,
                 emit_cand=emit_cand,
                 cand_cap=cand_cap,
+                emit_cand_bucketed=emit_cand_bucketed,
+                accept_cap=accept_cap,
             )
 
             compiled = cute.compile(
@@ -8422,6 +8448,8 @@ def _compile(cls,
                 seed_counts=seed_counts_fake,
                 cand=cand_fake,
                 cand_ctl=cand_ctl_fake,
+                cand_idx_t=cand_idx_fake,
+                cand_cur=cand_cur_fake,
                 options="--enable-tvm-ffi",
             )
             cls.kernel_cache[key] = compiled
@@ -8453,6 +8481,10 @@ def forward(
             emit_cand: bool = False,
             cand_out: Optional[torch.Tensor] = None,
             cand_ctl_out: Optional[torch.Tensor] = None,
+            emit_cand_bucketed: bool = False,
+            accept_cap: int = 8192,
+            cand_idx_out: Optional[torch.Tensor] = None,
+            cand_cur_out: Optional[torch.Tensor] = None,
         ) -> torch.Tensor:
             """Execute FP4 paged MQA logits kernel.
 
@@ -8628,15 +8660,54 @@ def forward(
                             "emit_cand requires a caller-zeroed cand_ctl_out "
                             "int32 [B*next_n, 2]")
                 cand_cap = cand_out.shape[1] // 2
+            elif emit_cand_bucketed:
+                assert emit_seed_counts, (
+                    "emit_cand_bucketed requires emit_seed_counts")
+                # SoA v5 contract: cand_out = fp32 VALUES [rows, 2*segA+capC],
+                # cand_idx_out = int32 positions (same width), cand_cur_out =
+                # int32 [rows, 4] cursors (caller-zeroed), cand_ctl_out =
+                # int32 [rows, 4] {n0, void, n1, n2} (caller-zeroed)
+                W = cand_out.shape[1]
+                assert (
+                    cand_out is not None and cand_out.dtype == torch.float32
+                    and cand_out.is_cuda and cand_out.is_contiguous()
+                    and cand_out.dim() == 2 and cand_out.shape[0] == B * next_n
+                    and W > 2 * accept_cap), (
+                        "bucketed requires cand_out fp32 [rows, 2*segA+capC]")
+                assert (cand_idx_out is not None
+                        and cand_idx_out.dtype == torch.int32
+                        and cand_idx_out.is_cuda
+                        and cand_idx_out.is_contiguous()
+                        and cand_idx_out.shape == cand_out.shape), (
+                            "bucketed requires cand_idx_out int32, same shape")
+                assert (cand_ctl_out is not None
+                        and cand_ctl_out.dtype == torch.int32
+                        and cand_ctl_out.is_cuda
+                        and cand_ctl_out.is_contiguous()
+                        and cand_ctl_out.shape == (B * next_n, 4)), (
+                            "bucketed requires caller-zeroed cand_ctl_out "
+                            "int32 [rows, 4]")
+                assert (cand_cur_out is not None
+                        and cand_cur_out.dtype == torch.int32
+                        and cand_cur_out.is_cuda
+                        and cand_cur_out.is_contiguous()
+                        and cand_cur_out.shape == (B * next_n, 4)), (
+                            "bucketed requires caller-zeroed cand_cur_out "
+                            "int32 [rows, 4]")
+                cand_cap = W - 2 * accept_cap
             else:
                 cand_out = None
                 cand_ctl_out = None
+            if not emit_cand_bucketed:
+                cand_idx_out = None
+                cand_cur_out = None
 
             # Compile if needed (fake tensors, no real data required)
             key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
                    num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
-                   emit_seed_counts, seed_packed, emit_cand, cand_cap)
+                   emit_seed_counts, seed_packed, emit_cand, cand_cap,
+                   emit_cand_bucketed, accept_cap)
             if key not in cls.kernel_cache:
                 cls._compile(
                     compute_block_kv,
@@ -8654,7 +8725,9 @@ def forward(
                     emit_seed_counts=emit_seed_counts,
                     seed_packed=seed_packed,
                     emit_cand=emit_cand,
-                    cand_cap=cand_cap)
+                    cand_cap=cand_cap,
+                    emit_cand_bucketed=emit_cand_bucketed,
+                    accept_cap=accept_cap)
             compiled = cls.kernel_cache[key]
 
             # TVM FFI: pass raw tensors, no dlpack/stream needed
@@ -8662,11 +8735,12 @@ def forward(
                 compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                          context_lens, schedule_meta, num_phys_blocks, B,
                          block_max_out, hit_stats_out, hit_bitmap, seed_thr,
-                         seed_counts_out, cand_out, cand_ctl_out)
+                         seed_counts_out, cand_out, cand_ctl_out, cand_idx_out,
+                         cand_cur_out)
                 return logits, block_max_out, hit_stats_out
             compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                      context_lens, schedule_meta, num_phys_blocks, B, None,
-                     None, None, None, None, None, None)
+                     None, None, None, None, None, None, None, None)
             return logits
 
     @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits",
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index c7b4fcaf1335..2481a5c5b8b3 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -472,6 +472,8 @@ def __init__(
         seed_packed: bool = False,
         emit_cand: bool = False,
         cand_cap: int = 5120,
+        emit_cand_bucketed: bool = False,
+        accept_cap: int = 8192,
     ):
         # Static FP4 invariants — see plan Sanity checklist.
         assert num_heads == 64, "FP4 kernel hardcodes num_heads=64 for TMEM/SMEM budget"
@@ -576,6 +578,21 @@ def __init__(
             raise ValueError("emit_cand requires emit_seed_counts (t_0 source)")
         self.emit_cand = emit_cand
         self.cand_cap = cand_cap
+        # emit_cand_bucketed (v5 list contract): three fixed SoA segments
+        # (A=[0,segA) holds >= t2, B=[segA,2segA) holds [t1,t2), C=
+        # [2segA,2segA+capC) holds [t0,t1)); a full segment spills to the
+        # next looser one. A/B use EXACT ballot claims (their prefixes
+        # must stay pad-free - the consumer's prefix math assumes it), C
+        # keeps the claim-window scheme (pads are legal there). Cursors
+        # live in caller-zeroed cand_cur [rows,4]; ctl [rows,4] carries
+        # {n0 incl C pads, void, n1, n2} with n1/n2 flushed from the L1
+        # seed counters.
+        if emit_cand_bucketed and not emit_seed_counts:
+            raise ValueError("emit_cand_bucketed requires emit_seed_counts")
+        if emit_cand_bucketed and emit_cand:
+            raise ValueError("emit_cand_bucketed and emit_cand are exclusive")
+        self.emit_cand_bucketed = emit_cand_bucketed
+        self.accept_cap = accept_cap
         # Per-warp claim window: one atomic claims (hits + CAND_WIN) slots;
         # subsequent hits consume the window latency-free. The unconsumed
         # tail is sentinel-filled (idx = -1) at q-transition/loop end, so
@@ -809,6 +826,8 @@ def __call__(
         seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 out, caller-zeroed
         cand: cute.Tensor = None,  # [num_rows, CAP*2] int32 {val bits, idx} pairs
         cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 {claimed, void}, zeroed
+        cand_idx_t: cute.Tensor = None,  # bucketed: [num_rows, 2*segA+capC] int32 SoA
+        cand_cur: cute.Tensor = None,  # bucketed: [num_rows, 4] int32 cursors, zeroed
     ):
         # Derive KV data and SF views from the fused uint8 buffer.
         # Fused layout per phys block: [data half_head_dim*phys_block_kv bytes]
@@ -1014,6 +1033,8 @@ class SharedStorage:
             seed_counts,
             cand,
             cand_ctl,
+            cand_idx_t,
+            cand_cur,
             self.cluster_layout_vmnk,
             self.a_smem_layout_staged,
             self.b_smem_layout_staged,
@@ -1075,7 +1096,7 @@ def _flush_hit_agg(
             hacc_cnt[t] = cutlass.Int32(0)
 
     @cute.jit
-    def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None):
+    def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None, cand_ctl=None):
         """Warp-redux the lane-local seed counters and fire one lane-0
         red.global.add per (t, threshold). Caller zero-initializes the
         count slots each step; cross-CTA totals accumulate atomically.
@@ -1099,6 +1120,15 @@ def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None):
                             cutlass.Int64(row) * cutlass.Int64(3) + cutlass.Int64(j)
                         ) * cutlass.Int64(4)
                         _red_global_add_s32(addr, w_cnt)
+                if cutlass.const_expr(self.emit_cand_bucketed):
+                    if j >= 1:
+                        # consumer contract: ctl = {n0, void, n1, n2}
+                        if meta_lane == cutlass.Int32(0):
+                            row_b = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t)
+                            ctl_a = cand_ctl.iterator.toint() + (
+                                cutlass.Int64(row_b) * cutlass.Int64(4) + cutlass.Int64(j + 1)
+                            ) * cutlass.Int64(4)
+                            _red_global_add_s32(ctl_a, w_cnt)
                 scnt[t * 3 + j] = cutlass.Int32(0)
         if cutlass.const_expr(self.seed_packed and spass is not None):
             # packed col 6: adaptive-skip pass count (lane0-accumulated,
@@ -1113,6 +1143,45 @@ def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None):
                     _red_global_add_f32(addr, cutlass.Float32(w_bp))
                 spass[t] = cutlass.Int32(0)
 
+    @cute.jit
+    def _flush_cand_window_bucketed(self, mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane):
+        """Sentinel-fill the unconsumed C-window tail in BOTH SoA columns
+        (score -inf, idx -1: the v5 consumer pads by score) and invalidate
+        the window. Segment C sits at base 2*segA in each row."""
+        next_n = cutlass.const_expr(self.next_n)
+        segA_f = cutlass.const_expr(self.accept_cap)
+        capC_f = cutlass.const_expr(self.cand_cap)
+        wtot_f = cutlass.const_expr(2 * self.accept_cap + self.cand_cap)
+        vbase_f = mCand.iterator.toint()
+        ibase_f = mCandIdx.iterator.toint()
+        for t in cutlass.range_constexpr(next_n):
+            if cwleft[t] > cutlass.Int32(0):
+                row_f = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t)
+                sl_f = cwbase[t] + meta_lane
+                if meta_lane < cwleft[t] and sl_f < cutlass.Int32(capC_f):
+                    off_f = (
+                        cutlass.Int64(row_f) * cutlass.Int64(wtot_f)
+                        + cutlass.Int64(2 * segA_f + sl_f)
+                    ) * cutlass.Int64(4)
+                    vp_f = cute.make_ptr(
+                        cutlass.Float32,
+                        vbase_f + off_f,
+                        cute.AddressSpace.gmem,
+                        assumed_align=4,
+                    )
+                    cute.make_tensor(vp_f, cute.make_layout((1,)))[0] = cutlass.Float32(
+                        _META_NEG_FLT_MAX
+                    )
+                    ip_f = cute.make_ptr(
+                        cutlass.Int32,
+                        ibase_f + off_f,
+                        cute.AddressSpace.gmem,
+                        assumed_align=4,
+                    )
+                    cute.make_tensor(ip_f, cute.make_layout((1,)))[0] = cutlass.Int32(-1)
+                cwbase[t] = cutlass.Int32(0)
+                cwleft[t] = cutlass.Int32(0)
+
     @cute.jit
     def _flush_cand_window(self, mCand, q_idx, cwbase, cwleft, meta_lane):
         """Sentinel-fill the unconsumed tail of each per-(warp, t) claim
@@ -1165,6 +1234,8 @@ def kernel(
         mSeedCounts: cute.Tensor,  # [num_rows, 3] int32 counts out (or None)
         mCand: cute.Tensor,  # [num_rows, CAP*2] int32 pair scatter (or None)
         mCandCtl: cute.Tensor,  # [num_rows, 2] int32 {claimed, void} (or None)
+        mCandIdx: cute.Tensor,  # bucketed: [num_rows, 2*segA+capC] int32 SoA (or None)
+        mCandCur: cute.Tensor,  # bucketed: [num_rows, 4] int32 cursors (or None)
         cluster_layout_vmnk: cute.Layout,
         a_smem_layout_staged: cute.ComposedLayout,
         b_smem_layout_staged: cute.ComposedLayout,
@@ -2176,7 +2247,7 @@ def kernel(
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
-                    if cutlass.const_expr(self.emit_cand):
+                    if cutlass.const_expr(self.emit_cand or self.emit_cand_bucketed):
                         cwbase = cute.make_fragment(next_n, cutlass.Int32)
                         cwleft = cute.make_fragment(next_n, cutlass.Int32)
                         for _i in cutlass.range_constexpr(next_n):
@@ -2249,7 +2320,12 @@ def kernel(
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
                                     self._flush_seed_counts(
-                                        mSeedCounts, q_idx_old, scnt, meta_lane, spass=spass
+                                        mSeedCounts,
+                                        q_idx_old,
+                                        scnt,
+                                        meta_lane,
+                                        spass=spass,
+                                        cand_ctl=mCandCtl,
                                     )
                                 # (re)load this q's thresholds - gated on
                                 # emit_seed_counts, NOT emit_cand: counts-
@@ -2263,6 +2339,11 @@ def kernel(
                                     self._flush_cand_window(
                                         mCand, q_idx_old, cwbase, cwleft, meta_lane
                                     )
+                            if cutlass.const_expr(self.emit_cand_bucketed):
+                                if q_idx_old < batch_size:
+                                    self._flush_cand_window_bucketed(
+                                        mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane
+                                    )
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 0 (kv_idx + 0)
@@ -2603,6 +2684,220 @@ def kernel(
                                         cute.make_tensor(iptr_c, cute.make_layout((1,)))[0] = kv_pos
                                     cwbase[t] = cwbase[t] + cnt_c
                                     cwleft[t] = cwleft[t] - cnt_c
+                            if cutlass.const_expr(self.emit_cand_bucketed):
+                                # v5 bucketed SoA: A/B EXACT ballot claims
+                                # (their prefixes must stay pad-free for
+                                # the consumer's prefix math), C keeps the
+                                # claim-window; a full segment spills to
+                                # the next looser one. Every warp
+                                # collective sits at the TOP level of this
+                                # warp-uniform bound gate - no collectives
+                                # inside nested dynamic branches (DSL).
+                                if r_bmax >= sthr[t * 3 + 0]:
+                                    segA_k = cutlass.const_expr(self.accept_cap)
+                                    capC_k = cutlass.const_expr(self.cand_cap)
+                                    wtot_k = cutlass.const_expr(2 * self.accept_cap + self.cand_cap)
+                                    row_k = q_idx * next_n + t
+                                    vb_k = mCand.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(wtot_k) * cutlass.Int64(4)
+                                    ib_k = mCandIdx.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(wtot_k) * cutlass.Int64(4)
+                                    cur_k = mCandCur.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(16)
+                                    ctl_k = mCandCtl.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(16)
+                                    lmk_k = (
+                                        cutlass.Uint32(1) << cutlass.Uint32(meta_lane)
+                                    ) - cutlass.Uint32(1)
+                                    # exclusive class predicates
+                                    pA_k = cutlass.Int32(0)
+                                    pB_k = cutlass.Int32(0)
+                                    pC_k = cutlass.Int32(0)
+                                    if meta_valid:
+                                        if f32_t >= sthr[t * 3 + 2]:
+                                            pA_k = cutlass.Int32(1)
+                                        if f32_t >= sthr[t * 3 + 1] and pA_k == cutlass.Int32(0):
+                                            pB_k = cutlass.Int32(1)
+                                        if (
+                                            f32_t >= sthr[t * 3 + 0]
+                                            and pA_k == cutlass.Int32(0)
+                                            and pB_k == cutlass.Int32(0)
+                                        ):
+                                            pC_k = cutlass.Int32(1)
+                                    # ---- A: exact claim ----
+                                    mA_k = cute.arch.vote_ballot_sync(pA_k != cutlass.Int32(0))
+                                    cntA_k = cutlass.Int32(cute.arch.popc(mA_k))
+                                    offA_k = cutlass.Int32(cute.arch.popc(mA_k & lmk_k))
+                                    baseA_k = cutlass.Int32(0)
+                                    if meta_lane == cutlass.Int32(0) and cntA_k > cutlass.Int32(0):
+                                        baseA_k = _atom_global_add_s32(cur_k, cntA_k)
+                                    baseA_k = cute.arch.shuffle_sync(baseA_k, cutlass.Int32(0))
+                                    slotA_k = baseA_k + offA_k
+                                    spA_k = cutlass.Int32(0)
+                                    if pA_k != cutlass.Int32(0) and slotA_k >= cutlass.Int32(
+                                        segA_k
+                                    ):
+                                        spA_k = cutlass.Int32(1)
+                                    if pA_k != cutlass.Int32(0) and slotA_k < cutlass.Int32(segA_k):
+                                        vp_k = cute.make_ptr(
+                                            cutlass.Float32,
+                                            vb_k + cutlass.Int64(slotA_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(vp_k, cute.make_layout((1,)))[0] = f32_t
+                                        ip_k = cute.make_ptr(
+                                            cutlass.Int32,
+                                            ib_k + cutlass.Int64(slotA_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(ip_k, cute.make_layout((1,)))[0] = kv_pos
+                                    # ---- B: exact claim (native + A spill) ----
+                                    pBe_k = cutlass.Int32(0)
+                                    if pB_k != cutlass.Int32(0) or spA_k != cutlass.Int32(0):
+                                        pBe_k = cutlass.Int32(1)
+                                    mB_k = cute.arch.vote_ballot_sync(pBe_k != cutlass.Int32(0))
+                                    cntB_k = cutlass.Int32(cute.arch.popc(mB_k))
+                                    offB_k = cutlass.Int32(cute.arch.popc(mB_k & lmk_k))
+                                    baseB_k = cutlass.Int32(0)
+                                    if meta_lane == cutlass.Int32(0) and cntB_k > cutlass.Int32(0):
+                                        baseB_k = _atom_global_add_s32(
+                                            cur_k + cutlass.Int64(4), cntB_k
+                                        )
+                                    baseB_k = cute.arch.shuffle_sync(baseB_k, cutlass.Int32(0))
+                                    slotB_k = baseB_k + offB_k
+                                    spB_k = cutlass.Int32(0)
+                                    if pBe_k != cutlass.Int32(0) and slotB_k >= cutlass.Int32(
+                                        segA_k
+                                    ):
+                                        spB_k = cutlass.Int32(1)
+                                    if pBe_k != cutlass.Int32(0) and slotB_k < cutlass.Int32(
+                                        segA_k
+                                    ):
+                                        vp2_k = cute.make_ptr(
+                                            cutlass.Float32,
+                                            vb_k
+                                            + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(vp2_k, cute.make_layout((1,)))[0] = f32_t
+                                        ip2_k = cute.make_ptr(
+                                            cutlass.Int32,
+                                            ib_k
+                                            + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(ip2_k, cute.make_layout((1,)))[0] = kv_pos
+                                    # n0 += exact placements in A and B
+                                    plc_k = (
+                                        cntA_k
+                                        - cutlass.Int32(
+                                            cute.arch.popc(
+                                                cute.arch.vote_ballot_sync(
+                                                    spA_k != cutlass.Int32(0)
+                                                )
+                                            )
+                                        )
+                                    ) + (
+                                        cntB_k
+                                        - cutlass.Int32(
+                                            cute.arch.popc(
+                                                cute.arch.vote_ballot_sync(
+                                                    spB_k != cutlass.Int32(0)
+                                                )
+                                            )
+                                        )
+                                    )
+                                    if meta_lane == cutlass.Int32(0) and plc_k > cutlass.Int32(0):
+                                        _atom_global_add_s32(ctl_k, plc_k)
+                                    # ---- C: claim window (native + B spill) ----
+                                    pCe_k = cutlass.Int32(0)
+                                    if pC_k != cutlass.Int32(0) or spB_k != cutlass.Int32(0):
+                                        pCe_k = cutlass.Int32(1)
+                                    mC_k = cute.arch.vote_ballot_sync(pCe_k != cutlass.Int32(0))
+                                    cntC_k = cutlass.Int32(cute.arch.popc(mC_k))
+                                    offC_k = cutlass.Int32(cute.arch.popc(mC_k & lmk_k))
+                                    if cntC_k > cwleft[t]:
+                                        # sentinel-fill the old window tail
+                                        # (BOTH columns: the consumer pads
+                                        # by score -inf, idx -1)
+                                        slo_k = cwbase[t] + meta_lane
+                                        if meta_lane < cwleft[t] and slo_k < cutlass.Int32(capC_k):
+                                            vpo_k = cute.make_ptr(
+                                                cutlass.Float32,
+                                                vb_k
+                                                + cutlass.Int64(2 * segA_k + slo_k)
+                                                * cutlass.Int64(4),
+                                                cute.AddressSpace.gmem,
+                                                assumed_align=4,
+                                            )
+                                            cute.make_tensor(vpo_k, cute.make_layout((1,)))[0] = (
+                                                cutlass.Float32(_META_NEG_FLT_MAX)
+                                            )
+                                            ipo_k = cute.make_ptr(
+                                                cutlass.Int32,
+                                                ib_k
+                                                + cutlass.Int64(2 * segA_k + slo_k)
+                                                * cutlass.Int64(4),
+                                                cute.AddressSpace.gmem,
+                                                assumed_align=4,
+                                            )
+                                            cute.make_tensor(ipo_k, cute.make_layout((1,)))[0] = (
+                                                cutlass.Int32(-1)
+                                            )
+                                        mC2_k = cntC_k + cutlass.Int32(self.CAND_WIN)
+                                        nbC_k = cutlass.Int32(0)
+                                        if meta_lane == cutlass.Int32(0):
+                                            nbC_k = _atom_global_add_s32(
+                                                cur_k + cutlass.Int64(8), mC2_k
+                                            )
+                                            _atom_global_add_s32(ctl_k, mC2_k)
+                                            if nbC_k + mC2_k > cutlass.Int32(
+                                                capC_k
+                                            ) and nbC_k <= cutlass.Int32(capC_k):
+                                                vdp_k = cute.make_ptr(
+                                                    cutlass.Int32,
+                                                    ctl_k + cutlass.Int64(4),
+                                                    cute.AddressSpace.gmem,
+                                                    assumed_align=4,
+                                                )
+                                                cute.make_tensor(vdp_k, cute.make_layout((1,)))[
+                                                    0
+                                                ] = cutlass.Int32(1)
+                                        nbC_k = cute.arch.shuffle_sync(nbC_k, cutlass.Int32(0))
+                                        cwbase[t] = nbC_k
+                                        cwleft[t] = mC2_k
+                                    slotC_k = cwbase[t] + offC_k
+                                    if pCe_k != cutlass.Int32(0) and slotC_k < cutlass.Int32(
+                                        capC_k
+                                    ):
+                                        vpc_k = cute.make_ptr(
+                                            cutlass.Float32,
+                                            vb_k
+                                            + cutlass.Int64(2 * segA_k + slotC_k)
+                                            * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(vpc_k, cute.make_layout((1,)))[0] = f32_t
+                                        ipc_k = cute.make_ptr(
+                                            cutlass.Int32,
+                                            ib_k
+                                            + cutlass.Int64(2 * segA_k + slotC_k)
+                                            * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(ipc_k, cute.make_layout((1,)))[0] = kv_pos
+                                    cwbase[t] = cwbase[t] + cntC_k
+                                    cwleft[t] = cwleft[t] - cntC_k
                             if cutlass.const_expr(self.emit_hit_stats):
                                 # Lane-local accumulation, fully BRANCHLESS
                                 # (data-dependent `if meta_hit` compiled to
@@ -2674,10 +2969,17 @@ def kernel(
                         )
                 if cutlass.const_expr(self.emit_seed_counts):
                     if q_idx < batch_size:
-                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane, spass=spass)
+                        self._flush_seed_counts(
+                            mSeedCounts, q_idx, scnt, meta_lane, spass=spass, cand_ctl=mCandCtl
+                        )
                 if cutlass.const_expr(self.emit_cand):
                     if q_idx < batch_size:
                         self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane)
+                if cutlass.const_expr(self.emit_cand_bucketed):
+                    if q_idx < batch_size:
+                        self._flush_cand_window_bucketed(
+                            mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane
+                        )
 
                 # Release last Q stage (WG 0)
                 if q_idx < batch_size:
@@ -2738,7 +3040,7 @@ def kernel(
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
-                    if cutlass.const_expr(self.emit_cand):
+                    if cutlass.const_expr(self.emit_cand or self.emit_cand_bucketed):
                         cwbase = cute.make_fragment(next_n, cutlass.Int32)
                         cwleft = cute.make_fragment(next_n, cutlass.Int32)
                         for _i in cutlass.range_constexpr(next_n):
@@ -2811,7 +3113,12 @@ def kernel(
                             if cutlass.const_expr(self.emit_seed_counts):
                                 if q_idx_old < batch_size:
                                     self._flush_seed_counts(
-                                        mSeedCounts, q_idx_old, scnt, meta_lane, spass=spass
+                                        mSeedCounts,
+                                        q_idx_old,
+                                        scnt,
+                                        meta_lane,
+                                        spass=spass,
+                                        cand_ctl=mCandCtl,
                                     )
                                 # (re)load this q's thresholds - gated on
                                 # emit_seed_counts, NOT emit_cand (see the
@@ -2824,6 +3131,11 @@ def kernel(
                                     self._flush_cand_window(
                                         mCand, q_idx_old, cwbase, cwleft, meta_lane
                                     )
+                            if cutlass.const_expr(self.emit_cand_bucketed):
+                                if q_idx_old < batch_size:
+                                    self._flush_cand_window_bucketed(
+                                        mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane
+                                    )
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 1 (kv_idx + 1)
@@ -3158,6 +3470,220 @@ def kernel(
                                         cute.make_tensor(iptr_c, cute.make_layout((1,)))[0] = kv_pos
                                     cwbase[t] = cwbase[t] + cnt_c
                                     cwleft[t] = cwleft[t] - cnt_c
+                            if cutlass.const_expr(self.emit_cand_bucketed):
+                                # v5 bucketed SoA: A/B EXACT ballot claims
+                                # (their prefixes must stay pad-free for
+                                # the consumer's prefix math), C keeps the
+                                # claim-window; a full segment spills to
+                                # the next looser one. Every warp
+                                # collective sits at the TOP level of this
+                                # warp-uniform bound gate - no collectives
+                                # inside nested dynamic branches (DSL).
+                                if r_bmax >= sthr[t * 3 + 0]:
+                                    segA_k = cutlass.const_expr(self.accept_cap)
+                                    capC_k = cutlass.const_expr(self.cand_cap)
+                                    wtot_k = cutlass.const_expr(2 * self.accept_cap + self.cand_cap)
+                                    row_k = q_idx * next_n + t
+                                    vb_k = mCand.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(wtot_k) * cutlass.Int64(4)
+                                    ib_k = mCandIdx.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(wtot_k) * cutlass.Int64(4)
+                                    cur_k = mCandCur.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(16)
+                                    ctl_k = mCandCtl.iterator.toint() + cutlass.Int64(
+                                        row_k
+                                    ) * cutlass.Int64(16)
+                                    lmk_k = (
+                                        cutlass.Uint32(1) << cutlass.Uint32(meta_lane)
+                                    ) - cutlass.Uint32(1)
+                                    # exclusive class predicates
+                                    pA_k = cutlass.Int32(0)
+                                    pB_k = cutlass.Int32(0)
+                                    pC_k = cutlass.Int32(0)
+                                    if meta_valid:
+                                        if f32_t >= sthr[t * 3 + 2]:
+                                            pA_k = cutlass.Int32(1)
+                                        if f32_t >= sthr[t * 3 + 1] and pA_k == cutlass.Int32(0):
+                                            pB_k = cutlass.Int32(1)
+                                        if (
+                                            f32_t >= sthr[t * 3 + 0]
+                                            and pA_k == cutlass.Int32(0)
+                                            and pB_k == cutlass.Int32(0)
+                                        ):
+                                            pC_k = cutlass.Int32(1)
+                                    # ---- A: exact claim ----
+                                    mA_k = cute.arch.vote_ballot_sync(pA_k != cutlass.Int32(0))
+                                    cntA_k = cutlass.Int32(cute.arch.popc(mA_k))
+                                    offA_k = cutlass.Int32(cute.arch.popc(mA_k & lmk_k))
+                                    baseA_k = cutlass.Int32(0)
+                                    if meta_lane == cutlass.Int32(0) and cntA_k > cutlass.Int32(0):
+                                        baseA_k = _atom_global_add_s32(cur_k, cntA_k)
+                                    baseA_k = cute.arch.shuffle_sync(baseA_k, cutlass.Int32(0))
+                                    slotA_k = baseA_k + offA_k
+                                    spA_k = cutlass.Int32(0)
+                                    if pA_k != cutlass.Int32(0) and slotA_k >= cutlass.Int32(
+                                        segA_k
+                                    ):
+                                        spA_k = cutlass.Int32(1)
+                                    if pA_k != cutlass.Int32(0) and slotA_k < cutlass.Int32(segA_k):
+                                        vp_k = cute.make_ptr(
+                                            cutlass.Float32,
+                                            vb_k + cutlass.Int64(slotA_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(vp_k, cute.make_layout((1,)))[0] = f32_t
+                                        ip_k = cute.make_ptr(
+                                            cutlass.Int32,
+                                            ib_k + cutlass.Int64(slotA_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(ip_k, cute.make_layout((1,)))[0] = kv_pos
+                                    # ---- B: exact claim (native + A spill) ----
+                                    pBe_k = cutlass.Int32(0)
+                                    if pB_k != cutlass.Int32(0) or spA_k != cutlass.Int32(0):
+                                        pBe_k = cutlass.Int32(1)
+                                    mB_k = cute.arch.vote_ballot_sync(pBe_k != cutlass.Int32(0))
+                                    cntB_k = cutlass.Int32(cute.arch.popc(mB_k))
+                                    offB_k = cutlass.Int32(cute.arch.popc(mB_k & lmk_k))
+                                    baseB_k = cutlass.Int32(0)
+                                    if meta_lane == cutlass.Int32(0) and cntB_k > cutlass.Int32(0):
+                                        baseB_k = _atom_global_add_s32(
+                                            cur_k + cutlass.Int64(4), cntB_k
+                                        )
+                                    baseB_k = cute.arch.shuffle_sync(baseB_k, cutlass.Int32(0))
+                                    slotB_k = baseB_k + offB_k
+                                    spB_k = cutlass.Int32(0)
+                                    if pBe_k != cutlass.Int32(0) and slotB_k >= cutlass.Int32(
+                                        segA_k
+                                    ):
+                                        spB_k = cutlass.Int32(1)
+                                    if pBe_k != cutlass.Int32(0) and slotB_k < cutlass.Int32(
+                                        segA_k
+                                    ):
+                                        vp2_k = cute.make_ptr(
+                                            cutlass.Float32,
+                                            vb_k
+                                            + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(vp2_k, cute.make_layout((1,)))[0] = f32_t
+                                        ip2_k = cute.make_ptr(
+                                            cutlass.Int32,
+                                            ib_k
+                                            + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(ip2_k, cute.make_layout((1,)))[0] = kv_pos
+                                    # n0 += exact placements in A and B
+                                    plc_k = (
+                                        cntA_k
+                                        - cutlass.Int32(
+                                            cute.arch.popc(
+                                                cute.arch.vote_ballot_sync(
+                                                    spA_k != cutlass.Int32(0)
+                                                )
+                                            )
+                                        )
+                                    ) + (
+                                        cntB_k
+                                        - cutlass.Int32(
+                                            cute.arch.popc(
+                                                cute.arch.vote_ballot_sync(
+                                                    spB_k != cutlass.Int32(0)
+                                                )
+                                            )
+                                        )
+                                    )
+                                    if meta_lane == cutlass.Int32(0) and plc_k > cutlass.Int32(0):
+                                        _atom_global_add_s32(ctl_k, plc_k)
+                                    # ---- C: claim window (native + B spill) ----
+                                    pCe_k = cutlass.Int32(0)
+                                    if pC_k != cutlass.Int32(0) or spB_k != cutlass.Int32(0):
+                                        pCe_k = cutlass.Int32(1)
+                                    mC_k = cute.arch.vote_ballot_sync(pCe_k != cutlass.Int32(0))
+                                    cntC_k = cutlass.Int32(cute.arch.popc(mC_k))
+                                    offC_k = cutlass.Int32(cute.arch.popc(mC_k & lmk_k))
+                                    if cntC_k > cwleft[t]:
+                                        # sentinel-fill the old window tail
+                                        # (BOTH columns: the consumer pads
+                                        # by score -inf, idx -1)
+                                        slo_k = cwbase[t] + meta_lane
+                                        if meta_lane < cwleft[t] and slo_k < cutlass.Int32(capC_k):
+                                            vpo_k = cute.make_ptr(
+                                                cutlass.Float32,
+                                                vb_k
+                                                + cutlass.Int64(2 * segA_k + slo_k)
+                                                * cutlass.Int64(4),
+                                                cute.AddressSpace.gmem,
+                                                assumed_align=4,
+                                            )
+                                            cute.make_tensor(vpo_k, cute.make_layout((1,)))[0] = (
+                                                cutlass.Float32(_META_NEG_FLT_MAX)
+                                            )
+                                            ipo_k = cute.make_ptr(
+                                                cutlass.Int32,
+                                                ib_k
+                                                + cutlass.Int64(2 * segA_k + slo_k)
+                                                * cutlass.Int64(4),
+                                                cute.AddressSpace.gmem,
+                                                assumed_align=4,
+                                            )
+                                            cute.make_tensor(ipo_k, cute.make_layout((1,)))[0] = (
+                                                cutlass.Int32(-1)
+                                            )
+                                        mC2_k = cntC_k + cutlass.Int32(self.CAND_WIN)
+                                        nbC_k = cutlass.Int32(0)
+                                        if meta_lane == cutlass.Int32(0):
+                                            nbC_k = _atom_global_add_s32(
+                                                cur_k + cutlass.Int64(8), mC2_k
+                                            )
+                                            _atom_global_add_s32(ctl_k, mC2_k)
+                                            if nbC_k + mC2_k > cutlass.Int32(
+                                                capC_k
+                                            ) and nbC_k <= cutlass.Int32(capC_k):
+                                                vdp_k = cute.make_ptr(
+                                                    cutlass.Int32,
+                                                    ctl_k + cutlass.Int64(4),
+                                                    cute.AddressSpace.gmem,
+                                                    assumed_align=4,
+                                                )
+                                                cute.make_tensor(vdp_k, cute.make_layout((1,)))[
+                                                    0
+                                                ] = cutlass.Int32(1)
+                                        nbC_k = cute.arch.shuffle_sync(nbC_k, cutlass.Int32(0))
+                                        cwbase[t] = nbC_k
+                                        cwleft[t] = mC2_k
+                                    slotC_k = cwbase[t] + offC_k
+                                    if pCe_k != cutlass.Int32(0) and slotC_k < cutlass.Int32(
+                                        capC_k
+                                    ):
+                                        vpc_k = cute.make_ptr(
+                                            cutlass.Float32,
+                                            vb_k
+                                            + cutlass.Int64(2 * segA_k + slotC_k)
+                                            * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(vpc_k, cute.make_layout((1,)))[0] = f32_t
+                                        ipc_k = cute.make_ptr(
+                                            cutlass.Int32,
+                                            ib_k
+                                            + cutlass.Int64(2 * segA_k + slotC_k)
+                                            * cutlass.Int64(4),
+                                            cute.AddressSpace.gmem,
+                                            assumed_align=4,
+                                        )
+                                        cute.make_tensor(ipc_k, cute.make_layout((1,)))[0] = kv_pos
+                                    cwbase[t] = cwbase[t] + cntC_k
+                                    cwleft[t] = cwleft[t] - cntC_k
                             if cutlass.const_expr(self.emit_hit_stats):
                                 # Lane-local accumulation, fully BRANCHLESS
                                 # (data-dependent `if meta_hit` compiled to
@@ -3229,10 +3755,17 @@ def kernel(
                         )
                 if cutlass.const_expr(self.emit_seed_counts):
                     if q_idx < batch_size:
-                        self._flush_seed_counts(mSeedCounts, q_idx, scnt, meta_lane, spass=spass)
+                        self._flush_seed_counts(
+                            mSeedCounts, q_idx, scnt, meta_lane, spass=spass, cand_ctl=mCandCtl
+                        )
                 if cutlass.const_expr(self.emit_cand):
                     if q_idx < batch_size:
                         self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane)
+                if cutlass.const_expr(self.emit_cand_bucketed):
+                    if q_idx < batch_size:
+                        self._flush_cand_window_bucketed(
+                            mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane
+                        )
 
                 # Release last Q stage (WG 1)
                 if q_idx < batch_size:
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
index 9ac8c880846a..3b7bb55b4bab 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
@@ -1579,3 +1579,174 @@ def dg_fn(data=data, dg_ctx_2d=dg_ctx_2d, q_fp4_dg=q_fp4_dg):
         varlen=args.varlen,
         block_kv=args.block_kv,
     )
+
+
+@skip_not_sm100
+@pytest.mark.parametrize("batch_size", [1, 4])
+@pytest.mark.parametrize("next_n", [1, 2])
+@pytest.mark.parametrize("avg_ctx", [4096, 4224])
+@pytest.mark.parametrize("phys_block_kv", [64, 128])
+@pytest.mark.parametrize("cap_mode", ["roomy", "tight"])
+def test_cute_dsl_fp4_paged_mqa_logits_cand_bucketed(
+    batch_size,
+    next_n,
+    avg_ctx,
+    phys_block_kv,
+    cap_mode,
+):
+    """emit_cand_bucketed (v5 SoA contract): three fixed segments with
+    pad-free A/B prefixes and spill-to-looser, ctl {n0, void, n1, n2}
+    with n1/n2 mirrored from the seed counters, C-window pads carrying
+    score -inf / idx -1. All invariants recomputed from the kernel's
+    own logits."""
+    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner
+
+    torch.manual_seed(23)
+    torch.cuda.manual_seed(23)
+    num_heads, head_dim = 64, 128
+    max_model_len = max(avg_ctx * 2, 2048)
+    device = "cuda"
+    context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device)
+    num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv)
+    num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2
+    max_blocks_per_seq = int(num_blocks_per_seq.max().item())
+    block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device)
+    pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32)
+    off = 0
+    for i, n_blks in enumerate(num_blocks_per_seq.tolist()):
+        block_table[i, :n_blks] = pool[off : off + n_blks]
+        off += n_blks
+    q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16)
+    kv_cache = torch.randn(
+        (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16
+    )
+    weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32)
+    q_packed, sf_q_packed = per_token_cast_to_fp4(
+        q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True
+    )
+    q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2)
+    sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads)
+    remove_online_sf_transpose = phys_block_kv == 128
+    kv_fused, _ = kv_cache_cast_to_fp4(
+        kv_cache, remove_online_sf_transpose=remove_online_sf_transpose
+    )
+    DG_METADATA_BLOCK_KV = 64
+    num_sms = deep_gemm.get_num_sms()
+    schedule_meta = deep_gemm.get_paged_mqa_logits_metadata(
+        context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms
+    )
+    aligned_max_ctx = align(max_model_len, 256)
+    nb_pad = aligned_max_ctx // 128
+    num_rows = batch_size * next_n
+    nan = float("nan")
+    block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device)
+    common = dict(
+        num_epi_subtiles=1,
+        epi_dtype=torch.float32,
+        output_dtype=torch.bfloat16,
+        remove_online_sf_transpose=remove_online_sf_transpose,
+    )
+    logits0, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward(
+        q_fp4,
+        sf_q,
+        kv_fused,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_model_len,
+        emit_block_meta=True,
+        emit_hit_stats=False,
+        block_max_out=block_max,
+        **common,
+    )
+    torch.cuda.synchronize()
+    lf0 = logits0.float()
+    seed_row = torch.zeros((num_rows, 8), dtype=torch.float32, device=device)
+    for row in range(num_rows):
+        ctx = int(context_lens[row // next_n].item())
+        vals = lf0[row, :ctx]
+        seed_row[row, 0] = torch.quantile(vals, 0.60)
+        seed_row[row, 1] = torch.quantile(vals, 0.90)
+        seed_row[row, 2] = torch.quantile(vals, 0.99)
+    # segment caps: roomy fits everything; tight forces A/B spill and a
+    # C-window void
+    if cap_mode == "roomy":
+        segA, capC = 2048, 4096
+    else:
+        segA, capC = 32, 128
+    W = 2 * segA + capC
+    cand_vals = torch.full((num_rows, W), nan, dtype=torch.float32, device=device)
+    cand_idx = torch.full((num_rows, W), -7, dtype=torch.int32, device=device)
+    cand_ctl = torch.zeros((num_rows, 4), dtype=torch.int32, device=device)
+    cand_cur = torch.zeros((num_rows, 4), dtype=torch.int32, device=device)
+    block_max.fill_(nan)
+    logits, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward(
+        q_fp4,
+        sf_q,
+        kv_fused,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_model_len,
+        emit_block_meta=True,
+        emit_hit_stats=False,
+        block_max_out=block_max,
+        emit_seed_counts=True,
+        seed_thr=seed_row,
+        emit_cand_bucketed=True,
+        accept_cap=segA,
+        cand_out=cand_vals,
+        cand_idx_out=cand_idx,
+        cand_ctl_out=cand_ctl,
+        cand_cur_out=cand_cur,
+        **common,
+    )
+    torch.cuda.synchronize()
+    lf = logits.float()
+    for row in range(num_rows):
+        ctx = int(context_lens[row // next_n].item())
+        t0, t1, t2 = (float(seed_row[row, j]) for j in range(3))
+        v = lf[row, :ctx]
+        n0_ref = int((v >= t0).sum())
+        n1_ref = int((v >= t1).sum())
+        n2_ref = int((v >= t2).sum())
+        n0c, voidc, n1c, n2c = (int(cand_ctl[row, j]) for j in range(4))
+        tag = f"row={row} caps=({segA},{capC}) refs=({n0_ref},{n1_ref},{n2_ref})"
+        assert n1c == n1_ref and n2c == n2_ref, f"n1/n2 mismatch {tag} got {n1c},{n2c}"
+        curA, curB, curC = (int(cand_cur[row, j]) for j in range(3))
+        lenA = min(n2_ref, segA)
+        lenB = min(n1_ref - n2_ref + max(n2_ref - segA, 0), segA)
+        assert min(curA, segA) >= lenA or curA == n2_ref, f"curA {curA} {tag}"
+        # A prefix: pad-free, every entry >= t2, positions valid + unique
+        pa = cand_idx[row, :lenA]
+        va = cand_vals[row, :lenA]
+        assert (pa >= 0).all() and (pa < ctx).all(), f"A idx {tag}"
+        assert (va >= t2).all(), f"A vals {tag}"
+        got_a = lf[row, pa.long()]
+        torch.testing.assert_close(got_a, va, atol=0.0, rtol=0.0)
+        # B prefix: pad-free, [t1, t2) or A-spill (>= t2)
+        pb = cand_idx[row, segA : segA + lenB]
+        vb = cand_vals[row, segA : segA + lenB]
+        assert (pb >= 0).all() and (pb < ctx).all(), f"B idx {tag}"
+        assert (vb >= t1).all(), f"B vals {tag}"
+        torch.testing.assert_close(lf[row, pb.long()], vb, atol=0.0, rtol=0.0)
+        if voidc == 0:
+            # full coverage: union of live entries == the >= t0 set
+            lenC = n0c - lenA - lenB
+            pc = cand_idx[row, 2 * segA : 2 * segA + lenC]
+            vc = cand_vals[row, 2 * segA : 2 * segA + lenC]
+            live = pc >= 0
+            assert (vc[live] >= t0).all(), f"C vals {tag}"
+            # pads carry -FLT_MAX (never ranks; the emu uses -inf, the
+            # kernel the finite sentinel - both satisfy the contract)
+            assert (vc[~live] <= -3e38).all(), f"C pads {tag}"
+            allp = torch.cat([pa, pb, pc[live]])
+            assert allp.unique().numel() == allp.numel() == n0_ref, (
+                f"coverage {tag}: {allp.unique().numel()} vs {n0_ref}"
+            )
+        else:
+            assert cap_mode == "tight", f"unexpected void {tag}"
+    if cap_mode == "tight":
+        assert int(cand_ctl[:, 1].sum()) > 0, "tight caps never voided"

From 321f8eb2c9a1e364027b0725125a7fc2d8d7a772 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 27 Jul 2026 12:06:17 -0700
Subject: [PATCH 051/117] [None][feat] GVR: emission-assisted tiers on the
 production top-k op

trtllm::cute_dsl_gvr_topk_decode grows the ext face the emission suite
feeds: packed seed row ([rows,8] -> counts tier; [rows,3] -> closed-
loop rungs tier), bucketed SoA candidate list (+ ctl) for the list
tier, block_max prefix, xstate closed-loop publish, plus num_threads /
accept_cap / kc_override knobs. The runner derives the mode from which
tensors arrive (mirrors gvr_routing.plan_emission), compiles the
matching kernel flavor (fakes and ctor mirror the validated harness
wrapper, incl. the 3-rung slot layout) and applies the small-K list
512-thread rule. Ext tiers are single-CTA/sort-path only (asserted
against LB). Face test: counts / list / rungs / counts+block_max all
bit-exact vs torch.topk on B200 with closed-loop state published.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  | 132 +++++++++++++++++-
 1 file changed, 129 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 263e805a8554..bf2fa3df843b 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -6716,12 +6716,22 @@ def _compile(
             return_output_values: bool,
             cluster_size: int,
             seqlen_sorted: bool,
+            enable_block_skip: bool = False,
+            use_ext_counts: bool = False,
+            emit_xstate: bool = False,
+            use_ext_cand: bool = False,
+            ext_rungs: bool = False,
+            cand_cap: int = 5120,
+            accept_cap: Optional[int] = None,
+            kc_override: Optional[int] = None,
         ) -> tuple:
             key = (dtype, top_k, next_n, enable_unroll_4, enable_phase3_unroll,
                    use_constant_hint, min_blocks_per_mp, use_256bit_load,
                    num_threads_per_block, enable_warp_parallel_reduce,
                    compress_ratio, return_output_values, cluster_size,
-                   seqlen_sorted)
+                   seqlen_sorted, enable_block_skip, use_ext_counts,
+                   emit_xstate, use_ext_cand, ext_rungs, cand_cap, accept_cap,
+                   kc_override)
             if key in cls.kernel_cache:
                 return key
             n_rows = cute.sym_int()
@@ -6754,6 +6764,32 @@ def _compile(
             order_row_fake = (cute.runtime.make_fake_compact_tensor(
                 cutlass.Int32,
                 (n_batch, ), stride_order=(0, )) if seqlen_sorted else None)
+            # emission-assisted tiers (list/counts/rungs, see
+            # gvr_routing): fake shapes mirror the harness wrapper
+            block_max_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Float32, (n_rows, cute.sym_int()),
+                stride_order=(1, 0),
+                assumed_align=16) if enable_block_skip else None)
+            seed_thr_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Float32, (n_rows, 8 if use_ext_counts else 3),
+                stride_order=(1, 0),
+                assumed_align=4) if (use_ext_counts or ext_rungs) else None)
+            xstate_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Float32, (n_rows, 8),
+                stride_order=(1, 0),
+                assumed_align=4) if emit_xstate else None)
+            cand_vals_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Float32, (n_rows, cand_cap),
+                stride_order=(1, 0),
+                assumed_align=4) if use_ext_cand else None)
+            cand_idx_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Int32, (n_rows, cand_cap),
+                stride_order=(1, 0),
+                assumed_align=4) if use_ext_cand else None)
+            cand_ctl_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Int32, (n_rows, 4),
+                stride_order=(1, 0),
+                assumed_align=8) if use_ext_cand else None)
             fake_stream = cute.runtime.make_fake_stream(
                 use_tvm_ffi_env_stream=True)
 
@@ -6772,6 +6808,18 @@ def _compile(
                 return_output_values=return_output_values,
                 cluster_size=cluster_size,
                 seqlen_sorted=seqlen_sorted,
+                enable_block_skip=enable_block_skip,
+                use_ext_counts=use_ext_counts,
+                emit_xstate=emit_xstate,
+                use_ext_cand=use_ext_cand,
+                ext_rungs=ext_rungs,
+                cand_cap=cand_cap,
+                accept_cap=accept_cap,
+                kc_override=kc_override,
+                # ext modes need 3 rung slots (M_thr == 3); the qfrac
+                # VALUES are irrelevant (P1b skipped), only slot count
+                r0_qfracs=((0.85, 0.35) if
+                           (use_ext_counts or ext_rungs) else None),
             )
             cls.kernel_cache[key] = cute.compile(
                 kernel,
@@ -6782,6 +6830,13 @@ def _compile(
                 out_indices_fake,
                 order_row_fake,
                 stream=fake_stream,
+                block_max=block_max_fake,
+                seed_thr=seed_thr_fake,
+                seed_counts=None,
+                xstate=xstate_fake,
+                cand_vals=cand_vals_fake,
+                cand_idx=cand_idx_fake,
+                cand_ctl=cand_ctl_fake,
                 options="--enable-tvm-ffi",
             )
             logger.debug(f"[compile cute_dsl gvr_topk_decode] {key}")
@@ -6891,6 +6946,15 @@ def forward(
             order_row: Optional[torch.Tensor] = None,
             counters: Optional[torch.Tensor] = None,
             max_batch_size: Optional[int] = None,
+            seed_thr: Optional[torch.Tensor] = None,
+            xstate: Optional[torch.Tensor] = None,
+            cand_vals: Optional[torch.Tensor] = None,
+            cand_idx: Optional[torch.Tensor] = None,
+            cand_ctl: Optional[torch.Tensor] = None,
+            block_max: Optional[torch.Tensor] = None,
+            num_threads: Optional[int] = None,
+            accept_cap: Optional[int] = None,
+            kc_override: Optional[int] = None,
         ) -> None:
             """Three paths, picked by ``(counters, order_row)``:
 
@@ -7013,6 +7077,31 @@ def forward(
                 ), ("order_row must be int32, CUDA, shape == seq_lens.shape "
                     f"(={tuple(seq_lens.shape)}); got dtype={order_row.dtype} "
                     f"shape={tuple(order_row.shape)}")
+            # emission-assisted tiers: mode from which ext tensors the
+            # caller handed in (see gvr_routing.plan_emission)
+            use_ext_counts = seed_thr is not None and seed_thr.shape[1] >= 6
+            ext_rungs = seed_thr is not None and seed_thr.shape[1] == 3
+            use_ext_cand = cand_vals is not None
+            enable_block_skip = block_max is not None
+            emit_xstate = xstate is not None
+            if use_ext_cand:
+                assert use_ext_counts, (
+                    "candidate list requires the packed seed row "
+                    "([rows, 8]: lines + counts)")
+                assert (cand_idx is not None and cand_ctl is not None
+                        and cand_vals.shape == cand_idx.shape
+                        and cand_ctl.shape == (num_rows, 4)), (
+                            "list tier needs cand_vals/cand_idx same shape "
+                            "+ cand_ctl [rows, 4]")
+            if (use_ext_counts or ext_rungs or use_ext_cand
+                    or enable_block_skip):
+                assert not lb_mode and order_row is None, (
+                    "ext tiers are single-CTA/sort-path only")
+            if num_threads is not None:
+                tuning = dict(tuning, num_threads_per_block=num_threads)
+            elif use_ext_cand and top_k <= 512:
+                # small-K list rule: hit rows do O(list) work
+                tuning = dict(tuning, num_threads_per_block=512)
             key = cls._compile(
                 cute_dtype,
                 top_k,
@@ -7021,10 +7110,20 @@ def forward(
                 return_output_values=return_output_values,
                 cluster_size=cluster_size,
                 seqlen_sorted=seqlen_sorted,
+                enable_block_skip=enable_block_skip,
+                use_ext_counts=use_ext_counts,
+                emit_xstate=emit_xstate,
+                use_ext_cand=use_ext_cand,
+                ext_rungs=ext_rungs,
+                cand_cap=(cand_vals.shape[1] if use_ext_cand else 5120),
+                accept_cap=accept_cap,
+                kc_override=kc_override,
                 **tuning,
             )
             cls.kernel_cache[key](logits, pre_idx, seq_lens, None,
-                                  output_indices, order_row)
+                                  output_indices, order_row, block_max,
+                                  seed_thr, None, xstate, cand_vals, cand_idx,
+                                  cand_ctl)
 
     # TODO(dsa.py): wire ``order_row = argsort(seq_lens, descending=True)``
     # (device-side, graph-safe) into the LJF row-reorder branch when
@@ -7033,7 +7132,7 @@ def forward(
     # swap. Below that threshold the win is noise / can regress a few
     # percent (B200 N∈{8K,16K,32K} sweep 2026-06-23).
     @torch.library.custom_op("trtllm::cute_dsl_gvr_topk_decode",
-                             mutates_args=("output_indices", ),
+                             mutates_args=("output_indices", "xstate"),
                              device_types="cuda")
     def cute_dsl_gvr_topk_decode(
         logits: torch.Tensor,
@@ -7048,6 +7147,15 @@ def cute_dsl_gvr_topk_decode(
         order_row: Optional[torch.Tensor] = None,
         counters: Optional[torch.Tensor] = None,
         max_batch_size: Optional[int] = None,
+        seed_thr: Optional[torch.Tensor] = None,
+        xstate: Optional[torch.Tensor] = None,
+        cand_vals: Optional[torch.Tensor] = None,
+        cand_idx: Optional[torch.Tensor] = None,
+        cand_ctl: Optional[torch.Tensor] = None,
+        block_max: Optional[torch.Tensor] = None,
+        num_threads: Optional[int] = None,
+        accept_cap: Optional[int] = None,
+        kc_override: Optional[int] = None,
     ) -> None:
         """CuTe DSL GVR (Guess-Verify-Refine) Top-K decode for Blackwell.
 
@@ -7118,6 +7226,15 @@ def cute_dsl_gvr_topk_decode(
             order_row=order_row,
             counters=counters,
             max_batch_size=max_batch_size,
+            seed_thr=seed_thr,
+            xstate=xstate,
+            cand_vals=cand_vals,
+            cand_idx=cand_idx,
+            cand_ctl=cand_ctl,
+            block_max=block_max,
+            num_threads=num_threads,
+            accept_cap=accept_cap,
+            kc_override=kc_override,
         )
 
     @torch.library.register_fake("trtllm::cute_dsl_gvr_topk_decode")
@@ -7134,6 +7251,15 @@ def _(
         order_row: Optional[torch.Tensor] = None,
         counters: Optional[torch.Tensor] = None,
         max_batch_size: Optional[int] = None,
+        seed_thr: Optional[torch.Tensor] = None,
+        xstate: Optional[torch.Tensor] = None,
+        cand_vals: Optional[torch.Tensor] = None,
+        cand_idx: Optional[torch.Tensor] = None,
+        cand_ctl: Optional[torch.Tensor] = None,
+        block_max: Optional[torch.Tensor] = None,
+        num_threads: Optional[int] = None,
+        accept_cap: Optional[int] = None,
+        kc_override: Optional[int] = None,
     ) -> None:
         return None
 

From 6d652af49e7ffd8b90c847666a22740f310dd336 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 28 Jul 2026 00:57:51 -0700
Subject: [PATCH 052/117] [None][feat] DSA: emission-assisted GVR decode wiring
 (env opt-in)

Wire the validated emission suite into the DSA decode path, gated on
TRTLLM_GVR_EXT=1 (plus the cute-dsl topk + FP4 DSL indexer flags) so
the default path stays byte-identical.

- gvr_ext.GvrExtState: persistent stable-address buffers (packed seed
  row, bucketed SoA list, block_max, xstate, prev-topk feedback - the
  same feedback-loop shape as heuristic_prev_topk), the device-side
  closed-loop seed-row update (pure tensor ops, graph-capturable,
  11/11-step replay-exact) and the per-step tier plan/route via
  gvr_routing. Cold-start rows publish non-finite lines so the kernel
  validity guard keeps exactness independent of host state.
- fp4 paged-MQA op grows the emission face (block_max / packed seed
  counts / bucketed candidate list as caller-owned mutates).
- dsa.py: the indexer call merges the planned emit kwargs; the decode
  top-k branch consumes the previous step's emission through
  trtllm::cute_dsl_gvr_topk_decode with warm-start from the ext
  feedback buffer. Scoped to next_n == 1 without atom-split; every
  other shape takes the unchanged paths.

Verified: glue-logic CPU unit (tier transitions, kwargs shapes, cold
start), fp4 DSL decode pytest regression with the flag off (4 passed),
op-level tiers + emission E2E + graph replay all bit-exact from the
prior commits. Real-model TP4 A/B on DeepSeek V4-Flash is the
remaining acceptance run.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  70 +++++++-
 .../attention_backend/sparse/gvr_ext.py       | 166 ++++++++++++++++++
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  37 +++-
 3 files changed, 268 insertions(+), 5 deletions(-)
 create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index fa71d2790e40..15cad0268b2f 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1811,6 +1811,16 @@ def __init__(self,
         self.use_cute_dsl_paged_mqa_logits = (
             sparse_params.use_cute_dsl_paged_mqa_logits
             and IS_CUTLASS_DSL_AVAILABLE)
+        # GVR emission-assisted decode (opt-in, experimental): the FP4
+        # indexer epilogue emits seed counts / the bucketed candidate
+        # list and the top-k consumes them (see gvr_ext / gvr_routing).
+        # Env-gated so the default path stays byte-identical.
+        self.use_gvr_ext = (os.environ.get("TRTLLM_GVR_EXT", "0") == "1"
+                            and self.use_cute_dsl_topk
+                            and self.use_cute_dsl_paged_mqa_logits
+                            and sparse_params.indexer_k_dtype == "fp4")
+        self._gvr_ext = None  # lazy GvrExtState (first decode step)
+        self._gvr_route = None
         self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5
 
         self._enable_heuristic_topk = (sparse_params.enable_heuristic_topk
@@ -2766,10 +2776,35 @@ def sparse_attn_indexer(
                         dsl_schedule_meta = (
                             metadata.scheduler_metadata_buffer_expanded)
 
+                    gvr_emit_kwargs = {}
+                    if (self.use_gvr_ext and next_n == 1
+                            and not dsl_atom_split):
+                        from .gvr_ext import GvrExtState
+                        if self._gvr_ext is None:
+                            self._gvr_ext = GvrExtState(
+                                max_rows=metadata.max_num_sequences,
+                                top_k=self.index_topk,
+                                device=q_fp8.device)
+                        st = self._gvr_ext
+                        n_comp = indexer_max_seq_len // max(
+                            self.compress_ratio, 1)
+                        emit_tier, self._gvr_route = st.plan(
+                            batch_size, n_comp,
+                            torch.cuda.get_device_properties(
+                                q_fp8.device).multi_processor_count)
+                        st.update_seed_rows(batch_size)
+                        gvr_emit_kwargs = st.indexer_emit_kwargs(
+                            emit_tier, batch_size)
+                        if self._gvr_route.attach_block_max or emit_tier in (
+                                "counts", "list"):
+                            gvr_emit_kwargs["block_max_out"] = (
+                                st.ensure_block_max(
+                                    indexer_max_seq_len //
+                                    max(self.compress_ratio, 1))[:batch_size])
                     logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits(
                         dsl_q, decode_q_scale, k_cache, weights_decode,
                         dsl_context_lens, dsl_block_table, dsl_schedule_meta,
-                        indexer_max_seq_len)
+                        indexer_max_seq_len, **gvr_emit_kwargs)
                 else:
                     # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}.
                     # Atom-split benefits small-batch / low-ntask configs by
@@ -2832,7 +2867,38 @@ def sparse_attn_indexer(
                             metadata.heuristic_scratch_values[
                                 :num_gen_tokens]
 
-                if self.use_cute_dsl_topk and self._enable_heuristic_topk:
+                # CuTE DSL top-k allocates O(num_gen_tokens * kv_len) global
+                # memory. Beyond 256 tokens the extra memory becomes significant,
+                # so we cap it at 256 for now and fall back to the CUDA C++
+                # indexer_topk_decode. This limit can be removed if GPU memory
+                # is not a bottleneck.
+                if (self.use_gvr_ext and self._gvr_ext is not None
+                        and self._gvr_route is not None and next_n == 1
+                        and num_gen_tokens <= 256):
+                    # emission-assisted GVR: consume what the indexer
+                    # epilogue emitted this step (packed seed row /
+                    # bucketed list / block_max per the picked route)
+                    st = self._gvr_ext
+                    out_slice = topk_indices_buffer[
+                        num_ctx_tokens:num_ctx_tokens + num_gen_tokens, :]
+                    seq_1d = (context_lens if self.compress_ratio > 1 else
+                              gen_kv_lens_cuda).reshape(-1)[:num_gen_tokens]
+                    ext_kw = st.topk_ext_kwargs(
+                        self._gvr_route, num_gen_tokens,
+                        st.block_max[:num_gen_tokens]
+                        if st.block_max is not None else None)
+                    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+                        logits_decode,
+                        st.prev_topk[:num_gen_tokens],
+                        seq_1d.contiguous(),
+                        out_slice,
+                        self.index_topk,
+                        next_n,
+                        self.compress_ratio,
+                        max_seq_len=indexer_max_seq_len,
+                        **ext_kw)
+                    st.prev_topk[:num_gen_tokens].copy_(out_slice)
+                elif self.use_cute_dsl_topk and self._enable_heuristic_topk:
                     # GVR DSL: supports all compress_ratio and next_n values.
                     torch.ops.trtllm.cute_dsl_gvr_topk_decode(
                         logits_decode,
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
new file mode 100644
index 000000000000..e1ca588cc3b5
--- /dev/null
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -0,0 +1,166 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION &
+# AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Emission-assisted GVR top-k state for the DSA decode path.
+
+Owns the persistent (graph-address-stable) buffers the emission tiers
+ride on, the device-side closed-loop seed-row update (pure tensor ops:
+CUDA-graph capturable, validated 11/11-step replay-exact against eager)
+and the per-step routing decision. All of it is opt-in: without the
+flag the DSA decode path is byte-identical to before.
+
+Tier semantics (see gvr_routing):
+  * this step's TOP-K consumes what the PREVIOUS step's indexer
+    epilogue emitted;
+  * this step's INDEXER emits what the routing planned for the NEXT
+    step. N changes by at most one slot per step, so tier flapping is
+    a non-issue.
+"""
+
+from typing import Optional
+
+import torch
+
+from ...cute_dsl_kernels.blackwell.top_k.gvr_routing import TopkRoute, pick_config, plan_emission
+
+# Bucketed list geometry (validated defaults: B* = 8192 segment cap,
+# 24576-entry C segment; see the f15/f17 sweeps).
+LIST_SEG_A = 8192
+LIST_CAP_C = 24576
+LIST_WIDTH = 2 * LIST_SEG_A + LIST_CAP_C
+
+# Closed-loop line derivation around the published k-th anchor: t1
+# hugs the k-th value from below, t0/t2 guard by the (anchor - kth)
+# span. Matches the graph_test.py-validated update.
+GUARD_LO = 2.0
+GUARD_HI = 0.5
+
+
+class GvrExtState:
+    """Per-attention-backend emission state (persistent buffers)."""
+
+    def __init__(
+        self, max_rows: int, top_k: int, device: torch.device, enable_list_tier: bool = True
+    ):
+        self.max_rows = max_rows
+        self.top_k = top_k
+        # packed seed row: lines at cols 0..2, counts (emission-filled)
+        # at 3..5, adaptive-skip pass count at 6
+        self.seed_row = torch.zeros((max_rows, 8), dtype=torch.float32, device=device)
+        self.xstate = torch.zeros((max_rows, 8), dtype=torch.float32, device=device)
+        self.cand_vals: Optional[torch.Tensor] = None
+        self.cand_idx: Optional[torch.Tensor] = None
+        self.cand_ctl: Optional[torch.Tensor] = None
+        self.cand_cur: Optional[torch.Tensor] = None
+        if enable_list_tier:
+            self.cand_vals = torch.zeros((max_rows, LIST_WIDTH), dtype=torch.float32, device=device)
+            self.cand_idx = torch.zeros((max_rows, LIST_WIDTH), dtype=torch.int32, device=device)
+            self.cand_ctl = torch.zeros((max_rows, 4), dtype=torch.int32, device=device)
+            self.cand_cur = torch.zeros((max_rows, 4), dtype=torch.int32, device=device)
+        # GVR warm-start feedback: this layer's previous-step top-k
+        # (same stable-address feedback-loop shape as
+        # heuristic_prev_topk; zero-init -> first step's pre_idx points
+        # at index 0, a valid benign candidate)
+        self.prev_topk = torch.zeros((max_rows, top_k), dtype=torch.int32, device=device)
+        # block_max prefix ([rows, nb_pad*4] fp32 warp-partials),
+        # allocated lazily once max_seq_len is known
+        self.block_max: Optional[torch.Tensor] = None
+        # tier the PREVIOUS indexer call emitted (what this step's
+        # top-k may consume); "rungs" until the first emission lands
+        self.emitted_tier = "rungs"
+
+    def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
+        nb4 = ((max_seq_len + 255) // 256 * 256) // 128 * 4
+        if self.block_max is None or self.block_max.shape[1] < nb4:
+            self.block_max = torch.zeros(
+                (self.max_rows, nb4), dtype=torch.float32, device=self.seed_row.device
+            )
+        return self.block_max
+
+    def plan(self, batch: int, n_comp: int, num_sms: int) -> tuple[str, TopkRoute]:
+        """Route this step: (tier to EMIT next, launch knobs to CONSUME
+        what was emitted last step)."""
+        emit_tier = plan_emission(batch, n_comp, self.top_k, have_epilogue=True)
+        route = pick_config(self.emitted_tier, batch, n_comp, self.top_k, num_sms)
+        return emit_tier, route
+
+    def update_seed_rows(self, num_rows: int) -> None:
+        """Device-side closed-loop line update from the last publish.
+
+        Pure tensor ops (graph-capturable). Rows whose xstate is not
+        valid (col 0 == 0, e.g. cold start) get non-finite lines, which
+        the kernel's validity guard routes to the stock path - the
+        closed loop never rides on host data quality.
+        """
+        s = self.seed_row[:num_rows]
+        x = self.xstate[:num_rows]
+        kth = x[:, 1]
+        anch = torch.maximum(x[:, 2], kth + 1e-5)
+        span = (anch - kth).clamp_min(1e-4)
+        valid = x[:, 0] > 0
+        inf = torch.full_like(kth, float("inf"))
+        s[:, 0] = torch.where(valid, kth - GUARD_LO * span, inf)
+        s[:, 1] = torch.where(valid, kth - 1e-6, inf)
+        s[:, 2] = torch.where(valid, kth + GUARD_HI * span, inf)
+        s[:, 3:8] = 0.0
+        if self.cand_ctl is not None:
+            self.cand_ctl[:num_rows].zero_()
+            self.cand_cur[:num_rows].zero_()
+
+    def indexer_emit_kwargs(self, emit_tier: str, num_rows: int) -> dict:
+        """kwargs for CuteDSLFP4PagedMQALogitsRunner.forward covering the
+        planned emission tier (caller merges into its call)."""
+        kw: dict = {}
+        if emit_tier in ("counts", "list"):
+            kw.update(emit_seed_counts=True, seed_thr=self.seed_row[:num_rows])
+        if emit_tier == "list":
+            kw.update(
+                emit_cand_bucketed=True,
+                accept_cap=LIST_SEG_A,
+                cand_out=self.cand_vals[:num_rows],
+                cand_idx_out=self.cand_idx[:num_rows],
+                cand_ctl_out=self.cand_ctl[:num_rows],
+                cand_cur_out=self.cand_cur[:num_rows],
+            )
+        self.emitted_tier = emit_tier
+        return kw
+
+    def topk_ext_kwargs(
+        self, route: TopkRoute, num_rows: int, block_max: Optional[torch.Tensor]
+    ) -> dict:
+        """kwargs for trtllm::cute_dsl_gvr_topk_decode consuming the
+        PREVIOUS step's emission per the picked route."""
+        kw: dict = {
+            "xstate": self.xstate[:num_rows],
+            "cluster_size": route.cluster_size,
+        }
+        if route.num_threads is not None:
+            kw["num_threads"] = route.num_threads
+        if route.tier in ("counts", "list"):
+            kw["seed_thr"] = self.seed_row[:num_rows]
+        # rungs tier (first step / no emission yet): pass no seed at
+        # all - cold-start xstate is invalid so the lines would be
+        # non-finite anyway; the plain stock path is the right fallback
+        # (a [rows, 3] column view of the packed row is non-contiguous
+        # and would trip the runner's contract assert)
+        if route.tier == "list":
+            kw.update(
+                cand_vals=self.cand_vals[:num_rows],
+                cand_idx=self.cand_idx[:num_rows],
+                cand_ctl=self.cand_ctl[:num_rows],
+            )
+        if route.attach_block_max and block_max is not None:
+            kw["block_max"] = block_max
+        return kw
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index bf2fa3df843b..0c2ee1664a67 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -8870,7 +8870,9 @@ def forward(
             return logits
 
     @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits",
-                             mutates_args=(),
+                             mutates_args=("block_max_out", "seed_thr",
+                                           "cand_out", "cand_idx_out",
+                                           "cand_ctl_out", "cand_cur_out"),
                              device_types="cuda")
     def cute_dsl_fp4_paged_mqa_logits(
         q: torch.Tensor,
@@ -8885,6 +8887,13 @@ def cute_dsl_fp4_paged_mqa_logits(
         epi_dtype: torch.dtype = torch.float32,
         output_dtype: torch.dtype = torch.float32,
         remove_online_sf_transpose: bool = False,
+        block_max_out: Optional[torch.Tensor] = None,
+        seed_thr: Optional[torch.Tensor] = None,
+        cand_out: Optional[torch.Tensor] = None,
+        cand_idx_out: Optional[torch.Tensor] = None,
+        cand_ctl_out: Optional[torch.Tensor] = None,
+        cand_cur_out: Optional[torch.Tensor] = None,
+        accept_cap: int = 8192,
     ) -> torch.Tensor:
         if not is_sm_100f():
             raise ValueError(
@@ -8910,7 +8919,7 @@ def cute_dsl_fp4_paged_mqa_logits(
             f"epi_dtype={epi_dtype} output_dtype={output_dtype}",
             key="cute_dsl_fp4_paged_mqa_logits_inputs",
         )
-        return CuteDSLFP4PagedMQALogitsRunner.forward(
+        ret = CuteDSLFP4PagedMQALogitsRunner.forward(
             q,
             sf_q,
             kv_fused,
@@ -8922,7 +8931,22 @@ def cute_dsl_fp4_paged_mqa_logits(
             num_epi_subtiles=num_epi_subtiles,
             epi_dtype=epi_dtype,
             output_dtype=output_dtype,
-            remove_online_sf_transpose=remove_online_sf_transpose)
+            remove_online_sf_transpose=remove_online_sf_transpose,
+            emit_block_meta=block_max_out is not None,
+            emit_hit_stats=False,
+            block_max_out=block_max_out,
+            emit_seed_counts=seed_thr is not None,
+            seed_thr=seed_thr,
+            emit_cand_bucketed=cand_out is not None,
+            accept_cap=accept_cap,
+            cand_out=cand_out,
+            cand_idx_out=cand_idx_out,
+            cand_ctl_out=cand_ctl_out,
+            cand_cur_out=cand_cur_out)
+        # with emission on, the runner returns (logits, block_max,
+        # hit_stats) - the emission buffers are caller-owned mutates,
+        # the op face stays logits-only
+        return ret[0] if isinstance(ret, tuple) else ret
 
     @torch.library.register_fake("trtllm::cute_dsl_fp4_paged_mqa_logits")
     def _(
@@ -8938,6 +8962,13 @@ def _(
         epi_dtype: torch.dtype = torch.float32,
         output_dtype: torch.dtype = torch.float32,
         remove_online_sf_transpose: bool = False,
+        block_max_out: Optional[torch.Tensor] = None,
+        seed_thr: Optional[torch.Tensor] = None,
+        cand_out: Optional[torch.Tensor] = None,
+        cand_idx_out: Optional[torch.Tensor] = None,
+        cand_ctl_out: Optional[torch.Tensor] = None,
+        cand_cur_out: Optional[torch.Tensor] = None,
+        accept_cap: int = 8192,
     ) -> torch.Tensor:
         B = q.shape[0]
         next_n = q.shape[1]

From 3a4f15d49bdfcc6a3dd6cb38b2d9c1efe5d86acc Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 28 Jul 2026 01:04:00 -0700
Subject: [PATCH 053/117] [None][fix] cute-dsl ops: keep optional emission
 tensors out of mutates_args

torch.library's in-place bookkeeping IndexErrors when a mutates_args
entry is an optional tensor left at its None default - the bare
default-path call of the fp4 paged-MQA op crashed outright. Follow the
existing precedent (heuristic_scratch on trtllm::indexer_topk_decode
is written but unlisted): drop the optional emission tensors and
xstate from mutates_args on both extended ops. Bare-op default path
verified bit-identical to the runner-direct call on B200.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py        | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 0c2ee1664a67..902932cf2393 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -7131,8 +7131,10 @@ def forward(
     # full SM-row's worth of CTAs so the sort has long-vs-short rows to
     # swap. Below that threshold the win is noise / can regress a few
     # percent (B200 N∈{8K,16K,32K} sweep 2026-06-23).
+    # xstate is written by the kernel but stays out of mutates_args
+    # (optional-mutate None-default IndexError; see the fp4 op note)
     @torch.library.custom_op("trtllm::cute_dsl_gvr_topk_decode",
-                             mutates_args=("output_indices", "xstate"),
+                             mutates_args=("output_indices", ),
                              device_types="cuda")
     def cute_dsl_gvr_topk_decode(
         logits: torch.Tensor,
@@ -8869,10 +8871,13 @@ def forward(
                      None, None, None, None, None, None, None, None)
             return logits
 
+    # NOTE: the optional emission tensors ARE written by the kernel but
+    # deliberately not in mutates_args - torch.library's in-place
+    # bookkeeping IndexErrors on optional mutates left at their None
+    # default (same precedent as heuristic_scratch on
+    # trtllm::indexer_topk_decode).
     @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits",
-                             mutates_args=("block_max_out", "seed_thr",
-                                           "cand_out", "cand_idx_out",
-                                           "cand_ctl_out", "cand_cur_out"),
+                             mutates_args=(),
                              device_types="cuda")
     def cute_dsl_fp4_paged_mqa_logits(
         q: torch.Tensor,

From 143453c2940f859fa53bf250650759aa08737f64 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 28 Jul 2026 03:20:01 -0700
Subject: [PATCH 054/117] [None][fix] GVR top-k: rescue degenerate preIdx seeds
 instead of identity output

The stock seed path derives the P2 refine bracket from the preIdx
gather; its exactness invariant (count(>= gather min) >= K) only holds
when the preIdx row carries K distinct in-range positions. A degenerate
gather (duplicate or invalid indices) produced an unusable bracket and
the old shortcut emitted identity indices [0, K) - NOT the top-K on
real data.

Production hit: the first decode step of every request feeds the
zero-init prev_topk feedback buffer (512 copies of index 0), so every
indexer layer returned identity on step 1. Real-model V4-Flash TP2
acceptance showed 42/231 dumped rows (= 21 layers x 2 sequences, first
step each) with a wrong score multiset and no boundary ties. Also
reachable via stale batch-slot reuse (all indices past the new row's
N_eff) and an all-tied gather (FP4 quantization).

Fix: P1r data reseed - when the bracket is degenerate and N > K,
rebuild it from the row itself (full-row min/max, cnt_lo = N), which
restores the P2 invariant unconditionally, then run the normal
pipeline. A still-degenerate bracket after the rescue means every
in-range value is identical (or N <= K), where identity output is
provably exact - the shortcut is kept for exactly those rows.
Non-degenerate rows pay nothing.

Validation: real-model acceptance re-run 231/231 rows score-multiset
exact (in-run selection, all layers/steps); 37-cell degenerate unit
battery green; ext tiers (counts/list/rungs/counts+bm) EXACT, CUDA
graph replay 11/11, bare-op default path identical. A new xfail test
documents the pre-existing beyond-kC tie-flood limitation (index dups
when the kth tie class alone exceeds kC; value multiset still exact;
reproduces on the unmodified upstream kernel).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        |  90 +++++++++++-
 .../sparse/test_cute_dsl_gvr_topk_decode.py   | 138 ++++++++++++++++++
 2 files changed, 226 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index c4c021dd314a..8c975806e2c3 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -1088,6 +1088,68 @@ def phase1_preidx_stats(
                 s_iscalars[4] = cutlass.Int32(0)
         cute.arch.barrier()
 
+    # ------------------------------------------------------------------
+    # P1r — degenerate-seed rescue: rebuild the refine bracket from the
+    # data itself. Runs only when the preIdx gather produced an unusable
+    # bracket (duplicate or invalid preIdx: cold-start zero-init slots,
+    # stale slots pointing past N, or an all-tied gather). A full-row
+    # min/max restores the P2 invariant count(>= v_lo) >= K, so the
+    # normal pipeline stays exact; the extra row scan is paid only by
+    # the (rare) degenerate rows. Bounds are clamped to +-FLT_MAX/2 so
+    # the secant range arithmetic stays finite against inf-laden rows.
+    # ------------------------------------------------------------------
+    @cute.jit
+    def phase1r_data_reseed(
+        self,
+        input_row,  # cute.Tensor [N] (row-major slice of the row)
+        N,  # runtime row length
+        smem_wmin_f32,  # cute.Tensor [NUM_WARPS] float32 (reused P1 buffer)
+        smem_wmax_f32,  # cute.Tensor [NUM_WARPS] float32 (reused P1 buffer)
+        s_thr,  # cute.Tensor [3] float32: [threshold, val_lo, val_hi]
+        s_iscalars,  # [cand_count, done, cnt_lo, cnt_hi, out_count, ...]
+        s_mt_thr,  # rung columns (r0_vseed parks the seed line here)
+        tidx,
+        warp_id,
+        lane,
+    ):
+        local_min = cutlass.Float32(self.FLT_MAX)
+        local_max = cutlass.Float32(self.NEG_FLT_MAX)
+        i = cutlass.Int32(tidx)
+        while i < N:
+            v = self._load_fp32(input_row, i)
+            local_max = cute.arch.fmax(local_max, v)
+            local_min = _fmin_f32_inline(local_min, v)
+            i = i + cutlass.Int32(self.num_threads)
+        wmin = self.warp_reduce_min_f32(local_min)
+        wmax = self.warp_reduce_max_f32(local_max)
+        if lane == 0:
+            smem_wmin_f32[warp_id] = wmin
+            smem_wmax_f32[warp_id] = wmax
+        cute.arch.barrier()
+        if tidx == 0:
+            rmin = cutlass.Float32(self.FLT_MAX)
+            rmax = cutlass.Float32(self.NEG_FLT_MAX)
+            for w in cutlass.range_constexpr(self.num_warps):
+                rmax = cute.arch.fmax(rmax, smem_wmax_f32[w])
+                rmin = _fmin_f32_inline(rmin, smem_wmin_f32[w])
+            # finite clamp keeps rng = val_hi - val_lo representable; rows
+            # with mass beyond +-FLT_MAX/2 are adversarial-only (production
+            # indexer scores are small finite values).
+            rmin = cute.arch.fmax(rmin, cutlass.Float32(self.NEG_FLT_MAX * 0.5))
+            rmax = _fmin_f32_inline(rmax, cutlass.Float32(self.FLT_MAX * 0.5))
+            mid = (rmin + rmax) * cutlass.Float32(0.5)
+            s_thr[0] = mid
+            s_thr[1] = rmin
+            s_thr[2] = rmax
+            if cutlass.const_expr(self.r0_vseed):
+                s_mt_thr[self.M_thr - 1] = mid
+            s_iscalars[0] = cutlass.Int32(0)  # cand_count
+            s_iscalars[1] = cutlass.Int32(0)  # done
+            s_iscalars[2] = N  # cnt_lo: count(>= row min) = N, truthful
+            s_iscalars[3] = cutlass.Int32(1)  # cnt_hi seed (same as P1)
+            s_iscalars[4] = cutlass.Int32(0)  # out_count
+        cute.arch.barrier()
+
     # ------------------------------------------------------------------
     # P1b — 256-bin SMEM histogram over the prev-topK gathered values
     # (band [v_lo, v_hi] = P1's pmin/pmax = s_thr[1]/s_thr[2]), then M
@@ -5632,8 +5694,32 @@ def _run_phases(
             )
 
         # Degenerate threshold init: val_hi <= -self.FLT_MAX or val_lo >= val_hi.
-        # When preIdx values produce an unusable bracket (e.g. all -inf or
-        # identical), skip Phase 2-4 and emit identity output instead.
+        # A duplicate/invalid preIdx gather (cold-start zero-init slots, stale
+        # slots pointing past N, an all-tied gather) produces an unusable
+        # bracket. When N > K real selection work remains, so rebuild the
+        # bracket from the data itself (P1r) and run the normal pipeline —
+        # the old identity shortcut here returned indices [0, K), which is
+        # NOT the top-K on real data (production hit: the first decode step
+        # of every request feeds the zero-init prev_topk feedback buffer).
+        # If the bracket is STILL degenerate after the rescue, every
+        # in-range value is identical (or N <= K), and identity output is
+        # then exact — keep the shortcut for exactly those rows.
+        v_lo = s_thr[1]
+        v_hi = s_thr[2]
+        if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi:
+            if N > cutlass.Int32(self.top_k):
+                self.phase1r_data_reseed(
+                    input_row,
+                    N,
+                    smem_wmin,
+                    smem_wmax,
+                    s_thr,
+                    s_iscalars,
+                    s_mt_thr,
+                    tidx,
+                    warp_id,
+                    lane,
+                )
         v_lo = s_thr[1]
         v_hi = s_thr[2]
         if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi:
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
index f8f7215eb6ee..7a5bae65a719 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
@@ -1135,3 +1135,141 @@ def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size)
     _tie_aware_check(out_sec, logits, seq_lens, top_k, next_n=1, compress_ratio=1)
     if dtype == torch.float32:
         assert torch.equal(out.sort(dim=-1).values, out_sec.sort(dim=-1).values)
+
+
+# ===========================================================================
+# Degenerate preIdx states (P1r data reseed).
+#
+# The stock seed path derives the P2 refine bracket from the preIdx gather;
+# its exactness invariant (count(>= gather min) >= K) holds only when the
+# preIdx row carries K DISTINCT in-range positions. Production-reachable
+# violations: the first decode step of a request feeds the zero-init
+# prev_topk feedback buffer (all-duplicate index 0); a reused batch slot can
+# carry stale indices past the new row's N_eff (all-invalid); FP4-quantized
+# logits can tie the whole gather (zero-width bracket). The old degenerate
+# shortcut emitted identity indices [0, K) — NOT the top-K on real data.
+# P1r rebuilds the bracket from the row itself, so all cases must pass the
+# strict multiset check.
+# ===========================================================================
+
+
+@skip_not_sm100
+@pytest.mark.parametrize(
+    "dtype,top_k",
+    [
+        (torch.bfloat16, 512),
+        (torch.float32, 2048),
+    ],
+)
+@pytest.mark.parametrize("compress_ratio", [1, 4])
+@pytest.mark.parametrize("pre_mode", ["zero", "dup", "oob"])
+@pytest.mark.parametrize("data_mode", ["random", "all_tied", "tie_flood"])
+def test_cute_dsl_gvr_topk_decode_degenerate_preidx(
+    dtype, top_k, compress_ratio, pre_mode, data_mode
+):
+    """Degenerate preIdx rows must still produce an exact top-K."""
+    N = 4096
+    num_rows = 4
+    torch.manual_seed(3)
+    torch.cuda.manual_seed(3)
+    if data_mode == "random":
+        logits = (torch.randn(num_rows, N, device="cuda") * 2.0).to(dtype)
+    elif data_mode == "all_tied":
+        # rescue re-degenerates (row min == max) -> identity output, which
+        # is exact here because every value is identical
+        logits = torch.full((num_rows, N), 5.0, device="cuda", dtype=dtype)
+    else:  # tie_flood: kth sits inside a large tie class (within kC —
+        # count(>= tie value) must stay under the candidate capacity; the
+        # beyond-kC flood is a known pre-existing limitation, see the
+        # xfail test below)
+        logits = (torch.rand(num_rows, N, device="cuda") * 0.5).to(dtype)
+        for r in range(num_rows):
+            perm = torch.randperm(N, device="cuda")
+            logits[r, perm[: top_k * 2]] = 1.0
+            logits[r, perm[top_k * 2 : top_k * 2 + top_k // 4]] = 2.0
+
+    if pre_mode == "zero":
+        pre_idx = torch.zeros(num_rows, top_k, dtype=torch.int32, device="cuda")
+    elif pre_mode == "dup":
+        pre_idx = torch.full((num_rows, top_k), 37, dtype=torch.int32, device="cuda")
+    else:  # oob: every slot past N_eff (stale-slot state, pcnt == 0)
+        pre_idx = torch.full((num_rows, top_k), N + 7, dtype=torch.int32, device="cuda")
+
+    seq_lens = torch.full((num_rows,), N * compress_ratio, dtype=torch.int32, device="cuda")
+    out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda")
+    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+        logits,
+        pre_idx,
+        seq_lens,
+        out_indices,
+        top_k=top_k,
+        next_n=1,
+        compress_ratio=compress_ratio,
+    )
+    torch.cuda.synchronize()
+    _tie_aware_check(out_indices, logits, seq_lens, top_k, 1, compress_ratio=compress_ratio)
+
+
+@skip_not_sm100
+def test_cute_dsl_gvr_topk_decode_degenerate_preidx_cs4():
+    """cs>1 rows run the rescue per-CTA (redundant full-row scan) — the
+    cluster path must stay exact for the zero-init cold-start state too."""
+    N = 65536
+    top_k = 512
+    num_rows = 2
+    torch.manual_seed(5)
+    logits = (torch.randn(num_rows, N, device="cuda") * 2.0).to(torch.bfloat16)
+    pre_idx = torch.zeros(num_rows, top_k, dtype=torch.int32, device="cuda")
+    seq_lens = torch.full((num_rows,), N, dtype=torch.int32, device="cuda")
+    out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda")
+    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+        logits,
+        pre_idx,
+        seq_lens,
+        out_indices,
+        top_k=top_k,
+        next_n=1,
+        compress_ratio=1,
+        cluster_size=4,
+    )
+    torch.cuda.synchronize()
+    _tie_aware_check(out_indices, logits, seq_lens, top_k, 1, compress_ratio=1)
+
+
+@skip_not_sm100
+@pytest.mark.xfail(
+    reason="known pre-existing limitation (upstream lineage, reproduces on "
+    "the PR-tip kernel unmodified): when the kth tie class alone exceeds "
+    "the candidate capacity kC, no threshold lands in [K, kC]; the selected "
+    "VALUE multiset is still exact but the index list can contain "
+    "duplicate / unwritten (-1) slots. Requires >kC exactly-equal scores "
+    "at the boundary — unreachable for real FP4 indexer logits observed "
+    "so far. Tracked as a follow-up; independent of the P1r rescue.",
+    strict=False,
+)
+def test_cute_dsl_gvr_topk_decode_tie_flood_beyond_capacity():
+    N = 4096
+    top_k = 512
+    num_rows = 4
+    torch.manual_seed(3)
+    torch.cuda.manual_seed(3)
+    logits = torch.ones(num_rows, N, device="cuda", dtype=torch.bfloat16)
+    for r in range(num_rows):
+        hot = torch.randperm(N, device="cuda")[: top_k // 4]
+        logits[r, hot] = 2.0
+    # healthiest possible pre (true previous top-k) — the flood defect is
+    # independent of preIdx quality
+    pre_idx = torch.topk(logits.float(), top_k, dim=-1).indices.int().contiguous()
+    seq_lens = torch.full((num_rows,), N, dtype=torch.int32, device="cuda")
+    out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda")
+    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+        logits,
+        pre_idx,
+        seq_lens,
+        out_indices,
+        top_k=top_k,
+        next_n=1,
+        compress_ratio=1,
+    )
+    torch.cuda.synchronize()
+    _tie_aware_check(out_indices, logits, seq_lens, top_k, 1, compress_ratio=1)

From 5390d7f083c7f2c269df69319cf147f169a0a09e Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 28 Jul 2026 03:20:46 -0700
Subject: [PATCH 055/117] [None][fix] DSA GVR ext wiring: live seq domains,
 block_max width, dump hook

Fixes found during real-model (DeepSeek V4-Flash TP2) bring-up of the
emission-assisted GVR decode path:

- The GVR op takes RAW-domain seq_lens; on the DSL indexer path the
  live compressed lengths are in gen_indexer_kv_lens_cuda_runtime
  (kv_lens_cuda_2d stays zero there), so multiply those by
  compress_ratio instead of passing context_lens.
- Same trap on the radix baseline branch: pass the live compressed
  lens reshaped to 1-D (the op requires ndim == 1; the old
  context_lens path handed it a [B, 1] view and cr > 1 runs died).
- ensure_block_max allocates the exact validated width (the runner
  asserts shape equality, a wider reused buffer trips it).
- topk_ext_kwargs emits tensors only; the Runner-level mode booleans
  are derived inside the op from tensor presence.
- TRTLLM_GVR_DUMP diagnostic hook (graph-free runs only): saves
  logits / raw seq lens / pre-op prev_topk / the in-run selection per
  layer so an offline pass can verify score-multiset exactness of the
  exact production path; guarded to skip warmup rows.

Acceptance with all of the above plus the P1r kernel rescue: 231/231
dumped rows (21 layers x all captured steps x 2 sequences) score-
multiset exact vs torch.topk on the in-run selection.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    | 53 ++++++++++++++++---
 .../attention_backend/sparse/gvr_ext.py       |  7 +--
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 15cad0268b2f..7d0dd651976f 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2799,8 +2799,7 @@ def sparse_attn_indexer(
                                 "counts", "list"):
                             gvr_emit_kwargs["block_max_out"] = (
                                 st.ensure_block_max(
-                                    indexer_max_seq_len //
-                                    max(self.compress_ratio, 1))[:batch_size])
+                                    indexer_max_seq_len)[:batch_size])
                     logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits(
                         dsl_q, decode_q_scale, k_cache, weights_decode,
                         dsl_context_lens, dsl_block_table, dsl_schedule_meta,
@@ -2881,8 +2880,18 @@ def sparse_attn_indexer(
                     st = self._gvr_ext
                     out_slice = topk_indices_buffer[
                         num_ctx_tokens:num_ctx_tokens + num_gen_tokens, :]
-                    seq_1d = (context_lens if self.compress_ratio > 1 else
-                              gen_kv_lens_cuda).reshape(-1)[:num_gen_tokens]
+                    # the GVR op takes RAW-domain seq_lens (its kernel
+                    # ceil-divides by compress_ratio internally). On the
+                    # DSL indexer path the live compressed lengths are in
+                    # gen_indexer_kv_lens_cuda_runtime (kv_lens_cuda_2d
+                    # stays zero there - same trap as the indexer call).
+                    if self.compress_ratio > 1:
+                        gvr_lens = metadata.gen_indexer_kv_lens_cuda_runtime
+                        assert gvr_lens is not None
+                        seq_1d = (gvr_lens.reshape(-1)[:num_gen_tokens] *
+                                  self.compress_ratio)
+                    else:
+                        seq_1d = gen_kv_lens_cuda.reshape(-1)[:num_gen_tokens]
                     ext_kw = st.topk_ext_kwargs(
                         self._gvr_route, num_gen_tokens,
                         st.block_max[:num_gen_tokens]
@@ -2897,6 +2906,27 @@ def sparse_attn_indexer(
                         self.compress_ratio,
                         max_seq_len=indexer_max_seq_len,
                         **ext_kw)
+                    # diagnostic dump (NO_GRAPH runs only): capture the
+                    # inputs AND the in-run selection so an offline pass
+                    # can check score-multiset equality vs torch.topk on
+                    # the exact production path. prev_topk still holds
+                    # the pre-op value here (copy-back is below). The
+                    # min-seq guard skips warmup/dummy rows (n < K).
+                    if (os.environ.get("TRTLLM_GVR_DUMP")
+                            and getattr(self, "_gvr_dumped", 0) < 6
+                            and int(seq_1d.min())
+                            >= self.index_topk * self.compress_ratio):
+                        self._gvr_dumped = getattr(self, "_gvr_dumped", 0) + 1
+                        torch.save(
+                            {
+                                "logits": logits_decode.clone().cpu(),
+                                "seq_raw": seq_1d.clone().cpu(),
+                                "pre":
+                                st.prev_topk[:num_gen_tokens].clone().cpu(),
+                                "sel": out_slice.clone().cpu(),
+                                "layer": self.layer_idx,
+                            }, f"/tmp/siyid_e2e/dump_l{self.layer_idx}_"
+                            f"{self._gvr_dumped}.pt")
                     st.prev_topk[:num_gen_tokens].copy_(out_slice)
                 elif self.use_cute_dsl_topk and self._enable_heuristic_topk:
                     # GVR DSL: supports all compress_ratio and next_n values.
@@ -2917,9 +2947,20 @@ def sparse_attn_indexer(
                 # significant, so we cap it at 256 and fall back to C++.
                 elif (self.use_cute_dsl_topk and num_gen_tokens <= 256
                       and (self.compress_ratio == 1 or next_n == 1)):
+                    # request-level seq_lens must be 1-D; on the DSL
+                    # indexer path the live compressed lengths are in
+                    # gen_indexer_kv_lens_cuda_runtime (kv_lens_cuda_2d
+                    # stays zero there)
+                    if self.compress_ratio > 1:
+                        radix_lens = (metadata.gen_indexer_kv_lens_cuda_runtime
+                                      if self.use_cute_dsl_paged_mqa_logits and
+                                      metadata.gen_indexer_kv_lens_cuda_runtime
+                                      is not None else context_lens)
+                        radix_lens = radix_lens.reshape(-1)
+                    else:
+                        radix_lens = gen_kv_lens_cuda
                     torch.ops.trtllm.cute_dsl_indexer_topk_decode(
-                        logits_decode, context_lens
-                        if self.compress_ratio > 1 else gen_kv_lens_cuda,
+                        logits_decode, radix_lens,
                         topk_indices_buffer[num_ctx_tokens:num_ctx_tokens +
                                             num_gen_tokens, :], self.index_topk,
                         next_n)
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
index e1ca588cc3b5..da2a5fff3beb 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -83,7 +83,9 @@ def __init__(
 
     def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
         nb4 = ((max_seq_len + 255) // 256 * 256) // 128 * 4
-        if self.block_max is None or self.block_max.shape[1] < nb4:
+        # exact width: the runner asserts shape == (rows, nrec), so a
+        # wider reused buffer would trip it
+        if self.block_max is None or self.block_max.shape[1] != nb4:
             self.block_max = torch.zeros(
                 (self.max_rows, nb4), dtype=torch.float32, device=self.seed_row.device
             )
@@ -124,10 +126,9 @@ def indexer_emit_kwargs(self, emit_tier: str, num_rows: int) -> dict:
         planned emission tier (caller merges into its call)."""
         kw: dict = {}
         if emit_tier in ("counts", "list"):
-            kw.update(emit_seed_counts=True, seed_thr=self.seed_row[:num_rows])
+            kw["seed_thr"] = self.seed_row[:num_rows]
         if emit_tier == "list":
             kw.update(
-                emit_cand_bucketed=True,
                 accept_cap=LIST_SEG_A,
                 cand_out=self.cand_vals[:num_rows],
                 cand_idx_out=self.cand_idx[:num_rows],

From af9df05683be3b07f99162e917f1aa78e0445abe Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 28 Jul 2026 08:18:45 -0700
Subject: [PATCH 056/117] [None][fix] GVR: xstate anchor visibility barrier +
 defensive contract checks

Review follow-ups (coderabbit round 1):

- List rows publish the closed-loop anchor by reading
  output_indices_row[K-1], written by peer threads in rank-scatter /
  tail repair; the eager-position path dropped the swap loop's trailing
  block barrier, so the read could race the last writes. Publish
  visibility explicitly with a barrier gated on list_used (uniform:
  admission is decided from shared control words). Counts/stock rows
  pay nothing.
- GvrExtState.plan() demotes the list tier to counts when the state was
  constructed without candidate buffers (enable_list_tier=False
  previously routed into a None subscript).
- ensure_block_max() raises instead of silently allocating inside CUDA
  graph capture (a captured dangling address would corrupt replays).
- The TRTLLM_GVR_DUMP diagnostic hook takes the dump directory from the
  env value instead of a hard-coded path.
- The op face rejects unexpected seed_thr widths (4/5 silently disabled
  both ext modes) and checks cand_out presence before reading its
  shape.
- Debug env knobs parse malformed values without raising at import.
- tie_flood unit case sizes its tie class as 1.5*K so it stays a
  genuine flood at top_k=2048 (2*K covered the whole N=4096 row and
  collapsed into all_tied).

Validation: unified smoke 30/30 exact (list path exercises the new
barrier), degenerate/flood battery 37 passed + 1 xfail, ext tiers
EXACT, CUDA graph replay 11/11, bare-op path identical, GvrExtState
CPU unit incl. the new demotion case.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    | 12 ++++++----
 .../attention_backend/sparse/gvr_ext.py       | 14 +++++++++++
 .../_torch/custom_ops/cute_dsl_custom_ops.py  | 13 ++++++++---
 .../blackwell/top_k/gvr_topk_decode.py        | 23 ++++++++++++++++---
 .../sparse/test_cute_dsl_gvr_topk_decode.py   |  9 ++++++--
 5 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 7d0dd651976f..d3ca287eba94 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2912,10 +2912,12 @@ def sparse_attn_indexer(
                     # the exact production path. prev_topk still holds
                     # the pre-op value here (copy-back is below). The
                     # min-seq guard skips warmup/dummy rows (n < K).
-                    if (os.environ.get("TRTLLM_GVR_DUMP")
-                            and getattr(self, "_gvr_dumped", 0) < 6
+                    # TRTLLM_GVR_DUMP holds the target directory.
+                    gvr_dump_dir = os.environ.get("TRTLLM_GVR_DUMP")
+                    if (gvr_dump_dir and getattr(self, "_gvr_dumped", 0) < 6
                             and int(seq_1d.min())
                             >= self.index_topk * self.compress_ratio):
+                        os.makedirs(gvr_dump_dir, exist_ok=True)
                         self._gvr_dumped = getattr(self, "_gvr_dumped", 0) + 1
                         torch.save(
                             {
@@ -2925,8 +2927,10 @@ def sparse_attn_indexer(
                                 st.prev_topk[:num_gen_tokens].clone().cpu(),
                                 "sel": out_slice.clone().cpu(),
                                 "layer": self.layer_idx,
-                            }, f"/tmp/siyid_e2e/dump_l{self.layer_idx}_"
-                            f"{self._gvr_dumped}.pt")
+                            },
+                            os.path.join(
+                                gvr_dump_dir, f"dump_l{self.layer_idx}_"
+                                f"{self._gvr_dumped}.pt"))
                     st.prev_topk[:num_gen_tokens].copy_(out_slice)
                 elif self.use_cute_dsl_topk and self._enable_heuristic_topk:
                     # GVR DSL: supports all compress_ratio and next_n values.
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
index da2a5fff3beb..0611bbd10df7 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -86,6 +86,16 @@ def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
         # exact width: the runner asserts shape == (rows, nrec), so a
         # wider reused buffer would trip it
         if self.block_max is None or self.block_max.shape[1] != nb4:
+            # max_seq_len is engine-static, so this allocates once on the
+            # first (eager warmup) step; allocating inside CUDA graph
+            # capture would bake a dangling address into the graph, so
+            # fail loudly instead of corrupting the capture.
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError(
+                    "GvrExtState.ensure_block_max: (re)allocation requested "
+                    "during CUDA graph capture; the block_max buffer must be "
+                    "created by a warmup step before capture"
+                )
             self.block_max = torch.zeros(
                 (self.max_rows, nb4), dtype=torch.float32, device=self.seed_row.device
             )
@@ -95,6 +105,10 @@ def plan(self, batch: int, n_comp: int, num_sms: int) -> tuple[str, TopkRoute]:
         """Route this step: (tier to EMIT next, launch knobs to CONSUME
         what was emitted last step)."""
         emit_tier = plan_emission(batch, n_comp, self.top_k, have_epilogue=True)
+        if emit_tier == "list" and self.cand_vals is None:
+            # constructed with enable_list_tier=False: no candidate
+            # buffers to emit into, demote to the counts tier
+            emit_tier = "counts"
         route = pick_config(self.emitted_tier, batch, n_comp, self.top_k, num_sms)
         return emit_tier, route
 
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 902932cf2393..a2d6a5e09a2b 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -7079,6 +7079,11 @@ def forward(
                     f"shape={tuple(order_row.shape)}")
             # emission-assisted tiers: mode from which ext tensors the
             # caller handed in (see gvr_routing.plan_emission)
+            if seed_thr is not None:
+                assert seed_thr.shape[1] == 3 or seed_thr.shape[1] >= 6, (
+                    "seed_thr width must be 3 (rungs) or >= 6 (packed "
+                    f"counts row); got {seed_thr.shape[1]} - refusing to "
+                    "silently ignore the seed")
             use_ext_counts = seed_thr is not None and seed_thr.shape[1] >= 6
             ext_rungs = seed_thr is not None and seed_thr.shape[1] == 3
             use_ext_cand = cand_vals is not None
@@ -8795,11 +8800,13 @@ def forward(
                 # cand_idx_out = int32 positions (same width), cand_cur_out =
                 # int32 [rows, 4] cursors (caller-zeroed), cand_ctl_out =
                 # int32 [rows, 4] {n0, void, n1, n2} (caller-zeroed)
+                assert cand_out is not None, (
+                    "emit_cand_bucketed requires cand_out")
                 W = cand_out.shape[1]
                 assert (
-                    cand_out is not None and cand_out.dtype == torch.float32
-                    and cand_out.is_cuda and cand_out.is_contiguous()
-                    and cand_out.dim() == 2 and cand_out.shape[0] == B * next_n
+                    cand_out.dtype == torch.float32 and cand_out.is_cuda
+                    and cand_out.is_contiguous() and cand_out.dim() == 2
+                    and cand_out.shape[0] == B * next_n
                     and W > 2 * accept_cap), (
                         "bucketed requires cand_out fp32 [rows, 2*segA+capC]")
                 assert (cand_idx_out is not None
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 8c975806e2c3..10de11addfc8 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -42,14 +42,21 @@
 from ..utils import TRTLLM_ENABLE_PDL, griddepcontrol_launch_dependents, griddepcontrol_wait
 from .block_scan import warp_scan
 
+
+def _env_flag(name: str) -> bool:
+    """Debug-knob parse that never raises at import time: any of
+    ""/0/false/off/no (case-insensitive) is False, everything else True."""
+    return os.environ.get(name, "0").strip().lower() not in ("", "0", "false", "off", "no")
+
+
 # Diagnostic knob: compile per-phase clock64 stamps of the list path
 # into the spare xstate slots (harness-side analysis). Off by default;
 # NEVER set in production.
-_P4_TAIL_DBG = bool(int(os.environ.get("GVR_P4_TAIL_DBG", "0")))
+_P4_TAIL_DBG = _env_flag("GVR_P4_TAIL_DBG")
 # P4 sub-phase clock64 breakdown -> xstate[1,2,4,5,6,7] (debug: clobbers
 # the closed-loop thr/anch publish; single-shot cells only, not chains)
-_P4_SUB_DBG = bool(int(os.environ.get("GVR_P4_SUB_DBG", "0")))
-_SKIP_DBG = bool(int(os.environ.get("GVR_SKIP_DBG", "0")))
+_P4_SUB_DBG = _env_flag("GVR_P4_SUB_DBG")
+_SKIP_DBG = _env_flag("GVR_SKIP_DBG")
 
 
 # ---------------------------------------------------------------------------
@@ -6960,6 +6967,16 @@ def _run_phases(
                     # tight lower bound of the true kth), [2] accepted
                     # threshold, [3] cand_count. The next step derives its
                     # seed rung group from these.
+                    if list_used == cutlass.Int32(1):
+                        # the anchor below reads output_indices_row[K-1],
+                        # written by peer threads in rank-scatter / tail
+                        # repair; not every exit of that phase ends in a
+                        # block barrier (the eager-position path dropped
+                        # the swap loop's trailing one), so publish
+                        # visibility explicitly. list_used is uniform
+                        # (admission is decided from shared control
+                        # words), so the barrier is block-safe.
+                        cute.arch.barrier()
                     if tidx == cutlass.Int32(0):
                         xstate_row[0] = cutlass.Float32(1.0)
                         thr_pub = s_thr[0]
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
index 7a5bae65a719..ea382f8b4145 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
@@ -1183,10 +1183,15 @@ def test_cute_dsl_gvr_topk_decode_degenerate_preidx(
         # beyond-kC flood is a known pre-existing limitation, see the
         # xfail test below)
         logits = (torch.rand(num_rows, N, device="cuda") * 0.5).to(dtype)
+        # tie class sized to 1.5*K so the case stays a genuine flood for
+        # every top_k parametrization (2*K would cover the whole row at
+        # top_k=2048/N=4096 and collapse into all_tied) while
+        # count(>= 1.0) = 1.75*K stays inside the candidate capacity
+        tie_n = top_k + top_k // 2
         for r in range(num_rows):
             perm = torch.randperm(N, device="cuda")
-            logits[r, perm[: top_k * 2]] = 1.0
-            logits[r, perm[top_k * 2 : top_k * 2 + top_k // 4]] = 2.0
+            logits[r, perm[:tie_n]] = 1.0
+            logits[r, perm[tie_n : tie_n + top_k // 4]] = 2.0
 
     if pre_mode == "zero":
         pre_idx = torch.zeros(num_rows, top_k, dtype=torch.int32, device="cuda")

From 58ebe349f14ef750ae5b08592aaf63dc6e01db35 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 01:50:13 -0700
Subject: [PATCH 057/117] [None][fix] test: seed-row col 6 is the adaptive-skip
 pass count, not a stray write

The packed seed-row contract grew a column after this assertion was
written: the emission epilogue publishes the adaptive-skip pass count
into col 6 (lane0-accumulated; the top-k consumer reads it when
block_max rides along), and this test always emits block_max, so col 6
is a legitimate output on every packed case. The stray-write guard now
covers only col 7, and col 6 gets a semantic bound check instead
(0 <= pass count <= block-max record count).

Fixes the 96 seed_counts[True-...] CI failures on B300/DGX_B200
(pipeline #50383). Full-file validation on B200: 1066 passed
(emission), 708 passed + 1 xfailed (top-k).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../sparse/test_cute_dsl_fp4_paged_mqa_logits.py       | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
index 3b7bb55b4bab..7f0142af1fb8 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
@@ -906,7 +906,15 @@ def test_cute_dsl_fp4_paged_mqa_logits_seed_counts(
         torch.testing.assert_close(lf[row, :ctx], lf0[row, :ctx], atol=0.0, rtol=0.0)
     if packed:
         assert torch.equal(seed_row[:, 0:3], seed_thr), "lines clobbered"
-        assert (seed_row[:, 6:8] == 0).all(), "stray write past counts"
+        # col 6 carries the adaptive-skip pass count (lane0-accumulated
+        # diagnostic; the top-k consumer reads it when block_max rides
+        # along), so it is a legitimate output here - bounded by the
+        # block-max record count. Only col 7 must stay untouched.
+        assert (seed_row[:, 7] == 0).all(), "stray write past counts"
+        nrec = block_max.shape[1]
+        assert ((seed_row[:, 6] >= 0) & (seed_row[:, 6] <= nrec)).all(), (
+            "adaptive-skip pass count out of range"
+        )
         seed_counts = seed_row[:, 3:6].to(torch.int32)
     for row in range(num_rows):
         ctx = int(context_lens[row // next_n].item())

From 8af7c59e9e87deb9e5a5c3f71793f788ef4fc3e4 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 04:32:47 -0700
Subject: [PATCH 058/117] [None][perf] Fused indexer->top-k handshake: PDL
 early release + per-row arrival gating

Eliminates the kernel-boundary tax between the FP4 indexer and the GVR
top-k (measured: launch staging + consumer cold prologue + the global
barrier that blocks per-row pipelining; PDL alone recovers only
0.2-1.5us of it because the indexer never triggered early).

Producer (fp4_paged_mqa_logits, emit_arrival mode):
- every CTA calls griddepcontrol.launch_dependents right after reading
  its schedule slice - the dependent top-k grid stages onto SMs as
  producer CTAs retire, and data safety moves entirely to the per-row
  arrival counters;
- each of the 8 math warps counts the splits it processed for a row
  and publishes the count at the existing q-transition / loop-exit
  flush points with sync_warp + fence.acq_rel.gpu + red.relaxed.gpu
  (warp-sync cumulativity orders every lane's logits stores and
  emission atomics before the counter add).

Consumer (gvr_topk_decode, handshake mode):
- after row-view setup, thread 0 spins on ld.global.acquire.gpu until
  the row's counter reaches 8 * ceil(num_kv_tiles / 2), then resets it
  with st.release for the next layer's reuse (the next indexer cannot
  race the reset: its launch is chained behind this kernel's trailing
  launch_dependents);
- MVP scope pins next_n == 1 and cluster_size == 1 (ctor-enforced).

Wiring: opt-in via TRTLLM_GVR_FUSE=1 on top of TRTLLM_GVR_EXT; the
arrival buffer lives in GvrExtState (zero-init once, consumer
self-resets); the DSA call site only engages the handshake when BOTH
sides of the step run an ext tier at cluster_size 1 - the first step
emits counts while its top-k still routes rungs, and emitting arrival
there would leave residue that un-gates the next step early.

Validation: three-cell functional battery PASS x4 iterations (exact vs
torch.topk, fused == unfused as index sets, counters self-reset);
CUDA-graph capture/replay 8/8; real-model DeepSeek V4-Flash TP2 with
TRTLLM_GVR_FUSE=1: 231/231 dumped rows score-multiset exact, graph
mode runs clean. Non-fused regression untouched: degenerate battery
37 passed + 1 xfailed, ext tiers EXACT, bare-op path identical.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    | 26 ++++++-
 .../attention_backend/sparse/gvr_ext.py       | 13 +++-
 .../_torch/custom_ops/cute_dsl_custom_ops.py  | 58 +++++++++++---
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 76 +++++++++++++++++++
 .../blackwell/top_k/gvr_topk_decode.py        | 29 +++++++
 5 files changed, 191 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index d3ca287eba94..6b520a5b19da 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1819,6 +1819,12 @@ def __init__(self,
                             and self.use_cute_dsl_topk
                             and self.use_cute_dsl_paged_mqa_logits
                             and sparse_params.indexer_k_dtype == "fp4")
+        # fused-handshake mode (TRTLLM_GVR_FUSE=1, requires GVR_EXT): the
+        # indexer emits per-row arrival counters and releases the top-k
+        # grid early via PDL; the top-k spin-gates each row on its counter
+        # instead of waiting for the whole indexer grid.
+        self.use_gvr_fuse = (os.environ.get("TRTLLM_GVR_FUSE", "0") == "1"
+                             and self.use_gvr_ext)
         self._gvr_ext = None  # lazy GvrExtState (first decode step)
         self._gvr_route = None
         self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5
@@ -2784,7 +2790,8 @@ def sparse_attn_indexer(
                             self._gvr_ext = GvrExtState(
                                 max_rows=metadata.max_num_sequences,
                                 top_k=self.index_topk,
-                                device=q_fp8.device)
+                                device=q_fp8.device,
+                                enable_fused_handshake=self.use_gvr_fuse)
                         st = self._gvr_ext
                         n_comp = indexer_max_seq_len // max(
                             self.compress_ratio, 1)
@@ -2800,6 +2807,19 @@ def sparse_attn_indexer(
                             gvr_emit_kwargs["block_max_out"] = (
                                 st.ensure_block_max(
                                     indexer_max_seq_len)[:batch_size])
+                        # fused handshake only when BOTH sides of this
+                        # step run an ext tier: the first step emits
+                        # counts but its top-k still routes "rungs" and
+                        # would never consume/reset the counters -
+                        # skipping emission keeps them clean
+                        self._gvr_emitted_arrival = (
+                            st.arrival is not None
+                            and emit_tier in ("counts", "list")
+                            and self._gvr_route.tier in ("counts", "list")
+                            and self._gvr_route.cluster_size == 1)
+                        if self._gvr_emitted_arrival:
+                            gvr_emit_kwargs["arrival_out"] = (
+                                st.arrival[:batch_size])
                     logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits(
                         dsl_q, decode_q_scale, k_cache, weights_decode,
                         dsl_context_lens, dsl_block_table, dsl_schedule_meta,
@@ -2896,6 +2916,10 @@ def sparse_attn_indexer(
                         self._gvr_route, num_gen_tokens,
                         st.block_max[:num_gen_tokens]
                         if st.block_max is not None else None)
+                    if getattr(self, "_gvr_emitted_arrival", False):
+                        # fused handshake: gate each row on this step's
+                        # indexer arrival counter (consumer resets it)
+                        ext_kw["arrival"] = st.arrival[:num_gen_tokens]
                     torch.ops.trtllm.cute_dsl_gvr_topk_decode(
                         logits_decode,
                         st.prev_topk[:num_gen_tokens],
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
index 0611bbd10df7..ed3b30e5a17b 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -52,7 +52,12 @@ class GvrExtState:
     """Per-attention-backend emission state (persistent buffers)."""
 
     def __init__(
-        self, max_rows: int, top_k: int, device: torch.device, enable_list_tier: bool = True
+        self,
+        max_rows: int,
+        top_k: int,
+        device: torch.device,
+        enable_list_tier: bool = True,
+        enable_fused_handshake: bool = False,
     ):
         self.max_rows = max_rows
         self.top_k = top_k
@@ -80,6 +85,12 @@ def __init__(
         # tier the PREVIOUS indexer call emitted (what this step's
         # top-k may consume); "rungs" until the first emission lands
         self.emitted_tier = "rungs"
+        # fused-handshake arrival counters (zero-once: the consumer
+        # resets each row right after gating on it, so steady-state
+        # layer-to-layer reuse needs no host-side zeroing)
+        self.arrival: Optional[torch.Tensor] = None
+        if enable_fused_handshake:
+            self.arrival = torch.zeros((max_rows,), dtype=torch.int32, device=device)
 
     def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
         nb4 = ((max_seq_len + 255) // 256 * 256) // 128 * 4
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index a2d6a5e09a2b..5e30e87a400b 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -6724,6 +6724,7 @@ def _compile(
             cand_cap: int = 5120,
             accept_cap: Optional[int] = None,
             kc_override: Optional[int] = None,
+            handshake: bool = False,
         ) -> tuple:
             key = (dtype, top_k, next_n, enable_unroll_4, enable_phase3_unroll,
                    use_constant_hint, min_blocks_per_mp, use_256bit_load,
@@ -6731,7 +6732,7 @@ def _compile(
                    compress_ratio, return_output_values, cluster_size,
                    seqlen_sorted, enable_block_skip, use_ext_counts,
                    emit_xstate, use_ext_cand, ext_rungs, cand_cap, accept_cap,
-                   kc_override)
+                   kc_override, handshake)
             if key in cls.kernel_cache:
                 return key
             n_rows = cute.sym_int()
@@ -6778,6 +6779,9 @@ def _compile(
                 cutlass.Float32, (n_rows, 8),
                 stride_order=(1, 0),
                 assumed_align=4) if emit_xstate else None)
+            arrival_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Int32, (n_rows, ), stride_order=(0, ), assumed_align=4)
+                            if handshake else None)
             cand_vals_fake = (cute.runtime.make_fake_compact_tensor(
                 cutlass.Float32, (n_rows, cand_cap),
                 stride_order=(1, 0),
@@ -6811,6 +6815,7 @@ def _compile(
                 enable_block_skip=enable_block_skip,
                 use_ext_counts=use_ext_counts,
                 emit_xstate=emit_xstate,
+                handshake=handshake,
                 use_ext_cand=use_ext_cand,
                 ext_rungs=ext_rungs,
                 cand_cap=cand_cap,
@@ -6834,6 +6839,7 @@ def _compile(
                 seed_thr=seed_thr_fake,
                 seed_counts=None,
                 xstate=xstate_fake,
+                arrival=arrival_fake,
                 cand_vals=cand_vals_fake,
                 cand_idx=cand_idx_fake,
                 cand_ctl=cand_ctl_fake,
@@ -6952,6 +6958,7 @@ def forward(
             cand_idx: Optional[torch.Tensor] = None,
             cand_ctl: Optional[torch.Tensor] = None,
             block_max: Optional[torch.Tensor] = None,
+            arrival: Optional[torch.Tensor] = None,
             num_threads: Optional[int] = None,
             accept_cap: Optional[int] = None,
             kc_override: Optional[int] = None,
@@ -7089,6 +7096,12 @@ def forward(
             use_ext_cand = cand_vals is not None
             enable_block_skip = block_max is not None
             emit_xstate = xstate is not None
+            handshake = arrival is not None
+            if handshake:
+                assert (arrival.dtype == torch.int32 and arrival.is_cuda
+                        and arrival.is_contiguous() and arrival.numel()
+                        >= num_rows), ("arrival must be int32 [>= num_rows]")
+                assert next_n == 1, "handshake requires next_n == 1"
             if use_ext_cand:
                 assert use_ext_counts, (
                     "candidate list requires the packed seed row "
@@ -7123,12 +7136,13 @@ def forward(
                 cand_cap=(cand_vals.shape[1] if use_ext_cand else 5120),
                 accept_cap=accept_cap,
                 kc_override=kc_override,
+                handshake=handshake,
                 **tuning,
             )
             cls.kernel_cache[key](logits, pre_idx, seq_lens, None,
                                   output_indices, order_row, block_max,
                                   seed_thr, None, xstate, cand_vals, cand_idx,
-                                  cand_ctl)
+                                  cand_ctl, arrival)
 
     # TODO(dsa.py): wire ``order_row = argsort(seq_lens, descending=True)``
     # (device-side, graph-safe) into the LJF row-reorder branch when
@@ -7156,6 +7170,7 @@ def cute_dsl_gvr_topk_decode(
         max_batch_size: Optional[int] = None,
         seed_thr: Optional[torch.Tensor] = None,
         xstate: Optional[torch.Tensor] = None,
+        arrival: Optional[torch.Tensor] = None,
         cand_vals: Optional[torch.Tensor] = None,
         cand_idx: Optional[torch.Tensor] = None,
         cand_ctl: Optional[torch.Tensor] = None,
@@ -7235,6 +7250,7 @@ def cute_dsl_gvr_topk_decode(
             max_batch_size=max_batch_size,
             seed_thr=seed_thr,
             xstate=xstate,
+            arrival=arrival,
             cand_vals=cand_vals,
             cand_idx=cand_idx,
             cand_ctl=cand_ctl,
@@ -7260,6 +7276,7 @@ def _(
         max_batch_size: Optional[int] = None,
         seed_thr: Optional[torch.Tensor] = None,
         xstate: Optional[torch.Tensor] = None,
+        arrival: Optional[torch.Tensor] = None,
         cand_vals: Optional[torch.Tensor] = None,
         cand_idx: Optional[torch.Tensor] = None,
         cand_ctl: Optional[torch.Tensor] = None,
@@ -8391,13 +8408,14 @@ def _compile(cls,
                      emit_cand=False,
                      cand_cap=5120,
                      emit_cand_bucketed=False,
-                     accept_cap=8192):
+                     accept_cap=8192,
+                     emit_arrival=False):
             """Compile kernel using fake tensors + TVM FFI."""
             key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
                    num_sms, num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
                    emit_seed_counts, seed_packed, emit_cand, cand_cap,
-                   emit_cand_bucketed, accept_cap)
+                   emit_cand_bucketed, accept_cap, emit_arrival)
             if key in cls.kernel_cache:
                 return
 
@@ -8511,6 +8529,12 @@ def _compile(cls,
                     cutlass.Int32, (cute.sym_int(), 4),
                     stride_order=(1, 0),
                     assumed_align=4)
+            arrival_fake = None
+            if emit_arrival:
+                arrival_fake = cute.runtime.make_fake_compact_tensor(
+                    cutlass.Int32, (cute.sym_int(), ),
+                    stride_order=(0, ),
+                    assumed_align=4)
             seed_thr_fake = None
             seed_counts_fake = None
             if emit_seed_counts:
@@ -8559,6 +8583,7 @@ def _compile(cls,
                 cand_cap=cand_cap,
                 emit_cand_bucketed=emit_cand_bucketed,
                 accept_cap=accept_cap,
+                emit_arrival=emit_arrival,
             )
 
             compiled = cute.compile(
@@ -8583,6 +8608,7 @@ def _compile(cls,
                 cand_ctl=cand_ctl_fake,
                 cand_idx_t=cand_idx_fake,
                 cand_cur=cand_cur_fake,
+                arrival=arrival_fake,
                 options="--enable-tvm-ffi",
             )
             cls.kernel_cache[key] = compiled
@@ -8618,6 +8644,7 @@ def forward(
             accept_cap: int = 8192,
             cand_idx_out: Optional[torch.Tensor] = None,
             cand_cur_out: Optional[torch.Tensor] = None,
+            arrival_out: Optional[torch.Tensor] = None,
         ) -> torch.Tensor:
             """Execute FP4 paged MQA logits kernel.
 
@@ -8838,11 +8865,20 @@ def forward(
                 cand_cur_out = None
 
             # Compile if needed (fake tensors, no real data required)
+            emit_arrival = arrival_out is not None
+            if emit_arrival:
+                assert (arrival_out.dtype == torch.int32 and arrival_out.is_cuda
+                        and arrival_out.is_contiguous()
+                        and arrival_out.numel() >= B * next_n), (
+                            "arrival_out must be int32 [>= B*next_n] "
+                            "(zero-initialized; consumer self-resets)")
+                assert emit_block_meta and next_n == 1, (
+                    "emit_arrival requires emit_block_meta and next_n == 1")
             key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
                    num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
                    emit_seed_counts, seed_packed, emit_cand, cand_cap,
-                   emit_cand_bucketed, accept_cap)
+                   emit_cand_bucketed, accept_cap, emit_arrival)
             if key not in cls.kernel_cache:
                 cls._compile(
                     compute_block_kv,
@@ -8862,7 +8898,8 @@ def forward(
                     emit_cand=emit_cand,
                     cand_cap=cand_cap,
                     emit_cand_bucketed=emit_cand_bucketed,
-                    accept_cap=accept_cap)
+                    accept_cap=accept_cap,
+                    emit_arrival=emit_arrival)
             compiled = cls.kernel_cache[key]
 
             # TVM FFI: pass raw tensors, no dlpack/stream needed
@@ -8871,11 +8908,11 @@ def forward(
                          context_lens, schedule_meta, num_phys_blocks, B,
                          block_max_out, hit_stats_out, hit_bitmap, seed_thr,
                          seed_counts_out, cand_out, cand_ctl_out, cand_idx_out,
-                         cand_cur_out)
+                         cand_cur_out, arrival_out)
                 return logits, block_max_out, hit_stats_out
             compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                      context_lens, schedule_meta, num_phys_blocks, B, None,
-                     None, None, None, None, None, None, None, None)
+                     None, None, None, None, None, None, None, None, None)
             return logits
 
     # NOTE: the optional emission tensors ARE written by the kernel but
@@ -8906,6 +8943,7 @@ def cute_dsl_fp4_paged_mqa_logits(
         cand_ctl_out: Optional[torch.Tensor] = None,
         cand_cur_out: Optional[torch.Tensor] = None,
         accept_cap: int = 8192,
+        arrival_out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         if not is_sm_100f():
             raise ValueError(
@@ -8954,7 +8992,8 @@ def cute_dsl_fp4_paged_mqa_logits(
             cand_out=cand_out,
             cand_idx_out=cand_idx_out,
             cand_ctl_out=cand_ctl_out,
-            cand_cur_out=cand_cur_out)
+            cand_cur_out=cand_cur_out,
+            arrival_out=arrival_out)
         # with emission on, the runner returns (logits, block_max,
         # hit_stats) - the emission buffers are caller-owned mutates,
         # the op face stays logits-only
@@ -8981,6 +9020,7 @@ def _(
         cand_ctl_out: Optional[torch.Tensor] = None,
         cand_cur_out: Optional[torch.Tensor] = None,
         accept_cap: int = 8192,
+        arrival_out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         B = q.shape[0]
         next_n = q.shape[1]
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index 2481a5c5b8b3..8bb3ac4a3190 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -60,6 +60,9 @@
 from cutlass.cutlass_dsl import T, dsl_user_op
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 
+from ..top_k.single_pass_multi_cta_radix_topk import red_release_gpu
+from ..utils import griddepcontrol_launch_dependents
+
 # CuTe DSL CUDA 13 validates rounding modes as string literals. The string
 # form is also accepted by older wrappers, so keep it version-independent.
 _RND_RN = "rn"
@@ -469,6 +472,7 @@ def __init__(
         emit_block_meta: bool = False,
         emit_hit_stats: bool = True,
         emit_seed_counts: bool = False,
+        emit_arrival: bool = False,
         seed_packed: bool = False,
         emit_cand: bool = False,
         cand_cap: int = 5120,
@@ -562,6 +566,15 @@ def __init__(
         if emit_seed_counts and not emit_block_meta:
             raise ValueError("emit_seed_counts requires emit_block_meta")
         self.emit_seed_counts = emit_seed_counts
+        # emit_arrival (fused-handshake mode): per-row arrival counters -
+        # every math warp publishes its split count for a finished row with
+        # a gpu-scope release; the dependent top-k kernel spin-gates each
+        # row on 8 * ceil(num_kv_tiles / 2) instead of waiting for the
+        # whole grid. Rides the q-transition flush points, so it needs the
+        # meta machinery.
+        if emit_arrival and not emit_block_meta:
+            raise ValueError("emit_arrival requires emit_block_meta")
+        self.emit_arrival = emit_arrival
         # seed_packed: single [num_rows, 8] fp32 seed row per the top-k
         # pre-packed contract - lines at cols 0..2, counts ACCUMULATED AS
         # FLOATS at cols 3..5 (exact to 2^24; red.global.add.f32). The
@@ -828,6 +841,7 @@ def __call__(
         cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 {claimed, void}, zeroed
         cand_idx_t: cute.Tensor = None,  # bucketed: [num_rows, 2*segA+capC] int32 SoA
         cand_cur: cute.Tensor = None,  # bucketed: [num_rows, 4] int32 cursors, zeroed
+        arrival: cute.Tensor = None,  # [num_rows] int32 arrival counters, zeroed
     ):
         # Derive KV data and SF views from the fused uint8 buffer.
         # Fused layout per phys block: [data half_head_dim*phys_block_kv bytes]
@@ -1035,6 +1049,7 @@ class SharedStorage:
             cand_ctl,
             cand_idx_t,
             cand_cur,
+            arrival,
             self.cluster_layout_vmnk,
             self.a_smem_layout_staged,
             self.b_smem_layout_staged,
@@ -1236,6 +1251,7 @@ def kernel(
         mCandCtl: cute.Tensor,  # [num_rows, 2] int32 {claimed, void} (or None)
         mCandIdx: cute.Tensor,  # bucketed: [num_rows, 2*segA+capC] int32 SoA (or None)
         mCandCur: cute.Tensor,  # bucketed: [num_rows, 4] int32 cursors (or None)
+        mArrival: cute.Tensor,  # [num_rows] int32 arrival counters (or None)
         cluster_layout_vmnk: cute.Layout,
         a_smem_layout_staged: cute.ComposedLayout,
         b_smem_layout_staged: cute.ComposedLayout,
@@ -1290,6 +1306,12 @@ def kernel(
         start_q_clamped = min(start_q, batch_size - 1)
         current_num_kv = (mContextLens[start_q_clamped] + self.block_kv - 1) // self.block_kv
 
+        if cutlass.const_expr(self.emit_arrival):
+            # fused handshake: release the dependent (top-k) grid NOW -
+            # data safety rides on the per-row arrival counters, PDL only
+            # lets consumer CTAs stage onto SMs as producer CTAs retire.
+            griddepcontrol_launch_dependents()
+
         if is_tma_warp:
             cpasync.prefetch_descriptor(tma_atom_a)
             cpasync.prefetch_descriptor(tma_atom_b)
@@ -2277,6 +2299,9 @@ def kernel(
                         meta_j = cutlass.Int32(0)
                         hitw_batch = cutlass.Int32(0)
 
+                # fused handshake: this warp's split count for the current row
+                arr_cnt = cutlass.Int32(0)
+
                 while has_work:
                     # fetch_next_task: commit next → current
                     q_idx_old = q_idx
@@ -2344,6 +2369,21 @@ def kernel(
                                     self._flush_cand_window_bucketed(
                                         mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane
                                     )
+                            if cutlass.const_expr(self.emit_arrival):
+                                if q_idx_old < batch_size:
+                                    # publish this warp's split count for
+                                    # the finished row: sync_warp + the
+                                    # fence inside red_release_gpu order
+                                    # every prior write of this warp's
+                                    # lanes (logits STGs + emission
+                                    # atomics) before the add lands
+                                    cute.arch.sync_warp()
+                                    if meta_lane == 0:
+                                        red_release_gpu(
+                                            mArrival.iterator + q_idx_old,
+                                            arr_cnt,
+                                        )
+                                arr_cnt = cutlass.Int32(0)
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 0 (kv_idx + 0)
@@ -2949,6 +2989,9 @@ def kernel(
                             out_row = q_idx * next_n + t
                             mLogits[(out_row, kv_pos)] = result_arr[t]
 
+                    if cutlass.const_expr(self.emit_arrival):
+                        arr_cnt = arr_cnt + cutlass.Int32(1)
+
                     # Advance: inline fetch_next_task
                     next_kv_idx = kv_idx + NUM_MATH_WG
                     if next_kv_idx >= num_kv:
@@ -2981,6 +3024,12 @@ def kernel(
                             mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane
                         )
 
+                if cutlass.const_expr(self.emit_arrival):
+                    if q_idx < batch_size:
+                        cute.arch.sync_warp()
+                        if meta_lane == 0:
+                            red_release_gpu(mArrival.iterator + q_idx, arr_cnt)
+
                 # Release last Q stage (WG 0)
                 if q_idx < batch_size:
                     q_pipeline.consumer_release(q_cons_state)
@@ -3070,6 +3119,9 @@ def kernel(
                         meta_j = cutlass.Int32(0)
                         hitw_batch = cutlass.Int32(0)
 
+                # fused handshake: this warp's split count for the current row
+                arr_cnt = cutlass.Int32(0)
+
                 while has_work:
                     # fetch_next_task: commit next → current
                     q_idx_old = q_idx
@@ -3136,6 +3188,21 @@ def kernel(
                                     self._flush_cand_window_bucketed(
                                         mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane
                                     )
+                            if cutlass.const_expr(self.emit_arrival):
+                                if q_idx_old < batch_size:
+                                    # publish this warp's split count for
+                                    # the finished row: sync_warp + the
+                                    # fence inside red_release_gpu order
+                                    # every prior write of this warp's
+                                    # lanes (logits STGs + emission
+                                    # atomics) before the add lands
+                                    cute.arch.sync_warp()
+                                    if meta_lane == 0:
+                                        red_release_gpu(
+                                            mArrival.iterator + q_idx_old,
+                                            arr_cnt,
+                                        )
+                                arr_cnt = cutlass.Int32(0)
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 1 (kv_idx + 1)
@@ -3735,6 +3802,9 @@ def kernel(
                             out_row = q_idx * next_n + t
                             mLogits[(out_row, kv_pos)] = result_arr[t]
 
+                    if cutlass.const_expr(self.emit_arrival):
+                        arr_cnt = arr_cnt + cutlass.Int32(1)
+
                     # Advance: inline fetch_next_task
                     next_kv_idx = kv_idx + NUM_MATH_WG
                     if next_kv_idx >= num_kv:
@@ -3767,6 +3837,12 @@ def kernel(
                             mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane
                         )
 
+                if cutlass.const_expr(self.emit_arrival):
+                    if q_idx < batch_size:
+                        cute.arch.sync_warp()
+                        if meta_lane == 0:
+                            red_release_gpu(mArrival.iterator + q_idx, arr_cnt)
+
                 # Release last Q stage (WG 1)
                 if q_idx < batch_size:
                     q_pipeline.consumer_release(q_cons_state)
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 10de11addfc8..9ac9dd8189ce 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -41,6 +41,7 @@
 
 from ..utils import TRTLLM_ENABLE_PDL, griddepcontrol_launch_dependents, griddepcontrol_wait
 from .block_scan import warp_scan
+from .single_pass_multi_cta_radix_topk import ld_acquire_gpu, st_release_gpu
 
 
 def _env_flag(name: str) -> bool:
@@ -274,6 +275,7 @@ def __init__(
         smem_cache_elems: int = 32768,
         seqlen_sorted: bool = False,
         kc_diet: Optional[bool] = None,
+        handshake: bool = False,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
@@ -527,6 +529,12 @@ def __init__(
         # kernel passes False for BOTH member instances so their SMEM layouts
         # stay byte-identical (the DSL sizes the launch from the last-traced
         # SmemAllocator only; see GvrTopKLBKernel).
+        # handshake (fused-op mode): spin-gate each row on the producer
+        # indexer's per-row arrival counter instead of relying on the
+        # grid-wide PDL wait. MVP scope pins next_n == 1 / cluster_size == 1.
+        self.handshake = bool(handshake)
+        if self.handshake and (next_n != 1 or cluster_size != 1):
+            raise ValueError("handshake requires next_n == 1 and cluster_size == 1")
         if kc_diet is None:
             kc_diet = cluster_size == 1
         if enable_r0 and top_k == 512 and kc_diet and self.kC > 3072:
@@ -4809,6 +4817,7 @@ def gvr_topk_kernel(
         cand_vals: cute.Tensor,  # [numRows, CAP] fp32 scores (or None)
         cand_idx: cute.Tensor,  # [numRows, CAP] int32 positions (or None)
         cand_ctl: cute.Tensor,  # [numRows, 2] int32 {claimed, void} (or None)
+        arrival: cute.Tensor,  # [numRows] int32 arrival counters (or None)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
 
@@ -4867,6 +4876,7 @@ def gvr_topk_kernel(
             cand_vals=cand_vals,
             cand_idx=cand_idx,
             cand_ctl=cand_ctl,
+            arrival=arrival,
         )
 
     @cute.jit
@@ -4885,6 +4895,7 @@ def run_one_row(
         cand_vals: cute.Tensor = None,  # [numRows, CAP] fp32 (ext cand)
         cand_idx: cute.Tensor = None,  # [numRows, CAP] int32 (ext cand)
         cand_ctl: cute.Tensor = None,  # [numRows, 2] int32 (ext cand)
+        arrival: cute.Tensor = None,  # [numRows] int32 arrival counters
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
 
@@ -5217,6 +5228,22 @@ def run_one_row(
             s_cluster_partial_m = None
             smem_gath = None
 
+        if cutlass.const_expr(self.handshake):
+            # fused handshake: gate this row on the producer indexer's
+            # arrival counter (8 math warps each publish their split count;
+            # total = 8 * ceil(num_kv_tiles / 2)), then reset the counter
+            # for the next layer's reuse. The reset cannot race the next
+            # indexer launch: its CTAs only start after this kernel's
+            # trailing launch_dependents (PDL stream chain).
+            nkv_hs = (N + cutlass.Int32(127)) // cutlass.Int32(128)
+            target_hs = cutlass.Int32(8) * ((nkv_hs + cutlass.Int32(1)) // cutlass.Int32(2))
+            arr_ptr_hs = arrival.iterator + row_idx
+            if tidx == cutlass.Int32(0):
+                while ld_acquire_gpu(arr_ptr_hs) < target_hs:
+                    pass
+                st_release_gpu(arr_ptr_hs, cutlass.Int32(0))
+            cute.arch.barrier()
+
         # ---- Per-row dispatch ----
         # Three branches:
         #   1. Degenerate (N <= top_k): no GVR work, leader emits identity.
@@ -7159,6 +7186,7 @@ def __call__(
         cand_vals: cute.Tensor = None,  # [num_rows, CAP] fp32 (ext cand)
         cand_idx: cute.Tensor = None,  # [num_rows, CAP] int32 (ext cand)
         cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 (ext cand)
+        arrival: cute.Tensor = None,  # [num_rows] int32 (handshake)
     ):
         num_rows = input_data.shape[0]
         cluster_size = cutlass.const_expr(self.cluster_size)
@@ -7187,6 +7215,7 @@ def __call__(
             cand_vals,
             cand_idx,
             cand_ctl,
+            arrival,
         ).launch(
             grid=(total_ctas, 1, 1),
             block=(self.num_threads, 1, 1),

From bc2621b1ba394b31284c8b1ab3e73754b7003ee9 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 18:24:38 -0700
Subject: [PATCH 059/117] Revert "[None][perf] Fused indexer->top-k handshake:
 PDL early release + per-row arrival gating"

This reverts commit 8af7c59e9e.

The handshake is measurably net-negative and the reason is structural,
not a tuning miss. Cold graph-replay A/B across six shapes put the
fused pair at 0.81-0.97x of the unfused one, and the large-batch
hypothesis - that rows become CTA-local and therefore finish early -
was falsified too (0.91-0.97x at B32-256). A three-arm ablation
attributes the loss to the protocol: the producer-side fence and
counter cost +3.2us at small batch and the consumer spin adds another
1.8-3.7us, against a measured boundary cost of only 2.2-2.6us.

Root cause: the DeepGEMM scheduler deliberately spreads each row
across many CTAs for load balance (256-token halves, split without
regard to row boundaries), which is antagonistic to per-row early
completion - a row is ready only when essentially the whole grid is,
so the gate buys nothing while the protocol still charges. With equal
row lengths no schedule can help, because every row finishes at the
same time.

The implementation is kept on perf/fused-handshake-falsified.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    | 26 +------
 .../attention_backend/sparse/gvr_ext.py       | 13 +---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  | 58 +++-----------
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 76 -------------------
 .../blackwell/top_k/gvr_topk_decode.py        | 29 -------
 5 files changed, 11 insertions(+), 191 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 6b520a5b19da..d3ca287eba94 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1819,12 +1819,6 @@ def __init__(self,
                             and self.use_cute_dsl_topk
                             and self.use_cute_dsl_paged_mqa_logits
                             and sparse_params.indexer_k_dtype == "fp4")
-        # fused-handshake mode (TRTLLM_GVR_FUSE=1, requires GVR_EXT): the
-        # indexer emits per-row arrival counters and releases the top-k
-        # grid early via PDL; the top-k spin-gates each row on its counter
-        # instead of waiting for the whole indexer grid.
-        self.use_gvr_fuse = (os.environ.get("TRTLLM_GVR_FUSE", "0") == "1"
-                             and self.use_gvr_ext)
         self._gvr_ext = None  # lazy GvrExtState (first decode step)
         self._gvr_route = None
         self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5
@@ -2790,8 +2784,7 @@ def sparse_attn_indexer(
                             self._gvr_ext = GvrExtState(
                                 max_rows=metadata.max_num_sequences,
                                 top_k=self.index_topk,
-                                device=q_fp8.device,
-                                enable_fused_handshake=self.use_gvr_fuse)
+                                device=q_fp8.device)
                         st = self._gvr_ext
                         n_comp = indexer_max_seq_len // max(
                             self.compress_ratio, 1)
@@ -2807,19 +2800,6 @@ def sparse_attn_indexer(
                             gvr_emit_kwargs["block_max_out"] = (
                                 st.ensure_block_max(
                                     indexer_max_seq_len)[:batch_size])
-                        # fused handshake only when BOTH sides of this
-                        # step run an ext tier: the first step emits
-                        # counts but its top-k still routes "rungs" and
-                        # would never consume/reset the counters -
-                        # skipping emission keeps them clean
-                        self._gvr_emitted_arrival = (
-                            st.arrival is not None
-                            and emit_tier in ("counts", "list")
-                            and self._gvr_route.tier in ("counts", "list")
-                            and self._gvr_route.cluster_size == 1)
-                        if self._gvr_emitted_arrival:
-                            gvr_emit_kwargs["arrival_out"] = (
-                                st.arrival[:batch_size])
                     logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits(
                         dsl_q, decode_q_scale, k_cache, weights_decode,
                         dsl_context_lens, dsl_block_table, dsl_schedule_meta,
@@ -2916,10 +2896,6 @@ def sparse_attn_indexer(
                         self._gvr_route, num_gen_tokens,
                         st.block_max[:num_gen_tokens]
                         if st.block_max is not None else None)
-                    if getattr(self, "_gvr_emitted_arrival", False):
-                        # fused handshake: gate each row on this step's
-                        # indexer arrival counter (consumer resets it)
-                        ext_kw["arrival"] = st.arrival[:num_gen_tokens]
                     torch.ops.trtllm.cute_dsl_gvr_topk_decode(
                         logits_decode,
                         st.prev_topk[:num_gen_tokens],
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
index ed3b30e5a17b..0611bbd10df7 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -52,12 +52,7 @@ class GvrExtState:
     """Per-attention-backend emission state (persistent buffers)."""
 
     def __init__(
-        self,
-        max_rows: int,
-        top_k: int,
-        device: torch.device,
-        enable_list_tier: bool = True,
-        enable_fused_handshake: bool = False,
+        self, max_rows: int, top_k: int, device: torch.device, enable_list_tier: bool = True
     ):
         self.max_rows = max_rows
         self.top_k = top_k
@@ -85,12 +80,6 @@ def __init__(
         # tier the PREVIOUS indexer call emitted (what this step's
         # top-k may consume); "rungs" until the first emission lands
         self.emitted_tier = "rungs"
-        # fused-handshake arrival counters (zero-once: the consumer
-        # resets each row right after gating on it, so steady-state
-        # layer-to-layer reuse needs no host-side zeroing)
-        self.arrival: Optional[torch.Tensor] = None
-        if enable_fused_handshake:
-            self.arrival = torch.zeros((max_rows,), dtype=torch.int32, device=device)
 
     def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
         nb4 = ((max_seq_len + 255) // 256 * 256) // 128 * 4
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 5e30e87a400b..a2d6a5e09a2b 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -6724,7 +6724,6 @@ def _compile(
             cand_cap: int = 5120,
             accept_cap: Optional[int] = None,
             kc_override: Optional[int] = None,
-            handshake: bool = False,
         ) -> tuple:
             key = (dtype, top_k, next_n, enable_unroll_4, enable_phase3_unroll,
                    use_constant_hint, min_blocks_per_mp, use_256bit_load,
@@ -6732,7 +6731,7 @@ def _compile(
                    compress_ratio, return_output_values, cluster_size,
                    seqlen_sorted, enable_block_skip, use_ext_counts,
                    emit_xstate, use_ext_cand, ext_rungs, cand_cap, accept_cap,
-                   kc_override, handshake)
+                   kc_override)
             if key in cls.kernel_cache:
                 return key
             n_rows = cute.sym_int()
@@ -6779,9 +6778,6 @@ def _compile(
                 cutlass.Float32, (n_rows, 8),
                 stride_order=(1, 0),
                 assumed_align=4) if emit_xstate else None)
-            arrival_fake = (cute.runtime.make_fake_compact_tensor(
-                cutlass.Int32, (n_rows, ), stride_order=(0, ), assumed_align=4)
-                            if handshake else None)
             cand_vals_fake = (cute.runtime.make_fake_compact_tensor(
                 cutlass.Float32, (n_rows, cand_cap),
                 stride_order=(1, 0),
@@ -6815,7 +6811,6 @@ def _compile(
                 enable_block_skip=enable_block_skip,
                 use_ext_counts=use_ext_counts,
                 emit_xstate=emit_xstate,
-                handshake=handshake,
                 use_ext_cand=use_ext_cand,
                 ext_rungs=ext_rungs,
                 cand_cap=cand_cap,
@@ -6839,7 +6834,6 @@ def _compile(
                 seed_thr=seed_thr_fake,
                 seed_counts=None,
                 xstate=xstate_fake,
-                arrival=arrival_fake,
                 cand_vals=cand_vals_fake,
                 cand_idx=cand_idx_fake,
                 cand_ctl=cand_ctl_fake,
@@ -6958,7 +6952,6 @@ def forward(
             cand_idx: Optional[torch.Tensor] = None,
             cand_ctl: Optional[torch.Tensor] = None,
             block_max: Optional[torch.Tensor] = None,
-            arrival: Optional[torch.Tensor] = None,
             num_threads: Optional[int] = None,
             accept_cap: Optional[int] = None,
             kc_override: Optional[int] = None,
@@ -7096,12 +7089,6 @@ def forward(
             use_ext_cand = cand_vals is not None
             enable_block_skip = block_max is not None
             emit_xstate = xstate is not None
-            handshake = arrival is not None
-            if handshake:
-                assert (arrival.dtype == torch.int32 and arrival.is_cuda
-                        and arrival.is_contiguous() and arrival.numel()
-                        >= num_rows), ("arrival must be int32 [>= num_rows]")
-                assert next_n == 1, "handshake requires next_n == 1"
             if use_ext_cand:
                 assert use_ext_counts, (
                     "candidate list requires the packed seed row "
@@ -7136,13 +7123,12 @@ def forward(
                 cand_cap=(cand_vals.shape[1] if use_ext_cand else 5120),
                 accept_cap=accept_cap,
                 kc_override=kc_override,
-                handshake=handshake,
                 **tuning,
             )
             cls.kernel_cache[key](logits, pre_idx, seq_lens, None,
                                   output_indices, order_row, block_max,
                                   seed_thr, None, xstate, cand_vals, cand_idx,
-                                  cand_ctl, arrival)
+                                  cand_ctl)
 
     # TODO(dsa.py): wire ``order_row = argsort(seq_lens, descending=True)``
     # (device-side, graph-safe) into the LJF row-reorder branch when
@@ -7170,7 +7156,6 @@ def cute_dsl_gvr_topk_decode(
         max_batch_size: Optional[int] = None,
         seed_thr: Optional[torch.Tensor] = None,
         xstate: Optional[torch.Tensor] = None,
-        arrival: Optional[torch.Tensor] = None,
         cand_vals: Optional[torch.Tensor] = None,
         cand_idx: Optional[torch.Tensor] = None,
         cand_ctl: Optional[torch.Tensor] = None,
@@ -7250,7 +7235,6 @@ def cute_dsl_gvr_topk_decode(
             max_batch_size=max_batch_size,
             seed_thr=seed_thr,
             xstate=xstate,
-            arrival=arrival,
             cand_vals=cand_vals,
             cand_idx=cand_idx,
             cand_ctl=cand_ctl,
@@ -7276,7 +7260,6 @@ def _(
         max_batch_size: Optional[int] = None,
         seed_thr: Optional[torch.Tensor] = None,
         xstate: Optional[torch.Tensor] = None,
-        arrival: Optional[torch.Tensor] = None,
         cand_vals: Optional[torch.Tensor] = None,
         cand_idx: Optional[torch.Tensor] = None,
         cand_ctl: Optional[torch.Tensor] = None,
@@ -8408,14 +8391,13 @@ def _compile(cls,
                      emit_cand=False,
                      cand_cap=5120,
                      emit_cand_bucketed=False,
-                     accept_cap=8192,
-                     emit_arrival=False):
+                     accept_cap=8192):
             """Compile kernel using fake tensors + TVM FFI."""
             key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
                    num_sms, num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
                    emit_seed_counts, seed_packed, emit_cand, cand_cap,
-                   emit_cand_bucketed, accept_cap, emit_arrival)
+                   emit_cand_bucketed, accept_cap)
             if key in cls.kernel_cache:
                 return
 
@@ -8529,12 +8511,6 @@ def _compile(cls,
                     cutlass.Int32, (cute.sym_int(), 4),
                     stride_order=(1, 0),
                     assumed_align=4)
-            arrival_fake = None
-            if emit_arrival:
-                arrival_fake = cute.runtime.make_fake_compact_tensor(
-                    cutlass.Int32, (cute.sym_int(), ),
-                    stride_order=(0, ),
-                    assumed_align=4)
             seed_thr_fake = None
             seed_counts_fake = None
             if emit_seed_counts:
@@ -8583,7 +8559,6 @@ def _compile(cls,
                 cand_cap=cand_cap,
                 emit_cand_bucketed=emit_cand_bucketed,
                 accept_cap=accept_cap,
-                emit_arrival=emit_arrival,
             )
 
             compiled = cute.compile(
@@ -8608,7 +8583,6 @@ def _compile(cls,
                 cand_ctl=cand_ctl_fake,
                 cand_idx_t=cand_idx_fake,
                 cand_cur=cand_cur_fake,
-                arrival=arrival_fake,
                 options="--enable-tvm-ffi",
             )
             cls.kernel_cache[key] = compiled
@@ -8644,7 +8618,6 @@ def forward(
             accept_cap: int = 8192,
             cand_idx_out: Optional[torch.Tensor] = None,
             cand_cur_out: Optional[torch.Tensor] = None,
-            arrival_out: Optional[torch.Tensor] = None,
         ) -> torch.Tensor:
             """Execute FP4 paged MQA logits kernel.
 
@@ -8865,20 +8838,11 @@ def forward(
                 cand_cur_out = None
 
             # Compile if needed (fake tensors, no real data required)
-            emit_arrival = arrival_out is not None
-            if emit_arrival:
-                assert (arrival_out.dtype == torch.int32 and arrival_out.is_cuda
-                        and arrival_out.is_contiguous()
-                        and arrival_out.numel() >= B * next_n), (
-                            "arrival_out must be int32 [>= B*next_n] "
-                            "(zero-initialized; consumer self-resets)")
-                assert emit_block_meta and next_n == 1, (
-                    "emit_arrival requires emit_block_meta and next_n == 1")
             key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
                    num_epi_subtiles, epi_dtype, output_dtype,
                    remove_online_sf_transpose, emit_block_meta, emit_hit_stats,
                    emit_seed_counts, seed_packed, emit_cand, cand_cap,
-                   emit_cand_bucketed, accept_cap, emit_arrival)
+                   emit_cand_bucketed, accept_cap)
             if key not in cls.kernel_cache:
                 cls._compile(
                     compute_block_kv,
@@ -8898,8 +8862,7 @@ def forward(
                     emit_cand=emit_cand,
                     cand_cap=cand_cap,
                     emit_cand_bucketed=emit_cand_bucketed,
-                    accept_cap=accept_cap,
-                    emit_arrival=emit_arrival)
+                    accept_cap=accept_cap)
             compiled = cls.kernel_cache[key]
 
             # TVM FFI: pass raw tensors, no dlpack/stream needed
@@ -8908,11 +8871,11 @@ def forward(
                          context_lens, schedule_meta, num_phys_blocks, B,
                          block_max_out, hit_stats_out, hit_bitmap, seed_thr,
                          seed_counts_out, cand_out, cand_ctl_out, cand_idx_out,
-                         cand_cur_out, arrival_out)
+                         cand_cur_out)
                 return logits, block_max_out, hit_stats_out
             compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table,
                      context_lens, schedule_meta, num_phys_blocks, B, None,
-                     None, None, None, None, None, None, None, None, None)
+                     None, None, None, None, None, None, None, None)
             return logits
 
     # NOTE: the optional emission tensors ARE written by the kernel but
@@ -8943,7 +8906,6 @@ def cute_dsl_fp4_paged_mqa_logits(
         cand_ctl_out: Optional[torch.Tensor] = None,
         cand_cur_out: Optional[torch.Tensor] = None,
         accept_cap: int = 8192,
-        arrival_out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         if not is_sm_100f():
             raise ValueError(
@@ -8992,8 +8954,7 @@ def cute_dsl_fp4_paged_mqa_logits(
             cand_out=cand_out,
             cand_idx_out=cand_idx_out,
             cand_ctl_out=cand_ctl_out,
-            cand_cur_out=cand_cur_out,
-            arrival_out=arrival_out)
+            cand_cur_out=cand_cur_out)
         # with emission on, the runner returns (logits, block_max,
         # hit_stats) - the emission buffers are caller-owned mutates,
         # the op face stays logits-only
@@ -9020,7 +8981,6 @@ def _(
         cand_ctl_out: Optional[torch.Tensor] = None,
         cand_cur_out: Optional[torch.Tensor] = None,
         accept_cap: int = 8192,
-        arrival_out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         B = q.shape[0]
         next_n = q.shape[1]
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index 8bb3ac4a3190..2481a5c5b8b3 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -60,9 +60,6 @@
 from cutlass.cutlass_dsl import T, dsl_user_op
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 
-from ..top_k.single_pass_multi_cta_radix_topk import red_release_gpu
-from ..utils import griddepcontrol_launch_dependents
-
 # CuTe DSL CUDA 13 validates rounding modes as string literals. The string
 # form is also accepted by older wrappers, so keep it version-independent.
 _RND_RN = "rn"
@@ -472,7 +469,6 @@ def __init__(
         emit_block_meta: bool = False,
         emit_hit_stats: bool = True,
         emit_seed_counts: bool = False,
-        emit_arrival: bool = False,
         seed_packed: bool = False,
         emit_cand: bool = False,
         cand_cap: int = 5120,
@@ -566,15 +562,6 @@ def __init__(
         if emit_seed_counts and not emit_block_meta:
             raise ValueError("emit_seed_counts requires emit_block_meta")
         self.emit_seed_counts = emit_seed_counts
-        # emit_arrival (fused-handshake mode): per-row arrival counters -
-        # every math warp publishes its split count for a finished row with
-        # a gpu-scope release; the dependent top-k kernel spin-gates each
-        # row on 8 * ceil(num_kv_tiles / 2) instead of waiting for the
-        # whole grid. Rides the q-transition flush points, so it needs the
-        # meta machinery.
-        if emit_arrival and not emit_block_meta:
-            raise ValueError("emit_arrival requires emit_block_meta")
-        self.emit_arrival = emit_arrival
         # seed_packed: single [num_rows, 8] fp32 seed row per the top-k
         # pre-packed contract - lines at cols 0..2, counts ACCUMULATED AS
         # FLOATS at cols 3..5 (exact to 2^24; red.global.add.f32). The
@@ -841,7 +828,6 @@ def __call__(
         cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 {claimed, void}, zeroed
         cand_idx_t: cute.Tensor = None,  # bucketed: [num_rows, 2*segA+capC] int32 SoA
         cand_cur: cute.Tensor = None,  # bucketed: [num_rows, 4] int32 cursors, zeroed
-        arrival: cute.Tensor = None,  # [num_rows] int32 arrival counters, zeroed
     ):
         # Derive KV data and SF views from the fused uint8 buffer.
         # Fused layout per phys block: [data half_head_dim*phys_block_kv bytes]
@@ -1049,7 +1035,6 @@ class SharedStorage:
             cand_ctl,
             cand_idx_t,
             cand_cur,
-            arrival,
             self.cluster_layout_vmnk,
             self.a_smem_layout_staged,
             self.b_smem_layout_staged,
@@ -1251,7 +1236,6 @@ def kernel(
         mCandCtl: cute.Tensor,  # [num_rows, 2] int32 {claimed, void} (or None)
         mCandIdx: cute.Tensor,  # bucketed: [num_rows, 2*segA+capC] int32 SoA (or None)
         mCandCur: cute.Tensor,  # bucketed: [num_rows, 4] int32 cursors (or None)
-        mArrival: cute.Tensor,  # [num_rows] int32 arrival counters (or None)
         cluster_layout_vmnk: cute.Layout,
         a_smem_layout_staged: cute.ComposedLayout,
         b_smem_layout_staged: cute.ComposedLayout,
@@ -1306,12 +1290,6 @@ def kernel(
         start_q_clamped = min(start_q, batch_size - 1)
         current_num_kv = (mContextLens[start_q_clamped] + self.block_kv - 1) // self.block_kv
 
-        if cutlass.const_expr(self.emit_arrival):
-            # fused handshake: release the dependent (top-k) grid NOW -
-            # data safety rides on the per-row arrival counters, PDL only
-            # lets consumer CTAs stage onto SMs as producer CTAs retire.
-            griddepcontrol_launch_dependents()
-
         if is_tma_warp:
             cpasync.prefetch_descriptor(tma_atom_a)
             cpasync.prefetch_descriptor(tma_atom_b)
@@ -2299,9 +2277,6 @@ def kernel(
                         meta_j = cutlass.Int32(0)
                         hitw_batch = cutlass.Int32(0)
 
-                # fused handshake: this warp's split count for the current row
-                arr_cnt = cutlass.Int32(0)
-
                 while has_work:
                     # fetch_next_task: commit next → current
                     q_idx_old = q_idx
@@ -2369,21 +2344,6 @@ def kernel(
                                     self._flush_cand_window_bucketed(
                                         mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane
                                     )
-                            if cutlass.const_expr(self.emit_arrival):
-                                if q_idx_old < batch_size:
-                                    # publish this warp's split count for
-                                    # the finished row: sync_warp + the
-                                    # fence inside red_release_gpu order
-                                    # every prior write of this warp's
-                                    # lanes (logits STGs + emission
-                                    # atomics) before the add lands
-                                    cute.arch.sync_warp()
-                                    if meta_lane == 0:
-                                        red_release_gpu(
-                                            mArrival.iterator + q_idx_old,
-                                            arr_cnt,
-                                        )
-                                arr_cnt = cutlass.Int32(0)
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 0 (kv_idx + 0)
@@ -2989,9 +2949,6 @@ def kernel(
                             out_row = q_idx * next_n + t
                             mLogits[(out_row, kv_pos)] = result_arr[t]
 
-                    if cutlass.const_expr(self.emit_arrival):
-                        arr_cnt = arr_cnt + cutlass.Int32(1)
-
                     # Advance: inline fetch_next_task
                     next_kv_idx = kv_idx + NUM_MATH_WG
                     if next_kv_idx >= num_kv:
@@ -3024,12 +2981,6 @@ def kernel(
                             mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane
                         )
 
-                if cutlass.const_expr(self.emit_arrival):
-                    if q_idx < batch_size:
-                        cute.arch.sync_warp()
-                        if meta_lane == 0:
-                            red_release_gpu(mArrival.iterator + q_idx, arr_cnt)
-
                 # Release last Q stage (WG 0)
                 if q_idx < batch_size:
                     q_pipeline.consumer_release(q_cons_state)
@@ -3119,9 +3070,6 @@ def kernel(
                         meta_j = cutlass.Int32(0)
                         hitw_batch = cutlass.Int32(0)
 
-                # fused handshake: this warp's split count for the current row
-                arr_cnt = cutlass.Int32(0)
-
                 while has_work:
                     # fetch_next_task: commit next → current
                     q_idx_old = q_idx
@@ -3188,21 +3136,6 @@ def kernel(
                                     self._flush_cand_window_bucketed(
                                         mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane
                                     )
-                            if cutlass.const_expr(self.emit_arrival):
-                                if q_idx_old < batch_size:
-                                    # publish this warp's split count for
-                                    # the finished row: sync_warp + the
-                                    # fence inside red_release_gpu order
-                                    # every prior write of this warp's
-                                    # lanes (logits STGs + emission
-                                    # atomics) before the add lands
-                                    cute.arch.sync_warp()
-                                    if meta_lane == 0:
-                                        red_release_gpu(
-                                            mArrival.iterator + q_idx_old,
-                                            arr_cnt,
-                                        )
-                                arr_cnt = cutlass.Int32(0)
                             ctx_cur = mContextLens[q_idx]
 
                     # Process KV block for group 1 (kv_idx + 1)
@@ -3802,9 +3735,6 @@ def kernel(
                             out_row = q_idx * next_n + t
                             mLogits[(out_row, kv_pos)] = result_arr[t]
 
-                    if cutlass.const_expr(self.emit_arrival):
-                        arr_cnt = arr_cnt + cutlass.Int32(1)
-
                     # Advance: inline fetch_next_task
                     next_kv_idx = kv_idx + NUM_MATH_WG
                     if next_kv_idx >= num_kv:
@@ -3837,12 +3767,6 @@ def kernel(
                             mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane
                         )
 
-                if cutlass.const_expr(self.emit_arrival):
-                    if q_idx < batch_size:
-                        cute.arch.sync_warp()
-                        if meta_lane == 0:
-                            red_release_gpu(mArrival.iterator + q_idx, arr_cnt)
-
                 # Release last Q stage (WG 1)
                 if q_idx < batch_size:
                     q_pipeline.consumer_release(q_cons_state)
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 9ac9dd8189ce..10de11addfc8 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -41,7 +41,6 @@
 
 from ..utils import TRTLLM_ENABLE_PDL, griddepcontrol_launch_dependents, griddepcontrol_wait
 from .block_scan import warp_scan
-from .single_pass_multi_cta_radix_topk import ld_acquire_gpu, st_release_gpu
 
 
 def _env_flag(name: str) -> bool:
@@ -275,7 +274,6 @@ def __init__(
         smem_cache_elems: int = 32768,
         seqlen_sorted: bool = False,
         kc_diet: Optional[bool] = None,
-        handshake: bool = False,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
@@ -529,12 +527,6 @@ def __init__(
         # kernel passes False for BOTH member instances so their SMEM layouts
         # stay byte-identical (the DSL sizes the launch from the last-traced
         # SmemAllocator only; see GvrTopKLBKernel).
-        # handshake (fused-op mode): spin-gate each row on the producer
-        # indexer's per-row arrival counter instead of relying on the
-        # grid-wide PDL wait. MVP scope pins next_n == 1 / cluster_size == 1.
-        self.handshake = bool(handshake)
-        if self.handshake and (next_n != 1 or cluster_size != 1):
-            raise ValueError("handshake requires next_n == 1 and cluster_size == 1")
         if kc_diet is None:
             kc_diet = cluster_size == 1
         if enable_r0 and top_k == 512 and kc_diet and self.kC > 3072:
@@ -4817,7 +4809,6 @@ def gvr_topk_kernel(
         cand_vals: cute.Tensor,  # [numRows, CAP] fp32 scores (or None)
         cand_idx: cute.Tensor,  # [numRows, CAP] int32 positions (or None)
         cand_ctl: cute.Tensor,  # [numRows, 2] int32 {claimed, void} (or None)
-        arrival: cute.Tensor,  # [numRows] int32 arrival counters (or None)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
 
@@ -4876,7 +4867,6 @@ def gvr_topk_kernel(
             cand_vals=cand_vals,
             cand_idx=cand_idx,
             cand_ctl=cand_ctl,
-            arrival=arrival,
         )
 
     @cute.jit
@@ -4895,7 +4885,6 @@ def run_one_row(
         cand_vals: cute.Tensor = None,  # [numRows, CAP] fp32 (ext cand)
         cand_idx: cute.Tensor = None,  # [numRows, CAP] int32 (ext cand)
         cand_ctl: cute.Tensor = None,  # [numRows, 2] int32 (ext cand)
-        arrival: cute.Tensor = None,  # [numRows] int32 arrival counters
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
 
@@ -5228,22 +5217,6 @@ def run_one_row(
             s_cluster_partial_m = None
             smem_gath = None
 
-        if cutlass.const_expr(self.handshake):
-            # fused handshake: gate this row on the producer indexer's
-            # arrival counter (8 math warps each publish their split count;
-            # total = 8 * ceil(num_kv_tiles / 2)), then reset the counter
-            # for the next layer's reuse. The reset cannot race the next
-            # indexer launch: its CTAs only start after this kernel's
-            # trailing launch_dependents (PDL stream chain).
-            nkv_hs = (N + cutlass.Int32(127)) // cutlass.Int32(128)
-            target_hs = cutlass.Int32(8) * ((nkv_hs + cutlass.Int32(1)) // cutlass.Int32(2))
-            arr_ptr_hs = arrival.iterator + row_idx
-            if tidx == cutlass.Int32(0):
-                while ld_acquire_gpu(arr_ptr_hs) < target_hs:
-                    pass
-                st_release_gpu(arr_ptr_hs, cutlass.Int32(0))
-            cute.arch.barrier()
-
         # ---- Per-row dispatch ----
         # Three branches:
         #   1. Degenerate (N <= top_k): no GVR work, leader emits identity.
@@ -7186,7 +7159,6 @@ def __call__(
         cand_vals: cute.Tensor = None,  # [num_rows, CAP] fp32 (ext cand)
         cand_idx: cute.Tensor = None,  # [num_rows, CAP] int32 (ext cand)
         cand_ctl: cute.Tensor = None,  # [num_rows, 2] int32 (ext cand)
-        arrival: cute.Tensor = None,  # [num_rows] int32 (handshake)
     ):
         num_rows = input_data.shape[0]
         cluster_size = cutlass.const_expr(self.cluster_size)
@@ -7215,7 +7187,6 @@ def __call__(
             cand_vals,
             cand_idx,
             cand_ctl,
-            arrival,
         ).launch(
             grid=(total_ctas, 1, 1),
             block=(self.num_threads, 1, 1),

From b3410e21aafaeb605cfe7e2292cc150336aa1a3e Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 07:37:51 -0700
Subject: [PATCH 060/117] [None][perf] GVR top-k: take the PDL wait after the
 prologue, not at entry

The kernel called griddepcontrol_wait() as its first statement, so a
dependent CTA blocked before doing any of its own setup. Everything
above the first read of producer-written data - row resolution, the
whole SMEM allocation, config folding, and (on a cold call) the
instruction fetch for all of it - can just as well run while the
producer kernel drains: dependent CTAs stage onto SMs as producer CTAs
retire. Instruction-fetch starvation is 37-44% of this kernel's stall
cycles on the small-N cells (ncu, cache-control all), so warming that
under the producer's shadow is worth real time.

Safety: nothing before the new wait site touches producer output.
seq_lens and pre_idx come from host-side metadata and the feedback
buffer; logits / block_max / seed_thr / cand are first read in the
dispatch below it.

The placement is a ctor knob (pdl_wait_late, default on) so both
variants compile into one process for A/B - cross-run cold drift on a
shared node is +-3us, larger than the effect, and only interleaved
in-process comparison resolves it.

Interleaved cold graph-replay A/B of [indexer + top-k], 24 reps per
arm, median (B x compressed-N):
  2x8k   28.51 -> 28.03us (1.017x)
  8x8k   32.03 -> 30.82us (1.040x)
  64x8k  39.30 -> 38.69us (1.016x)
  2x64k  35.97 -> 35.62us (1.010x)
  8x64k  40.13 -> 40.16us (1.000x)

Exactness unchanged (same computation, later sync): unified smoke
30/30, degenerate battery 37 passed + 1 xfailed, ext tiers EXACT,
CUDA-graph replay 11/11, bare-op path identical.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 24 ++++++++++++++++++-
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  4 ++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 10de11addfc8..156da70a9c10 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -274,6 +274,7 @@ def __init__(
         smem_cache_elems: int = 32768,
         seqlen_sorted: bool = False,
         kc_diet: Optional[bool] = None,
+        pdl_wait_late: bool = True,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
@@ -527,6 +528,13 @@ def __init__(
         # kernel passes False for BOTH member instances so their SMEM layouts
         # stay byte-identical (the DSL sizes the launch from the last-traced
         # SmemAllocator only; see GvrTopKLBKernel).
+        # pdl_wait_late: place the PDL wait after the prologue (row
+        # resolution + SMEM allocation + config folding) instead of at
+        # kernel entry, so that work - and its cold instruction fetch -
+        # overlaps the producer's tail. Kept as a knob so both
+        # placements can be A/B'd inside one process (cross-run cold
+        # drift on shared nodes is larger than the effect).
+        self.pdl_wait_late = bool(pdl_wait_late)
         if kc_diet is None:
             kc_diet = cluster_size == 1
         if enable_r0 and top_k == 512 and kc_diet and self.kC > 3072:
@@ -4990,7 +4998,8 @@ def run_one_row(
         output_indices_row = output_indices[row_idx, None]
         pre_idx_count = pre_idx.shape[1]
 
-        griddepcontrol_wait()
+        if cutlass.const_expr(not self.pdl_wait_late):
+            griddepcontrol_wait()
 
         # ---- Shared memory allocation ----
         smem = SmemAllocator()
@@ -5217,6 +5226,19 @@ def run_one_row(
             s_cluster_partial_m = None
             smem_gath = None
 
+        # PDL wait placed as late as possible: everything above (row
+        # resolution, SMEM allocation, config folding - and, on a cold
+        # call, the instruction fetch for all of it) overlaps with the
+        # producer indexer's tail, because dependent CTAs stage onto SMs
+        # as producer CTAs retire. Nothing above reads producer-written
+        # data: seq_lens/pre_idx come from the host-side metadata and the
+        # feedback buffer, while logits / block_max / seed_thr / cand are
+        # first touched below. Instruction-fetch starvation is 37-44% of
+        # this kernel's stall cycles on the small-N cells, so warming it
+        # under the producer's shadow is the point.
+        if cutlass.const_expr(self.pdl_wait_late):
+            griddepcontrol_wait()
+
         # ---- Per-row dispatch ----
         # Three branches:
         #   1. Degenerate (N <= top_k): no GVR work, leader emits identity.
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 58ca7336b6c7..24e928474d5f 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -66,6 +66,7 @@ def _compile(
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
     enable_block_skip: bool = False,
+    pdl_wait_late: bool = True,
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
     use_ext_cand: bool = False,
@@ -203,6 +204,7 @@ def _compile(
         p4_warp_redundant=p4_warp_redundant,
         p2_warp_redundant=p2_warp_redundant,
         enable_block_skip=enable_block_skip,
+        pdl_wait_late=pdl_wait_late,
         use_ext_counts=use_ext_counts,
         emit_xstate=emit_xstate,
         use_ext_cand=use_ext_cand,
@@ -664,6 +666,7 @@ def gvr_topk_decode(
     order_row: Optional[torch.Tensor] = None,
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
+    pdl_wait_late: bool = True,
     block_max: Optional[torch.Tensor] = None,
     skip_min_n: Optional[int] = 200_000,
     seed_thr: Optional[torch.Tensor] = None,
@@ -950,6 +953,7 @@ def gvr_topk_decode(
         p4_warp_redundant,
         p2_warp_redundant,
         enable_block_skip,
+        pdl_wait_late,
         use_ext_counts,
         emit_xstate,
         use_ext_cand,

From 6949d7fc48b9eb78e2b29147c9a6958d24ef8d60 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 10:47:18 -0700
Subject: [PATCH 061/117] [None][perf] GVR P4: filter the fine recursion by
 value range, not by bin recompute

Phase 4's fine recursion exists only to locate the handful of
candidates that land in the straddling coarse bin, yet it recomputed
every candidate's coarse bin to find them: a subtract, a multiply and
two clamps per candidate, then a compare against b*. Membership in
bin b* is exactly 'value in [f_lo, f_hi)', so two compares do the same
job. The clamped ends of the binning fold out-of-range values into bin
0 and bin kBins-1, so those two bins drop the matching side of the
range test and stay bit-identical.

Phase-4 sub-phase stamps (clock64, small-N cells) put the fine
recursion at 44% of the phase - 1.7us of 3.9us - which is what made
this the first cut.

Interleaved cold graph-replay A/B of [indexer + top-k], K=512 cr=4
(V4-Flash production shape), 24 reps per arm, median:
  B2 x 8k    26.24 -> 26.08us (1.006x)
  B8 x 8k    31.74 -> 30.02us (1.058x)
  B64 x 8k   34.46 -> 34.56us (0.997x)
  B2 x 64k   43.84 -> 42.37us (1.035x)
  B8 x 64k   44.80 -> 44.70us (1.002x)

Also measured and rejected on the way: caching each thread's
candidates in registers during the coarse build so the fine pass skips
the SMEM sweep entirely. It removes more work but the extra live
registers push the fp32 kernel past its budget - 0.95-0.99x on 4 of 5
cells. The knob keeps the old path available for A/B.

Exactness: full top-k unit file 708 passed + 1 xfailed, unified smoke
30/30, ext tiers EXACT, CUDA-graph replay 11/11, bare-op identical.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 87 ++++++++++++++-----
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  4 +
 2 files changed, 71 insertions(+), 20 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 156da70a9c10..9a529523b915 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -56,6 +56,9 @@ def _env_flag(name: str) -> bool:
 # P4 sub-phase clock64 breakdown -> xstate[1,2,4,5,6,7] (debug: clobbers
 # the closed-loop thr/anch publish; single-shot cells only, not chains)
 _P4_SUB_DBG = _env_flag("GVR_P4_SUB_DBG")
+# GVR_P4_SUB_DBG=2: publish the P4 HEAD triple (minmax / histogram build /
+# coarse search) instead of the tail triple, so the phase budget adds up.
+_P4_SUB_HEAD = os.environ.get("GVR_P4_SUB_DBG", "0").strip() == "2"
 _SKIP_DBG = _env_flag("GVR_SKIP_DBG")
 
 
@@ -275,6 +278,7 @@ def __init__(
         seqlen_sorted: bool = False,
         kc_diet: Optional[bool] = None,
         pdl_wait_late: bool = True,
+        p4_fine_rangetest: Optional[bool] = None,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
@@ -553,6 +557,15 @@ def __init__(
         # measured as a wash under the same protocol and stay stock.
         if enable_r0 and top_k == 2048 and self.kNumBins > 512:
             self.kNumBins = 512
+        # p4_fine_rangetest: the fine recursion exists only to locate
+        # the handful of candidates inside the straddling coarse bin, so
+        # filter them with a value-range compare instead of recomputing
+        # each candidate's bin (subtract + multiply + two clamps). It is
+        # 44% of Phase 4 on the small-N cells. Caching the candidates in
+        # registers during the coarse build was measured and rejected:
+        # the extra live registers push the fp32 kernel past its budget
+        # (0.95-0.99x on 4 of 5 cells).
+        self.p4_fine_rangetest = True if p4_fine_rangetest is None else bool(p4_fine_rangetest)
         self.r0_qfracs = tuple(float(q) for q in r0_qfracs) if r0_qfracs else ()
         if self.r0_qfracs:
             assert all(0.0 < q < 1.0 for q in self.r0_qfracs), self.r0_qfracs
@@ -3649,22 +3662,50 @@ def phase4_rank_scatter(
                     smem_hist[iz] = cutlass.Int32(0)
                     iz = iz + cutlass.Int32(num_threads)
                 cute.arch.barrier()
-                ifb = tidx
-                while ifb < cand_count:
-                    vf = smem_keys[ifb]
-                    cb = cutlass.Int32((vf - bmin_r) * inv1)
-                    if cb < cutlass.Int32(0):
-                        cb = cutlass.Int32(0)
-                    if cb > cutlass.Int32(kBins - 1):
-                        cb = cutlass.Int32(kBins - 1)
-                    if cb == b_star:
-                        sb = cutlass.Int32((vf - f_lo) * finv)
-                        if sb < cutlass.Int32(0):
-                            sb = cutlass.Int32(0)
-                        if sb > cutlass.Int32(fbins - 1):
-                            sb = cutlass.Int32(fbins - 1)
-                        atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1))
-                    ifb = ifb + cutlass.Int32(num_threads)
+                if cutlass.const_expr(self.p4_fine_rangetest):
+                    # A candidate belongs to bin b* exactly when its value
+                    # lies in [f_lo, f_hi), so the filter is two compares -
+                    # the bin recompute (subtract + multiply + two clamps)
+                    # per candidate is redundant work. The clamped ends of
+                    # the binning fold out-of-range values INTO bin 0 and
+                    # bin kBins-1, so those two bins drop the matching side
+                    # of the range test to stay bit-identical.
+                    f_hi = f_lo + cutlass.Float32(1.0) / inv1
+                    lo_edge = b_star == cutlass.Int32(0)
+                    hi_edge = b_star == cutlass.Int32(kBins - 1)
+                    ifb = tidx
+                    while ifb < cand_count:
+                        vf = smem_keys[ifb]
+                        inb = vf >= f_lo and vf < f_hi
+                        if lo_edge:
+                            inb = vf < f_hi
+                        if hi_edge:
+                            inb = vf >= f_lo
+                        if inb:
+                            sb = cutlass.Int32((vf - f_lo) * finv)
+                            if sb < cutlass.Int32(0):
+                                sb = cutlass.Int32(0)
+                            if sb > cutlass.Int32(fbins - 1):
+                                sb = cutlass.Int32(fbins - 1)
+                            atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1))
+                        ifb = ifb + cutlass.Int32(num_threads)
+                else:
+                    ifb = tidx
+                    while ifb < cand_count:
+                        vfo = smem_keys[ifb]
+                        cbo = cutlass.Int32((vfo - bmin_r) * inv1)
+                        if cbo < cutlass.Int32(0):
+                            cbo = cutlass.Int32(0)
+                        if cbo > cutlass.Int32(kBins - 1):
+                            cbo = cutlass.Int32(kBins - 1)
+                        if cbo == b_star:
+                            sbo = cutlass.Int32((vfo - f_lo) * finv)
+                            if sbo < cutlass.Int32(0):
+                                sbo = cutlass.Int32(0)
+                            if sbo > cutlass.Int32(fbins - 1):
+                                sbo = cutlass.Int32(fbins - 1)
+                            atomicAdd(smem_hist.iterator + sbo, cutlass.Int32(1))
+                        ifb = ifb + cutlass.Int32(num_threads)
                 cute.arch.barrier()
                 # fine 3-step search seeded at rank_above (over fbins bins)
                 fws = cutlass.Int32(0)
@@ -7035,10 +7076,16 @@ def _run_phases(
                             # (minmax/hist/coarse, wcnt[8..10]) are not
                             # published.
                             xstate_row[1] = s_thr[1]
-                            xstate_row[4] = cutlass.Float32(smem_wcnt[11])
-                            xstate_row[5] = cutlass.Float32(smem_wcnt[12])
-                            xstate_row[6] = cutlass.Float32(smem_wcnt[13])
-                            xstate_row[7] = cutlass.Float32(cutlass.Int32(ck_sw1 - ck_sw0))
+                            if cutlass.const_expr(_P4_SUB_HEAD):
+                                xstate_row[4] = cutlass.Float32(smem_wcnt[8])
+                                xstate_row[5] = cutlass.Float32(smem_wcnt[9])
+                                xstate_row[6] = cutlass.Float32(smem_wcnt[10])
+                                xstate_row[7] = cutlass.Float32(smem_wcnt[11])
+                            else:
+                                xstate_row[4] = cutlass.Float32(smem_wcnt[11])
+                                xstate_row[5] = cutlass.Float32(smem_wcnt[12])
+                                xstate_row[6] = cutlass.Float32(smem_wcnt[13])
+                                xstate_row[7] = cutlass.Float32(cutlass.Int32(ck_sw1 - ck_sw0))
                         # cand_count_p4 = pre-P4 snapshot (P4 repurposes
                         # the s_iscalars slots).
                         xstate_row[3] = cutlass.Float32(cand_count_p4)
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 24e928474d5f..23285c5a333f 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -67,6 +67,7 @@ def _compile(
     p2_warp_redundant: bool = True,
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
+    p4_fine_rangetest: "bool | None" = None,
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
     use_ext_cand: bool = False,
@@ -205,6 +206,7 @@ def _compile(
         p2_warp_redundant=p2_warp_redundant,
         enable_block_skip=enable_block_skip,
         pdl_wait_late=pdl_wait_late,
+        p4_fine_rangetest=p4_fine_rangetest,
         use_ext_counts=use_ext_counts,
         emit_xstate=emit_xstate,
         use_ext_cand=use_ext_cand,
@@ -667,6 +669,7 @@ def gvr_topk_decode(
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
+    p4_fine_rangetest: Optional[bool] = None,
     block_max: Optional[torch.Tensor] = None,
     skip_min_n: Optional[int] = 200_000,
     seed_thr: Optional[torch.Tensor] = None,
@@ -954,6 +957,7 @@ def gvr_topk_decode(
         p2_warp_redundant,
         enable_block_skip,
         pdl_wait_late,
+        p4_fine_rangetest,
         use_ext_counts,
         emit_xstate,
         use_ext_cand,

From d66fa12dfeb06cd073d070923d7325fe8d1c083f Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 11:55:08 -0700
Subject: [PATCH 062/117] [None][perf] GVR P4: lane-parallel warp bin sums in
 both bin searches

Both bin searches opened with each warp summing its slice of the
histogram, but the loop indexed by warp only - all 32 lanes walked the
same bins_per_warp bins and built the same running sum. That is a
dependency chain as long as the segment (64 adds at kBins=1024, 16
warps) with the whole warp doing identical work. Spreading the segment
across the lanes and closing with one warp reduction turns it into 2
loads plus 5 shuffles. Same fix in the coarse search and in the fine
recursion's search.

Measurement (this is the methodology the numbers below rely on):
interleaved cold graph replay of [indexer + top-k], both variants
compiled in one process, arms alternated inside each rep, and the
PAIRED difference taken per rep so slow drift cancels. 48 pairs per
cell, K=512 cr=4 (V4-Flash production shape). Comparing per-arm
medians instead - the earlier approach - buries a 3-5% effect under
+-1.3us of per-cell noise; the paired sign test resolves it. Cells with
no effect land at ~50% wins, which is the built-in control.

  B2 x 8k    +1.54us  28/48 wins   (inside noise)
  B8 x 8k    +1.76us  41/48 wins
  B64 x 8k   +0.13us  24/48 wins   (neutral)
  B2 x 64k   +1.54us  40/48 wins
  B8 x 64k   -0.10us  22/48 wins   (neutral)

Exactness: full top-k unit file 708 passed + 1 xfailed, unified smoke
30/30, ext tiers EXACT, CUDA-graph replay 11/11, bare-op identical.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py    | 7 +++++++
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py       | 4 ++++
 2 files changed, 11 insertions(+)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 9a529523b915..938f72a4b587 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -279,6 +279,7 @@ def __init__(
         kc_diet: Optional[bool] = None,
         pdl_wait_late: bool = True,
         p4_fine_rangetest: Optional[bool] = None,
+        p4_lane_binsum: bool = True,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
@@ -566,6 +567,12 @@ def __init__(
         # the extra live registers push the fp32 kernel past its budget
         # (0.95-0.99x on 4 of 5 cells).
         self.p4_fine_rangetest = True if p4_fine_rangetest is None else bool(p4_fine_rangetest)
+        # p4_lane_binsum: the per-warp bin sums that open both bin
+        # searches indexed by warp only, so all 32 lanes walked the same
+        # bins_per_warp bins - a dependency chain as long as the segment,
+        # with the whole warp doing identical work. Spread the segment
+        # across the lanes and close with one warp reduction instead.
+        self.p4_lane_binsum = bool(p4_lane_binsum)
         self.r0_qfracs = tuple(float(q) for q in r0_qfracs) if r0_qfracs else ()
         if self.r0_qfracs:
             assert all(0.0 < q < 1.0 for q in self.r0_qfracs), self.r0_qfracs
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 23285c5a333f..7910c5b3b7fe 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -68,6 +68,7 @@ def _compile(
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
     p4_fine_rangetest: "bool | None" = None,
+    p4_lane_binsum: bool = True,
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
     use_ext_cand: bool = False,
@@ -207,6 +208,7 @@ def _compile(
         enable_block_skip=enable_block_skip,
         pdl_wait_late=pdl_wait_late,
         p4_fine_rangetest=p4_fine_rangetest,
+        p4_lane_binsum=p4_lane_binsum,
         use_ext_counts=use_ext_counts,
         emit_xstate=emit_xstate,
         use_ext_cand=use_ext_cand,
@@ -670,6 +672,7 @@ def gvr_topk_decode(
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
     p4_fine_rangetest: Optional[bool] = None,
+    p4_lane_binsum: bool = True,
     block_max: Optional[torch.Tensor] = None,
     skip_min_n: Optional[int] = 200_000,
     seed_thr: Optional[torch.Tensor] = None,
@@ -958,6 +961,7 @@ def gvr_topk_decode(
         enable_block_skip,
         pdl_wait_late,
         p4_fine_rangetest,
+        p4_lane_binsum,
         use_ext_counts,
         emit_xstate,
         use_ext_cand,

From bdeabeb6a6fb51137a1e4f84d731f242d9e67650 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 13:58:13 -0700
Subject: [PATCH 063/117] [None][perf] GVR P4: classify the scatter by value
 range too

The scatter walks every candidate to split it into above / inside /
below the straddling bin, and recomputed the candidate's bin index to
decide - the third O(candidates) loop in Phase 4 paying a subtract, a
multiply and two clamps for a question the value-range compare already
answers. The bin-b* edge constants now live above the fine recursion so
both loops share them. bin_i is only ever compared against b*, so the
class is encoded as b*-1 / b* / b*+1 and every branch below stays as
it was.

Paired cold graph-replay A/B (48 pairs per cell, K=512 cr=4,
production V4-Flash shape; arms alternate inside each rep and the
per-rep difference is taken, so slow drift cancels):
  B2 x 8k    -0.16us  22/48 wins  (neutral)
  B8 x 8k    +1.28us  36/48 wins
  B64 x 8k    0.00us  23/48 wins  (neutral)
  B2 x 64k   +1.44us  43/48 wins
  B8 x 64k   +0.03us  24/48 wins  (neutral)

With the two earlier cuts (fine-recursion filter, lane-parallel bin
sums) the three together are worth ~4.8us on B8 x 8k and ~4.4us on
B2 x 64k. All three share one root cause: Phase 4's three
O(candidates) loops each re-derived a bin index for a question that is
a comparison against the straddling bin's value range.

Exactness: full top-k unit file 708 passed + 1 xfailed, unified smoke
30/30, ext tiers EXACT, CUDA-graph replay 11/11, bare-op identical.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 47 +++++++++++++++----
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  4 ++
 2 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 938f72a4b587..ecfc165bad48 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -280,6 +280,7 @@ def __init__(
         pdl_wait_late: bool = True,
         p4_fine_rangetest: Optional[bool] = None,
         p4_lane_binsum: bool = True,
+        p4_scat_rangetest: bool = True,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
@@ -573,6 +574,12 @@ def __init__(
         # with the whole warp doing identical work. Spread the segment
         # across the lanes and close with one warp reduction instead.
         self.p4_lane_binsum = bool(p4_lane_binsum)
+        # p4_scat_rangetest: the scatter classifies each candidate as
+        # above / inside / below the straddling bin and recomputed the
+        # bin index to do it; the same value-range compare the fine
+        # recursion uses answers it without the subtract, multiply and
+        # two clamps per candidate.
+        self.p4_scat_rangetest = bool(p4_scat_rangetest)
         self.r0_qfracs = tuple(float(q) for q in r0_qfracs) if r0_qfracs else ()
         if self.r0_qfracs:
             assert all(0.0 < q < 1.0 for q in self.r0_qfracs), self.r0_qfracs
@@ -3663,6 +3670,12 @@ def phase4_rank_scatter(
                 # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1)
                 f_lo = bmin_r + cutlass.Float32(b_star) / inv1
                 finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1
+                # bin b* spans [f_lo, f_hi) under the coarse binning; the
+                # clamped ends fold out-of-range values INTO bin 0 and bin
+                # kBins-1, so those two bins drop the matching side.
+                f_hi = f_lo + cutlass.Float32(1.0) / inv1
+                lo_edge = b_star == cutlass.Int32(0)
+                hi_edge = b_star == cutlass.Int32(kBins - 1)
                 # re-zero (only fbins slots) + build fine sub-hist of bin-b* cands
                 iz = tidx
                 while iz < cutlass.Int32(fbins):
@@ -3677,9 +3690,6 @@ def phase4_rank_scatter(
                     # the binning fold out-of-range values INTO bin 0 and
                     # bin kBins-1, so those two bins drop the matching side
                     # of the range test to stay bit-identical.
-                    f_hi = f_lo + cutlass.Float32(1.0) / inv1
-                    lo_edge = b_star == cutlass.Int32(0)
-                    hi_edge = b_star == cutlass.Int32(kBins - 1)
                     ifb = tidx
                     while ifb < cand_count:
                         vf = smem_keys[ifb]
@@ -3782,11 +3792,32 @@ def phase4_rank_scatter(
                 isc = tidx
                 while isc < cand_count:
                     v = smem_keys[isc]
-                    bin_i = cutlass.Int32((v - bmin_r) * inv1)
-                    if bin_i < cutlass.Int32(0):
-                        bin_i = cutlass.Int32(0)
-                    if bin_i > cutlass.Int32(kBins - 1):
-                        bin_i = cutlass.Int32(kBins - 1)
+                    if cutlass.const_expr(self.p4_scat_rangetest):
+                        # same three-way split as the bin recompute, by value:
+                        # above b* <=> v >= f_hi (impossible at the top bin,
+                        # which absorbs everything higher), inside b* <=> v in
+                        # [f_lo, f_hi) with the edge bins dropping their
+                        # absorbed side. bin_i is only ever compared against
+                        # b*, so encoding the class as b*-1 / b* / b*+1 keeps
+                        # the branches below unchanged.
+                        abv = v >= f_hi
+                        inb2 = v >= f_lo and v < f_hi
+                        if lo_edge:
+                            inb2 = v < f_hi
+                        if hi_edge:
+                            abv = False
+                            inb2 = v >= f_lo
+                        bin_i = b_star - cutlass.Int32(1)
+                        if abv:
+                            bin_i = b_star + cutlass.Int32(1)
+                        if inb2:
+                            bin_i = b_star
+                    else:
+                        bin_i = cutlass.Int32((v - bmin_r) * inv1)
+                        if bin_i < cutlass.Int32(0):
+                            bin_i = cutlass.Int32(0)
+                        if bin_i > cutlass.Int32(kBins - 1):
+                            bin_i = cutlass.Int32(kBins - 1)
                     if bin_i > b_star:
                         pos = atomicAdd(s_iscalars.iterator + cutlass.Int32(4), cutlass.Int32(1))
                         if pos < cutlass.Int32(kK):
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 7910c5b3b7fe..0420bb2a926f 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -69,6 +69,7 @@ def _compile(
     pdl_wait_late: bool = True,
     p4_fine_rangetest: "bool | None" = None,
     p4_lane_binsum: bool = True,
+    p4_scat_rangetest: bool = True,
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
     use_ext_cand: bool = False,
@@ -209,6 +210,7 @@ def _compile(
         pdl_wait_late=pdl_wait_late,
         p4_fine_rangetest=p4_fine_rangetest,
         p4_lane_binsum=p4_lane_binsum,
+        p4_scat_rangetest=p4_scat_rangetest,
         use_ext_counts=use_ext_counts,
         emit_xstate=emit_xstate,
         use_ext_cand=use_ext_cand,
@@ -673,6 +675,7 @@ def gvr_topk_decode(
     pdl_wait_late: bool = True,
     p4_fine_rangetest: Optional[bool] = None,
     p4_lane_binsum: bool = True,
+    p4_scat_rangetest: bool = True,
     block_max: Optional[torch.Tensor] = None,
     skip_min_n: Optional[int] = 200_000,
     seed_thr: Optional[torch.Tensor] = None,
@@ -962,6 +965,7 @@ def gvr_topk_decode(
         pdl_wait_late,
         p4_fine_rangetest,
         p4_lane_binsum,
+        p4_scat_rangetest,
         use_ext_counts,
         emit_xstate,
         use_ext_cand,

From bba4dd41311b39cbdcd1f7284b6171a7cdda6b0c Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 20:52:30 -0700
Subject: [PATCH 064/117] [None][fix] disable the two P4 range-test fast paths
 (not fp32-exact)

p4_fine_rangetest and p4_scat_rangetest replaced a candidate's bin
recompute, floor((v - bmin) * inv1) == b_star, with a value-range
compare against the bin's [lo, hi) edges. The two agree in exact
arithmetic but not in fp32, and each was applied to only one of the two
passes that must classify a candidate identically: the fine histogram
and the scatter. A candidate the passes disagree on is counted by one
and placed by the other.

Both failure directions were observed on a captured 256k decode row
(N=65537, K=1024, one tie at the K-th value):

  - fine only:  one output slot left unwritten (1023 of 1024 filled,
                last slot -1)
  - scat only:  the scatter writes past the rank it reserved, so
                out-of-range indices reach the output; a downstream
                gather on those indices trips a device-side assert

Default both off, restoring the bit-exact bin recompute. The captured
row is exact on 8/8 replays with block skip on and off. p4_lane_binsum
is unaffected - it changes only how the per-warp bin sums are reduced,
not how a candidate is classified - and stays on.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 21 ++++++++++++-------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  4 ++--
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index ecfc165bad48..5fb2727a74e7 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -280,7 +280,7 @@ def __init__(
         pdl_wait_late: bool = True,
         p4_fine_rangetest: Optional[bool] = None,
         p4_lane_binsum: bool = True,
-        p4_scat_rangetest: bool = True,
+        p4_scat_rangetest: bool = False,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
         kc_override: "int | None" = None,
@@ -562,12 +562,15 @@ def __init__(
         # p4_fine_rangetest: the fine recursion exists only to locate
         # the handful of candidates inside the straddling coarse bin, so
         # filter them with a value-range compare instead of recomputing
-        # each candidate's bin (subtract + multiply + two clamps). It is
-        # 44% of Phase 4 on the small-N cells. Caching the candidates in
-        # registers during the coarse build was measured and rejected:
-        # the extra live registers push the fp32 kernel past its budget
-        # (0.95-0.99x on 4 of 5 cells).
-        self.p4_fine_rangetest = True if p4_fine_rangetest is None else bool(p4_fine_rangetest)
+        # each candidate's bin (subtract + multiply + two clamps).
+        # DEFAULT OFF - not bit-equivalent to the bin recompute it
+        # replaces. v in [f_lo, f_hi) and floor((v - bmin) * inv1) ==
+        # b_star agree in exact arithmetic but not in fp32, and the
+        # scatter below still classifies by bin recompute. A candidate
+        # the two passes disagree on is counted by one and placed by the
+        # other, leaving one output slot unwritten (observed on a real
+        # 256k chain row: 1023 of 1024 filled, last slot -1).
+        self.p4_fine_rangetest = False if p4_fine_rangetest is None else bool(p4_fine_rangetest)
         # p4_lane_binsum: the per-warp bin sums that open both bin
         # searches indexed by warp only, so all 32 lanes walked the same
         # bins_per_warp bins - a dependency chain as long as the segment,
@@ -579,6 +582,10 @@ def __init__(
         # bin index to do it; the same value-range compare the fine
         # recursion uses answers it without the subtract, multiply and
         # two clamps per candidate.
+        # DEFAULT OFF - same fp32 non-equivalence as p4_fine_rangetest,
+        # in the opposite direction: the range test admits a candidate
+        # the histogram binned elsewhere, so the scatter writes past the
+        # rank it reserved and out-of-range indices reach the output.
         self.p4_scat_rangetest = bool(p4_scat_rangetest)
         self.r0_qfracs = tuple(float(q) for q in r0_qfracs) if r0_qfracs else ()
         if self.r0_qfracs:
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 0420bb2a926f..773a4d214549 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -69,7 +69,7 @@ def _compile(
     pdl_wait_late: bool = True,
     p4_fine_rangetest: "bool | None" = None,
     p4_lane_binsum: bool = True,
-    p4_scat_rangetest: bool = True,
+    p4_scat_rangetest: bool = False,
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
     use_ext_cand: bool = False,
@@ -675,7 +675,7 @@ def gvr_topk_decode(
     pdl_wait_late: bool = True,
     p4_fine_rangetest: Optional[bool] = None,
     p4_lane_binsum: bool = True,
-    p4_scat_rangetest: bool = True,
+    p4_scat_rangetest: bool = False,
     block_max: Optional[torch.Tensor] = None,
     skip_min_n: Optional[int] = 200_000,
     seed_thr: Optional[torch.Tensor] = None,

From 66cc0320d0e516158e386c9762edfd632f2baa85 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 22:15:48 -0700
Subject: [PATCH 065/117] [None][perf] GVR routing: stay on the stock kernel
 below 8k context

The assist tiers have a fixed-cost floor (~12us at K=1024, ~15us at
K=512) that does not shrink with N. Once the row is short enough that
the stock kernel finishes under that floor, no tier can win, and the
emission tax makes it a straight loss.

Worst-step grid over all layers and all decode steps (162 cells, one
B200, emission tax charged), grouped by selectivity n_comp/K:

  n_comp/K   cells  geomean  worst  losing
  1.0-1.5        9    0.793  0.717     9/9
  1.5-3         18    1.004  0.908    8/18
  3-6           18    1.031  0.910    2/18
  6-20          36    1.094  0.943    5/36
  20-80         36    1.174  0.941    2/36
  80+           45    1.721  0.880    2/45

The 1.0-1.5 band is V4-Pro at 4k context: 1027 candidates for K=1024,
a top-k that selects nearly everything. Every cell there loses.

plan_emission returns "none" below n_comp 4096 (~8k raw context at
compress_ratio 4), including on the zero-emission rungs path - it is
the same kernel, so it has the same floor - and the dsa call site
falls through to the stock GVR branch. The ext state is left untouched
below the gate, which is the cold-start case its closed loop already
handles: xstate stays invalid, so the first step above the gate seeds
from the stock path.

Cost is flash 4k/8k at geomean 1.03, inside the +-3% run-to-run band.
Grid moves from geomean 1.218 / worst 0.717 / 28 losing cells to
1.231 / 0.880 / 10.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  9 ++++++-
 .../blackwell/top_k/gvr_routing.py            | 26 +++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 45281f2cb6d4..b903659e411e 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2919,8 +2919,15 @@ def sparse_attn_indexer(
                 # so we cap it at 256 for now and fall back to the CUDA C++
                 # indexer_topk_decode. This limit can be removed if GPU memory
                 # is not a bottleneck.
+                # tier "none" = the router judged this shape too short for
+                # any assist to pay (gvr_routing.ASSIST_MIN_N_COMP); fall
+                # through to the stock GVR branch below. The ext state is
+                # left untouched, which is exactly the cold-start case its
+                # closed loop already handles: xstate stays invalid, so the
+                # first step above the gate seeds from the stock path.
                 if (self.use_gvr_ext and self._gvr_ext is not None
-                        and self._gvr_route is not None and next_n == 1
+                        and self._gvr_route is not None
+                        and self._gvr_route.tier != "none" and next_n == 1
                         and num_gen_tokens <= 256):
                     # emission-assisted GVR: consume what the indexer
                     # epilogue emitted this step (packed seed row /
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index a3bd54ca6256..588203b4202b 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -42,6 +42,27 @@
 
 # ---- measured thresholds (B200, f15 grid) --------------------------------
 
+# Below this compressed length no assist tier can pay for itself. The
+# kernel has a fixed-cost floor (~12us at K=1024, ~15us at K=512) that
+# does not shrink with N, and by n_comp/K ~ 1 the stock kernel finishes
+# under that floor - there is nothing left to win.
+#
+# f17 worst-step grid (2026-07-30, 162 cells, one B200, emission tax
+# charged) grouped by selectivity n_comp/K:
+#
+#   n_comp/K   cells  geomean  worst  losing
+#   1.0-1.5        9    0.793  0.717     9/9   <- pro 4k: 1027 candidates
+#   1.5-3         18    1.004  0.908    8/18      for K=1024, a top-k that
+#   3-6           18    1.031  0.910    2/18      selects nearly everything
+#   6-20          36    1.094  0.943    5/36
+#   20-80         36    1.174  0.941    2/36
+#   80+           45    1.721  0.880    2/45
+#
+# Gating here costs flash 4k/8k (geomean 1.03, inside the +-3%
+# run-to-run band) and takes the grid from geomean 1.218 / worst 0.717
+# / 28 losing cells to 1.231 / 0.880 / 10.
+ASSIST_MIN_N_COMP = 4096  # ~8k raw context at compress_ratio 4
+
 # Block-skip prefix pays only when whole-row reads dominate.
 SKIP_MIN_N_COUNTS = 65536  # va: attach block_max unconditionally here up
 SKIP_MIN_N_RUNGS_FLASH = 131072  # vb (flash): bm pays from here
@@ -84,6 +105,11 @@ def plan_emission(batch: int, n_comp: int, k: int, have_epilogue: bool) -> str:
     top-k kernel's N. Returns the tier name; the epilogue emits the
     matching buffers and the next top-k launch routes on them.
     """
+    if n_comp < ASSIST_MIN_N_COMP:
+        # short rows: the stock kernel is already under our fixed cost,
+        # and this holds for the zero-emission rungs tier too - it is
+        # the same kernel, so the floor is the same
+        return "none"
     if not have_epilogue:
         return "rungs"  # closed-loop lines cost nothing to carry
     if batch <= LIST_EMIT_MAX_B and n_comp >= LIST_EMIT_MIN_N:

From d2ab3c1c3e0b8ac23150f75efd5550fa2ff5759a Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 22:57:21 -0700
Subject: [PATCH 066/117] [None][fix] GVR: keep the stock kernel path on
 upstream behaviour

The emission work landed several changes that were not behind a
compile-time gate, so a deployment that never enables the ext path
still ran them. Measured on the worst-step protocol, our stock path
against upstream over 36 short-context cells: geomean 0.955, worst
0.825. Split by K it is entirely the P4 tail repair - flash (K=512,
the fast path is off) is 1.017, pro (K=1024, on) is 0.896 with 17 of
18 cells below parity.

  - p4_tail_v3 (new, default off): the compacted-class repair now sits
    behind a const_expr and upstream's thread0 serial select is the
    default branch again. Two colliding locals in the restored copy are
    suffixed so the DSL's function-scope analysis cannot mix the two
    siblings.
  - pdl_wait_late defaults to False: the entry-point wait is what
    upstream emits. (Measured separately: this knob was not the source
    of the regression, 0.955 -> 0.954.)
  - the ext min/max fast path moves inside its const_expr instead of
    being selected by a runtime compare, so the stock build never
    traces it.
  - fp4 indexer __call__: block_max / hit_stats / hit_bitmap move after
    stream and take defaults, restoring the upstream positional
    signature for callers that pass ten arguments.
  - p4_lane_binsum is removed. It gated nothing: the lane-parallel bin
    sum it described is upstream's own code, so the knob was dead and
    the A/B that "validated" it compared identical machine code.

Deliberately NOT gated: the P1r degenerate-bracket rescue. Upstream
returns identity [0, K) when prev_topk is degenerate, which is wrong on
real data (42 of 231 rows on a V4-Flash TP2 capture, every request's
first decode step). Gating it would put the defect back.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  10 +-
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  |   8 +-
 .../blackwell/top_k/gvr_routing.py            |  22 +-
 .../blackwell/top_k/gvr_topk_decode.py        | 807 ++++++++++++------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |   8 +-
 5 files changed, 549 insertions(+), 306 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index b903659e411e..00953c7dba6e 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2919,12 +2919,10 @@ def sparse_attn_indexer(
                 # so we cap it at 256 for now and fall back to the CUDA C++
                 # indexer_topk_decode. This limit can be removed if GPU memory
                 # is not a bottleneck.
-                # tier "none" = the router judged this shape too short for
-                # any assist to pay (gvr_routing.ASSIST_MIN_N_COMP); fall
-                # through to the stock GVR branch below. The ext state is
-                # left untouched, which is exactly the cold-start case its
-                # closed loop already handles: xstate stays invalid, so the
-                # first step above the gate seeds from the stock path.
+                # tier "none" (gvr_routing.ASSIST_MIN_N_COMP): too short for
+                # any assist to pay - fall through to the stock branch. The
+                # untouched ext state reads as cold start, which its closed
+                # loop already handles.
                 if (self.use_gvr_ext and self._gvr_ext is not None
                         and self._gvr_route is not None
                         and self._gvr_route.tier != "none" and next_n == 1
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index 2481a5c5b8b3..dd587afc330c 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -818,10 +818,12 @@ def __call__(
         schedule_meta: cute.Tensor,  # [num_sms+1, 2] int32
         num_phys_blocks: cutlass.Int32,
         batch_size: cutlass.Int32,
-        block_max: cute.Tensor,  # [num_rows, nb_pad*4] fp32 warp-partials (or None)
-        hit_stats: cute.Tensor,  # [num_rows, 4] fp32 hit aggregate (or None unless emit_hit_stats)
-        hit_bitmap: cute.Tensor,  # [batch, nb_pad*4] int32 (or None unless emit_block_meta)
         stream: cuda.CUstream,
+        # everything below is emission-only and defaulted, so the
+        # positional signature stays the one callers already use
+        block_max: cute.Tensor = None,  # [num_rows, nb_pad*4] fp32 warp-partials
+        hit_stats: cute.Tensor = None,  # [num_rows, 4] fp32 (emit_hit_stats)
+        hit_bitmap: cute.Tensor = None,  # [batch, nb_pad*4] int32 (emit_block_meta)
         seed_thr: cute.Tensor = None,  # [num_rows, 3] fp32 (emit_seed_counts)
         seed_counts: cute.Tensor = None,  # [num_rows, 3] int32 out, caller-zeroed
         cand: cute.Tensor = None,  # [num_rows, CAP*2] int32 {val bits, idx} pairs
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index 588203b4202b..40453abe330f 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -42,25 +42,9 @@
 
 # ---- measured thresholds (B200, f15 grid) --------------------------------
 
-# Below this compressed length no assist tier can pay for itself. The
-# kernel has a fixed-cost floor (~12us at K=1024, ~15us at K=512) that
-# does not shrink with N, and by n_comp/K ~ 1 the stock kernel finishes
-# under that floor - there is nothing left to win.
-#
-# f17 worst-step grid (2026-07-30, 162 cells, one B200, emission tax
-# charged) grouped by selectivity n_comp/K:
-#
-#   n_comp/K   cells  geomean  worst  losing
-#   1.0-1.5        9    0.793  0.717     9/9   <- pro 4k: 1027 candidates
-#   1.5-3         18    1.004  0.908    8/18      for K=1024, a top-k that
-#   3-6           18    1.031  0.910    2/18      selects nearly everything
-#   6-20          36    1.094  0.943    5/36
-#   20-80         36    1.174  0.941    2/36
-#   80+           45    1.721  0.880    2/45
-#
-# Gating here costs flash 4k/8k (geomean 1.03, inside the +-3%
-# run-to-run band) and takes the grid from geomean 1.218 / worst 0.717
-# / 28 losing cells to 1.231 / 0.880 / 10.
+# The kernel's fixed cost does not shrink with N (~12us at K=1024), so
+# once the stock kernel finishes under that floor no tier can win. The
+# whole n_comp/K < 1.5 band measures below parity.
 ASSIST_MIN_N_COMP = 4096  # ~8k raw context at compress_ratio 4
 
 # Block-skip prefix pays only when whole-row reads dominate.
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 5fb2727a74e7..76aab8f6a4e0 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -277,9 +277,8 @@ def __init__(
         smem_cache_elems: int = 32768,
         seqlen_sorted: bool = False,
         kc_diet: Optional[bool] = None,
-        pdl_wait_late: bool = True,
+        pdl_wait_late: bool = False,
         p4_fine_rangetest: Optional[bool] = None,
-        p4_lane_binsum: bool = True,
         p4_scat_rangetest: bool = False,
         enable_r0: bool = True,
         accept_cap: "int | None" = None,
@@ -296,6 +295,7 @@ def __init__(
         enable_p4_rank_scatter_exact: Optional[bool] = None,
         p4_exact_tail: Optional[bool] = None,
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
+        p4_tail_v3: bool = False,
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
         enable_block_skip: bool = False,
@@ -534,12 +534,9 @@ def __init__(
         # kernel passes False for BOTH member instances so their SMEM layouts
         # stay byte-identical (the DSL sizes the launch from the last-traced
         # SmemAllocator only; see GvrTopKLBKernel).
-        # pdl_wait_late: place the PDL wait after the prologue (row
-        # resolution + SMEM allocation + config folding) instead of at
-        # kernel entry, so that work - and its cold instruction fetch -
-        # overlaps the producer's tail. Kept as a knob so both
-        # placements can be A/B'd inside one process (cross-run cold
-        # drift on shared nodes is larger than the effect).
+        # pdl_wait_late: move the PDL wait past the prologue so that work
+        # overlaps the producer's tail. Off by default: the entry-point
+        # wait is what upstream emits.
         self.pdl_wait_late = bool(pdl_wait_late)
         if kc_diet is None:
             kc_diet = cluster_size == 1
@@ -571,12 +568,6 @@ def __init__(
         # other, leaving one output slot unwritten (observed on a real
         # 256k chain row: 1023 of 1024 filled, last slot -1).
         self.p4_fine_rangetest = False if p4_fine_rangetest is None else bool(p4_fine_rangetest)
-        # p4_lane_binsum: the per-warp bin sums that open both bin
-        # searches indexed by warp only, so all 32 lanes walked the same
-        # bins_per_warp bins - a dependency chain as long as the segment,
-        # with the whole warp doing identical work. Spread the segment
-        # across the lanes and close with one warp reduction instead.
-        self.p4_lane_binsum = bool(p4_lane_binsum)
         # p4_scat_rangetest: the scatter classifies each candidate as
         # above / inside / below the straddling bin and recomputed the
         # bin index to do it; the same value-range compare the fine
@@ -769,6 +760,11 @@ def __init__(
         if p4_tail_fast is None:  # [p4tt]
             p4_tail_fast = self.p4_exact_tail and top_k >= 1024
         self.p4_tail_fast = bool(p4_tail_fast) and self.p4_exact_tail  # [p4tt]
+        # p4_tail_v3: compacted-class repair (block-parallel radix +
+        # pure-tie pre-check) in place of the stock thread0 serial
+        # select. Off by default - the stock body is what upstream
+        # emits, and this rewrite only pays at K>=1024.
+        self.p4_tail_v3 = bool(p4_tail_v3)
 
     # ------------------------------------------------------------------
     # SMEM slice cache loader. Streams this CTA's slice GMEM → SMEM via
@@ -3545,20 +3541,20 @@ def phase4_rank_scatter(
             use_ext_r = cutlass.Int32(0)
             if cutlass.const_expr(ext_range_flag is not None):
                 use_ext_r = ext_range_flag
-            if use_ext_r == cutlass.Int32(1):
-                # list rows: the take walk pre-zeroed the hist and staged
-                # per-warp maxima in smem_wcnt (its end barrier orders
-                # them); min := cut line by construction. The minmax
-                # scan, the zero pass and their three barriers vanish.
-                if cutlass.const_expr(ext_min is not None):
-                    bmin_r = ext_min
-                for w in cutlass.range_constexpr(self.num_warps):
-                    vmax = cutlass.Float32(
-                        llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value())
-                    )
-                    bmax_r = cute.arch.fmax(bmax_r, vmax)
-                if bmax_r <= bmin_r:
-                    bmax_r = bmin_r + cutlass.Float32(1e-6)
+                if use_ext_r == cutlass.Int32(1):
+                    # list rows: the take walk pre-zeroed the hist and staged
+                    # per-warp maxima in smem_wcnt (its end barrier orders
+                    # them); min := cut line by construction. The minmax
+                    # scan, the zero pass and their three barriers vanish.
+                    if cutlass.const_expr(ext_min is not None):
+                        bmin_r = ext_min
+                    for w in cutlass.range_constexpr(self.num_warps):
+                        vmax = cutlass.Float32(
+                            llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value())
+                        )
+                        bmax_r = cute.arch.fmax(bmax_r, vmax)
+                    if bmax_r <= bmin_r:
+                        bmax_r = bmin_r + cutlass.Float32(1e-6)
             if use_ext_r == cutlass.Int32(0):
                 # ---- block min/max over candidates ----
                 local_cmin = cutlass.Float32(self.FLT_MAX)
@@ -3889,269 +3885,532 @@ def phase4_rank_scatter(
                 # (need ~100 x class ~100). Classes beyond capc take the
                 # UNMODIFIED full-candidate radix below (verbatim copy).
                 if cutlass.const_expr(self.p4_exact_tail and self.p4_tail_fast):  # [p4tt]
-                    need0 = cutlass.Int32(kK) - rank_above_fine
-                    # [p4tt-v3] per-thread compact buffers, bounded by the
-                    # strided trip count over the candidate array
-                    nbuf7 = cutlass.const_expr((self.kC + self.num_threads - 1) // self.num_threads)
-                    rv7 = cute.make_fragment((nbuf7,), cutlass.Float32)
-                    ri7 = cute.make_fragment((nbuf7,), cutlass.Int32)
-                    if cutlass.const_expr(_P4_TAIL_DBG or _P4_SUB_DBG):
-                        if tidx == cutlass.Int32(0):
-                            s_thr[1] = cutlass.Float32(cnt_strad)
-                            s_thr[2] = cutlass.Float32(need0)
-                    fast_done = cutlass.Int32(1)
-                    if cnt_strad > need0 and need0 > cutlass.Int32(0):
-                        fast_done = cutlass.Int32(0)
-                        # [p4tt-v3] block-wide pure-tie check, ANY class
-                        # size: min/max order key over the (b*, sb*) class.
-                        # A pure-tie class needs NO repair — the scatter's
-                        # arrival fill of bit-equal values is already
-                        # value-set exact. Real fp8-lineage logits tie in
-                        # the thousands, which used to take the full radix.
-                        # Staging mirrors the head min/max (wcnt + hist
-                        # slots [0..31], both dead here; pairs live at
-                        # 260+).
-                        kmn6 = cutlass.Int32(2147483647)
-                        kmx6 = cutlass.Int32(-2147483648)
-                        it6 = tidx
-                        while it6 < cand_count:
-                            v6 = smem_keys[it6]
-                            b6 = cutlass.Int32((v6 - bmin_r) * inv1)
-                            if b6 < cutlass.Int32(0):
-                                b6 = cutlass.Int32(0)
-                            if b6 > cutlass.Int32(kBins - 1):
-                                b6 = cutlass.Int32(kBins - 1)
-                            if b6 == b_star:
-                                s6 = cutlass.Int32((v6 - f_lo) * finv)
-                                if s6 < cutlass.Int32(0):
-                                    s6 = cutlass.Int32(0)
-                                if s6 > cutlass.Int32(fbins - 1):
-                                    s6 = cutlass.Int32(fbins - 1)
-                                if s6 == sb_star:
-                                    k6 = f32_order_key(v6) ^ cutlass.Int32(-2147483648)
-                                    if k6 < kmn6:
-                                        kmn6 = k6
-                                    if k6 > kmx6:
-                                        kmx6 = k6
-                            it6 = it6 + cutlass.Int32(num_threads)
-                        kmn6 = cute.arch.warp_redux_sync(kmn6, "min")
-                        kmx6 = cute.arch.warp_redux_sync(kmx6, "max")
-                        if lane == cutlass.Int32(0):
-                            smem_wcnt[warp_id] = kmn6
-                            smem_hist[warp_id] = kmx6
-                        cute.arch.barrier()
-                        kmn7 = cutlass.Int32(2147483647)
-                        kmx7 = cutlass.Int32(-2147483648)
-                        for w8 in cutlass.range_constexpr(self.num_warps):
-                            pa8 = smem_wcnt[w8]
-                            pb8 = smem_hist[w8]
-                            if pa8 < kmn7:
-                                kmn7 = pa8
-                            if pb8 > kmx7:
-                                kmx7 = pb8
-                        if kmn7 == kmx7:
-                            fast_done = cutlass.Int32(1)
-                        if fast_done == cutlass.Int32(0):
-                            # [p4tt-v3] mixed class: compact it IN PLACE
-                            # into smem_keys/vals[0..cnt_strad) with a
-                            # register-buffered two-phase pass (every
-                            # thread reads its strided candidates first,
-                            # ONE barrier, then claimed compact writes —
-                            # no read/write overlap by construction). The
-                            # candidate array has no readers after the
-                            # tail, and compaction makes the repair cost a
-                            # function of the CLASS size only, for ANY
-                            # class size up to cand_count (the old full-
-                            # candidate radix fallback is gone).
+                    if cutlass.const_expr(self.p4_tail_v3):
+                        need0 = cutlass.Int32(kK) - rank_above_fine
+                        # [p4tt-v3] per-thread compact buffers, bounded by the
+                        # strided trip count over the candidate array
+                        nbuf7 = cutlass.const_expr(
+                            (self.kC + self.num_threads - 1) // self.num_threads
+                        )
+                        rv7 = cute.make_fragment((nbuf7,), cutlass.Float32)
+                        ri7 = cute.make_fragment((nbuf7,), cutlass.Int32)
+                        if cutlass.const_expr(_P4_TAIL_DBG or _P4_SUB_DBG):
                             if tidx == cutlass.Int32(0):
-                                s_iscalars[0] = cutlass.Int32(0)
-                            nh7 = cutlass.Int32(0)
-                            it7 = tidx
-                            while it7 < cand_count:
-                                v7 = smem_keys[it7]
-                                b7 = cutlass.Int32((v7 - bmin_r) * inv1)
-                                if b7 < cutlass.Int32(0):
-                                    b7 = cutlass.Int32(0)
-                                if b7 > cutlass.Int32(kBins - 1):
-                                    b7 = cutlass.Int32(kBins - 1)
-                                if b7 == b_star:
-                                    s7 = cutlass.Int32((v7 - f_lo) * finv)
-                                    if s7 < cutlass.Int32(0):
-                                        s7 = cutlass.Int32(0)
-                                    if s7 > cutlass.Int32(fbins - 1):
-                                        s7 = cutlass.Int32(fbins - 1)
-                                    if s7 == sb_star:
-                                        # static predicated fragment write
-                                        # (dodges dynamic register indexing)
-                                        for sl7 in cutlass.range_constexpr(nbuf7):
-                                            if cutlass.Int32(sl7) == nh7:
-                                                rv7[sl7] = v7
-                                                ri7[sl7] = smem_vals[it7]
-                                        nh7 = nh7 + cutlass.Int32(1)
-                                it7 = it7 + cutlass.Int32(num_threads)
-                            cute.arch.barrier()
-                            # warp-aggregated claim: intra-warp exclusive
-                            # prefix via shfl scan + ONE atomic per warp
-                            # (a thousand same-address claims serialize
-                            # and scale with the class size)
-                            pf7 = nh7
-                            for so3 in cutlass.range_constexpr(5):
-                                oth3 = cute.arch.shuffle_sync_up(
-                                    pf7, cutlass.Int32(1 << so3), mask_and_clamp=0
-                                )
-                                if lane >= cutlass.Int32(1 << so3):
-                                    pf7 = pf7 + oth3
-                            tot7 = cute.arch.shuffle_sync(pf7, cutlass.Int32(31))
-                            wb7 = cutlass.Int32(0)
-                            if lane == cutlass.Int32(31):
-                                if tot7 > cutlass.Int32(0):
-                                    wb7 = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), tot7)
-                            wb7 = cute.arch.shuffle_sync(wb7, cutlass.Int32(31))
-                            bs7 = wb7 + pf7 - nh7
-                            for sl8 in cutlass.range_constexpr(nbuf7):
-                                if cutlass.Int32(sl8) < nh7:
-                                    smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8]
-                                    smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8]
+                                s_thr[1] = cutlass.Float32(cnt_strad)
+                                s_thr[2] = cutlass.Float32(need0)
+                        fast_done = cutlass.Int32(1)
+                        if cnt_strad > need0 and need0 > cutlass.Int32(0):
+                            fast_done = cutlass.Int32(0)
+                            # [p4tt-v3] block-wide pure-tie check, ANY class
+                            # size: min/max order key over the (b*, sb*) class.
+                            # A pure-tie class needs NO repair — the scatter's
+                            # arrival fill of bit-equal values is already
+                            # value-set exact. Real fp8-lineage logits tie in
+                            # the thousands, which used to take the full radix.
+                            # Staging mirrors the head min/max (wcnt + hist
+                            # slots [0..31], both dead here; pairs live at
+                            # 260+).
+                            kmn6 = cutlass.Int32(2147483647)
+                            kmx6 = cutlass.Int32(-2147483648)
+                            it6 = tidx
+                            while it6 < cand_count:
+                                v6 = smem_keys[it6]
+                                b6 = cutlass.Int32((v6 - bmin_r) * inv1)
+                                if b6 < cutlass.Int32(0):
+                                    b6 = cutlass.Int32(0)
+                                if b6 > cutlass.Int32(kBins - 1):
+                                    b6 = cutlass.Int32(kBins - 1)
+                                if b6 == b_star:
+                                    s6 = cutlass.Int32((v6 - f_lo) * finv)
+                                    if s6 < cutlass.Int32(0):
+                                        s6 = cutlass.Int32(0)
+                                    if s6 > cutlass.Int32(fbins - 1):
+                                        s6 = cutlass.Int32(fbins - 1)
+                                    if s6 == sb_star:
+                                        k6 = f32_order_key(v6) ^ cutlass.Int32(-2147483648)
+                                        if k6 < kmn6:
+                                            kmn6 = k6
+                                        if k6 > kmx6:
+                                            kmx6 = k6
+                                it6 = it6 + cutlass.Int32(num_threads)
+                            kmn6 = cute.arch.warp_redux_sync(kmn6, "min")
+                            kmx6 = cute.arch.warp_redux_sync(kmx6, "max")
+                            if lane == cutlass.Int32(0):
+                                smem_wcnt[warp_id] = kmn6
+                                smem_hist[warp_id] = kmx6
                             cute.arch.barrier()
-                            if cnt_strad <= cutlass.Int32(128):
-                                # warp0 exact pairwise rank (rank = #{key
-                                # greater} + #{key equal, earlier slot} is
-                                # unique in [0, class)) rewrites every
-                                # winner slot in [raf, raf + need0) once.
-                                if warp_id == cutlass.Int32(0):
-                                    ie5 = lane
-                                    while ie5 < cnt_strad:
-                                        vi5 = smem_keys[ie5]
-                                        ki5 = f32_order_key(vi5) ^ cutlass.Int32(-2147483648)
-                                        r5 = cutlass.Int32(0)
-                                        j5 = cutlass.Int32(0)
-                                        while j5 < cnt_strad:
-                                            vj5 = smem_keys[j5]
-                                            kj5 = f32_order_key(vj5) ^ cutlass.Int32(-2147483648)
-                                            if kj5 > ki5:
-                                                r5 = r5 + cutlass.Int32(1)
-                                            elif kj5 == ki5 and j5 < ie5:
-                                                r5 = r5 + cutlass.Int32(1)
-                                            j5 = j5 + cutlass.Int32(1)
-                                        if r5 < need0:
-                                            pos = rank_above_fine + r5
+                            kmn7 = cutlass.Int32(2147483647)
+                            kmx7 = cutlass.Int32(-2147483648)
+                            for w8 in cutlass.range_constexpr(self.num_warps):
+                                pa8 = smem_wcnt[w8]
+                                pb8 = smem_hist[w8]
+                                if pa8 < kmn7:
+                                    kmn7 = pa8
+                                if pb8 > kmx7:
+                                    kmx7 = pb8
+                            if kmn7 == kmx7:
+                                fast_done = cutlass.Int32(1)
+                            if fast_done == cutlass.Int32(0):
+                                # [p4tt-v3] mixed class: compact it IN PLACE
+                                # into smem_keys/vals[0..cnt_strad) with a
+                                # register-buffered two-phase pass (every
+                                # thread reads its strided candidates first,
+                                # ONE barrier, then claimed compact writes —
+                                # no read/write overlap by construction). The
+                                # candidate array has no readers after the
+                                # tail, and compaction makes the repair cost a
+                                # function of the CLASS size only, for ANY
+                                # class size up to cand_count (the old full-
+                                # candidate radix fallback is gone).
+                                if tidx == cutlass.Int32(0):
+                                    s_iscalars[0] = cutlass.Int32(0)
+                                nh7 = cutlass.Int32(0)
+                                it7 = tidx
+                                while it7 < cand_count:
+                                    v7 = smem_keys[it7]
+                                    b7 = cutlass.Int32((v7 - bmin_r) * inv1)
+                                    if b7 < cutlass.Int32(0):
+                                        b7 = cutlass.Int32(0)
+                                    if b7 > cutlass.Int32(kBins - 1):
+                                        b7 = cutlass.Int32(kBins - 1)
+                                    if b7 == b_star:
+                                        s7 = cutlass.Int32((v7 - f_lo) * finv)
+                                        if s7 < cutlass.Int32(0):
+                                            s7 = cutlass.Int32(0)
+                                        if s7 > cutlass.Int32(fbins - 1):
+                                            s7 = cutlass.Int32(fbins - 1)
+                                        if s7 == sb_star:
+                                            # static predicated fragment write
+                                            # (dodges dynamic register indexing)
+                                            for sl7 in cutlass.range_constexpr(nbuf7):
+                                                if cutlass.Int32(sl7) == nh7:
+                                                    rv7[sl7] = v7
+                                                    ri7[sl7] = smem_vals[it7]
+                                            nh7 = nh7 + cutlass.Int32(1)
+                                    it7 = it7 + cutlass.Int32(num_threads)
+                                cute.arch.barrier()
+                                # warp-aggregated claim: intra-warp exclusive
+                                # prefix via shfl scan + ONE atomic per warp
+                                # (a thousand same-address claims serialize
+                                # and scale with the class size)
+                                pf7 = nh7
+                                for so3 in cutlass.range_constexpr(5):
+                                    oth3 = cute.arch.shuffle_sync_up(
+                                        pf7, cutlass.Int32(1 << so3), mask_and_clamp=0
+                                    )
+                                    if lane >= cutlass.Int32(1 << so3):
+                                        pf7 = pf7 + oth3
+                                tot7 = cute.arch.shuffle_sync(pf7, cutlass.Int32(31))
+                                wb7 = cutlass.Int32(0)
+                                if lane == cutlass.Int32(31):
+                                    if tot7 > cutlass.Int32(0):
+                                        wb7 = atomicAdd(
+                                            s_iscalars.iterator + cutlass.Int32(0), tot7
+                                        )
+                                wb7 = cute.arch.shuffle_sync(wb7, cutlass.Int32(31))
+                                bs7 = wb7 + pf7 - nh7
+                                for sl8 in cutlass.range_constexpr(nbuf7):
+                                    if cutlass.Int32(sl8) < nh7:
+                                        smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8]
+                                        smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8]
+                                cute.arch.barrier()
+                                if cnt_strad <= cutlass.Int32(128):
+                                    # warp0 exact pairwise rank (rank = #{key
+                                    # greater} + #{key equal, earlier slot} is
+                                    # unique in [0, class)) rewrites every
+                                    # winner slot in [raf, raf + need0) once.
+                                    if warp_id == cutlass.Int32(0):
+                                        ie5 = lane
+                                        while ie5 < cnt_strad:
+                                            vi5 = smem_keys[ie5]
+                                            ki5 = f32_order_key(vi5) ^ cutlass.Int32(-2147483648)
+                                            r5 = cutlass.Int32(0)
+                                            j5 = cutlass.Int32(0)
+                                            while j5 < cnt_strad:
+                                                vj5 = smem_keys[j5]
+                                                kj5 = f32_order_key(vj5) ^ cutlass.Int32(
+                                                    -2147483648
+                                                )
+                                                if kj5 > ki5:
+                                                    r5 = r5 + cutlass.Int32(1)
+                                                elif kj5 == ki5 and j5 < ie5:
+                                                    r5 = r5 + cutlass.Int32(1)
+                                                j5 = j5 + cutlass.Int32(1)
+                                            if r5 < need0:
+                                                pos = rank_above_fine + r5
+                                                if pos < cutlass.Int32(kK):
+                                                    if cutlass.const_expr(
+                                                        self.return_output_values
+                                                    ):
+                                                        output_values_row[pos] = self.dtype(vi5)
+                                                    output_indices_row[pos] = smem_vals[ie5]
+                                            ie5 = ie5 + cutlass.Int32(32)
+                                    cute.arch.barrier()
+                                else:
+                                    # block-parallel 4-level MSB radix over the
+                                    # compacted class (scans touch class pairs
+                                    # only; warp0 shuffle-scan digit search — 3
+                                    # block barriers per level instead of 5).
+                                    if tidx == cutlass.Int32(0):
+                                        smem_hist[256] = cutlass.Int32(0)
+                                        smem_hist[257] = need0
+                                        smem_hist[258] = cutlass.Int32(0)
+                                    cute.arch.barrier()
+                                    for lvl2 in cutlass.range_constexpr(4):
+                                        shift2 = cutlass.const_expr(24 - 8 * lvl2)
+                                        iz3 = tidx
+                                        while iz3 < cutlass.Int32(256):
+                                            smem_hist[iz3] = cutlass.Int32(0)
+                                            iz3 = iz3 + cutlass.Int32(num_threads)
+                                        cute.arch.barrier()
+                                        uthr_c2 = smem_hist[256]
+                                        ic2 = tidx
+                                        while ic2 < cnt_strad:
+                                            uk3 = f32_order_key(smem_keys[ic2])
+                                            pm2 = cutlass.Int32(1)
+                                            if cutlass.const_expr(lvl2 > 0):
+                                                if (uk3 >> cutlass.Int32(shift2 + 8)) != (
+                                                    uthr_c2 >> cutlass.Int32(shift2 + 8)
+                                                ):
+                                                    pm2 = cutlass.Int32(0)
+                                            if pm2 == cutlass.Int32(1):
+                                                dg2 = (
+                                                    uk3 >> cutlass.Int32(shift2)
+                                                ) & cutlass.Int32(0xFF)
+                                                atomicAdd(
+                                                    smem_hist.iterator + dg2, cutlass.Int32(1)
+                                                )
+                                            ic2 = ic2 + cutlass.Int32(num_threads)
+                                        cute.arch.barrier()
+                                        if warp_id == cutlass.Int32(0):
+                                            ws3 = cutlass.Int32(0)
+                                            for jd3 in cutlass.range_constexpr(8):
+                                                di3 = (
+                                                    cutlass.Int32(255)
+                                                    - lane * cutlass.Int32(8)
+                                                    - cutlass.Int32(jd3)
+                                                )
+                                                ws3 = ws3 + smem_hist[di3]
+                                            pre6 = ws3
+                                            for so2 in cutlass.range_constexpr(5):
+                                                oth2 = cute.arch.shuffle_sync_up(
+                                                    pre6,
+                                                    cutlass.Int32(1 << so2),
+                                                    mask_and_clamp=0,
+                                                )
+                                                if lane >= cutlass.Int32(1 << so2):
+                                                    pre6 = pre6 + oth2
+                                            needl3 = smem_hist[257]
+                                            if pre6 >= needl3 and (pre6 - ws3) < needl3:
+                                                base5 = pre6 - ws3
+                                                dstar2 = cutlass.Int32(0)
+                                                above5 = base5
+                                                sd5 = cutlass.Int32(0)
+                                                for jd4 in cutlass.range_constexpr(8):
+                                                    di4 = (
+                                                        cutlass.Int32(255)
+                                                        - lane * cutlass.Int32(8)
+                                                        - cutlass.Int32(jd4)
+                                                    )
+                                                    ra5 = base5
+                                                    base5 = base5 + smem_hist[di4]
+                                                    if base5 >= needl3 and sd5 == cutlass.Int32(0):
+                                                        dstar2 = di4
+                                                        above5 = ra5
+                                                        sd5 = cutlass.Int32(1)
+                                                smem_hist[256] = uthr_c2 | (
+                                                    dstar2 << cutlass.Int32(shift2)
+                                                )
+                                                smem_hist[257] = needl3 - above5
+                                                smem_hist[258] = smem_hist[258] + above5
+                                        cute.arch.barrier()
+                                    u_thr2 = smem_hist[256]
+                                    cnt_ab2 = smem_hist[258]
+                                    need_eq2 = smem_hist[257]
+                                    kthr2 = u_thr2 ^ cutlass.Int32(-2147483648)
+                                    if tidx == cutlass.Int32(0):
+                                        s_iscalars[4] = cutlass.Int32(0)
+                                        s_iscalars[0] = cutlass.Int32(0)
+                                    cute.arch.barrier()
+                                    ir3 = tidx
+                                    while ir3 < cnt_strad:
+                                        vv3 = smem_keys[ir3]
+                                        uk4 = f32_order_key(vv3)
+                                        ks4 = uk4 ^ cutlass.Int32(-2147483648)
+                                        if ks4 > kthr2:
+                                            o4 = atomicAdd(
+                                                s_iscalars.iterator + cutlass.Int32(4),
+                                                cutlass.Int32(1),
+                                            )
+                                            pos = rank_above_fine + o4
                                             if pos < cutlass.Int32(kK):
                                                 if cutlass.const_expr(self.return_output_values):
-                                                    output_values_row[pos] = self.dtype(vi5)
-                                                output_indices_row[pos] = smem_vals[ie5]
-                                        ie5 = ie5 + cutlass.Int32(32)
+                                                    output_values_row[pos] = self.dtype(vv3)
+                                                output_indices_row[pos] = smem_vals[ir3]
+                                        elif ks4 == kthr2:
+                                            q4 = atomicAdd(
+                                                s_iscalars.iterator + cutlass.Int32(0),
+                                                cutlass.Int32(1),
+                                            )
+                                            if q4 < need_eq2:
+                                                pos = rank_above_fine + cnt_ab2 + q4
+                                                if pos < cutlass.Int32(kK):
+                                                    if cutlass.const_expr(
+                                                        self.return_output_values
+                                                    ):
+                                                        output_values_row[pos] = self.dtype(vv3)
+                                                    output_indices_row[pos] = smem_vals[ir3]
+                                        ir3 = ir3 + cutlass.Int32(num_threads)
+                                    cute.arch.barrier()
+                    else:
+                        need0_s = cutlass.Int32(kK) - rank_above_fine
+                        if cnt_strad > need0_s and need0_s > cutlass.Int32(0):
+                            if cnt_strad <= cutlass.Int32(128):
+                                # [p4tt] SMEM: (value_bits, cand_idx) pairs at
+                                # smem_hist[2*o]/[2*o+1], o < 128 (slots 0..255).
+                                # The 256 digit bins are dead here (the fast path
+                                # replaces the radix levels that used them); the
+                                # sb_star/ra staging in slots 2/3 was read by
+                                # every thread before the pre-scatter barrier.
+                                # Persistent radix scalars [256..258] untouched.
+                                # Collect counter = s_iscalars[0] (dead after the
+                                # scatter; same reuse as the radix rewrite pass).
+                                if tidx == cutlass.Int32(0):
+                                    s_iscalars[0] = cutlass.Int32(0)
+                                cute.arch.barrier()
+                                itc = tidx
+                                while itc < cand_count:
+                                    tv = smem_keys[itc]
+                                    tb = cutlass.Int32((tv - bmin_r) * inv1)
+                                    if tb < cutlass.Int32(0):
+                                        tb = cutlass.Int32(0)
+                                    if tb > cutlass.Int32(kBins - 1):
+                                        tb = cutlass.Int32(kBins - 1)
+                                    if tb == b_star:
+                                        ts = cutlass.Int32((tv - f_lo) * finv)
+                                        if ts < cutlass.Int32(0):
+                                            ts = cutlass.Int32(0)
+                                        if ts > cutlass.Int32(fbins - 1):
+                                            ts = cutlass.Int32(fbins - 1)
+                                        if ts == sb_star:
+                                            to = atomicAdd(
+                                                s_iscalars.iterator + cutlass.Int32(0),
+                                                cutlass.Int32(1),
+                                            )
+                                            if to < cutlass.Int32(128):
+                                                smem_hist[to + to] = float_as_int32(tv)
+                                                smem_hist[to + to + cutlass.Int32(1)] = smem_vals[
+                                                    itc
+                                                ]
+                                    itc = itc + cutlass.Int32(num_threads)
+                                cute.arch.barrier()
+                                # [p4tt] thread0 exact top-need0_s select rewriting
+                                # positions [rank_above_fine, kK). Consumed flag =
+                                # the cand_idx slot set to -1 (indices are always
+                                # >= 0), so a genuine -FLT_MAX value in the class
+                                # remains selectable (no value sentinel). Ties
+                                # (bit-equal values) pick arbitrarily: value-set
+                                # exact.
+                                if tidx == cutlass.Int32(0):
+                                    tj = cutlass.Int32(0)
+                                    while tj < need0_s:
+                                        tbv = cutlass.Float32(self.NEG_FLT_MAX)
+                                        tbi = cutlass.Int32(-1)
+                                        ti = cutlass.Int32(0)
+                                        while ti < cnt_strad:
+                                            tvi = smem_hist[ti + ti + cutlass.Int32(1)]
+                                            if tvi >= cutlass.Int32(0):
+                                                tvb = smem_hist[ti + ti]
+                                                tvv = cutlass.Float32(
+                                                    llvm.bitcast(
+                                                        cutlass.Float32.mlir_type,
+                                                        tvb.ir_value(),
+                                                    )
+                                                )
+                                                take = cutlass.Int32(0)
+                                                if tbi < cutlass.Int32(0):
+                                                    take = cutlass.Int32(1)
+                                                elif tvv > tbv:
+                                                    take = cutlass.Int32(1)
+                                                if take == cutlass.Int32(1):
+                                                    tbv = tvv
+                                                    tbi = ti
+                                            ti = ti + cutlass.Int32(1)
+                                        pos_s = rank_above_fine + tj
+                                        if cutlass.const_expr(self.return_output_values):
+                                            output_values_row[pos_s] = self.dtype(tbv)
+                                        output_indices_row[pos_s] = smem_hist[
+                                            tbi + tbi + cutlass.Int32(1)
+                                        ]
+                                        smem_hist[tbi + tbi + cutlass.Int32(1)] = cutlass.Int32(-1)
+                                        tj = tj + cutlass.Int32(1)
                                 cute.arch.barrier()
                             else:
-                                # block-parallel 4-level MSB radix over the
-                                # compacted class (scans touch class pairs
-                                # only; warp0 shuffle-scan digit search — 3
-                                # block barriers per level instead of 5).
+                                # Persistent scalars live above the 256 digit bins
+                                # (kNumBins >= 512 always): [256] key prefix (chosen
+                                # digits, remaining bits 0), [257] slots still to fill
+                                # inside the current equal-prefix set, [258] ties
+                                # strictly above the prefix (their slots precede it).
                                 if tidx == cutlass.Int32(0):
                                     smem_hist[256] = cutlass.Int32(0)
-                                    smem_hist[257] = need0
+                                    smem_hist[257] = need0_s
                                     smem_hist[258] = cutlass.Int32(0)
                                 cute.arch.barrier()
-                                for lvl2 in cutlass.range_constexpr(4):
-                                    shift2 = cutlass.const_expr(24 - 8 * lvl2)
-                                    iz3 = tidx
-                                    while iz3 < cutlass.Int32(256):
-                                        smem_hist[iz3] = cutlass.Int32(0)
-                                        iz3 = iz3 + cutlass.Int32(num_threads)
+                                for lvl in cutlass.range_constexpr(4):
+                                    shift = cutlass.const_expr(24 - 8 * lvl)
+                                    iz2 = tidx
+                                    while iz2 < cutlass.Int32(256):
+                                        smem_hist[iz2] = cutlass.Int32(0)
+                                        iz2 = iz2 + cutlass.Int32(num_threads)
                                     cute.arch.barrier()
-                                    uthr_c2 = smem_hist[256]
-                                    ic2 = tidx
-                                    while ic2 < cnt_strad:
-                                        uk3 = f32_order_key(smem_keys[ic2])
-                                        pm2 = cutlass.Int32(1)
-                                        if cutlass.const_expr(lvl2 > 0):
-                                            if (uk3 >> cutlass.Int32(shift2 + 8)) != (
-                                                uthr_c2 >> cutlass.Int32(shift2 + 8)
-                                            ):
-                                                pm2 = cutlass.Int32(0)
-                                        if pm2 == cutlass.Int32(1):
-                                            dg2 = (uk3 >> cutlass.Int32(shift2)) & cutlass.Int32(
-                                                0xFF
-                                            )
-                                            atomicAdd(smem_hist.iterator + dg2, cutlass.Int32(1))
-                                        ic2 = ic2 + cutlass.Int32(num_threads)
+                                    uthr_cur = smem_hist[256]
+                                    it2 = tidx
+                                    while it2 < cand_count:
+                                        vt = smem_keys[it2]
+                                        bt = cutlass.Int32((vt - bmin_r) * inv1)
+                                        if bt < cutlass.Int32(0):
+                                            bt = cutlass.Int32(0)
+                                        if bt > cutlass.Int32(kBins - 1):
+                                            bt = cutlass.Int32(kBins - 1)
+                                        if bt == b_star:
+                                            st2 = cutlass.Int32((vt - f_lo) * finv)
+                                            if st2 < cutlass.Int32(0):
+                                                st2 = cutlass.Int32(0)
+                                            if st2 > cutlass.Int32(fbins - 1):
+                                                st2 = cutlass.Int32(fbins - 1)
+                                            if st2 == sb_star:
+                                                uk = f32_order_key(vt)
+                                                pmatch = cutlass.Int32(1)
+                                                if cutlass.const_expr(lvl > 0):
+                                                    if (uk >> cutlass.Int32(shift + 8)) != (
+                                                        uthr_cur >> cutlass.Int32(shift + 8)
+                                                    ):
+                                                        pmatch = cutlass.Int32(0)
+                                                if pmatch == cutlass.Int32(1):
+                                                    dg = (
+                                                        uk >> cutlass.Int32(shift)
+                                                    ) & cutlass.Int32(0xFF)
+                                                    atomicAdd(
+                                                        smem_hist.iterator + dg, cutlass.Int32(1)
+                                                    )
+                                        it2 = it2 + cutlass.Int32(num_threads)
                                     cute.arch.barrier()
-                                    if warp_id == cutlass.Int32(0):
-                                        ws3 = cutlass.Int32(0)
-                                        for jd3 in cutlass.range_constexpr(8):
-                                            di3 = (
+                                    # Two-stage descending digit scan (mirrors the
+                                    # fine 3-step search): per-warp partial sums,
+                                    # thread0 picks the target warp, its lane0 walks
+                                    # the warp's digit range — 2*num_warps serial
+                                    # steps instead of 256.
+                                    fdw = cutlass.const_expr(256 // self.num_warps)
+                                    wsum2 = cutlass.Int32(0)
+                                    for jd in cutlass.range_constexpr(fdw):
+                                        dix = (
+                                            cutlass.Int32(255)
+                                            - warp_id * cutlass.Int32(fdw)
+                                            - cutlass.Int32(jd)
+                                        )
+                                        wsum2 = wsum2 + smem_hist[dix]
+                                    if lane == cutlass.Int32(0):
+                                        smem_wcnt[warp_id] = wsum2
+                                    cute.arch.barrier()
+                                    if tidx == cutlass.Int32(0):
+                                        needl = smem_hist[257]
+                                        cw = cutlass.Int32(0)
+                                        tw3 = cutlass.Int32(num_warps - 1)
+                                        f3 = cutlass.Int32(0)
+                                        for w4 in cutlass.range_constexpr(self.num_warps):
+                                            cw = cw + smem_wcnt[w4]
+                                            if cw >= needl and f3 == cutlass.Int32(0):
+                                                tw3 = cutlass.Int32(w4)
+                                                f3 = cutlass.Int32(1)
+                                        pre3 = cutlass.Int32(0)
+                                        for w5 in cutlass.range_constexpr(self.num_warps):
+                                            if cutlass.Int32(w5) < tw3:
+                                                pre3 = pre3 + smem_wcnt[w5]
+                                        s_iscalars[4] = pre3  # prefix above target warp
+                                        s_iscalars[0] = tw3  # target warp
+                                    cute.arch.barrier()
+                                    pre4 = s_iscalars[4]
+                                    tw4 = s_iscalars[0]
+                                    if warp_id == tw4 and lane == cutlass.Int32(0):
+                                        needl2 = smem_hist[257]
+                                        base4 = pre4
+                                        dstar = cutlass.Int32(0)
+                                        above_d = pre4
+                                        sd4 = cutlass.Int32(0)
+                                        for jd2 in cutlass.range_constexpr(fdw):
+                                            dix2 = (
                                                 cutlass.Int32(255)
-                                                - lane * cutlass.Int32(8)
-                                                - cutlass.Int32(jd3)
+                                                - tw4 * cutlass.Int32(fdw)
+                                                - cutlass.Int32(jd2)
                                             )
-                                            ws3 = ws3 + smem_hist[di3]
-                                        pre6 = ws3
-                                        for so2 in cutlass.range_constexpr(5):
-                                            oth2 = cute.arch.shuffle_sync_up(
-                                                pre6,
-                                                cutlass.Int32(1 << so2),
-                                                mask_and_clamp=0,
-                                            )
-                                            if lane >= cutlass.Int32(1 << so2):
-                                                pre6 = pre6 + oth2
-                                        needl3 = smem_hist[257]
-                                        if pre6 >= needl3 and (pre6 - ws3) < needl3:
-                                            base5 = pre6 - ws3
-                                            dstar2 = cutlass.Int32(0)
-                                            above5 = base5
-                                            sd5 = cutlass.Int32(0)
-                                            for jd4 in cutlass.range_constexpr(8):
-                                                di4 = (
-                                                    cutlass.Int32(255)
-                                                    - lane * cutlass.Int32(8)
-                                                    - cutlass.Int32(jd4)
-                                                )
-                                                ra5 = base5
-                                                base5 = base5 + smem_hist[di4]
-                                                if base5 >= needl3 and sd5 == cutlass.Int32(0):
-                                                    dstar2 = di4
-                                                    above5 = ra5
-                                                    sd5 = cutlass.Int32(1)
-                                            smem_hist[256] = uthr_c2 | (
-                                                dstar2 << cutlass.Int32(shift2)
-                                            )
-                                            smem_hist[257] = needl3 - above5
-                                            smem_hist[258] = smem_hist[258] + above5
+                                            ra4 = base4
+                                            base4 = base4 + smem_hist[dix2]
+                                            if base4 >= needl2 and sd4 == cutlass.Int32(0):
+                                                dstar = dix2
+                                                above_d = ra4
+                                                sd4 = cutlass.Int32(1)
+                                        smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift))
+                                        smem_hist[257] = needl2 - above_d
+                                        smem_hist[258] = smem_hist[258] + above_d
                                     cute.arch.barrier()
-                                u_thr2 = smem_hist[256]
-                                cnt_ab2 = smem_hist[258]
-                                need_eq2 = smem_hist[257]
-                                kthr2 = u_thr2 ^ cutlass.Int32(-2147483648)
+                                # Rewrite the tie slot range: ties with key > u_thr
+                                # first (there are exactly cnt_ab of them), then the
+                                # first need_eq bitwise-equal-to-u_thr ties in arrival
+                                # order (value-exact by construction). Signed compare
+                                # needs the top bit flipped (unsigned-monotonic key).
+                                u_thr = smem_hist[256]
+                                cnt_ab = smem_hist[258]
+                                need_eq = smem_hist[257]
+                                ks_thr = u_thr ^ cutlass.Int32(-2147483648)
                                 if tidx == cutlass.Int32(0):
-                                    s_iscalars[4] = cutlass.Int32(0)
-                                    s_iscalars[0] = cutlass.Int32(0)
+                                    s_iscalars[4] = cutlass.Int32(0)  # above-writer ctr
+                                    s_iscalars[0] = cutlass.Int32(0)  # equal-writer ctr
                                 cute.arch.barrier()
-                                ir3 = tidx
-                                while ir3 < cnt_strad:
-                                    vv3 = smem_keys[ir3]
-                                    uk4 = f32_order_key(vv3)
-                                    ks4 = uk4 ^ cutlass.Int32(-2147483648)
-                                    if ks4 > kthr2:
-                                        o4 = atomicAdd(
-                                            s_iscalars.iterator + cutlass.Int32(4),
-                                            cutlass.Int32(1),
-                                        )
-                                        pos = rank_above_fine + o4
-                                        if pos < cutlass.Int32(kK):
-                                            if cutlass.const_expr(self.return_output_values):
-                                                output_values_row[pos] = self.dtype(vv3)
-                                            output_indices_row[pos] = smem_vals[ir3]
-                                    elif ks4 == kthr2:
-                                        q4 = atomicAdd(
-                                            s_iscalars.iterator + cutlass.Int32(0),
-                                            cutlass.Int32(1),
-                                        )
-                                        if q4 < need_eq2:
-                                            pos = rank_above_fine + cnt_ab2 + q4
-                                            if pos < cutlass.Int32(kK):
-                                                if cutlass.const_expr(self.return_output_values):
-                                                    output_values_row[pos] = self.dtype(vv3)
-                                                output_indices_row[pos] = smem_vals[ir3]
-                                    ir3 = ir3 + cutlass.Int32(num_threads)
+                                ir2 = tidx
+                                while ir2 < cand_count:
+                                    vr = smem_keys[ir2]
+                                    br = cutlass.Int32((vr - bmin_r) * inv1)
+                                    if br < cutlass.Int32(0):
+                                        br = cutlass.Int32(0)
+                                    if br > cutlass.Int32(kBins - 1):
+                                        br = cutlass.Int32(kBins - 1)
+                                    if br == b_star:
+                                        sr = cutlass.Int32((vr - f_lo) * finv)
+                                        if sr < cutlass.Int32(0):
+                                            sr = cutlass.Int32(0)
+                                        if sr > cutlass.Int32(fbins - 1):
+                                            sr = cutlass.Int32(fbins - 1)
+                                        if sr == sb_star:
+                                            uk2 = f32_order_key(vr)
+                                            ks2 = uk2 ^ cutlass.Int32(-2147483648)
+                                            if ks2 > ks_thr:
+                                                o2 = atomicAdd(
+                                                    s_iscalars.iterator + cutlass.Int32(4),
+                                                    cutlass.Int32(1),
+                                                )
+                                                pos_s = rank_above_fine + o2
+                                                if pos_s < cutlass.Int32(kK):
+                                                    if cutlass.const_expr(
+                                                        self.return_output_values
+                                                    ):
+                                                        output_values_row[pos_s] = self.dtype(vr)
+                                                    output_indices_row[pos_s] = smem_vals[ir2]
+                                            elif ks2 == ks_thr:
+                                                q2 = atomicAdd(
+                                                    s_iscalars.iterator + cutlass.Int32(0),
+                                                    cutlass.Int32(1),
+                                                )
+                                                if q2 < need_eq:
+                                                    pos_s = rank_above_fine + cnt_ab + q2
+                                                    if pos_s < cutlass.Int32(kK):
+                                                        if cutlass.const_expr(
+                                                            self.return_output_values
+                                                        ):
+                                                            output_values_row[pos_s] = self.dtype(
+                                                                vr
+                                                            )
+                                                        output_indices_row[pos_s] = smem_vals[ir2]
+                                    ir2 = ir2 + cutlass.Int32(num_threads)
                                 cute.arch.barrier()
                 elif cutlass.const_expr(self.p4_exact_tail):  # [p4tt] if->elif only
                     need0 = cutlass.Int32(kK) - rank_above_fine
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 773a4d214549..30ba825ec8ee 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -67,8 +67,8 @@ def _compile(
     p2_warp_redundant: bool = True,
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
+    p4_tail_v3: bool = False,
     p4_fine_rangetest: "bool | None" = None,
-    p4_lane_binsum: bool = True,
     p4_scat_rangetest: bool = False,
     use_ext_counts: bool = False,
     emit_xstate: bool = False,
@@ -208,8 +208,8 @@ def _compile(
         p2_warp_redundant=p2_warp_redundant,
         enable_block_skip=enable_block_skip,
         pdl_wait_late=pdl_wait_late,
+        p4_tail_v3=p4_tail_v3,
         p4_fine_rangetest=p4_fine_rangetest,
-        p4_lane_binsum=p4_lane_binsum,
         p4_scat_rangetest=p4_scat_rangetest,
         use_ext_counts=use_ext_counts,
         emit_xstate=emit_xstate,
@@ -673,8 +673,8 @@ def gvr_topk_decode(
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
+    p4_tail_v3: bool = False,
     p4_fine_rangetest: Optional[bool] = None,
-    p4_lane_binsum: bool = True,
     p4_scat_rangetest: bool = False,
     block_max: Optional[torch.Tensor] = None,
     skip_min_n: Optional[int] = 200_000,
@@ -963,8 +963,8 @@ def gvr_topk_decode(
         p2_warp_redundant,
         enable_block_skip,
         pdl_wait_late,
+        p4_tail_v3,
         p4_fine_rangetest,
-        p4_lane_binsum,
         p4_scat_rangetest,
         use_ext_counts,
         emit_xstate,

From c65dc31834c12f4151611d229b8f5f31134f9397 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 23:03:50 -0700
Subject: [PATCH 067/117] [None][chore] GVR: condense the narrative comment
 blocks

Drop the version tags and the measurement logs from the kernel
comments; keep the contract and the DSL traps.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 153 +++++++-----------
 1 file changed, 62 insertions(+), 91 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 76aab8f6a4e0..3d5d29c3e0af 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -458,17 +458,13 @@ def __init__(
         self.enable_r0 = bool(enable_r0)
         self.mt_unroll = int(mt_unroll)
         self.fb_fix = bool(fb_fix)
-        # enable_block_skip: gate the R0 M-ary count pass and the Phase-3
-        # stream-write on per-32-position upper bounds emitted by the
-        # indexer epilogue (``block_max [num_rows, nb_pad*4]`` fp32, record
-        # r = exact max of positions [r*32, r*32+32) of the POST-CONVERSION
-        # stored logits — contract in workspace/epilogue_topk_interface.md).
-        # Lossless: the active-block list is built at the LOOSEST rung, so a
-        # skipped block cannot contain any element >= any rung and every
-        # rung count (and the collect) equals its dense value. The port
-        # ships the measured-optimal configuration only: grain 32, int16
-        # list entries (16KB SMEM, 3 CTA/SM at T512), strided coalesced
-        # O(1)-barrier build, UN=2 software-pipelined compact scan.
+        # enable_block_skip: gate the R0 count pass and the Phase-3
+        # stream-write on per-32-position upper bounds from the indexer
+        # epilogue (block_max [num_rows, nb_pad*4] fp32; record r = max
+        # over positions [r*32, r*32+32) of the stored logits). Lossless:
+        # the active list is built at the loosest rung, so a skipped
+        # block holds nothing >= any rung and every count equals its
+        # dense value.
         self.enable_block_skip = bool(enable_block_skip)
         self.SKIP_BLOCK = 32
         self.SKIP_BLOCK_LOG2 = 5
@@ -556,17 +552,11 @@ def __init__(
         # measured as a wash under the same protocol and stay stock.
         if enable_r0 and top_k == 2048 and self.kNumBins > 512:
             self.kNumBins = 512
-        # p4_fine_rangetest: the fine recursion exists only to locate
-        # the handful of candidates inside the straddling coarse bin, so
-        # filter them with a value-range compare instead of recomputing
-        # each candidate's bin (subtract + multiply + two clamps).
-        # DEFAULT OFF - not bit-equivalent to the bin recompute it
-        # replaces. v in [f_lo, f_hi) and floor((v - bmin) * inv1) ==
-        # b_star agree in exact arithmetic but not in fp32, and the
-        # scatter below still classifies by bin recompute. A candidate
-        # the two passes disagree on is counted by one and placed by the
-        # other, leaving one output slot unwritten (observed on a real
-        # 256k chain row: 1023 of 1024 filled, last slot -1).
+        # p4_fine_rangetest: filter the fine recursion by value range
+        # instead of recomputing each candidate's bin. OFF - the two are
+        # not fp32-equivalent, and the scatter still classifies by bin
+        # recompute, so a candidate they disagree on is counted by one
+        # pass and placed by the other, leaving an output slot unwritten.
         self.p4_fine_rangetest = False if p4_fine_rangetest is None else bool(p4_fine_rangetest)
         # p4_scat_rangetest: the scatter classifies each candidate as
         # above / inside / below the straddling bin and recomputed the
@@ -634,18 +624,14 @@ def __init__(
         # segment that can void).
         self.list_cap = max(0, int(cand_cap) - 2 * self.accept_cap)
         self.cand_rung = int(cand_rung)
-        # self_scan (fused self-contained mode): the kernel itself streams
-        # the row ONCE against the three closed-loop lines, bucketing
-        # VALUES into on-chip segments (A >= t2 / B [t1,t2) / C [t0,t1),
-        # bases 0 / accept_cap / 2*accept_cap inside an enlarged
-        # smem_keys) and POSITIONS into the cand_idx tensor (write-only
-        # until the deferred K-gather). No external emitter, no candidate
-        # value column in gmem. A line cut compacts the winning segment
-        # runs to the smem_keys prefix and fills smem_vals with each
-        # entry's SEGMENT COORDINATE — from there on the v5 consumer
-        # (P4, tail repair, deferred gather via cand_idx[coord]) runs
-        # unchanged. Ineligible rows take the stock fallback, whose
-        # P3/P4 use smem_keys[:kC]/smem_vals verbatim.
+        # self_scan (fused self-contained mode): the kernel streams the
+        # row once against the three closed-loop lines, bucketing values
+        # into on-chip segments (A >= t2 / B [t1,t2) / C [t0,t1) at bases
+        # 0 / accept_cap / 2*accept_cap) and positions into cand_idx. A
+        # line cut compacts the winning segment to the smem_keys prefix
+        # and fills smem_vals with segment coordinates, after which the
+        # list consumer runs unchanged. Ineligible rows take the stock
+        # fallback.
         self.self_scan = bool(self_scan)
         if self.self_scan:
             if not use_ext_counts:
@@ -743,7 +729,7 @@ def __init__(
         if p4_exact_tail is None:
             p4_exact_tail = self.enable_p4_rank_scatter_exact and dtype == cutlass.Float32
         self.p4_exact_tail = bool(p4_exact_tail) and self.enable_p4_rank_scatter_exact
-        # [p4tt] p4_tail_fast: tiny-tie COLLECT+SELECT fast path inside the
+        # p4_tail_fast: tiny-tie COLLECT+SELECT fast path inside the
         # exact-tail fire branch. When the (b*, sb*) tie class holds <= 128
         # entries (the real firing cells have 2), ONE candidate pass collects
         # (value_bits, cand_idx) pairs into SMEM and thread0 selects the
@@ -1437,7 +1423,7 @@ def phase0_scan_bucket(
             # just the one 32B bmax vector).
             bm_addr = block_max_row.iterator.toint()
             nb0 = (N + cutlass.Int32(31)) >> cutlass.Int32(5)
-            # v9 two-pass skip: (1) DENSE-scan the bmax array itself (it
+            # Two-pass skip: (1) DENSE-scan the bmax array itself (it
             # is 1/32 of the row) with the tuned vector loop, compacting
             # PASSING BLOCK IDS into the idle C segment (single-band mode
             # never fills C; ids < 2^23 store exactly as floats);
@@ -1614,7 +1600,7 @@ def phase0_scan_bucket(
         cpw = cutlass.const_expr(4)  # cp.async caps at 16B per copy
         step1 = cutlass.const_expr(num_threads * cpw)
         st2log = cutlass.const_expr((2 * step1).bit_length() - 1)
-        # v14 pair-step cp.async pipeline: the scan is instruction-issue
+        # Pair-step cp.async pipeline: the scan is instruction-issue
         # bound (pcsamp: no_instructions + wait dominate; long_scoreboard
         # is 6%), so each step processes TWO 16B vectors per thread —
         # loop/wait/commit/address overhead amortizes over 8 elements
@@ -2249,17 +2235,14 @@ def block_count_ge_multi(
                         cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vh >= thr_frag[m])
                     hh = hh + cutlass.Int32(num_threads)
                 # Rung-tightening build (cs==1): a rung whose active list
-                # exceeds CAP blocks is provably or near-provably
-                # unacceptable (count >= list length; on real low-hit-rate
-                # rows the loosest sample-quantile rung retains 60%+ of the
-                # blocks and destroys the skip). DROP it — a dropped rung is
-                # merely an unmeasured probe (recorded in the mask at
-                # s_active_cnt[2]; classify and the fallback seeding skip
-                # it) — and rebuild at the next tighter threshold. Bounded
-                # by M-1 extra builds (~2-5us each at nb=8192). At cs>1 the
-                # per-CTA list lengths differ, so the drop decision would
-                # diverge across the cluster: keep the plain loosest-rung
-                # build there.
+                # exceeds CAP blocks cannot be accepted (count >= list
+                # length), so drop it and rebuild at the next tighter
+                # threshold. A dropped rung is only an unmeasured probe,
+                # recorded in the mask at s_active_cnt[2] and skipped by
+                # classify and by the fallback seeding. Bounded by M-1
+                # extra builds. At cs>1 the per-CTA list lengths differ
+                # and the drop decision would diverge across the cluster,
+                # so keep the plain loosest-rung build there.
                 CAP_BLOCKS = cutlass.const_expr(3 * self.kC // 4)
                 if cutlass.const_expr(cluster_size == 1):
                     build_done = cutlass.Int32(0)
@@ -3874,7 +3857,7 @@ def phase4_rank_scatter(
                 # slot range [rank_above_fine, kK). Unambiguous rows (the
                 # overwhelming majority) pay two scalar compares; the counters
                 # and the fine histogram are reused, so SMEM does not grow.
-                # [p4tt] boundary-class repair: collect the (b*, sb*) tie
+                # boundary-class repair: collect the (b*, sb*) tie
                 # class compactly (ONE candidate pass), then select inside
                 # it. Tiny jobs (need x class <= 512) keep the thread0
                 # serial select (cheapest at that size). Bigger classes up
@@ -3887,7 +3870,7 @@ def phase4_rank_scatter(
                 if cutlass.const_expr(self.p4_exact_tail and self.p4_tail_fast):  # [p4tt]
                     if cutlass.const_expr(self.p4_tail_v3):
                         need0 = cutlass.Int32(kK) - rank_above_fine
-                        # [p4tt-v3] per-thread compact buffers, bounded by the
+                        # per-thread compact buffers, bounded by the
                         # strided trip count over the candidate array
                         nbuf7 = cutlass.const_expr(
                             (self.kC + self.num_threads - 1) // self.num_threads
@@ -3901,7 +3884,7 @@ def phase4_rank_scatter(
                         fast_done = cutlass.Int32(1)
                         if cnt_strad > need0 and need0 > cutlass.Int32(0):
                             fast_done = cutlass.Int32(0)
-                            # [p4tt-v3] block-wide pure-tie check, ANY class
+                            # block-wide pure-tie check, ANY class
                             # size: min/max order key over the (b*, sb*) class.
                             # A pure-tie class needs NO repair — the scatter's
                             # arrival fill of bit-equal values is already
@@ -3951,7 +3934,7 @@ def phase4_rank_scatter(
                             if kmn7 == kmx7:
                                 fast_done = cutlass.Int32(1)
                             if fast_done == cutlass.Int32(0):
-                                # [p4tt-v3] mixed class: compact it IN PLACE
+                                # mixed class: compact it IN PLACE
                                 # into smem_keys/vals[0..cnt_strad) with a
                                 # register-buffered two-phase pass (every
                                 # thread reads its strided candidates first,
@@ -4166,7 +4149,7 @@ def phase4_rank_scatter(
                         need0_s = cutlass.Int32(kK) - rank_above_fine
                         if cnt_strad > need0_s and need0_s > cutlass.Int32(0):
                             if cnt_strad <= cutlass.Int32(128):
-                                # [p4tt] SMEM: (value_bits, cand_idx) pairs at
+                                # SMEM: (value_bits, cand_idx) pairs at
                                 # smem_hist[2*o]/[2*o+1], o < 128 (slots 0..255).
                                 # The 256 digit bins are dead here (the fast path
                                 # replaces the radix levels that used them); the
@@ -4204,7 +4187,7 @@ def phase4_rank_scatter(
                                                 ]
                                     itc = itc + cutlass.Int32(num_threads)
                                 cute.arch.barrier()
-                                # [p4tt] thread0 exact top-need0_s select rewriting
+                                # thread0 exact top-need0_s select rewriting
                                 # positions [rank_above_fine, kK). Consumed flag =
                                 # the cand_idx slot set to -1 (indices are always
                                 # >= 0), so a genuine -FLT_MAX value in the class
@@ -4412,7 +4395,7 @@ def phase4_rank_scatter(
                                                         output_indices_row[pos_s] = smem_vals[ir2]
                                     ir2 = ir2 + cutlass.Int32(num_threads)
                                 cute.arch.barrier()
-                elif cutlass.const_expr(self.p4_exact_tail):  # [p4tt] if->elif only
+                elif cutlass.const_expr(self.p4_exact_tail):  # if->elif only
                     need0 = cutlass.Int32(kK) - rank_above_fine
                     if cnt_strad > need0 and need0 > cutlass.Int32(0):
                         # Persistent scalars live above the 256 digit bins
@@ -6125,24 +6108,17 @@ def _run_phases(
                     if cutlass.const_expr(self.emit_xstate):
                         xstate_row[0] = cutlass.Float32(0.0)  # degenerate
         else:
-            # ---- List path v4: known-counts admission ----
-            # The emitter wrote the SoA list (score column + position
-            # column, sentinel score -inf) collected at t0 = seed_thr[0]
-            # and COUNTED the two tighter lines on the way out: the
-            # control words carry {n0, void, n1, n2} with n_i = #(>= t_i),
-            # n0 >= n1 >= n2. Admission and cut selection are pure scalar
-            # lookups - no in-kernel counting, no staging, no gamble:
-            #   1. some n_i lands in the acceptance band [K, B*]  ->
-            #      cut at the TIGHTEST such line, ONE filtered pass.
-            #   2. the band is straddled or overshot by every line ->
-            #      a histogram over the gmem list CLAMPED between the
-            #      two known bracket lines finds an in-band edge (the
-            #      narrow domain kills the long-tail bin collapse).
-            #   3. void, or n0 < K + 64 (64 = emitter sentinel bound,
-            #      so live coverage of K is proven)  ->  fallback.
-            # The K+64 slack also lets every accepted cut load run
-            # WITHOUT any overflow net: counts and load predicates are
-            # the same comparison on the same data.
+            # ---- List path: known-counts admission ----
+            # The emitter wrote an SoA list collected at t0 = seed_thr[0]
+            # and counted the two tighter lines on the way out, so the
+            # control words carry {n0, void, n1, n2} with n_i = #(>= t_i).
+            # Admission is then a scalar lookup:
+            #   1. some n_i in [K, B*] -> cut at the tightest, one pass
+            #   2. every line straddles/overshoots -> histogram over the
+            #      list clamped between the two known bracket lines
+            #   3. void, or n0 < K + 64 (emitter sentinel bound) -> fall
+            #      back. The slack also lets accepted cuts load without
+            #      an overflow net: count and load are the same compare.
             # cs>1 and 16-bit dtypes keep the plain fallback.
             take_cand = cutlass.Int32(0)
             list_used = cutlass.Int32(0)  # list path taken (xstate publish)
@@ -6153,7 +6129,7 @@ def _run_phases(
                 and cluster_size == 1
                 and self.dtype == cutlass.Float32
             ):
-                # ---- List path v5: BUCKETED segments ----
+                # ---- List path: bucketed segments ----
                 # The emitter classifies each entry by the tightest line
                 # it passes and appends into one of three fixed segments
                 # (A = [0, segA) holds >= t2, B = [segA, 2*segA) holds
@@ -6767,23 +6743,18 @@ def _run_phases(
                     # preIdx stats are full-row).
                     if cutlass.const_expr(self.use_ext_counts):
                         if ext_row == cutlass.Int32(1):
-                            # ---- Waterfall L1 admission (ext rungs, v2a) ----
-                            # Rung thresholds arrive from the indexer epilogue:
-                            # ONLY P1b is skipped. The stock M-ary count pass runs
-                            # on the ext rungs so the block-skip list build, rung
-                            # tightening, per-thread hand-off and classify all
-                            # compose unchanged (v1 routed through the dense
-                            # refine and forfeited the compact-walk win: flash 1M
-                            # ext 34.7us vs skipR0 15.6us cold).
-                            # v2b: when an ext count is already in [K, kC], park
-                            # THE ADMITTED THRESHOLD IN ALL RUNG SLOTS — the M-ary
-                            # pass degenerates to one compact single-threshold
-                            # count (+ list build at that threshold) and classify
-                            # admits it; a full miss keeps the 3 distinct rungs as
-                            # measured brackets for the seeded refine.
-                            # rung parking + single-path staging happen in
-                            # the P1-init thread0 block (one barrier for
-                            # the whole admission prologue).
+                            # ---- Waterfall L1 admission (ext rungs) ----
+                            # Rung thresholds arrive from the indexer epilogue,
+                            # so only P1b is skipped: the stock M-ary count pass
+                            # runs on the ext rungs and the block-skip list
+                            # build, rung tightening, hand-off and classify all
+                            # compose unchanged. When an ext count already lies
+                            # in [K, kC] the admitted threshold is parked in all
+                            # rung slots, degenerating the pass to one compact
+                            # single-threshold count; a full miss keeps the
+                            # three distinct rungs as brackets for the refine.
+                            # Parking and staging happen in the P1-init thread0
+                            # block, one barrier for the whole prologue.
                             if cutlass.const_expr(not self.enable_block_skip):
                                 # v3: the parked M-ary pass counted the SAME
                                 # threshold in all three columns (3x compare

From d697e3d6d054d8846b4530d4de759ec18ebb429c Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 23:11:52 -0700
Subject: [PATCH 068/117] [None][fix] GVR: drop the const-flag branch shells
 from the stock path

Five sites set a routing flag to a compile-time Int32(0), let a
const_expr-gated feature block possibly set it to 1, then branched on
it with a runtime compare. With the feature off the predicate is a
constant, but the compare is still traced, so the stock build carries
an scf.if region upstream does not have and every value written inside
it becomes a region yield. The largest wrapped all of Phase 2 and
Phase 3 - 470 lines - in one region.

Each site now computes a Python bool under the const_expr guard that
already exists a few lines above, so the stock configuration traces the
body straight-line exactly as upstream does:

  dense_ok        block_count_ge_multi, two shells (block skip)
  park_cursors    phase3_collect cursor parking (block skip)
  run_stock_range phase4 block min/max (ext candidate range)
  run_stock_p23   the Phase 2 + Phase 3 pipeline (list / self-scan)
  run_mary        the M-ary count pass (ext counts)

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 40 +++++++++++--------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 3d5d29c3e0af..f1f45834b6ef 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -2202,6 +2202,7 @@ def block_count_ge_multi(
         # here separately (per-thread order contract: head elements FIRST,
         # then list entries — Phase 3's compact write replays the same).
         skip_ok = cutlass.Int32(0)
+        dense_ok = True  # Python bool: no scf.if when block skip is off
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             skip_ok = cutlass.Int32(1)
             # Capacity/id-width guard: the active list holds at most
@@ -2220,6 +2221,7 @@ def block_count_ge_multi(
                 skip_ok = cutlass.Int32(0)
             if blk_hi_g > cutlass.Int32(32767):
                 skip_ok = cutlass.Int32(0)
+            dense_ok = skip_ok == cutlass.Int32(0)
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             if skip_ok == cutlass.Int32(1):
                 head_end = (
@@ -2360,7 +2362,7 @@ def block_count_ge_multi(
                                     jj = jj + cutlass.Int32(1)
                     li = li + cutlass.Int32(stride_un)
 
-        if self.enable_unroll_4 and skip_ok == cutlass.Int32(0):
+        if self.enable_unroll_4 and dense_ok:
             rng_frag = cute.make_fragment((vec_w,), self.dtype)
             big_iters = cutlass.Int32(0)
             if slice_end > i + cutlass.Int32(vec_w - 1):
@@ -2387,7 +2389,7 @@ def block_count_ge_multi(
             i = i + big_iters * cutlass.Int32(step_elem)
 
         tail_frag = cute.make_fragment((vec_w,), self.dtype)
-        if skip_ok == cutlass.Int32(0):
+        if dense_ok:
             while i + cutlass.Int32(vec_w - 1) < slice_end:
                 src_ptr = cute.make_ptr(
                     self.dtype,
@@ -2932,9 +2934,11 @@ def phase3_collect_candidates(
         # Only taken when the list is CURRENT (s_active_cnt[1] == 1, set by
         # the build; cleared on any dense fallback re-count).
         skip_wr = cutlass.Int32(0)
+        park_cursors = False  # Python bool: no scf.if when block skip is off
         if cutlass.const_expr(self.enable_block_skip and smem_active is not None):
             if s_active_cnt[1] == cutlass.Int32(1):
                 skip_wr = cutlass.Int32(1)
+            park_cursors = skip_wr == cutlass.Int32(1)
         if cutlass.const_expr(self.enable_block_skip and smem_active is not None):
             if skip_wr == cutlass.Int32(1):
                 # head region first — same per-thread order as the count pass
@@ -3001,7 +3005,7 @@ def phase3_collect_candidates(
         # When the compact write ran, park the dense cursors at the end so
         # all three dense loops below (4-way, vec tail, scalar tail) fall
         # through without re-indenting them.
-        if skip_wr == cutlass.Int32(1):
+        if park_cursors:
             ic = N_local
             n_aligned = N_local
 
@@ -3521,9 +3525,13 @@ def phase4_rank_scatter(
                 sc0 = cute.arch.clock64()
             bmin_r = cutlass.Float32(self.FLT_MAX)
             bmax_r = cutlass.Float32(self.NEG_FLT_MAX)
-            use_ext_r = cutlass.Int32(0)
+            # a Python bool here, so with the feature off the stock body
+            # below traces straight-line instead of inside an scf.if whose
+            # predicate is a compile-time constant
+            run_stock_range = True
             if cutlass.const_expr(ext_range_flag is not None):
                 use_ext_r = ext_range_flag
+                run_stock_range = use_ext_r == cutlass.Int32(0)
                 if use_ext_r == cutlass.Int32(1):
                     # list rows: the take walk pre-zeroed the hist and staged
                     # per-warp maxima in smem_wcnt (its end barrier orders
@@ -3538,7 +3546,7 @@ def phase4_rank_scatter(
                         bmax_r = cute.arch.fmax(bmax_r, vmax)
                     if bmax_r <= bmin_r:
                         bmax_r = bmin_r + cutlass.Float32(1e-6)
-            if use_ext_r == cutlass.Int32(0):
+            if run_stock_range:
                 # ---- block min/max over candidates ----
                 local_cmin = cutlass.Float32(self.FLT_MAX)
                 local_cmax = cutlass.Float32(self.NEG_FLT_MAX)
@@ -6121,6 +6129,10 @@ def _run_phases(
             #      an overflow net: count and load are the same compare.
             # cs>1 and 16-bit dtypes keep the plain fallback.
             take_cand = cutlass.Int32(0)
+            # Python bool: with every list/scan feature off, the stock
+            # Phase 2 + Phase 3 below trace straight-line rather than as
+            # one 470-line scf.if region with a large yield list.
+            run_stock_p23 = True
             list_used = cutlass.Int32(0)  # list path taken (xstate publish)
             claimed_c = cutlass.Int32(0)
             if cutlass.const_expr(
@@ -6715,7 +6727,9 @@ def _run_phases(
                 if cutlass.const_expr(_P4_TAIL_DBG):
                     ck1 = cute.arch.clock64()
 
-            if take_cand == cutlass.Int32(0):
+            if cutlass.const_expr(self.use_ext_cand or self.use_ext_counts or self.self_scan):
+                run_stock_p23 = take_cand == cutlass.Int32(0)
+            if run_stock_p23:
                 # Stage this CTA's slice into SMEM once before Phase 2's
                 # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses
                 # scatter-loads OUTSIDE this slice, so it stays on GMEM.
@@ -6880,6 +6894,7 @@ def _run_phases(
                                 lane,
                             )
                     r0_par = cutlass.Int32(0)
+                    run_mary = True  # Python bool: no scf.if without ext counts
                     if cutlass.const_expr(self.use_ext_counts and not self.enable_block_skip):
                         # single-column fast path accepted: the parked count
                         # is done and admitted; the M-ary pass, argmin,
@@ -6889,15 +6904,8 @@ def _run_phases(
                             1
                         ] == cutlass.Int32(1):
                             r0_par = cutlass.Int32(1)
-                    if r0_par == cutlass.Int32(0):
-                        # NOTE: a loosest-line pass-rate veto (packed seed
-                        # col 6) was measured here and REMOVED: the rung-
-                        # tightening build below already salvages fat
-                        # loosest-rung rows by rebuilding at a tighter
-                        # line, so any veto threshold (3/8 and 7/8 both
-                        # tried, cold, 30 layers) only preempts that and
-                        # costs 3-18% across the board. Col 6 stays as a
-                        # zero-cost diagnostic for host-side routing.
+                        run_mary = r0_par == cutlass.Int32(0)
+                    if run_mary:
                         self.block_count_ge_multi(
                             input_row,
                             slice_start,
@@ -6917,7 +6925,7 @@ def _run_phases(
                             s_active_cnt=s_active_cnt,
                         )
                     cute.arch.barrier()
-                    if tidx == 0 and r0_par == cutlass.Int32(0):
+                    if run_mary and tidx == 0:
                         # tightest admissible rung = SMALLEST count in [K, kC].
                         # (Explicit argmin: with r0_vseed the pmean column is not
                         # sorted into the rung order; for sorted rungs this is

From e3db787e2695eba7a13584c831999aee12cabc83 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 23:39:12 -0700
Subject: [PATCH 069/117] [None][chore] GVR: put the P1r degenerate rescue
 behind a const_expr

Default ON - upstream's identity shortcut is wrong on real data and
gating it off would put that defect back. The knob makes the cost
measurable, and it is worth stating:

  degenerate prev_topk (the rescue fires):  +7.8 / +8.2 us
  real prev_topk (it does not):             -0.26 / -0.28 us, 17 and
                                            16 of 40 paired reps slower,
                                            i.e. noise

So the rescue is free on steady-state steps and costs ~8 us on the
first decode step of a request, which is where the zero-init feedback
buffer makes the bracket degenerate. That buys a correct first step -
without it 42 of 231 rows on a V4-Flash TP2 capture are wrong.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 40 +++++++++++--------
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  4 ++
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index f1f45834b6ef..da0af580032b 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -296,6 +296,7 @@ def __init__(
         p4_exact_tail: Optional[bool] = None,
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
         p4_tail_v3: bool = False,
+        p1r_rescue: bool = True,
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
         enable_block_skip: bool = False,
@@ -751,6 +752,12 @@ def __init__(
         # select. Off by default - the stock body is what upstream
         # emits, and this rewrite only pays at K>=1024.
         self.p4_tail_v3 = bool(p4_tail_v3)
+        # p1r_rescue: rebuild the refine bracket from the row when the
+        # seed bracket is degenerate. ON by default - upstream's identity
+        # shortcut is wrong on real data (every request's first decode
+        # step feeds a zero-init prev_topk). The knob exists to measure
+        # its cost, not to ship it off.
+        self.p1r_rescue = bool(p1r_rescue)
 
     # ------------------------------------------------------------------
     # SMEM slice cache loader. Streams this CTA's slice GMEM → SMEM via
@@ -6069,22 +6076,23 @@ def _run_phases(
         # If the bracket is STILL degenerate after the rescue, every
         # in-range value is identical (or N <= K), and identity output is
         # then exact — keep the shortcut for exactly those rows.
-        v_lo = s_thr[1]
-        v_hi = s_thr[2]
-        if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi:
-            if N > cutlass.Int32(self.top_k):
-                self.phase1r_data_reseed(
-                    input_row,
-                    N,
-                    smem_wmin,
-                    smem_wmax,
-                    s_thr,
-                    s_iscalars,
-                    s_mt_thr,
-                    tidx,
-                    warp_id,
-                    lane,
-                )
+        if cutlass.const_expr(self.p1r_rescue):
+            v_lo = s_thr[1]
+            v_hi = s_thr[2]
+            if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi:
+                if N > cutlass.Int32(self.top_k):
+                    self.phase1r_data_reseed(
+                        input_row,
+                        N,
+                        smem_wmin,
+                        smem_wmax,
+                        s_thr,
+                        s_iscalars,
+                        s_mt_thr,
+                        tidx,
+                        warp_id,
+                        lane,
+                    )
         v_lo = s_thr[1]
         v_hi = s_thr[2]
         if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi:
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 30ba825ec8ee..f0bbc8b8a2a6 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -68,6 +68,7 @@ def _compile(
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
     p4_tail_v3: bool = False,
+    p1r_rescue: bool = True,
     p4_fine_rangetest: "bool | None" = None,
     p4_scat_rangetest: bool = False,
     use_ext_counts: bool = False,
@@ -209,6 +210,7 @@ def _compile(
         enable_block_skip=enable_block_skip,
         pdl_wait_late=pdl_wait_late,
         p4_tail_v3=p4_tail_v3,
+        p1r_rescue=p1r_rescue,
         p4_fine_rangetest=p4_fine_rangetest,
         p4_scat_rangetest=p4_scat_rangetest,
         use_ext_counts=use_ext_counts,
@@ -674,6 +676,7 @@ def gvr_topk_decode(
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
     p4_tail_v3: bool = False,
+    p1r_rescue: bool = True,
     p4_fine_rangetest: Optional[bool] = None,
     p4_scat_rangetest: bool = False,
     block_max: Optional[torch.Tensor] = None,
@@ -964,6 +967,7 @@ def gvr_topk_decode(
         enable_block_skip,
         pdl_wait_late,
         p4_tail_v3,
+        p1r_rescue,
         p4_fine_rangetest,
         p4_scat_rangetest,
         use_ext_counts,

From 1fe0feee0649b91e780ad7e7284911f0e56c5f9e Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 29 Jul 2026 23:47:57 -0700
Subject: [PATCH 070/117] [None][fix] GVR: keep range-test and debug scalars
 off the stock path

f_hi / lo_edge / hi_edge feed only the two range-test arms, both off by
default, but were computed unconditionally - a divide and two compares
upstream never emits. The P4 sub-phase and tail stopwatches were
likewise declared outside their debug guards, which promotes them to
region yields when a neighbouring branch is traced.

Both now sit under the const_expr that already gates their only
readers. Four-path exactness (stock / stock+v3 / counts / counts+bm)
unchanged.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 33 ++++++++++---------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index da0af580032b..62713e5f1982 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3521,14 +3521,13 @@ def phase4_rank_scatter(
                 output_indices_row[i4] = smem_vals[i4]
                 i4 = i4 + cutlass.Int32(num_threads)
         elif cand_count > cutlass.Int32(kK):
-            sc0 = cutlass.Int64(0)
-            sc1 = cutlass.Int64(0)
-            sc2 = cutlass.Int64(0)
-            sc3 = cutlass.Int64(0)
-            sc4 = cutlass.Int64(0)
-            sc5 = cutlass.Int64(0)
-            sc6 = cutlass.Int64(0)
             if cutlass.const_expr(_P4_SUB_DBG):
+                sc1 = cutlass.Int64(0)
+                sc2 = cutlass.Int64(0)
+                sc3 = cutlass.Int64(0)
+                sc4 = cutlass.Int64(0)
+                sc5 = cutlass.Int64(0)
+                sc6 = cutlass.Int64(0)
                 sc0 = cute.arch.clock64()
             bmin_r = cutlass.Float32(self.FLT_MAX)
             bmax_r = cutlass.Float32(self.NEG_FLT_MAX)
@@ -3671,12 +3670,14 @@ def phase4_rank_scatter(
                 # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1)
                 f_lo = bmin_r + cutlass.Float32(b_star) / inv1
                 finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1
-                # bin b* spans [f_lo, f_hi) under the coarse binning; the
-                # clamped ends fold out-of-range values INTO bin 0 and bin
-                # kBins-1, so those two bins drop the matching side.
-                f_hi = f_lo + cutlass.Float32(1.0) / inv1
-                lo_edge = b_star == cutlass.Int32(0)
-                hi_edge = b_star == cutlass.Int32(kBins - 1)
+                # bin b* spans [f_lo, f_hi); the clamped ends fold
+                # out-of-range values into bin 0 and bin kBins-1, so those
+                # two drop the matching side. Only the range-test arms read
+                # these, so upstream's build must not compute them.
+                if cutlass.const_expr(self.p4_fine_rangetest or self.p4_scat_rangetest):
+                    f_hi = f_lo + cutlass.Float32(1.0) / inv1
+                    lo_edge = b_star == cutlass.Int32(0)
+                    hi_edge = b_star == cutlass.Int32(kBins - 1)
                 # re-zero (only fbins slots) + build fine sub-hist of bin-b* cands
                 iz = tidx
                 while iz < cutlass.Int32(fbins):
@@ -5856,10 +5857,10 @@ def _run_phases(
         # misses (pro-1M cold: stock-skip 21us vs ext-miss 51us). All
         # threads read the same control words, so the predicate is
         # CTA-uniform and the dynamic branches below stay convergent.
-        ck0 = cutlass.Int64(0)
-        ck1 = cutlass.Int64(0)
-        ckE = cutlass.Int64(0)
         if cutlass.const_expr(_P4_TAIL_DBG):
+            ck0 = cutlass.Int64(0)
+            ck1 = cutlass.Int64(0)
+            ckE = cutlass.Int64(0)
             ckE = cute.arch.clock64()  # row-phase entry (device-residency ref)
         ext_row = cutlass.Int32(0)
         if cutlass.const_expr(self.use_ext_counts):

From 84f8a36ca3e655f61cb853d38c1ee4c21741e20d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 00:49:33 -0700
Subject: [PATCH 071/117] [None][fix] GVR: let the assist path keep the
 compacted tail repair

Defaulting p4_tail_v3 to False isolated the stock path but also took
the rewrite away from our own tiers, since nothing on the production
side passes the flag. Derive it from the configuration instead: a stock
kernel gets upstream's thread0 serial select, any emission-assisted or
block-skip kernel gets the compacted repair. No plumbing through the
router, the ext state or the op face.

Stock-path constants were diffed against origin/main for K=1024 and
K=512: every constant upstream also has now resolves to the same value,
and the 22 extras are all new facilities sitting at their off state
(p1r_rescue is the one deliberate exception, see its own commit).

Four-path exactness unchanged: stock / stock+v3 / counts / counts+bm.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py                 | 10 +++++++---
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py   |  4 ++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 62713e5f1982..cf12941cb2db 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -295,7 +295,7 @@ def __init__(
         enable_p4_rank_scatter_exact: Optional[bool] = None,
         p4_exact_tail: Optional[bool] = None,
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
-        p4_tail_v3: bool = False,
+        p4_tail_v3: Optional[bool] = None,
         p1r_rescue: bool = True,
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
@@ -749,8 +749,12 @@ def __init__(
         self.p4_tail_fast = bool(p4_tail_fast) and self.p4_exact_tail  # [p4tt]
         # p4_tail_v3: compacted-class repair (block-parallel radix +
         # pure-tie pre-check) in place of the stock thread0 serial
-        # select. Off by default - the stock body is what upstream
-        # emits, and this rewrite only pays at K>=1024.
+        # select. Default follows the configuration: a stock kernel gets
+        # upstream's body, any emission-assisted one gets the rewrite.
+        if p4_tail_v3 is None:
+            p4_tail_v3 = bool(
+                use_ext_counts or use_ext_cand or ext_rungs or self_scan or enable_block_skip
+            )
         self.p4_tail_v3 = bool(p4_tail_v3)
         # p1r_rescue: rebuild the refine bracket from the row when the
         # seed bracket is degenerate. ON by default - upstream's identity
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index f0bbc8b8a2a6..44585e2c210d 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -67,7 +67,7 @@ def _compile(
     p2_warp_redundant: bool = True,
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
-    p4_tail_v3: bool = False,
+    p4_tail_v3: "bool | None" = None,
     p1r_rescue: bool = True,
     p4_fine_rangetest: "bool | None" = None,
     p4_scat_rangetest: bool = False,
@@ -675,7 +675,7 @@ def gvr_topk_decode(
     p4_warp_redundant: bool = True,
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
-    p4_tail_v3: bool = False,
+    p4_tail_v3: "bool | None" = None,
     p1r_rescue: bool = True,
     p4_fine_rangetest: Optional[bool] = None,
     p4_scat_rangetest: bool = False,

From cf87bf0f8b1a29691d7d1c9c586857573a067076 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 01:36:30 -0700
Subject: [PATCH 072/117] [None][perf] GVR routing: pick the tier from the
 shape-dependent tax

The emission tax is a fraction of the indexer's own time - roughly
9-14% for the bucketed list, ~2% for the packed counts - so in absolute
terms it grows with B*N while what the top-k saves does not. The old
rule charged a flat cost, which made the counts tier look free at large
batch and kept the list tier out of the 32k-64k band where it already
pays.

Re-derived over the 126-cell worst-step grid with the tax modelled as a
fraction of the indexer time, evaluated at both ends of the measured
range (9% and 14%):

  rule                                   9%      11.5%    14%
  old (list >=65536 & B<=16, else counts)  1.193   1.149   1.110
  new                                      1.248   1.225   1.203
  per-cell optimum                         1.286   1.253   1.228

Losing cells go from 14/20/29 to 9/11/13. The new rule keeps ~85-90% of
what per-cell selection would buy, and holds its ranking across the
whole tax range rather than at one point in it.

  list   n_comp >= 32768 and B <= 16   (was 65536)
  rungs  B >= 32                        (new: emit nothing, no tax)
  counts otherwise

Lifting the batch cap was tested and is not viable: the list tier at
B>=32 drops the grid to geomean 1.022 and worst 0.411, because its tax
scales past the saving. The cap stays.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_routing.py            | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index 40453abe330f..fe47d17124de 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -52,10 +52,17 @@
 SKIP_MIN_N_RUNGS_FLASH = 131072  # vb (flash): bm pays from here
 SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split on top of bm from here
 
-# Emission tax ~ B*N on the GEMM side: the candidate list is only worth
-# emitting for latency-bound shapes (small B, long rows).
+# The emission tax is a fraction of the indexer's own time (list ~9-14%,
+# counts ~2%), so in absolute terms it grows with B*N while the top-k
+# saving does not. That gives three regimes, measured over the 126-cell
+# worst-step grid and checked at both ends of the tax range:
+#   long rows, small batch  -> the list repays its tax several times over
+#   large batch             -> emit nothing; the rungs tier has no tax
+#   otherwise               -> counts, whose tax is small either way
 LIST_EMIT_MAX_B = 16
-LIST_EMIT_MIN_N = 65536
+LIST_EMIT_MIN_N = 32768
+# Above this batch even the counts tax outruns what it buys.
+RUNGS_MIN_B = 32
 
 # rungs-tier block_max pays only at small K: with K=1024 the tight-line
 # pass rate runs too high and the prefix read is pure overhead.
@@ -97,8 +104,10 @@ def plan_emission(batch: int, n_comp: int, k: int, have_epilogue: bool) -> str:
     if not have_epilogue:
         return "rungs"  # closed-loop lines cost nothing to carry
     if batch <= LIST_EMIT_MAX_B and n_comp >= LIST_EMIT_MIN_N:
-        return "list"  # latency-bound long rows: list pays big
-    return "counts"  # near-free tax, wins almost everywhere
+        return "list"
+    if batch >= RUNGS_MIN_B:
+        return "rungs"  # throughput regime: emitting anything is a loss
+    return "counts"
 
 
 def pick_config(tier: str, batch: int, n_comp: int, k: int, num_sms: int) -> TopkRoute:

From 52973e6d39730a15d1f06108f6c8a964f06cc666 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 01:47:27 -0700
Subject: [PATCH 073/117] [None][fix] GVR routing: cap the list tier at batch
 1, measured

The previous two commits set the list-tier batch cap from a tax model,
first flat and then proportional to the indexer time. Both were wrong.
Measured on the FP4 indexer itself - same kernel, emission outputs
attached, wall delta - at ctx 256k and 1M:

  batch      1       2       4       8      32     128
  list   +13.8%  +64.3%  +60.7%  +62.0% +122.9% +150.7%
  counts  +9.4%   +3.5%   +0.9%   +1.1%   +0.9%   +2.6%

The list tier clears its own tax only at batch 1. From batch 2 the tax
is 20-60us against a top-k saving of a few, and it keeps growing. The
proportional model underestimated batch 128 by an order of magnitude
(45us modelled, 589us measured).

Re-scoring the 126-cell grid with the measured tax:

  rule                        geomean  worst  losing
  list B<=16, N>=65536          0.969  0.159   43
  list B<=16, N>=32768          0.926  0.159   41
  list B==1,  N>=32768          1.155  0.758   17
  per-cell optimum              1.207  0.763   14

So the cap goes to 1. Note the first row: the tier was already a net
loss as shipped before this series - the 3.4us flat tax it was sized
against only holds at batch 1.

The counts tax stays inside ~2.6us at every batch measured, so the
counts/rungs split from the previous commit is unaffected.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_routing.py            | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index fe47d17124de..0bbe6f52c0ff 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -52,16 +52,20 @@
 SKIP_MIN_N_RUNGS_FLASH = 131072  # vb (flash): bm pays from here
 SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split on top of bm from here
 
-# The emission tax is a fraction of the indexer's own time (list ~9-14%,
-# counts ~2%), so in absolute terms it grows with B*N while the top-k
-# saving does not. That gives three regimes, measured over the 126-cell
-# worst-step grid and checked at both ends of the tax range:
-#   long rows, small batch  -> the list repays its tax several times over
-#   large batch             -> emit nothing; the rungs tier has no tax
-#   otherwise               -> counts, whose tax is small either way
-LIST_EMIT_MAX_B = 16
+# Emission tax measured on the FP4 indexer itself (ctx 256k and 1M,
+# batch 1..128), as the wall delta of the same kernel with the emission
+# outputs attached:
+#
+#   batch    1      2      4      8     32     128
+#   list  +13.8% +64.3% +60.7% +62.0% +122.9% +150.7%
+#   counts +9.4%  +3.5%  +0.9%  +1.1%   +0.9%   +2.6%
+#
+# The list tier only clears its own tax at batch 1; from batch 2 the tax
+# is tens of microseconds against a top-k saving of a few. The counts
+# tax stays inside ~2.6us throughout, and at large batch even that is
+# not repaid, so the zero-emission rungs tier takes over.
+LIST_EMIT_MAX_B = 1
 LIST_EMIT_MIN_N = 32768
-# Above this batch even the counts tax outruns what it buys.
 RUNGS_MIN_B = 32
 
 # rungs-tier block_max pays only at small K: with K=1024 the tight-line

From a048dfa32b514235c9b689669f662a5a3a47a4eb Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 02:04:30 -0700
Subject: [PATCH 074/117] [None][perf] GVR routing: gate the list tier on
 emitted volume

Measuring the indexer across batch 1..1024 and ctx 32k..1M shows the
list tier's tax tracks B*N - the volume it emits - rather than batch:

  B*N (raw tokens)   <=0.5M     1M       >=2M
  list                +11-14%  +55-71%  +150-270%
  counts               +1-3%    +1-3%     +0-3%

Volume is the main term but not the only one: the same 1M tokens costs
+21% as 1M x B1 and +71% as 128k x B8, so there is a per-row component
too. Both cheap regions are covered by "a single row, or under ~0.75M
tokens", which lets 128k rows carry the list to batch 4 and 256k to
batch 2 - where the previous commit's flat B==1 cap gave them up.

Scored on the 126-cell grid with the measured tax:

  rule                       geomean  worst  losing
  pre-series (B<=16, N>=64k)   0.894  0.078   46
  previous (B==1)              1.145  0.769   19
  this (B==1 or <=0.75M)       1.162  0.786   17
  per-cell optimum             1.182  0.786   14

The counts tax stays inside 3% at every batch measured, so it remains
the default and the rungs handover at B>=32 is unchanged.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  6 ++--
 .../attention_backend/sparse/gvr_ext.py       |  8 +++--
 .../blackwell/top_k/gvr_routing.py            | 30 ++++++++++++-------
 3 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 00953c7dba6e..1f03b25d454a 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2837,9 +2837,11 @@ def sparse_attn_indexer(
                         n_comp = indexer_max_seq_len // max(
                             self.compress_ratio, 1)
                         emit_tier, self._gvr_route = st.plan(
-                            batch_size, n_comp,
+                            batch_size,
+                            n_comp,
                             torch.cuda.get_device_properties(
-                                q_fp8.device).multi_processor_count)
+                                q_fp8.device).multi_processor_count,
+                            compress_ratio=max(self.compress_ratio, 1))
                         st.update_seed_rows(batch_size)
                         gvr_emit_kwargs = st.indexer_emit_kwargs(
                             emit_tier, batch_size)
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
index 0611bbd10df7..296ec4a5ed86 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -101,10 +101,14 @@ def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
             )
         return self.block_max
 
-    def plan(self, batch: int, n_comp: int, num_sms: int) -> tuple[str, TopkRoute]:
+    def plan(
+        self, batch: int, n_comp: int, num_sms: int, compress_ratio: int = 4
+    ) -> tuple[str, TopkRoute]:
         """Route this step: (tier to EMIT next, launch knobs to CONSUME
         what was emitted last step)."""
-        emit_tier = plan_emission(batch, n_comp, self.top_k, have_epilogue=True)
+        emit_tier = plan_emission(
+            batch, n_comp, self.top_k, have_epilogue=True, compress_ratio=compress_ratio
+        )
         if emit_tier == "list" and self.cand_vals is None:
             # constructed with enable_list_tier=False: no candidate
             # buffers to emit into, demote to the counts tier
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index 0bbe6f52c0ff..f09d6aa895bb 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -56,15 +56,21 @@
 # batch 1..128), as the wall delta of the same kernel with the emission
 # outputs attached:
 #
-#   batch    1      2      4      8     32     128
-#   list  +13.8% +64.3% +60.7% +62.0% +122.9% +150.7%
-#   counts +9.4%  +3.5%  +0.9%  +1.1%   +0.9%   +2.6%
+# What the list tier costs tracks B*N - the total emitted volume - not
+# batch on its own:
 #
-# The list tier only clears its own tax at batch 1; from batch 2 the tax
-# is tens of microseconds against a top-k saving of a few. The counts
-# tax stays inside ~2.6us throughout, and at large batch even that is
-# not repaid, so the zero-emission rungs tier takes over.
-LIST_EMIT_MAX_B = 1
+#   B*N (raw tokens)   <=0.5M    1M      >=2M
+#   list                +11-14%  +55-71%  +150-270%
+#   counts               +1-3%    +1-3%     +0-3%
+#
+# Volume is the main term but not the only one: at the same 1M tokens
+# the tax is +21% as 1M x B1 and +71% as 128k x B8, so there is a
+# per-row cost on top. Both measured-cheap regions are covered by
+# "half a megatoken, or a single row". The counts tax stays inside 3%
+# everywhere, which is why it remains the default; past RUNGS_MIN_B
+# even that is not repaid and the zero-emission rungs tier takes over.
+LIST_EMIT_MAX_TOKENS = 786432  # B * raw length; between the measured
+# cheap band (<=0.52M) and the first expensive point (1.05M)
 LIST_EMIT_MIN_N = 32768
 RUNGS_MIN_B = 32
 
@@ -93,7 +99,9 @@ class TopkRoute:
     attach_block_max: bool = False
 
 
-def plan_emission(batch: int, n_comp: int, k: int, have_epilogue: bool) -> str:
+def plan_emission(
+    batch: int, n_comp: int, k: int, have_epilogue: bool, compress_ratio: int = 4
+) -> str:
     """Which assist tier the indexer epilogue should emit this step.
 
     ``n_comp``: compressed row length (post compress_ratio) - the
@@ -107,7 +115,9 @@ def plan_emission(batch: int, n_comp: int, k: int, have_epilogue: bool) -> str:
         return "none"
     if not have_epilogue:
         return "rungs"  # closed-loop lines cost nothing to carry
-    if batch <= LIST_EMIT_MAX_B and n_comp >= LIST_EMIT_MIN_N:
+    if n_comp >= LIST_EMIT_MIN_N and (
+        batch == 1 or batch * n_comp * compress_ratio <= LIST_EMIT_MAX_TOKENS
+    ):
         return "list"
     if batch >= RUNGS_MIN_B:
         return "rungs"  # throughput regime: emitting anything is a loss

From dc10a1ecf6d8d4f2df88e104571de4ddef938c9c Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 04:35:21 -0700
Subject: [PATCH 075/117] [None][perf] GVR routing: widen the list tier, hand
 over to rungs sooner

Scored on the full mean-over-layers-and-steps grid (length 8k-1024k,
batch 1-1024, all powers of two, measured emission tax):

  rule                          geomean  worst  losing
  list N>=32768, rungs B>=32      1.133  0.802  20/154
  list N>=16384                   1.140  0.802  16/154
  rungs B>=16                     1.133  0.802  19/154
  both                            1.140  0.802  15/154
  per-cell optimum                1.147  0.802  13/154

64k rows (n_comp 16387) were below the length gate but the list wins
there by ~2us at batch 2-8 once the volume gate lets it through, and at
batch 16 several cells prefer the zero-emission tier to counts.

The remaining band - 64k-512k at batch 2-16 - is not a routing problem:
the list tier is far worse there (0.65 down to 0.09; the emitted volume
crosses a megatoken and the tax takes the kernel from 17us to 100-200),
and counts is already the best of the three.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index f09d6aa895bb..83d3e1ff0e19 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -71,8 +71,8 @@
 # even that is not repaid and the zero-emission rungs tier takes over.
 LIST_EMIT_MAX_TOKENS = 786432  # B * raw length; between the measured
 # cheap band (<=0.52M) and the first expensive point (1.05M)
-LIST_EMIT_MIN_N = 32768
-RUNGS_MIN_B = 32
+LIST_EMIT_MIN_N = 16384
+RUNGS_MIN_B = 16
 
 # rungs-tier block_max pays only at small K: with K=1024 the tight-line
 # pass rate runs too high and the prefix read is pure overhead.

From 75f89e6214433435a252c2f4130bbf8bb5945e0e Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 05:06:55 -0700
Subject: [PATCH 076/117] [None][perf] GVR: turn the block-skip prefix on one
 doubling later

The counts tier attached block_max from n_comp 65536, but the prefix
does not earn its setup there. Attaching it to the same shape and
diffing, 30 interleaved cold reps:

  n_comp    batch 1   batch 4   batch 16
   65538      +0.0      -0.1      +0.6   (9-16 of 30 reps faster)
  131075      -4.6      -5.0      -4.5   (30 of 30)
  262127     -13.8     -14.2     -14.0   (30 of 30)

At 65538 it is a wash and a small loss at batch 16; the win starts at
131072. That threshold was where the mean-metric grid lost worst:
flash 256k (n_comp 65538) measured 0.802 / 0.834 / 0.859 at batch
4 / 8 / 16 against the baseline. With the prefix off those cells become
1.086 / 1.150 / 1.199 - the loss was the prefix's setup, not the tier.

The time model says the same thing: the counts tier's slope is already
better than the baseline's (28-35 vs 32-52 us per million compressed
elements) and its fixed cost is the same, so a 3-4us step at one
length was the whole gap.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/blackwell/top_k/gvr_routing.py      | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index 83d3e1ff0e19..409a269f0254 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -47,8 +47,13 @@
 # whole n_comp/K < 1.5 band measures below parity.
 ASSIST_MIN_N_COMP = 4096  # ~8k raw context at compress_ratio 4
 
-# Block-skip prefix pays only when whole-row reads dominate.
-SKIP_MIN_N_COUNTS = 65536  # va: attach block_max unconditionally here up
+# Block-skip prefix pays only when whole-row reads dominate. Measured
+# by attaching block_max to the same shape and diffing (30 interleaved
+# cold reps): at n_comp 65538 it is a wash (+0.0 / -0.1 / +0.6us at
+# batch 1/4/16, 9-16 of 30 reps faster), at 131075 it is -4.5 to -5.0us
+# and wins 30 of 30, at 262127 it is -14us. So the prefix earns its
+# setup one doubling later than it was switched on.
+SKIP_MIN_N_COUNTS = 131072  # va: attach block_max from here up
 SKIP_MIN_N_RUNGS_FLASH = 131072  # vb (flash): bm pays from here
 SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split on top of bm from here
 

From cda9b0390469852d05b56a3f520f8dddfbad5f13 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 05:20:05 -0700
Subject: [PATCH 077/117] [None][chore] GVR: make the coarse histogram width a
 knob

Default unchanged. Added to test whether the P4 head's bin search sets
the short-row floor - it does not: at pro n=4099 and n=8195, batch 1,
widths 128/256/512/1024 land within 0.5us of each other and wider is
marginally faster, so the search is not the bottleneck. Kept because it
is the only way to ask that question again without patching the file.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py      | 5 ++++-
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py         | 4 ++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index cf12941cb2db..b409780e2bbc 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -297,6 +297,7 @@ def __init__(
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
         p4_tail_v3: Optional[bool] = None,
         p1r_rescue: bool = True,
+        num_bins: Optional[int] = None,
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
         enable_block_skip: bool = False,
@@ -415,7 +416,9 @@ def __init__(
 
         params = GvrParams.get(self._dtype_name, top_k, self.compress_ratio)
         self.kC = params.kC
-        self.kNumBins = params.kNumBins
+        # num_bins: the coarse histogram width. The bin search walks it
+        # per warp, so it sets the barrier count of the P4 head.
+        self.kNumBins = params.kNumBins if num_bins is None else int(num_bins)
         self.kFTarget = params.kFTarget
 
         # Kernel-wide constants.
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 44585e2c210d..4880a592da9d 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -69,6 +69,7 @@ def _compile(
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
     p1r_rescue: bool = True,
+    num_bins: "int | None" = None,
     p4_fine_rangetest: "bool | None" = None,
     p4_scat_rangetest: bool = False,
     use_ext_counts: bool = False,
@@ -211,6 +212,7 @@ def _compile(
         pdl_wait_late=pdl_wait_late,
         p4_tail_v3=p4_tail_v3,
         p1r_rescue=p1r_rescue,
+        num_bins=num_bins,
         p4_fine_rangetest=p4_fine_rangetest,
         p4_scat_rangetest=p4_scat_rangetest,
         use_ext_counts=use_ext_counts,
@@ -677,6 +679,7 @@ def gvr_topk_decode(
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
     p1r_rescue: bool = True,
+    num_bins: "int | None" = None,
     p4_fine_rangetest: Optional[bool] = None,
     p4_scat_rangetest: bool = False,
     block_max: Optional[torch.Tensor] = None,
@@ -968,6 +971,7 @@ def gvr_topk_decode(
         pdl_wait_late,
         p4_tail_v3,
         p1r_rescue,
+        num_bins,
         p4_fine_rangetest,
         p4_scat_rangetest,
         use_ext_counts,

From 738ce6837e38abfa1c4db5ebfb9e80973a7ed1d5 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 30 Jul 2026 23:55:48 -0700
Subject: [PATCH 078/117] [None][perf] GVR routing: fix the block-skip and
 cluster-split gates

The previous commit raised the counts tier's block_max threshold to
131072 on a synthetic-Gaussian A/B. On captured V4 rows the prefix
already pays at 65536 (flash n_comp 65537, nsys kernel-only over 5
layers x all decode steps: 16.1us with vs 17.3us without at batch
4/8/16), so put it back and note that this threshold must be set from
captures.

Raising it also silently enabled cluster split for the rungs tier at
65537, because the cs=1 clamp it used to hit was conditioned on
block_max being attached. Cluster split is a loss for the assist tiers
at that length whether or not the prefix is on (rungs at batch 4:
17.9 / 19.0 / 20.4us at cs 1 / 4 / 8; cs8 spills to 31.5us at batch
16), so gate the split on its own threshold instead.

Worst mean-metric cell on the 154-cell grid: 0.794 -> 0.833.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_routing.py            | 24 ++++++++++---------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index 409a269f0254..200db91861d3 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -47,15 +47,19 @@
 # whole n_comp/K < 1.5 band measures below parity.
 ASSIST_MIN_N_COMP = 4096  # ~8k raw context at compress_ratio 4
 
-# Block-skip prefix pays only when whole-row reads dominate. Measured
-# by attaching block_max to the same shape and diffing (30 interleaved
-# cold reps): at n_comp 65538 it is a wash (+0.0 / -0.1 / +0.6us at
-# batch 1/4/16, 9-16 of 30 reps faster), at 131075 it is -4.5 to -5.0us
-# and wins 30 of 30, at 262127 it is -14us. So the prefix earns its
-# setup one doubling later than it was switched on.
-SKIP_MIN_N_COUNTS = 131072  # va: attach block_max from here up
+# Block-skip prefix pays for the counts tier from 65536 up, but not for
+# the zero-emission rungs tier below 131072. Measured on captured V4
+# rows (5 layers x all decode steps, nsys kernel-only, flash n_comp
+# 65537): counts 16.1us with the prefix vs 17.3 without at batch
+# 4/8/16, rungs 18.5 with vs 17.9 without. A synthetic-Gaussian A/B put
+# the counts break-even a doubling later - the block-max distribution
+# of real rows decides the pass rate, so set this from captures only.
+SKIP_MIN_N_COUNTS = 65536  # va: attach block_max from here up
 SKIP_MIN_N_RUNGS_FLASH = 131072  # vb (flash): bm pays from here
-SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split on top of bm from here
+# Cluster split is a loss for the assist tiers below this point,
+# block_max or not: at n_comp 65537 rungs measures 17.9 / 19.0 / 20.4us
+# at cs 1 / 4 / 8 (batch 4) and cs8 spills to 31.5us at batch 16.
+SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split from here up
 
 # Emission tax measured on the FP4 indexer itself (ctx 256k and 1M,
 # batch 1..128), as the wall delta of the same kernel with the emission
@@ -147,13 +151,11 @@ def pick_config(tier: str, batch: int, n_comp: int, k: int, num_sms: int) -> Top
     # rungs (vb)
     if k <= RUNGS_BM_MAX_K and n_comp >= SKIP_MIN_N_RUNGS_FLASH:
         r.attach_block_max = True
-    if n_comp >= 65536:
+    if n_comp >= SKIP_CS_MIN_N_RUNGS:
         if batch * 8 <= num_sms // CS8_HALF_DEVICE:
             r.cluster_size = 8
         elif batch * 4 <= (num_sms * CS_HEADROOM_NUM) // CS_HEADROOM_DEN:
             r.cluster_size = 4
         elif batch * 2 <= (num_sms * CS_HEADROOM_NUM) // CS_HEADROOM_DEN:
             r.cluster_size = 2
-    if r.attach_block_max and n_comp < SKIP_CS_MIN_N_RUNGS:
-        r.cluster_size = 1  # bm without cs below the split point
     return r

From 78227d5397428f5c53eb0a2e48cfa3c1ded8645f Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 31 Jul 2026 00:51:20 -0700
Subject: [PATCH 079/117] [None][fix] GVR: keep the P1 refine rescue off the
 stock path

p1r_rescue defaulted to True, so the single-operator path compiled it
too: 14016 vs 13570 PTX instructions against upstream for the same
launch (flash 256k B4, both trees pick cs=4 / 512 threads), and 3.5-6%
slower on the flash 128k/256k cells of the captured grid.

Derive it from the assist tiers instead, the way p4_tail_v3 already is.
The stock build is then instruction-identical to upstream (13570 / 14498
lines on both trees) and the tiers - which are the ones seeding from a
possibly zero-init prev_topk - keep the rescue. Exactness unchanged on
counts / rungs / list at flash 256k.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py              | 13 ++++++++-----
 .../scripts/cute_dsl_kernels/top_k/run_gvr_topk.py  |  4 ++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index b409780e2bbc..c511468181f9 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -296,7 +296,7 @@ def __init__(
         p4_exact_tail: Optional[bool] = None,
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
         p4_tail_v3: Optional[bool] = None,
-        p1r_rescue: bool = True,
+        p1r_rescue: Optional[bool] = None,
         num_bins: Optional[int] = None,
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
@@ -760,10 +760,13 @@ def __init__(
             )
         self.p4_tail_v3 = bool(p4_tail_v3)
         # p1r_rescue: rebuild the refine bracket from the row when the
-        # seed bracket is degenerate. ON by default - upstream's identity
-        # shortcut is wrong on real data (every request's first decode
-        # step feeds a zero-init prev_topk). The knob exists to measure
-        # its cost, not to ship it off.
+        # seed bracket is degenerate - the assisted tiers seed from the
+        # previous step's topk, and a request's first decode step feeds a
+        # zero-init prev_topk, so they need it. Defaults to the tiers
+        # only: on the stock path it would add ~450 PTX instructions
+        # (3.3%) for a case that path cannot reach the same way.
+        if p1r_rescue is None:
+            p1r_rescue = bool(use_ext_counts or use_ext_cand or ext_rungs or self_scan)
         self.p1r_rescue = bool(p1r_rescue)
 
     # ------------------------------------------------------------------
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 4880a592da9d..2f93d9738dd4 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -68,7 +68,7 @@ def _compile(
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
-    p1r_rescue: bool = True,
+    p1r_rescue: "bool | None" = None,
     num_bins: "int | None" = None,
     p4_fine_rangetest: "bool | None" = None,
     p4_scat_rangetest: bool = False,
@@ -678,7 +678,7 @@ def gvr_topk_decode(
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
-    p1r_rescue: bool = True,
+    p1r_rescue: "bool | None" = None,
     num_bins: "int | None" = None,
     p4_fine_rangetest: Optional[bool] = None,
     p4_scat_rangetest: bool = False,

From 1e70d74f526ce5cbe1b4dda71fb01d832f895d9e Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 31 Jul 2026 07:10:57 -0700
Subject: [PATCH 080/117] [None][perf] GVR routing: hand the mid-row band back
 to the stock kernel

There is a band of row lengths where the stock kernel splits the row
across a cluster and its split grid still fits one wave. The assist
tiers cannot follow there - splitting a row costs them more than the
scan it saves, with or without the block-skip prefix - so the stock
kernel wins the cell outright and the epilogue should emit nothing.

Measured against the stock kernel over the 154-cell captured grid
(2 models x 7 context lengths x 11 batches, nsys kernel-only, emission
tax included): without the band 15 cells run below stock, worst 0.863
at flash 512k batch 16; with it, 2 cells at 0.99. Against the reference
baseline the grid geomean goes 1.139 -> 1.145 and the worst cell
0.833 -> 0.912. Batch 1-2 keep the list tier inside the band, where it
still beats stock by 1.5x, so the list check runs first.

The larger K is, the less row the tiers need before they pay, hence the
two upper bounds.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  5 ++--
 .../blackwell/top_k/gvr_routing.py            | 29 ++++++++++++++++---
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 1f03b25d454a..160942091d06 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2921,8 +2921,9 @@ def sparse_attn_indexer(
                 # so we cap it at 256 for now and fall back to the CUDA C++
                 # indexer_topk_decode. This limit can be removed if GPU memory
                 # is not a bottleneck.
-                # tier "none" (gvr_routing.ASSIST_MIN_N_COMP): too short for
-                # any assist to pay - fall through to the stock branch. The
+                # tier "none": too short for any assist to pay, or inside
+                # the mid-row band where the stock kernel's split grid wins
+                # (see gvr_routing) - fall through to the stock branch. The
                 # untouched ext state reads as cold start, which its closed
                 # loop already handles.
                 if (self.use_gvr_ext and self._gvr_ext is not None
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index 200db91861d3..514e194f06ea 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -83,6 +83,22 @@
 LIST_EMIT_MIN_N = 16384
 RUNGS_MIN_B = 16
 
+# Mid-row weak band: rows long enough that the stock kernel splits them
+# across a cluster, but short enough (and at a small enough batch) that
+# its split grid still fits one wave. The assist tiers cannot follow -
+# splitting a row costs them more than the scan it saves - so the stock
+# kernel wins outright and the epilogue should emit nothing. Measured
+# against the stock kernel on the captured grid: without this band 15 of
+# 154 cells run below stock (worst 0.86 at flash 512k batch 16); with it,
+# 2 cells at 0.99. Longer K needs less row before the tiers pay, hence
+# the two upper bounds. Batch 1-2 still take the list tier (checked
+# first): there the emitted list beats stock by 1.5x even inside the band.
+ASSIST_WEAK_MIN_N = 49152
+ASSIST_WEAK_MAX_N_SMALL_K = 196608  # k <= ASSIST_WEAK_K
+ASSIST_WEAK_MAX_N_LARGE_K = 98304
+ASSIST_WEAK_K = 512
+ASSIST_WEAK_MAX_B = 32
+
 # rungs-tier block_max pays only at small K: with K=1024 the tight-line
 # pass rate runs too high and the prefix read is pure overhead.
 RUNGS_BM_MAX_K = 512
@@ -122,12 +138,17 @@ def plan_emission(
         # and this holds for the zero-emission rungs tier too - it is
         # the same kernel, so the floor is the same
         return "none"
-    if not have_epilogue:
-        return "rungs"  # closed-loop lines cost nothing to carry
-    if n_comp >= LIST_EMIT_MIN_N and (
-        batch == 1 or batch * n_comp * compress_ratio <= LIST_EMIT_MAX_TOKENS
+    if (
+        have_epilogue
+        and n_comp >= LIST_EMIT_MIN_N
+        and (batch == 1 or batch * n_comp * compress_ratio <= LIST_EMIT_MAX_TOKENS)
     ):
         return "list"
+    weak_max = ASSIST_WEAK_MAX_N_SMALL_K if k <= ASSIST_WEAK_K else ASSIST_WEAK_MAX_N_LARGE_K
+    if batch <= ASSIST_WEAK_MAX_B and ASSIST_WEAK_MIN_N <= n_comp < weak_max:
+        return "none"  # stock's split grid wins this band outright
+    if not have_epilogue:
+        return "rungs"  # closed-loop lines cost nothing to carry
     if batch >= RUNGS_MIN_B:
         return "rungs"  # throughput regime: emitting anything is a loss
     return "counts"

From 0982dc19db3b5800511a986fa727990271a42e22 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 1 Aug 2026 02:49:56 -0700
Subject: [PATCH 081/117] Revert "[None][fix] GVR: keep the P1 refine rescue
 off the stock path"

This reverts commit 78227d5397.

The reverted commit read the rescue's presence on the stock path as an
isolation leak. It is not: `143453c294` added it deliberately as a
correctness fix for that exact path - the first decode step of a request
feeds a zero-init prev_topk, and the old degenerate shortcut answers it
with identity indices [0, K) instead of the top-K. The tests that pin
this (test_cute_dsl_gvr_topk_decode_degenerate_preidx and its cs4
variant) call the stock custom op directly and do not exist on main, so
deriving the knob from the assist tiers turned the fix off exactly where
it is needed and broke them.

The +446 PTX instructions and the 3.5-6% it costs the stock path on the
flash 128k/256k cells are the price of that fix, not leakage from the
emission tiers. The tiers stay isolated the way they already were -
every emission knob is still compile-time gated off the stock build.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py              | 13 +++++--------
 .../scripts/cute_dsl_kernels/top_k/run_gvr_topk.py  |  4 ++--
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index c511468181f9..b409780e2bbc 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -296,7 +296,7 @@ def __init__(
         p4_exact_tail: Optional[bool] = None,
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
         p4_tail_v3: Optional[bool] = None,
-        p1r_rescue: Optional[bool] = None,
+        p1r_rescue: bool = True,
         num_bins: Optional[int] = None,
         p4_warp_redundant: bool = True,
         p2_warp_redundant: bool = True,
@@ -760,13 +760,10 @@ def __init__(
             )
         self.p4_tail_v3 = bool(p4_tail_v3)
         # p1r_rescue: rebuild the refine bracket from the row when the
-        # seed bracket is degenerate - the assisted tiers seed from the
-        # previous step's topk, and a request's first decode step feeds a
-        # zero-init prev_topk, so they need it. Defaults to the tiers
-        # only: on the stock path it would add ~450 PTX instructions
-        # (3.3%) for a case that path cannot reach the same way.
-        if p1r_rescue is None:
-            p1r_rescue = bool(use_ext_counts or use_ext_cand or ext_rungs or self_scan)
+        # seed bracket is degenerate. ON by default - upstream's identity
+        # shortcut is wrong on real data (every request's first decode
+        # step feeds a zero-init prev_topk). The knob exists to measure
+        # its cost, not to ship it off.
         self.p1r_rescue = bool(p1r_rescue)
 
     # ------------------------------------------------------------------
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 2f93d9738dd4..4880a592da9d 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -68,7 +68,7 @@ def _compile(
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
-    p1r_rescue: "bool | None" = None,
+    p1r_rescue: bool = True,
     num_bins: "int | None" = None,
     p4_fine_rangetest: "bool | None" = None,
     p4_scat_rangetest: bool = False,
@@ -678,7 +678,7 @@ def gvr_topk_decode(
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
-    p1r_rescue: "bool | None" = None,
+    p1r_rescue: bool = True,
     num_bins: "int | None" = None,
     p4_fine_rangetest: Optional[bool] = None,
     p4_scat_rangetest: bool = False,

From 15f1c458a08ffba8021a41d714317dcbb3768286 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 2 Aug 2026 20:14:15 -0700
Subject: [PATCH 082/117] [None][perf] GVR routing: refit the tiers on
 kernel-only emission cost

The emission cost was fitted against a WALL-clock indexer time, which
carries a ~22us launch floor, while every top-k number it was compared
against is kernel-only. Re-measuring the indexer kernel itself (nsys
kernel-only, ABBA-interleaved NVTX blocks, batch 1..64 x ctx 32k..512k)
changes the shape of the counts tier's cost completely:

  batch            1     2     4     8    16    64
  counts (us)    +5.6  +4.4  +2.9  +2.0  +0.4   0.0

It is a fixed reduction-latency chain that depends on batch alone and
hides once there are enough rows to overlap it - not a fraction of the
indexer. The old model charged it 3% of the wall indexer time, i.e. 12us
at batch 64 / 512k where the true cost is zero, and that is what put the
counts tier behind RUNGS_MIN_B. On the bare kernels the counts tier beats
rungs on 84 of 98 grid cells, by up to 4x.

Refit on consistent units:

  * counts is the default whenever the epilogue can emit and the step is
    at least half a megatoken (below that its latency chain is exposed);
  * rungs, which emits nothing, is the fallback rather than a
    large-batch special case, and keeps the short-row band it wins;
  * list now needs a long row AND a single request - measured against
    the indexer kernel it almost never repays its emission. An
    unbucketed single-segment list is 1.5-5.8x cheaper to emit, but
    rescoring the grid with it moves the geomean +0.5%, so the contract
    is left alone.

Scored over the 154-cell captured grid by the shipped plan_emission:

                        vs this PR's stock        vs the reference
  before               1.251, worst 0.558        1.169
  after                1.349, worst 1.000        1.260, best 5.40x

Zero cells below stock, and 0.6% off the per-cell optimum (1.357).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_routing.py            | 93 ++++++++++---------
 1 file changed, 48 insertions(+), 45 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index 514e194f06ea..abfdeeb5e2f8 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -40,12 +40,17 @@
 from dataclasses import dataclass
 from typing import Optional
 
-# ---- measured thresholds (B200, f15 grid) --------------------------------
+# ---- measured thresholds (B200) ------------------------------------------
+#
+# UNITS: every threshold here is fitted on KERNEL-ONLY time for both the
+# indexer and the top-k. An earlier fit compared a wall-clock indexer
+# (which carries a ~22us launch floor) against kernel-only top-k times;
+# that inflated the apparent emission budget at small shapes and is what
+# put the counts tier behind a batch gate it never needed.
 
-# The kernel's fixed cost does not shrink with N (~12us at K=1024), so
-# once the stock kernel finishes under that floor no tier can win. The
-# whole n_comp/K < 1.5 band measures below parity.
-ASSIST_MIN_N_COMP = 4096  # ~8k raw context at compress_ratio 4
+# The kernel's fixed cost does not shrink with N, so once the stock
+# kernel finishes under that floor no tier can win.
+ASSIST_MIN_N_COMP = 2048
 
 # Block-skip prefix pays for the counts tier from 65536 up, but not for
 # the zero-emission rungs tier below 131072. Measured on captured V4
@@ -61,43 +66,43 @@
 # at cs 1 / 4 / 8 (batch 4) and cs8 spills to 31.5us at batch 16.
 SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split from here up
 
-# Emission tax measured on the FP4 indexer itself (ctx 256k and 1M,
-# batch 1..128), as the wall delta of the same kernel with the emission
-# outputs attached:
-#
-# What the list tier costs tracks B*N - the total emitted volume - not
-# batch on its own:
+# Emission cost measured on the FP4 indexer KERNEL (nsys kernel-only,
+# ABBA-interleaved NVTX blocks, batch 1..64 x ctx 32k..512k) as the delta
+# of the same kernel with the emission outputs attached:
 #
-#   B*N (raw tokens)   <=0.5M    1M      >=2M
-#   list                +11-14%  +55-71%  +150-270%
-#   counts               +1-3%    +1-3%     +0-3%
+#   batch                1     2     4     8    16    64
+#   counts (us)        +5.6  +4.4  +2.9  +2.0  +0.4   0.0
+#   list, bucketed     +6.2  +6.2  +6.8 +10.5 +25.0  +63.1  (ctx 32k)
+#                     +10.3 +14.7 +25.6 +44.5 +85.7 +417.8  (ctx 512k)
 #
-# Volume is the main term but not the only one: at the same 1M tokens
-# the tax is +21% as 1M x B1 and +71% as 128k x B8, so there is a
-# per-row cost on top. Both measured-cheap regions are covered by
-# "half a megatoken, or a single row". The counts tax stays inside 3%
-# everywhere, which is why it remains the default; past RUNGS_MIN_B
-# even that is not repaid and the zero-emission rungs tier takes over.
-LIST_EMIT_MAX_TOKENS = 786432  # B * raw length; between the measured
-# cheap band (<=0.52M) and the first expensive point (1.05M)
-LIST_EMIT_MIN_N = 16384
-RUNGS_MIN_B = 16
+# The counts cost is a fixed reduction-latency chain that depends on
+# BATCH ONLY and hides once there are enough rows to overlap it - it is
+# not a fraction of the indexer, so it cannot price the tier out at
+# scale. An earlier fit modelled it as 3% of a WALL-clock indexer time
+# (which carries a ~22us launch floor) and so charged 12us at batch 64 /
+# 512k, where the true cost is zero; that is what put the counts tier
+# behind a batch gate. The list cost is real per-emitted-entry work and
+# grows with batch and context alike; an unbucketed single-segment list
+# measures 1.5-5.8x cheaper, but rescoring the grid with it moves the
+# geomean by +0.5%, so the bucketed contract stays as is.
+LIST_EMIT_MIN_N = 65536  # shorter rows: the emission outweighs the saving
+LIST_EMIT_MAX_B = 1  # past one row the list never repays its emission
+COUNTS_MIN_TOKENS = 524288  # B * raw length; below this the counts
+# latency chain is exposed and the zero-emission rungs tier wins
+RUNGS_ONLY_MIN_N = 16384  # short-row band where rungs also beats counts
+RUNGS_ONLY_MAX_N = 49152
 
 # Mid-row weak band: rows long enough that the stock kernel splits them
 # across a cluster, but short enough (and at a small enough batch) that
 # its split grid still fits one wave. The assist tiers cannot follow -
-# splitting a row costs them more than the scan it saves - so the stock
-# kernel wins outright and the epilogue should emit nothing. Measured
-# against the stock kernel on the captured grid: without this band 15 of
-# 154 cells run below stock (worst 0.86 at flash 512k batch 16); with it,
-# 2 cells at 0.99. Longer K needs less row before the tiers pay, hence
-# the two upper bounds. Batch 1-2 still take the list tier (checked
-# first): there the emitted list beats stock by 1.5x even inside the band.
+# splitting a row costs them more than the scan it saves, with or without
+# the block-skip prefix - so the stock kernel wins outright and the
+# epilogue should emit nothing.
 ASSIST_WEAK_MIN_N = 49152
-ASSIST_WEAK_MAX_N_SMALL_K = 196608  # k <= ASSIST_WEAK_K
-ASSIST_WEAK_MAX_N_LARGE_K = 98304
+ASSIST_WEAK_MAX_N_SMALL_K = 262144  # k <= ASSIST_WEAK_K
+ASSIST_WEAK_MAX_N_LARGE_K = 196608
 ASSIST_WEAK_K = 512
-ASSIST_WEAK_MAX_B = 32
+ASSIST_WEAK_MAX_B = 8
 
 # rungs-tier block_max pays only at small K: with K=1024 the tight-line
 # pass rate runs too high and the prefix read is pure overhead.
@@ -138,20 +143,18 @@ def plan_emission(
         # and this holds for the zero-emission rungs tier too - it is
         # the same kernel, so the floor is the same
         return "none"
-    if (
-        have_epilogue
-        and n_comp >= LIST_EMIT_MIN_N
-        and (batch == 1 or batch * n_comp * compress_ratio <= LIST_EMIT_MAX_TOKENS)
-    ):
-        return "list"
     weak_max = ASSIST_WEAK_MAX_N_SMALL_K if k <= ASSIST_WEAK_K else ASSIST_WEAK_MAX_N_LARGE_K
     if batch <= ASSIST_WEAK_MAX_B and ASSIST_WEAK_MIN_N <= n_comp < weak_max:
         return "none"  # stock's split grid wins this band outright
-    if not have_epilogue:
-        return "rungs"  # closed-loop lines cost nothing to carry
-    if batch >= RUNGS_MIN_B:
-        return "rungs"  # throughput regime: emitting anything is a loss
-    return "counts"
+    if have_epilogue and n_comp >= LIST_EMIT_MIN_N and batch <= LIST_EMIT_MAX_B:
+        return "list"
+    if (
+        have_epilogue
+        and batch * n_comp * compress_ratio >= COUNTS_MIN_TOKENS
+        and not (RUNGS_ONLY_MIN_N <= n_comp < RUNGS_ONLY_MAX_N)
+    ):
+        return "counts"
+    return "rungs"  # closed-loop lines cost nothing to carry
 
 
 def pick_config(tier: str, batch: int, n_comp: int, k: int, num_sms: int) -> TopkRoute:

From 445bb7d5044c00c6b3fb13e0a9c219e06bbab9e0 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 2 Aug 2026 21:39:17 -0700
Subject: [PATCH 083/117] [None][perf] GVR: park the tight lines for the list
 tier

The bucketed list emission spends most of its cost on the two tight
segments: they claim their slots with exact per-warp ballots so their
prefixes stay pad-free, while the loosest segment claims through an
amortised per-warp window. Measured on the indexer kernel (nsys
kernel-only, ABBA-interleaved NVTX blocks), the exact claims are 2.7-5x
the window's cost - at batch 16 / ctx 512k, +88.7us against +15.0us.

Park the two tight lines above the score range when the planned tier is
list, so every admitted entry lands in the window-claimed segment. The
consumer needs no change: it reads the segment counts from the control
row, finds both tight segments empty, and takes the loosest line as its
cut - the window's sentinel pads score -inf and never rank. Exact on all
3927 (layer, step, batch) combinations of the captured grid.

That makes the list tier affordable past a single request, so its gate
widens to four, and it now runs ahead of the mid-row weak band - that
band is about the stock kernel out-scanning us, and a list hit never
scans the row.

Scored over the 154-cell grid by the shipped plan_emission, with the
list arm re-measured under the new contract:

                        vs this PR's stock        vs the reference
  before                1.349, worst 1.000        1.260, worst 0.856
  after                 1.387, worst 1.000        1.296, worst 0.900

Still zero cells below stock, 0.5% off the per-cell optimum, and the
cells that had to fall back to the stock kernel drop from 20 to 4.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  2 +-
 .../attention_backend/sparse/gvr_ext.py       | 23 ++++++++++++++++---
 .../blackwell/top_k/gvr_routing.py            | 18 +++++++++------
 3 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 160942091d06..52dc9534f86d 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2842,7 +2842,7 @@ def sparse_attn_indexer(
                             torch.cuda.get_device_properties(
                                 q_fp8.device).multi_processor_count,
                             compress_ratio=max(self.compress_ratio, 1))
-                        st.update_seed_rows(batch_size)
+                        st.update_seed_rows(batch_size, emit_tier)
                         gvr_emit_kwargs = st.indexer_emit_kwargs(
                             emit_tier, batch_size)
                         if self._gvr_route.attach_block_max or emit_tier in (
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
index 296ec4a5ed86..525025eb35cf 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -47,6 +47,18 @@
 GUARD_LO = 2.0
 GUARD_HI = 0.5
 
+# List tier only: park the two tight lines above any score so every
+# admitted entry lands in the loosest segment, which is the one that
+# already claims its slots through a per-warp window instead of an
+# exact ballot. Measured on the indexer kernel (nsys kernel-only): the
+# emission cost drops 2.7-5x (e.g. batch 16 / ctx 512k, +88.7us ->
+# +15.0us) and the top-k side stays at parity, because the consumer
+# reads the segment counts from the control row and finds the two tight
+# segments empty. Any finite value above the score range works; the
+# kernel's eligibility check only needs the three lines to be
+# increasing and the loosest one finite.
+LIST_PARK_LINE = 1.0e30
+
 
 class GvrExtState:
     """Per-attention-backend emission state (persistent buffers)."""
@@ -116,7 +128,7 @@ def plan(
         route = pick_config(self.emitted_tier, batch, n_comp, self.top_k, num_sms)
         return emit_tier, route
 
-    def update_seed_rows(self, num_rows: int) -> None:
+    def update_seed_rows(self, num_rows: int, emit_tier: str = "counts") -> None:
         """Device-side closed-loop line update from the last publish.
 
         Pure tensor ops (graph-capturable). Rows whose xstate is not
@@ -132,8 +144,13 @@ def update_seed_rows(self, num_rows: int) -> None:
         valid = x[:, 0] > 0
         inf = torch.full_like(kth, float("inf"))
         s[:, 0] = torch.where(valid, kth - GUARD_LO * span, inf)
-        s[:, 1] = torch.where(valid, kth - 1e-6, inf)
-        s[:, 2] = torch.where(valid, kth + GUARD_HI * span, inf)
+        if emit_tier == "list":
+            park = torch.full_like(kth, LIST_PARK_LINE)
+            s[:, 1] = torch.where(valid, park, inf)
+            s[:, 2] = torch.where(valid, park + park, inf)
+        else:
+            s[:, 1] = torch.where(valid, kth - 1e-6, inf)
+            s[:, 2] = torch.where(valid, kth + GUARD_HI * span, inf)
         s[:, 3:8] = 0.0
         if self.cand_ctl is not None:
             self.cand_ctl[:num_rows].zero_()
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index abfdeeb5e2f8..f96c4ab2d957 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -82,11 +82,13 @@
 # (which carries a ~22us launch floor) and so charged 12us at batch 64 /
 # 512k, where the true cost is zero; that is what put the counts tier
 # behind a batch gate. The list cost is real per-emitted-entry work and
-# grows with batch and context alike; an unbucketed single-segment list
-# measures 1.5-5.8x cheaper, but rescoring the grid with it moves the
-# geomean by +0.5%, so the bucketed contract stays as is.
+# grows with batch and context alike. Parking the two tight lines above
+# the score range (see gvr_ext.LIST_PARK_LINE) drops it 2.7-5x - every
+# entry then lands in the one segment that claims through a per-warp
+# window instead of an exact ballot - at top-k parity, which is what
+# makes the list tier affordable past a single row.
 LIST_EMIT_MIN_N = 65536  # shorter rows: the emission outweighs the saving
-LIST_EMIT_MAX_B = 1  # past one row the list never repays its emission
+LIST_EMIT_MAX_B = 4  # past four rows the list stops repaying its emission
 COUNTS_MIN_TOKENS = 524288  # B * raw length; below this the counts
 # latency chain is exposed and the zero-emission rungs tier wins
 RUNGS_ONLY_MIN_N = 16384  # short-row band where rungs also beats counts
@@ -99,7 +101,7 @@
 # the block-skip prefix - so the stock kernel wins outright and the
 # epilogue should emit nothing.
 ASSIST_WEAK_MIN_N = 49152
-ASSIST_WEAK_MAX_N_SMALL_K = 262144  # k <= ASSIST_WEAK_K
+ASSIST_WEAK_MAX_N_SMALL_K = 196608  # k <= ASSIST_WEAK_K
 ASSIST_WEAK_MAX_N_LARGE_K = 196608
 ASSIST_WEAK_K = 512
 ASSIST_WEAK_MAX_B = 8
@@ -143,11 +145,13 @@ def plan_emission(
         # and this holds for the zero-emission rungs tier too - it is
         # the same kernel, so the floor is the same
         return "none"
+    if have_epilogue and n_comp >= LIST_EMIT_MIN_N and batch <= LIST_EMIT_MAX_B:
+        # checked before the weak band below: that band is about the stock
+        # kernel out-scanning us, and a list hit never scans the row
+        return "list"
     weak_max = ASSIST_WEAK_MAX_N_SMALL_K if k <= ASSIST_WEAK_K else ASSIST_WEAK_MAX_N_LARGE_K
     if batch <= ASSIST_WEAK_MAX_B and ASSIST_WEAK_MIN_N <= n_comp < weak_max:
         return "none"  # stock's split grid wins this band outright
-    if have_epilogue and n_comp >= LIST_EMIT_MIN_N and batch <= LIST_EMIT_MAX_B:
-        return "list"
     if (
         have_epilogue
         and batch * n_comp * compress_ratio >= COUNTS_MIN_TOKENS

From f8e687158a4e92822c73370e88bb8a702476860d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 3 Aug 2026 00:34:33 -0700
Subject: [PATCH 084/117] [None][fix] GVR: stop dividing the routing length by
 the compress ratio twice

indexer_max_seq_len is already the compressed length -
get_indexer_max_seq_len divides by the compress ratio, and the same value
goes to the top-k as max_seq_len a few lines below. The emission planner
was handed that value divided by the ratio a second time, so every
routing threshold acted on a length two doublings too small.

What production actually ran, against what the thresholds were fitted
for, over the 154-cell captured grid:

                 rungs  counts  list  none
  as shipped        74      50     0    30
  as intended       62      70    18     4

The list tier was unreachable and 30 shapes fell back to the stock
kernel instead of 4. Rescoring the grid with the measured arm costs:

                 vs this PR's stock        vs the reference
  before          1.221, worst 0.788       1.141, worst 0.766
  after           1.387, worst 1.000       1.296, worst 0.900

The two ratios are the same value for every layer that reaches this
path, and where the module's ratio is 1 the old and new expressions are
identical, so this only changes the compressed-KV layers - which are the
ones it was wrong for.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 52dc9534f86d..6f4d679f9306 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2834,8 +2834,14 @@ def sparse_attn_indexer(
                                 top_k=self.index_topk,
                                 device=q_fp8.device)
                         st = self._gvr_ext
-                        n_comp = indexer_max_seq_len // max(
-                            self.compress_ratio, 1)
+                        # indexer_max_seq_len is ALREADY the compressed
+                        # length - get_indexer_max_seq_len divides by the
+                        # compress ratio, and it is what goes to the top-k
+                        # as max_seq_len below. Dividing again shifted every
+                        # routing threshold by two doublings: the list tier
+                        # became unreachable and 30 of the 154 grid shapes
+                        # fell back to the stock kernel instead of 4.
+                        n_comp = indexer_max_seq_len
                         emit_tier, self._gvr_route = st.plan(
                             batch_size,
                             n_comp,

From 428c25c152d4e3d6bbeacd1575732da3d473f8ec Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 3 Aug 2026 01:08:21 -0700
Subject: [PATCH 085/117] [None][perf] GVR: resolve P4's coarse bin search
 lane-parallel

The fused rank-and-scatter path looks for the straddling histogram bin
with a serial walk: every thread re-sums its warp's whole bin slice,
thread 0 alone scans the warp totals, then one lane of the target warp
walks that slice bin by bin, with three publish barriers in between.
Measured, that search costs 1.04us and does not move when the candidate
count changes by 3.6x - it is a dependent-latency chain, not work.

The same search already had a lane-parallel form in this file
(_kth_bin_search_rw, used by the snap path): stage the slice sums once,
then let every warp resolve the answer from them with an idx-shuffle
scan and a ballot. Give the rank-scatter path the same treatment.
Integer sums are associative, so every warp lands on the same answer
and no leader or publish barrier is needed; the serial form stays as the
compile-time alternative under p4_warp_redundant.

Phase 4, counts tier, flash 256k batch 4, mean over 5 layers x all
decode steps: 4.36us -> 3.76us, with the coarse search 1.04 -> 0.63 and
the fine search 1.32 -> 1.13 (it now starts from registers instead of a
shared-memory round trip).

Whole kernel, 27 (tier, shape, batch) cells across list / counts / rungs:
faster on 27 of 27, mean +1.47us, median +1.51us - the list tier goes
8.01 -> 5.90us at flash 256k batch 4. Exact on every cell.

This one does touch the single-operator path, which also runs this
search: it measures +1.07us faster there and 848 PTX instructions
smaller, with bit-identical output.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 227 ++++++++++++++----
 1 file changed, 174 insertions(+), 53 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index b409780e2bbc..9ba4c7cdb199 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3492,6 +3492,118 @@ def _kth_bin_search_rw(self, smem_hist, smem_wcnt, lo, binw, tidx, warp_id, lane
             sel_out = cute.arch.shuffle_sync(sel_loc, src)
         return thr_out, sel_out
 
+    # ------------------------------------------------------------------
+    # _p4_coarse_rw - redundant-warp coarse bin search for the fused
+    # rank-and-scatter path. Same result as the high->low walk it
+    # replaces (the straddling bin and the count strictly above it),
+    # but staged once and then resolved lane-parallel on every warp:
+    # an idx-shuffle scan + ballot locate the target slice, a second
+    # scan + the unique crossing test locate the bin inside it. Two
+    # publish barriers and a bins_per_warp-deep serial LDS+IADD chain in
+    # a single lane disappear; integer sums are associative so every
+    # warp lands on the same answer bit-for-bit. Mirrors
+    # _kth_bin_search_rw, which does the same for the snap path.
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _p4_coarse_rw(self, smem_hist, smem_wcnt, warp_id, lane):
+        kK = cutlass.const_expr(self.top_k)
+        kBins = cutlass.const_expr(self.kNumBins)
+        bins_per_warp = cutlass.const_expr(kBins // self.num_warps)
+
+        warp_bin_sum = cutlass.Int32(0)
+        if cutlass.const_expr(bins_per_warp % self.WARP_SIZE == 0):
+            for jm in cutlass.range_constexpr(bins_per_warp // self.WARP_SIZE):
+                bidx_s = (
+                    cutlass.Int32(kBins - 1)
+                    - warp_id * cutlass.Int32(bins_per_warp)
+                    - (lane + cutlass.Int32(jm * self.WARP_SIZE))
+                )
+                warp_bin_sum = warp_bin_sum + smem_hist[bidx_s]
+            warp_bin_sum = self.warp_reduce_sum_i32(warp_bin_sum)
+        else:
+            for jb in cutlass.range_constexpr(bins_per_warp):
+                bidx_s = (
+                    cutlass.Int32(kBins - 1)
+                    - warp_id * cutlass.Int32(bins_per_warp)
+                    - cutlass.Int32(jb)
+                )
+                warp_bin_sum = warp_bin_sum + smem_hist[bidx_s]
+        if lane == cutlass.Int32(0):
+            smem_wcnt[warp_id] = warp_bin_sum
+        cute.arch.barrier()
+
+        # locate the target slice (lane w holds slot w)
+        v_s = cutlass.Int32(0)
+        if lane < cutlass.Int32(self.num_warps):
+            v_s = smem_wcnt[lane]
+        run2 = v_s
+        for d2 in cutlass.range_constexpr(5):
+            off2 = cutlass.const_expr(1 << d2)
+            src2 = lane - cutlass.Int32(off2)
+            if src2 < cutlass.Int32(0):
+                src2 = cutlass.Int32(0)
+            up2 = cute.arch.shuffle_sync(run2, src2)
+            if lane >= cutlass.Int32(off2):
+                run2 = run2 + up2
+        m2 = cute.arch.vote_ballot_sync(run2 >= cutlass.Int32(kK))
+        tw = cutlass.Int32(self.num_warps - 1)
+        if m2 != cutlass.Uint32(0):
+            low2 = m2 & (cutlass.Uint32(0) - m2)
+            tw = cutlass.Int32(cute.arch.popc(low2 - cutlass.Uint32(1)))
+        incl_tw = cute.arch.shuffle_sync(run2, tw)
+        slot_tw = cute.arch.shuffle_sync(v_s, tw)
+        prefix = incl_tw - slot_tw
+
+        # locate the bin inside the target slice
+        ppl = cutlass.const_expr((bins_per_warp + self.WARP_SIZE - 1) // self.WARP_SIZE)
+        cnt_frag = cute.make_fragment((ppl,), cutlass.Int32)
+        my_sum = cutlass.Int32(0)
+        for j3 in cutlass.range_constexpr(ppl):
+            pos = lane * cutlass.Int32(ppl) + cutlass.Int32(j3)
+            cnt_j = cutlass.Int32(0)
+            if pos < cutlass.Int32(bins_per_warp):
+                bidx3 = cutlass.Int32(kBins - 1) - tw * cutlass.Int32(bins_per_warp) - pos
+                cnt_j = smem_hist[bidx3]
+            cnt_frag[j3] = cnt_j
+            my_sum = my_sum + cnt_j
+        run3 = my_sum
+        for d3 in cutlass.range_constexpr(5):
+            off3 = cutlass.const_expr(1 << d3)
+            src3 = lane - cutlass.Int32(off3)
+            if src3 < cutlass.Int32(0):
+                src3 = cutlass.Int32(0)
+            up3 = cute.arch.shuffle_sync(run3, src3)
+            if lane >= cutlass.Int32(off3):
+                run3 = run3 + up3
+        base3 = prefix + (run3 - my_sum)
+
+        b_loc = cutlass.Int32(kBins - 1)
+        ra_loc = prefix
+        hit = cutlass.Int32(0)
+        r3 = base3
+        for j4 in cutlass.range_constexpr(ppl):
+            pos4 = lane * cutlass.Int32(ppl) + cutlass.Int32(j4)
+            cnt4 = cnt_frag[j4]
+            if (
+                pos4 < cutlass.Int32(bins_per_warp)
+                and r3 < cutlass.Int32(kK)
+                and r3 + cnt4 >= cutlass.Int32(kK)
+                and hit == cutlass.Int32(0)
+            ):
+                b_loc = cutlass.Int32(kBins - 1) - tw * cutlass.Int32(bins_per_warp) - pos4
+                ra_loc = r3
+                hit = cutlass.Int32(1)
+            r3 = r3 + cnt4
+        mask3 = cute.arch.vote_ballot_sync(hit != cutlass.Int32(0))
+        b_out = cutlass.Int32(kBins - 1)
+        ra_out = prefix
+        if mask3 != cutlass.Uint32(0):
+            low = mask3 & (cutlass.Uint32(0) - mask3)
+            src = cutlass.Int32(cute.arch.popc(low - cutlass.Uint32(1)))
+            b_out = cute.arch.shuffle_sync(b_loc, src)
+            ra_out = cute.arch.shuffle_sync(ra_loc, src)
+        return b_out, ra_out
+
     # ------------------------------------------------------------------
     # Phase 4 (alt): op#7 fused rank-and-scatter (enable_p4_rank_scatter).
     # Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py.
@@ -3610,61 +3722,70 @@ def phase4_rank_scatter(
             cute.arch.barrier()
             if cutlass.const_expr(_P4_SUB_DBG):
                 sc2 = cute.arch.clock64()
-            # ---- 3-step high→low bin search → straddling bin b* + rank_above ----
-            warp_bin_sum = cutlass.Int32(0)
-            for jb in cutlass.range_constexpr(bins_per_warp):
-                bidx_s = (
-                    cutlass.Int32(kBins - 1)
-                    - warp_id * cutlass.Int32(bins_per_warp)
-                    - cutlass.Int32(jb)
-                )
-                warp_bin_sum = warp_bin_sum + smem_hist[bidx_s]
-            if lane == cutlass.Int32(0):
-                smem_wcnt[warp_id] = warp_bin_sum
-            cute.arch.barrier()
-            if tidx == cutlass.Int32(0):
-                cum = cutlass.Int32(0)
-                tw = cutlass.Int32(num_warps - 1)
-                found = cutlass.Int32(0)
-                for w2 in cutlass.range_constexpr(self.num_warps):
-                    cum = cum + smem_wcnt[w2]
-                    if cum >= cutlass.Int32(kK) and found == cutlass.Int32(0):
-                        tw = cutlass.Int32(w2)
-                        found = cutlass.Int32(1)
-                cum2 = cutlass.Int32(0)
-                for w3 in cutlass.range_constexpr(self.num_warps):
-                    if cutlass.Int32(w3) < tw:
-                        cum2 = cum2 + smem_wcnt[w3]
-                s_iscalars[2] = cum2  # prefix-count before target warp
-                s_iscalars[3] = tw
-            cute.arch.barrier()
-            target_warp = s_iscalars[3]
-            if warp_id == target_warp and lane == cutlass.Int32(0):
-                base_cum = s_iscalars[2]
-                b_star = cutlass.Int32(kBins - 1)
-                rank_above = base_cum
-                set_d = cutlass.Int32(0)
-                for jb2 in cutlass.range_constexpr(bins_per_warp):
-                    bidx2 = (
+            # ---- high→low bin search → straddling bin b* + rank_above ----
+            if cutlass.const_expr(self.p4_warp_redundant):
+                b_star, rank_above = self._p4_coarse_rw(smem_hist, smem_wcnt, warp_id, lane)
+                if tidx == cutlass.Int32(0):
+                    s_iscalars[4] = cutlass.Int32(0)  # cnt_above
+                    s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
+                cute.arch.barrier()
+                if cutlass.const_expr(_P4_SUB_DBG):
+                    sc3 = cute.arch.clock64()
+            else:
+                warp_bin_sum = cutlass.Int32(0)
+                for jb in cutlass.range_constexpr(bins_per_warp):
+                    bidx_s = (
                         cutlass.Int32(kBins - 1)
-                        - target_warp * cutlass.Int32(bins_per_warp)
-                        - cutlass.Int32(jb2)
+                        - warp_id * cutlass.Int32(bins_per_warp)
+                        - cutlass.Int32(jb)
                     )
-                    ra_before = base_cum
-                    base_cum = base_cum + smem_hist[bidx2]
-                    if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0):
-                        b_star = bidx2
-                        rank_above = ra_before  # count in bins strictly above b*
-                        set_d = cutlass.Int32(1)
-                s_iscalars[2] = rank_above
-                s_iscalars[3] = b_star
-                s_iscalars[4] = cutlass.Int32(0)  # cnt_above
-                s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
-            cute.arch.barrier()
-            if cutlass.const_expr(_P4_SUB_DBG):
-                sc3 = cute.arch.clock64()
-            b_star = s_iscalars[3]
-            rank_above = s_iscalars[2]
+                    warp_bin_sum = warp_bin_sum + smem_hist[bidx_s]
+                if lane == cutlass.Int32(0):
+                    smem_wcnt[warp_id] = warp_bin_sum
+                cute.arch.barrier()
+                if tidx == cutlass.Int32(0):
+                    cum = cutlass.Int32(0)
+                    tw = cutlass.Int32(num_warps - 1)
+                    found = cutlass.Int32(0)
+                    for w2 in cutlass.range_constexpr(self.num_warps):
+                        cum = cum + smem_wcnt[w2]
+                        if cum >= cutlass.Int32(kK) and found == cutlass.Int32(0):
+                            tw = cutlass.Int32(w2)
+                            found = cutlass.Int32(1)
+                    cum2 = cutlass.Int32(0)
+                    for w3 in cutlass.range_constexpr(self.num_warps):
+                        if cutlass.Int32(w3) < tw:
+                            cum2 = cum2 + smem_wcnt[w3]
+                    s_iscalars[2] = cum2  # prefix-count before target warp
+                    s_iscalars[3] = tw
+                cute.arch.barrier()
+                target_warp = s_iscalars[3]
+                if warp_id == target_warp and lane == cutlass.Int32(0):
+                    base_cum = s_iscalars[2]
+                    b_star_s = cutlass.Int32(kBins - 1)
+                    rank_above_s = base_cum
+                    set_d = cutlass.Int32(0)
+                    for jb2 in cutlass.range_constexpr(bins_per_warp):
+                        bidx2 = (
+                            cutlass.Int32(kBins - 1)
+                            - target_warp * cutlass.Int32(bins_per_warp)
+                            - cutlass.Int32(jb2)
+                        )
+                        ra_before = base_cum
+                        base_cum = base_cum + smem_hist[bidx2]
+                        if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0):
+                            b_star_s = bidx2
+                            rank_above_s = ra_before  # count in bins strictly above b*
+                            set_d = cutlass.Int32(1)
+                    s_iscalars[2] = rank_above_s
+                    s_iscalars[3] = b_star_s
+                    s_iscalars[4] = cutlass.Int32(0)  # cnt_above
+                    s_iscalars[1] = cutlass.Int32(0)  # cnt_straddle
+                cute.arch.barrier()
+                if cutlass.const_expr(_P4_SUB_DBG):
+                    sc3 = cute.arch.clock64()
+                b_star = s_iscalars[3]
+                rank_above = s_iscalars[2]
 
             # ---- EXACT: one fine-histogram recursion on the straddling bin b* ----
             if cutlass.const_expr(self.enable_p4_rank_scatter_exact):

From ed8ea661a999b100b610f575f0bf241775cb7088 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 3 Aug 2026 02:14:08 -0700
Subject: [PATCH 086/117] [None][fix] GVR: pass the stream where the FP4
 indexer's __call__ expects it

The kernel's __call__ takes the stream right after batch_size, with the
emission tensors appended after it - upstream's positional order with our
slots on the end. The compile-time argument list still used the older
order and put the fake stream after hit_bitmap, so every argument from
block_max onwards landed one slot early:

  DSLRuntimeError: expects argument #15 (hit_bitmap) to be Tensor,
                   but got _FakeStream

This takes down the whole of
tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
at compile time - 1066 of 1066 cases, reproduced on B200 - and with it
the unittest/_torch/attention group. It is why that file has been in
every CI failure list since 07-31; the two runs before this one also hit
a harness fault, which is what the failure analysis latched onto.

The runtime call a few lines below was already right: it drops the
stream (the TVM FFI env stream is used) and is otherwise the same
sequence. Only the compile list was stale. The FP8 runner already passes
the stream in the correct slot and the other compile sites pass it by
keyword, so this was the only one.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index fc58140dc55c..ac543c7c28c3 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -9097,10 +9097,15 @@ def _compile(cls,
                 sm_fake,
                 cutlass.Int32(1),
                 cutlass.Int32(1),
+                # stream sits before the emission tensors in __call__ -
+                # upstream's positional order, with the emission slots
+                # appended after it. Keep this list in that order: the
+                # runtime call below drops the stream (the TVM FFI env
+                # stream is used) and is otherwise the same sequence.
+                fake_stream,
                 block_max_fake,
                 hit_stats_fake,
                 hit_bitmap_fake,
-                fake_stream,
                 seed_thr=seed_thr_fake,
                 seed_counts=seed_counts_fake,
                 cand=cand_fake,

From 8c147a724507547adf83375728555c854fa681f4 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 3 Aug 2026 20:27:38 -0700
Subject: [PATCH 087/117] [None][perf] GVR: give K=512 the compacted
 boundary-class repair

p4_tail_fast was gated on top_k >= 1024, so V4-Flash never took the
compacted path. Its boundary-class repair instead ran the original
4-level MSB radix over the WHOLE candidate array: every level re-zeroes
256 histogram bins and re-walks all the candidates, for a class that is
usually a few hundred entries. The compacted path pays for the class.

That repair only fires when the tie set in the straddling fine bin
overfills the remaining slots - about 4% of decode steps - but when it
fires it doubles the kernel, and it is the single largest source of
worst-case regression. Pairing every step against the reference kernel
found 74 (model, context, layer, step) triggers, 53 of them K=512.

Reproduced on the captured rows (flash 256k layer 20 step 4, batch 4):
the repair costs 5.04us against 0.17us on the neighbouring steps, and
2.17us with the compacted path enabled.

Over the four trigger-heavy flash units (8k / 32k / 64k / 256k, 11781
paired steps per tier):

                        rungs            counts
  worst step        0.478 -> 0.612   0.544 -> 0.629
  burst steps        284 -> 44        282 -> 35
  steps below par   3331 -> 2092       511 -> 594
  geomean           1.063 -> 1.089   1.185 -> 1.193

Exact on all 1848 (layer, step, batch) checks.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py                 | 10 +++++++++-
 tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py   |  8 ++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 9ba4c7cdb199..28eb73981783 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -748,7 +748,15 @@ def __init__(
         # fire census (pro/512k bench + 9 per-layer fixture cells) contains
         # NO K512 cell — so K512 keeps the original byte-identical kernel.
         if p4_tail_fast is None:  # [p4tt]
-            p4_tail_fast = self.p4_exact_tail and top_k >= 1024
+            # Was gated on top_k >= 1024, which left K=512 on the original
+            # full-candidate 4-level radix: every level re-zeroes 256 bins
+            # and re-walks the WHOLE candidate array, so a step whose
+            # boundary class is large costs ~5us against ~0.2us for its
+            # neighbours. The compacted path pays for the class only.
+            # Measured on the captured flash rows that trigger it
+            # (256k layer 20 step 4, K=512): 5.04us -> 2.17us, with the
+            # neighbouring steps unchanged at ~0.2us.
+            p4_tail_fast = self.p4_exact_tail
         self.p4_tail_fast = bool(p4_tail_fast) and self.p4_exact_tail  # [p4tt]
         # p4_tail_v3: compacted-class repair (block-parallel radix +
         # pure-tie pre-check) in place of the stock thread0 serial
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 4880a592da9d..81b621279fec 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -68,6 +68,8 @@ def _compile(
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
+    p4_exact_tail: "bool | None" = None,
+    p4_tail_fast: "bool | None" = None,
     p1r_rescue: bool = True,
     num_bins: "int | None" = None,
     p4_fine_rangetest: "bool | None" = None,
@@ -211,6 +213,8 @@ def _compile(
         enable_block_skip=enable_block_skip,
         pdl_wait_late=pdl_wait_late,
         p4_tail_v3=p4_tail_v3,
+        p4_exact_tail=p4_exact_tail,
+        p4_tail_fast=p4_tail_fast,
         p1r_rescue=p1r_rescue,
         num_bins=num_bins,
         p4_fine_rangetest=p4_fine_rangetest,
@@ -678,6 +682,8 @@ def gvr_topk_decode(
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
+    p4_exact_tail: "bool | None" = None,
+    p4_tail_fast: "bool | None" = None,
     p1r_rescue: bool = True,
     num_bins: "int | None" = None,
     p4_fine_rangetest: Optional[bool] = None,
@@ -970,6 +976,8 @@ def gvr_topk_decode(
         enable_block_skip,
         pdl_wait_late,
         p4_tail_v3,
+        p4_exact_tail,
+        p4_tail_fast,
         p1r_rescue,
         num_bins,
         p4_fine_rangetest,

From d1c4cd553d2c8734e9be1f4e516112b5597cb22a Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 3 Aug 2026 21:16:50 -0700
Subject: [PATCH 088/117] [None][perf] GVR: one candidate pass and a
 lane-parallel fold in the tail

The boundary-class repair walked the candidate array twice: once for the
pure-tie min/max over the class, once to compact the class into shared
memory. Both passes select on the same predicate and nothing writes the
candidate array in between, so the first pass can buffer the members it
already reads and the second walk - and its barrier - go away. The
staging barrier that publishes the per-warp min/max already orders those
reads against the compact writes.

The cross-warp fold that follows had every thread walk all num_warps
slots of two shared arrays with dependent reads. Lane w holds slot w and
one warp reduce settles it, the same shape the coarse bin search already
uses. Integer min/max are associative, so every warp still lands on the
same answer and the fold stays leaderless.

Trigger step (flash 256k layer 20 step 4, batch 4), tail repair:
  two passes + serial fold   2.33us   (pure-tie 0.96 | compact 0.66)
  one pass  + serial fold    2.05us   (pure-tie 1.04 | compact 0.33)
  one pass  + parallel fold  1.54us   (pure-tie 0.47 | compact 0.37)
against 0.34-0.42us on the neighbouring steps.

Over the four trigger-heavy flash units (11781 paired steps per tier),
measured against the reference kernel and counting the previous commit:

                        rungs            counts
  worst step        0.478 -> 0.610   0.544 -> 0.682
  burst steps        284 -> 12        282 -> 34
  steps below par   3331 -> 2226       511 -> 459

Exact on all 1848 checks.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 80 ++++++++-----------
 1 file changed, 33 insertions(+), 47 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 28eb73981783..fa81e16c9431 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -4045,6 +4045,9 @@ def phase4_rank_scatter(
                             # Staging mirrors the head min/max (wcnt + hist
                             # slots [0..31], both dead here; pairs live at
                             # 260+).
+                            if tidx == cutlass.Int32(0):
+                                s_iscalars[0] = cutlass.Int32(0)
+                            nh7 = cutlass.Int32(0)
                             kmn6 = cutlass.Int32(2147483647)
                             kmx6 = cutlass.Int32(-2147483648)
                             it6 = tidx
@@ -4067,6 +4070,14 @@ def phase4_rank_scatter(
                                             kmn6 = k6
                                         if k6 > kmx6:
                                             kmx6 = k6
+                                        # same predicate as the compaction
+                                        # pass used to re-derive: buffer the
+                                        # member here so that pass can go
+                                        for sl7 in cutlass.range_constexpr(nbuf7):
+                                            if cutlass.Int32(sl7) == nh7:
+                                                rv7[sl7] = v6
+                                                ri7[sl7] = smem_vals[it6]
+                                        nh7 = nh7 + cutlass.Int32(1)
                                 it6 = it6 + cutlass.Int32(num_threads)
                             kmn6 = cute.arch.warp_redux_sync(kmn6, "min")
                             kmx6 = cute.arch.warp_redux_sync(kmx6, "max")
@@ -4074,56 +4085,31 @@ def phase4_rank_scatter(
                                 smem_wcnt[warp_id] = kmn6
                                 smem_hist[warp_id] = kmx6
                             cute.arch.barrier()
-                            kmn7 = cutlass.Int32(2147483647)
-                            kmx7 = cutlass.Int32(-2147483648)
-                            for w8 in cutlass.range_constexpr(self.num_warps):
-                                pa8 = smem_wcnt[w8]
-                                pb8 = smem_hist[w8]
-                                if pa8 < kmn7:
-                                    kmn7 = pa8
-                                if pb8 > kmx7:
-                                    kmx7 = pb8
+                            # lane-parallel cross-warp fold: lane w holds
+                            # slot w and one warp reduce settles it, instead
+                            # of every thread walking all num_warps slots of
+                            # two arrays with dependent SMEM reads. Same
+                            # inputs in the same order on every warp, so the
+                            # result stays bit-identical and leaderless.
+                            pa8 = cutlass.Int32(2147483647)
+                            pb8 = cutlass.Int32(-2147483648)
+                            if lane < cutlass.Int32(self.num_warps):
+                                pa8 = smem_wcnt[lane]
+                                pb8 = smem_hist[lane]
+                            kmn7 = cute.arch.warp_redux_sync(pa8, "min")
+                            kmx7 = cute.arch.warp_redux_sync(pb8, "max")
                             if kmn7 == kmx7:
                                 fast_done = cutlass.Int32(1)
                             if fast_done == cutlass.Int32(0):
-                                # mixed class: compact it IN PLACE
-                                # into smem_keys/vals[0..cnt_strad) with a
-                                # register-buffered two-phase pass (every
-                                # thread reads its strided candidates first,
-                                # ONE barrier, then claimed compact writes —
-                                # no read/write overlap by construction). The
-                                # candidate array has no readers after the
-                                # tail, and compaction makes the repair cost a
-                                # function of the CLASS size only, for ANY
-                                # class size up to cand_count (the old full-
-                                # candidate radix fallback is gone).
-                                if tidx == cutlass.Int32(0):
-                                    s_iscalars[0] = cutlass.Int32(0)
-                                nh7 = cutlass.Int32(0)
-                                it7 = tidx
-                                while it7 < cand_count:
-                                    v7 = smem_keys[it7]
-                                    b7 = cutlass.Int32((v7 - bmin_r) * inv1)
-                                    if b7 < cutlass.Int32(0):
-                                        b7 = cutlass.Int32(0)
-                                    if b7 > cutlass.Int32(kBins - 1):
-                                        b7 = cutlass.Int32(kBins - 1)
-                                    if b7 == b_star:
-                                        s7 = cutlass.Int32((v7 - f_lo) * finv)
-                                        if s7 < cutlass.Int32(0):
-                                            s7 = cutlass.Int32(0)
-                                        if s7 > cutlass.Int32(fbins - 1):
-                                            s7 = cutlass.Int32(fbins - 1)
-                                        if s7 == sb_star:
-                                            # static predicated fragment write
-                                            # (dodges dynamic register indexing)
-                                            for sl7 in cutlass.range_constexpr(nbuf7):
-                                                if cutlass.Int32(sl7) == nh7:
-                                                    rv7[sl7] = v7
-                                                    ri7[sl7] = smem_vals[it7]
-                                            nh7 = nh7 + cutlass.Int32(1)
-                                    it7 = it7 + cutlass.Int32(num_threads)
-                                cute.arch.barrier()
+                                # mixed class: compact the members buffered by
+                                # the pure-tie pass above into
+                                # smem_keys/vals[0..cnt_strad). The buffering
+                                # pass already read every candidate it needs,
+                                # and the staging barrier above orders those
+                                # reads before these writes, so the second
+                                # full-candidate walk (and its barrier) is
+                                # gone. Repair cost is a function of the CLASS
+                                # size only, for any class size.
                                 # warp-aggregated claim: intra-warp exclusive
                                 # prefix via shfl scan + ONE atomic per warp
                                 # (a thousand same-address claims serialize

From 97639c276d9652a951e171fed7ee312f8746a47d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 3 Aug 2026 21:38:24 -0700
Subject: [PATCH 089/117] [None][perf] GVR routing: narrow the weak band after
 the phase-4 work

The band was fitted before phase 4 got its coarse bin search resolved
lane-parallel and its boundary-class repair compacted. Those two gained
the assist tiers about 1.5us, which is enough to take the upper half of
the band back off the stock kernel.

Rescored over the 154-cell grid with the post-fix arm times, by the
shipped plan_emission: geomean against this PR's stock 1.370 -> 1.372,
against the reference 1.364 -> 1.365, still zero cells below stock. The
cells that fall through to the stock kernel drop from 4 to 2, and the
counts tier picks them up (72 cells, was 70). The routing is 0.5% off
the per-cell optimum given current arm costs.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/blackwell/top_k/gvr_routing.py      | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index f96c4ab2d957..bf7f72eb64b2 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -99,10 +99,13 @@
 # its split grid still fits one wave. The assist tiers cannot follow -
 # splitting a row costs them more than the scan it saves, with or without
 # the block-skip prefix - so the stock kernel wins outright and the
-# epilogue should emit nothing.
+# epilogue should emit nothing. Narrowed after phase 4 got its coarse
+# search and its boundary-class repair back: the tiers gained about
+# 1.5us there, which is enough to take the upper half of the band back
+# off the stock kernel (2 cells fall through now, down from 4).
 ASSIST_WEAK_MIN_N = 49152
-ASSIST_WEAK_MAX_N_SMALL_K = 196608  # k <= ASSIST_WEAK_K
-ASSIST_WEAK_MAX_N_LARGE_K = 196608
+ASSIST_WEAK_MAX_N_SMALL_K = 98304  # k <= ASSIST_WEAK_K
+ASSIST_WEAK_MAX_N_LARGE_K = 98304
 ASSIST_WEAK_K = 512
 ASSIST_WEAK_MAX_B = 8
 

From 2eb7d92781f68cfe17bee9485a9d13cc9782aaaf Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Mon, 3 Aug 2026 23:00:43 -0700
Subject: [PATCH 090/117] [None][perf] GVR: drain the last two serial folds out
 of phase 4

Phase 4 still had two places where every thread walked all num_warps
slots of a shared array with dependent LDS reads, or where one lane
walked a whole bin slice while the other 1023 threads waited on a
barrier. Both now use the redundant-warp form already proven on the
coarse search: stage the per-warp partials once, then let every warp
resolve the answer with a shuffle scan and a ballot. Integer sums and
min/max both reassociate, so all warps land on the same value, the
result is bit-identical, and the publish barriers go away.

  * the min/max cross-warp fold, measured on flash 256k layer 20
    batch 4: 0.87us -> 0.31us per step. This runs on every step of
    every tier except list (which skips it via ext_range_flag), and
    is where the whole gain below comes from.
  * the fine sub-bin search (_p4_fine_rw). On its own this measured
    flat - that window is dominated by the histogram build pass over
    all 1024 candidates, not by the search - but it removes two
    barriers and the fbins/num_warps-deep single-lane chain, and
    hands sb_star/rank_above_fine back in registers instead of
    parking them in smem_hist[2]/[3].

Same node, same session, three trees on the va tier (kernel-only,
nsys, us/step):

  cell                pre-P4-work    committed    this tree
  flash 256k  B4         13.35         12.91        11.88
  flash 1024k B4         14.74         13.68        12.78
  flash 8k    B4          6.93          6.62         6.57

Exactness: 2409/2409 steps over flash 8k/64k/256k/1024k and pro
1024k, both assist tiers.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 186 +++++++++++-------
 1 file changed, 120 insertions(+), 66 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index fa81e16c9431..5381dece8651 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3612,6 +3612,105 @@ def _p4_coarse_rw(self, smem_hist, smem_wcnt, warp_id, lane):
             ra_out = cute.arch.shuffle_sync(ra_loc, src)
         return b_out, ra_out
 
+    # ------------------------------------------------------------------
+    # _p4_fine_rw - redundant-warp variant of the fine sub-bin search,
+    # the same transformation _p4_coarse_rw applies one level up. The
+    # serial form stages per-warp slice sums, has thread 0 walk the warp
+    # totals, then has one lane of the target warp walk that slice bin by
+    # bin, with three publish barriers. Here every warp resolves it from
+    # the staged sums with an idx-shuffle scan and a ballot, so two
+    # barriers and the fbins/num_warps-deep serial LDS chain disappear.
+    # Integer sums are associative: every warp lands on the same answer.
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _p4_fine_rw(self, smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane):
+        kK = cutlass.const_expr(self.top_k)
+        fbpw = cutlass.const_expr(fbins // self.num_warps)
+
+        ws = cutlass.Int32(0)
+        if cutlass.const_expr(fbpw <= self.WARP_SIZE):
+            if lane < cutlass.Int32(fbpw):
+                bif = cutlass.Int32(fbins - 1) - warp_id * cutlass.Int32(fbpw) - lane
+                ws = smem_hist[bif]
+            ws = self.warp_reduce_sum_i32(ws)
+        else:
+            for jm in cutlass.range_constexpr(fbpw):
+                bif = cutlass.Int32(fbins - 1) - warp_id * cutlass.Int32(fbpw) - cutlass.Int32(jm)
+                ws = ws + smem_hist[bif]
+        if lane == cutlass.Int32(0):
+            smem_wcnt[warp_id] = ws
+        cute.arch.barrier()
+
+        v_s = cutlass.Int32(0)
+        if lane < cutlass.Int32(self.num_warps):
+            v_s = smem_wcnt[lane]
+        run2 = v_s
+        for d2 in cutlass.range_constexpr(5):
+            off2 = cutlass.const_expr(1 << d2)
+            src2 = lane - cutlass.Int32(off2)
+            if src2 < cutlass.Int32(0):
+                src2 = cutlass.Int32(0)
+            up2 = cute.arch.shuffle_sync(run2, src2)
+            if lane >= cutlass.Int32(off2):
+                run2 = run2 + up2
+        m2 = cute.arch.vote_ballot_sync(rank_above + run2 >= cutlass.Int32(kK))
+        tw = cutlass.Int32(self.num_warps - 1)
+        if m2 != cutlass.Uint32(0):
+            low2 = m2 & (cutlass.Uint32(0) - m2)
+            tw = cutlass.Int32(cute.arch.popc(low2 - cutlass.Uint32(1)))
+        incl_tw = cute.arch.shuffle_sync(run2, tw)
+        slot_tw = cute.arch.shuffle_sync(v_s, tw)
+        prefix = rank_above + (incl_tw - slot_tw)
+
+        ppl = cutlass.const_expr((fbpw + self.WARP_SIZE - 1) // self.WARP_SIZE)
+        cnt_frag = cute.make_fragment((ppl,), cutlass.Int32)
+        my_sum = cutlass.Int32(0)
+        for j3 in cutlass.range_constexpr(ppl):
+            pos = lane * cutlass.Int32(ppl) + cutlass.Int32(j3)
+            cj = cutlass.Int32(0)
+            if pos < cutlass.Int32(fbpw):
+                sbi = cutlass.Int32(fbins - 1) - tw * cutlass.Int32(fbpw) - pos
+                cj = smem_hist[sbi]
+            cnt_frag[j3] = cj
+            my_sum = my_sum + cj
+        run3 = my_sum
+        for d3 in cutlass.range_constexpr(5):
+            off3 = cutlass.const_expr(1 << d3)
+            src3 = lane - cutlass.Int32(off3)
+            if src3 < cutlass.Int32(0):
+                src3 = cutlass.Int32(0)
+            up3 = cute.arch.shuffle_sync(run3, src3)
+            if lane >= cutlass.Int32(off3):
+                run3 = run3 + up3
+        base3 = prefix + (run3 - my_sum)
+
+        sb_loc = cutlass.Int32(fbins - 1)
+        ra_loc = prefix
+        hit = cutlass.Int32(0)
+        r3 = base3
+        for j4 in cutlass.range_constexpr(ppl):
+            pos4 = lane * cutlass.Int32(ppl) + cutlass.Int32(j4)
+            c4 = cnt_frag[j4]
+            if (
+                pos4 < cutlass.Int32(fbpw)
+                and r3 < cutlass.Int32(kK)
+                and r3 + c4 >= cutlass.Int32(kK)
+                and hit == cutlass.Int32(0)
+            ):
+                sb_loc = cutlass.Int32(fbins - 1) - tw * cutlass.Int32(fbpw) - pos4
+                ra_loc = r3
+                hit = cutlass.Int32(1)
+            r3 = r3 + c4
+        mask3 = cute.arch.vote_ballot_sync(hit != cutlass.Int32(0))
+        sb_out = cutlass.Int32(fbins - 1)
+        ra_out = prefix
+        if mask3 != cutlass.Uint32(0):
+            low = mask3 & (cutlass.Uint32(0) - mask3)
+            src = cutlass.Int32(cute.arch.popc(low - cutlass.Uint32(1)))
+            sb_out = cute.arch.shuffle_sync(sb_loc, src)
+            ra_out = cute.arch.shuffle_sync(ra_loc, src)
+        return sb_out, ra_out
+
     # ------------------------------------------------------------------
     # Phase 4 (alt): op#7 fused rank-and-scatter (enable_p4_rank_scatter).
     # Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py.
@@ -3695,15 +3794,22 @@ def phase4_rank_scatter(
                     smem_wcnt[warp_id] = float_as_uint32(cmin)
                     smem_hist[warp_id] = float_as_uint32(cmax)
                 cute.arch.barrier()
-                for w in cutlass.range_constexpr(self.num_warps):
-                    vmin = cutlass.Float32(
-                        llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value())
+                # lane-parallel cross-warp fold: lane w holds slot w and one
+                # warp reduce settles it, instead of every thread walking all
+                # num_warps slots of two arrays with dependent SMEM reads.
+                # min/max reassociate freely, so the result is bit-identical
+                # and every warp still lands on it without a leader.
+                pmn = cutlass.Float32(self.FLT_MAX)
+                pmx = cutlass.Float32(self.NEG_FLT_MAX)
+                if lane < cutlass.Int32(self.num_warps):
+                    pmn = cutlass.Float32(
+                        llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[lane].ir_value())
                     )
-                    vmax = cutlass.Float32(
-                        llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[w].ir_value())
+                    pmx = cutlass.Float32(
+                        llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[lane].ir_value())
                     )
-                    bmin_r = _fmin_f32_inline(bmin_r, vmin)
-                    bmax_r = cute.arch.fmax(bmax_r, vmax)
+                bmin_r = _fmin_f32_inline(bmin_r, self.warp_reduce_min_f32(pmn))
+                bmax_r = cute.arch.fmax(bmax_r, self.warp_reduce_max_f32(pmx))
                 if bmax_r <= bmin_r:
                     bmax_r = bmin_r + cutlass.Float32(1e-6)
                 cute.arch.barrier()
@@ -3802,7 +3908,6 @@ def phase4_rank_scatter(
                 # sub-bins over bin b* gives kNumBins×256 effective resolution,
                 # enough to resolve the straddling bin to ≤1 distinct value.
                 fbins = cutlass.const_expr(256)
-                fbpw = cutlass.const_expr(256 // self.num_warps)
                 # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1)
                 f_lo = bmin_r + cutlass.Float32(b_star) / inv1
                 finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1
@@ -3862,62 +3967,13 @@ def phase4_rank_scatter(
                             atomicAdd(smem_hist.iterator + sbo, cutlass.Int32(1))
                         ifb = ifb + cutlass.Int32(num_threads)
                 cute.arch.barrier()
-                # fine 3-step search seeded at rank_above (over fbins bins)
-                fws = cutlass.Int32(0)
-                for jbf in cutlass.range_constexpr(fbpw):
-                    bif = (
-                        cutlass.Int32(fbins - 1)
-                        - warp_id * cutlass.Int32(fbpw)
-                        - cutlass.Int32(jbf)
-                    )
-                    fws = fws + smem_hist[bif]
-                if lane == cutlass.Int32(0):
-                    smem_wcnt[warp_id] = fws
-                cute.arch.barrier()
-                if tidx == cutlass.Int32(0):
-                    cumf = rank_above
-                    twf = cutlass.Int32(num_warps - 1)
-                    fnd = cutlass.Int32(0)
-                    for w2 in cutlass.range_constexpr(self.num_warps):
-                        cumf = cumf + smem_wcnt[w2]
-                        if cumf >= cutlass.Int32(kK) and fnd == cutlass.Int32(0):
-                            twf = cutlass.Int32(w2)
-                            fnd = cutlass.Int32(1)
-                    pre = rank_above
-                    for w3 in cutlass.range_constexpr(self.num_warps):
-                        if cutlass.Int32(w3) < twf:
-                            pre = pre + smem_wcnt[w3]
-                    # Stage prefix/target-warp metadata in spare s_iscalars
-                    # slots, NOT smem_hist[0]/[1]: the last fine warp's reverse
-                    # scan below walks fine bins down to 0/1, so reusing those
-                    # histogram bins as scratch would corrupt sb_star/ra_fine
-                    # when twf2 == num_warps-1. Slots [4]/[1] are dead here
-                    # (re-zeroed at the cnt_above/cnt_strad reset below).
-                    s_iscalars[4] = pre  # prefix into target fine warp
-                    s_iscalars[1] = twf  # target fine warp
-                cute.arch.barrier()
-                pre_f = s_iscalars[4]
-                twf2 = s_iscalars[1]
-                if warp_id == twf2 and lane == cutlass.Int32(0):
-                    base_f = pre_f
-                    sb_star = cutlass.Int32(fbins - 1)
-                    ra_fine = base_f
-                    sd = cutlass.Int32(0)
-                    for jb3 in cutlass.range_constexpr(fbpw):
-                        sbi = (
-                            cutlass.Int32(fbins - 1)
-                            - twf2 * cutlass.Int32(fbpw)
-                            - cutlass.Int32(jb3)
-                        )
-                        ra_b = base_f
-                        base_f = base_f + smem_hist[sbi]
-                        if base_f >= cutlass.Int32(kK) and sd == cutlass.Int32(0):
-                            sb_star = sbi
-                            ra_fine = ra_b
-                            sd = cutlass.Int32(1)
-                    smem_hist[2] = sb_star
-                    smem_hist[3] = ra_fine
-                cute.arch.barrier()
+                # fine sub-bin search, resolved lane-parallel on every
+                # warp (see _p4_fine_rw): the two publish barriers and the
+                # fbins/num_warps-deep single-lane walk that used to sit
+                # here are gone, and the answer comes back in registers.
+                sb_star, rank_above_fine = self._p4_fine_rw(
+                    smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane
+                )
                 if tidx == cutlass.Int32(0):
                     s_iscalars[4] = cutlass.Int32(0)  # cnt_above
                     s_iscalars[0] = cutlass.Int32(0)  # cnt_mid (b*, sub>sb*)
@@ -3925,8 +3981,6 @@ def phase4_rank_scatter(
                 cute.arch.barrier()
                 if cutlass.const_expr(_P4_SUB_DBG):
                     sc4 = cute.arch.clock64()
-                sb_star = smem_hist[2]
-                rank_above_fine = smem_hist[3]
                 isc = tidx
                 while isc < cand_count:
                     v = smem_keys[isc]

From 26f4bc6b2605047a8c93a7878ee37a5146961a2e Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 4 Aug 2026 04:09:46 -0700
Subject: [PATCH 091/117] [None][perf] GVR: stop paying the big-class tail
 machinery on a 2-way tie

The exact-tail repair fires when the straddling fine bin holds more
candidates than there are output slots left. On the captured rows it
fires with a class of TWO and one slot to fill, and it cost 3.6us to
break that tie - which is where essentially all of the remaining
per-step regressions live. Per-step pairing over the routed grid put
5.7% of pro/flash 8k-32k steps below the stock kernel; the phase clock
shows head/histogram/coarse/fine identical on those steps and the whole
delta in the tail window.

Two things were being bought for a class of two:

  * a whole extra candidate walk for the pure-tie pre-check. That check
    only exists to SKIP the repair when the class is bit-uniform, but
    the small route below ranks by (key, arrival) and rewrites exactly
    the need0 winner slots, which is already value-exact on a pure tie.
    So for a small class the walk is bought for nothing.
  * a per-thread staging array nbuf7 = kC/num_threads deep, whose every
    store is an unrolled dynamic-index search over it. It exists to save
    the compaction its own walk - but a candidate walk is ~0.3us and the
    array measured 1.2us in the compaction alone.

A small mixed class now takes its own route ahead of both: one candidate
walk claiming straight into a pair buffer parked in the (dead) digit
bins, then the same warp0 pairwise rank. Nothing rides in registers, and
smem_keys stays readable while the pass runs, so no staging barrier is
needed. The large-class route keeps the pre-check, the staging array and
the radix - it is the one that really wants to avoid a second walk.

Measured on the pro 8k capture (30 layers, 420 cold steps, phase-4
clock, us/step):

  step class            before   after
  fires (class 2)        6.77     5.72
  does not fire          3.14     3.14

Exactness: 30/30 layers on the firing capture (B200); 1947/1947 records
over flash 8k/64k/256k/1024k and pro 1024k on both assist tiers (B300);
the full decode test file passes on both B200 and B300 (708 passed,
144 skipped, 1 xfailed).

About 2.5us of the fire-step delta is still unaccounted for and is not
in these two passes - the remaining route is one walk, a few barriers
and a rank over two elements. That is the next thing to chase.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 365 +++++++++++-------
 1 file changed, 218 insertions(+), 147 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 5381dece8651..707acf0085c5 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -4105,7 +4105,15 @@ def phase4_rank_scatter(
                             kmn6 = cutlass.Int32(2147483647)
                             kmx6 = cutlass.Int32(-2147483648)
                             it6 = tidx
-                            while it6 < cand_count:
+                            # The pure-tie pre-check only exists to SKIP the
+                            # repair when the class is bit-uniform. The small
+                            # route below ranks by (key, arrival) and rewrites
+                            # exactly the need0 winner slots, which is already
+                            # value-exact on a pure tie - so for a small class
+                            # the check is a whole candidate walk bought for
+                            # nothing. Only the large-class route, whose radix
+                            # really is worth avoiding, still pays for it.
+                            while it6 < cand_count and cnt_strad > cutlass.Int32(128):
                                 v6 = smem_keys[it6]
                                 b6 = cutlass.Int32((v6 - bmin_r) * inv1)
                                 if b6 < cutlass.Int32(0):
@@ -4124,13 +4132,25 @@ def phase4_rank_scatter(
                                             kmn6 = k6
                                         if k6 > kmx6:
                                             kmx6 = k6
-                                        # same predicate as the compaction
-                                        # pass used to re-derive: buffer the
-                                        # member here so that pass can go
-                                        for sl7 in cutlass.range_constexpr(nbuf7):
-                                            if cutlass.Int32(sl7) == nh7:
-                                                rv7[sl7] = v6
-                                                ri7[sl7] = smem_vals[it6]
+                                        # Buffering the member here saves the
+                                        # compaction its own walk, but the
+                                        # per-thread array is nbuf7 =
+                                        # kC/num_threads deep and every store
+                                        # is an unrolled dynamic-index search
+                                        # over it. Measured on pro 8k (class
+                                        # 2, one slot to fill): this pass
+                                        # 2.4us + compaction 1.2us, against a
+                                        # second candidate walk that costs
+                                        # ~0.3us. So buffer ONLY for classes
+                                        # too large for the re-walk route
+                                        # below; cnt_strad is block-uniform,
+                                        # so small classes skip the unrolled
+                                        # stores entirely.
+                                        if cnt_strad > cutlass.Int32(128):
+                                            for sl7 in cutlass.range_constexpr(nbuf7):
+                                                if cutlass.Int32(sl7) == nh7:
+                                                    rv7[sl7] = v6
+                                                    ri7[sl7] = smem_vals[it6]
                                         nh7 = nh7 + cutlass.Int32(1)
                                 it6 = it6 + cutlass.Int32(num_threads)
                             kmn6 = cute.arch.warp_redux_sync(kmn6, "min")
@@ -4154,7 +4174,92 @@ def phase4_rank_scatter(
                             kmx7 = cute.arch.warp_redux_sync(pb8, "max")
                             if kmn7 == kmx7:
                                 fast_done = cutlass.Int32(1)
-                            if fast_done == cutlass.Int32(0):
+                            if cnt_strad <= cutlass.Int32(128):
+                                # Small mixed class: walk the candidates once
+                                # more and claim straight into a pair buffer
+                                # in the (dead) digit bins, so nothing has to
+                                # ride in registers and smem_keys stays
+                                # readable while the pass runs. 128 pairs fit
+                                # under the 256 bins the radix route would
+                                # otherwise zero.
+                                if tidx == cutlass.Int32(0):
+                                    s_iscalars[0] = cutlass.Int32(0)
+                                cute.arch.barrier()
+                                itc9 = tidx
+                                while itc9 < cand_count:
+                                    tv9 = smem_keys[itc9]
+                                    tb9 = cutlass.Int32((tv9 - bmin_r) * inv1)
+                                    if tb9 < cutlass.Int32(0):
+                                        tb9 = cutlass.Int32(0)
+                                    if tb9 > cutlass.Int32(kBins - 1):
+                                        tb9 = cutlass.Int32(kBins - 1)
+                                    if tb9 == b_star:
+                                        ts9 = cutlass.Int32((tv9 - f_lo) * finv)
+                                        if ts9 < cutlass.Int32(0):
+                                            ts9 = cutlass.Int32(0)
+                                        if ts9 > cutlass.Int32(fbins - 1):
+                                            ts9 = cutlass.Int32(fbins - 1)
+                                        if ts9 == sb_star:
+                                            to9 = atomicAdd(
+                                                s_iscalars.iterator + cutlass.Int32(0),
+                                                cutlass.Int32(1),
+                                            )
+                                            if to9 < cutlass.Int32(128):
+                                                smem_hist[to9 + to9] = float_as_int32(tv9)
+                                                smem_hist[to9 + to9 + cutlass.Int32(1)] = smem_vals[
+                                                    itc9
+                                                ]
+                                    itc9 = itc9 + cutlass.Int32(num_threads)
+                                cute.arch.barrier()
+                                # warp0 exact pairwise rank over the pairs
+                                # (rank = #{key greater} + #{key equal,
+                                # earlier slot} is unique in [0, class)),
+                                # rewriting each winner slot once.
+                                if warp_id == cutlass.Int32(0):
+                                    ie9 = lane
+                                    while ie9 < cnt_strad:
+                                        bi9 = smem_hist[ie9 + ie9]
+                                        ki9 = f32_order_key(
+                                            cutlass.Float32(
+                                                llvm.bitcast(
+                                                    cutlass.Float32.mlir_type, bi9.ir_value()
+                                                )
+                                            )
+                                        ) ^ cutlass.Int32(-2147483648)
+                                        r9 = cutlass.Int32(0)
+                                        j9 = cutlass.Int32(0)
+                                        while j9 < cnt_strad:
+                                            bj9 = smem_hist[j9 + j9]
+                                            kj9 = f32_order_key(
+                                                cutlass.Float32(
+                                                    llvm.bitcast(
+                                                        cutlass.Float32.mlir_type, bj9.ir_value()
+                                                    )
+                                                )
+                                            ) ^ cutlass.Int32(-2147483648)
+                                            if kj9 > ki9:
+                                                r9 = r9 + cutlass.Int32(1)
+                                            elif kj9 == ki9 and j9 < ie9:
+                                                r9 = r9 + cutlass.Int32(1)
+                                            j9 = j9 + cutlass.Int32(1)
+                                        if r9 < need0:
+                                            pos9 = rank_above_fine + r9
+                                            if pos9 < cutlass.Int32(kK):
+                                                if cutlass.const_expr(self.return_output_values):
+                                                    output_values_row[pos9] = self.dtype(
+                                                        cutlass.Float32(
+                                                            llvm.bitcast(
+                                                                cutlass.Float32.mlir_type,
+                                                                bi9.ir_value(),
+                                                            )
+                                                        )
+                                                    )
+                                                output_indices_row[pos9] = smem_hist[
+                                                    ie9 + ie9 + cutlass.Int32(1)
+                                                ]
+                                        ie9 = ie9 + cutlass.Int32(32)
+                                cute.arch.barrier()
+                            elif fast_done == cutlass.Int32(0):
                                 # mixed class: compact the members buffered by
                                 # the pure-tie pass above into
                                 # smem_keys/vals[0..cnt_strad). The buffering
@@ -4189,154 +4294,120 @@ def phase4_rank_scatter(
                                         smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8]
                                         smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8]
                                 cute.arch.barrier()
-                                if cnt_strad <= cutlass.Int32(128):
-                                    # warp0 exact pairwise rank (rank = #{key
-                                    # greater} + #{key equal, earlier slot} is
-                                    # unique in [0, class)) rewrites every
-                                    # winner slot in [raf, raf + need0) once.
-                                    if warp_id == cutlass.Int32(0):
-                                        ie5 = lane
-                                        while ie5 < cnt_strad:
-                                            vi5 = smem_keys[ie5]
-                                            ki5 = f32_order_key(vi5) ^ cutlass.Int32(-2147483648)
-                                            r5 = cutlass.Int32(0)
-                                            j5 = cutlass.Int32(0)
-                                            while j5 < cnt_strad:
-                                                vj5 = smem_keys[j5]
-                                                kj5 = f32_order_key(vj5) ^ cutlass.Int32(
-                                                    -2147483648
-                                                )
-                                                if kj5 > ki5:
-                                                    r5 = r5 + cutlass.Int32(1)
-                                                elif kj5 == ki5 and j5 < ie5:
-                                                    r5 = r5 + cutlass.Int32(1)
-                                                j5 = j5 + cutlass.Int32(1)
-                                            if r5 < need0:
-                                                pos = rank_above_fine + r5
-                                                if pos < cutlass.Int32(kK):
-                                                    if cutlass.const_expr(
-                                                        self.return_output_values
-                                                    ):
-                                                        output_values_row[pos] = self.dtype(vi5)
-                                                    output_indices_row[pos] = smem_vals[ie5]
-                                            ie5 = ie5 + cutlass.Int32(32)
+                                # Large class only (the small one is handled
+                                # above, ahead of the pure-tie pass, so it
+                                # never reaches here): block-parallel 4-level
+                                # MSB radix over the compacted class (scans
+                                # touch class pairs only; warp0 shuffle-scan
+                                # digit search - 3 block barriers per level
+                                # instead of 5).
+                                if tidx == cutlass.Int32(0):
+                                    smem_hist[256] = cutlass.Int32(0)
+                                    smem_hist[257] = need0
+                                    smem_hist[258] = cutlass.Int32(0)
+                                cute.arch.barrier()
+                                for lvl2 in cutlass.range_constexpr(4):
+                                    shift2 = cutlass.const_expr(24 - 8 * lvl2)
+                                    iz3 = tidx
+                                    while iz3 < cutlass.Int32(256):
+                                        smem_hist[iz3] = cutlass.Int32(0)
+                                        iz3 = iz3 + cutlass.Int32(num_threads)
                                     cute.arch.barrier()
-                                else:
-                                    # block-parallel 4-level MSB radix over the
-                                    # compacted class (scans touch class pairs
-                                    # only; warp0 shuffle-scan digit search — 3
-                                    # block barriers per level instead of 5).
-                                    if tidx == cutlass.Int32(0):
-                                        smem_hist[256] = cutlass.Int32(0)
-                                        smem_hist[257] = need0
-                                        smem_hist[258] = cutlass.Int32(0)
+                                    uthr_c2 = smem_hist[256]
+                                    ic2 = tidx
+                                    while ic2 < cnt_strad:
+                                        uk3 = f32_order_key(smem_keys[ic2])
+                                        pm2 = cutlass.Int32(1)
+                                        if cutlass.const_expr(lvl2 > 0):
+                                            if (uk3 >> cutlass.Int32(shift2 + 8)) != (
+                                                uthr_c2 >> cutlass.Int32(shift2 + 8)
+                                            ):
+                                                pm2 = cutlass.Int32(0)
+                                        if pm2 == cutlass.Int32(1):
+                                            dg2 = (uk3 >> cutlass.Int32(shift2)) & cutlass.Int32(
+                                                0xFF
+                                            )
+                                            atomicAdd(smem_hist.iterator + dg2, cutlass.Int32(1))
+                                        ic2 = ic2 + cutlass.Int32(num_threads)
                                     cute.arch.barrier()
-                                    for lvl2 in cutlass.range_constexpr(4):
-                                        shift2 = cutlass.const_expr(24 - 8 * lvl2)
-                                        iz3 = tidx
-                                        while iz3 < cutlass.Int32(256):
-                                            smem_hist[iz3] = cutlass.Int32(0)
-                                            iz3 = iz3 + cutlass.Int32(num_threads)
-                                        cute.arch.barrier()
-                                        uthr_c2 = smem_hist[256]
-                                        ic2 = tidx
-                                        while ic2 < cnt_strad:
-                                            uk3 = f32_order_key(smem_keys[ic2])
-                                            pm2 = cutlass.Int32(1)
-                                            if cutlass.const_expr(lvl2 > 0):
-                                                if (uk3 >> cutlass.Int32(shift2 + 8)) != (
-                                                    uthr_c2 >> cutlass.Int32(shift2 + 8)
-                                                ):
-                                                    pm2 = cutlass.Int32(0)
-                                            if pm2 == cutlass.Int32(1):
-                                                dg2 = (
-                                                    uk3 >> cutlass.Int32(shift2)
-                                                ) & cutlass.Int32(0xFF)
-                                                atomicAdd(
-                                                    smem_hist.iterator + dg2, cutlass.Int32(1)
-                                                )
-                                            ic2 = ic2 + cutlass.Int32(num_threads)
-                                        cute.arch.barrier()
-                                        if warp_id == cutlass.Int32(0):
-                                            ws3 = cutlass.Int32(0)
-                                            for jd3 in cutlass.range_constexpr(8):
-                                                di3 = (
+                                    if warp_id == cutlass.Int32(0):
+                                        ws3 = cutlass.Int32(0)
+                                        for jd3 in cutlass.range_constexpr(8):
+                                            di3 = (
+                                                cutlass.Int32(255)
+                                                - lane * cutlass.Int32(8)
+                                                - cutlass.Int32(jd3)
+                                            )
+                                            ws3 = ws3 + smem_hist[di3]
+                                        pre6 = ws3
+                                        for so2 in cutlass.range_constexpr(5):
+                                            oth2 = cute.arch.shuffle_sync_up(
+                                                pre6,
+                                                cutlass.Int32(1 << so2),
+                                                mask_and_clamp=0,
+                                            )
+                                            if lane >= cutlass.Int32(1 << so2):
+                                                pre6 = pre6 + oth2
+                                        needl3 = smem_hist[257]
+                                        if pre6 >= needl3 and (pre6 - ws3) < needl3:
+                                            base5 = pre6 - ws3
+                                            dstar2 = cutlass.Int32(0)
+                                            above5 = base5
+                                            sd5 = cutlass.Int32(0)
+                                            for jd4 in cutlass.range_constexpr(8):
+                                                di4 = (
                                                     cutlass.Int32(255)
                                                     - lane * cutlass.Int32(8)
-                                                    - cutlass.Int32(jd3)
-                                                )
-                                                ws3 = ws3 + smem_hist[di3]
-                                            pre6 = ws3
-                                            for so2 in cutlass.range_constexpr(5):
-                                                oth2 = cute.arch.shuffle_sync_up(
-                                                    pre6,
-                                                    cutlass.Int32(1 << so2),
-                                                    mask_and_clamp=0,
-                                                )
-                                                if lane >= cutlass.Int32(1 << so2):
-                                                    pre6 = pre6 + oth2
-                                            needl3 = smem_hist[257]
-                                            if pre6 >= needl3 and (pre6 - ws3) < needl3:
-                                                base5 = pre6 - ws3
-                                                dstar2 = cutlass.Int32(0)
-                                                above5 = base5
-                                                sd5 = cutlass.Int32(0)
-                                                for jd4 in cutlass.range_constexpr(8):
-                                                    di4 = (
-                                                        cutlass.Int32(255)
-                                                        - lane * cutlass.Int32(8)
-                                                        - cutlass.Int32(jd4)
-                                                    )
-                                                    ra5 = base5
-                                                    base5 = base5 + smem_hist[di4]
-                                                    if base5 >= needl3 and sd5 == cutlass.Int32(0):
-                                                        dstar2 = di4
-                                                        above5 = ra5
-                                                        sd5 = cutlass.Int32(1)
-                                                smem_hist[256] = uthr_c2 | (
-                                                    dstar2 << cutlass.Int32(shift2)
+                                                    - cutlass.Int32(jd4)
                                                 )
-                                                smem_hist[257] = needl3 - above5
-                                                smem_hist[258] = smem_hist[258] + above5
-                                        cute.arch.barrier()
-                                    u_thr2 = smem_hist[256]
-                                    cnt_ab2 = smem_hist[258]
-                                    need_eq2 = smem_hist[257]
-                                    kthr2 = u_thr2 ^ cutlass.Int32(-2147483648)
-                                    if tidx == cutlass.Int32(0):
-                                        s_iscalars[4] = cutlass.Int32(0)
-                                        s_iscalars[0] = cutlass.Int32(0)
-                                    cute.arch.barrier()
-                                    ir3 = tidx
-                                    while ir3 < cnt_strad:
-                                        vv3 = smem_keys[ir3]
-                                        uk4 = f32_order_key(vv3)
-                                        ks4 = uk4 ^ cutlass.Int32(-2147483648)
-                                        if ks4 > kthr2:
-                                            o4 = atomicAdd(
-                                                s_iscalars.iterator + cutlass.Int32(4),
-                                                cutlass.Int32(1),
+                                                ra5 = base5
+                                                base5 = base5 + smem_hist[di4]
+                                                if base5 >= needl3 and sd5 == cutlass.Int32(0):
+                                                    dstar2 = di4
+                                                    above5 = ra5
+                                                    sd5 = cutlass.Int32(1)
+                                            smem_hist[256] = uthr_c2 | (
+                                                dstar2 << cutlass.Int32(shift2)
                                             )
-                                            pos = rank_above_fine + o4
+                                            smem_hist[257] = needl3 - above5
+                                            smem_hist[258] = smem_hist[258] + above5
+                                    cute.arch.barrier()
+                                u_thr2 = smem_hist[256]
+                                cnt_ab2 = smem_hist[258]
+                                need_eq2 = smem_hist[257]
+                                kthr2 = u_thr2 ^ cutlass.Int32(-2147483648)
+                                if tidx == cutlass.Int32(0):
+                                    s_iscalars[4] = cutlass.Int32(0)
+                                    s_iscalars[0] = cutlass.Int32(0)
+                                cute.arch.barrier()
+                                ir3 = tidx
+                                while ir3 < cnt_strad:
+                                    vv3 = smem_keys[ir3]
+                                    uk4 = f32_order_key(vv3)
+                                    ks4 = uk4 ^ cutlass.Int32(-2147483648)
+                                    if ks4 > kthr2:
+                                        o4 = atomicAdd(
+                                            s_iscalars.iterator + cutlass.Int32(4),
+                                            cutlass.Int32(1),
+                                        )
+                                        pos = rank_above_fine + o4
+                                        if pos < cutlass.Int32(kK):
+                                            if cutlass.const_expr(self.return_output_values):
+                                                output_values_row[pos] = self.dtype(vv3)
+                                            output_indices_row[pos] = smem_vals[ir3]
+                                    elif ks4 == kthr2:
+                                        q4 = atomicAdd(
+                                            s_iscalars.iterator + cutlass.Int32(0),
+                                            cutlass.Int32(1),
+                                        )
+                                        if q4 < need_eq2:
+                                            pos = rank_above_fine + cnt_ab2 + q4
                                             if pos < cutlass.Int32(kK):
                                                 if cutlass.const_expr(self.return_output_values):
                                                     output_values_row[pos] = self.dtype(vv3)
                                                 output_indices_row[pos] = smem_vals[ir3]
-                                        elif ks4 == kthr2:
-                                            q4 = atomicAdd(
-                                                s_iscalars.iterator + cutlass.Int32(0),
-                                                cutlass.Int32(1),
-                                            )
-                                            if q4 < need_eq2:
-                                                pos = rank_above_fine + cnt_ab2 + q4
-                                                if pos < cutlass.Int32(kK):
-                                                    if cutlass.const_expr(
-                                                        self.return_output_values
-                                                    ):
-                                                        output_values_row[pos] = self.dtype(vv3)
-                                                    output_indices_row[pos] = smem_vals[ir3]
-                                        ir3 = ir3 + cutlass.Int32(num_threads)
-                                    cute.arch.barrier()
+                                    ir3 = ir3 + cutlass.Int32(num_threads)
+                                cute.arch.barrier()
                     else:
                         need0_s = cutlass.Int32(kK) - rank_above_fine
                         if cnt_strad > need0_s and need0_s > cutlass.Int32(0):

From f378bf87fd666d3b1abf9b1d3119326461a1f156 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 4 Aug 2026 19:59:01 -0700
Subject: [PATCH 092/117] [None][perf] GVR: let the scatter hand the tie class
 to the tail repair

The exact-tail repair fires on a class of two candidates contending for
one slot, and even after the previous commit it still cost 1.5us more
than a step that does not fire. Attribution by clock stamp, on the
pro 8k capture (phase-4 window, us/step): the rank itself 0.33, a
candidate walk to collect the class 0.61, and 0.55 of machinery that
runs whether or not it is needed.

All three go away:

  * The scatter already classifies every candidate and hands out the
    intra-class ordinal, so it now parks each member as a (value bits,
    index) pair above the digit bins. The repair reads them directly -
    no collection walk, no staging barrier.
  * The pure-tie pre-check, its cross-warp fold and the publish barrier
    that fold needs exist only to decide whether the radix can be
    skipped. The small-class route needs none of it and used to run it
    anyway; it now lives in the large-class branch, which is the one
    with a radix to avoid.
  * The rank loop is forced rolled (unroll=1). Its trip count is 1-2,
    but an unrolled body plus its remainder ladder is code volume that
    ONE warp walks cold while every other warp waits at the barrier
    below, so the fetch latency has nothing to hide behind.

The pair buffer's capacity is DERIVED from the histogram size, not
assumed: K=2048 with R0 shrinks the bin count to 512, where a fixed 128
pairs would run four ints past the end of the allocation and into the
per-thread count buffer that follows it. It is rounded down to a power
of two so the scatter can wrap the ordinal with a mask instead of a
bounds branch, and a class past the cap takes the large-class route,
which does not read the buffer.

Phase-4 window on the firing capture (pro 8k, 30 layers, 420 cold
steps, us/step):

  step class          before this   after
  fires (class 2)        6.77        3.42
  does not fire          3.14        3.16

Per-step paired against the same baseline, same node, same session,
over the six cells that carried the regressions (6804 steps, previous
tree vs this one):

  steps slower than baseline   92 -> 27
  steps below 0.9              29 -> 18
  worst single step          0.642 -> 0.675
  mean                       1.254 -> 1.248

Exactness: 30/30 layers on the firing capture; 1023/1023 records over
flash 8k/64k/256k and pro 1024k on both assist tiers; the full decode
test file passes on B200 and B300 (708 passed, 144 skipped, 1 xfailed).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 558 +++++++++---------
 1 file changed, 290 insertions(+), 268 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 707acf0085c5..07882ed648b6 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -59,6 +59,28 @@ def _env_flag(name: str) -> bool:
 # GVR_P4_SUB_DBG=2: publish the P4 HEAD triple (minmax / histogram build /
 # coarse search) instead of the tail triple, so the phase budget adds up.
 _P4_SUB_HEAD = os.environ.get("GVR_P4_SUB_DBG", "0").strip() == "2"
+# Exact-tail small-class pair buffer. The scatter parks each member of the
+# straddling tie class as (value bits, index) here, above the 256 digit bins
+# the large-class radix zeroes and above its [256..258] scalars, so the repair
+# never has to re-walk the candidates. The capacity is DERIVED from the bin
+# count, never assumed: K=2048 with R0 shrinks the histogram to 512 bins (see
+# the kNumBins override in __init__), where 260 + 2*128 would run 4 ints past
+# the end of the allocation - into the per-thread count buffer that follows it.
+# Rounded down to a power of two so the scatter can wrap the ordinal with a
+# mask instead of a bounds branch; a class past the cap takes the large-class
+# route, which does not use this buffer.
+_PAIR_BASE = 260
+_PAIR_MAX = 128
+
+
+def _pair_cap_for(n_bins: int) -> int:
+    """Largest power-of-two pair count that fits [_PAIR_BASE, n_bins)."""
+    room = (n_bins - _PAIR_BASE) // 2
+    if room < 1:
+        return 0
+    return min(_PAIR_MAX, 1 << (room.bit_length() - 1))
+
+
 _SKIP_DBG = _env_flag("GVR_SKIP_DBG")
 
 
@@ -3735,6 +3757,7 @@ def phase4_rank_scatter(
     ):
         kK = cutlass.const_expr(self.top_k)
         kBins = cutlass.const_expr(self.kNumBins)
+        pair_cap = cutlass.const_expr(_pair_cap_for(self.kNumBins))
         num_threads = cutlass.const_expr(self.num_threads)
         num_warps = cutlass.const_expr(self.num_warps)
         bins_per_warp = cutlass.const_expr(kBins // self.num_warps)
@@ -4031,6 +4054,28 @@ def phase4_rank_scatter(
                                 output_indices_row[pos] = smem_vals[isc]
                         elif sb == sb_star:
                             o = atomicAdd(s_iscalars.iterator + cutlass.Int32(1), cutlass.Int32(1))
+                            if cutlass.const_expr(
+                                self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3
+                            ):
+                                # The tie class is exactly what the exact-tail
+                                # repair needs, and this pass has already
+                                # classified every candidate AND handed out the
+                                # intra-class ordinal. Park (value bits, index)
+                                # here so the repair does not have to walk the
+                                # candidates again: measured 0.61us of its
+                                # remaining cost was that second walk. Pairs
+                                # live above the digit bins the radix route
+                                # zeroes; small classes are the only consumer.
+                                # Unconditional, index wrapped: the buffer is
+                                # only ever READ when the class fits it, and
+                                # then the wrap is a no-op. Dropping the bound
+                                # branch keeps this off the scatter's critical
+                                # path, which every step pays.
+                                ow = (o & cutlass.Int32(pair_cap - 1)) * cutlass.Int32(2)
+                                smem_hist[cutlass.Int32(_PAIR_BASE) + ow] = float_as_int32(v)
+                                smem_hist[cutlass.Int32(_PAIR_BASE) + ow + cutlass.Int32(1)] = (
+                                    smem_vals[isc]
+                                )
                             pos = rank_above_fine + o
                             if pos < cutlass.Int32(kK):
                                 if cutlass.const_expr(self.return_output_values):
@@ -4090,135 +4135,24 @@ def phase4_rank_scatter(
                         fast_done = cutlass.Int32(1)
                         if cnt_strad > need0 and need0 > cutlass.Int32(0):
                             fast_done = cutlass.Int32(0)
-                            # block-wide pure-tie check, ANY class
-                            # size: min/max order key over the (b*, sb*) class.
-                            # A pure-tie class needs NO repair — the scatter's
-                            # arrival fill of bit-equal values is already
-                            # value-set exact. Real fp8-lineage logits tie in
-                            # the thousands, which used to take the full radix.
-                            # Staging mirrors the head min/max (wcnt + hist
-                            # slots [0..31], both dead here; pairs live at
-                            # 260+).
-                            if tidx == cutlass.Int32(0):
-                                s_iscalars[0] = cutlass.Int32(0)
-                            nh7 = cutlass.Int32(0)
-                            kmn6 = cutlass.Int32(2147483647)
-                            kmx6 = cutlass.Int32(-2147483648)
-                            it6 = tidx
-                            # The pure-tie pre-check only exists to SKIP the
-                            # repair when the class is bit-uniform. The small
-                            # route below ranks by (key, arrival) and rewrites
-                            # exactly the need0 winner slots, which is already
-                            # value-exact on a pure tie - so for a small class
-                            # the check is a whole candidate walk bought for
-                            # nothing. Only the large-class route, whose radix
-                            # really is worth avoiding, still pays for it.
-                            while it6 < cand_count and cnt_strad > cutlass.Int32(128):
-                                v6 = smem_keys[it6]
-                                b6 = cutlass.Int32((v6 - bmin_r) * inv1)
-                                if b6 < cutlass.Int32(0):
-                                    b6 = cutlass.Int32(0)
-                                if b6 > cutlass.Int32(kBins - 1):
-                                    b6 = cutlass.Int32(kBins - 1)
-                                if b6 == b_star:
-                                    s6 = cutlass.Int32((v6 - f_lo) * finv)
-                                    if s6 < cutlass.Int32(0):
-                                        s6 = cutlass.Int32(0)
-                                    if s6 > cutlass.Int32(fbins - 1):
-                                        s6 = cutlass.Int32(fbins - 1)
-                                    if s6 == sb_star:
-                                        k6 = f32_order_key(v6) ^ cutlass.Int32(-2147483648)
-                                        if k6 < kmn6:
-                                            kmn6 = k6
-                                        if k6 > kmx6:
-                                            kmx6 = k6
-                                        # Buffering the member here saves the
-                                        # compaction its own walk, but the
-                                        # per-thread array is nbuf7 =
-                                        # kC/num_threads deep and every store
-                                        # is an unrolled dynamic-index search
-                                        # over it. Measured on pro 8k (class
-                                        # 2, one slot to fill): this pass
-                                        # 2.4us + compaction 1.2us, against a
-                                        # second candidate walk that costs
-                                        # ~0.3us. So buffer ONLY for classes
-                                        # too large for the re-walk route
-                                        # below; cnt_strad is block-uniform,
-                                        # so small classes skip the unrolled
-                                        # stores entirely.
-                                        if cnt_strad > cutlass.Int32(128):
-                                            for sl7 in cutlass.range_constexpr(nbuf7):
-                                                if cutlass.Int32(sl7) == nh7:
-                                                    rv7[sl7] = v6
-                                                    ri7[sl7] = smem_vals[it6]
-                                        nh7 = nh7 + cutlass.Int32(1)
-                                it6 = it6 + cutlass.Int32(num_threads)
-                            kmn6 = cute.arch.warp_redux_sync(kmn6, "min")
-                            kmx6 = cute.arch.warp_redux_sync(kmx6, "max")
-                            if lane == cutlass.Int32(0):
-                                smem_wcnt[warp_id] = kmn6
-                                smem_hist[warp_id] = kmx6
-                            cute.arch.barrier()
-                            # lane-parallel cross-warp fold: lane w holds
-                            # slot w and one warp reduce settles it, instead
-                            # of every thread walking all num_warps slots of
-                            # two arrays with dependent SMEM reads. Same
-                            # inputs in the same order on every warp, so the
-                            # result stays bit-identical and leaderless.
-                            pa8 = cutlass.Int32(2147483647)
-                            pb8 = cutlass.Int32(-2147483648)
-                            if lane < cutlass.Int32(self.num_warps):
-                                pa8 = smem_wcnt[lane]
-                                pb8 = smem_hist[lane]
-                            kmn7 = cute.arch.warp_redux_sync(pa8, "min")
-                            kmx7 = cute.arch.warp_redux_sync(pb8, "max")
-                            if kmn7 == kmx7:
-                                fast_done = cutlass.Int32(1)
-                            if cnt_strad <= cutlass.Int32(128):
-                                # Small mixed class: walk the candidates once
-                                # more and claim straight into a pair buffer
-                                # in the (dead) digit bins, so nothing has to
-                                # ride in registers and smem_keys stays
-                                # readable while the pass runs. 128 pairs fit
-                                # under the 256 bins the radix route would
-                                # otherwise zero.
-                                if tidx == cutlass.Int32(0):
-                                    s_iscalars[0] = cutlass.Int32(0)
-                                cute.arch.barrier()
-                                itc9 = tidx
-                                while itc9 < cand_count:
-                                    tv9 = smem_keys[itc9]
-                                    tb9 = cutlass.Int32((tv9 - bmin_r) * inv1)
-                                    if tb9 < cutlass.Int32(0):
-                                        tb9 = cutlass.Int32(0)
-                                    if tb9 > cutlass.Int32(kBins - 1):
-                                        tb9 = cutlass.Int32(kBins - 1)
-                                    if tb9 == b_star:
-                                        ts9 = cutlass.Int32((tv9 - f_lo) * finv)
-                                        if ts9 < cutlass.Int32(0):
-                                            ts9 = cutlass.Int32(0)
-                                        if ts9 > cutlass.Int32(fbins - 1):
-                                            ts9 = cutlass.Int32(fbins - 1)
-                                        if ts9 == sb_star:
-                                            to9 = atomicAdd(
-                                                s_iscalars.iterator + cutlass.Int32(0),
-                                                cutlass.Int32(1),
-                                            )
-                                            if to9 < cutlass.Int32(128):
-                                                smem_hist[to9 + to9] = float_as_int32(tv9)
-                                                smem_hist[to9 + to9 + cutlass.Int32(1)] = smem_vals[
-                                                    itc9
-                                                ]
-                                    itc9 = itc9 + cutlass.Int32(num_threads)
-                                cute.arch.barrier()
-                                # warp0 exact pairwise rank over the pairs
-                                # (rank = #{key greater} + #{key equal,
-                                # earlier slot} is unique in [0, class)),
-                                # rewriting each winner slot once.
+                            if cnt_strad <= cutlass.Int32(pair_cap):
+                                # Small mixed class: the scatter already parked
+                                # every member as a (value bits, index) pair,
+                                # so there is nothing to collect - go straight
+                                # to the rank. No extra candidate walk, no
+                                # staging barrier, nothing riding in registers.
+                                # unroll=1 on the rank loop: the class is tiny,
+                                # so the trip count is 1-2, but an unrolled
+                                # body (plus its remainder ladder) is code
+                                # volume that ONE warp walks cold while every
+                                # other warp waits at the barrier below - the
+                                # fetch latency has nothing to hide behind.
+                                # Keeping it rolled trades a branch per trip
+                                # for a body that fits a cache line or two.
                                 if warp_id == cutlass.Int32(0):
                                     ie9 = lane
                                     while ie9 < cnt_strad:
-                                        bi9 = smem_hist[ie9 + ie9]
+                                        bi9 = smem_hist[_PAIR_BASE + ie9 + ie9]
                                         ki9 = f32_order_key(
                                             cutlass.Float32(
                                                 llvm.bitcast(
@@ -4227,9 +4161,8 @@ def phase4_rank_scatter(
                                             )
                                         ) ^ cutlass.Int32(-2147483648)
                                         r9 = cutlass.Int32(0)
-                                        j9 = cutlass.Int32(0)
-                                        while j9 < cnt_strad:
-                                            bj9 = smem_hist[j9 + j9]
+                                        for j9 in cutlass.range(0, cnt_strad, 1, unroll=1):
+                                            bj9 = smem_hist[_PAIR_BASE + j9 + j9]
                                             kj9 = f32_order_key(
                                                 cutlass.Float32(
                                                     llvm.bitcast(
@@ -4241,7 +4174,6 @@ def phase4_rank_scatter(
                                                 r9 = r9 + cutlass.Int32(1)
                                             elif kj9 == ki9 and j9 < ie9:
                                                 r9 = r9 + cutlass.Int32(1)
-                                            j9 = j9 + cutlass.Int32(1)
                                         if r9 < need0:
                                             pos9 = rank_above_fine + r9
                                             if pos9 < cutlass.Int32(kK):
@@ -4255,159 +4187,249 @@ def phase4_rank_scatter(
                                                         )
                                                     )
                                                 output_indices_row[pos9] = smem_hist[
-                                                    ie9 + ie9 + cutlass.Int32(1)
+                                                    _PAIR_BASE + ie9 + ie9 + cutlass.Int32(1)
                                                 ]
                                         ie9 = ie9 + cutlass.Int32(32)
                                 cute.arch.barrier()
-                            elif fast_done == cutlass.Int32(0):
-                                # mixed class: compact the members buffered by
-                                # the pure-tie pass above into
-                                # smem_keys/vals[0..cnt_strad). The buffering
-                                # pass already read every candidate it needs,
-                                # and the staging barrier above orders those
-                                # reads before these writes, so the second
-                                # full-candidate walk (and its barrier) is
-                                # gone. Repair cost is a function of the CLASS
-                                # size only, for any class size.
-                                # warp-aggregated claim: intra-warp exclusive
-                                # prefix via shfl scan + ONE atomic per warp
-                                # (a thousand same-address claims serialize
-                                # and scale with the class size)
-                                pf7 = nh7
-                                for so3 in cutlass.range_constexpr(5):
-                                    oth3 = cute.arch.shuffle_sync_up(
-                                        pf7, cutlass.Int32(1 << so3), mask_and_clamp=0
-                                    )
-                                    if lane >= cutlass.Int32(1 << so3):
-                                        pf7 = pf7 + oth3
-                                tot7 = cute.arch.shuffle_sync(pf7, cutlass.Int32(31))
-                                wb7 = cutlass.Int32(0)
-                                if lane == cutlass.Int32(31):
-                                    if tot7 > cutlass.Int32(0):
-                                        wb7 = atomicAdd(
-                                            s_iscalars.iterator + cutlass.Int32(0), tot7
-                                        )
-                                wb7 = cute.arch.shuffle_sync(wb7, cutlass.Int32(31))
-                                bs7 = wb7 + pf7 - nh7
-                                for sl8 in cutlass.range_constexpr(nbuf7):
-                                    if cutlass.Int32(sl8) < nh7:
-                                        smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8]
-                                        smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8]
-                                cute.arch.barrier()
-                                # Large class only (the small one is handled
-                                # above, ahead of the pure-tie pass, so it
-                                # never reaches here): block-parallel 4-level
-                                # MSB radix over the compacted class (scans
-                                # touch class pairs only; warp0 shuffle-scan
-                                # digit search - 3 block barriers per level
-                                # instead of 5).
+                            else:
+                                # Large class only. Everything below -
+                                # the pure-tie pre-check walk, its
+                                # cross-warp fold and the publish
+                                # barrier that fold needs - exists to
+                                # decide whether the radix can be
+                                # skipped. The small route above needs
+                                # none of it, and used to run it anyway
+                                # for a class of two.
+                                # block-wide pure-tie check, ANY class
+                                # size: min/max order key over the (b*, sb*) class.
+                                # A pure-tie class needs NO repair — the scatter's
+                                # arrival fill of bit-equal values is already
+                                # value-set exact. Real fp8-lineage logits tie in
+                                # the thousands, which used to take the full radix.
+                                # Staging mirrors the head min/max (wcnt + hist
+                                # slots [0..31], both dead here; pairs live at
+                                # 260+).
                                 if tidx == cutlass.Int32(0):
-                                    smem_hist[256] = cutlass.Int32(0)
-                                    smem_hist[257] = need0
-                                    smem_hist[258] = cutlass.Int32(0)
+                                    s_iscalars[0] = cutlass.Int32(0)
+                                nh7 = cutlass.Int32(0)
+                                kmn6 = cutlass.Int32(2147483647)
+                                kmx6 = cutlass.Int32(-2147483648)
+                                it6 = tidx
+                                # The pure-tie pre-check only exists to SKIP the
+                                # repair when the class is bit-uniform. The small
+                                # route below ranks by (key, arrival) and rewrites
+                                # exactly the need0 winner slots, which is already
+                                # value-exact on a pure tie - so for a small class
+                                # the check is a whole candidate walk bought for
+                                # nothing. Only the large-class route, whose radix
+                                # really is worth avoiding, still pays for it.
+                                while it6 < cand_count:
+                                    v6 = smem_keys[it6]
+                                    b6 = cutlass.Int32((v6 - bmin_r) * inv1)
+                                    if b6 < cutlass.Int32(0):
+                                        b6 = cutlass.Int32(0)
+                                    if b6 > cutlass.Int32(kBins - 1):
+                                        b6 = cutlass.Int32(kBins - 1)
+                                    if b6 == b_star:
+                                        s6 = cutlass.Int32((v6 - f_lo) * finv)
+                                        if s6 < cutlass.Int32(0):
+                                            s6 = cutlass.Int32(0)
+                                        if s6 > cutlass.Int32(fbins - 1):
+                                            s6 = cutlass.Int32(fbins - 1)
+                                        if s6 == sb_star:
+                                            k6 = f32_order_key(v6) ^ cutlass.Int32(-2147483648)
+                                            if k6 < kmn6:
+                                                kmn6 = k6
+                                            if k6 > kmx6:
+                                                kmx6 = k6
+                                            # Buffer the member so the compaction
+                                            # below needs no walk of its own. The
+                                            # array is nbuf7 = kC/num_threads
+                                            # deep and every store is an unrolled
+                                            # dynamic-index search over it, which
+                                            # is why only the large class - whose
+                                            # alternative is a 4-level radix -
+                                            # pays for it.
+                                            for sl7 in cutlass.range_constexpr(nbuf7):
+                                                if cutlass.Int32(sl7) == nh7:
+                                                    rv7[sl7] = v6
+                                                    ri7[sl7] = smem_vals[it6]
+                                            nh7 = nh7 + cutlass.Int32(1)
+                                    it6 = it6 + cutlass.Int32(num_threads)
+                                kmn6 = cute.arch.warp_redux_sync(kmn6, "min")
+                                kmx6 = cute.arch.warp_redux_sync(kmx6, "max")
+                                if lane == cutlass.Int32(0):
+                                    smem_wcnt[warp_id] = kmn6
+                                    smem_hist[warp_id] = kmx6
                                 cute.arch.barrier()
-                                for lvl2 in cutlass.range_constexpr(4):
-                                    shift2 = cutlass.const_expr(24 - 8 * lvl2)
-                                    iz3 = tidx
-                                    while iz3 < cutlass.Int32(256):
-                                        smem_hist[iz3] = cutlass.Int32(0)
-                                        iz3 = iz3 + cutlass.Int32(num_threads)
-                                    cute.arch.barrier()
-                                    uthr_c2 = smem_hist[256]
-                                    ic2 = tidx
-                                    while ic2 < cnt_strad:
-                                        uk3 = f32_order_key(smem_keys[ic2])
-                                        pm2 = cutlass.Int32(1)
-                                        if cutlass.const_expr(lvl2 > 0):
-                                            if (uk3 >> cutlass.Int32(shift2 + 8)) != (
-                                                uthr_c2 >> cutlass.Int32(shift2 + 8)
-                                            ):
-                                                pm2 = cutlass.Int32(0)
-                                        if pm2 == cutlass.Int32(1):
-                                            dg2 = (uk3 >> cutlass.Int32(shift2)) & cutlass.Int32(
-                                                0xFF
+                                # lane-parallel cross-warp fold: lane w holds
+                                # slot w and one warp reduce settles it, instead
+                                # of every thread walking all num_warps slots of
+                                # two arrays with dependent SMEM reads. Same
+                                # inputs in the same order on every warp, so the
+                                # result stays bit-identical and leaderless.
+                                pa8 = cutlass.Int32(2147483647)
+                                pb8 = cutlass.Int32(-2147483648)
+                                if lane < cutlass.Int32(self.num_warps):
+                                    pa8 = smem_wcnt[lane]
+                                    pb8 = smem_hist[lane]
+                                kmn7 = cute.arch.warp_redux_sync(pa8, "min")
+                                kmx7 = cute.arch.warp_redux_sync(pb8, "max")
+                                if kmn7 == kmx7:
+                                    fast_done = cutlass.Int32(1)
+                                if fast_done == cutlass.Int32(0):
+                                    # mixed class: compact the members buffered by
+                                    # the pure-tie pass above into
+                                    # smem_keys/vals[0..cnt_strad). The buffering
+                                    # pass already read every candidate it needs,
+                                    # and the staging barrier above orders those
+                                    # reads before these writes, so the second
+                                    # full-candidate walk (and its barrier) is
+                                    # gone. Repair cost is a function of the CLASS
+                                    # size only, for any class size.
+                                    # warp-aggregated claim: intra-warp exclusive
+                                    # prefix via shfl scan + ONE atomic per warp
+                                    # (a thousand same-address claims serialize
+                                    # and scale with the class size)
+                                    pf7 = nh7
+                                    for so3 in cutlass.range_constexpr(5):
+                                        oth3 = cute.arch.shuffle_sync_up(
+                                            pf7, cutlass.Int32(1 << so3), mask_and_clamp=0
+                                        )
+                                        if lane >= cutlass.Int32(1 << so3):
+                                            pf7 = pf7 + oth3
+                                    tot7 = cute.arch.shuffle_sync(pf7, cutlass.Int32(31))
+                                    wb7 = cutlass.Int32(0)
+                                    if lane == cutlass.Int32(31):
+                                        if tot7 > cutlass.Int32(0):
+                                            wb7 = atomicAdd(
+                                                s_iscalars.iterator + cutlass.Int32(0), tot7
                                             )
-                                            atomicAdd(smem_hist.iterator + dg2, cutlass.Int32(1))
-                                        ic2 = ic2 + cutlass.Int32(num_threads)
+                                    wb7 = cute.arch.shuffle_sync(wb7, cutlass.Int32(31))
+                                    bs7 = wb7 + pf7 - nh7
+                                    for sl8 in cutlass.range_constexpr(nbuf7):
+                                        if cutlass.Int32(sl8) < nh7:
+                                            smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8]
+                                            smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8]
                                     cute.arch.barrier()
-                                    if warp_id == cutlass.Int32(0):
-                                        ws3 = cutlass.Int32(0)
-                                        for jd3 in cutlass.range_constexpr(8):
-                                            di3 = (
-                                                cutlass.Int32(255)
-                                                - lane * cutlass.Int32(8)
-                                                - cutlass.Int32(jd3)
-                                            )
-                                            ws3 = ws3 + smem_hist[di3]
-                                        pre6 = ws3
-                                        for so2 in cutlass.range_constexpr(5):
-                                            oth2 = cute.arch.shuffle_sync_up(
-                                                pre6,
-                                                cutlass.Int32(1 << so2),
-                                                mask_and_clamp=0,
-                                            )
-                                            if lane >= cutlass.Int32(1 << so2):
-                                                pre6 = pre6 + oth2
-                                        needl3 = smem_hist[257]
-                                        if pre6 >= needl3 and (pre6 - ws3) < needl3:
-                                            base5 = pre6 - ws3
-                                            dstar2 = cutlass.Int32(0)
-                                            above5 = base5
-                                            sd5 = cutlass.Int32(0)
-                                            for jd4 in cutlass.range_constexpr(8):
-                                                di4 = (
+                                    # Large class only (the small one is handled
+                                    # above, ahead of the pure-tie pass, so it
+                                    # never reaches here): block-parallel 4-level
+                                    # MSB radix over the compacted class (scans
+                                    # touch class pairs only; warp0 shuffle-scan
+                                    # digit search - 3 block barriers per level
+                                    # instead of 5).
+                                    if tidx == cutlass.Int32(0):
+                                        smem_hist[256] = cutlass.Int32(0)
+                                        smem_hist[257] = need0
+                                        smem_hist[258] = cutlass.Int32(0)
+                                    cute.arch.barrier()
+                                    for lvl2 in cutlass.range_constexpr(4):
+                                        shift2 = cutlass.const_expr(24 - 8 * lvl2)
+                                        iz3 = tidx
+                                        while iz3 < cutlass.Int32(256):
+                                            smem_hist[iz3] = cutlass.Int32(0)
+                                            iz3 = iz3 + cutlass.Int32(num_threads)
+                                        cute.arch.barrier()
+                                        uthr_c2 = smem_hist[256]
+                                        ic2 = tidx
+                                        while ic2 < cnt_strad:
+                                            uk3 = f32_order_key(smem_keys[ic2])
+                                            pm2 = cutlass.Int32(1)
+                                            if cutlass.const_expr(lvl2 > 0):
+                                                if (uk3 >> cutlass.Int32(shift2 + 8)) != (
+                                                    uthr_c2 >> cutlass.Int32(shift2 + 8)
+                                                ):
+                                                    pm2 = cutlass.Int32(0)
+                                            if pm2 == cutlass.Int32(1):
+                                                dg2 = (
+                                                    uk3 >> cutlass.Int32(shift2)
+                                                ) & cutlass.Int32(0xFF)
+                                                atomicAdd(
+                                                    smem_hist.iterator + dg2, cutlass.Int32(1)
+                                                )
+                                            ic2 = ic2 + cutlass.Int32(num_threads)
+                                        cute.arch.barrier()
+                                        if warp_id == cutlass.Int32(0):
+                                            ws3 = cutlass.Int32(0)
+                                            for jd3 in cutlass.range_constexpr(8):
+                                                di3 = (
                                                     cutlass.Int32(255)
                                                     - lane * cutlass.Int32(8)
-                                                    - cutlass.Int32(jd4)
+                                                    - cutlass.Int32(jd3)
                                                 )
-                                                ra5 = base5
-                                                base5 = base5 + smem_hist[di4]
-                                                if base5 >= needl3 and sd5 == cutlass.Int32(0):
-                                                    dstar2 = di4
-                                                    above5 = ra5
-                                                    sd5 = cutlass.Int32(1)
-                                            smem_hist[256] = uthr_c2 | (
-                                                dstar2 << cutlass.Int32(shift2)
-                                            )
-                                            smem_hist[257] = needl3 - above5
-                                            smem_hist[258] = smem_hist[258] + above5
+                                                ws3 = ws3 + smem_hist[di3]
+                                            pre6 = ws3
+                                            for so2 in cutlass.range_constexpr(5):
+                                                oth2 = cute.arch.shuffle_sync_up(
+                                                    pre6,
+                                                    cutlass.Int32(1 << so2),
+                                                    mask_and_clamp=0,
+                                                )
+                                                if lane >= cutlass.Int32(1 << so2):
+                                                    pre6 = pre6 + oth2
+                                            needl3 = smem_hist[257]
+                                            if pre6 >= needl3 and (pre6 - ws3) < needl3:
+                                                base5 = pre6 - ws3
+                                                dstar2 = cutlass.Int32(0)
+                                                above5 = base5
+                                                sd5 = cutlass.Int32(0)
+                                                for jd4 in cutlass.range_constexpr(8):
+                                                    di4 = (
+                                                        cutlass.Int32(255)
+                                                        - lane * cutlass.Int32(8)
+                                                        - cutlass.Int32(jd4)
+                                                    )
+                                                    ra5 = base5
+                                                    base5 = base5 + smem_hist[di4]
+                                                    if base5 >= needl3 and sd5 == cutlass.Int32(0):
+                                                        dstar2 = di4
+                                                        above5 = ra5
+                                                        sd5 = cutlass.Int32(1)
+                                                smem_hist[256] = uthr_c2 | (
+                                                    dstar2 << cutlass.Int32(shift2)
+                                                )
+                                                smem_hist[257] = needl3 - above5
+                                                smem_hist[258] = smem_hist[258] + above5
+                                        cute.arch.barrier()
+                                    u_thr2 = smem_hist[256]
+                                    cnt_ab2 = smem_hist[258]
+                                    need_eq2 = smem_hist[257]
+                                    kthr2 = u_thr2 ^ cutlass.Int32(-2147483648)
+                                    if tidx == cutlass.Int32(0):
+                                        s_iscalars[4] = cutlass.Int32(0)
+                                        s_iscalars[0] = cutlass.Int32(0)
                                     cute.arch.barrier()
-                                u_thr2 = smem_hist[256]
-                                cnt_ab2 = smem_hist[258]
-                                need_eq2 = smem_hist[257]
-                                kthr2 = u_thr2 ^ cutlass.Int32(-2147483648)
-                                if tidx == cutlass.Int32(0):
-                                    s_iscalars[4] = cutlass.Int32(0)
-                                    s_iscalars[0] = cutlass.Int32(0)
-                                cute.arch.barrier()
-                                ir3 = tidx
-                                while ir3 < cnt_strad:
-                                    vv3 = smem_keys[ir3]
-                                    uk4 = f32_order_key(vv3)
-                                    ks4 = uk4 ^ cutlass.Int32(-2147483648)
-                                    if ks4 > kthr2:
-                                        o4 = atomicAdd(
-                                            s_iscalars.iterator + cutlass.Int32(4),
-                                            cutlass.Int32(1),
-                                        )
-                                        pos = rank_above_fine + o4
-                                        if pos < cutlass.Int32(kK):
-                                            if cutlass.const_expr(self.return_output_values):
-                                                output_values_row[pos] = self.dtype(vv3)
-                                            output_indices_row[pos] = smem_vals[ir3]
-                                    elif ks4 == kthr2:
-                                        q4 = atomicAdd(
-                                            s_iscalars.iterator + cutlass.Int32(0),
-                                            cutlass.Int32(1),
-                                        )
-                                        if q4 < need_eq2:
-                                            pos = rank_above_fine + cnt_ab2 + q4
+                                    ir3 = tidx
+                                    while ir3 < cnt_strad:
+                                        vv3 = smem_keys[ir3]
+                                        uk4 = f32_order_key(vv3)
+                                        ks4 = uk4 ^ cutlass.Int32(-2147483648)
+                                        if ks4 > kthr2:
+                                            o4 = atomicAdd(
+                                                s_iscalars.iterator + cutlass.Int32(4),
+                                                cutlass.Int32(1),
+                                            )
+                                            pos = rank_above_fine + o4
                                             if pos < cutlass.Int32(kK):
                                                 if cutlass.const_expr(self.return_output_values):
                                                     output_values_row[pos] = self.dtype(vv3)
                                                 output_indices_row[pos] = smem_vals[ir3]
-                                    ir3 = ir3 + cutlass.Int32(num_threads)
-                                cute.arch.barrier()
+                                        elif ks4 == kthr2:
+                                            q4 = atomicAdd(
+                                                s_iscalars.iterator + cutlass.Int32(0),
+                                                cutlass.Int32(1),
+                                            )
+                                            if q4 < need_eq2:
+                                                pos = rank_above_fine + cnt_ab2 + q4
+                                                if pos < cutlass.Int32(kK):
+                                                    if cutlass.const_expr(
+                                                        self.return_output_values
+                                                    ):
+                                                        output_values_row[pos] = self.dtype(vv3)
+                                                    output_indices_row[pos] = smem_vals[ir3]
+                                        ir3 = ir3 + cutlass.Int32(num_threads)
+                                    cute.arch.barrier()
                     else:
                         need0_s = cutlass.Int32(kK) - rank_above_fine
                         if cnt_strad > need0_s and need0_s > cutlass.Int32(0):

From 7f823dea65f2aeeeccaf6fcc51e3f92c7f653a56 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 5 Aug 2026 21:57:32 -0700
Subject: [PATCH 093/117] [None][perf] GVR: rank the tail class with the block,
 not one warp

The exact-tail repair's small-class route ranked the parked tie class on
warp 0: each lane took a stride of the class and walked the WHOLE class
per element, re-reading each partner out of SMEM and re-deriving its
order key every comparison. Fifteen warps sat at the barrier below while
that happened.

Now each warp owns a stride of the class and its 32 lanes split the
comparisons, so the per-lane trip count is ceil(n/num_warps) *
ceil(n/32) plus one warp reduce - one trip for the class of two that
fires on the captured rows, eight for a class of fifty. Same rank
definition (#{key greater} + #{key equal, earlier slot}), so the
selection is unchanged.

Phase 4 on the pro capture (30 layers, cold steps, device clock,
us/step):

  unit          before   after
  pro 8k         3.15     2.99
  pro 128k       3.22     3.10
  pro 512k       3.88     3.56

Exactness: 30/30 layers on the firing capture; 1947/1947 records over
flash 8k/64k/256k/1024k and pro 1024k on both assist tiers; the full
decode test file passes (708 passed, 144 skipped, 1 xfailed).

This also removes the reason the fine histogram level could not be
dropped: with the old warp-0 rank, collapsing the sub-binning and
handing the whole coarse class to the repair cost more in the tail than
it saved in the fine level (+0.59 against -0.28 measured). With the
block rank the two now net out even, so the fine level stays for
another reason - the coarse class it avoids ranking is 11 at the median
and 49 at the max over 520 captured steps - but the ranking is no
longer what decides it.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 87 ++++++++++---------
 1 file changed, 45 insertions(+), 42 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 07882ed648b6..f6ff8b974ee4 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -4141,55 +4141,58 @@ def phase4_rank_scatter(
                                 # so there is nothing to collect - go straight
                                 # to the rank. No extra candidate walk, no
                                 # staging barrier, nothing riding in registers.
-                                # unroll=1 on the rank loop: the class is tiny,
-                                # so the trip count is 1-2, but an unrolled
-                                # body (plus its remainder ladder) is code
-                                # volume that ONE warp walks cold while every
-                                # other warp waits at the barrier below - the
-                                # fetch latency has nothing to hide behind.
-                                # Keeping it rolled trades a branch per trip
-                                # for a body that fits a cache line or two.
-                                if warp_id == cutlass.Int32(0):
-                                    ie9 = lane
-                                    while ie9 < cnt_strad:
-                                        bi9 = smem_hist[_PAIR_BASE + ie9 + ie9]
-                                        ki9 = f32_order_key(
+                                # Rank the class with the WHOLE BLOCK, not
+                                # one warp. Each warp owns a stride of the
+                                # class and its 32 lanes split the
+                                # comparisons, so the cost is
+                                # ceil(n/num_warps) * ceil(n/32) trips per
+                                # lane plus one warp reduce - 1 trip for the
+                                # class of two that fires today, 8 for a class
+                                # of fifty. The old form had a single warp walk
+                                # n^2/32 steps while the other fifteen waited
+                                # at the barrier below, which is what made a
+                                # larger class unaffordable.
+                                e9 = warp_id
+                                while e9 < cnt_strad:
+                                    be9 = smem_hist[_PAIR_BASE + e9 + e9]
+                                    ke9 = f32_order_key(
+                                        cutlass.Float32(
+                                            llvm.bitcast(cutlass.Float32.mlir_type, be9.ir_value())
+                                        )
+                                    ) ^ cutlass.Int32(-2147483648)
+                                    c9 = cutlass.Int32(0)
+                                    j9 = lane
+                                    while j9 < cnt_strad:
+                                        bj9 = smem_hist[_PAIR_BASE + j9 + j9]
+                                        kj9 = f32_order_key(
                                             cutlass.Float32(
                                                 llvm.bitcast(
-                                                    cutlass.Float32.mlir_type, bi9.ir_value()
+                                                    cutlass.Float32.mlir_type, bj9.ir_value()
                                                 )
                                             )
                                         ) ^ cutlass.Int32(-2147483648)
-                                        r9 = cutlass.Int32(0)
-                                        for j9 in cutlass.range(0, cnt_strad, 1, unroll=1):
-                                            bj9 = smem_hist[_PAIR_BASE + j9 + j9]
-                                            kj9 = f32_order_key(
-                                                cutlass.Float32(
-                                                    llvm.bitcast(
-                                                        cutlass.Float32.mlir_type, bj9.ir_value()
-                                                    )
-                                                )
-                                            ) ^ cutlass.Int32(-2147483648)
-                                            if kj9 > ki9:
-                                                r9 = r9 + cutlass.Int32(1)
-                                            elif kj9 == ki9 and j9 < ie9:
-                                                r9 = r9 + cutlass.Int32(1)
-                                        if r9 < need0:
-                                            pos9 = rank_above_fine + r9
-                                            if pos9 < cutlass.Int32(kK):
-                                                if cutlass.const_expr(self.return_output_values):
-                                                    output_values_row[pos9] = self.dtype(
-                                                        cutlass.Float32(
-                                                            llvm.bitcast(
-                                                                cutlass.Float32.mlir_type,
-                                                                bi9.ir_value(),
-                                                            )
+                                        if kj9 > ke9:
+                                            c9 = c9 + cutlass.Int32(1)
+                                        elif kj9 == ke9 and j9 < e9:
+                                            c9 = c9 + cutlass.Int32(1)
+                                        j9 = j9 + cutlass.Int32(32)
+                                    r9 = self.warp_reduce_sum_i32(c9)
+                                    if lane == cutlass.Int32(0) and r9 < need0:
+                                        pos9 = rank_above_fine + r9
+                                        if pos9 < cutlass.Int32(kK):
+                                            if cutlass.const_expr(self.return_output_values):
+                                                output_values_row[pos9] = self.dtype(
+                                                    cutlass.Float32(
+                                                        llvm.bitcast(
+                                                            cutlass.Float32.mlir_type,
+                                                            be9.ir_value(),
                                                         )
                                                     )
-                                                output_indices_row[pos9] = smem_hist[
-                                                    _PAIR_BASE + ie9 + ie9 + cutlass.Int32(1)
-                                                ]
-                                        ie9 = ie9 + cutlass.Int32(32)
+                                                )
+                                            output_indices_row[pos9] = smem_hist[
+                                                _PAIR_BASE + e9 + e9 + cutlass.Int32(1)
+                                            ]
+                                    e9 = e9 + cutlass.Int32(num_warps)
                                 cute.arch.barrier()
                             else:
                                 # Large class only. Everything below -

From d17a58ac35c9014acba9ff7eddae8f9b1b5f9604 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 8 Aug 2026 09:13:20 -0700
Subject: [PATCH 094/117] [None][perf] GVR routing: the weak band no longer
 holds at n_comp 65536

The weak band hands a mid-length row back to the stock kernel on the
premise that stock's split grid wins there outright. At the top of the
band that premise is now false. Measured on the full grid (kernel-only,
us/step, 176 cells x 61677 steps per arm):

  cell              stock   counts   rungs    list
  flash 256k B8     14.03    11.84   14.67    5.85
  pro   256k B8     15.48    13.10   16.01    7.44

Stock is not the fastest arm there - counts beats it by 2.19 and 2.38.
The list arm is faster still, but it does not pay for itself at this
batch: re-measured on the indexer kernel with the tight lines parked,
the list emission costs +10.75us at batch 8 / ctx 256k against an 8.11us
kernel saving, so LIST_EMIT_MAX_B stays at 4. The counts emission at the
same point costs +0.27us (+0.27 at B4 is 3.02, at B16 it is free) - the
table this band was fitted against charged +2.0us at batch 8, before the
parking made the emission 2.7-5x cheaper. Net +1.9 to +2.1us per step.

Replaying the routed grid over the same measurements:

  cell             before   after
  flash 256k B8     1.119   1.327   (+19%)
  pro   256k B8     1.112   1.314   (+18%)
  whole grid        1.702   1.704

Only the top of the band moves. The interior (49152 <= n_comp < 65536)
is untouched - no capture unit lands there, so it keeps the stock kernel
until someone measures it rather than inheriting a premise that has now
been shown to expire.

The full decode test file passes (708 passed, 144 skipped, 1 xfailed).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/blackwell/top_k/gvr_routing.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index bf7f72eb64b2..ed081c5e7c84 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -104,8 +104,17 @@
 # 1.5us there, which is enough to take the upper half of the band back
 # off the stock kernel (2 cells fall through now, down from 4).
 ASSIST_WEAK_MIN_N = 49152
-ASSIST_WEAK_MAX_N_SMALL_K = 98304  # k <= ASSIST_WEAK_K
-ASSIST_WEAK_MAX_N_LARGE_K = 98304
+# Narrowed again from 98304 to 65536: at n_comp 65536 the band's premise no
+# longer holds. Measured on the full grid (kernel-only, us/step), the stock
+# kernel is NOT the fastest arm there - counts beats it by 2.19 (flash) and
+# 2.38 (pro), and the counts emission at that batch costs only +0.27us on the
+# indexer (re-measured after the tight-line parking; the old table charged
+# +2.0us at batch 8, which is what kept the band this wide). Net +1.9 to
+# +2.1us per step, i.e. those cells go 1.11 -> ~1.31 against the baseline.
+# The interior of the band (49152 <= n_comp < 65536) is untouched: no grid
+# unit lands there, so it stays on the stock kernel until someone measures it.
+ASSIST_WEAK_MAX_N_SMALL_K = 65536  # k <= ASSIST_WEAK_K
+ASSIST_WEAK_MAX_N_LARGE_K = 65536
 ASSIST_WEAK_K = 512
 ASSIST_WEAK_MAX_B = 8
 

From 9466f59150d00768b705605752feaf7ce7964af5 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 8 Aug 2026 10:22:05 -0700
Subject: [PATCH 095/117] [None][perf] GVR: drop phase 4's fine level, at
 compile time

Phase 4's cost is a chain of sequential stages, not work: its total
barely moves when the candidate count triples (575 -> 2048 cost 0.2us)
and it is flat from n_comp 2051 to 131072. So the only way to cut it is
to remove a stage, and the fine 256-bin level is the one that can go:
the straddling COARSE bin holds 11 candidates at the median and 49 at
the max over 520 captured steps, which the tail repair now absorbs
directly since it ranks with the whole block.

Collapsing the sub-binning (finv 0 sends every member of the coarse
class to sub-bin 0 == sb*) makes the scatter park the whole class and
the tail rank it. A class past the pair buffer still falls into the
tail's radix, which handles any size.

The removal has to be at COMPILE time. Gating the same three pieces
behind a runtime branch was measured twice and is a wash - the fine
window only fell 0.89 -> 0.64us because the branch itself costs the
stage. Compiled out, it falls to 0.10us:

  phase 4 window   before   after
  min/max            0.22    0.21
  histogram          0.44    0.43
  coarse search      0.51    0.51
  fine level         0.89    0.10
  scatter            0.62    0.61
  tail repair        0.37    0.70
  phase 4 total      2.99    2.66

Across shapes (device clock, us/step; kernel total in brackets):

  pro 8k      2.99 -> 2.66   [6.32 -> 6.29]
  pro 128k    3.10 -> 2.85   [11.49 -> 10.41]
  pro 512k    3.56 -> 3.19   [22.26 -> 21.69]
  flash 256k       -> 2.74   [      -> 14.40]

Gated to the configuration that owns the small-class tail route, so the
stock kernel's codegen is untouched; p4_no_fine=False restores the old
body for A/B.

Exactness: 30/30 layers on four capture units; 1947/1947 records over
flash 8k/64k/256k/1024k and pro 1024k on both assist tiers; the full
decode test file passes (708 passed, 144 skipped, 1 xfailed).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 140 +++++++++++-------
 1 file changed, 85 insertions(+), 55 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index f6ff8b974ee4..83090f5fba2a 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -318,6 +318,7 @@ def __init__(
         p4_exact_tail: Optional[bool] = None,
         p4_tail_fast: Optional[bool] = None,  # [p4tt]
         p4_tail_v3: Optional[bool] = None,
+        p4_no_fine: Optional[bool] = None,
         p1r_rescue: bool = True,
         num_bins: Optional[int] = None,
         p4_warp_redundant: bool = True,
@@ -789,6 +790,22 @@ def __init__(
                 use_ext_counts or use_ext_cand or ext_rungs or self_scan or enable_block_skip
             )
         self.p4_tail_v3 = bool(p4_tail_v3)
+        # p4_no_fine: drop the 256-bin fine level from phase 4 outright and
+        # let the tail repair rank the whole straddling COARSE bin instead.
+        # Phase 4's cost is a chain of sequential stages, not work: its total
+        # barely moves when the candidate count triples, so the only way to
+        # cut it is to remove a stage. The coarse bin holds 11 candidates at
+        # the median and 49 at the max over 520 captured steps - inside the
+        # tail's pair buffer - and the tail now ranks with the whole block,
+        # so it absorbs the class cheaply. A class past the buffer falls into
+        # the tail's radix, which handles any size. Gated to the same
+        # configuration as that small-class route; the stock kernel keeps the
+        # fine level so its codegen is untouched.
+        if p4_no_fine is None:
+            p4_no_fine = bool(
+                self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3
+            ) and not (self.p4_fine_rangetest or self.p4_scat_rangetest)
+        self.p4_no_fine = bool(p4_no_fine)
         # p1r_rescue: rebuild the refine bracket from the row when the
         # seed bracket is degenerate. ON by default - upstream's identity
         # shortcut is wrong on real data (every request's first decode
@@ -3934,6 +3951,18 @@ def phase4_rank_scatter(
                 # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1)
                 f_lo = bmin_r + cutlass.Float32(b_star) / inv1
                 finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1
+                if cutlass.const_expr(self.p4_no_fine):
+                    # Sub-binning collapsed: finv 0 sends every member of the
+                    # straddling coarse bin to sub-bin 0 == sb*, so the
+                    # scatter parks the whole coarse class and the tail ranks
+                    # it. The re-zero, the build and the search below are not
+                    # traced at all - this is a compile-time removal, not a
+                    # runtime branch, because a runtime branch still costs the
+                    # stage (measured: 0.64us for a fine window whose work was
+                    # gated off).
+                    finv = cutlass.Float32(0.0)
+                    sb_star = cutlass.Int32(0)
+                    rank_above_fine = rank_above
                 # bin b* spans [f_lo, f_hi); the clamped ends fold
                 # out-of-range values into bin 0 and bin kBins-1, so those
                 # two drop the matching side. Only the range-test arms read
@@ -3942,61 +3971,62 @@ def phase4_rank_scatter(
                     f_hi = f_lo + cutlass.Float32(1.0) / inv1
                     lo_edge = b_star == cutlass.Int32(0)
                     hi_edge = b_star == cutlass.Int32(kBins - 1)
-                # re-zero (only fbins slots) + build fine sub-hist of bin-b* cands
-                iz = tidx
-                while iz < cutlass.Int32(fbins):
-                    smem_hist[iz] = cutlass.Int32(0)
-                    iz = iz + cutlass.Int32(num_threads)
-                cute.arch.barrier()
-                if cutlass.const_expr(self.p4_fine_rangetest):
-                    # A candidate belongs to bin b* exactly when its value
-                    # lies in [f_lo, f_hi), so the filter is two compares -
-                    # the bin recompute (subtract + multiply + two clamps)
-                    # per candidate is redundant work. The clamped ends of
-                    # the binning fold out-of-range values INTO bin 0 and
-                    # bin kBins-1, so those two bins drop the matching side
-                    # of the range test to stay bit-identical.
-                    ifb = tidx
-                    while ifb < cand_count:
-                        vf = smem_keys[ifb]
-                        inb = vf >= f_lo and vf < f_hi
-                        if lo_edge:
-                            inb = vf < f_hi
-                        if hi_edge:
-                            inb = vf >= f_lo
-                        if inb:
-                            sb = cutlass.Int32((vf - f_lo) * finv)
-                            if sb < cutlass.Int32(0):
-                                sb = cutlass.Int32(0)
-                            if sb > cutlass.Int32(fbins - 1):
-                                sb = cutlass.Int32(fbins - 1)
-                            atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1))
-                        ifb = ifb + cutlass.Int32(num_threads)
-                else:
-                    ifb = tidx
-                    while ifb < cand_count:
-                        vfo = smem_keys[ifb]
-                        cbo = cutlass.Int32((vfo - bmin_r) * inv1)
-                        if cbo < cutlass.Int32(0):
-                            cbo = cutlass.Int32(0)
-                        if cbo > cutlass.Int32(kBins - 1):
-                            cbo = cutlass.Int32(kBins - 1)
-                        if cbo == b_star:
-                            sbo = cutlass.Int32((vfo - f_lo) * finv)
-                            if sbo < cutlass.Int32(0):
-                                sbo = cutlass.Int32(0)
-                            if sbo > cutlass.Int32(fbins - 1):
-                                sbo = cutlass.Int32(fbins - 1)
-                            atomicAdd(smem_hist.iterator + sbo, cutlass.Int32(1))
-                        ifb = ifb + cutlass.Int32(num_threads)
-                cute.arch.barrier()
-                # fine sub-bin search, resolved lane-parallel on every
-                # warp (see _p4_fine_rw): the two publish barriers and the
-                # fbins/num_warps-deep single-lane walk that used to sit
-                # here are gone, and the answer comes back in registers.
-                sb_star, rank_above_fine = self._p4_fine_rw(
-                    smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane
-                )
+                if cutlass.const_expr(not self.p4_no_fine):
+                    # re-zero (only fbins slots) + build fine sub-hist of bin-b* cands
+                    iz = tidx
+                    while iz < cutlass.Int32(fbins):
+                        smem_hist[iz] = cutlass.Int32(0)
+                        iz = iz + cutlass.Int32(num_threads)
+                    cute.arch.barrier()
+                    if cutlass.const_expr(self.p4_fine_rangetest):
+                        # A candidate belongs to bin b* exactly when its value
+                        # lies in [f_lo, f_hi), so the filter is two compares -
+                        # the bin recompute (subtract + multiply + two clamps)
+                        # per candidate is redundant work. The clamped ends of
+                        # the binning fold out-of-range values INTO bin 0 and
+                        # bin kBins-1, so those two bins drop the matching side
+                        # of the range test to stay bit-identical.
+                        ifb = tidx
+                        while ifb < cand_count:
+                            vf = smem_keys[ifb]
+                            inb = vf >= f_lo and vf < f_hi
+                            if lo_edge:
+                                inb = vf < f_hi
+                            if hi_edge:
+                                inb = vf >= f_lo
+                            if inb:
+                                sb = cutlass.Int32((vf - f_lo) * finv)
+                                if sb < cutlass.Int32(0):
+                                    sb = cutlass.Int32(0)
+                                if sb > cutlass.Int32(fbins - 1):
+                                    sb = cutlass.Int32(fbins - 1)
+                                atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1))
+                            ifb = ifb + cutlass.Int32(num_threads)
+                    else:
+                        ifb = tidx
+                        while ifb < cand_count:
+                            vfo = smem_keys[ifb]
+                            cbo = cutlass.Int32((vfo - bmin_r) * inv1)
+                            if cbo < cutlass.Int32(0):
+                                cbo = cutlass.Int32(0)
+                            if cbo > cutlass.Int32(kBins - 1):
+                                cbo = cutlass.Int32(kBins - 1)
+                            if cbo == b_star:
+                                sbo = cutlass.Int32((vfo - f_lo) * finv)
+                                if sbo < cutlass.Int32(0):
+                                    sbo = cutlass.Int32(0)
+                                if sbo > cutlass.Int32(fbins - 1):
+                                    sbo = cutlass.Int32(fbins - 1)
+                                atomicAdd(smem_hist.iterator + sbo, cutlass.Int32(1))
+                            ifb = ifb + cutlass.Int32(num_threads)
+                    cute.arch.barrier()
+                    # fine sub-bin search, resolved lane-parallel on every
+                    # warp (see _p4_fine_rw): the two publish barriers and the
+                    # fbins/num_warps-deep single-lane walk that used to sit
+                    # here are gone, and the answer comes back in registers.
+                    sb_star, rank_above_fine = self._p4_fine_rw(
+                        smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane
+                    )
                 if tidx == cutlass.Int32(0):
                     s_iscalars[4] = cutlass.Int32(0)  # cnt_above
                     s_iscalars[0] = cutlass.Int32(0)  # cnt_mid (b*, sub>sb*)

From 53836075066be258b60fff1195ffc0f64e809688 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sat, 8 Aug 2026 11:31:56 -0700
Subject: [PATCH 096/117] [None][perf] GVR: collect the dead weight the fine
 level left behind

With the fine level compiled out, three things it fed are dead but were
still being computed every step:

  * the scatter's sub-bin recompute - a subtract, a multiply and two
    clamps per candidate that lands in the straddling bin, whose result
    is a constant 0 == sb* once the sub-binning is collapsed;
  * f_lo, whose only consumer is that recompute, and which costs an
    fp32 divide;
  * the min side of phase 4's range scan. The accepted threshold is an
    EXACT lower bound - every candidate got there by comparing >=
    against it - so the scan only has to find the max. That drops half
    the walk, one warp reduce and one arm of the cross-warp fold. The
    window it widens is at the BOTTOM, and the k-th lives at the top.

Phase 4 on the pro 8k capture (device clock, us/step):

  window          fine kept   fine dropped   + this
  min/max              0.22           0.21     0.15
  histogram            0.44           0.43     0.33
  coarse search        0.51           0.51     0.52
  fine level           0.89           0.10     0.04
  scatter              0.62           0.61     0.60
  tail repair          0.37           0.70     0.65
  total                2.99           2.66     2.38

Same node, same binary, p4_no_fine off vs on (device clock, us/step):

  unit           fine kept   dropped
  pro 8k              3.05      2.38
  pro 128k            3.14      3.23
  pro 512k            3.58      3.31
  flash 256k          3.11      2.75

pro 128k is the one shape that does not want it: its collection line
sits at rank 2048 rather than 1536, and with the sub-binning collapsed
every member of the straddling class takes the same atomic claim, so
the scatter goes 0.68 -> 1.01 and the tail 0.34 -> 0.73 there. The other
three shapes pay neither. Averaged over the four it is -0.31us/step;
the knob (p4_no_fine) is there if a shape-aware default turns out to be
worth it.

Exactness: 30/30 layers on four capture units; 2277/2277 records over
flash 8k/64k/256k and pro 1024k on both assist tiers; the full decode
test file passes (708 passed, 144 skipped, 1 xfailed).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 56 ++++++++++++++-----
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  4 ++
 2 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 83090f5fba2a..766d8a272a97 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3820,18 +3820,31 @@ def phase4_rank_scatter(
                         bmax_r = bmin_r + cutlass.Float32(1e-6)
             if run_stock_range:
                 # ---- block min/max over candidates ----
+                # The accepted threshold is an EXACT lower bound: every
+                # candidate got here by comparing >= against it. So the min
+                # side of this scan buys nothing but a slightly tighter
+                # window - and the window only has to be tight where the
+                # k-th lives, which is the TOP of it. Taking the threshold
+                # instead drops half the scan, one warp reduce and one arm
+                # of the cross-warp fold. Only the assist tiers have a
+                # published threshold, hence the gate.
+                use_thr_min = cutlass.const_expr(self.p4_no_fine)
                 local_cmin = cutlass.Float32(self.FLT_MAX)
                 local_cmax = cutlass.Float32(self.NEG_FLT_MAX)
                 i5 = tidx
                 while i5 < cand_count:
                     v = smem_keys[i5]
-                    local_cmin = _fmin_f32_inline(local_cmin, v)
+                    if cutlass.const_expr(not use_thr_min):
+                        local_cmin = _fmin_f32_inline(local_cmin, v)
                     local_cmax = cute.arch.fmax(local_cmax, v)
                     i5 = i5 + cutlass.Int32(num_threads)
-                cmin = self.warp_reduce_min_f32(local_cmin)
+                cmin = cutlass.Float32(0.0)
+                if cutlass.const_expr(not use_thr_min):
+                    cmin = self.warp_reduce_min_f32(local_cmin)
                 cmax = self.warp_reduce_max_f32(local_cmax)
                 if lane == cutlass.Int32(0):
-                    smem_wcnt[warp_id] = float_as_uint32(cmin)
+                    if cutlass.const_expr(not use_thr_min):
+                        smem_wcnt[warp_id] = float_as_uint32(cmin)
                     smem_hist[warp_id] = float_as_uint32(cmax)
                 cute.arch.barrier()
                 # lane-parallel cross-warp fold: lane w holds slot w and one
@@ -3842,13 +3855,17 @@ def phase4_rank_scatter(
                 pmn = cutlass.Float32(self.FLT_MAX)
                 pmx = cutlass.Float32(self.NEG_FLT_MAX)
                 if lane < cutlass.Int32(self.num_warps):
-                    pmn = cutlass.Float32(
-                        llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[lane].ir_value())
-                    )
+                    if cutlass.const_expr(not use_thr_min):
+                        pmn = cutlass.Float32(
+                            llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[lane].ir_value())
+                        )
                     pmx = cutlass.Float32(
                         llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[lane].ir_value())
                     )
-                bmin_r = _fmin_f32_inline(bmin_r, self.warp_reduce_min_f32(pmn))
+                if cutlass.const_expr(use_thr_min):
+                    bmin_r = s_thr[0]
+                else:
+                    bmin_r = _fmin_f32_inline(bmin_r, self.warp_reduce_min_f32(pmn))
                 bmax_r = cute.arch.fmax(bmax_r, self.warp_reduce_max_f32(pmx))
                 if bmax_r <= bmin_r:
                     bmax_r = bmin_r + cutlass.Float32(1e-6)
@@ -3949,8 +3966,11 @@ def phase4_rank_scatter(
                 # enough to resolve the straddling bin to ≤1 distinct value.
                 fbins = cutlass.const_expr(256)
                 # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1)
-                f_lo = bmin_r + cutlass.Float32(b_star) / inv1
-                finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1
+                f_lo = cutlass.Float32(0.0)
+                finv = cutlass.Float32(0.0)
+                if cutlass.const_expr(not self.p4_no_fine):
+                    f_lo = bmin_r + cutlass.Float32(b_star) / inv1
+                    finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1
                 if cutlass.const_expr(self.p4_no_fine):
                     # Sub-binning collapsed: finv 0 sends every member of the
                     # straddling coarse bin to sub-bin 0 == sb*, so the
@@ -4070,11 +4090,19 @@ def phase4_rank_scatter(
                                 output_values_row[pos] = self.dtype(v)
                             output_indices_row[pos] = smem_vals[isc]
                     elif bin_i == b_star:
-                        sb = cutlass.Int32((v - f_lo) * finv)
-                        if sb < cutlass.Int32(0):
-                            sb = cutlass.Int32(0)
-                        if sb > cutlass.Int32(fbins - 1):
-                            sb = cutlass.Int32(fbins - 1)
+                        # With the fine level compiled out the sub-bin is a
+                        # constant 0 == sb*, so the whole three-way split
+                        # below collapses to the park arm and the per-
+                        # candidate recompute (a subtract, a multiply and two
+                        # clamps) is dead. Keep it a compile-time constant so
+                        # the scatter does not carry it.
+                        sb = cutlass.Int32(0)
+                        if cutlass.const_expr(not self.p4_no_fine):
+                            sb = cutlass.Int32((v - f_lo) * finv)
+                            if sb < cutlass.Int32(0):
+                                sb = cutlass.Int32(0)
+                            if sb > cutlass.Int32(fbins - 1):
+                                sb = cutlass.Int32(fbins - 1)
                         if sb > sb_star:
                             o = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), cutlass.Int32(1))
                             pos = rank_above + o
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 81b621279fec..45891a0f3a05 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -68,6 +68,7 @@ def _compile(
     enable_block_skip: bool = False,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
+    p4_no_fine: "bool | None" = None,
     p4_exact_tail: "bool | None" = None,
     p4_tail_fast: "bool | None" = None,
     p1r_rescue: bool = True,
@@ -213,6 +214,7 @@ def _compile(
         enable_block_skip=enable_block_skip,
         pdl_wait_late=pdl_wait_late,
         p4_tail_v3=p4_tail_v3,
+        p4_no_fine=p4_no_fine,
         p4_exact_tail=p4_exact_tail,
         p4_tail_fast=p4_tail_fast,
         p1r_rescue=p1r_rescue,
@@ -682,6 +684,7 @@ def gvr_topk_decode(
     p2_warp_redundant: bool = True,
     pdl_wait_late: bool = True,
     p4_tail_v3: "bool | None" = None,
+    p4_no_fine: "bool | None" = None,
     p4_exact_tail: "bool | None" = None,
     p4_tail_fast: "bool | None" = None,
     p1r_rescue: bool = True,
@@ -976,6 +979,7 @@ def gvr_topk_decode(
         enable_block_skip,
         pdl_wait_late,
         p4_tail_v3,
+        p4_no_fine,
         p4_exact_tail,
         p4_tail_fast,
         p1r_rescue,

From 1bd86c026e4fe61f59a1636d74f958aeaa20de3d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 9 Aug 2026 03:58:22 -0700
Subject: [PATCH 097/117] [None][test] give the GVR decode unittest entry 120
 minutes on B300

The whole test file runs as ONE integration item
(test_unittests_v2[unittest/...test_cute_dsl_gvr_topk_decode.py]), and
an item without a TIMEOUT directive gets the harness default of 3600s
with --timeout-method=thread - which kills the whole pytest process on
breach. The file compiles dozens of CuTe DSL kernel variants and takes
~43 minutes on a dedicated B300; on the flex CI runners (fewer CPUs, so
slower kernel compilation) it crosses the hour.

That is exactly the observed failure shape, three CI runs in a row:
precisely TWO "Test terminated unexpectedly" results - the file-level
item plus whichever inner case the clock happened to land on
(p4_exact_tail_ties[...512-16384-1] on 08/03,
launch_autoconfig[dtype2-1024-131072-2] on 08/06,
launch_autoconfig[dtype0-2048-32768-1] on 08/09) - with no assertion
failures, and the stage total at 84 minutes: ~24 minutes for the other
entries plus the kill at the 3600s mark. The same file passes 708/144/1
in ~43 minutes on a dedicated B300 node, three times across different
commits.

TIMEOUT (120) follows the existing precedent in the same list
(thop/parallel TIMEOUT (90)). B200 lists do not carry this entry, which
is why only B300 stages ever showed the failure.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tests/integration/test_lists/test-db/l0_b300.yml     | 2 +-
 tests/integration/test_lists/test-db/l0_dgx_b300.yml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml
index 95938a31cd98..c7b935df83fc 100644
--- a/tests/integration/test_lists/test-db/l0_b300.yml
+++ b/tests/integration/test_lists/test-db/l0_b300.yml
@@ -19,7 +19,7 @@ l0_b300:
   - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
   - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py
   - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
-  - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+  - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120)
   - unittest/_torch/thop/parallel TIMEOUT (90)
   - unittest/_torch/thop/serial
   - unittest/_torch/executor # 250s
diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300.yml b/tests/integration/test_lists/test-db/l0_dgx_b300.yml
index 736319d8409f..0dac8a4fd41c 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b300.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b300.yml
@@ -18,7 +18,7 @@ l0_dgx_b300:
   - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
   - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py
   - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
-  - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+  - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120)
   - unittest/_torch/executor
   # ------------- modules (multi-GPU) ---------------
   - unittest/_torch/modules/test_mla_helix.py

From fda4f39ac5885982ccf3cc927a440e445d1b7faf Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 9 Aug 2026 10:48:51 -0700
Subject: [PATCH 098/117] [None][fix] correct GVR ext list-tier consumption and
 gate its state

The production list tier (TRTLLM_GVR_EXT=1, batch <= 4, n_comp >= 64k)
was silently broken end to end; the new ext-tier unit tests run the
emitter contract verbatim and caught four defects:

- the consumer kwargs never passed accept_cap, so the kernel derived
  segment bases from kC while the emitter wrote at LIST_SEG_A bases
- the kernel clamped segment bases/extents by min(accept_cap, kC); the
  emitter geometry is accept_cap (K512's kC=3072 diet read bases
  0/3072/6144 against writes at 0/8192/16384). Only the admission
  bound keeps the kC clamp
- the Phase-1 skip preview admitted parked seed rows (t1/t2 = 1e30);
  with no in-band count the raw lines reached the threshold scratch
  and the defensive rerun spun on the poisoned bracket, emitting -1
  for the whole row. The preview now requires real lines; parked rows
  keep Phase 1 and the list take is unaffected
- the clamped-histogram fallback bracketed segment C by a parked upper
  line; it now reuses the need_max segment-max pass instead

Also route on the tier emitted THIS step (emission and consumption
happen inside the same forward; the routes-on-previous-tier member is
gone), skip the ~5GB candidate buffers when the engine's static length
cannot reach the list tier, skip seed-row maintenance on rungs steps,
and drop the TRTLLM_GVR_DUMP diagnostic block from the decode hot path.

Why nothing caught this before: the harness always ran three live
lines with B* <= kC, and e2e at batch 128 never fires the list tier
(LIST_EMIT_MAX_B=4). New tests: packed seed row (K 512/1024/2048 x
band/fat/miss/inf), candidate list (hit/pads/hist/void/bucketed at
production geometry), 32-grain block-max records (exact/pad_inf).
Full decode suite on B200: 727 passed / 144 skipped / 1 xfailed.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../attention_backend/sparse/dsa/indexer.py   |  83 +---
 .../attention_backend/sparse/gvr_ext.py       |  87 ++--
 .../blackwell/top_k/gvr_topk_decode.py        | 453 +++++++-----------
 .../sparse/test_cute_dsl_gvr_topk_decode.py   | 223 +++++++++
 4 files changed, 458 insertions(+), 388 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index 340bf32927d3..f3e0c32ca9e3 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -683,10 +683,8 @@ def __init__(
         self.use_cute_dsl_paged_mqa_logits = (
             sparse_params.use_cute_dsl_paged_mqa_logits and IS_CUTLASS_DSL_AVAILABLE
         )
-        # GVR emission-assisted decode (opt-in, experimental): the FP4
-        # indexer epilogue emits seed counts / the bucketed candidate
-        # list and the top-k consumes them (see gvr_ext / gvr_routing).
-        # Env-gated so the default path stays byte-identical.
+        # GVR emission-assisted decode (opt-in, experimental): the FP4 indexer
+        # epilogue emits candidates for the top-k (see gvr_ext / gvr_routing).
         self.use_gvr_ext = (
             os.environ.get("TRTLLM_GVR_EXT", "0") == "1"
             and self.use_cute_dsl_topk
@@ -1716,31 +1714,32 @@ def sparse_attn_indexer(
 
                     gvr_emit_kwargs = {}
                     if self.use_gvr_ext and next_n == 1 and not dsl_atom_split:
-                        from ..gvr_ext import GvrExtState
+                        from ..gvr_ext import LIST_EMIT_MIN_N, GvrExtState
 
+                        # indexer_max_seq_len is already the compressed length
+                        # (get_indexer_max_seq_len divides); do not divide again.
+                        n_comp = indexer_max_seq_len
                         if self._gvr_ext is None:
                             self._gvr_ext = GvrExtState(
                                 max_rows=metadata.max_num_sequences,
                                 top_k=self.index_topk,
                                 device=q_fp8.device,
+                                # the list tier is only reachable at
+                                # n_comp >= LIST_EMIT_MIN_N; skip its large
+                                # candidate buffers when the engine's static
+                                # length cannot get there
+                                enable_list_tier=n_comp >= LIST_EMIT_MIN_N,
                             )
                         st = self._gvr_ext
-                        # indexer_max_seq_len is ALREADY the compressed
-                        # length - get_indexer_max_seq_len divides by the
-                        # compress ratio, and it is what goes to the top-k
-                        # as max_seq_len below. Dividing again shifted every
-                        # routing threshold by two doublings: the list tier
-                        # became unreachable and 30 of the 154 grid shapes
-                        # fell back to the stock kernel instead of 4.
-                        n_comp = indexer_max_seq_len
                         emit_tier, self._gvr_route = st.plan(
                             batch_size,
                             n_comp,
                             torch.cuda.get_device_properties(q_fp8.device).multi_processor_count,
                             compress_ratio=max(self.compress_ratio, 1),
                         )
-                        st.update_seed_rows(batch_size, emit_tier)
-                        gvr_emit_kwargs = st.indexer_emit_kwargs(emit_tier, batch_size)
+                        if emit_tier in ("counts", "list"):
+                            st.update_seed_rows(batch_size, emit_tier)
+                            gvr_emit_kwargs = st.indexer_emit_kwargs(emit_tier, batch_size)
                         if self._gvr_route.attach_block_max or emit_tier in (
                             "counts",
                             "list",
@@ -1827,11 +1826,8 @@ def sparse_attn_indexer(
                     if not metadata.use_cute_dsl_topk:
                         heuristic_scratch = metadata.heuristic_scratch_values[:num_gen_tokens]
 
-                # tier "none": too short for any assist to pay, or inside
-                # the mid-row band where the stock kernel's split grid wins
-                # (see gvr_routing) - fall through to the stock branch. The
-                # untouched ext state reads as cold start, which its closed
-                # loop already handles.
+                # tier "none" (see gvr_routing) falls through to the stock
+                # branch; the untouched ext state reads as a cold start there.
                 if (
                     self.use_gvr_ext
                     and self._gvr_ext is not None
@@ -1840,16 +1836,13 @@ def sparse_attn_indexer(
                     and next_n == 1
                     and num_gen_tokens <= 256
                 ):
-                    # emission-assisted GVR: consume what the indexer
-                    # epilogue emitted this step (packed seed row /
-                    # bucketed list / block_max per the picked route)
+                    # emission-assisted GVR: consume what the indexer epilogue
+                    # emitted this step (seed row / list / block_max per route)
                     st = self._gvr_ext
                     out_slice = topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :]
-                    # the GVR op takes RAW-domain seq_lens (its kernel
-                    # ceil-divides by compress_ratio internally). On the
-                    # DSL indexer path the live compressed lengths are in
-                    # gen_indexer_kv_lens_cuda_runtime (kv_lens_cuda_2d
-                    # stays zero there - same trap as the indexer call).
+                    # GVR op takes RAW-domain seq_lens (kernel ceil-divides by
+                    # compress_ratio); on the DSL path the live compressed lens
+                    # are gen_indexer_kv_lens_cuda_runtime, kv_lens_cuda_2d is 0.
                     if self.compress_ratio > 1:
                         gvr_lens = metadata.gen_indexer_kv_lens_cuda_runtime
                         assert gvr_lens is not None
@@ -1872,34 +1865,6 @@ def sparse_attn_indexer(
                         max_seq_len=indexer_max_seq_len,
                         **ext_kw,
                     )
-                    # diagnostic dump (NO_GRAPH runs only): capture the
-                    # inputs AND the in-run selection so an offline pass
-                    # can check score-multiset equality vs torch.topk on
-                    # the exact production path. prev_topk still holds
-                    # the pre-op value here (copy-back is below). The
-                    # min-seq guard skips warmup/dummy rows (n < K).
-                    # TRTLLM_GVR_DUMP holds the target directory.
-                    gvr_dump_dir = os.environ.get("TRTLLM_GVR_DUMP")
-                    if (
-                        gvr_dump_dir
-                        and getattr(self, "_gvr_dumped", 0) < 6
-                        and int(seq_1d.min()) >= self.index_topk * self.compress_ratio
-                    ):
-                        os.makedirs(gvr_dump_dir, exist_ok=True)
-                        self._gvr_dumped = getattr(self, "_gvr_dumped", 0) + 1
-                        torch.save(
-                            {
-                                "logits": logits_decode.clone().cpu(),
-                                "seq_raw": seq_1d.clone().cpu(),
-                                "pre": st.prev_topk[:num_gen_tokens].clone().cpu(),
-                                "sel": out_slice.clone().cpu(),
-                                "layer": self.layer_idx,
-                            },
-                            os.path.join(
-                                gvr_dump_dir,
-                                f"dump_l{self.layer_idx}_{self._gvr_dumped}.pt",
-                            ),
-                        )
                     st.prev_topk[:num_gen_tokens].copy_(out_slice)
                 elif self.use_cute_dsl_topk and self._enable_heuristic_topk:
                     torch.ops.trtllm.cute_dsl_gvr_topk_decode(
@@ -1914,10 +1879,8 @@ def sparse_attn_indexer(
                         order_row=metadata.kv_lens_row_reorder,
                     )
                 elif self.use_cute_dsl_topk and (self.compress_ratio == 1 or next_n == 1):
-                    # request-level seq_lens must be 1-D; on the DSL
-                    # indexer path the live compressed lengths are in
-                    # gen_indexer_kv_lens_cuda_runtime (kv_lens_cuda_2d
-                    # stays zero there)
+                    # seq_lens must be 1-D; on the DSL path the live compressed
+                    # lens are gen_indexer_kv_lens_cuda_runtime (2d buf is zero)
                     if self.compress_ratio > 1:
                         radix_lens = (
                             metadata.gen_indexer_kv_lens_cuda_runtime
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
index 525025eb35cf..034b2d3a212f 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
@@ -16,47 +16,45 @@
 """Emission-assisted GVR top-k state for the DSA decode path.
 
 Owns the persistent (graph-address-stable) buffers the emission tiers
-ride on, the device-side closed-loop seed-row update (pure tensor ops:
-CUDA-graph capturable, validated 11/11-step replay-exact against eager)
-and the per-step routing decision. All of it is opt-in: without the
-flag the DSA decode path is byte-identical to before.
+ride on, the device-side closed-loop seed-row update (pure tensor ops,
+CUDA-graph capturable) and the per-step routing decision. Opt-in:
+without the flag the DSA decode path is unchanged.
 
 Tier semantics (see gvr_routing):
   * this step's TOP-K consumes what the PREVIOUS step's indexer
     epilogue emitted;
   * this step's INDEXER emits what the routing planned for the NEXT
-    step. N changes by at most one slot per step, so tier flapping is
-    a non-issue.
+    step.
 """
 
 from typing import Optional
 
 import torch
 
-from ...cute_dsl_kernels.blackwell.top_k.gvr_routing import TopkRoute, pick_config, plan_emission
+from ...cute_dsl_kernels.blackwell.top_k.gvr_routing import (
+    LIST_EMIT_MIN_N,
+    TopkRoute,
+    pick_config,
+    plan_emission,
+)
 
-# Bucketed list geometry (validated defaults: B* = 8192 segment cap,
-# 24576-entry C segment; see the f15/f17 sweeps).
+# Bucketed candidate-list geometry: two tight segments of LIST_SEG_A
+# entries plus a LIST_CAP_C-entry loose segment.
 LIST_SEG_A = 8192
 LIST_CAP_C = 24576
 LIST_WIDTH = 2 * LIST_SEG_A + LIST_CAP_C
 
-# Closed-loop line derivation around the published k-th anchor: t1
-# hugs the k-th value from below, t0/t2 guard by the (anchor - kth)
-# span. Matches the graph_test.py-validated update.
+# Closed-loop lines around the published k-th anchor: t1 hugs the k-th
+# value from below; t0/t2 guard by the (anchor - kth) span.
+__all__ = ["GvrExtState", "LIST_EMIT_MIN_N", "LIST_PARK_LINE"]
+
 GUARD_LO = 2.0
 GUARD_HI = 0.5
 
 # List tier only: park the two tight lines above any score so every
-# admitted entry lands in the loosest segment, which is the one that
-# already claims its slots through a per-warp window instead of an
-# exact ballot. Measured on the indexer kernel (nsys kernel-only): the
-# emission cost drops 2.7-5x (e.g. batch 16 / ctx 512k, +88.7us ->
-# +15.0us) and the top-k side stays at parity, because the consumer
-# reads the segment counts from the control row and finds the two tight
-# segments empty. Any finite value above the score range works; the
-# kernel's eligibility check only needs the three lines to be
-# increasing and the loosest one finite.
+# admitted entry lands in the loosest segment. Any finite value above
+# the score range works; the kernel's eligibility check only needs the
+# three lines increasing and the loosest one finite.
 LIST_PARK_LINE = 1.0e30
 
 
@@ -81,27 +79,20 @@ def __init__(
             self.cand_idx = torch.zeros((max_rows, LIST_WIDTH), dtype=torch.int32, device=device)
             self.cand_ctl = torch.zeros((max_rows, 4), dtype=torch.int32, device=device)
             self.cand_cur = torch.zeros((max_rows, 4), dtype=torch.int32, device=device)
-        # GVR warm-start feedback: this layer's previous-step top-k
-        # (same stable-address feedback-loop shape as
-        # heuristic_prev_topk; zero-init -> first step's pre_idx points
-        # at index 0, a valid benign candidate)
+        # previous-step top-k feedback (address-stable; zero-init ->
+        # first step's pre_idx points at index 0, a benign candidate)
         self.prev_topk = torch.zeros((max_rows, top_k), dtype=torch.int32, device=device)
         # block_max prefix ([rows, nb_pad*4] fp32 warp-partials),
         # allocated lazily once max_seq_len is known
         self.block_max: Optional[torch.Tensor] = None
-        # tier the PREVIOUS indexer call emitted (what this step's
-        # top-k may consume); "rungs" until the first emission lands
-        self.emitted_tier = "rungs"
 
     def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
         nb4 = ((max_seq_len + 255) // 256 * 256) // 128 * 4
         # exact width: the runner asserts shape == (rows, nrec), so a
         # wider reused buffer would trip it
         if self.block_max is None or self.block_max.shape[1] != nb4:
-            # max_seq_len is engine-static, so this allocates once on the
-            # first (eager warmup) step; allocating inside CUDA graph
-            # capture would bake a dangling address into the graph, so
-            # fail loudly instead of corrupting the capture.
+            # allocating during CUDA graph capture would bake a dangling
+            # address into the graph, so fail loudly instead
             if torch.cuda.is_current_stream_capturing():
                 raise RuntimeError(
                     "GvrExtState.ensure_block_max: (re)allocation requested "
@@ -116,8 +107,8 @@ def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
     def plan(
         self, batch: int, n_comp: int, num_sms: int, compress_ratio: int = 4
     ) -> tuple[str, TopkRoute]:
-        """Route this step: (tier to EMIT next, launch knobs to CONSUME
-        what was emitted last step)."""
+        """Route this step: (tier the epilogue emits, launch knobs the
+        top-k consumes it with)."""
         emit_tier = plan_emission(
             batch, n_comp, self.top_k, have_epilogue=True, compress_ratio=compress_ratio
         )
@@ -125,16 +116,17 @@ def plan(
             # constructed with enable_list_tier=False: no candidate
             # buffers to emit into, demote to the counts tier
             emit_tier = "counts"
-        route = pick_config(self.emitted_tier, batch, n_comp, self.top_k, num_sms)
+        # emission and consumption happen inside the SAME forward (zero,
+        # emit, consume), so the consumer routes on this step's tier
+        route = pick_config(emit_tier, batch, n_comp, self.top_k, num_sms)
         return emit_tier, route
 
     def update_seed_rows(self, num_rows: int, emit_tier: str = "counts") -> None:
         """Device-side closed-loop line update from the last publish.
 
-        Pure tensor ops (graph-capturable). Rows whose xstate is not
-        valid (col 0 == 0, e.g. cold start) get non-finite lines, which
-        the kernel's validity guard routes to the stock path - the
-        closed loop never rides on host data quality.
+        Pure tensor ops (graph-capturable). Rows with invalid xstate
+        (col 0 == 0, e.g. cold start) get non-finite lines, which the
+        kernel's validity guard routes to the stock path.
         """
         s = self.seed_row[:num_rows]
         x = self.xstate[:num_rows]
@@ -170,14 +162,13 @@ def indexer_emit_kwargs(self, emit_tier: str, num_rows: int) -> dict:
                 cand_ctl_out=self.cand_ctl[:num_rows],
                 cand_cur_out=self.cand_cur[:num_rows],
             )
-        self.emitted_tier = emit_tier
         return kw
 
     def topk_ext_kwargs(
         self, route: TopkRoute, num_rows: int, block_max: Optional[torch.Tensor]
     ) -> dict:
-        """kwargs for trtllm::cute_dsl_gvr_topk_decode consuming the
-        PREVIOUS step's emission per the picked route."""
+        """kwargs for trtllm::cute_dsl_gvr_topk_decode consuming this
+        step's emission per the picked route."""
         kw: dict = {
             "xstate": self.xstate[:num_rows],
             "cluster_size": route.cluster_size,
@@ -186,16 +177,18 @@ def topk_ext_kwargs(
             kw["num_threads"] = route.num_threads
         if route.tier in ("counts", "list"):
             kw["seed_thr"] = self.seed_row[:num_rows]
-        # rungs tier (first step / no emission yet): pass no seed at
-        # all - cold-start xstate is invalid so the lines would be
-        # non-finite anyway; the plain stock path is the right fallback
-        # (a [rows, 3] column view of the packed row is non-contiguous
-        # and would trip the runner's contract assert)
+        # rungs tier: pass no seed at all (a [rows, 3] column view of the
+        # packed row is non-contiguous and would trip the runner's assert)
         if route.tier == "list":
+            # accept_cap must match the emitter's segment geometry: the
+            # buffers are laid out at bases 0 / LIST_SEG_A / 2*LIST_SEG_A,
+            # and the consumer derives the C capacity from the tensor
+            # width minus 2*accept_cap.
             kw.update(
                 cand_vals=self.cand_vals[:num_rows],
                 cand_idx=self.cand_idx[:num_rows],
                 cand_ctl=self.cand_ctl[:num_rows],
+                accept_cap=LIST_SEG_A,
             )
         if route.attach_block_max and block_max is not None:
             kw["block_max"] = block_max
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 766d8a272a97..042d716ae8c8 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -52,13 +52,13 @@ def _env_flag(name: str) -> bool:
 # Diagnostic knob: compile per-phase clock64 stamps of the list path
 # into the spare xstate slots (harness-side analysis). Off by default;
 # NEVER set in production.
-_P4_TAIL_DBG = _env_flag("GVR_P4_TAIL_DBG")
+_P4_TAIL_DBG = _env_flag("TRTLLM_GVR_P4_TAIL_DBG")
 # P4 sub-phase clock64 breakdown -> xstate[1,2,4,5,6,7] (debug: clobbers
 # the closed-loop thr/anch publish; single-shot cells only, not chains)
-_P4_SUB_DBG = _env_flag("GVR_P4_SUB_DBG")
-# GVR_P4_SUB_DBG=2: publish the P4 HEAD triple (minmax / histogram build /
+_P4_SUB_DBG = _env_flag("TRTLLM_GVR_P4_SUB_DBG")
+# TRTLLM_GVR_P4_SUB_DBG=2: publish the P4 HEAD triple (minmax / histogram build /
 # coarse search) instead of the tail triple, so the phase budget adds up.
-_P4_SUB_HEAD = os.environ.get("GVR_P4_SUB_DBG", "0").strip() == "2"
+_P4_SUB_HEAD = os.environ.get("TRTLLM_GVR_P4_SUB_DBG", "0").strip() == "2"
 # Exact-tail small-class pair buffer. The scatter parks each member of the
 # straddling tie class as (value bits, index) here, above the 256 digit bins
 # the large-class radix zeroes and above its [256..258] scalars, so the repair
@@ -81,7 +81,7 @@ def _pair_cap_for(n_bins: int) -> int:
     return min(_PAIR_MAX, 1 << (room.bit_length() - 1))
 
 
-_SKIP_DBG = _env_flag("GVR_SKIP_DBG")
+_SKIP_DBG = _env_flag("TRTLLM_GVR_SKIP_DBG")
 
 
 # ---------------------------------------------------------------------------
@@ -635,8 +635,7 @@ def __init__(
             # physical candidate-buffer capacity override (B* search)
             self.kC = int(kc_override)
         # acceptance band top B*: a cut whose count fits [K, B*] goes
-        # straight to Phase 4. Cost-bounded (refine ~1.9us/1k cands vs
-        # ~21-30us fallback), physically bounded by kC.
+        # straight to Phase 4. Physically bounded by kC.
         self.accept_cap = int(accept_cap) if accept_cap is not None else self.kC
         self.cand_cap = int(cand_cap)
         # list path: the score column is staged into a DEDICATED smem
@@ -757,28 +756,13 @@ def __init__(
             p4_exact_tail = self.enable_p4_rank_scatter_exact and dtype == cutlass.Float32
         self.p4_exact_tail = bool(p4_exact_tail) and self.enable_p4_rank_scatter_exact
         # p4_tail_fast: tiny-tie COLLECT+SELECT fast path inside the
-        # exact-tail fire branch. When the (b*, sb*) tie class holds <= 128
-        # entries (the real firing cells have 2), ONE candidate pass collects
-        # (value_bits, cand_idx) pairs into SMEM and thread0 selects the
-        # top-need exactly, replacing the 4 unconditional radix passes
-        # (~5.3us -> ~1 pass on pro/512k). Larger tie classes fall through to
-        # the existing radix select. Pure optimization (the radix backstop
-        # keeps exactness identical either way); False compiles the original
-        # text (byte-identical PTX modulo kernel name) for A/B.
-        # Default gate = p4_exact_tail AND top_k >= 1024: the non-firing
-        # codegen tax concentrates at K512 cs=1 mid-N (flash 64k/128k
-        # -6.6/-9.1%, cross-GPU reproducible, 2026-07-20 b200-035) while the
-        # fire census (pro/512k bench + 9 per-layer fixture cells) contains
-        # NO K512 cell — so K512 keeps the original byte-identical kernel.
+        # exact-tail fire branch: when the boundary tie class holds few
+        # enough entries to buffer, one candidate pass replaces the radix
+        # passes; larger classes fall through to the radix backstop, so
+        # exactness is identical either way.
         if p4_tail_fast is None:  # [p4tt]
-            # Was gated on top_k >= 1024, which left K=512 on the original
-            # full-candidate 4-level radix: every level re-zeroes 256 bins
-            # and re-walks the WHOLE candidate array, so a step whose
-            # boundary class is large costs ~5us against ~0.2us for its
-            # neighbours. The compacted path pays for the class only.
-            # Measured on the captured flash rows that trigger it
-            # (256k layer 20 step 4, K=512): 5.04us -> 2.17us, with the
-            # neighbouring steps unchanged at ~0.2us.
+            # default follows p4_exact_tail for every K (the compacted
+            # path pays for the boundary class only)
             p4_tail_fast = self.p4_exact_tail
         self.p4_tail_fast = bool(p4_tail_fast) and self.p4_exact_tail  # [p4tt]
         # p4_tail_v3: compacted-class repair (block-parallel radix +
@@ -790,27 +774,30 @@ def __init__(
                 use_ext_counts or use_ext_cand or ext_rungs or self_scan or enable_block_skip
             )
         self.p4_tail_v3 = bool(p4_tail_v3)
-        # p4_no_fine: drop the 256-bin fine level from phase 4 outright and
-        # let the tail repair rank the whole straddling COARSE bin instead.
-        # Phase 4's cost is a chain of sequential stages, not work: its total
-        # barely moves when the candidate count triples, so the only way to
-        # cut it is to remove a stage. The coarse bin holds 11 candidates at
-        # the median and 49 at the max over 520 captured steps - inside the
-        # tail's pair buffer - and the tail now ranks with the whole block,
-        # so it absorbs the class cheaply. A class past the buffer falls into
-        # the tail's radix, which handles any size. Gated to the same
-        # configuration as that small-class route; the stock kernel keeps the
-        # fine level so its codegen is untouched.
+        # p4_no_fine: drop the 256-bin fine level from phase 4 and let the
+        # tail repair rank the whole straddling COARSE bin instead. A class
+        # past the tail's pair buffer falls into its radix, which handles
+        # any size. Gated to the same configuration as that small-class
+        # route; the stock kernel keeps the fine level so its codegen is
+        # untouched.
         if p4_no_fine is None:
             p4_no_fine = bool(
                 self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3
             ) and not (self.p4_fine_rangetest or self.p4_scat_rangetest)
         self.p4_no_fine = bool(p4_no_fine)
+        if self.p4_no_fine:
+            if not (self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3):
+                raise ValueError("p4_no_fine requires the exact-tail repair chain")
+            if self.p4_fine_rangetest or self.p4_scat_rangetest:
+                raise ValueError("p4_no_fine is incompatible with the range-test arms")
+        if self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3:
+            if _pair_cap_for(self.kNumBins) < 1:
+                raise ValueError(
+                    f"kNumBins={self.kNumBins} leaves no room for the tail pair buffer"
+                )
         # p1r_rescue: rebuild the refine bracket from the row when the
-        # seed bracket is degenerate. ON by default - upstream's identity
-        # shortcut is wrong on real data (every request's first decode
-        # step feeds a zero-init prev_topk). The knob exists to measure
-        # its cost, not to ship it off.
+        # seed bracket is degenerate (e.g. the zero-init prev_topk every
+        # request's first decode step feeds). ON by default.
         self.p1r_rescue = bool(p1r_rescue)
 
     # ------------------------------------------------------------------
@@ -1440,16 +1427,13 @@ def phase0_scan_bucket(
         uncapped), so {n0, void, n1, n2} fall out for free — the same
         contract the v5 emitter produced externally.
 
-        Perf shape (v13): the dense scan is a cp.async pipeline — each
-        thread streams one 16B vector per step into its private
-        slot-major smem staging slot (LDGSTS: no data registers, no
-        scoreboard stall until the wait), keeping ``stage_slots`` steps
-        in flight; classification reads the staged values and claims
-        passers with per-element direct smem atomics (they don't
-        synchronize the warp and hide under the async copy stream).
-        Segment overflow is resolved in-claim by spilling to the next
-        looser segment (a segment overflows at most once per row, and
-        divergent scalar atomics need no warp coordination)."""
+        The dense scan is a cp.async pipeline: each thread streams one
+        16B vector per step into its private slot-major smem staging
+        slot, keeping ``stage_slots`` steps in flight; classification
+        reads the staged values and claims passers with per-element
+        direct smem atomics (they don't synchronize the warp). Segment
+        overflow is resolved in-claim by spilling to the next looser
+        segment (a segment overflows at most once per row)."""
         num_threads = cutlass.const_expr(self.num_threads)
         segA = cutlass.const_expr(self.accept_cap)
         capC = cutlass.const_expr(self.cap_c)
@@ -1475,13 +1459,10 @@ def phase0_scan_bucket(
         # are the same non-synchronizing per-element atomics as the
         # dense loop - so the block loop needs no uniform trip counts.
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
-            # v8.2: 8 blocks per warp iteration. Every lane vector-loads
-            # the SAME 8 bmax values (L1 broadcast - the pass decisions
-            # are warp-uniform registers), then the passing blocks are
-            # loaded back-to-back (independent 128B coalesced loads, so
-            # the memory level parallelism the naive one-block loop
-            # lacked is restored; at high skip rates an iteration is
-            # just the one 32B bmax vector).
+            # 8 blocks per warp iteration: every lane vector-loads the
+            # SAME 8 bmax values (L1 broadcast - the pass decisions are
+            # warp-uniform registers), then the passing blocks are
+            # loaded back-to-back as independent 128B coalesced loads.
             bm_addr = block_max_row.iterator.toint()
             nb0 = (N + cutlass.Int32(31)) >> cutlass.Int32(5)
             # Two-pass skip: (1) DENSE-scan the bmax array itself (it
@@ -1489,10 +1470,8 @@ def phase0_scan_bucket(
             # PASSING BLOCK IDS into the idle C segment (single-band mode
             # never fills C; ids < 2^23 store exactly as floats);
             # (2) walk the compact list, 8 blocks per warp round issued
-            # unguarded back-to-back - every element read is useful and
-            # the loads pipeline. If the list overflows capC the row
-            # falls back to the dense full scan (routing should have
-            # sent it there anyway).
+            # unguarded back-to-back. If the list overflows capC the row
+            # falls back to the dense full scan.
             # pass-1 vectors: 128-bit (the bmax row base is only 16B
             # aligned: nb_pad %% 4)
             pass1_atom = cute.make_copy_atom(
@@ -1661,17 +1640,13 @@ def phase0_scan_bucket(
         cpw = cutlass.const_expr(4)  # cp.async caps at 16B per copy
         step1 = cutlass.const_expr(num_threads * cpw)
         st2log = cutlass.const_expr((2 * step1).bit_length() - 1)
-        # Pair-step cp.async pipeline: the scan is instruction-issue
-        # bound (pcsamp: no_instructions + wait dominate; long_scoreboard
-        # is 6%), so each step processes TWO 16B vectors per thread —
-        # loop/wait/commit/address overhead amortizes over 8 elements
-        # instead of 4 while the in-flight byte count stays put (2 pairs
-        # x 32B across the 4 staging slots). One commit group per pair;
-        # wait_group(1) pops the oldest pair. The staging buffer aliases
-        # smem_vals (only written after phase 0); every non-empty group
-        # is drained inside the loop, so nothing is in flight once the
-        # alias is read. FULL-pair steps only in the hot loop; the
-        # remainder takes the scalar tail below.
+        # Pair-step cp.async pipeline: each step processes TWO 16B
+        # vectors per thread (2 pairs x 32B across the 4 staging slots).
+        # One commit group per pair; wait_group(1) pops the oldest pair.
+        # The staging buffer aliases smem_vals (only written after phase
+        # 0); every non-empty group is drained inside the loop, so
+        # nothing is in flight once the alias is read. FULL-pair steps
+        # only in the hot loop; the remainder takes the scalar tail below.
         nfull = N >> st2log
         if cutlass.const_expr(self.enable_block_skip and block_max_row is not None):
             nfull = cutlass.Int32(0)
@@ -2110,9 +2085,8 @@ def block_count_ge(
     # columns (the accepted rung's column seeds Phase 3 per CTA).
     # ------------------------------------------------------------------
 
-    # ---- block-skip machinery (enable_block_skip; ported from the
-    # skip-finegrain development chain, measured-optimal configuration
-    # only: grain 32, int16 list, grouped strided build, UN=2 scan) ----
+    # ---- block-skip machinery (enable_block_skip): grain 32, int16
+    # list, grouped strided build, UN=2 scan ----
 
     @cute.jit
     def _list_ld(self, smem_active, idx):
@@ -3541,15 +3515,12 @@ def _kth_bin_search_rw(self, smem_hist, smem_wcnt, lo, binw, tidx, warp_id, lane
 
     # ------------------------------------------------------------------
     # _p4_coarse_rw - redundant-warp coarse bin search for the fused
-    # rank-and-scatter path. Same result as the high->low walk it
-    # replaces (the straddling bin and the count strictly above it),
-    # but staged once and then resolved lane-parallel on every warp:
-    # an idx-shuffle scan + ballot locate the target slice, a second
-    # scan + the unique crossing test locate the bin inside it. Two
-    # publish barriers and a bins_per_warp-deep serial LDS+IADD chain in
-    # a single lane disappear; integer sums are associative so every
-    # warp lands on the same answer bit-for-bit. Mirrors
-    # _kth_bin_search_rw, which does the same for the snap path.
+    # rank-and-scatter path: returns the straddling bin and the count
+    # strictly above it, resolved lane-parallel on every warp (an
+    # idx-shuffle scan + ballot locate the target slice, a second scan
+    # + the unique crossing test locate the bin inside it). Integer
+    # sums are associative, so every warp lands on the same answer
+    # bit-for-bit. Mirrors _kth_bin_search_rw (snap path).
     # ------------------------------------------------------------------
     @cute.jit
     def _p4_coarse_rw(self, smem_hist, smem_wcnt, warp_id, lane):
@@ -3653,13 +3624,10 @@ def _p4_coarse_rw(self, smem_hist, smem_wcnt, warp_id, lane):
 
     # ------------------------------------------------------------------
     # _p4_fine_rw - redundant-warp variant of the fine sub-bin search,
-    # the same transformation _p4_coarse_rw applies one level up. The
-    # serial form stages per-warp slice sums, has thread 0 walk the warp
-    # totals, then has one lane of the target warp walk that slice bin by
-    # bin, with three publish barriers. Here every warp resolves it from
-    # the staged sums with an idx-shuffle scan and a ballot, so two
-    # barriers and the fbins/num_warps-deep serial LDS chain disappear.
-    # Integer sums are associative: every warp lands on the same answer.
+    # the same transformation _p4_coarse_rw applies one level up: every
+    # warp resolves it from the staged per-warp sums with an idx-shuffle
+    # scan and a ballot. Integer sums are associative: every warp lands
+    # on the same answer.
     # ------------------------------------------------------------------
     @cute.jit
     def _p4_fine_rw(self, smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane):
@@ -3751,8 +3719,7 @@ def _p4_fine_rw(self, smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane):
         return sb_out, ra_out
 
     # ------------------------------------------------------------------
-    # Phase 4 (alt): op#7 fused rank-and-scatter (enable_p4_rank_scatter).
-    # Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py.
+    # Phase 4 (alt): fused rank-and-scatter (enable_p4_rank_scatter).
     # ------------------------------------------------------------------
     @cute.jit
     def phase4_rank_scatter(
@@ -3807,8 +3774,7 @@ def phase4_rank_scatter(
                 if use_ext_r == cutlass.Int32(1):
                     # list rows: the take walk pre-zeroed the hist and staged
                     # per-warp maxima in smem_wcnt (its end barrier orders
-                    # them); min := cut line by construction. The minmax
-                    # scan, the zero pass and their three barriers vanish.
+                    # them); min := cut line by construction.
                     if cutlass.const_expr(ext_min is not None):
                         bmin_r = ext_min
                     for w in cutlass.range_constexpr(self.num_warps):
@@ -3820,14 +3786,9 @@ def phase4_rank_scatter(
                         bmax_r = bmin_r + cutlass.Float32(1e-6)
             if run_stock_range:
                 # ---- block min/max over candidates ----
-                # The accepted threshold is an EXACT lower bound: every
-                # candidate got here by comparing >= against it. So the min
-                # side of this scan buys nothing but a slightly tighter
-                # window - and the window only has to be tight where the
-                # k-th lives, which is the TOP of it. Taking the threshold
-                # instead drops half the scan, one warp reduce and one arm
-                # of the cross-warp fold. Only the assist tiers have a
-                # published threshold, hence the gate.
+                # The accepted threshold is an EXACT lower bound on every
+                # candidate, so it can stand in for the min. Only the
+                # assist tiers have a published threshold, hence the gate.
                 use_thr_min = cutlass.const_expr(self.p4_no_fine)
                 local_cmin = cutlass.Float32(self.FLT_MAX)
                 local_cmax = cutlass.Float32(self.NEG_FLT_MAX)
@@ -3848,10 +3809,8 @@ def phase4_rank_scatter(
                     smem_hist[warp_id] = float_as_uint32(cmax)
                 cute.arch.barrier()
                 # lane-parallel cross-warp fold: lane w holds slot w and one
-                # warp reduce settles it, instead of every thread walking all
-                # num_warps slots of two arrays with dependent SMEM reads.
-                # min/max reassociate freely, so the result is bit-identical
-                # and every warp still lands on it without a leader.
+                # warp reduce settles it. min/max reassociate freely, so the
+                # result is bit-identical on every warp without a leader.
                 pmn = cutlass.Float32(self.FLT_MAX)
                 pmx = cutlass.Float32(self.NEG_FLT_MAX)
                 if lane < cutlass.Int32(self.num_warps):
@@ -3960,8 +3919,7 @@ def phase4_rank_scatter(
 
             # ---- EXACT: one fine-histogram recursion on the straddling bin b* ----
             if cutlass.const_expr(self.enable_p4_rank_scatter_exact):
-                # FIXED small fine-bin count (independent of kNumBins) — cuts the
-                # re-zero + 3-step cost (esp. K=2048 where kNumBins=2048); 256
+                # FIXED small fine-bin count (independent of kNumBins): 256
                 # sub-bins over bin b* gives kNumBins×256 effective resolution,
                 # enough to resolve the straddling bin to ≤1 distinct value.
                 fbins = cutlass.const_expr(256)
@@ -3975,11 +3933,8 @@ def phase4_rank_scatter(
                     # Sub-binning collapsed: finv 0 sends every member of the
                     # straddling coarse bin to sub-bin 0 == sb*, so the
                     # scatter parks the whole coarse class and the tail ranks
-                    # it. The re-zero, the build and the search below are not
-                    # traced at all - this is a compile-time removal, not a
-                    # runtime branch, because a runtime branch still costs the
-                    # stage (measured: 0.64us for a fine window whose work was
-                    # gated off).
+                    # it. MUST stay a compile-time removal (the re-zero, build
+                    # and search below untraced), not a runtime branch.
                     finv = cutlass.Float32(0.0)
                     sb_star = cutlass.Int32(0)
                     rank_above_fine = rank_above
@@ -4000,12 +3955,10 @@ def phase4_rank_scatter(
                     cute.arch.barrier()
                     if cutlass.const_expr(self.p4_fine_rangetest):
                         # A candidate belongs to bin b* exactly when its value
-                        # lies in [f_lo, f_hi), so the filter is two compares -
-                        # the bin recompute (subtract + multiply + two clamps)
-                        # per candidate is redundant work. The clamped ends of
-                        # the binning fold out-of-range values INTO bin 0 and
-                        # bin kBins-1, so those two bins drop the matching side
-                        # of the range test to stay bit-identical.
+                        # lies in [f_lo, f_hi). The clamped ends of the binning
+                        # fold out-of-range values INTO bin 0 and bin kBins-1,
+                        # so those two bins must drop the matching side of the
+                        # range test to stay bit-identical.
                         ifb = tidx
                         while ifb < cand_count:
                             vf = smem_keys[ifb]
@@ -4041,9 +3994,8 @@ def phase4_rank_scatter(
                             ifb = ifb + cutlass.Int32(num_threads)
                     cute.arch.barrier()
                     # fine sub-bin search, resolved lane-parallel on every
-                    # warp (see _p4_fine_rw): the two publish barriers and the
-                    # fbins/num_warps-deep single-lane walk that used to sit
-                    # here are gone, and the answer comes back in registers.
+                    # warp (see _p4_fine_rw); the answer comes back in
+                    # registers.
                     sb_star, rank_above_fine = self._p4_fine_rw(
                         smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane
                     )
@@ -4115,20 +4067,13 @@ def phase4_rank_scatter(
                             if cutlass.const_expr(
                                 self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3
                             ):
-                                # The tie class is exactly what the exact-tail
-                                # repair needs, and this pass has already
-                                # classified every candidate AND handed out the
-                                # intra-class ordinal. Park (value bits, index)
-                                # here so the repair does not have to walk the
-                                # candidates again: measured 0.61us of its
-                                # remaining cost was that second walk. Pairs
+                                # Park (value bits, index) so the exact-tail
+                                # repair never re-walks the candidates. Pairs
                                 # live above the digit bins the radix route
                                 # zeroes; small classes are the only consumer.
                                 # Unconditional, index wrapped: the buffer is
                                 # only ever READ when the class fits it, and
-                                # then the wrap is a no-op. Dropping the bound
-                                # branch keeps this off the scatter's critical
-                                # path, which every step pays.
+                                # then the wrap is a no-op.
                                 ow = (o & cutlass.Int32(pair_cap - 1)) * cutlass.Int32(2)
                                 smem_hist[cutlass.Int32(_PAIR_BASE) + ow] = float_as_int32(v)
                                 smem_hist[cutlass.Int32(_PAIR_BASE) + ow + cutlass.Int32(1)] = (
@@ -4166,16 +4111,10 @@ def phase4_rank_scatter(
                 # slot range [rank_above_fine, kK). Unambiguous rows (the
                 # overwhelming majority) pay two scalar compares; the counters
                 # and the fine histogram are reused, so SMEM does not grow.
-                # boundary-class repair: collect the (b*, sb*) tie
-                # class compactly (ONE candidate pass), then select inside
-                # it. Tiny jobs (need x class <= 512) keep the thread0
-                # serial select (cheapest at that size). Bigger classes up
-                # to capc get a 4-level MSB radix over the COLLECTED class
-                # (each level scans <= capc pairs with all threads instead
-                # of re-scanning every candidate) — the old serial select
-                # was O(need x class) and measured 8.5us on real chain rows
-                # (need ~100 x class ~100). Classes beyond capc take the
-                # UNMODIFIED full-candidate radix below (verbatim copy).
+                # boundary-class repair: pure-tie classes (one key value)
+                # exit on a warp-reduce precheck; classes parked in the
+                # pair buffer rank block-parallel over the parked pairs;
+                # anything larger takes the full-candidate radix below.
                 if cutlass.const_expr(self.p4_exact_tail and self.p4_tail_fast):  # [p4tt]
                     if cutlass.const_expr(self.p4_tail_v3):
                         need0 = cutlass.Int32(kK) - rank_above_fine
@@ -4197,19 +4136,9 @@ def phase4_rank_scatter(
                                 # Small mixed class: the scatter already parked
                                 # every member as a (value bits, index) pair,
                                 # so there is nothing to collect - go straight
-                                # to the rank. No extra candidate walk, no
-                                # staging barrier, nothing riding in registers.
-                                # Rank the class with the WHOLE BLOCK, not
-                                # one warp. Each warp owns a stride of the
-                                # class and its 32 lanes split the
-                                # comparisons, so the cost is
-                                # ceil(n/num_warps) * ceil(n/32) trips per
-                                # lane plus one warp reduce - 1 trip for the
-                                # class of two that fires today, 8 for a class
-                                # of fifty. The old form had a single warp walk
-                                # n^2/32 steps while the other fifteen waited
-                                # at the barrier below, which is what made a
-                                # larger class unaffordable.
+                                # to the rank. Rank with the WHOLE BLOCK: each
+                                # warp owns a stride of the class and its 32
+                                # lanes split the comparisons.
                                 e9 = warp_id
                                 while e9 < cnt_strad:
                                     be9 = smem_hist[_PAIR_BASE + e9 + e9]
@@ -4253,20 +4182,12 @@ def phase4_rank_scatter(
                                     e9 = e9 + cutlass.Int32(num_warps)
                                 cute.arch.barrier()
                             else:
-                                # Large class only. Everything below -
-                                # the pure-tie pre-check walk, its
-                                # cross-warp fold and the publish
-                                # barrier that fold needs - exists to
-                                # decide whether the radix can be
-                                # skipped. The small route above needs
-                                # none of it, and used to run it anyway
-                                # for a class of two.
-                                # block-wide pure-tie check, ANY class
-                                # size: min/max order key over the (b*, sb*) class.
+                                # Large class only: block-wide pure-tie check
+                                # (min/max order key over the (b*, sb*) class)
+                                # decides whether the radix can be skipped.
                                 # A pure-tie class needs NO repair — the scatter's
                                 # arrival fill of bit-equal values is already
-                                # value-set exact. Real fp8-lineage logits tie in
-                                # the thousands, which used to take the full radix.
+                                # value-set exact.
                                 # Staging mirrors the head min/max (wcnt + hist
                                 # slots [0..31], both dead here; pairs live at
                                 # 260+).
@@ -4276,14 +4197,10 @@ def phase4_rank_scatter(
                                 kmn6 = cutlass.Int32(2147483647)
                                 kmx6 = cutlass.Int32(-2147483648)
                                 it6 = tidx
-                                # The pure-tie pre-check only exists to SKIP the
-                                # repair when the class is bit-uniform. The small
-                                # route below ranks by (key, arrival) and rewrites
-                                # exactly the need0 winner slots, which is already
-                                # value-exact on a pure tie - so for a small class
-                                # the check is a whole candidate walk bought for
-                                # nothing. Only the large-class route, whose radix
-                                # really is worth avoiding, still pays for it.
+                                # The pure-tie pre-check SKIPs the repair when
+                                # the class is bit-uniform; only the large-class
+                                # route runs it (the small route is already
+                                # value-exact on a pure tie).
                                 while it6 < cand_count:
                                     v6 = smem_keys[it6]
                                     b6 = cutlass.Int32((v6 - bmin_r) * inv1)
@@ -4292,11 +4209,15 @@ def phase4_rank_scatter(
                                     if b6 > cutlass.Int32(kBins - 1):
                                         b6 = cutlass.Int32(kBins - 1)
                                     if b6 == b_star:
-                                        s6 = cutlass.Int32((v6 - f_lo) * finv)
-                                        if s6 < cutlass.Int32(0):
-                                            s6 = cutlass.Int32(0)
-                                        if s6 > cutlass.Int32(fbins - 1):
-                                            s6 = cutlass.Int32(fbins - 1)
+                                        # with the fine level compiled out the
+                                        # whole coarse bin IS the class
+                                        s6 = cutlass.Int32(0)
+                                        if cutlass.const_expr(not self.p4_no_fine):
+                                            s6 = cutlass.Int32((v6 - f_lo) * finv)
+                                            if s6 < cutlass.Int32(0):
+                                                s6 = cutlass.Int32(0)
+                                            if s6 > cutlass.Int32(fbins - 1):
+                                                s6 = cutlass.Int32(fbins - 1)
                                         if s6 == sb_star:
                                             k6 = f32_order_key(v6) ^ cutlass.Int32(-2147483648)
                                             if k6 < kmn6:
@@ -4304,13 +4225,9 @@ def phase4_rank_scatter(
                                             if k6 > kmx6:
                                                 kmx6 = k6
                                             # Buffer the member so the compaction
-                                            # below needs no walk of its own. The
+                                            # below needs no walk of its own; the
                                             # array is nbuf7 = kC/num_threads
-                                            # deep and every store is an unrolled
-                                            # dynamic-index search over it, which
-                                            # is why only the large class - whose
-                                            # alternative is a 4-level radix -
-                                            # pays for it.
+                                            # deep (unrolled dynamic-index store).
                                             for sl7 in cutlass.range_constexpr(nbuf7):
                                                 if cutlass.Int32(sl7) == nh7:
                                                     rv7[sl7] = v6
@@ -4341,17 +4258,12 @@ def phase4_rank_scatter(
                                 if fast_done == cutlass.Int32(0):
                                     # mixed class: compact the members buffered by
                                     # the pure-tie pass above into
-                                    # smem_keys/vals[0..cnt_strad). The buffering
-                                    # pass already read every candidate it needs,
-                                    # and the staging barrier above orders those
-                                    # reads before these writes, so the second
-                                    # full-candidate walk (and its barrier) is
-                                    # gone. Repair cost is a function of the CLASS
-                                    # size only, for any class size.
+                                    # smem_keys/vals[0..cnt_strad). The staging
+                                    # barrier above orders the buffered reads
+                                    # before these writes.
                                     # warp-aggregated claim: intra-warp exclusive
                                     # prefix via shfl scan + ONE atomic per warp
-                                    # (a thousand same-address claims serialize
-                                    # and scale with the class size)
+                                    # (same-address claims would serialize)
                                     pf7 = nh7
                                     for so3 in cutlass.range_constexpr(5):
                                         oth3 = cute.arch.shuffle_sync_up(
@@ -4373,13 +4285,11 @@ def phase4_rank_scatter(
                                             smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8]
                                             smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8]
                                     cute.arch.barrier()
-                                    # Large class only (the small one is handled
-                                    # above, ahead of the pure-tie pass, so it
-                                    # never reaches here): block-parallel 4-level
+                                    # Large class only (the small one never
+                                    # reaches here): block-parallel 4-level
                                     # MSB radix over the compacted class (scans
                                     # touch class pairs only; warp0 shuffle-scan
-                                    # digit search - 3 block barriers per level
-                                    # instead of 5).
+                                    # digit search).
                                     if tidx == cutlass.Int32(0):
                                         smem_hist[256] = cutlass.Int32(0)
                                         smem_hist[257] = need0
@@ -4627,8 +4537,7 @@ def phase4_rank_scatter(
                                     # Two-stage descending digit scan (mirrors the
                                     # fine 3-step search): per-warp partial sums,
                                     # thread0 picks the target warp, its lane0 walks
-                                    # the warp's digit range — 2*num_warps serial
-                                    # steps instead of 256.
+                                    # the warp's digit range.
                                     fdw = cutlass.const_expr(256 // self.num_warps)
                                     wsum2 = cutlass.Int32(0)
                                     for jd in cutlass.range_constexpr(fdw):
@@ -5629,8 +5538,7 @@ def run_one_row(
             block_max_row = None
         if cutlass.const_expr(self.use_ext_counts and seed_thr is not None):
             # packed seed row [>=6] fp32: [0..2] lines, [3..5] counts as
-            # floats (exact to 2^24) - ONE 32B sector serves both, halving
-            # the serial cold loads of the admission preview
+            # floats (exact to 2^24) - ONE 32B sector serves both
             seed_thr_row = seed_thr[row_idx, None]
             seed_counts_row = None
         elif cutlass.const_expr(self.ext_rungs and seed_thr is not None):
@@ -5900,16 +5808,11 @@ def run_one_row(
             s_cluster_partial_m = None
             smem_gath = None
 
-        # PDL wait placed as late as possible: everything above (row
-        # resolution, SMEM allocation, config folding - and, on a cold
-        # call, the instruction fetch for all of it) overlaps with the
-        # producer indexer's tail, because dependent CTAs stage onto SMs
-        # as producer CTAs retire. Nothing above reads producer-written
-        # data: seq_lens/pre_idx come from the host-side metadata and the
-        # feedback buffer, while logits / block_max / seed_thr / cand are
-        # first touched below. Instruction-fetch starvation is 37-44% of
-        # this kernel's stall cycles on the small-N cells, so warming it
-        # under the producer's shadow is the point.
+        # PDL wait placed as late as possible so the prologue overlaps
+        # the producer indexer's tail. INVARIANT: nothing above may read
+        # producer-written data - seq_lens/pre_idx come from host-side
+        # metadata and the feedback buffer; logits / block_max /
+        # seed_thr / cand are first touched below.
         if cutlass.const_expr(self.pdl_wait_late):
             griddepcontrol_wait()
 
@@ -6183,10 +6086,8 @@ def _run_phases(
         # Use the epilogue rungs ONLY when the row is valid (finite t_0,
         # xstate contract) AND some rung count already lies in [K, kC].
         # A miss/invalid row runs the full stock path (P1 + P1b + vseed +
-        # count): real data shows stock beats ext-bracket refine on
-        # misses (pro-1M cold: stock-skip 21us vs ext-miss 51us). All
-        # threads read the same control words, so the predicate is
-        # CTA-uniform and the dynamic branches below stay convergent.
+        # count). All threads read the same control words, so the
+        # predicate is CTA-uniform and the branches below stay convergent.
         if cutlass.const_expr(_P4_TAIL_DBG):
             ck0 = cutlass.Int64(0)
             ck1 = cutlass.Int64(0)
@@ -6194,12 +6095,11 @@ def _run_phases(
             ckE = cute.arch.clock64()  # row-phase entry (device-residency ref)
         ext_row = cutlass.Int32(0)
         if cutlass.const_expr(self.use_ext_counts):
-            # line validity mirrors ext_rungs: ALL THREE lines finite and
-            # strictly ascending. The old t0-only guard let a NaN in
-            # t1/t2 get parked into the refine brackets (the same failure
-            # mode that broke ext_rungs exactness on-chain); invalid rows
-            # fall to the stock path, exactness never rides on the host
-            # loop's line quality.
+            # line validity mirrors ext_rungs: ALL THREE lines must be
+            # finite and strictly ascending (a NaN in t1/t2 must not
+            # reach the refine brackets); invalid rows fall to the stock
+            # path, so exactness never rides on the host loop's line
+            # quality.
             if (
                 seed_thr_row[0] < cutlass.Float32(1e37)
                 and seed_thr_row[0] > cutlass.Float32(-1e37)
@@ -6222,6 +6122,10 @@ def _run_phases(
             ):
                 claimed_p = cutlass.Int32(cand_ctl_row[0])
                 void_p = cutlass.Int32(cand_ctl_row[1])
+                # real (non-parked) lines only: the skip stages the raw
+                # lines into the threshold scratch, and a parked line
+                # (1e30) would poison every later bracket. Parked rows
+                # keep Phase 1; the list take below is independent.
                 if (
                     void_p == cutlass.Int32(0)
                     and claimed_p >= cutlass.Int32(self.top_k + 64)
@@ -6230,6 +6134,7 @@ def _run_phases(
                     and seed_thr_row[0] > cutlass.Float32(-1e37)
                     and seed_thr_row[1] > seed_thr_row[0]
                     and seed_thr_row[2] > seed_thr_row[1]
+                    and seed_thr_row[2] < cutlass.Float32(1e29)
                 ):
                     ext_row = cutlass.Int32(1)
             # ---- self_scan phase 0: fused scan-bucket ----
@@ -6400,13 +6305,11 @@ def _run_phases(
         # A duplicate/invalid preIdx gather (cold-start zero-init slots, stale
         # slots pointing past N, an all-tied gather) produces an unusable
         # bracket. When N > K real selection work remains, so rebuild the
-        # bracket from the data itself (P1r) and run the normal pipeline —
-        # the old identity shortcut here returned indices [0, K), which is
-        # NOT the top-K on real data (production hit: the first decode step
-        # of every request feeds the zero-init prev_topk feedback buffer).
-        # If the bracket is STILL degenerate after the rescue, every
-        # in-range value is identical (or N <= K), and identity output is
-        # then exact — keep the shortcut for exactly those rows.
+        # bracket from the data itself (P1r) and run the normal pipeline;
+        # an identity shortcut here would NOT be the top-K. If the bracket
+        # is STILL degenerate after the rescue, every in-range value is
+        # identical (or N <= K), and identity output is then exact — keep
+        # the shortcut for exactly those rows.
         if cutlass.const_expr(self.p1r_rescue):
             v_lo = s_thr[1]
             v_hi = s_thr[2]
@@ -6496,8 +6399,11 @@ def _run_phases(
                 void_c = cutlass.Int32(cand_ctl_row[1])
                 n1_c = cutlass.Int32(cand_ctl_row[2])
                 n2_c = cutlass.Int32(cand_ctl_row[3])
-                segA = cutlass.const_expr(min(self.accept_cap, self.kC))
-                bstar = segA
+                # segment bases/extents follow the EMITTER geometry
+                # (accept_cap); the admission bound is additionally
+                # clamped by the physical candidate capacity kC.
+                segA = cutlass.const_expr(self.accept_cap)
+                bstar = cutlass.const_expr(min(self.accept_cap, self.kC))
                 if cutlass.const_expr(_P4_TAIL_DBG):
                     ck0 = cute.arch.clock64()
                     ck1 = ck0
@@ -6585,6 +6491,9 @@ def _run_phases(
                             hs_len = lenA
                             hs_band = n2_c
                             need_max = cutlass.Int32(1)
+                        if b_hi >= cutlass.Float32(1e29):
+                            # parked upper line: bracket by the segment max
+                            need_max = cutlass.Int32(1)
                         if hs_band < cutlass.Int32(1):
                             hs_band = cutlass.Int32(1)
                         samp_f = (cutlass.Float32(1.0) * hs_len) / hs_band
@@ -6623,9 +6532,8 @@ def _run_phases(
                         # by segment/band); the descend base stays in
                         # sample units too - the post-load exact-count
                         # net absorbs the sampling error.
-                        # fire target = 1.25x the K-need: the sampled
-                        # estimate carries ~1-2% noise and firing at the
-                        # band's bottom edge would demote half the time
+                        # fire target = 1.25x the K-need (headroom over
+                        # the sampling noise)
                         kneedS = cutlass.Int32(
                             (cutlass.Float32(1.25) * (kK_l - base_c)) * samp_f
                             + cutlass.Float32(0.5)
@@ -6733,10 +6641,8 @@ def _run_phases(
                     lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1)
                     # fused P4 prologue: zero the coarse hist here and
                     # accumulate the candidate max INSIDE the cut walk
-                    # (per-fragment fmax is free ILP; min := cut line by
-                    # construction). P4's minmax scan + zero pass and
-                    # their three barriers disappear for list rows - the
-                    # staging rides this walk's own end barrier.
+                    # (min := cut line by construction); the staging
+                    # rides this walk's own end barrier.
                     izh_c = tidx
                     while izh_c < cutlass.Int32(self.kNumBins):
                         smem_hist[izh_c] = cutlass.Int32(0)
@@ -6897,12 +6803,8 @@ def _run_phases(
                 and self.dtype == cutlass.Float32
             ):
                 if ext_row == cutlass.Int32(1) and N < cutlass.Int32(16384):
-                    # short rows only: the fused walk wins ~1us below
-                    # ~16k (pass fixed costs dominate there); at 16k+
-                    # the claim-dense cells (pro 9.4% band) lose to
-                    # count-then-place and the rest sit at parity, so
-                    # the two-pass path keeps them (measured both ways,
-                    # loncheng cells; captures all sit above the gate)
+                    # short rows only (< 16k): fused claim-collect;
+                    # longer rows keep the two-pass count-then-place path
                     cut_p = s_thr[0]  # parked line (staged at P1-init)
                     rbase = input_row.iterator.toint()
                     vw_p = cutlass.const_expr(self.vec_bits // self.dtype.width)
@@ -7083,10 +6985,8 @@ def _run_phases(
 
                 # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast
                 # path) or the secant threshold search ----
-                # enable_r0 gates to cluster_size==1 for now: op#26's R0 scans the
-                # full row in one CTA. The slice-parallel + cluster count-merge
-                # variant that lets R0 cover the cs>1 long-row branch lands in a
-                # later commit; until then cs>1 keeps the secant path.
+                # enable_r0 gates to cluster_size==1: R0 scans the full row in
+                # one CTA; cs>1 keeps the secant path.
                 if cutlass.const_expr(self.enable_r0):
                     # P1b rung placement -> ONE M-ary R0 count pass -> accept the
                     # tightest rung with count in [K, kC]. On a miss, fall back to
@@ -7109,12 +7009,10 @@ def _run_phases(
                             # Parking and staging happen in the P1-init thread0
                             # block, one barrier for the whole prologue.
                             if cutlass.const_expr(not self.enable_block_skip):
-                                # v3: the parked M-ary pass counted the SAME
-                                # threshold in all three columns (3x compare
-                                # + 3 ptcnt columns for identical values).
-                                # Count it ONCE with the refine primitive -
-                                # same per-thread ptcnt cache and cluster
-                                # merge P3 consumes - and accept in place.
+                                # parked admission: count the parked threshold
+                                # ONCE with the refine primitive - same
+                                # per-thread ptcnt cache and cluster merge P3
+                                # consumes - and accept in place.
                                 if s_r0col[0] == cutlass.Int32(self.M_qf):
                                     self.block_count_ge(
                                         input_row,
@@ -7265,12 +7163,10 @@ def _run_phases(
                         )
                     cute.arch.barrier()
                     if run_mary and tidx == 0:
-                        # tightest admissible rung = SMALLEST count in [K, kC].
-                        # (Explicit argmin: with r0_vseed the pmean column is not
-                        # sorted into the rung order; for sorted rungs this is
-                        # equivalent to the old "last m in window" rule.)
-                        # Dropped rungs (block-skip rung tightening) hold PARTIAL
-                        # counts — never admissible.
+                        # tightest admissible rung = SMALLEST count in [K, kC]
+                        # (explicit argmin: with r0_vseed the pmean column is not
+                        # sorted into the rung order). Dropped rungs (block-skip
+                        # rung tightening) hold PARTIAL counts — never admissible.
                         dmask_c = cutlass.Int32(0)
                         if cutlass.const_expr(self.enable_block_skip):
                             dmask_c = s_active_cnt[2]
@@ -7314,13 +7210,11 @@ def _run_phases(
                         smem_ptcnt[tidx] = smem_ptcnt_multi[bc * cutlass.Int32(num_threads) + tidx]
                     cute.arch.barrier()
                     # ---- R0 miss: SEEDED bounded log-falsi refine ----
-                    # At large N the M2D rungs straddle [K, kC]; the refine must
-                    # find a threshold with count in [K, kC] between the measured
-                    # rungs. SEED the loop with the rung bracket AND its known
-                    # counts (clo/chi) so it does log-count regula-falsi from
-                    # iter 0 with no re-measure and no separate R1 shot -> ~2-3
-                    # count passes (op#26 efficiency) instead of ~6. done=1 on
-                    # accept so Phase 3 skips its retry-shrink.
+                    # The refine must find a threshold with count in [K, kC]
+                    # between the measured rungs. SEED the loop with the rung
+                    # bracket AND its known counts (clo/chi) so it does
+                    # log-count regula-falsi from iter 0 with no re-measure.
+                    # done=1 on accept so Phase 3 skips its retry-shrink.
                     if bc < cutlass.Int32(0):
                         if cutlass.const_expr(self.enable_block_skip):
                             if tidx == cutlass.Int32(0):
@@ -7957,14 +7851,11 @@ def pick_config(
         if n_row < 65536:
             cluster_size = 1
         elif has_block_max and n_row >= 200_000:
-            # Block-skip sweet spot is cs == 1 with a large per-CTA slice:
-            # the compact list + rung tightening (cs1-only) beat the
-            # row-split configs outright once the bounds prune the scan
-            # (cold protocol, real data: BS1 1.18x, BS64 2.14x, BS1024
-            # 4.16x vs this policy's stock picks; splitting shrinks each
-            # CTA's slice below the skip break-even and disables
-            # tightening). Below 200k the wrapper drops block_max anyway
-            # (skip_min_n gate) and the stock picks apply.
+            # Block-skip requires cs == 1 with a large per-CTA slice
+            # (splitting shrinks each CTA's slice below the skip
+            # break-even and disables rung tightening). Below 200k the
+            # wrapper drops block_max anyway (skip_min_n gate) and the
+            # stock picks apply.
             cluster_size = 1
         elif num_rows <= 4 and n_row >= 131072:
             cluster_size = 8
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
index ea382f8b4145..95ccff38ee87 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
@@ -1278,3 +1278,226 @@ def test_cute_dsl_gvr_topk_decode_tie_flood_beyond_capacity():
     )
     torch.cuda.synchronize()
     _tie_aware_check(out_indices, logits, seq_lens, top_k, 1, compress_ratio=1)
+
+
+# ---------------------------------------------------------------------------
+# Emission-assisted (ext) tiers: packed seed row / candidate list / block max.
+# Inputs emulate the indexer epilogue host-side against the layout contracts
+# in ``gvr_ext`` (segments at bases 0 / LIST_SEG_A / 2*LIST_SEG_A, packed row
+# = lines at [0..2] + exact counts at [3..5] + skip pass count at [6]).
+# ---------------------------------------------------------------------------
+
+_FLT_MAX = 3.4028234663852886e38
+
+
+def _lines_at_counts(logits_f32, n_eff, targets):
+    """Per-row threshold lines placed at exact counts (descending targets
+    -> ascending line values). count(logits >= line[j]) == targets[j]."""
+    num_rows = logits_f32.shape[0]
+    lines = torch.empty((num_rows, len(targets)), dtype=torch.float32, device=logits_f32.device)
+    for r in range(num_rows):
+        ne = int(n_eff[r])
+        row = logits_f32[r, :ne]
+        for j, c in enumerate(targets):
+            c = min(int(c), ne)
+            lines[r, j] = torch.kthvalue(row, ne - c + 1).values
+    return lines
+
+
+def _pack_seed_row(logits_f32, n_eff, lines, block_max=None):
+    """[rows, 8] fp32 packed seed row: lines + exact counts (+ skip count)."""
+    num_rows, N = logits_f32.shape
+    pos = torch.arange(N, device=logits_f32.device)[None, :]
+    valid = pos < n_eff[:, None]
+    counts = torch.stack(
+        [((logits_f32 >= lines[:, j : j + 1]) & valid).sum(-1) for j in range(3)], 1
+    ).int()
+    pack = torch.zeros((num_rows, 8), dtype=torch.float32, device=logits_f32.device)
+    pack[:, 0:3] = lines
+    pack[:, 3:6] = counts.float()
+    if block_max is not None:
+        pack[:, 6] = (block_max >= lines[:, 0:1]).sum(dim=1).float()
+    return pack.contiguous()
+
+
+@skip_not_sm100
+@pytest.mark.parametrize("top_k", [512, 1024, 2048])
+@pytest.mark.parametrize("mode", ["band", "fat", "miss", "inf"])
+def test_cute_dsl_gvr_topk_decode_ext_counts(top_k, mode):
+    """Packed seed row ([rows, 8]: lines + exact counts) consumption.
+
+    band: one line's count sits inside the admission band (direct path).
+    fat:  every count overshoots the candidate capacity -> full fallback.
+    miss: lines above the row max (count 0) -> seed rejected.
+    inf:  non-finite lines (production cold start) -> validity guard.
+    """
+    N, batch = 131072, 4
+    logits, pre_idx, seq_lens = _make_inputs(batch, N, top_k, torch.float32, 1, seed=7, varlen=True)
+    n_eff = seq_lens.to(device=logits.device, dtype=torch.long)
+    if mode == "band":
+        lines = _lines_at_counts(logits, n_eff, (4 * top_k, 2 * top_k, top_k + top_k // 4))
+    elif mode == "fat":
+        lines = _lines_at_counts(logits, n_eff, (32768, 24576, 16384))
+    elif mode == "miss":
+        pos = torch.arange(N, device=logits.device)[None, :]
+        rowmax = torch.where(pos < n_eff[:, None], logits, float("-inf")).amax(-1)
+        lines = rowmax[:, None] + torch.tensor([1.0, 2.0, 3.0], device=logits.device)
+    else:
+        lines = torch.full((batch, 3), float("inf"), device=logits.device)
+    seed_row = _pack_seed_row(logits, n_eff, lines)
+    xstate = torch.zeros((batch, 8), dtype=torch.float32, device=logits.device)
+    out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda")
+    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+        logits,
+        pre_idx,
+        seq_lens,
+        out_indices,
+        top_k=top_k,
+        cluster_size=1,
+        seed_thr=seed_row,
+        xstate=xstate,
+    )
+    torch.cuda.synchronize()
+    _tie_aware_check(out_indices, logits, seq_lens, top_k, 1)
+    if mode == "band":
+        # the closed loop must republish valid state for the next step
+        assert bool((xstate[:, 0] > 0).all().item())
+
+
+@skip_not_sm100
+@pytest.mark.parametrize("mode", ["hit", "pads", "hist", "void", "bucketed"])
+def test_cute_dsl_gvr_topk_decode_ext_list(mode):
+    """Candidate-list tier at the production geometry (accept_cap =
+    LIST_SEG_A, width = LIST_WIDTH).
+
+    hit:      parked lines (production shape), claimed inside [K+64, B*]
+              -> line cut, single mapped load.
+    pads:     same + interleaved idx=-1 window sentinels (claimed counts
+              them; the K+64 slack absorbs them).
+    hist:     claimed past B* but list complete -> clamped-histogram
+              fallback over segment C.
+    void:     collection overflows LIST_CAP_C -> void=1 -> full scan.
+    bucketed: three live lines spread across segments A/B/C, cut at the
+              tightest line inside the band.
+    """
+    from tensorrt_llm._torch.attention_backend.sparse.gvr_ext import (
+        LIST_CAP_C,
+        LIST_PARK_LINE,
+        LIST_SEG_A,
+        LIST_WIDTH,
+    )
+
+    top_k, N, batch = 512, 131072, 2
+    logits, pre_idx, seq_lens = _make_inputs(
+        batch, N, top_k, torch.float32, 1, seed=11, varlen=True
+    )
+    dev = logits.device
+    n_eff = seq_lens.to(device=dev, dtype=torch.long)
+    if mode == "bucketed":
+        lines = _lines_at_counts(logits, n_eff, (20000, 4000, 600))
+    else:
+        n0 = {"hit": 4096, "pads": 4096, "hist": 12000, "void": 30000}[mode]
+        l0 = _lines_at_counts(logits, n_eff, (n0,))
+        lines = torch.cat(
+            [l0, torch.full_like(l0, LIST_PARK_LINE), torch.full_like(l0, 2 * LIST_PARK_LINE)], 1
+        )
+    seed_row = _pack_seed_row(logits, n_eff, lines)
+    cand_vals = torch.full((batch, LIST_WIDTH), float("-inf"), dtype=torch.float32, device=dev)
+    cand_idx = torch.full((batch, LIST_WIDTH), -1, dtype=torch.int32, device=dev)
+    cand_ctl = torch.zeros((batch, 4), dtype=torch.int32, device=dev)
+    pads = 64 if mode == "pads" else 0
+    for r in range(batch):
+        ne = int(n_eff[r])
+        row = logits[r, :ne]
+        hits = torch.nonzero(row >= lines[r, 0], as_tuple=False).flatten()
+        cand_ctl[r, 2] = int((row >= lines[r, 1]).sum())
+        cand_ctl[r, 3] = int((row >= lines[r, 2]).sum())
+        # emission order is value-blind: shuffle, then classify by the
+        # tightest line passed; a full segment spills to the looser one
+        perm = hits[torch.randperm(hits.numel(), device=dev)]
+        v = row[perm]
+        seg = torch.where(v >= lines[r, 2], 0, torch.where(v >= lines[r, 1], 1, 2))
+        in_a = seg == 0
+        ord_a = torch.cumsum(in_a.int(), 0)
+        stay_a = in_a & (ord_a <= LIST_SEG_A)
+        in_b = (seg == 1) | (in_a & ~stay_a)
+        ord_b = torch.cumsum(in_b.int(), 0)
+        stay_b = in_b & (ord_b <= LIST_SEG_A)
+        in_c = (seg == 2) | (in_b & ~stay_b)
+        ord_c = torch.cumsum(in_c.int(), 0)
+        stay_c = in_c & (ord_c <= LIST_CAP_C - pads)
+        slot = torch.full_like(seg, -1)
+        slot[stay_a] = ord_a[stay_a] - 1
+        slot[stay_b] = LIST_SEG_A + (ord_b[stay_b] - 1)
+        slot[stay_c] = 2 * LIST_SEG_A + (ord_c[stay_c] - 1)
+        if pads:
+            # sentinels displace C entries later in emission order: the
+            # kept ordinals shift by how many sentinels landed before them
+            pad_slots = torch.randperm(int(stay_c.sum()) + pads, device=dev)[:pads]
+            keep = slot[stay_c] - 2 * LIST_SEG_A
+            shift = (pad_slots[None, :] <= keep[:, None]).sum(-1)
+            slot[stay_c] = 2 * LIST_SEG_A + keep + shift
+        live = slot >= 0
+        cand_vals[r, slot[live].long()] = v[live]
+        cand_idx[r, slot[live].long()] = perm[live].int()
+        cand_ctl[r, 0] = int(hits.numel()) + pads
+        cand_ctl[r, 1] = 1 if int(in_c.sum()) > LIST_CAP_C - pads else 0
+    out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda")
+    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+        logits,
+        pre_idx,
+        seq_lens,
+        out_indices,
+        top_k=top_k,
+        cluster_size=1,
+        seed_thr=seed_row,
+        cand_vals=cand_vals,
+        cand_idx=cand_idx,
+        cand_ctl=cand_ctl,
+        accept_cap=LIST_SEG_A,
+        num_threads=512,
+    )
+    torch.cuda.synchronize()
+    _tie_aware_check(out_indices, logits, seq_lens, top_k, 1)
+
+
+@skip_not_sm100
+@pytest.mark.parametrize("tail_mode", ["exact", "pad_inf"])
+def test_cute_dsl_gvr_topk_decode_ext_block_max(tail_mode):
+    """32-grain positional upper-bound records + packed seed row.
+
+    The skip walk may only skip units whose bound clears no line, so it
+    must be exact under both legal tail bounds: the tight max over valid
+    positions and the worst legal inflation (+FLT_MAX on a partially
+    valid record).
+    """
+    top_k, N, batch = 1024, 262144, 2
+    seq_lens = torch.tensor([N, N - 37], dtype=torch.int32, device="cuda")
+    logits, pre_idx, seq_lens = _make_inputs(
+        batch, N, top_k, torch.float32, 1, seed=13, seq_lens=seq_lens
+    )
+    dev = logits.device
+    n_eff = seq_lens.to(device=dev, dtype=torch.long)
+    pos = torch.arange(N, device=dev)[None, :]
+    masked = torch.where(pos < n_eff[:, None], logits, float("-inf"))
+    records = masked.view(batch, N // 32, 32).amax(-1)
+    if tail_mode == "pad_inf":
+        rec_start = torch.arange(N // 32, device=dev)[None, :] * 32
+        partial = (rec_start < n_eff[:, None]) & (rec_start + 32 > n_eff[:, None])
+        records = torch.where(partial, torch.full_like(records, _FLT_MAX), records)
+    records = records.contiguous()
+    lines = _lines_at_counts(logits, n_eff, (4 * top_k, 2 * top_k, top_k + top_k // 4))
+    seed_row = _pack_seed_row(logits, n_eff, lines, block_max=records)
+    out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda")
+    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+        logits,
+        pre_idx,
+        seq_lens,
+        out_indices,
+        top_k=top_k,
+        cluster_size=1,
+        seed_thr=seed_row,
+        block_max=records,
+    )
+    torch.cuda.synchronize()
+    _tie_aware_check(out_indices, logits, seq_lens, top_k, 1)

From fb9643ece80839c07635cbc4f7d27e9891da550c Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 9 Aug 2026 10:49:54 -0700
Subject: [PATCH 099/117] [None][chore] GVR routing exclusivity, entry guards,
 comment slimming

- rungs tier: the skip prefix and the cluster split are mutually
  exclusive; the cluster split is decided first and the prefix attaches
  only when the row keeps cluster_size 1 (the K<=512, n>=196608 overlap
  previously set both)
- merge the two identical weak-band constants into ASSIST_WEAK_MAX_N
- kernel ctor rejects p4_no_fine without the exact-tail chain or with
  the range-test arms, and kNumBins that leave no tail pair buffer
- op entry rejects ext tensors with non-fp32 logits or next_n > 1 (the
  ext tiers are compiled for fp32 single-token rows only)
- pure-tie precheck: the fine-bin recompute is compiled out under
  p4_no_fine, matching the scatter
- slim narrative comments across the GVR files (verified equivalent by
  docstring-stripped AST compare); fix stale docs (p4_tail_fast default
  claim, the tail repair description, the FP4 runner emit-meta shapes);
  debug envs renamed under the TRTLLM_GVR_ prefix

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  54 ++--
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 280 +++++++-----------
 .../blackwell/top_k/gvr_routing.py            | 110 ++-----
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    |  66 ++---
 4 files changed, 176 insertions(+), 334 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 293b2276449c..0ca818b3b6f5 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -7430,8 +7430,7 @@ def _compile(
             order_row_fake = (cute.runtime.make_fake_compact_tensor(
                 cutlass.Int32,
                 (n_batch, ), stride_order=(0, )) if seqlen_sorted else None)
-            # emission-assisted tiers (list/counts/rungs, see
-            # gvr_routing): fake shapes mirror the harness wrapper
+            # emission-assisted tier fake tensors (list/counts/rungs)
             block_max_fake = (cute.runtime.make_fake_compact_tensor(
                 cutlass.Float32, (n_rows, cute.sym_int()),
                 stride_order=(1, 0),
@@ -7482,8 +7481,8 @@ def _compile(
                 cand_cap=cand_cap,
                 accept_cap=accept_cap,
                 kc_override=kc_override,
-                # ext modes need 3 rung slots (M_thr == 3); the qfrac
-                # VALUES are irrelevant (P1b skipped), only slot count
+                # ext modes need 3 rung slots (M_thr == 3); only the
+                # slot count matters, not the qfrac values (P1b skipped)
                 r0_qfracs=((0.85, 0.35) if
                            (use_ext_counts or ext_rungs) else None),
             )
@@ -7768,6 +7767,9 @@ def forward(
                     or enable_block_skip):
                 assert not lb_mode and order_row is None, (
                     "ext tiers are single-CTA/sort-path only")
+                assert logits.dtype == torch.float32 and next_n == 1, (
+                    "ext tiers are compiled for fp32 logits and next_n==1; "
+                    f"got dtype={logits.dtype} next_n={next_n}")
             if num_threads is not None:
                 tuning = dict(tuning, num_threads_per_block=num_threads)
             elif use_ext_cand and top_k <= 512:
@@ -9039,10 +9041,8 @@ def _compile(cls,
             seed_counts_fake = None
             if emit_seed_counts:
                 if seed_packed:
-                    # single [rows, 8] fp32 packed seed row: lines at
-                    # cols 0..2 (kernel reads (row, j<=2) unchanged),
-                    # counts accumulate as fp32 at cols 3..5; the same
-                    # tensor is bound to both params.
+                    # [rows, 8] fp32 packed seed row: lines at cols 0..2,
+                    # counts at cols 3..5; same tensor bound to both params.
                     seed_thr_fake = cute.runtime.make_fake_compact_tensor(
                         cutlass.Float32, (cute.sym_int(), 8),
                         stride_order=(1, 0),
@@ -9097,11 +9097,8 @@ def _compile(cls,
                 sm_fake,
                 cutlass.Int32(1),
                 cutlass.Int32(1),
-                # stream sits before the emission tensors in __call__ -
-                # upstream's positional order, with the emission slots
-                # appended after it. Keep this list in that order: the
-                # runtime call below drops the stream (the TVM FFI env
-                # stream is used) and is otherwise the same sequence.
+                # keep __call__'s argument order: stream, then emission
+                # slots (the runtime call drops the stream, same sequence)
                 fake_stream,
                 block_max_fake,
                 hit_stats_fake,
@@ -9164,10 +9161,12 @@ def forward(
                 output_dtype: output logits dtype
                 emit_block_meta: also emit per-128-block metadata for the
                     fused GVR top-k. Requires ``hit_bitmap``
-                    [B, nb_pad*4] int32 (1 bit per compressed kv position,
-                    request-level). ``block_max_out``/``hit_stats_out``
-                    ([B*next_n, nb_pad] fp32 / [B*next_n, nb_pad, 4] fp32)
-                    are allocated when not supplied.
+                    [B, >= nb_pad*4] int32 (1 bit per compressed kv
+                    position, request-level). ``block_max_out``
+                    [B*next_n, nb_pad*4] fp32 (4 warp-partial records per
+                    block) is allocated when not supplied;
+                    ``hit_stats_out`` [B*next_n, 4] fp32 must be supplied
+                    when ``emit_hit_stats`` is set.
             Returns:
                 logits [B*next_n, max_context_len]; with emit_block_meta,
                 the tuple (logits, block_max, hit_stats).
@@ -9265,10 +9264,9 @@ def forward(
                     "emit_seed_counts requires emit_block_meta")
                 if seed_counts_out is None:
                     # Packed contract: seed_thr IS the [rows, 8] fp32 seed
-                    # row (top-k pre-packed layout). Lines at cols 0..2;
-                    # counts accumulate as fp32 at cols 3..5 - the caller
-                    # zeroes cols 3..7 and writes lines each step. The
-                    # same buffer then feeds the top-k launch directly.
+                    # row; lines at cols 0..2, counts accumulate as fp32 at
+                    # cols 3..5. Caller zeroes cols 3..7 and writes lines
+                    # each step.
                     seed_packed = True
                     assert (
                         seed_thr is not None and seed_thr.dtype == torch.float32
@@ -9325,7 +9323,7 @@ def forward(
             elif emit_cand_bucketed:
                 assert emit_seed_counts, (
                     "emit_cand_bucketed requires emit_seed_counts")
-                # SoA v5 contract: cand_out = fp32 VALUES [rows, 2*segA+capC],
+                # SoA contract: cand_out = fp32 VALUES [rows, 2*segA+capC],
                 # cand_idx_out = int32 positions (same width), cand_cur_out =
                 # int32 [rows, 4] cursors (caller-zeroed), cand_ctl_out =
                 # int32 [rows, 4] {n0, void, n1, n2} (caller-zeroed)
@@ -9407,11 +9405,8 @@ def forward(
                      None, None, None, None, None, None, None, None)
             return logits
 
-    # NOTE: the optional emission tensors ARE written by the kernel but
-    # deliberately not in mutates_args - torch.library's in-place
-    # bookkeeping IndexErrors on optional mutates left at their None
-    # default (same precedent as heuristic_scratch on
-    # trtllm::indexer_topk_decode).
+    # NOTE: the optional emission tensors ARE written by the kernel but must
+    # stay out of mutates_args (torch.library IndexErrors on None defaults).
     @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits",
                              mutates_args=(),
                              device_types="cuda")
@@ -9484,9 +9479,8 @@ def cute_dsl_fp4_paged_mqa_logits(
             cand_idx_out=cand_idx_out,
             cand_ctl_out=cand_ctl_out,
             cand_cur_out=cand_cur_out)
-        # with emission on, the runner returns (logits, block_max,
-        # hit_stats) - the emission buffers are caller-owned mutates,
-        # the op face stays logits-only
+        # with emission on the runner returns a tuple; the op returns
+        # logits only (emission buffers are caller-owned)
         return ret[0] if isinstance(ret, tuple) else ret
 
     @torch.library.register_fake("trtllm::cute_dsl_fp4_paged_mqa_logits")
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index dd587afc330c..b4aac136177b 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -64,21 +64,16 @@
 # form is also accepted by older wrappers, so keep it version-independent.
 _RND_RN = "rn"
 
-# Reduction identities for the emit_block_meta epilogue path; match the GVR
-# kernel's FLT_MAX/NEG_FLT_MAX sentinels (gvr_topk_decode.py) so the fused
-# Phase 1's degenerate checks behave identically to the gather path.
+# Reduction identities for emit_block_meta; must match the GVR kernel's
+# FLT_MAX/NEG_FLT_MAX sentinels (gvr_topk_decode.py).
 _META_FLT_MAX = 3.4028235e38
 _META_NEG_FLT_MAX = -3.4028235e38
 
 
-# ---------------------------------------------------------------------------
 # Global-memory reductions for the per-row hit aggregate (emit_hit_stats).
-# fp32 has no native atomic min/max; we use the standard order-preserving
-# int encoding enc(f) = bits(f) >= 0 ? bits(f) : bits(f) ^ 0x7FFFFFFF
-# (an involution; signed-int order == float order, -0.0 quirk harmless for
-# min/max seeding) and red.global.{min,max}.s32. Sum uses red.global.add.f32
-# (order-nondeterministic — perturbs only the heuristic mean seed).
-# ---------------------------------------------------------------------------
+# fp32 min/max use the order-preserving int encoding
+# enc(f) = bits(f) >= 0 ? bits(f) : bits(f) ^ 0x7FFFFFFF (an involution)
+# with red.global.{min,max}.s32; sum uses red.global.add.f32.
 @dsl_user_op
 def _red_global_fmin_ordered(addr_i64, fval, *, loc=None, ip=None):
     llvm.inline_asm(
@@ -519,84 +514,69 @@ def __init__(
         # one contiguous LSU phase after the for-t loop (epilogue micro-opt).
         self.use_batched_store = use_batched_store
         # When True, the epilogue additionally emits per-128-token-block
-        # metadata consumed by the fused GVR top-k (gvr_topk_decode.py —
-        # fused_preidx_stats / enable_block_skip). Emission is fully
-        # WARP-AUTONOMOUS: each of the WG's 4 warps writes one partial
-        # record per tile per t (record index = tile*4 + warp); the GVR
-        # consumer folds the 4 partials per block. No cross-warp barrier —
-        # a per-tile named-barrier fold costs +53% indexer wall-clock.
-        #   block_max [num_rows, nb_pad*4]    fp32 — warp-partial max of
-        #     f32(stored logit) over valid positions (kv_pos < ctx);
-        #     fold(4) is the lossless block-skip upper bound. Computed on
-        #     the POST-conversion value so it bounds what GVR reads back,
-        #     bit-exactly.
-        #   hit_agg [num_rows, 4] fp32 — PER-ROW aggregate
+        # metadata consumed by the fused GVR top-k (gvr_topk_decode.py).
+        # Emission is warp-autonomous: each of the WG's 4 warps writes one
+        # partial record per tile per t (record index = tile*4 + warp);
+        # the GVR consumer folds the 4 partials per block. No cross-warp
+        # barrier.
+        #   block_max [num_rows, nb_pad*4] fp32 — warp-partial max of
+        #     f32(stored logit) over valid positions (kv_pos < ctx),
+        #     computed on the POST-conversion value so it bounds what GVR
+        #     reads back bit-exactly.
+        #   hit_agg [num_rows, 4] fp32 — per-row aggregate
         #     {enc_min, enc_max, sum, cnt} of stored logits at positions
         #     flagged in hit_bitmap; min/max slots hold the
         #     order-preserving int encoding (see _red_global_fmin_ordered).
         #     Buffer must be pre-initialized to
-        #     {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step. Replaces GVR
-        #     Phase 1's random gather.
-        #   hit_bitmap [batch, nb_pad*4]      int32 — 1 bit per kv
-        #     position (request-level; from the previous step's top-k).
-        # Cost per tile per t: 1 bitmap LDG/thread + 1 warp redux + 1
-        # lane-0 STG (block_max) + <= 4 LANE-LOCAL accumulator ops (hit
-        # stats); the hit accumulators flush ONCE per q-transition via
-        # atomics (_flush_hit_agg). A per-tile warp-redux emission of the
-        # hit stats was tried first and cost +72% indexer wall-clock.
+        #     {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step.
+        #   hit_bitmap [batch, nb_pad*4] int32 — 1 bit per kv position
+        #     (request-level; from the previous step's top-k).
+        # Hit accumulators are lane-local and flush once per q-transition
+        # via atomics (_flush_hit_agg).
         #
         # emit_hit_stats sub-knob (only meaningful with emit_block_meta):
-        # False emits block_max ONLY (no bitmap read, no hit aggregate) —
-        # sufficient for GVR block-skip without the fused Phase 1.
+        # False emits block_max ONLY (no bitmap read, no hit aggregate).
         self.emit_block_meta = emit_block_meta
         self.emit_hit_stats = emit_hit_stats
-        # emit_seed_counts (L1 of the epilogue suite; requires
-        # emit_block_meta): per row, count stored logits >= each of the
-        # T=3 caller-provided thresholds (cross-step seed thresholds from
-        # the GVR xstate). Lane-local accumulation across tiles + one
-        # warp-redux + lane-0 red.global.add.s32 per threshold per row
-        # transition — the hit-stats cost profile. Counts are computed on
-        # the POST-conversION value over valid positions only, matching
-        # the top-k consumer's verification semantics
-        # (workspace/epilogue_topk_interface.md).
+        # emit_seed_counts (requires emit_block_meta): per row, count
+        # stored logits >= each of the T=3 caller-provided thresholds.
+        # Counts are computed on the POST-conversion value over valid
+        # positions only.
         if emit_seed_counts and not emit_block_meta:
             raise ValueError("emit_seed_counts requires emit_block_meta")
         self.emit_seed_counts = emit_seed_counts
         # seed_packed: single [num_rows, 8] fp32 seed row per the top-k
-        # pre-packed contract - lines at cols 0..2, counts ACCUMULATED AS
-        # FLOATS at cols 3..5 (exact to 2^24; red.global.add.f32). The
-        # caller zeroes cols 3..7 and writes the lines each step; the
-        # same buffer feeds the top-k launch with no host repack.
+        # pre-packed contract - lines at cols 0..2, counts accumulated as
+        # floats at cols 3..5 (exact to 2^24; red.global.add.f32). The
+        # caller zeroes cols 3..7 and writes the lines each step.
         if seed_packed and not emit_seed_counts:
             raise ValueError("seed_packed requires emit_seed_counts")
         self.seed_packed = seed_packed
-        # emit_cand (L2 of the epilogue suite): unordered pre-collect of all
-        # (value, index) pairs >= the t_0 seed threshold via warp ballot +
-        # lane0 batch atomic claim. claimed >= K certifies the candidate
-        # set covers the true top-K (contract in epilogue_topk_interface.md).
+        # emit_cand: unordered pre-collect of all (value, index) pairs >=
+        # the t_0 seed threshold. claimed >= K certifies the candidate set
+        # covers the true top-K.
         if emit_cand and not emit_seed_counts:
             raise ValueError("emit_cand requires emit_seed_counts (t_0 source)")
         self.emit_cand = emit_cand
         self.cand_cap = cand_cap
-        # emit_cand_bucketed (v5 list contract): three fixed SoA segments
-        # (A=[0,segA) holds >= t2, B=[segA,2segA) holds [t1,t2), C=
-        # [2segA,2segA+capC) holds [t0,t1)); a full segment spills to the
-        # next looser one. A/B use EXACT ballot claims (their prefixes
-        # must stay pad-free - the consumer's prefix math assumes it), C
-        # keeps the claim-window scheme (pads are legal there). Cursors
-        # live in caller-zeroed cand_cur [rows,4]; ctl [rows,4] carries
-        # {n0 incl C pads, void, n1, n2} with n1/n2 flushed from the L1
-        # seed counters.
+        # emit_cand_bucketed: three fixed SoA segments (A=[0,segA) holds
+        # >= t2, B=[segA,2segA) holds [t1,t2), C=[2segA,2segA+capC) holds
+        # [t0,t1)); a full segment spills to the next looser one. A/B use
+        # EXACT ballot claims (their prefixes must stay pad-free - the
+        # consumer's prefix math assumes it), C keeps the claim-window
+        # scheme (pads are legal there). Cursors live in caller-zeroed
+        # cand_cur [rows,4]; ctl [rows,4] carries {n0 incl C pads, void,
+        # n1, n2} with n1/n2 flushed from the seed counters.
         if emit_cand_bucketed and not emit_seed_counts:
             raise ValueError("emit_cand_bucketed requires emit_seed_counts")
         if emit_cand_bucketed and emit_cand:
             raise ValueError("emit_cand_bucketed and emit_cand are exclusive")
         self.emit_cand_bucketed = emit_cand_bucketed
         self.accept_cap = accept_cap
-        # Per-warp claim window: one atomic claims (hits + CAND_WIN) slots;
-        # subsequent hits consume the window latency-free. The unconsumed
-        # tail is sentinel-filled (idx = -1) at q-transition/loop end, so
-        # `claimed` over-approximates the true count (counts[r][0] exact).
+        # Per-warp claim window: one atomic claims (hits + CAND_WIN) slots.
+        # The unconsumed tail is sentinel-filled (idx = -1) at
+        # q-transition/loop end, so `claimed` over-approximates the true
+        # count (counts[r][0] exact).
         self.CAND_WIN = 8
         # epi_bytes covers fp16 and bf16 (FP8 only handled fp16).
         self.epi_bytes = 2 if epi_dtype in (cutlass.Float16, cutlass.BFloat16) else 4
@@ -819,8 +799,8 @@ def __call__(
         num_phys_blocks: cutlass.Int32,
         batch_size: cutlass.Int32,
         stream: cuda.CUstream,
-        # everything below is emission-only and defaulted, so the
-        # positional signature stays the one callers already use
+        # emission-only tensors; defaulted so the positional signature
+        # stays the one callers already use
         block_max: cute.Tensor = None,  # [num_rows, nb_pad*4] fp32 warp-partials
         hit_stats: cute.Tensor = None,  # [num_rows, 4] fp32 (emit_hit_stats)
         hit_bitmap: cute.Tensor = None,  # [batch, nb_pad*4] int32 (emit_block_meta)
@@ -1073,10 +1053,8 @@ def _flush_hit_agg(
         """Warp-reduce the per-lane hit accumulators and merge them into the
         per-row global aggregate via one set of atomics (encoded-int
         min/max + fp32 adds), then reset to identities. Called once per
-        q-transition per warp — atomic traffic is negligible. The
-        aggregate buffer must be pre-initialized to
-        {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step (prototype: test
-        harness; production: folded into the bitmap prepare kernel)."""
+        q-transition per warp. The aggregate buffer must be pre-initialized
+        to {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step."""
         next_n = cutlass.const_expr(self.next_n)
         base_addr = mHitAgg.iterator.toint()
         for t in cutlass.range_constexpr(next_n):
@@ -1133,8 +1111,7 @@ def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None, ca
                             _red_global_add_s32(ctl_a, w_cnt)
                 scnt[t * 3 + j] = cutlass.Int32(0)
         if cutlass.const_expr(self.seed_packed and spass is not None):
-            # packed col 6: adaptive-skip pass count (lane0-accumulated,
-            # so the warp redux is exactly the warp's record total)
+            # packed col 6: adaptive-skip pass count (lane0-accumulated)
             for t in cutlass.range_constexpr(next_n):
                 w_bp = cute.arch.warp_redux_sync(spass[t], "add")
                 if meta_lane == cutlass.Int32(0):
@@ -1148,7 +1125,7 @@ def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None, ca
     @cute.jit
     def _flush_cand_window_bucketed(self, mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane):
         """Sentinel-fill the unconsumed C-window tail in BOTH SoA columns
-        (score -inf, idx -1: the v5 consumer pads by score) and invalidate
+        (score -inf, idx -1: the consumer pads by score) and invalidate
         the window. Segment C sits at base 2*segA in each row."""
         next_n = cutlass.const_expr(self.next_n)
         segA_f = cutlass.const_expr(self.accept_cap)
@@ -1486,12 +1463,9 @@ def kernel(
             layout=sf_kv_smem_layout_staged,
             byte_alignment=128,
         )
-        # Block-meta emission is fully warp-autonomous (each warp's lane 0
-        # writes its own warp-partial record straight to GMEM; the GVR
-        # side folds 4 partials per block) — no SMEM scratch, no named
-        # barrier. A cross-warp SMEM fold + per-tile named barrier was
-        # tried first and cost +53% indexer wall-clock by serializing the
-        # epilogue's warp pipelining.
+        # Block-meta emission is warp-autonomous: each warp's lane 0 writes
+        # its own warp-partial record straight to GMEM; the GVR side folds
+        # 4 partials per block. No SMEM scratch, no named barrier.
 
         a_mcast_mask = cpasync.create_tma_multicast_mask(
             cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
@@ -2214,18 +2188,14 @@ def kernel(
                 else:  # fp32, 4-byte weights
                     MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64
                 if cutlass.const_expr(self.emit_block_meta):
-                    # Free ~8 registers for the meta accumulators/fragments
-                    # — the epilogue's weight cache is tuned to the spill
-                    # edge (see MAX_NUM_W_IN_REG SASS notes above).
+                    # Free ~8 registers for the meta accumulators/fragments;
+                    # the epilogue's weight cache sits at the spill edge.
                     MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
                     if cutlass.const_expr(self.emit_hit_stats):
                         # Hit accumulators + bitmap word add ~6 more live
                         # registers across the tile loop.
                         MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
-                    # emit_seed_counts: no budget cut — ncu shows the cost
-                    # is epilogue ALU/issue, not registers (occupancy and
-                    # block limits identical with or without a -8 cut; the
-                    # cut itself measured neutral-to-slower).
+                    # emit_seed_counts needs no extra budget cut.
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the
@@ -2256,10 +2226,9 @@ def kernel(
                             cwbase[_i] = cutlass.Int32(0)
                             cwleft[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_hit_stats):
-                        # Per-lane hit accumulators, carried across ALL
+                        # Per-lane hit accumulators, carried across all
                         # tiles of the same q and flushed once per
-                        # q-transition — zero warp-wide ops per tile (the
-                        # per-tile redux version cost +72% wall-clock).
+                        # q-transition — no warp-wide ops per tile.
                         hacc_min = cute.make_fragment(next_n, cutlass.Float32)
                         hacc_max = cute.make_fragment(next_n, cutlass.Float32)
                         hacc_sum = cute.make_fragment(next_n, cutlass.Float32)
@@ -2269,13 +2238,10 @@ def kernel(
                             hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX)
                             hacc_sum[_t] = cutlass.Float32(0.0)
                             hacc_cnt[_t] = cutlass.Int32(0)
-                        # Batched bitmap read state: a per-tile LDG's
-                        # ~300cy L2 trip lands on the critical path every
-                        # tile (isolated at +42% kernel time; one-ahead
-                        # prefetch did not help). All 32 lanes of a warp
-                        # need the SAME word per tile, so instead lane l
-                        # loads the word for tile j+l once per 32 tiles
-                        # and each tile takes its word via one shuffle.
+                        # Batched bitmap read state: all 32 lanes of a warp
+                        # need the SAME word per tile, so lane l loads the
+                        # word for tile j+l once per 32 tiles and each tile
+                        # takes its word via one shuffle.
                         meta_j = cutlass.Int32(0)
                         hitw_batch = cutlass.Int32(0)
 
@@ -2559,18 +2525,15 @@ def kernel(
                             mLogits[(out_row, kv_pos)] = stored_t
                         if cutlass.const_expr(self.emit_block_meta):
                             # Meta reduction on the POST-conversion value so
-                            # block_max bounds what GVR reads back bit-exactly
-                            # (pre-conversion fp32 max could round up past a
-                            # stored logit's converted value's block max).
+                            # block_max bounds what GVR reads back bit-exactly.
                             f32_t = cutlass.Float32(stored_t)
                             bmax_v = cutlass.Float32(_META_NEG_FLT_MAX)
                             if meta_valid:
                                 bmax_v = f32_t
                             r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax")
                             # Warp-autonomous store: record index =
-                            # tile*4 + warp; the GVR consumer folds the 4
-                            # warp-partials per block (fold of partials ==
-                            # block stat — associative + identity-padded).
+                            # tile*4 + warp; the GVR consumer folds the
+                            # 4 warp-partials per block.
                             if meta_lane == cutlass.Int32(0):
                                 out_row_m = q_idx * next_n + t
                                 rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp
@@ -2588,26 +2551,20 @@ def kernel(
                                     # adaptive-skip pass count: one record
                                     # per (tile, warp); r_bmax is warp-
                                     # uniform so lane0 alone accumulates
-                                    # (the flush redux then sums warps)
                                     if meta_lane == cutlass.Int32(0):
                                         spass[t] = spass[t] + cutlass.Int32(
                                             r_bmax >= sthr[t * 3 + 0]
                                         )
                             if cutlass.const_expr(self.emit_cand):
-                                # L2 pre-collect at t_0 with per-warp claim
-                                # WINDOWS: refills claim (hits + CAND_WIN)
-                                # slots in ONE atomic; between refills the
-                                # warp consumes its window latency-free (a
-                                # per-hit atomic round-trip stalled the
-                                # epilogue 2x at 131k ctx). Unconsumed tail
-                                # is sentinel-filled on flush; counts[r][0]
-                                # stays the exact count. Gated on the
-                                # already-computed warp-uniform 32-position
-                                # bound: r_bmax < t_0 proves zero hits, so
-                                # the common no-hit iteration costs ONE fp
-                                # compare (no ballot, no select). bound >=
-                                # t_0 guarantees a nonzero ballot (exact
-                                # per-lane max, invalid -> -FLT_MAX).
+                                # Candidate pre-collect at t_0 with per-warp
+                                # claim windows: one atomic claims
+                                # (hits + CAND_WIN) slots per refill.
+                                # Unconsumed tail is sentinel-filled on
+                                # flush; counts[r][0] stays the exact count.
+                                # Gated on the warp-uniform 32-position
+                                # bound: r_bmax < t_0 proves zero hits;
+                                # bound >= t_0 guarantees a nonzero ballot
+                                # (exact per-lane max, invalid -> -FLT_MAX).
                                 if r_bmax >= sthr[t * 3 + 0]:
                                     pred_c = cutlass.Int32(0)
                                     if meta_valid:
@@ -2687,7 +2644,7 @@ def kernel(
                                     cwbase[t] = cwbase[t] + cnt_c
                                     cwleft[t] = cwleft[t] - cnt_c
                             if cutlass.const_expr(self.emit_cand_bucketed):
-                                # v5 bucketed SoA: A/B EXACT ballot claims
+                                # Bucketed SoA: A/B EXACT ballot claims
                                 # (their prefixes must stay pad-free for
                                 # the consumer's prefix math), C keeps the
                                 # claim-window; a full segment spills to
@@ -2901,14 +2858,12 @@ def kernel(
                                     cwbase[t] = cwbase[t] + cntC_k
                                     cwleft[t] = cwleft[t] - cntC_k
                             if cutlass.const_expr(self.emit_hit_stats):
-                                # Lane-local accumulation, fully BRANCHLESS
-                                # (data-dependent `if meta_hit` compiled to
-                                # real divergent branches whose condition
-                                # waits on the bitmap LDG — 43% of all warp
-                                # stall samples). Bit-mask select is also
-                                # NaN-safe for OOB-tile garbage logits. No
-                                # valid-mask needed: the bitmap contract
-                                # only sets bits inside [0, ctx).
+                                # Lane-local accumulation; keep it branchless
+                                # (`if meta_hit` compiles to real divergent
+                                # branches). Bit-mask select is NaN-safe for
+                                # OOB-tile garbage logits. No valid-mask
+                                # needed: the bitmap contract only sets bits
+                                # inside [0, ctx).
                                 meta_hit = (
                                     hit_word >> (kv_pos & cutlass.Int32(31))
                                 ) & cutlass.Int32(1)
@@ -3007,18 +2962,14 @@ def kernel(
                 else:
                     MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64
                 if cutlass.const_expr(self.emit_block_meta):
-                    # Free ~8 registers for the meta accumulators/fragments
-                    # — the epilogue's weight cache is tuned to the spill
-                    # edge (see MAX_NUM_W_IN_REG SASS notes above).
+                    # Free ~8 registers for the meta accumulators/fragments;
+                    # the epilogue's weight cache sits at the spill edge.
                     MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
                     if cutlass.const_expr(self.emit_hit_stats):
                         # Hit accumulators + bitmap word add ~6 more live
                         # registers across the tile loop.
                         MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8
-                    # emit_seed_counts: no budget cut — ncu shows the cost
-                    # is epilogue ALU/issue, not registers (occupancy and
-                    # block limits identical with or without a -8 cut; the
-                    # cut itself measured neutral-to-slower).
+                    # emit_seed_counts needs no extra budget cut.
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
                 # Batched STG: hold reduced result per t in register; the
@@ -3049,10 +3000,9 @@ def kernel(
                             cwbase[_i] = cutlass.Int32(0)
                             cwleft[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_hit_stats):
-                        # Per-lane hit accumulators, carried across ALL
+                        # Per-lane hit accumulators, carried across all
                         # tiles of the same q and flushed once per
-                        # q-transition — zero warp-wide ops per tile (the
-                        # per-tile redux version cost +72% wall-clock).
+                        # q-transition — no warp-wide ops per tile.
                         hacc_min = cute.make_fragment(next_n, cutlass.Float32)
                         hacc_max = cute.make_fragment(next_n, cutlass.Float32)
                         hacc_sum = cute.make_fragment(next_n, cutlass.Float32)
@@ -3062,13 +3012,10 @@ def kernel(
                             hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX)
                             hacc_sum[_t] = cutlass.Float32(0.0)
                             hacc_cnt[_t] = cutlass.Int32(0)
-                        # Batched bitmap read state: a per-tile LDG's
-                        # ~300cy L2 trip lands on the critical path every
-                        # tile (isolated at +42% kernel time; one-ahead
-                        # prefetch did not help). All 32 lanes of a warp
-                        # need the SAME word per tile, so instead lane l
-                        # loads the word for tile j+l once per 32 tiles
-                        # and each tile takes its word via one shuffle.
+                        # Batched bitmap read state: all 32 lanes of a warp
+                        # need the SAME word per tile, so lane l loads the
+                        # word for tile j+l once per 32 tiles and each tile
+                        # takes its word via one shuffle.
                         meta_j = cutlass.Int32(0)
                         hitw_batch = cutlass.Int32(0)
 
@@ -3345,18 +3292,15 @@ def kernel(
                             mLogits[(out_row, kv_pos)] = stored_t
                         if cutlass.const_expr(self.emit_block_meta):
                             # Meta reduction on the POST-conversion value so
-                            # block_max bounds what GVR reads back bit-exactly
-                            # (pre-conversion fp32 max could round up past a
-                            # stored logit's converted value's block max).
+                            # block_max bounds what GVR reads back bit-exactly.
                             f32_t = cutlass.Float32(stored_t)
                             bmax_v = cutlass.Float32(_META_NEG_FLT_MAX)
                             if meta_valid:
                                 bmax_v = f32_t
                             r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax")
                             # Warp-autonomous store: record index =
-                            # tile*4 + warp; the GVR consumer folds the 4
-                            # warp-partials per block (fold of partials ==
-                            # block stat — associative + identity-padded).
+                            # tile*4 + warp; the GVR consumer folds the
+                            # 4 warp-partials per block.
                             if meta_lane == cutlass.Int32(0):
                                 out_row_m = q_idx * next_n + t
                                 rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp
@@ -3374,26 +3318,20 @@ def kernel(
                                     # adaptive-skip pass count: one record
                                     # per (tile, warp); r_bmax is warp-
                                     # uniform so lane0 alone accumulates
-                                    # (the flush redux then sums warps)
                                     if meta_lane == cutlass.Int32(0):
                                         spass[t] = spass[t] + cutlass.Int32(
                                             r_bmax >= sthr[t * 3 + 0]
                                         )
                             if cutlass.const_expr(self.emit_cand):
-                                # L2 pre-collect at t_0 with per-warp claim
-                                # WINDOWS: refills claim (hits + CAND_WIN)
-                                # slots in ONE atomic; between refills the
-                                # warp consumes its window latency-free (a
-                                # per-hit atomic round-trip stalled the
-                                # epilogue 2x at 131k ctx). Unconsumed tail
-                                # is sentinel-filled on flush; counts[r][0]
-                                # stays the exact count. Gated on the
-                                # already-computed warp-uniform 32-position
-                                # bound: r_bmax < t_0 proves zero hits, so
-                                # the common no-hit iteration costs ONE fp
-                                # compare (no ballot, no select). bound >=
-                                # t_0 guarantees a nonzero ballot (exact
-                                # per-lane max, invalid -> -FLT_MAX).
+                                # Candidate pre-collect at t_0 with per-warp
+                                # claim windows: one atomic claims
+                                # (hits + CAND_WIN) slots per refill.
+                                # Unconsumed tail is sentinel-filled on
+                                # flush; counts[r][0] stays the exact count.
+                                # Gated on the warp-uniform 32-position
+                                # bound: r_bmax < t_0 proves zero hits;
+                                # bound >= t_0 guarantees a nonzero ballot
+                                # (exact per-lane max, invalid -> -FLT_MAX).
                                 if r_bmax >= sthr[t * 3 + 0]:
                                     pred_c = cutlass.Int32(0)
                                     if meta_valid:
@@ -3473,7 +3411,7 @@ def kernel(
                                     cwbase[t] = cwbase[t] + cnt_c
                                     cwleft[t] = cwleft[t] - cnt_c
                             if cutlass.const_expr(self.emit_cand_bucketed):
-                                # v5 bucketed SoA: A/B EXACT ballot claims
+                                # Bucketed SoA: A/B EXACT ballot claims
                                 # (their prefixes must stay pad-free for
                                 # the consumer's prefix math), C keeps the
                                 # claim-window; a full segment spills to
@@ -3687,14 +3625,12 @@ def kernel(
                                     cwbase[t] = cwbase[t] + cntC_k
                                     cwleft[t] = cwleft[t] - cntC_k
                             if cutlass.const_expr(self.emit_hit_stats):
-                                # Lane-local accumulation, fully BRANCHLESS
-                                # (data-dependent `if meta_hit` compiled to
-                                # real divergent branches whose condition
-                                # waits on the bitmap LDG — 43% of all warp
-                                # stall samples). Bit-mask select is also
-                                # NaN-safe for OOB-tile garbage logits. No
-                                # valid-mask needed: the bitmap contract
-                                # only sets bits inside [0, ctx).
+                                # Lane-local accumulation; keep it branchless
+                                # (`if meta_hit` compiles to real divergent
+                                # branches). Bit-mask select is NaN-safe for
+                                # OOB-tile garbage logits. No valid-mask
+                                # needed: the bitmap contract only sets bits
+                                # inside [0, ctx).
                                 meta_hit = (
                                     hit_word >> (kv_pos & cutlass.Int32(31))
                                 ) & cutlass.Int32(1)
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
index ed081c5e7c84..2bc3992e0b15 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py
@@ -30,104 +30,51 @@
 NEXT step (the emission tax is shape-dependent); ``pick_config`` maps
 (tier, B, N, K) to concrete launch knobs for THIS step.
 
-All thresholds are measured on B200 (f15 layer-complete grid,
-2026-07-27, cold-L2 kernel-only protocol, validated on both the shared
-grid dataset and first-party captures). They are deployment defaults,
-not universal truths - keep them in one place so retuning is a
-constant edit, not a logic edit.
+All thresholds are B200 deployment defaults, kept in one place so
+retuning is a constant edit, not a logic edit.
 """
 
 from dataclasses import dataclass
 from typing import Optional
 
-# ---- measured thresholds (B200) ------------------------------------------
+# ---- thresholds (B200 deployment defaults) --------------------------------
 #
-# UNITS: every threshold here is fitted on KERNEL-ONLY time for both the
-# indexer and the top-k. An earlier fit compared a wall-clock indexer
-# (which carries a ~22us launch floor) against kernel-only top-k times;
-# that inflated the apparent emission budget at small shapes and is what
-# put the counts tier behind a batch gate it never needed.
-
-# The kernel's fixed cost does not shrink with N, so once the stock
-# kernel finishes under that floor no tier can win.
+# UNITS: every threshold is fitted on KERNEL-ONLY time for both the
+# indexer and the top-k; do not mix in wall-clock numbers when retuning.
+
+# Below this the stock kernel already runs at the fixed-cost floor.
 ASSIST_MIN_N_COMP = 2048
 
-# Block-skip prefix pays for the counts tier from 65536 up, but not for
-# the zero-emission rungs tier below 131072. Measured on captured V4
-# rows (5 layers x all decode steps, nsys kernel-only, flash n_comp
-# 65537): counts 16.1us with the prefix vs 17.3 without at batch
-# 4/8/16, rungs 18.5 with vs 17.9 without. A synthetic-Gaussian A/B put
-# the counts break-even a doubling later - the block-max distribution
-# of real rows decides the pass rate, so set this from captures only.
+# Block-skip prefix break-evens. Retune these from real captures only:
+# the real block-max distribution decides the pass rate.
 SKIP_MIN_N_COUNTS = 65536  # va: attach block_max from here up
 SKIP_MIN_N_RUNGS_FLASH = 131072  # vb (flash): bm pays from here
-# Cluster split is a loss for the assist tiers below this point,
-# block_max or not: at n_comp 65537 rungs measures 17.9 / 19.0 / 20.4us
-# at cs 1 / 4 / 8 (batch 4) and cs8 spills to 31.5us at batch 16.
 SKIP_CS_MIN_N_RUNGS = 196608  # vb: cluster split from here up
 
-# Emission cost measured on the FP4 indexer KERNEL (nsys kernel-only,
-# ABBA-interleaved NVTX blocks, batch 1..64 x ctx 32k..512k) as the delta
-# of the same kernel with the emission outputs attached:
-#
-#   batch                1     2     4     8    16    64
-#   counts (us)        +5.6  +4.4  +2.9  +2.0  +0.4   0.0
-#   list, bucketed     +6.2  +6.2  +6.8 +10.5 +25.0  +63.1  (ctx 32k)
-#                     +10.3 +14.7 +25.6 +44.5 +85.7 +417.8  (ctx 512k)
-#
-# The counts cost is a fixed reduction-latency chain that depends on
-# BATCH ONLY and hides once there are enough rows to overlap it - it is
-# not a fraction of the indexer, so it cannot price the tier out at
-# scale. An earlier fit modelled it as 3% of a WALL-clock indexer time
-# (which carries a ~22us launch floor) and so charged 12us at batch 64 /
-# 512k, where the true cost is zero; that is what put the counts tier
-# behind a batch gate. The list cost is real per-emitted-entry work and
-# grows with batch and context alike. Parking the two tight lines above
-# the score range (see gvr_ext.LIST_PARK_LINE) drops it 2.7-5x - every
-# entry then lands in the one segment that claims through a per-warp
-# window instead of an exact ballot - at top-k parity, which is what
-# makes the list tier affordable past a single row.
+# Emission-cost model: counts emission is a batch-only latency chain
+# (hides at large batch); list emission grows with batch and context.
 LIST_EMIT_MIN_N = 65536  # shorter rows: the emission outweighs the saving
 LIST_EMIT_MAX_B = 4  # past four rows the list stops repaying its emission
-COUNTS_MIN_TOKENS = 524288  # B * raw length; below this the counts
-# latency chain is exposed and the zero-emission rungs tier wins
+COUNTS_MIN_TOKENS = 524288  # B * raw length; below this the rungs tier wins
 RUNGS_ONLY_MIN_N = 16384  # short-row band where rungs also beats counts
 RUNGS_ONLY_MAX_N = 49152
 
-# Mid-row weak band: rows long enough that the stock kernel splits them
-# across a cluster, but short enough (and at a small enough batch) that
-# its split grid still fits one wave. The assist tiers cannot follow -
-# splitting a row costs them more than the scan it saves, with or without
-# the block-skip prefix - so the stock kernel wins outright and the
-# epilogue should emit nothing. Narrowed after phase 4 got its coarse
-# search and its boundary-class repair back: the tiers gained about
-# 1.5us there, which is enough to take the upper half of the band back
-# off the stock kernel (2 cells fall through now, down from 4).
+# Mid-row weak band: the stock kernel's split grid fits one wave here
+# and out-scans every assist tier, so the epilogue emits nothing.
 ASSIST_WEAK_MIN_N = 49152
-# Narrowed again from 98304 to 65536: at n_comp 65536 the band's premise no
-# longer holds. Measured on the full grid (kernel-only, us/step), the stock
-# kernel is NOT the fastest arm there - counts beats it by 2.19 (flash) and
-# 2.38 (pro), and the counts emission at that batch costs only +0.27us on the
-# indexer (re-measured after the tight-line parking; the old table charged
-# +2.0us at batch 8, which is what kept the band this wide). Net +1.9 to
-# +2.1us per step, i.e. those cells go 1.11 -> ~1.31 against the baseline.
-# The interior of the band (49152 <= n_comp < 65536) is untouched: no grid
-# unit lands there, so it stays on the stock kernel until someone measures it.
-ASSIST_WEAK_MAX_N_SMALL_K = 65536  # k <= ASSIST_WEAK_K
-ASSIST_WEAK_MAX_N_LARGE_K = 65536
-ASSIST_WEAK_K = 512
+# The band interior is unmeasured; it stays on the stock kernel.
+ASSIST_WEAK_MAX_N = 65536
 ASSIST_WEAK_MAX_B = 8
 
-# rungs-tier block_max pays only at small K: with K=1024 the tight-line
-# pass rate runs too high and the prefix read is pure overhead.
+# rungs block_max pays only at small K; at large K the prefix is overhead.
 RUNGS_BM_MAX_K = 512
 
 # 512-thread build wins for list-hit rows at small K (work is O(list)).
 SMALL_K_LIST_THREADS = 512
 SMALL_K_MAX = 512
 
-# GPC packing: cs=8 only while all row-clusters fit half the device
-# (B=16 x cs8 wave-spill regression); cs4/2 keep a 10% headroom.
+# GPC packing: cs=8 only while all row-clusters fit half the device;
+# cs4/2 keep a 10% headroom.
 CS8_HALF_DEVICE = 2
 CS_HEADROOM_NUM = 9
 CS_HEADROOM_DEN = 10
@@ -153,16 +100,12 @@ def plan_emission(
     matching buffers and the next top-k launch routes on them.
     """
     if n_comp < ASSIST_MIN_N_COMP:
-        # short rows: the stock kernel is already under our fixed cost,
-        # and this holds for the zero-emission rungs tier too - it is
-        # the same kernel, so the floor is the same
+        # short rows: the stock kernel is already under the fixed cost
         return "none"
     if have_epilogue and n_comp >= LIST_EMIT_MIN_N and batch <= LIST_EMIT_MAX_B:
-        # checked before the weak band below: that band is about the stock
-        # kernel out-scanning us, and a list hit never scans the row
+        # must stay ahead of the weak-band gate: a list hit never scans the row
         return "list"
-    weak_max = ASSIST_WEAK_MAX_N_SMALL_K if k <= ASSIST_WEAK_K else ASSIST_WEAK_MAX_N_LARGE_K
-    if batch <= ASSIST_WEAK_MAX_B and ASSIST_WEAK_MIN_N <= n_comp < weak_max:
+    if batch <= ASSIST_WEAK_MAX_B and ASSIST_WEAK_MIN_N <= n_comp < ASSIST_WEAK_MAX_N:
         return "none"  # stock's split grid wins this band outright
     if (
         have_epilogue
@@ -181,16 +124,13 @@ def pick_config(tier: str, batch: int, n_comp: int, k: int, num_sms: int) -> Top
     if tier == "list":
         if k <= SMALL_K_MAX:
             r.num_threads = SMALL_K_LIST_THREADS
-        # list + block_max: miss rows fall back to a skip-walk instead
-        # of a dense re-scan (measured -19% on pro long chains).
+        # block_max: miss rows take a skip-walk instead of a dense re-scan
         r.attach_block_max = n_comp >= SKIP_MIN_N_COUNTS
         return r
     if tier == "counts":
         r.attach_block_max = n_comp >= SKIP_MIN_N_COUNTS
         return r
-    # rungs (vb)
-    if k <= RUNGS_BM_MAX_K and n_comp >= SKIP_MIN_N_RUNGS_FLASH:
-        r.attach_block_max = True
+    # rungs (vb); the skip prefix and cluster split are mutually exclusive
     if n_comp >= SKIP_CS_MIN_N_RUNGS:
         if batch * 8 <= num_sms // CS8_HALF_DEVICE:
             r.cluster_size = 8
@@ -198,4 +138,6 @@ def pick_config(tier: str, batch: int, n_comp: int, k: int, num_sms: int) -> Top
             r.cluster_size = 4
         elif batch * 2 <= (num_sms * CS_HEADROOM_NUM) // CS_HEADROOM_DEN:
             r.cluster_size = 2
+    if r.cluster_size == 1 and k <= RUNGS_BM_MAX_K and n_comp >= SKIP_MIN_N_RUNGS_FLASH:
+        r.attach_block_max = True
     return r
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 45891a0f3a05..041706d6fab5 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -230,9 +230,8 @@ def _compile(
         kc_override=kc_override,
         self_scan=self_scan,
         cap_c=cap_c,
-        # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed. The
-        # qfrac VALUES are irrelevant on this path (P1b is skipped) — only
-        # the slot count matters.
+        # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed;
+        # the qfrac values are unused here — only the slot count matters.
         r0_qfracs=(0.85, 0.35) if (use_ext_counts or ext_rungs) else None,
     )
     return cute.compile(
@@ -289,15 +288,11 @@ def emu_block_max(
       "pad_inf": a partially-valid tail unit is forced to +FLT_MAX, the
                  worst legal inflation (the indexer masks by request-level
                  ctx >= N_eff, so tail positions can inflate the bound).
-                 Tests use this to prove inflated bounds only *disable*
-                 skipping, never break correctness.
     records:
       "rotate":     fold-correctness fixture — the 128-block max lands in
                     ONE slot rotated by blk % 4, the other 3 hold
-                    -FLT_MAX. Any 4 values folding to the block max are
-                    legal; the rotation makes a consumer that drops or
-                    mis-reads partial slots fail the tests. Valid ONLY
-                    for grain-128 consumers (slots are NOT positional).
+                    -FLT_MAX. Valid ONLY for grain-128 consumers (slots
+                    are NOT positional).
       "positional": production semantics — record r is the exact max of
                     positions [r*32, r*32+32) (the indexer's TMEM T2R
                     partition gives warp w of a tile the contiguous
@@ -338,9 +333,7 @@ def emu_block_max(
 
 
 # Rung offsets below each 32-position record's max for the meta-seed
-# metadata (L6_geo.25 — offline-validated: 1.00 real scans on synth +
-# real captures for both V4 models with the kernel's S=16 x 2-pass
-# quantized search).
+# metadata.
 _META_DELTAS = (0.25, 0.5, 1.0, 2.0, 4.0, 8.0)
 
 
@@ -359,8 +352,6 @@ def emu_block_meta(
     upper bound 32). Positions beyond the row's effective length are
     excluded. Layout matches ``block_max``: ``[num_rows, nrec]`` int32
     with ``nrec = ceil(N/128)*4`` (one record per 32 positions).
-    On the indexer side this is ~L ballots+popcounts per record in the
-    epilogue, same shape as the +1.3% block_max emission.
     """
     assert next_n == 1, "emu_block_meta: next_n == 1 only"
     x = logits.float()
@@ -394,8 +385,7 @@ def derive_seed_rungs(
     Estimates the per-row local slope of log2(count) vs threshold from the
     PREVIOUS step's 3 rung measurements and places the next step's guard
     rungs ``count_octaves`` octaves away from the mid rung (= the previous
-    accepted threshold). Real-chain validation (V4-Pro/Flash captures):
-    in-band admission 0.958/0.975 vs 0.82-0.97 for any fixed spread.
+    accepted threshold).
 
     Args:
         prev_thr: [rows] previous accepted threshold (xstate[:, 2]).
@@ -418,10 +408,8 @@ def derive_seed_rungs(
             count_octaves / slope.clamp(min=0.05),
             torch.full_like(slope, fallback_spread),
         ).clamp(0.1, 4.0)
-        # undershoot hysteresis: a row whose PREVIOUS loose rung caught
-        # fewer than K cannot use its list (fallback ~26us); widen only
-        # that row's next down-guard by +2 octaves. Rows in band keep the
-        # tight spread (a fat list taxes every step's walk).
+        # undershoot hysteresis: a row whose previous loose rung caught
+        # fewer than K widens only its next down-guard by +2 octaves.
         oct_lo = torch.full_like(slope, count_octaves)
         if top_k is not None:
             oct_lo = torch.where(prev_counts[:, 0].float() < float(top_k), oct_lo + 2.0, oct_lo)
@@ -770,17 +758,12 @@ def gvr_topk_decode(
     cute_dtype = _DTYPE_TORCH_TO_CUTE[logits.dtype]
 
     num_rows = logits.shape[0]
-    # Host dispatch gate: the compact walk only wins when the per-row
-    # compressed length is large (protocol: >= 2x at N=262k, 4-14% LOSS at
-    # N <= 131k). Below skip_min_n (compressed-index space, shape-based so
-    # no device sync) drop block_max and run the dense arms. None disables
-    # the gate (A/B probes).
+    # Host dispatch gate: below skip_min_n (compressed-index space,
+    # shape-based so no device sync) drop block_max; None disables the gate.
     if block_max is not None and skip_min_n is not None and logits.shape[1] < skip_min_n:
         block_max = None
-    # K > 512 at tiny batch: the acceptance band is proportionally tighter
-    # (kC/K = 6 vs 10), the bounds prune less, and the row-split configs
-    # win outright (cold protocol, pro 262k BS1: skip 21.3-21.6us at
-    # cs1/cs8 vs stock cs8 19.7us) -> keep the stock path.
+    # K > 512 at tiny batch: drop block_max (stock path) unless the admitted
+    # line is known up front (ext counts / packed seed / self_scan).
     if (
         block_max is not None
         and num_rows < 8
@@ -789,20 +772,12 @@ def gvr_topk_decode(
         and not (seed_thr is not None and seed_counts is not None)
         and not (seed_thr is not None and seed_thr.shape[1] >= 6)
     ):
-        # Stock-path small-batch gate, tuned in the loose-rung era (the
-        # sample-quantile list kept 60%+ of the blocks at K>512). It does
-        # NOT apply when the admitted line is known up front: ext counts
-        # park the tight line in every rung slot, so the list forms at
-        # the accepted threshold and skips by the real band density
-        # (pro-1M: 12.7% pass -> the list is live again). self_scan owns
-        # its own skip economics likewise.
         block_max = None
     if self_scan:
         # fused self-contained mode: kernel scans/buckets the row itself.
         # Inputs: seed_thr (three closed-loop lines) + a write-only gmem
-        # POSITION column passed through the cand_idx slot. seed_counts is
-        # a dummy (the ext-counts preview reads it but zeros never pass
-        # the [K, kC] admission, so routing is owned by the phase-0 gate).
+        # POSITION column passed through the cand_idx slot; seed_counts is
+        # a dummy (zeros never pass the [K, kC] admission).
         assert seed_thr is not None, "self_scan requires seed_thr"
         assert cand_vals is None and cand_ctl is None, (
             "self_scan excludes external candidate values/control"
@@ -906,17 +881,12 @@ def gvr_topk_decode(
     N_dec = max_seq_len if max_seq_len is not None else N_cols
     if num_threads_per_block is None:
         if use_ext_cand and top_k <= 512:
-            # list-hit rows do O(list) work (~2-4K entries at K<=512),
-            # not O(N): the N-keyed 1024-thread pick pays barrier-heavy
-            # phases for no parallel gain (flash-512k v5: 8.4 -> 7.1us
-            # at 512 threads; K=1024 lists are big enough to keep the
-            # N-keyed pick).
+            # list-hit rows do O(list) work, not O(N): use 512 threads
+            # (K=1024 lists are big enough to keep the N-keyed pick).
             num_threads_per_block = 512
         elif self_scan:
-            # self_scan owns the whole row scan in one CTA: the phase-0
-            # cp.async pipeline scales with warp count at every N (the
-            # 512-thread short-row heuristic below is tuned for the
-            # stock multi-pass kernel and costs ~5us/cell here).
+            # self_scan scans the whole row in one CTA; the phase-0
+            # cp.async pipeline scales with warp count at every N.
             num_threads_per_block = 1024
         else:
             if max_seq_len is not None and logits.dtype != torch.float32:

From 1ec6f2f640e208044c75005b7250772fec02322a Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 9 Aug 2026 16:27:40 -0700
Subject: [PATCH 100/117] [None][test] split the heavy CuTe DSL files out of
 the B200 attention entry

The bare directory entry ran the whole attention tree under the default
3600s wrapper timeout. The GVR decode file alone runs ~90 minutes, so
the wrapper died mid-file on DGX_B200 (both the wrapper item and the
in-flight case report 'Test terminated unexpectedly'). Mirror the B300
layout: --ignore the three CuTe DSL files in the directory entry and
list them separately, GVR decode with TIMEOUT (120).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tests/integration/test_lists/test-db/l0_b200.yml | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index b3a07d3b526b..5d9a1e0b8435 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_b200.yml
@@ -72,7 +72,10 @@ l0_b200:
   - test_e2e.py::test_ptp_quickstart_advanced_ngram[Llama-3.1-8B-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct]
   - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-False-False]
   - test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b]
-  - unittest/_torch/attention
+  - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+  - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py
+  - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+  - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120)
   - unittest/_torch/compilation
   - unittest/_torch/debugger
   - unittest/_torch/kv_cache_compression

From 06207209c26724a5da46571bd82c52e43061dbeb Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 11 Aug 2026 04:07:53 -0700
Subject: [PATCH 101/117] [None][fix] address GVR emission review findings

Review round 2026-08-11 (brnguyen2, longcheng-nv), all inline findings:

- list tier line-cut: the mapped-prefix copy re-measures the
  non-sentinel candidate count in flight and demotes to the stock path
  below K - a pad-inflated claim can no longer admit a starved list.
  Regression test ext_list[starved].
- update_seed_rows: slope-fit of log2(count) vs threshold from the
  previous step's lines+counts (harness derive_seed_lines_v4
  construction) at K-relative target counts; list tier two-point-fits
  through the published exact k-th; no-fit rows take multiplicative
  guards (~9% of kth) instead of the degenerate 1e-4 span. Chained
  test ext_closed_loop (3 steps, kth drift 0.05, all three tiers).
- rungs tier reachable: persistent contiguous [rows, 3] seed mirror
  passed for the rungs route; counts telemetry from the kernel's
  xstate[4..6] publish.
- one shared per-step gate (batch <= 256, next_n == 1) for planning,
  emission and consumption; batch > 256 pays no emission tax and a
  stale route cannot be consumed.
- prefill->decode handoff seeds prev_topk and zeroes xstate for new
  requests (same slot convention as heuristic_prev_topk); positional
  identity assumption documented at the consumption site.
- seed width assert == 3 or == 8 with matching use_ext_counts
  derivation; ext tiers require enable_r0; list-capacity ctor reject;
  trace-time flag/tensor contract errors; block-skip single-band void
  contract; cand_ctl width comments; SKIP_MAX_BLOCKS budget comment;
  radix_lens comment reconciled with metadata (declared latent-bug
  fix); mutates_args limitation re-verified on the pinned torch
  (IndexError when a declared-mutable Optional is None at call) and
  documented precisely in both op docstrings.
- rename per review naming feedback: TRTLLM_GVR_EXT ->
  TRTLLM_GVR_EMISSION, gvr_ext.py -> gvr_emission.py, GvrExtState ->
  GvrEmissionState.

Validation on B200: full decode unit file 730 passed / 0 failed / 144
skipped / 1 xfailed; new ext_list[starved] and 3-tier ext_closed_loop
chained tests pass; mutates_args declaration attempt reverted after
uniform stock-path failures (None-case IndexError).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../attention_backend/sparse/dsa/indexer.py   | 100 ++++++++++++------
 .../sparse/{gvr_ext.py => gvr_emission.py}    |  98 +++++++++++++----
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  28 ++++-
 .../blackwell/top_k/gvr_topk_decode.py        |  58 ++++++++--
 .../sparse/test_cute_dsl_gvr_topk_decode.py   |  91 +++++++++++++++-
 5 files changed, 306 insertions(+), 69 deletions(-)
 rename tensorrt_llm/_torch/attention_backend/sparse/{gvr_ext.py => gvr_emission.py} (64%)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index f3e0c32ca9e3..34c04ef3dcc2 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -684,14 +684,14 @@ def __init__(
             sparse_params.use_cute_dsl_paged_mqa_logits and IS_CUTLASS_DSL_AVAILABLE
         )
         # GVR emission-assisted decode (opt-in, experimental): the FP4 indexer
-        # epilogue emits candidates for the top-k (see gvr_ext / gvr_routing).
-        self.use_gvr_ext = (
-            os.environ.get("TRTLLM_GVR_EXT", "0") == "1"
+        # epilogue emits candidates for the top-k (see gvr_emission / gvr_routing).
+        self.use_gvr_emission = (
+            os.environ.get("TRTLLM_GVR_EMISSION", "0") == "1"
             and self.use_cute_dsl_topk
             and self.use_cute_dsl_paged_mqa_logits
             and self.use_fp4
         )
-        self._gvr_ext = None  # lazy GvrExtState (first decode step)
+        self._gvr_emission = None  # lazy GvrEmissionState (first decode step)
         self._gvr_route = None
         self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5
 
@@ -1342,6 +1342,24 @@ def _call_paged_mqa_logits(
             max_seq_len,
         )
 
+    def _ensure_gvr_emission(self, metadata: DSAtrtllmAttentionMetadata, device: torch.device):
+        """Lazy per-layer GVR ext state (see gvr_emission.GvrEmissionState)."""
+        if self._gvr_emission is None:
+            from ..gvr_emission import LIST_EMIT_MIN_N, GvrEmissionState
+
+            # get_indexer_max_seq_len already returns the compressed length
+            n_comp = metadata.get_indexer_max_seq_len()
+            self._gvr_emission = GvrEmissionState(
+                max_rows=metadata.max_num_sequences,
+                top_k=self.index_topk,
+                device=device,
+                # the list tier is only reachable at n_comp >= LIST_EMIT_MIN_N;
+                # skip its large candidate buffers when the engine's static
+                # length cannot get there
+                enable_list_tier=n_comp >= LIST_EMIT_MIN_N,
+            )
+        return self._gvr_emission
+
     def sparse_attn_indexer(
         self,
         metadata: DSAtrtllmAttentionMetadata,
@@ -1589,14 +1607,29 @@ def sparse_attn_indexer(
         # bottom of the decode block): new gens from finishing prefill
         # append after currently-active gens, i.e., slots
         # [num_generations : num_generations + num_contexts].
-        if self._enable_heuristic_topk and has_prefill and not metadata.skip_indexer_for_ctx_reqs:
-            local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx]
+        if (
+            has_prefill
+            and not metadata.skip_indexer_for_ctx_reqs
+            and (self._enable_heuristic_topk or self.use_gvr_emission)
+        ):
             ctx_seq_lens = metadata.seq_lens[:num_contexts]
             # Per-sequence last context-token offset (exclusive cumsum minus 1).
             last_ctx_idx = (torch.cumsum(ctx_seq_lens, dim=0) - 1).to(dtype=torch.long)
-            metadata.heuristic_prev_topk[
-                local_layer, num_generations : num_generations + num_contexts
-            ].copy_(topk_indices_buffer[last_ctx_idx, :])
+            last_ctx_topk = topk_indices_buffer[last_ctx_idx, :]
+            if self._enable_heuristic_topk:
+                local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx]
+                metadata.heuristic_prev_topk[
+                    local_layer, num_generations : num_generations + num_contexts
+                ].copy_(last_ctx_topk)
+            if self.use_gvr_emission:
+                # Ext state is positional like heuristic_prev_topk (no churn
+                # signal exists to key on): zeroed xstate cold-starts new gens
+                # onto the stock path; stale rows after churn only mis-place
+                # lines - counts are re-measured in-kernel, so exactness holds.
+                st = self._ensure_gvr_emission(metadata, topk_indices_buffer.device)
+                new_gens = slice(num_generations, num_generations + num_contexts)
+                st.prev_topk[new_gens].copy_(last_ctx_topk[:, : st.prev_topk.shape[1]])
+                st.xstate[new_gens].zero_()
 
         reuse_topk = (
             self.mtp_index_share
@@ -1610,6 +1643,7 @@ def sparse_attn_indexer(
                 metadata.shared_topk_indices[:num_generations, :]
             )
         elif has_decode and not metadata.skip_indexer_for_gen_reqs:
+            gvr_step_ok = False
             # Get decode lengths per request (from seq_lens) for validation
             gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
             max_decode_len = gen_seq_lens.max().item()
@@ -1713,32 +1747,29 @@ def sparse_attn_indexer(
                         dsl_schedule_meta = metadata.scheduler_metadata_buffer_expanded
 
                     gvr_emit_kwargs = {}
-                    if self.use_gvr_ext and next_n == 1 and not dsl_atom_split:
-                        from ..gvr_ext import LIST_EMIT_MIN_N, GvrExtState
-
+                    # one step gate shared with the consume branch below:
+                    # emitting for a step the top-k cannot consume only
+                    # churns the closed-loop state
+                    gvr_step_ok = (
+                        self.use_gvr_emission
+                        and use_custom_topk
+                        and next_n == 1
+                        and not dsl_atom_split
+                        and num_gen_tokens <= 256
+                    )
+                    if gvr_step_ok:
+                        st = self._ensure_gvr_emission(metadata, q_fp8.device)
                         # indexer_max_seq_len is already the compressed length
                         # (get_indexer_max_seq_len divides); do not divide again.
-                        n_comp = indexer_max_seq_len
-                        if self._gvr_ext is None:
-                            self._gvr_ext = GvrExtState(
-                                max_rows=metadata.max_num_sequences,
-                                top_k=self.index_topk,
-                                device=q_fp8.device,
-                                # the list tier is only reachable at
-                                # n_comp >= LIST_EMIT_MIN_N; skip its large
-                                # candidate buffers when the engine's static
-                                # length cannot get there
-                                enable_list_tier=n_comp >= LIST_EMIT_MIN_N,
-                            )
-                        st = self._gvr_ext
                         emit_tier, self._gvr_route = st.plan(
                             batch_size,
-                            n_comp,
+                            indexer_max_seq_len,
                             torch.cuda.get_device_properties(q_fp8.device).multi_processor_count,
                             compress_ratio=max(self.compress_ratio, 1),
                         )
-                        if emit_tier in ("counts", "list"):
+                        if emit_tier in ("counts", "list", "rungs"):
                             st.update_seed_rows(batch_size, emit_tier)
+                        if emit_tier in ("counts", "list"):
                             gvr_emit_kwargs = st.indexer_emit_kwargs(emit_tier, batch_size)
                         if self._gvr_route.attach_block_max or emit_tier in (
                             "counts",
@@ -1829,16 +1860,15 @@ def sparse_attn_indexer(
                 # tier "none" (see gvr_routing) falls through to the stock
                 # branch; the untouched ext state reads as a cold start there.
                 if (
-                    self.use_gvr_ext
-                    and self._gvr_ext is not None
+                    self.use_gvr_emission
+                    and gvr_step_ok
+                    and self._gvr_emission is not None
                     and self._gvr_route is not None
                     and self._gvr_route.tier != "none"
-                    and next_n == 1
-                    and num_gen_tokens <= 256
                 ):
                     # emission-assisted GVR: consume what the indexer epilogue
                     # emitted this step (seed row / list / block_max per route)
-                    st = self._gvr_ext
+                    st = self._gvr_emission
                     out_slice = topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :]
                     # GVR op takes RAW-domain seq_lens (kernel ceil-divides by
                     # compress_ratio); on the DSL path the live compressed lens
@@ -1879,8 +1909,10 @@ def sparse_attn_indexer(
                         order_row=metadata.kv_lens_row_reorder,
                     )
                 elif self.use_cute_dsl_topk and (self.compress_ratio == 1 or next_n == 1):
-                    # seq_lens must be 1-D; on the DSL path the live compressed
-                    # lens are gen_indexer_kv_lens_cuda_runtime (2d buf is zero)
+                    # latent-bug fix: the op takes 1-D request-level lens, but
+                    # context_lens is the 2-D kv_lens_cuda_2d slice; on the
+                    # FP4-DSL path the live compressed lens are
+                    # gen_indexer_kv_lens_cuda_runtime
                     if self.compress_ratio > 1:
                         radix_lens = (
                             metadata.gen_indexer_kv_lens_cuda_runtime
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
similarity index 64%
rename from tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
rename to tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
index 034b2d3a212f..770ce37647dd 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_ext.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
@@ -27,6 +27,7 @@
     step.
 """
 
+import math
 from typing import Optional
 
 import torch
@@ -44,10 +45,21 @@
 LIST_CAP_C = 24576
 LIST_WIDTH = 2 * LIST_SEG_A + LIST_CAP_C
 
-# Closed-loop lines around the published k-th anchor: t1 hugs the k-th
-# value from below; t0/t2 guard by the (anchor - kth) span.
-__all__ = ["GvrExtState", "LIST_EMIT_MIN_N", "LIST_PARK_LINE"]
-
+__all__ = ["GvrEmissionState", "LIST_EMIT_MIN_N", "LIST_PARK_LINE"]
+
+# Closed-loop line placement: fit the slope of log2(count) vs threshold
+# from the previous step's (lines, counts) and place the new lines at
+# these K-relative target counts (t0 loosest .. t2 tightest).
+LINE_TARGETS = (8.0, 5.0, 2.0)
+LIST_T0_TARGET = 2.5  # list tier: single collect-line target (xK)
+LIST_T0_COUNT_MAX = 6144.0  # keep n0 inside the [K+64, segA] admission band
+SLOPE_MIN = 0.05
+SLOPE_MAX = 64.0
+
+# No-fit fallback: multiplicative guards around the published k-th value
+# (t1 hugs it from below; t0/t2 guard by GUARD_LO/GUARD_HI spans).
+FALLBACK_REL = 2.0**0.125 - 1.0
+FALLBACK_ABS = 1e-3
 GUARD_LO = 2.0
 GUARD_HI = 0.5
 
@@ -58,7 +70,7 @@
 LIST_PARK_LINE = 1.0e30
 
 
-class GvrExtState:
+class GvrEmissionState:
     """Per-attention-backend emission state (persistent buffers)."""
 
     def __init__(
@@ -69,6 +81,9 @@ def __init__(
         # packed seed row: lines at cols 0..2, counts (emission-filled)
         # at 3..5, adaptive-skip pass count at 6
         self.seed_row = torch.zeros((max_rows, 8), dtype=torch.float32, device=device)
+        # contiguous alias of the three lines for the rungs tier (a
+        # [rows, 3] column view of the packed row is non-contiguous)
+        self.seed_rungs = torch.zeros((max_rows, 3), dtype=torch.float32, device=device)
         self.xstate = torch.zeros((max_rows, 8), dtype=torch.float32, device=device)
         self.cand_vals: Optional[torch.Tensor] = None
         self.cand_idx: Optional[torch.Tensor] = None
@@ -95,7 +110,7 @@ def ensure_block_max(self, max_seq_len: int) -> torch.Tensor:
             # address into the graph, so fail loudly instead
             if torch.cuda.is_current_stream_capturing():
                 raise RuntimeError(
-                    "GvrExtState.ensure_block_max: (re)allocation requested "
+                    "GvrEmissionState.ensure_block_max: (re)allocation requested "
                     "during CUDA graph capture; the block_max buffer must be "
                     "created by a warmup step before capture"
                 )
@@ -124,26 +139,70 @@ def plan(
     def update_seed_rows(self, num_rows: int, emit_tier: str = "counts") -> None:
         """Device-side closed-loop line update from the last publish.
 
-        Pure tensor ops (graph-capturable). Rows with invalid xstate
-        (col 0 == 0, e.g. cold start) get non-finite lines, which the
-        kernel's validity guard routes to the stock path.
+        Slope-fits log2(count) vs threshold from the previous step's
+        (lines, counts) and places the new lines at K-relative target
+        counts. Counts come from the packed row (counts/list emission)
+        or from the kernel's rung-count publish in xstate cols 4..6
+        (rungs tier). Rows without a usable fit get multiplicative
+        guards around the published k-th value; rows with invalid
+        xstate (col 0 == 0, e.g. cold start) get non-finite lines,
+        which the kernel's validity guard routes to the stock path.
+        Pure tensor ops (graph-capturable).
         """
         s = self.seed_row[:num_rows]
         x = self.xstate[:num_rows]
-        kth = x[:, 1]
-        anch = torch.maximum(x[:, 2], kth + 1e-5)
-        span = (anch - kth).clamp_min(1e-4)
         valid = x[:, 0] > 0
+        kth = x[:, 1]
+        anchor = x[:, 2]
+        t_prev0 = s[:, 0]
+        t_prev2 = s[:, 2]
+        cnts = x[:, 4:7] if emit_tier == "rungs" else s[:, 3:6]
+        k = float(self.top_k)
         inf = torch.full_like(kth, float("inf"))
-        s[:, 0] = torch.where(valid, kth - GUARD_LO * span, inf)
+        d_fb = kth.abs() * FALLBACK_REL + FALLBACK_ABS
         if emit_tier == "list":
+            # two-point fit (t0_prev, n0) / (kth, K): kth is the exact
+            # k-th boundary on list rows
+            n0 = cnts[:, 0].clamp_min(1.0)
+            dthr = (kth - t_prev0).clamp_min(1e-3)
+            slope = ((torch.log2(n0) - math.log2(k)) / dthr).clamp(SLOPE_MIN, SLOPE_MAX)
+            tgt0 = min(LIST_T0_TARGET * k, LIST_T0_COUNT_MAX)
+            t0 = kth - math.log2(tgt0 / k) / slope
+            fit_ok = torch.isfinite(t_prev0) & (n0 > k)
+            t0 = torch.where(fit_ok, t0, kth - GUARD_LO * d_fb)
             park = torch.full_like(kth, LIST_PARK_LINE)
-            s[:, 1] = torch.where(valid, park, inf)
-            s[:, 2] = torch.where(valid, park + park, inf)
+            new0 = torch.where(valid, t0, inf)
+            new1 = torch.where(valid, park, inf)
+            new2 = torch.where(valid, park + park, inf)
         else:
-            s[:, 1] = torch.where(valid, kth - 1e-6, inf)
-            s[:, 2] = torch.where(valid, kth + GUARD_HI * span, inf)
+            c0 = cnts[:, 0].clamp_min(1.0)
+            c2 = cnts[:, 2].clamp_min(1.0)
+            dthr = (t_prev2 - t_prev0).clamp_min(1e-3)
+            slope = ((torch.log2(c0) - torch.log2(c2)) / dthr).clamp(SLOPE_MIN, SLOPE_MAX)
+            # anchor count estimate: slide the anchor onto the prev line fit
+            anch_c = (c2 * torch.exp2(-(anchor - t_prev2) * slope)).clamp(1.0, 1e6)
+            t0 = anchor + torch.log2(anch_c / (LINE_TARGETS[0] * k)) / slope
+            t1 = anchor + torch.log2(anch_c / (LINE_TARGETS[1] * k)) / slope
+            t2 = anchor + torch.log2(anch_c / (LINE_TARGETS[2] * k)) / slope
+            # t_prev2 < 1e29 also rejects a parked line left by a tier flip
+            fit_ok = torch.isfinite(t_prev0) & (t_prev2 < 1e29) & (c0 > c2)
+            t0 = torch.where(fit_ok, t0, kth - GUARD_LO * d_fb)
+            t1 = torch.where(fit_ok, t1, kth - 1e-6)
+            t2 = torch.where(fit_ok, t2, kth + GUARD_HI * d_fb)
+            # strictly ascending (kernel line-validity contract)
+            t1 = torch.maximum(t1, t0 + 1e-4)
+            t2 = torch.maximum(t2, t1 + 1e-4)
+            new0 = torch.where(valid, t0, inf)
+            new1 = torch.where(valid, t1, inf)
+            new2 = torch.where(valid, t2, inf)
+        s[:, 0] = new0
+        s[:, 1] = new1
+        s[:, 2] = new2
         s[:, 3:8] = 0.0
+        rungs = self.seed_rungs[:num_rows]
+        rungs[:, 0] = new0
+        rungs[:, 1] = new1
+        rungs[:, 2] = new2
         if self.cand_ctl is not None:
             self.cand_ctl[:num_rows].zero_()
             self.cand_cur[:num_rows].zero_()
@@ -177,8 +236,9 @@ def topk_ext_kwargs(
             kw["num_threads"] = route.num_threads
         if route.tier in ("counts", "list"):
             kw["seed_thr"] = self.seed_row[:num_rows]
-        # rungs tier: pass no seed at all (a [rows, 3] column view of the
-        # packed row is non-contiguous and would trip the runner's assert)
+        elif route.tier == "rungs":
+            # [rows, 3] seed selects the op's ext_rungs variant
+            kw["seed_thr"] = self.seed_rungs[:num_rows]
         if route.tier == "list":
             # accept_cap must match the emitter's segment geometry: the
             # buffers are laid out at bases 0 / LIST_SEG_A / 2*LIST_SEG_A,
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 0ca818b3b6f5..6067a39966dd 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -7745,11 +7745,11 @@ def forward(
             # emission-assisted tiers: mode from which ext tensors the
             # caller handed in (see gvr_routing.plan_emission)
             if seed_thr is not None:
-                assert seed_thr.shape[1] == 3 or seed_thr.shape[1] >= 6, (
-                    "seed_thr width must be 3 (rungs) or >= 6 (packed "
-                    f"counts row); got {seed_thr.shape[1]} - refusing to "
-                    "silently ignore the seed")
-            use_ext_counts = seed_thr is not None and seed_thr.shape[1] >= 6
+                assert seed_thr.shape[1] == 3 or seed_thr.shape[1] == 8, (
+                    "seed_thr must be [rows, 3] (rungs) or [rows, 8] "
+                    "(packed lines + counts row, the width the kernel "
+                    f"is compiled for); got width {seed_thr.shape[1]}")
+            use_ext_counts = seed_thr is not None and seed_thr.shape[1] == 8
             ext_rungs = seed_thr is not None and seed_thr.shape[1] == 3
             use_ext_cand = cand_vals is not None
             enable_block_skip = block_max is not None
@@ -7864,6 +7864,16 @@ def cute_dsl_gvr_topk_decode(
             max_batch_size: Required with ``counters``; ignored
                 otherwise. Power of 2 in ``[64, 1024]``, must match
                 the value passed to LB prepare.
+
+        Of the optional hint tensors the kernel WRITES ``xstate`` (the
+        closed-loop publish); it cannot be declared in ``mutates_args``:
+        torch.library raises IndexError when a declared-mutable Optional
+        arg is None at call time (re-verified on the pinned torch), and
+        most calls pass no hints. Under torch.compile/functionalization
+        the undeclared write is invisible, so the hint path is eager /
+        CUDA-graph only. ``TRTLLM_GVR_EMISSION=1`` gates the
+        emission-assisted wiring that feeds these tensors (opt-in,
+        experimental).
         """
         if not is_sm_100f():
             raise ValueError(
@@ -9170,6 +9180,14 @@ def forward(
             Returns:
                 logits [B*next_n, max_context_len]; with emit_block_meta,
                 the tuple (logits, block_max, hit_stats).
+
+            The optional emission tensors (``block_max_out`` /
+            ``seed_thr`` / ``cand_*``) are written by the kernel but
+            cannot be declared in ``mutates_args``: torch.library raises
+            IndexError when a declared-mutable Optional arg is None at
+            call time (re-verified on the pinned torch), and plain
+            logits-only calls pass none of them. Emission is therefore
+            eager / CUDA-graph only.
             """
             B, next_n, H, half_D = q.shape
             N = next_n * H
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 042d716ae8c8..e99116be9677 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -495,7 +495,9 @@ def __init__(
         self.enable_block_skip = bool(enable_block_skip)
         self.SKIP_BLOCK = 32
         self.SKIP_BLOCK_LOG2 = 5
-        self.SKIP_MAX_BLOCKS = 8192  # covers N up to 262144 at grain 32
+        self.SKIP_MAX_BLOCKS = 8192  # smem active-list budget (32KB of local
+        # ids); bounds block-skip to N_local <= 262144 at grain 32 - longer
+        # rows run the dense fallback and owe nothing to block-skip
         self.SKIP_UNROLL = 2
         self.skip_order = "grouped"
         if enable_block_skip and num_threads not in (512, 1024):
@@ -692,6 +694,14 @@ def __init__(
             self.cap_c = 0
         if use_ext_cand and not use_ext_counts:
             raise ValueError("use_ext_cand requires use_ext_counts")
+        if (use_ext_counts or ext_rungs or use_ext_cand) and not enable_r0:
+            # the effective flags below are and-ed with enable_r0; reject
+            # instead of silently compiling the stock path
+            raise ValueError("ext tiers require enable_r0")
+        if use_ext_cand and self.list_cap <= 0:
+            raise ValueError(
+                f"cand_cap={cand_cap} leaves no C segment past 2*accept_cap={2 * self.accept_cap}"
+            )
         # ext_rungs (two-pass variant B): closed-loop rung THRESHOLDS come
         # from the host (previous-step xstep lines); the kernel counts them
         # itself via the stock R0 multi-count and admits the tightest rung
@@ -1793,7 +1803,11 @@ def phase0_scan_bucket(
                 # or the fallback (under) - exactly the v5 state machine
                 # fed with n0 == n1 == n2
                 curT0 = s_seg[0]
+                # claims past segA were dropped by the walk: honor the
+                # "void==0 means nothing was dropped" contract
                 s_seg[3] = cutlass.Int32(0)
+                if curT0 > cutlass.Int32(segA):
+                    s_seg[3] = cutlass.Int32(1)
                 s_seg[4] = curT0
                 s_seg[5] = curT0
                 s_seg[6] = curT0
@@ -5399,7 +5413,7 @@ def gvr_topk_kernel(
         xstate: cute.Tensor,  # [numRows, 8] fp32 closed-loop state (or None)
         cand_vals: cute.Tensor,  # [numRows, CAP] fp32 scores (or None)
         cand_idx: cute.Tensor,  # [numRows, CAP] int32 positions (or None)
-        cand_ctl: cute.Tensor,  # [numRows, 2] int32 {claimed, void} (or None)
+        cand_ctl: cute.Tensor,  # [numRows, 4] int32 {n0, void, n1, n2} (or None)
     ):
         """Thin entry: bidx → row_idx → run_one_row.
 
@@ -5475,7 +5489,7 @@ def run_one_row(
         xstate: cute.Tensor = None,  # [numRows, 8] fp32 (emit_xstate)
         cand_vals: cute.Tensor = None,  # [numRows, CAP] fp32 (ext cand)
         cand_idx: cute.Tensor = None,  # [numRows, CAP] int32 (ext cand)
-        cand_ctl: cute.Tensor = None,  # [numRows, 2] int32 (ext cand)
+        cand_ctl: cute.Tensor = None,  # [numRows, 4] int32 (ext cand)
     ):
         """Dispatch: compute per-row slice + cluster sync mode, call _run_phases.
 
@@ -5532,6 +5546,18 @@ def run_one_row(
         # Slice per-row views.
         input_row = input_data[row_idx, None]
         pre_idx_row = pre_idx[pre_idx_row_idx, None]
+        # trace-time contract checks: a feature flag without its tensor
+        # must fail HERE, not as a NoneType subscript deep in the phases
+        if cutlass.const_expr((self.use_ext_counts or self.ext_rungs) and seed_thr is None):
+            raise ValueError("use_ext_counts/ext_rungs kernels require seed_thr")
+        if cutlass.const_expr(self.emit_xstate and xstate is None):
+            raise ValueError("emit_xstate kernels require xstate")
+        if cutlass.const_expr(
+            self.use_ext_cand and (cand_vals is None or cand_idx is None or cand_ctl is None)
+        ):
+            raise ValueError("use_ext_cand kernels require cand_vals/cand_idx/cand_ctl")
+        if cutlass.const_expr(self.enable_block_skip and block_max is None):
+            raise ValueError("enable_block_skip kernels require block_max")
         if cutlass.const_expr(self.enable_block_skip and block_max is not None):
             block_max_row = block_max[row_idx, None]
         else:
@@ -6635,6 +6661,7 @@ def _run_phases(
                     list_used = cutlass.Int32(1)
                     if tidx == cutlass.Int32(0):
                         s_iscalars[0] = cutlass.Int32(0)
+                        s_iscalars[2] = cutlass.Int32(0)  # non-sentinel count
                         s_thr[0] = anch_t
                         s_iscalars[1] = cutlass.Int32(1)  # done
                     cute.arch.barrier()
@@ -6654,6 +6681,7 @@ def _run_phases(
                         # no atomics - every read is a winner candidate.
                         if tidx == cutlass.Int32(0):
                             s_iscalars[0] = cut_n
+                        nreal_c = cutlass.Int32(0)
                         i_c = tidx
                         while i_c < cut_n:
                             for _ju in cutlass.range_constexpr(4):
@@ -6683,14 +6711,32 @@ def _run_phases(
                                         cute.AddressSpace.gmem,
                                         assumed_align=4,
                                     )
-                                    smem_vals[j_c] = cute.make_tensor(ip_c, cute.make_layout((1,)))[
-                                        0
-                                    ]
+                                    iv_c = cute.make_tensor(ip_c, cute.make_layout((1,)))[0]
+                                    smem_vals[j_c] = iv_c
+                                    # sentinel pads carry idx -1; cut_n
+                                    # (= claimed n0) counts them, so the
+                                    # REAL candidate count must be
+                                    # re-measured during the copy
+                                    if iv_c >= cutlass.Int32(0):
+                                        nreal_c = nreal_c + cutlass.Int32(1)
                             i_c = i_c + cutlass.Int32(4 * num_threads)
+                        wsum_c = self.warp_reduce_sum_i32(nreal_c)
+                        if lane_c == cutlass.Int32(0):
+                            atomicAdd(s_iscalars.iterator + cutlass.Int32(2), wsum_c)
                         wmax_w = self.warp_reduce_max_f32(wmax_acc)
                         if lane_c == cutlass.Int32(0):
                             smem_wcnt[tidx // cutlass.Int32(32)] = float_as_uint32(wmax_w)
                         cute.arch.barrier()
+                        # demote when the pad-inflated claim admitted a
+                        # list that holds fewer than K real candidates;
+                        # the stock path below recovers exactly
+                        real_c = s_iscalars[2]
+                        if real_c < cutlass.Int32(self.top_k):
+                            take_cand = cutlass.Int32(0)
+                            list_used = cutlass.Int32(0)
+                            if tidx == cutlass.Int32(0):
+                                s_iscalars[1] = cutlass.Int32(0)
+                            cute.arch.barrier()
                     if line_cut == cutlass.Int32(0):
                         # ---- histogram-edge cut: value-filtered mapped
                         # walk with merged-ballot claims (float edges
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
index 95ccff38ee87..5c557c5a032a 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
@@ -1283,7 +1283,7 @@ def test_cute_dsl_gvr_topk_decode_tie_flood_beyond_capacity():
 # ---------------------------------------------------------------------------
 # Emission-assisted (ext) tiers: packed seed row / candidate list / block max.
 # Inputs emulate the indexer epilogue host-side against the layout contracts
-# in ``gvr_ext`` (segments at bases 0 / LIST_SEG_A / 2*LIST_SEG_A, packed row
+# in ``gvr_emission`` (segments at bases 0 / LIST_SEG_A / 2*LIST_SEG_A, packed row
 # = lines at [0..2] + exact counts at [3..5] + skip pass count at [6]).
 # ---------------------------------------------------------------------------
 
@@ -1365,7 +1365,7 @@ def test_cute_dsl_gvr_topk_decode_ext_counts(top_k, mode):
 
 
 @skip_not_sm100
-@pytest.mark.parametrize("mode", ["hit", "pads", "hist", "void", "bucketed"])
+@pytest.mark.parametrize("mode", ["hit", "pads", "hist", "void", "bucketed", "starved"])
 def test_cute_dsl_gvr_topk_decode_ext_list(mode):
     """Candidate-list tier at the production geometry (accept_cap =
     LIST_SEG_A, width = LIST_WIDTH).
@@ -1379,8 +1379,11 @@ def test_cute_dsl_gvr_topk_decode_ext_list(mode):
     void:     collection overflows LIST_CAP_C -> void=1 -> full scan.
     bucketed: three live lines spread across segments A/B/C, cut at the
               tightest line inside the band.
+    starved:  fewer than K real candidates, but sentinel pads lift the
+              claim into the admission band -> the line-cut copy must
+              re-measure and demote (exactness regression).
     """
-    from tensorrt_llm._torch.attention_backend.sparse.gvr_ext import (
+    from tensorrt_llm._torch.attention_backend.sparse.gvr_emission import (
         LIST_CAP_C,
         LIST_PARK_LINE,
         LIST_SEG_A,
@@ -1396,7 +1399,7 @@ def test_cute_dsl_gvr_topk_decode_ext_list(mode):
     if mode == "bucketed":
         lines = _lines_at_counts(logits, n_eff, (20000, 4000, 600))
     else:
-        n0 = {"hit": 4096, "pads": 4096, "hist": 12000, "void": 30000}[mode]
+        n0 = {"hit": 4096, "pads": 4096, "hist": 12000, "void": 30000, "starved": 400}[mode]
         l0 = _lines_at_counts(logits, n_eff, (n0,))
         lines = torch.cat(
             [l0, torch.full_like(l0, LIST_PARK_LINE), torch.full_like(l0, 2 * LIST_PARK_LINE)], 1
@@ -1405,7 +1408,7 @@ def test_cute_dsl_gvr_topk_decode_ext_list(mode):
     cand_vals = torch.full((batch, LIST_WIDTH), float("-inf"), dtype=torch.float32, device=dev)
     cand_idx = torch.full((batch, LIST_WIDTH), -1, dtype=torch.int32, device=dev)
     cand_ctl = torch.zeros((batch, 4), dtype=torch.int32, device=dev)
-    pads = 64 if mode == "pads" else 0
+    pads = {"pads": 64, "starved": 200}.get(mode, 0)
     for r in range(batch):
         ne = int(n_eff[r])
         row = logits[r, :ne]
@@ -1501,3 +1504,81 @@ def test_cute_dsl_gvr_topk_decode_ext_block_max(tail_mode):
     )
     torch.cuda.synchronize()
     _tie_aware_check(out_indices, logits, seq_lens, top_k, 1)
+
+
+def _emulate_emission(logits, n_eff, st, tier, top_k):
+    """Host-side stand-in for the indexer epilogue: fill the packed-row
+    counts (and the candidate list on the list tier) against the CURRENT
+    seed lines, exactly as the production emitter would."""
+    from tensorrt_llm._torch.attention_backend.sparse.gvr_emission import LIST_CAP_C, LIST_SEG_A
+
+    batch, N = logits.shape
+    dev = logits.device
+    lines = st.seed_row[:batch, 0:3]
+    pos = torch.arange(N, device=dev)[None, :]
+    valid = pos < n_eff[:, None]
+    finite = torch.isfinite(lines[:, 0])
+    counts = torch.stack(
+        [((logits >= lines[:, j : j + 1]) & valid).sum(-1) for j in range(3)], 1
+    ).float()
+    st.seed_row[:batch, 3:6] = torch.where(finite[:, None], counts, torch.zeros_like(counts))
+    if tier == "list" and st.cand_vals is not None:
+        st.cand_vals[:batch].fill_(float("-inf"))
+        st.cand_idx[:batch].fill_(-1)
+        st.cand_ctl[:batch].zero_()
+        for r in range(batch):
+            if not bool(finite[r]):
+                continue
+            ne = int(n_eff[r])
+            row = logits[r, :ne]
+            hits = torch.nonzero(row >= lines[r, 0], as_tuple=False).flatten()
+            cnt = int(hits.numel())
+            nwr = min(cnt, LIST_CAP_C)
+            base = 2 * LIST_SEG_A
+            st.cand_idx[r, base : base + nwr] = hits[:nwr].int()
+            st.cand_vals[r, base : base + nwr] = row[hits[:nwr]]
+            st.cand_ctl[r, 0] = cnt
+            st.cand_ctl[r, 1] = 1 if cnt > LIST_CAP_C else 0
+
+
+@skip_not_sm100
+@pytest.mark.parametrize(
+    "tier_shape", [("list", 2, 131072), ("counts", 8, 131072), ("rungs", 1, 32768)]
+)
+def test_cute_dsl_gvr_topk_decode_ext_closed_loop(tier_shape):
+    """Chained multi-step closed loop: kernel xstate publish ->
+    update_seed_rows -> next step's emission and admission.
+
+    Step 2's logits shift the k-th value far beyond any fixed guard
+    width, so line placement must come from the fitted slope; every step
+    must stay exact regardless of which internal path admission picks.
+    """
+    from tensorrt_llm._torch.attention_backend.sparse.gvr_emission import GvrEmissionState
+
+    want_tier, batch, N = tier_shape
+    top_k = 512
+    num_sms = torch.cuda.get_device_properties(0).multi_processor_count
+    logits0, pre_idx, seq_lens = _make_inputs(
+        batch, N, top_k, torch.float32, 1, seed=17, varlen=False
+    )
+    dev = logits0.device
+    n_eff = seq_lens.to(device=dev, dtype=torch.long)
+    st = GvrEmissionState(max_rows=batch, top_k=top_k, device=dev, enable_list_tier=True)
+    # k-th drift per step: step1 -> step2 rises by ~0.05 (>> any fixed
+    # guard), step3 falls back near step1's level
+    shifts = (0.0, 0.05, -0.03)
+    for step, shift in enumerate(shifts):
+        logits = (logits0 + shift).contiguous()
+        tier, route = st.plan(batch, N, num_sms, compress_ratio=1)
+        assert tier == want_tier
+        st.update_seed_rows(batch, tier)
+        _emulate_emission(logits, n_eff, st, tier, top_k)
+        pre = torch.topk(logits.float(), top_k, dim=-1).indices.int().contiguous()
+        out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda")
+        kw = st.topk_ext_kwargs(route, batch, None)
+        torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+            logits, pre, seq_lens, out_indices, top_k=top_k, **kw
+        )
+        torch.cuda.synchronize()
+        _tie_aware_check(out_indices, logits, seq_lens, top_k, 1)
+        assert bool((st.xstate[:batch, 0] > 0).all().item()), f"step {step}: publish missing"

From c48cacf424789fac2e43cb222f626f693fbfcc48 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 11 Aug 2026 04:10:43 -0700
Subject: [PATCH 102/117] [None][chore] commit the GVR per-step perf grid
 driver (provenance)

Per review: the driver behind the PR's performance numbers, in-tree.
Per-step paired protocol (one NVTX range per arm/model/isl/B/layer,
cold reps cycle all decode steps), nsys wrapper, per-step extraction
from the sqlite, pairing + production-routing aggregation into the
B x N mean/min tables. Capture root and output dir parameterized via
GVR_CAP_ROOT / GVR_BENCH_OUT.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../cute_dsl_kernels/top_k/perf/README.md     |  19 +
 .../cute_dsl_kernels/top_k/perf/ab_steps.py   | 403 ++++++++++++++++++
 .../cute_dsl_kernels/top_k/perf/f58_an.py     |  93 ++++
 .../cute_dsl_kernels/top_k/perf/f58_grid.sh   |  29 ++
 .../cute_dsl_kernels/top_k/perf/perstep.py    |  75 ++++
 5 files changed, 619 insertions(+)
 create mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/README.md
 create mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
 create mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
 create mode 100755 tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
 create mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py

diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/README.md b/tests/scripts/cute_dsl_kernels/top_k/perf/README.md
new file mode 100644
index 000000000000..02b6fcc147bc
--- /dev/null
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/README.md
@@ -0,0 +1,19 @@
+# GVR top-k per-step perf grid (provenance for the PR numbers)
+
+Per-step paired protocol: each (arm, model, ISL, B, layer) is one NVTX
+range; the cold reps inside cycle every usable decode step (L2 evicted,
+batch refilled before eviction), so the k-th kernel instance of a range
+IS decode step k and arms pair index-by-index.
+
+- `ab_steps.py` - the driver (env: ARM = pr|st|va|vb|wf, UNITS =
+  "pro:64k,...", BS_LIST, OUT; GVR_CAP_ROOT points at the capture data,
+  HARNESS_ROOT at the repo).
+- `f58_grid.sh <gpu> "<units>"` - nsys wrapper per arm x unit
+  (GVR_BENCH_OUT selects the output dir).
+- `perstep.py` - per-step extraction from the nsys sqlite.
+- `f58_an.py` - pairing + production routing (plan_emission) + the B x N
+  mean/min tables. `f58_regr.py`-style regression stats derive from the
+  same pickle.
+
+Requirements: B200 (SM100), cutlass-dsl 4.5.x on PYTHONPATH, exclusive
+GPU (no concurrent process - it pollutes event timing), nsys >= 2025.3.
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py b/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
new file mode 100644
index 000000000000..ba1c1ec5f9c7
--- /dev/null
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
@@ -0,0 +1,403 @@
+#!/usr/bin/env python3
+# ruff: noqa
+# Measurement harness committed verbatim for provenance; bench idioms
+# (loop-scoped buffers, del/rebind) trip static analysis.
+# Layer x step complete grid. One NVTX range per (arm, model, isl, B,
+# layer); the COLD reps cycle through ALL usable decode steps (batch
+# refilled from the step's row BEFORE the eviction, so the row itself
+# stays cold) -> the per-range cold mean IS the all-steps mean for that
+# layer. Tables aggregate mean-over-layers; per-layer means are the
+# per-range values (saved separately).
+#   env: ARM, UNITS ("flash:4k,pro:512k,..."), OUT, BS_LIST
+import json
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+CAP = Path(
+    os.environ.get(
+        "GVR_CAP_ROOT",
+        "/home/scratch.loncheng_gpu/workspace/perf/workloads/DSV4/"
+        "E2E_exp/indexer_decode_capture/data",
+    )
+)
+LAYERS = {"flash": list(range(2, 43, 2)), "pro": list(range(2, 61, 2)), "v32": list(range(61))}
+KCFG = {"flash": (512, 4), "pro": (1024, 4), "v32": (2048, 1)}
+ALIGN = 64
+FMIN = torch.finfo(torch.float32).min
+
+ARM = os.environ["ARM"]
+OUT = os.environ["OUT"]
+_LSUB = os.environ.get("LAYER_SUBSET")
+if _LSUB:
+    _ls = [int(x) for x in _LSUB.split(",")]
+    LAYERS = {k: [x for x in v if x in _ls] for k, v in LAYERS.items()}
+UNITS = [u.split(":") for u in os.environ["UNITS"].split(",")]
+BS_LIST = [
+    int(x) for x in os.environ.get("BS_LIST", "1,2,4,8,16,32,64,128,256,512,1024").split(",")
+]
+_fb = os.environ.get("FORCE_BM")
+FORCE_BM = None if _fb is None else int(_fb)
+FORCE_CS = int(os.environ.get("FORCE_CS", "0"))
+CONLY = os.environ.get("CONLY", "0") == "1"  # codespell:ignore
+CONLY_RANK = int(os.environ.get("CONLY_RANK", "0"))
+_pw = os.environ.get("FORCE_PWR")
+FORCE_PWR = None if _pw is None else (_pw == "1")
+_tf = os.environ.get("FORCE_TAILFAST")
+FORCE_TAILFAST = None if not _tf else (_tf == "1")
+_nf = os.environ.get("FORCE_NOFINE")
+FORCE_NOFINE = None if not _nf else (_nf == "1")
+_nb = os.environ.get("FORCE_BINS")
+FORCE_BINS = int(_nb) if _nb else None
+_frt = os.environ.get("FORCE_FINERT")
+FORCE_FINERT = None if not _frt else (_frt == "1")
+_srt = os.environ.get("FORCE_SCATRT")
+FORCE_SCATRT = None if not _srt else (_srt == "1")
+_tm = os.environ.get("TIGHT_MULT")
+TIGHT_MULT = float(_tm) if _tm else None
+_nt = os.environ.get("FORCE_THREADS")
+FORCE_THREADS = int(_nt) if _nt else None
+_v3 = os.environ.get("FORCE_TAILV3")
+FORCE_TAILV3 = None if not _v3 else (_v3 == "1")
+_et = os.environ.get("FORCE_EXACTTAIL")
+FORCE_EXACTTAIL = None if not _et else (_et == "1")
+XSTATE = os.environ.get("XSTATE", "0") == "1"
+_ft = os.environ.get("FORCE_TIGHTEN")
+FORCE_TIGHTEN = None if _ft is None else bool(int(_ft))
+sms = torch.cuda.get_device_properties(0).multi_processor_count
+_EVICT = torch.empty(512 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda")
+
+if ARM == "pr":
+    sys.path.insert(0, "/home/scratch.siyid_coreai/workspace")
+    from prpkg16457.top_k.gvr_topk_decode import GvrTopKKernel  # noqa: E402
+elif ARM == "bx":
+    # loncheng PR #16877 head (BSX multi-tier): guard -> bsx dispatch,
+    # guard-fail -> that branch's own in-tree kernel (its production path)
+    sys.path.insert(0, "/home/scratch.siyid_coreai/wt-bsx/tensorrt_llm/_torch/cute_dsl_kernels")
+    from blackwell.top_k.gvr_topk_decode import GvrTopKKernel  # noqa: E402
+    from blackwell.top_k.gvr_topk_decode_bsx_dispatch import (  # noqa: E402
+        bsx_topk,
+        is_bsx_supported,
+    )
+else:
+    _root = os.environ.get("HARNESS_ROOT", str(Path(__file__).resolve().parents[5]))
+    sys.path.insert(0, _root + "/tests/scripts/cute_dsl_kernels/top_k")
+    sys.path.insert(0, _root)
+    import run_gvr_topk as rg  # noqa: E402
+    from run_gvr_topk import GvrTopKKernel  # noqa: E402
+
+    if ARM != "st":
+        from run_gvr_topk import (
+            emu_block_max,  # noqa: E402
+            emu_cand_bucketed,
+            emu_seed_counts,
+            pack_seed,
+        )
+
+
+def wf_targets(model, K, N):
+    if model == "flash":
+        T = (4096, 2048, 1024)
+    elif model == "v32":
+        T = (16384, 8192, 3072)
+    else:
+        T = (4096, 3584, 1536) if N <= 40960 else (12288, 5120, 2048)
+    return [min(t, N) for t in T]
+
+
+def main():
+    f = open(OUT, "w")
+    for model, isl in UNITS:
+        K, cr = KCFG[model]
+        Ls = LAYERS[model]
+        # ---- load per-layer step dicts (logits + captured topk) ----
+        lgs, pks = {}, {}
+        for L in Ls:
+            d = CAP / model / f"ISL_{isl}" / f"layer_{L:02d}"
+            lgs[L] = torch.load(d / "decode.logits.in.pt", map_location="cpu", weights_only=False)
+            pks[L] = torch.load(d / "decode.topk.out.pt", map_location="cpu", weights_only=False)
+        steps_all = sorted(pks[Ls[0]].keys())
+        # warmup steps carry -1 sentinel topk (and even wrong-width
+        # logits buffers): valid = captured topk present; a measured
+        # step also needs the PREVIOUS step valid (its topk is the
+        # preIdx seed)
+        valid = [s for s in steps_all if int(pks[Ls[0]][s].max()) >= 0]
+        vset = set(valid)
+        usable = [s for s in valid if (s - 1) in vset]
+        assert usable, f"no usable steps for {model}:{isl}"
+        NS = {}
+        for s in valid:
+            NS[s] = max(int(pks[L][s].max()) + 1 for L in Ls)
+        N_label = NS[valid[-1]]
+        # unit-constant pad width: block_max is validated against the
+        # BATCH buffer width (logits.shape[1]), so per-step bm must be
+        # built on rows already padded to the unit max
+        Npad_u = max((NS[s] + ALIGN - 1) // ALIGN * ALIGN for s in usable)
+        # ---- per (layer, step) prep on GPU ----
+        # row (padded), preIdx, lines, ref values, arm extras (1-row)
+        prep = {}
+        for L in Ls:
+            for s in usable:
+                Ns = NS[s]
+                Npad = Npad_u
+                row = torch.full((1, Npad), FMIN, dtype=torch.float32, device="cuda")
+                src = lgs[L][s][0]
+                row[0, :Ns] = src[:Ns].float().cuda()
+                pre = pks[L][s - 1].flatten().to(torch.int32).view(1, K)
+                pre = pre.cuda().contiguous()
+                srt = torch.sort(row[0, :Ns], descending=True).values
+                ref = srt[:K].clone()
+                ks = wf_targets(model, K, Ns)
+                t = torch.empty((1, 3), dtype=torch.float32, device="cuda")
+                for j, kc in enumerate(ks):
+                    t[0, j] = srt[kc - 1]
+                if TIGHT_MULT:
+                    # the closed loop derives its tight line from the
+                    # PREVIOUS step's k-th value, i.e. it sits just above
+                    # rank K - not at the fixed rank wf_targets uses. Model
+                    # that here so the tiers are measured on the contract
+                    # they actually ship with.
+                    t[0, 2] = srt[min(int(TIGHT_MULT * K), Ns) - 1]
+                for j in (1, 2):
+                    t[0, j] = torch.maximum(t[0, j], t[0, j - 1] + 1e-6)
+                if CONLY:  # codespell:ignore
+                    # push both tight lines above the row max so every
+                    # admitted entry lands in the loosest segment (the
+                    # one that already uses the cheap claim window).
+                    # CONLY_RANK tightens the collection line so the
+                    # claimed count stays under the consumer's line-cut
+                    # limit (and less volume is emitted).
+                    if CONLY_RANK:
+                        t[0, 0] = srt[min(CONLY_RANK * K, Ns) - 1]
+                    hi = float(row[0, :Ns].max()) + 1e4
+                    t[0, 1] = hi
+                    t[0, 2] = hi + 1.0
+                e = dict(Ns=Ns, Npad=Npad, row=row, pre=pre, ref=ref, sthr=t)
+                if ARM not in ("pr", "st", "bx"):
+                    sl1 = torch.tensor([Ns * cr], dtype=torch.int32, device="cuda")
+                    scnt = emu_seed_counts(row, sl1, t, compress_ratio=cr)
+                    if ARM == "wf":
+                        cv, ci, ctl = emu_cand_bucketed(
+                            row, sl1, t, 24576, seg_cap=8192, compress_ratio=cr, sentinel_pad=64
+                        )
+                        e.update(cv=cv, ci=ci, ctl=ctl)
+                    # thresholds mirror gvr_routing (real-capture A/B)
+                    nb = (ARM == "va" and Ns >= 65536) or (
+                        ARM == "vb" and model == "flash" and Ns >= 131072
+                    )
+                    if FORCE_BM is not None:
+                        nb = bool(FORCE_BM)
+                    if nb:
+                        e["bm"] = emu_block_max(row, sl1, compress_ratio=cr, tail_mode="exact")
+                    # col 6 = adaptive-skip pass count when bm attached
+                    e["spk"] = pack_seed(t, scnt, block_max=e.get("bm"))
+                prep[(L, s)] = e
+        Npad_max = max(e["Npad"] for e in prep.values())
+        for B in BS_LIST:
+            # reusable batch buffers (filled per rep BEFORE eviction)
+            lg_b = torch.full((B, Npad_max), FMIN, dtype=torch.float32, device="cuda")
+            pre_b = torch.zeros((B, K), dtype=torch.int32, device="cuda")
+            sl_b = torch.zeros((B,), dtype=torch.int32, device="cuda")
+            thr_b = torch.zeros((B, 3), dtype=torch.float32, device="cuda")
+            xs_b = torch.zeros((B, 8), dtype=torch.float32, device="cuda") if XSTATE else None
+            spk_b = torch.zeros((B, 8), dtype=torch.float32, device="cuda")
+            bufs = {}
+            if ARM in ("pr", "st", "bx"):
+                bufs["outb"] = torch.empty(B, K, dtype=torch.int32, device="cuda")
+            if ARM == "wf":
+                # widths from the emu contract (2*seg_cap+cap for the
+                # bucketed list), NOT the nominal cap
+                e0 = next(iter(prep.values()))
+                CW = e0["cv"].shape[1]
+                bufs["cv"] = torch.zeros((B, CW), dtype=torch.float32, device="cuda")
+                bufs["ci"] = torch.zeros((B, CW), dtype=torch.int32, device="cuda")
+                bufs["ctl"] = torch.zeros((B, e0["ctl"].shape[1]), dtype=torch.int32, device="cuda")
+            bm_b = None
+            if ARM in ("va", "vb"):
+                nbp = max((e["bm"].shape[1] for e in prep.values() if "bm" in e), default=0)
+                if nbp:
+                    bm_b = torch.zeros((B, nbp), dtype=torch.float32, device="cuda")
+            state = {}
+
+            def fill(L, s):
+                e = prep[(L, s)]
+                state.update(e=e)
+                lg_b[:, : e["Npad"]].copy_(e["row"].expand(B, e["Npad"]))
+                if e["Npad"] < Npad_max:
+                    lg_b[:, e["Npad"] :].fill_(FMIN)
+                pre_b.copy_(e["pre"].expand(B, K))
+                sl_b.fill_(e["Ns"] * cr)
+                thr_b.copy_(e["sthr"].expand(B, 3))
+                if ARM not in ("pr", "st", "bx"):
+                    spk_b.copy_(e["spk"].expand(B, 8))
+                if ARM == "wf":
+                    bufs["cv"].copy_(e["cv"].expand_as(bufs["cv"]))
+                    bufs["ci"].copy_(e["ci"].expand_as(bufs["ci"]))
+                    bufs["ctl"].copy_(e["ctl"].expand_as(bufs["ctl"]))
+                if bm_b is not None and "bm" in e:
+                    # unit-constant Npad -> bm width == bm_b width
+                    bm_b.copy_(e["bm"].expand_as(bm_b))
+                    state["bmc"] = bm_b
+                torch.cuda.synchronize()
+
+            def akw():
+                e = state["e"]
+                Ns = e["Ns"]
+                kw = dict(next_n=1, compress_ratio=cr, num_sms=sms, return_output_values=False)
+                if xs_b is not None:
+                    kw["xstate"] = xs_b
+                if FORCE_TAILFAST is not None:
+                    kw["p4_tail_fast"] = FORCE_TAILFAST
+                if FORCE_EXACTTAIL is not None:
+                    kw["p4_exact_tail"] = FORCE_EXACTTAIL
+                if FORCE_TAILV3 is not None:
+                    kw["p4_tail_v3"] = FORCE_TAILV3
+                if FORCE_NOFINE is not None:
+                    kw["p4_no_fine"] = FORCE_NOFINE
+                if FORCE_BINS is not None:
+                    kw["num_bins"] = FORCE_BINS
+                if FORCE_FINERT is not None:
+                    kw["p4_fine_rangetest"] = FORCE_FINERT
+                if FORCE_SCATRT is not None:
+                    kw["p4_scat_rangetest"] = FORCE_SCATRT
+                if FORCE_THREADS is not None:
+                    kw["num_threads_per_block"] = FORCE_THREADS
+                if ARM == "sw":
+                    return dict(kw)  # wrapper stock: no seed, no xstate
+                if ARM == "wf":
+                    a = dict(
+                        seed_thr=spk_b,
+                        cand_vals=bufs["cv"],
+                        cand_idx=bufs["ci"],
+                        cand_ctl=bufs["ctl"],
+                        cluster_size=1,
+                        **kw,
+                    )
+                    if FORCE_PWR is not None:
+                        a["p4_warp_redundant"] = FORCE_PWR
+                    return a
+                if ARM == "va":
+                    a = dict(seed_thr=spk_b, cluster_size=FORCE_CS or 1, **kw)
+                    if "bm" in e:
+                        a.update(block_max=state["bmc"], skip_min_n=None)
+                    if FORCE_PWR is not None:
+                        a["p4_warp_redundant"] = FORCE_PWR
+                    return a
+                # vb
+                cs = 1
+                if Ns >= 196608:
+                    if B * 8 <= sms // 2:
+                        cs = 8
+                    elif B * 4 <= (sms * 9) // 10:
+                        cs = 4
+                    elif B * 2 <= (sms * 9) // 10:
+                        cs = 2
+                a = dict(seed_thr=thr_b, cluster_size=cs, **kw)
+                if "bm" in e:
+                    a.update(block_max=state["bmc"], skip_min_n=None)
+                if FORCE_CS:
+                    a["cluster_size"] = FORCE_CS
+                if FORCE_PWR is not None:
+                    a["p4_warp_redundant"] = FORCE_PWR
+                return a
+
+            def call():
+                if ARM == "bx":
+                    if "bx_ok" not in state:
+                        state["bx_ok"] = is_bsx_supported(
+                            lg_b, pre_b, sl_b, bufs["outb"], K, 1, cr, None, None
+                        )
+                    if state["bx_ok"]:
+                        bsx_topk(lg_b, pre_b, sl_b, bufs["outb"], K, 1, cr)
+                    else:
+                        GvrTopKKernel.launch(lg_b, pre_b, sl_b, bufs["outb"], K, compress_ratio=cr)
+                    return
+                if ARM in ("pr", "st"):
+                    ov = {}
+                    if ARM == "st" and FORCE_PWR is not None:
+                        ov["p4_warp_redundant"] = FORCE_PWR
+                    GvrTopKKernel.launch(
+                        lg_b, pre_b, sl_b, bufs["outb"], K, compress_ratio=cr, **ov
+                    )
+                else:
+                    rg.gvr_topk_decode(lg_b, pre_b, sl_b, K, **akw())
+
+            def out_idx():
+                if ARM in ("pr", "st", "bx"):
+                    call()
+                    torch.cuda.synchronize()
+                    return bufs["outb"]
+                _, o = rg.gvr_topk_decode(lg_b, pre_b, sl_b, K, **akw())
+                torch.cuda.synchronize()
+                return o
+
+            for L in Ls:
+                base = f"{ARM}|{model}|{isl}|N{N_label}|B{B}|L{L:02d}"
+                rec = dict(
+                    model=model, isl=isl, N=N_label, K=K, B=B, arm=ARM, layer=L, steps=len(usable)
+                )
+                try:
+                    # exactness: every step at B==min; spot at larger B
+                    chk = usable if B == BS_LIST[0] else usable[:1]
+                    ok = True
+                    for s in chk:
+                        fill(L, s)
+                        o = out_idx()
+                        e = prep[(L, s)]
+                        for i in (0, B - 1) if B > 1 else (0,):
+                            gid = o[i, :K].long()
+                            # gather with an out-of-range index raises a
+                            # device-side assert that poisons the context and
+                            # kills the whole worker -> screen first
+                            if int(gid.min()) < 0 or int(gid.max()) >= e["Ns"]:
+                                ok = False
+                                rec["oob"] = rec.get("oob", 0) + 1
+                                continue
+                            got = lg_b[i, gid].sort(descending=True).values
+                            if gid.unique().numel() != K or not torch.equal(got, e["ref"]):
+                                ok = False
+                        if not ok:
+                            break
+                    rec["exact"] = ok
+                    # warmup
+                    for s in usable[:4]:
+                        fill(L, s)
+                        call()
+                    torch.cuda.synchronize()
+                    # COLD: one rep per decode step -> range mean = all-
+                    # steps mean for this layer
+                    xrec = []
+                    for s in usable:
+                        fill(L, s)
+                        if xs_b is not None:
+                            xs_b.zero_()
+                        _EVICT.uniform_(0, 1)
+                        torch.cuda.synchronize()
+                        torch.cuda.nvtx.range_push(f"c|{base}")
+                        call()
+                        torch.cuda.synchronize()
+                        torch.cuda.nvtx.range_pop()
+                        if xs_b is not None:
+                            xrec.append([round(v, 1) for v in xs_b[0].tolist()])
+                    if xs_b is not None:
+                        rec["x"] = xrec
+                    torch.cuda.synchronize()
+                except Exception as ex:  # noqa: BLE001
+                    rec["error"] = f"{type(ex).__name__}: {str(ex)[:120]}"
+                f.write(json.dumps(rec) + "\n")
+                f.flush()
+            del lg_b, pre_b, sl_b, thr_b, spk_b, bufs, bm_b, xs_b
+            torch.cuda.empty_cache()
+        prep.clear()
+        lgs.clear()
+        pks.clear()
+        torch.cuda.empty_cache()
+        print(f"[{ARM}] {model} {isl} done", flush=True)
+    f.close()
+    print(f"{ARM}_STEPS_DONE")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
new file mode 100644
index 000000000000..a4b21a0f683b
--- /dev/null
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+# ruff: noqa
+# Measurement harness committed verbatim for provenance; bench idioms
+# (loop-scoped buffers, del/rebind) trip static analysis.
+# f58: B×N 四张表(flash/pro × 算数均值/最小值),逐步配对 vs PR16457,
+# 我们的臂按线上路由(plan_emission)选取。
+import glob
+import os
+import pickle
+import sys
+
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent))
+from perstep import per_step  # noqa: E402
+
+D = os.environ.get("GVR_BENCH_OUT", "./bench_results")
+ARMS = {
+    "pr": "f58_*_pr.sqlite",
+    "st": "f58_*_st.sqlite",
+    "va": "f58_*_va.sqlite",
+    "vb": "f58_*_vb.sqlite",
+    "wf": "f58w_*.sqlite",
+}
+PKL = "./f58.pkl"
+if os.path.exists(PKL) and os.environ.get("REUSE", "1") == "1":
+    data = pickle.load(open(PKL, "rb"))
+else:
+    data = {}
+    for a, pat in ARMS.items():
+        agg = {}
+        for f in sorted(glob.glob(os.path.join(D, pat))):
+            for k, v in per_step(f).items():
+                agg.setdefault(k, []).extend(v)
+        data[a] = agg
+        print(f"  {a}: {len(agg)} 个(格,层)  {sum(len(v) for v in agg.values())} 步", flush=True)
+    pickle.dump(data, open(PKL, "wb"))
+
+sys.path.insert(
+    0,
+    str(
+        __import__("pathlib").Path(__file__).parents[5]
+        / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k"
+    ),
+)
+from gvr_routing import plan_emission  # noqa: E402
+
+KC = {"flash": 512, "pro": 1024, "v32": 2048}
+CR = {"flash": 4, "pro": 4, "v32": 1}
+TIER2ARM = {"list": "wf", "counts": "va", "rungs": "vb", "none": "st"}
+NS = ["8k", "16k", "32k", "64k", "128k", "256k"]
+BS = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
+
+cell = {}
+for m, isl, b in sorted({(m, i, b) for (m, i, b, l) in data["pr"]}):
+    n = int(isl[:-1]) * 1024 // CR[m]
+    tier = plan_emission(b, n, KC[m], True)
+    arm = TIER2ARM[tier]
+    ratios = []
+    for (mm, ii, bb, l), pr in data["pr"].items():
+        if (mm, ii, bb) != (m, isl, b):
+            continue
+        ou = data[arm].get((mm, ii, bb, l), [])
+        for j in range(min(len(pr), len(ou))):
+            ratios.append(pr[j] / ou[j])
+    if ratios:
+        cell[(m, isl, b)] = dict(
+            tier=tier, n=len(ratios), mean=sum(ratios) / len(ratios), mn=min(ratios)
+        )
+
+
+def table(m, key):
+    print(f"\n### {m} — {'算数均值' if key == 'mean' else '最小值'}(对 PR16457,全层×全步)")
+    print("| B \\ N | " + " | ".join(NS) + " |")
+    print("|---" * (len(NS) + 1) + "|")
+    for b in BS:
+        row = [f"**{b}**"]
+        for isl in NS:
+            c = cell.get((m, isl, b))
+            row.append(f"{c[key]:.3f}" if c else "—")
+        print("| " + " | ".join(row) + " |")
+
+
+for m in ("flash", "pro", "v32"):
+    for key in ("mean", "mn"):
+        table(m, key)
+
+tot = sum(c["n"] for c in cell.values())
+print(
+    f"\n共 {len(cell)} 格 / {tot} 逐步配对;各格档位:",
+    {
+        t: sum(1 for c in cell.values() if c["tier"] == t)
+        for t in ("list", "counts", "rungs", "none")
+    },
+)
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
new file mode 100755
index 000000000000..26ca3b196945
--- /dev/null
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+# Full-grid re-run on the current tree (post P4 fold work), fresh f58 prefix.
+BR=${GVR_BENCH_OUT:-./bench_results}
+cd "$(dirname "$0")"
+# PYTHONPATH must carry cutlass-dsl 4.5.x and the repo root; see README.md
+export TMPDIR=/tmp/nsys_g$1; mkdir -p $TMPDIR
+for UNIT in $2; do
+  TAG=$(echo $UNIT | tr ':' '_')
+  for ARM in ${ARMS:-pr st va vb}; do
+    rm -f $BR/f58_${TAG}_${ARM}*.csv $BR/f58_${TAG}_${ARM}.sqlite
+    CUDA_VISIBLE_DEVICES=$1 env ARM=$ARM UNITS=$UNIT \
+      GVR_BSTAR=8192 GVR_KC=8192 GVR_CAPC=16384 OUT=$BR/f58_${TAG}_${ARM}.jsonl \
+      nsys profile --trace=cuda,nvtx --force-overwrite true -o $BR/f58_${TAG}_${ARM} \
+        python3 ab_steps.py >> $BR/f58_${TAG}.log 2>&1
+    nsys stats -r nvtx_kern_sum --format csv -o $BR/f58_${TAG}_${ARM} \
+      $BR/f58_${TAG}_${ARM}.nsys-rep >> $BR/f58_${TAG}.log 2>&1
+    echo "f58 done ${TAG}_${ARM}" >> $BR/f58.log
+  done
+  # list tier (wf) under the C-only contract, same dataset
+  rm -f $BR/f58w_${TAG}*.csv $BR/f58w_${TAG}.sqlite
+  CUDA_VISIBLE_DEVICES=$1 env ARM=wf UNITS=$UNIT CONLY=1 CONLY_RANK=2 \  # codespell:ignore
+    GVR_BSTAR=8192 GVR_KC=8192 GVR_CAPC=16384 OUT=$BR/f58w_${TAG}.jsonl \
+    nsys profile --trace=cuda,nvtx --force-overwrite true -o $BR/f58w_${TAG} \
+      python3 ab_steps.py >> $BR/f58_${TAG}.log 2>&1
+  nsys stats -r nvtx_kern_sum --format csv -o $BR/f58w_${TAG} \
+    $BR/f58w_${TAG}.nsys-rep >> $BR/f58_${TAG}.log 2>&1
+  echo "f58 done ${TAG}_wf" >> $BR/f58.log
+done
+echo "f58 GPU$1 ALL DONE" >> $BR/f58.log
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py b/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py
new file mode 100644
index 000000000000..6f950dc07624
--- /dev/null
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py
@@ -0,0 +1,75 @@
+# ruff: noqa
+# Measurement harness committed verbatim for provenance; bench idioms
+# (loop-scoped buffers, del/rebind) trip static analysis.
+# Per-(layer, step) pairing straight out of the nsys sqlite.
+#
+# One NVTX range = one (arm, model, isl, N, B, layer); inside it the
+# harness runs one kernel per decode step, in step order. So the k-th
+# kernel of a range IS step k, and the arms can be paired index by
+# index. This is the only way to get a true per-step ratio - the CSV
+# summary only carries Avg/Min/Max per range, which cannot be paired.
+import glob
+import os
+import re
+import sqlite3
+import sys
+from collections import defaultdict
+
+D = os.environ.get("GVR_BENCH_OUT", "./bench_results")
+
+
+def per_step(path):
+    """{(model, isl, B, layer): [step0_us, step1_us, ...]} for one file."""
+    out = {}
+    con = sqlite3.connect(path)
+    rng = con.execute(
+        "SELECT start, end, text FROM NVTX_EVENTS WHERE text LIKE 'c|%' ORDER BY start"
+    ).fetchall()
+    ker = con.execute(
+        "SELECT start, end - start FROM CUPTI_ACTIVITY_KIND_KERNEL ORDER BY start"
+    ).fetchall()
+    con.close()
+    # ranges do not overlap (one layer at a time), so a single merge walk
+    # assigns every kernel to the range containing it
+    i = 0
+    for rs, re_, text in rng:
+        while i < len(ker) and ker[i][0] < rs:
+            i += 1
+        p = text.lstrip(":").split("|")
+        if len(p) != 7:
+            while i < len(ker) and ker[i][0] < re_:
+                i += 1
+            continue
+        _, arm, m, isl, _n, b, l = p
+        acc = out.setdefault((m, isl, int(b[1:]), l), [])
+        while i < len(ker) and ker[i][0] < re_:
+            acc.append(ker[i][1] / 1000.0)
+            i += 1
+    return out
+
+
+def collect(pattern):
+    agg = {}
+    for f in sorted(glob.glob(os.path.join(D, pattern))):
+        for k, v in per_step(f).items():
+            agg.setdefault(k, []).extend(v) if k in agg else agg.setdefault(k, v)
+    return agg
+
+
+if __name__ == "__main__":
+    arms = {
+        "pr": "f22_*_pr.sqlite",
+        "st": "f22_*_st.sqlite",
+        "va": "f22_*_va.sqlite",
+        "vb": "f22_*_vb.sqlite",
+        "wf": "f30_*.sqlite",
+    }
+    data = {a: collect(p) for a, p in arms.items()}
+    for a, d in data.items():
+        n = sum(len(v) for v in d.values())
+        print(f"  {a}: {len(d)} 个(格,层)  {n} 个步", flush=True)
+    import pickle
+
+    with open("/home/scratch.siyid_coreai/workspace/perstep.pkl", "wb") as f:
+        pickle.dump(data, f)
+    print("saved perstep.pkl")

From c4de505ac74258e89df698b4e092be1c51a37829 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 11 Aug 2026 05:39:13 -0700
Subject: [PATCH 103/117] [None][chore] size list-tier candidate buffers by the
 routing cap

Review touch-ups (approved round): the wide candidate buffers only
serve the list tier, which plan_emission picks at
batch <= LIST_EMIT_MAX_B - allocate that many rows instead of
max_rows (~0.33 MB/row/layer, ~5 GB -> ~80 MB at max_rows=256 x 61
layers); clamp the per-step control-word zeroing accordingly. Also
state the engine-static routing-N assumption at the plan call site.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../attention_backend/sparse/dsa/indexer.py   |  5 +++++
 .../attention_backend/sparse/gvr_emission.py  | 20 +++++++++++++------
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index 34c04ef3dcc2..686c0ddd37b9 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -1761,6 +1761,11 @@ def sparse_attn_indexer(
                         st = self._ensure_gvr_emission(metadata, q_fp8.device)
                         # indexer_max_seq_len is already the compressed length
                         # (get_indexer_max_seq_len divides); do not divide again.
+                        # routing N is the ENGINE-STATIC max (graph capture
+                        # bakes the tier in); short live rows in a long-max
+                        # engine run assist machinery the planner would refuse
+                        # at their true length - exactness holds via the
+                        # in-kernel guards, the tax is routed pessimistically
                         emit_tier, self._gvr_route = st.plan(
                             batch_size,
                             indexer_max_seq_len,
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
index 770ce37647dd..1aa77f3cc1a5 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
@@ -33,6 +33,7 @@
 import torch
 
 from ...cute_dsl_kernels.blackwell.top_k.gvr_routing import (
+    LIST_EMIT_MAX_B,
     LIST_EMIT_MIN_N,
     TopkRoute,
     pick_config,
@@ -90,10 +91,16 @@ def __init__(
         self.cand_ctl: Optional[torch.Tensor] = None
         self.cand_cur: Optional[torch.Tensor] = None
         if enable_list_tier:
-            self.cand_vals = torch.zeros((max_rows, LIST_WIDTH), dtype=torch.float32, device=device)
-            self.cand_idx = torch.zeros((max_rows, LIST_WIDTH), dtype=torch.int32, device=device)
-            self.cand_ctl = torch.zeros((max_rows, 4), dtype=torch.int32, device=device)
-            self.cand_cur = torch.zeros((max_rows, 4), dtype=torch.int32, device=device)
+            # the routing only ever picks the list tier at
+            # batch <= LIST_EMIT_MAX_B, so the wide candidate buffers
+            # need that many rows, not max_rows (~0.33 MB/row/layer)
+            cand_rows = min(max_rows, LIST_EMIT_MAX_B)
+            self.cand_vals = torch.zeros(
+                (cand_rows, LIST_WIDTH), dtype=torch.float32, device=device
+            )
+            self.cand_idx = torch.zeros((cand_rows, LIST_WIDTH), dtype=torch.int32, device=device)
+            self.cand_ctl = torch.zeros((cand_rows, 4), dtype=torch.int32, device=device)
+            self.cand_cur = torch.zeros((cand_rows, 4), dtype=torch.int32, device=device)
         # previous-step top-k feedback (address-stable; zero-init ->
         # first step's pre_idx points at index 0, a benign candidate)
         self.prev_topk = torch.zeros((max_rows, top_k), dtype=torch.int32, device=device)
@@ -204,8 +211,9 @@ def update_seed_rows(self, num_rows: int, emit_tier: str = "counts") -> None:
         rungs[:, 1] = new1
         rungs[:, 2] = new2
         if self.cand_ctl is not None:
-            self.cand_ctl[:num_rows].zero_()
-            self.cand_cur[:num_rows].zero_()
+            nc = min(num_rows, self.cand_ctl.shape[0])
+            self.cand_ctl[:nc].zero_()
+            self.cand_cur[:nc].zero_()
 
     def indexer_emit_kwargs(self, emit_tier: str, num_rows: int) -> dict:
         """kwargs for CuteDSLFP4PagedMQALogitsRunner.forward covering the

From ba1d698f814048ed947422515987e36dc010c550 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 14 Aug 2026 02:46:14 -0700
Subject: [PATCH 104/117] [None][fix] restore the plateau terminal dropped in
 the main merge

The merge of main (which brought in #16877) resolved the shared kernel
file toward this branch's emission work and unintentionally dropped the
plateau-terminal feature added there (4642d92643, e382f9814d): when a
bitwise-equal tie plateau wider than the candidate buffer straddles the
K boundary, the bracket admits no threshold and the row previously fell
through to the legacy give-up, leaving -1 pads (CI: 10/10
plateau_terminal params failed at the merge head).

Port both commits onto the current drivers:
- leader driver: keep this branch's slope-fit retry loop and append the
  budget-exhausted bisection collapse behind it (coherent
  undershoot-overflow guard; the retry's bracket widening marks a side
  stale with -1 and fails the guard) -> adjacent-float bracket sets
  done = 3 and a recount at the terminal threshold feeds Phase 3;
- phase2_secant_search: same post-loop collapse ahead of the legacy
  give-up;
- register-resident redundant driver: adjacent-float terminal inside the
  refine loop plus the post-loop collapse, warp-uniform by replay;
- Phase 4: plateau fill from the tie class (ticket in the dedicated
  s_iscalars[7], seeded from the pre-P4 cand_count_p4 snapshot - the
  s_iscalars[0] slot is radix scratch by then), pad guard keyed on the
  captured s_iscalars[6] flag, for both the cs=1 and cs>1 leader paths.

B200: plateau_terminal 10/10; full gvr decode file 744 passed /
1 xpassed / 0 failed; tiers file 95 passed / 0 failed.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 384 +++++++++++++++++-
 1 file changed, 363 insertions(+), 21 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 904c56eac1f1..e15cb0be6fbe 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -2613,8 +2613,18 @@ def phase2_secant_search(
                 if nv == vlo_r or nv == vhi_r:
                     nv = (vlo_r + vhi_r) * cutlass.Float32(0.5)
                     if nv == vlo_r or nv == vhi_r:
-                        thr_r = vlo_r
-                        done_r = cutlass.Int32(2)
+                        # ADJACENT-FLOAT bracket, same terminal as the
+                        # leader path: a low side over the candidate
+                        # buffer plus a high side under K means the
+                        # boundary sits inside a bitwise-equal plateau
+                        # wider than kC. Keep the sure-winner threshold
+                        # and let Phase 4's plateau fill finish the row.
+                        if clo_r > cutlass.Int32(kCC) and chi_r < cutlass.Int32(kK):
+                            thr_r = vhi_r
+                            done_r = cutlass.Int32(3)
+                        else:
+                            thr_r = vlo_r
+                            done_r = cutlass.Int32(2)
                 if done_r == cutlass.Int32(0):
                     thr_r = nv
                     par_r = par_r ^ cutlass.Int32(1)
@@ -2644,6 +2654,75 @@ def phase2_secant_search(
                         vhi_r = thr_r
                         chi_r = cnt_r
                 it = it + cutlass.Int32(1)
+            # ---- Budget-exhausted plateau collapse (mirrors the leader
+            # path): the refine budget can run out while the bracket is
+            # still wide because a tie plateau wider than kC admits no
+            # threshold. On exactly that signature, bisect to adjacent
+            # floats so the plateau terminal is exact. Every thread
+            # replays this from identical registers, so the branch stays
+            # warp-uniform and block_count_ge keeps its barrier cadence.
+            if (
+                done_r == cutlass.Int32(0)
+                and clo_r > cutlass.Int32(kCC)
+                and chi_r >= cutlass.Int32(0)
+                and chi_r < cutlass.Int32(kK)
+            ):
+                itc = cutlass.Int32(0)
+                while itc < cutlass.Int32(64) and done_r == cutlass.Int32(0):
+                    mid_c = (vlo_r + vhi_r) * cutlass.Float32(0.5)
+                    if mid_c == vlo_r or mid_c == vhi_r:
+                        thr_r = vhi_r
+                        done_r = cutlass.Int32(3)
+                    else:
+                        thr_r = mid_c
+                        par_r = par_r ^ cutlass.Int32(1)
+                        cnt_r = self.block_count_ge(
+                            input_row,
+                            slice_start,
+                            slice_end,
+                            thr_r,
+                            smem_ptcnt,
+                            smem_wcnt,
+                            s_iscalars,
+                            s_cluster_partial,
+                            tidx,
+                            warp_id,
+                            lane,
+                            cutlass.Boolean(False),  # do_cluster_sync (cs==1)
+                            smem_input=smem_input,
+                            redundant=True,
+                            wcnt_off=par_r * cutlass.Int32(nwp2),
+                        )
+                        if cnt_r >= cutlass.Int32(kK) and cnt_r <= cutlass.Int32(kCC):
+                            done_r = cutlass.Int32(1)
+                        elif cnt_r > cutlass.Int32(kCC):
+                            vlo_r = thr_r
+                            clo_r = cnt_r
+                        else:
+                            vhi_r = thr_r
+                            chi_r = cnt_r
+                    itc = itc + cutlass.Int32(1)
+                if done_r == cutlass.Int32(3):
+                    # recount at the terminal threshold so Phase 3 sees
+                    # per-thread counts for the sure-winner set.
+                    par_r = par_r ^ cutlass.Int32(1)
+                    cnt_r = self.block_count_ge(
+                        input_row,
+                        slice_start,
+                        slice_end,
+                        thr_r,
+                        smem_ptcnt,
+                        smem_wcnt,
+                        s_iscalars,
+                        s_cluster_partial,
+                        tidx,
+                        warp_id,
+                        lane,
+                        cutlass.Boolean(False),  # do_cluster_sync (cs==1)
+                        smem_input=smem_input,
+                        redundant=True,
+                        wcnt_off=par_r * cutlass.Int32(nwp2),
+                    )
             if done_r == cutlass.Int32(0):
                 if clo_r <= cutlass.Int32(kCC * 2):
                     thr_r = vlo_r
@@ -2729,11 +2808,22 @@ def phase2_secant_search(
                     nv = vhi - rng * cutlass.Float32(0.05)
 
                 if nv == vlo or nv == vhi:
-                    # Bracket exhausted — try midpoint, else give up.
+                    # Bracket exhausted — try midpoint, else terminal.
                     nv = (vlo + vhi) * cutlass.Float32(0.5)
                     if nv == vlo or nv == vhi:
-                        s_thr[0] = vlo
-                        s_iscalars[1] = cutlass.Int32(2)  # done = 2 (give up)
+                        # ADJACENT-FLOAT bracket: every value in
+                        # [vlo, vhi) is bitwise-equal to vlo. Low side
+                        # overflowing the candidate buffer AND high side
+                        # undershooting K means the boundary sits inside
+                        # a bitwise-equal plateau wider than kC — record
+                        # the plateau terminal (done = 3) and keep the
+                        # sure-winner threshold vhi.
+                        if clo > cutlass.Int32(kCC) and chi < cutlass.Int32(kK):
+                            s_thr[0] = vhi
+                            s_iscalars[1] = cutlass.Int32(3)  # done = 3 (plateau)
+                        else:
+                            s_thr[0] = vlo
+                            s_iscalars[1] = cutlass.Int32(2)  # done = 2 (give up)
                     else:
                         s_thr[0] = nv
                 else:
@@ -2773,6 +2863,84 @@ def phase2_secant_search(
                 cute.arch.barrier()
             it = it + cutlass.Int32(1)
 
+        # ---- Budget-exhausted plateau collapse ----
+        # The refine budget can run out while the bracket is still wide: the
+        # secant step keeps making progress (the bracket shrinks every
+        # iteration) but a tie plateau wider than kC admits no threshold, so
+        # the count never lands in [kK, kCC]. In exactly that signature -
+        # count(>= v_lo) > kCC AND count(>= v_hi) < kK, both counts current -
+        # collapse the bracket by pure bisection until the ends are ADJACENT
+        # floats; every value in [v_lo, v_hi) is then bitwise-equal, so the
+        # plateau terminal (done = 3) is exact and Phase 4 completes the row
+        # from that tie class. A count landing in [kK, kCC] mid-collapse
+        # converges normally. Anything else keeps the legacy give-up below.
+        if (
+            s_iscalars[1] == cutlass.Int32(0)
+            and s_iscalars[2] > cutlass.Int32(kCC)
+            and s_iscalars[3] >= cutlass.Int32(0)
+            and s_iscalars[3] < cutlass.Int32(kK)
+        ):
+            itc = cutlass.Int32(0)
+            while itc < cutlass.Int32(64) and s_iscalars[1] == cutlass.Int32(0):
+                if tidx == 0:
+                    vlo_c = s_thr[1]
+                    vhi_c = s_thr[2]
+                    mid_c = (vlo_c + vhi_c) * cutlass.Float32(0.5)
+                    if mid_c == vlo_c or mid_c == vhi_c:
+                        s_thr[0] = vhi_c
+                        s_iscalars[1] = cutlass.Int32(3)  # plateau terminal
+                    else:
+                        s_thr[0] = mid_c
+                cute.arch.barrier()
+                if s_iscalars[1] == cutlass.Int32(0):
+                    self.block_count_ge(
+                        input_row,
+                        slice_start,
+                        slice_end,
+                        s_thr[0],
+                        smem_ptcnt,
+                        smem_wcnt,
+                        s_iscalars,
+                        s_cluster_partial,
+                        tidx,
+                        warp_id,
+                        lane,
+                        smem_input=smem_input,
+                        do_cluster_sync=do_cluster_sync,
+                    )
+                    if tidx == 0:
+                        c_c = s_iscalars[0]
+                        t_c = s_thr[0]
+                        if c_c >= cutlass.Int32(kK) and c_c <= cutlass.Int32(kCC):
+                            s_iscalars[1] = cutlass.Int32(1)
+                        elif c_c > cutlass.Int32(kCC):
+                            s_thr[1] = t_c
+                            s_iscalars[2] = c_c
+                        else:
+                            s_thr[2] = t_c
+                            s_iscalars[3] = c_c
+                    cute.arch.barrier()
+                itc = itc + cutlass.Int32(1)
+            if s_iscalars[1] == cutlass.Int32(3):
+                # recount at the terminal threshold so Phase 3's cached
+                # per-thread counts describe the sure-winner set.
+                self.block_count_ge(
+                    input_row,
+                    slice_start,
+                    slice_end,
+                    s_thr[0],
+                    smem_ptcnt,
+                    smem_wcnt,
+                    s_iscalars,
+                    s_cluster_partial,
+                    tidx,
+                    warp_id,
+                    lane,
+                    smem_input=smem_input,
+                    do_cluster_sync=do_cluster_sync,
+                )
+                cute.arch.barrier()
+
         # ---- Post-loop fallback: if still not done, force threshold ----
         if tidx == 0:
             if s_iscalars[1] == cutlass.Int32(0):
@@ -4876,12 +5044,13 @@ def phase4_rank_scatter(
                     output_values_row[i10] = self.dtype(smem_keys[i10])
                 output_indices_row[i10] = smem_vals[i10]
                 i10 = i10 + cutlass.Int32(num_threads)
-            i11 = cand_count + tidx
-            while i11 < cutlass.Int32(kK):
-                if cutlass.const_expr(self.return_output_values):
-                    output_values_row[i11] = self.dtype(self.NEG_FLT_MAX)
-                output_indices_row[i11] = cutlass.Int32(-1)
-                i11 = i11 + cutlass.Int32(num_threads)
+            if s_iscalars[6] == cutlass.Int32(0):  # plateau fill completes done=3
+                i11 = cand_count + tidx
+                while i11 < cutlass.Int32(kK):
+                    if cutlass.const_expr(self.return_output_values):
+                        output_values_row[i11] = self.dtype(self.NEG_FLT_MAX)
+                    output_indices_row[i11] = cutlass.Int32(-1)
+                    i11 = i11 + cutlass.Int32(num_threads)
 
     # ------------------------------------------------------------------
     # Phase 4: Histogram-based k-th selection + two-pass writeback
@@ -5387,12 +5556,13 @@ def phase4_histogram_snap(
                     )
                 output_indices_row[i10] = self._smem_ld(cutlass.Int32, vals_base, i10)
                 i10 = i10 + cutlass.Int32(num_threads)
-            i11 = cand_count + tidx
-            while i11 < cutlass.Int32(kK):
-                if cutlass.const_expr(self.return_output_values):
-                    output_values_row[i11] = self.dtype(self.NEG_FLT_MAX)
-                output_indices_row[i11] = cutlass.Int32(-1)
-                i11 = i11 + cutlass.Int32(num_threads)
+            if s_iscalars[6] == cutlass.Int32(0):  # plateau fill completes done=3
+                i11 = cand_count + tidx
+                while i11 < cutlass.Int32(kK):
+                    if cutlass.const_expr(self.return_output_values):
+                        output_values_row[i11] = self.dtype(self.NEG_FLT_MAX)
+                    output_indices_row[i11] = cutlass.Int32(-1)
+                    i11 = i11 + cutlass.Int32(num_threads)
 
     # ------------------------------------------------------------------
     # Main kernel — one CTA per row
@@ -5724,9 +5894,13 @@ def run_one_row(
         #   [4] out_count
         #   [5] local cand_count  (per-CTA snapshot before cluster all-reduce;
         #                          consumed by the kernel-level cluster handoff)
+        #   [6] plateau terminal flag, captured from [1] BEFORE Phase 4
+        #       (Phase 4 REUSES [1] as radix scratch, so the terminal must
+        #       never be re-read from it afterwards)
+        #   [7] plateau fill ticket (done == 3 only)
         s_iscalars = smem.allocate_tensor(
             element_type=cutlass.Int32,
-            layout=cute.make_ordered_layout((6,), order=(0,)),
+            layout=cute.make_ordered_layout((8,), order=(0,)),
             byte_alignment=16,
         )
         # Per-CTA DSMEM scratch for the cluster all-reduce of cand_count:
@@ -7391,9 +7565,96 @@ def _run_phases(
                                             s_iscalars[2] = cutlass.Int32(-1)
                                 cute.arch.barrier()
                                 rs = rs + cutlass.Int32(1)
-                            if s_iscalars[1] != cutlass.Int32(1):
-                                # tie-plateau fail-soft: land on the measured
-                                # undershoot side (count <= kC => no overflow).
+                        if s_iscalars[1] != cutlass.Int32(1):
+                            # The retry budget could not land in [K, kC].
+                            # ONLY the coherent undershoot-overflow corner
+                            # (count(>= lo) > kC AND 0 <= count(>= hi) < K,
+                            # both counts CURRENT — the retry's bracket
+                            # widening marks a side stale with -1 and thus
+                            # fails this guard) collapses the bracket by
+                            # pure bisection to ADJACENT floats, where the
+                            # plateau terminal (done = 3, threshold = hi)
+                            # is exact: Phase 4 emits the sure winners and
+                            # the plateau fill completes the row from the
+                            # tie class. A mid-collapse count landing in
+                            # [K, kC] converges normally; anything else
+                            # (incl. an exhausted collapse budget) falls
+                            # through to the fail-soft terminal below.
+                            it4 = cutlass.Int32(0)
+                            if (
+                                s_iscalars[2] <= cutlass.Int32(self.kC)
+                                or s_iscalars[3] < cutlass.Int32(0)
+                                or s_iscalars[3] >= cutlass.Int32(self.top_k)
+                            ):
+                                it4 = cutlass.Int32(40)  # guard: skip collapse
+                            while it4 < cutlass.Int32(40) and s_iscalars[1] == cutlass.Int32(0):
+                                if tidx == cutlass.Int32(0):
+                                    lo4 = s_thr[1]
+                                    hi4 = s_thr[2]
+                                    mid4 = (lo4 + hi4) * cutlass.Float32(0.5)
+                                    if mid4 == lo4 or mid4 == hi4:
+                                        s_thr[0] = hi4
+                                        s_iscalars[1] = cutlass.Int32(3)
+                                    else:
+                                        s_thr[0] = mid4
+                                cute.arch.barrier()
+                                if s_iscalars[1] == cutlass.Int32(0):
+                                    self.block_count_ge(
+                                        input_row,
+                                        slice_start,
+                                        slice_end,
+                                        s_thr[0],
+                                        smem_ptcnt,
+                                        smem_wcnt,
+                                        s_iscalars,
+                                        s_cluster_partial,
+                                        tidx,
+                                        warp_id,
+                                        lane,
+                                        do_cluster_sync=do_cluster_sync,
+                                        smem_input=smem_input,
+                                    )
+                                    cute.arch.barrier()
+                                    if tidx == cutlass.Int32(0):
+                                        c4 = s_iscalars[0]
+                                        t4 = s_thr[0]
+                                        if c4 >= cutlass.Int32(self.top_k) and c4 <= cutlass.Int32(
+                                            self.kC
+                                        ):
+                                            s_iscalars[1] = cutlass.Int32(1)
+                                        elif c4 > cutlass.Int32(self.kC):
+                                            s_thr[1] = t4
+                                            s_iscalars[2] = c4
+                                        else:
+                                            s_thr[2] = t4
+                                            s_iscalars[3] = c4
+                                    cute.arch.barrier()
+                                it4 = it4 + cutlass.Int32(1)
+                            if s_iscalars[1] == cutlass.Int32(3):
+                                # recount at the terminal threshold so P3's
+                                # cached per-thread counts describe the
+                                # sure-winner set the fill completes.
+                                self.block_count_ge(
+                                    input_row,
+                                    slice_start,
+                                    slice_end,
+                                    s_thr[0],
+                                    smem_ptcnt,
+                                    smem_wcnt,
+                                    s_iscalars,
+                                    s_cluster_partial,
+                                    tidx,
+                                    warp_id,
+                                    lane,
+                                    do_cluster_sync=do_cluster_sync,
+                                    smem_input=smem_input,
+                                )
+                                cute.arch.barrier()
+                            elif s_iscalars[1] != cutlass.Int32(1):
+                                # fail-soft (non-plateau): land on the
+                                # measured undershoot side (count <= kC =>
+                                # no overflow; -1 pad stays the documented
+                                # non-convergence encoding).
                                 self.block_count_ge(
                                     input_row,
                                     slice_start,
@@ -7497,6 +7758,13 @@ def _run_phases(
             cand_count_p4 = cutlass.Int32(0)
             if cutlass.const_expr(cluster_size == 1):
                 # cs=1: the single CTA per row IS the leader.
+                # Capture the P2 terminal BEFORE Phase 4: P4 reuses
+                # s_iscalars[1] as radix scratch.
+                if tidx == cutlass.Int32(0):
+                    s_iscalars[6] = cutlass.Int32(0)
+                    if s_iscalars[1] == cutlass.Int32(3):
+                        s_iscalars[6] = cutlass.Int32(1)
+                cute.arch.barrier()
                 cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC))
                 if cutlass.const_expr(self.enable_p4_rank_scatter):
                     if cutlass.const_expr(
@@ -7551,6 +7819,39 @@ def _run_phases(
                         warp_id,
                         lane,
                     )
+                # ---- plateau fill (done == 3): complete the row from the
+                # bitwise-equal plateau class. The terminal is only set on an
+                # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0])
+                # is bitwise-equal; Phase 4 has already emitted the
+                # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of
+                # the tie class is a valid tie-aware completion. Ticket counter
+                # lives in the DEDICATED s_iscalars[7].
+                if s_iscalars[6] == cutlass.Int32(1):
+                    pv_lo = s_thr[1]
+                    pv_hi = s_thr[0]
+                    if tidx == cutlass.Int32(0):
+                        # cand_count_p4 was captured BEFORE Phase 4;
+                        # s_iscalars[0] is radix scratch by now (same
+                        # hazard as the flag).
+                        s_iscalars[7] = cand_count_p4
+                    cute.arch.barrier()
+                    ifp = tidx
+                    while ifp < N:
+                        vfp = cutlass.Float32(0.0)
+                        if cutlass.const_expr(self.dtype == cutlass.Float32):
+                            vfp = input_row[ifp]
+                        else:
+                            vfp = cutlass.Float32(input_row[ifp])
+                        if vfp >= pv_lo and vfp < pv_hi:
+                            pfill = atomicAdd(
+                                s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1)
+                            )
+                            if pfill < cutlass.Int32(self.top_k):
+                                if cutlass.const_expr(self.return_output_values):
+                                    output_values_row[pfill] = self.dtype(vfp)
+                                output_indices_row[pfill] = ifp
+                        ifp = ifp + cutlass.Int32(self.num_threads)
+                    cute.arch.barrier()
                 ck_sw0 = cutlass.Int64(0)
                 ck_sw1 = cutlass.Int64(0)
                 if cutlass.const_expr(_P4_SUB_DBG):
@@ -7712,6 +8013,13 @@ def _run_phases(
                     # smem_keys/smem_vals (no peers to gather from).
 
                     # ---- Phase 4: histogram snap + writeback ----
+                    # Capture the P2 terminal BEFORE Phase 4: P4
+                    # reuses s_iscalars[1] as radix scratch.
+                    if tidx == cutlass.Int32(0):
+                        s_iscalars[6] = cutlass.Int32(0)
+                        if s_iscalars[1] == cutlass.Int32(3):
+                            s_iscalars[6] = cutlass.Int32(1)
+                    cute.arch.barrier()
                     cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC))
                     if cutlass.const_expr(self.enable_p4_rank_scatter):
                         self.phase4_rank_scatter(
@@ -7761,6 +8069,40 @@ def _run_phases(
                                 xstate_row[5] = cutlass.Float32(s_mt_cnt[1])
                                 xstate_row[6] = cutlass.Float32(s_mt_cnt[2])
 
+                    # ---- plateau fill (done == 3): complete the row from the
+                    # bitwise-equal plateau class. The terminal is only set on an
+                    # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0])
+                    # is bitwise-equal; Phase 4 has already emitted the
+                    # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of
+                    # the tie class is a valid tie-aware completion. Ticket counter
+                    # lives in the DEDICATED s_iscalars[7].
+                    if s_iscalars[6] == cutlass.Int32(1):
+                        pv_lo = s_thr[1]
+                        pv_hi = s_thr[0]
+                        if tidx == cutlass.Int32(0):
+                            # cand_count_p4 was captured BEFORE Phase 4;
+                            # s_iscalars[0] is radix scratch by now (same
+                            # hazard as the flag).
+                            s_iscalars[7] = cand_count_p4
+                        cute.arch.barrier()
+                        ifp = tidx
+                        while ifp < N:
+                            vfp = cutlass.Float32(0.0)
+                            if cutlass.const_expr(self.dtype == cutlass.Float32):
+                                vfp = input_row[ifp]
+                            else:
+                                vfp = cutlass.Float32(input_row[ifp])
+                            if vfp >= pv_lo and vfp < pv_hi:
+                                pfill = atomicAdd(
+                                    s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1)
+                                )
+                                if pfill < cutlass.Int32(self.top_k):
+                                    if cutlass.const_expr(self.return_output_values):
+                                        output_values_row[pfill] = self.dtype(vfp)
+                                    output_indices_row[pfill] = ifp
+                            ifp = ifp + cutlass.Int32(self.num_threads)
+                        cute.arch.barrier()
+
         # Final cluster barrier: keep peer CTAs (and their SMEM) alive
         # until the leader's gather + Phase 4 finish. Skipped at
         # do_cluster_sync=False (no peers; short-row degrade non-leaders

From 6c581a71008d80fb2605a10b06fe34ba1c699b89 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 20 Aug 2026 08:15:52 -0700
Subject: [PATCH 105/117] [None][fix] fp4 paged MQA logits: make_fragment ->
 make_rmem_tensor

The B300 lane installs a CuTe DSL without make_fragment; the emission
unittests this PR adds run there and hit the removed API
(AttributeError across test_cute_dsl_fp4_paged_mqa_logits_block_meta/
seed_counts/cand). Renamed all uses in the file to make_rmem_tensor,
matching the DSL 4.6.1 API main already uses in gvr_topk_decode.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../paged_mqa_logits/fp4_paged_mqa_logits.py  | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
index 37434d104a7c..89f32856d77c 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py
@@ -2211,17 +2211,17 @@ def kernel(
                     meta_warp = local_tidx // 32
                     meta_lane = local_tidx % 32
                     if cutlass.const_expr(self.emit_seed_counts):
-                        sthr = cute.make_fragment(next_n * 3, cutlass.Float32)
-                        scnt = cute.make_fragment(next_n * 3, cutlass.Int32)
-                        spass = cute.make_fragment(next_n, cutlass.Int32)
+                        sthr = cute.make_rmem_tensor(next_n * 3, cutlass.Float32)
+                        scnt = cute.make_rmem_tensor(next_n * 3, cutlass.Int32)
+                        spass = cute.make_rmem_tensor(next_n, cutlass.Int32)
                         for _i in cutlass.range_constexpr(next_n):
                             spass[_i] = cutlass.Int32(0)
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_cand or self.emit_cand_bucketed):
-                        cwbase = cute.make_fragment(next_n, cutlass.Int32)
-                        cwleft = cute.make_fragment(next_n, cutlass.Int32)
+                        cwbase = cute.make_rmem_tensor(next_n, cutlass.Int32)
+                        cwleft = cute.make_rmem_tensor(next_n, cutlass.Int32)
                         for _i in cutlass.range_constexpr(next_n):
                             cwbase[_i] = cutlass.Int32(0)
                             cwleft[_i] = cutlass.Int32(0)
@@ -2229,10 +2229,10 @@ def kernel(
                         # Per-lane hit accumulators, carried across all
                         # tiles of the same q and flushed once per
                         # q-transition — no warp-wide ops per tile.
-                        hacc_min = cute.make_fragment(next_n, cutlass.Float32)
-                        hacc_max = cute.make_fragment(next_n, cutlass.Float32)
-                        hacc_sum = cute.make_fragment(next_n, cutlass.Float32)
-                        hacc_cnt = cute.make_fragment(next_n, cutlass.Int32)
+                        hacc_min = cute.make_rmem_tensor(next_n, cutlass.Float32)
+                        hacc_max = cute.make_rmem_tensor(next_n, cutlass.Float32)
+                        hacc_sum = cute.make_rmem_tensor(next_n, cutlass.Float32)
+                        hacc_cnt = cute.make_rmem_tensor(next_n, cutlass.Int32)
                         for _t in cutlass.range_constexpr(next_n):
                             hacc_min[_t] = cutlass.Float32(_META_FLT_MAX)
                             hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX)
@@ -2985,17 +2985,17 @@ def kernel(
                     meta_warp = local_tidx // 32
                     meta_lane = local_tidx % 32
                     if cutlass.const_expr(self.emit_seed_counts):
-                        sthr = cute.make_fragment(next_n * 3, cutlass.Float32)
-                        scnt = cute.make_fragment(next_n * 3, cutlass.Int32)
-                        spass = cute.make_fragment(next_n, cutlass.Int32)
+                        sthr = cute.make_rmem_tensor(next_n * 3, cutlass.Float32)
+                        scnt = cute.make_rmem_tensor(next_n * 3, cutlass.Int32)
+                        spass = cute.make_rmem_tensor(next_n, cutlass.Int32)
                         for _i in cutlass.range_constexpr(next_n):
                             spass[_i] = cutlass.Int32(0)
                         for _i in cutlass.range_constexpr(next_n * 3):
                             sthr[_i] = cutlass.Float32(_META_FLT_MAX)
                             scnt[_i] = cutlass.Int32(0)
                     if cutlass.const_expr(self.emit_cand or self.emit_cand_bucketed):
-                        cwbase = cute.make_fragment(next_n, cutlass.Int32)
-                        cwleft = cute.make_fragment(next_n, cutlass.Int32)
+                        cwbase = cute.make_rmem_tensor(next_n, cutlass.Int32)
+                        cwleft = cute.make_rmem_tensor(next_n, cutlass.Int32)
                         for _i in cutlass.range_constexpr(next_n):
                             cwbase[_i] = cutlass.Int32(0)
                             cwleft[_i] = cutlass.Int32(0)
@@ -3003,10 +3003,10 @@ def kernel(
                         # Per-lane hit accumulators, carried across all
                         # tiles of the same q and flushed once per
                         # q-transition — no warp-wide ops per tile.
-                        hacc_min = cute.make_fragment(next_n, cutlass.Float32)
-                        hacc_max = cute.make_fragment(next_n, cutlass.Float32)
-                        hacc_sum = cute.make_fragment(next_n, cutlass.Float32)
-                        hacc_cnt = cute.make_fragment(next_n, cutlass.Int32)
+                        hacc_min = cute.make_rmem_tensor(next_n, cutlass.Float32)
+                        hacc_max = cute.make_rmem_tensor(next_n, cutlass.Float32)
+                        hacc_sum = cute.make_rmem_tensor(next_n, cutlass.Float32)
+                        hacc_cnt = cute.make_rmem_tensor(next_n, cutlass.Int32)
                         for _t in cutlass.range_constexpr(next_n):
                             hacc_min[_t] = cutlass.Float32(_META_FLT_MAX)
                             hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX)

From f24be00b1997a5db1fb82bda0ed5bd244aab6276 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 21 Aug 2026 00:17:51 -0700
Subject: [PATCH 106/117] [None][chore] restore the .claude subtree clobbered
 in the main merge

A staged-file cleanup during the 2a0c11f044 merge reset the .claude
index entries to the pre-merge state, silently reverting main's
updates for that subtree (85 files). Restore it to the merged main's
(c30f71365c) content; no code paths affected.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .claude/agents/exec-local-slurm.md            |   31 +
 .claude/agents/exec-remote-slurm.md           |   40 +
 .claude/agents/trtllm-test-specialist.md      |   18 +
 .claude/skills/exec-env-check/SKILL.md        |  142 ++
 .claude/skills/exec-local-docker/SKILL.md     |  121 ++
 .claude/skills/exec-local-slurm/SKILL.md      |  443 ++++++
 .claude/skills/exec-remote-slurm/SKILL.md     |  970 ++++++++++++++
 .../skills/exec-remote-slurm/scripts/build.sh |   42 +
 .../exec-remote-slurm/scripts/build.slurm     |   41 +
 .claude/skills/exec-slurm-compile/SKILL.md    |   14 +-
 .../exec-slurm-compile/scripts/compile.slurm  |    6 +-
 .../scripts/submit_compile.sh                 |    6 +-
 .../references/api-core.md                    |    3 +-
 .../references/concepts-tensors.md            |    3 +-
 .claude/skills/perf-analysis/SKILL.md         |    6 +
 .../perf-optimization-casebook/SKILL.md       |  293 ++++
 .../data/aliases.yaml                         |   53 +
 .../data/patterns.yaml                        |  406 ++++++
 .../perf-optimization-casebook/data/tags.yaml |  116 ++
 .../references/case-template.md               |  211 +++
 .../references/communication/deepep.md        |   37 +
 .../references/communication/index.md         |   74 ++
 .../communication/low-precision-dispatch.md   |   38 +
 .../communication/mnnvl-twoshot-allreduce.md  |   35 +
 .../shape-aware-allreduce-autotune.md         |   33 +
 .../userbuffers-symmetric-memory.md           |   37 +
 .../fold-scale-swizzle-into-kernel.md         |   33 +
 .../kernel-and-fusion/fp8-mla-kv-cache.md     |   38 +
 .../kernel-and-fusion/fuse-add-norm-quant.md  |   39 +
 .../kernel-and-fusion/fuse-ar-epilogue.md     |   38 +
 .../fuse-datamovement-into-quantize.md        |   85 ++
 .../fuse-moe-routing-kernel.md                |  115 ++
 .../kernel-and-fusion/fuse-qk-norm-rope.md    |   36 +
 .../hw-matched-lowprec-moe-gemm.md            |   36 +
 .../references/kernel-and-fusion/index.md     |  158 +++
 .../mega-fuse-moe-deepgemm.md                 |   41 +
 .../ranking-only-precision-tf32.md            |   87 ++
 .../reevaluate-fusion-boundary-per-dtype.md   |   87 ++
 .../kernel-and-fusion/relax-tl-constexpr.md   |   69 +
 .../sparse-mla-topk-attention.md              |  117 ++
 .../specialize-topk-selection-kernel.md       |  103 ++
 .../split-mla-reduction-kernel.md             |   35 +
 .../kernel-and-fusion/triton-to-cpp-op.md     |   81 ++
 .../trtllm-gen-fp4-moe-backend.md             |   36 +
 .../runtime-execution/attention-dp-padding.md |   33 +
 .../auxiliary-cache-in-kv-manager.md          |   94 ++
 .../cache-step-invariant-per-layer.md         |   85 ++
 .../chunked-prefill-aligned-auxiliary.md      |   86 ++
 .../runtime-execution/cuda-graph-padding.md   |   33 +
 .../free-mla-intermediates.md                 |   33 +
 .../hoist-torch-compile-closures.md           |   77 ++
 .../references/runtime-execution/index.md     |  165 +++
 .../runtime-execution/mla-kv-cache-reuse.md   |   35 +
 .../move-bookkeeping-into-cpp-op.md           |   36 +
 .../multi-stream-shared-routed-expert.md      |   37 +
 .../overlap-mla-rope-uk-bgemm.md              |   35 +
 .../runtime-execution/overlap-online-eplb.md  |   35 +
 .../runtime-execution/overlap-scheduler.md    |   41 +
 .../references/runtime-execution/pdl.md       |   35 +
 .../runtime-execution/piecewise-cuda-graph.md |   35 +
 .../pybind-wrapper-pure-python.md             |   93 ++
 .../relaxed-mtp-acceptance.md                 |   37 +
 .../skip-sparse-path-when-degenerate.md       |   90 ++
 .../split-custom-op-for-piecewise-capture.md  |   93 ++
 .../runtime-execution/two-model-mtp-eagle.md  |   35 +
 .claude/skills/perf-optimization/SKILL.md     |   61 +-
 .claude/skills/trtllm-case-executor/SKILL.md  |  440 ++++++
 .../scripts/detect_slurm_env.sh               |  449 +++++++
 .../trtllm-model-onboard-multimodal/SKILL.md  |   14 +-
 .../trtllm-test-script-builder/SKILL.md       |  385 ++++++
 .../references/trtllm_test_template.md        |  520 ++++++++
 .../scripts/build_slurm_script.py             |  643 +++++++++
 .../scripts/slurm_run_custom.sh               |   33 +
 .../skills/trtllm-test-specialist/SKILL.md    |  681 ++++++++++
 .../references/agg_config_template.yaml       |  257 ++++
 .../references/benchmark_config_template.yaml |  133 ++
 .../references/disagg_config_template.yaml    |   63 +
 .../references/smoke_test_config_template.yml |   19 +
 .../references/test_config_template.yaml      |  227 ++++
 .../trtllm_test_fix_recommendations.md        |   14 +
 .../scripts/build_test_command.py             |  768 +++++++++++
 .../scripts/extract_test_markers.py           |  200 +++
 .../scripts/generate_benchmark_config.py      | 1184 +++++++++++++++++
 .../scripts/generate_report.py                |  297 +++++
 .../scripts/parse_config.py                   |  237 ++++
 85 files changed, 12247 insertions(+), 74 deletions(-)
 create mode 100644 .claude/agents/exec-local-slurm.md
 create mode 100644 .claude/agents/exec-remote-slurm.md
 create mode 100644 .claude/agents/trtllm-test-specialist.md
 create mode 100644 .claude/skills/exec-env-check/SKILL.md
 create mode 100644 .claude/skills/exec-local-docker/SKILL.md
 create mode 100644 .claude/skills/exec-local-slurm/SKILL.md
 create mode 100644 .claude/skills/exec-remote-slurm/SKILL.md
 create mode 100644 .claude/skills/exec-remote-slurm/scripts/build.sh
 create mode 100644 .claude/skills/exec-remote-slurm/scripts/build.slurm
 create mode 100644 .claude/skills/perf-optimization-casebook/SKILL.md
 create mode 100644 .claude/skills/perf-optimization-casebook/data/aliases.yaml
 create mode 100644 .claude/skills/perf-optimization-casebook/data/patterns.yaml
 create mode 100644 .claude/skills/perf-optimization-casebook/data/tags.yaml
 create mode 100644 .claude/skills/perf-optimization-casebook/references/case-template.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/deepep.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/index.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/mnnvl-twoshot-allreduce.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/shape-aware-allreduce-autotune.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/userbuffers-symmetric-memory.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fp8-mla-kv-cache.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-add-norm-quant.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-ar-epilogue.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-datamovement-into-quantize.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-qk-norm-rope.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/index.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/ranking-only-precision-tf32.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/reevaluate-fusion-boundary-per-dtype.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/relax-tl-constexpr.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/sparse-mla-topk-attention.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/specialize-topk-selection-kernel.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/split-mla-reduction-kernel.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/triton-to-cpp-op.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/trtllm-gen-fp4-moe-backend.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/auxiliary-cache-in-kv-manager.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/cache-step-invariant-per-layer.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/chunked-prefill-aligned-auxiliary.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/cuda-graph-padding.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/free-mla-intermediates.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/hoist-torch-compile-closures.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/index.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/mla-kv-cache-reuse.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/move-bookkeeping-into-cpp-op.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/multi-stream-shared-routed-expert.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-mla-rope-uk-bgemm.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-scheduler.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/pdl.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/piecewise-cuda-graph.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/pybind-wrapper-pure-python.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/relaxed-mtp-acceptance.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/skip-sparse-path-when-degenerate.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/split-custom-op-for-piecewise-capture.md
 create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/two-model-mtp-eagle.md
 create mode 100644 .claude/skills/trtllm-case-executor/SKILL.md
 create mode 100755 .claude/skills/trtllm-case-executor/scripts/detect_slurm_env.sh
 create mode 100644 .claude/skills/trtllm-test-script-builder/SKILL.md
 create mode 100644 .claude/skills/trtllm-test-script-builder/references/trtllm_test_template.md
 create mode 100755 .claude/skills/trtllm-test-script-builder/scripts/build_slurm_script.py
 create mode 100755 .claude/skills/trtllm-test-script-builder/scripts/slurm_run_custom.sh
 create mode 100644 .claude/skills/trtllm-test-specialist/SKILL.md
 create mode 100644 .claude/skills/trtllm-test-specialist/references/agg_config_template.yaml
 create mode 100644 .claude/skills/trtllm-test-specialist/references/benchmark_config_template.yaml
 create mode 100644 .claude/skills/trtllm-test-specialist/references/disagg_config_template.yaml
 create mode 100644 .claude/skills/trtllm-test-specialist/references/smoke_test_config_template.yml
 create mode 100644 .claude/skills/trtllm-test-specialist/references/test_config_template.yaml
 create mode 100644 .claude/skills/trtllm-test-specialist/references/trtllm_test_fix_recommendations.md
 create mode 100644 .claude/skills/trtllm-test-specialist/scripts/build_test_command.py
 create mode 100644 .claude/skills/trtllm-test-specialist/scripts/extract_test_markers.py
 create mode 100644 .claude/skills/trtllm-test-specialist/scripts/generate_benchmark_config.py
 create mode 100755 .claude/skills/trtllm-test-specialist/scripts/generate_report.py
 create mode 100644 .claude/skills/trtllm-test-specialist/scripts/parse_config.py

diff --git a/.claude/agents/exec-local-slurm.md b/.claude/agents/exec-local-slurm.md
new file mode 100644
index 000000000000..795f5751441c
--- /dev/null
+++ b/.claude/agents/exec-local-slurm.md
@@ -0,0 +1,31 @@
+---
+name: exec-local-slurm
+description: >
+  Execute a TensorRT-LLM workload on a local Slurm cluster. Supports persistent
+  allocation (allocate once via nohup salloc, reuse across runs) and one-shot
+  sbatch. Workflow-agnostic — handles pytest, eval, benchmark, and custom
+  scripts identically. The orchestrator (typically trtllm-case-executor) writes
+  a job spec to <work_dir>/job_spec.json and invokes this agent to run it.
+tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"]
+model: sonnet
+license: Apache-2.0
+---
+
+You are the local Slurm executor agent. Load the `exec-local-slurm` skill (`trtllm-agent-toolkit:exec-local-slurm`) and follow its procedure exactly. The skill is the single source of truth for the execution flow; this agent file only adds the contract with the caller and the invariants that must hold across every run.
+
+## Input
+
+The caller passes a path to a job spec — typically `<work_dir>/job_spec.json` — plus a short summary of the fields used in this run. **Read `job_spec.json` first.** The skill's "Input (from orchestrator prompt)" section enumerates every field it consumes (`script_path`, `work_dir`, `model_name`, `workflow_type`, `success_patterns`, `failure_patterns`, `log_file_pattern`, `monitor_timeout_seconds`, `persistent_mode`, `release_allocation`, `alloc_time_limit`, `docker_image`, `container_name`, `container_mounts`, `repo_root`, `slurm_params`).
+
+Do not re-derive any field that `trtllm-case-executor` already wrote into the spec.
+
+## Invariants
+
+- **Hang detection.** Poll the log periodically; on a case-insensitive `hang detected` match, kill the process group and report `HANG_DETECTED`. Implementation lives in the skill.
+- **Wall-clock limit.** Honor `monitor_timeout_seconds` (default `3600`). On timeout, kill and report `TIMEOUT`.
+- **Persistent allocation lifecycle.** Default `persistent_mode=true`; reuse an existing allocation when present and valid. Only release when `release_allocation=true` is explicitly set — never auto-release.
+- **Single source of truth.** `node_count`, `job_name`, `container_image`, and `slurm_params` come from `job_spec.json`. Never recompute them or re-grep `current_image_tags.properties`.
+
+## Output
+
+Return a single report to the caller with: task type, status (`PASSED` / `FAILED` / `TIMEOUT` / `HANG_DETECTED` / `OUT_OF_MEMORY` / `CANCELLED` / `ERROR` / `BUILD_FAILED`), Slurm job id (and allocation id when persistent), log file path, summary, and any error excerpts (last ~100 lines on build/job failure). Do not perform follow-up actions beyond what the skill prescribes.
diff --git a/.claude/agents/exec-remote-slurm.md b/.claude/agents/exec-remote-slurm.md
new file mode 100644
index 000000000000..496493d0aa56
--- /dev/null
+++ b/.claude/agents/exec-remote-slurm.md
@@ -0,0 +1,40 @@
+---
+name: exec-remote-slurm
+description: >
+  Execute a TensorRT-LLM workload on a remote Slurm cluster via SSH. Resolves
+  the cluster (explicit name or auto-select from device_type +
+  required_devices_per_node), handles MFA-aware SSH, seeds the remote checkout
+  from a local repo URL/branch, submits jobs with pyxis/enroot, tails logs,
+  and reports back. The orchestrator (typically trtllm-case-executor) writes a
+  job spec to <work_dir>/job_spec.json and invokes this agent to run it.
+tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"]
+model: sonnet
+license: Apache-2.0
+---
+
+You are the remote Slurm executor agent. Load the `exec-remote-slurm` skill (`trtllm-agent-toolkit:exec-remote-slurm`) and follow its procedure exactly. The skill is the single source of truth for the execution flow; this agent file only adds the contract with the caller and the invariants that must hold across every run.
+
+## Input
+
+The caller passes a path to a job spec — typically `<work_dir>/job_spec.json` — plus a short summary of the fields used in this run. **Read `job_spec.json` first.** The skill's "Input" section enumerates every field it consumes (`script_path`, `script_name`, `work_dir`, `model_name`, `workflow_type`, `success_patterns`, `failure_patterns`, `log_file_pattern`, `slurm_cluster`, `ssh_host`, `slurm_user`, `remote_cwd`, `remote_work_dir`, `slurm_password`, `extra_files`, `repo_url`, `repo_branch`, `device_type`, `total_required_devices`, `required_devices_per_node`, `container_image`, `node_count`, `job_name`, `monitor_timeout_seconds`).
+
+Do not re-derive any field that `trtllm-case-executor` already wrote into the spec.
+
+## Cluster resolution
+
+**Precondition — optional dependency.** Before either bullet below, check whether the `skills/internal-env-info/` directory exists in the toolkit. If it does **not**, do not attempt to load any reference file or invoke the skill. Proceed using only the cluster fields the orchestrator wrote into `job_spec.json` (`ssh_host`, `slurm_user`, `remote_cwd`, `partition`, `account`, `container_image`, `mounts`, `gpus_per_node`, `mfa_style`). If a required field is also absent, stop and ask the user to supply it — do **not** report the missing skill as an error.
+
+- **Explicit cluster** (`slurm_cluster` is set) → invoke the `internal-env-info` skill in single-cluster mode to fetch per-cluster info (`mfa_style`, `default_models_repo`, `default_user_root_dir`, `gpus_per_node`). Connection fields (`ssh_host`, `slurm_user`, `remote_cwd`, `account`, `partition`, `mounts`, `container_image`) come from `job_spec.json` — there is no per-cluster connection-config file to parse.
+- **Auto-select** (only hardware constraints are present) → invoke the `internal-env-info` skill in constraint-based mode, passing `device_type` and `required_devices_per_node` (and optionally `total_required_devices`). Never re-implement the constraint filter.
+
+## Invariants
+
+- **`mfa_style` decides the SSH path.** `false` → direct; `true` → MFA flow; `null` → probe direct first, then fall back, and ask the user to update the cluster reference. Do **not** probe-and-fall-back on the SSH error string when `mfa_style` is known.
+- **Hang detection + `monitor_timeout_seconds`.** Same semantics as local execution; implementation lives in the skill.
+- **Pass-through fields.** `container_image`, `node_count`, and `job_name` come from `job_spec.json` and are used verbatim. Apply transport-specific URL rewrites (e.g., enroot `/` → `#`) at use-time. Never re-grep `current_image_tags.properties` or reconstruct `job_name` from `account` / `subproject` / `detail`.
+- **Remote repo bootstrap.** Use `repo_url` and `repo_branch` from `job_spec.json` to ensure the remote checkout matches the local one before submission.
+- **`node_count` is authoritative.** Use it directly as `--nodes`; never recompute from totals.
+
+## Output
+
+Return a single report to the caller with: task type, status (`PASSED` / `FAILED` / `TIMEOUT` / `HANG_DETECTED` / `OUT_OF_MEMORY` / `CANCELLED` / `ERROR` / `BUILD_FAILED`), remote Slurm job id, remote log path (and a local copy when synced back), summary, and any error excerpts (last ~100 lines on build/job failure). Do not perform follow-up actions beyond what the skill prescribes.
diff --git a/.claude/agents/trtllm-test-specialist.md b/.claude/agents/trtllm-test-specialist.md
new file mode 100644
index 000000000000..0952a24f01c2
--- /dev/null
+++ b/.claude/agents/trtllm-test-specialist.md
@@ -0,0 +1,18 @@
+---
+name: trtllm-test-specialist
+description: >
+  Runs model-level and module-level tests for TensorRT-LLM. Classifies the test
+  scope (module test or model test), builds the appropriate test commands, and
+  delegates execution to trtllm-case-executor. Supports functionality/smoke
+  tests, benchmarks, and evaluations. Writes structured test reports to a
+  caller-specified path.
+tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"]
+license: Apache-2.0
+---
+
+Role: dispatch TRT-LLM model-level and module-level test requests.
+
+Load the `trtllm-agent-toolkit:trtllm-test-specialist` skill, pass the caller's parameters through verbatim, and return its result.
+
+- If the caller supplies `report_file`, write the report to that exact path — do not substitute the skill's default (`./<MODEL_NAME>-auto-test-report.md`) or invent your own.
+- Return the skill's status and final report path verbatim; do not re-summarize or rename.
diff --git a/.claude/skills/exec-env-check/SKILL.md b/.claude/skills/exec-env-check/SKILL.md
new file mode 100644
index 000000000000..3b91216a74b1
--- /dev/null
+++ b/.claude/skills/exec-env-check/SKILL.md
@@ -0,0 +1,142 @@
+---
+name: exec-env-check
+description: >-
+  Check the local execution environment for GPU availability, Docker support,
+  and Slurm access. Returns the execution scenario (`satisfied, local, docker`,
+  `satisfied, local, direct`, `satisfied, slurm, local`, or `not_satisfied`),
+  the number of available GPUs, and the GPU type. On Slurm login nodes
+  without local GPUs, the cluster is identified by delegating the hostname
+  to internal-env-info (hostname-based mode), which owns the
+  hostname → cluster_name patterns; GPU type and gpus_per_node then come
+  from that skill's reference files. If internal-env-info is not
+  installed, the scenario falls back to `not_satisfied` without probing
+  compute nodes via srun.
+tags: [infrastructure, slurm, environment]
+license: Apache-2.0
+metadata:
+  author: NVIDIA Corporation
+---
+
+# TensorRT-LLM Environment Check
+
+Detect whether the current machine can run a GPU workload locally (Docker) or via Slurm, and report hardware details.
+
+## Input
+
+| Field | Description | Required |
+|-------|-------------|----------|
+| `required_devices` | Minimum number of GPUs needed | Yes |
+| `account` | Slurm account (unused for GPU probing, kept for compatibility) | No |
+
+## Procedure
+
+### 1. Check local GPUs
+
+```bash
+timeout 5 nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null
+```
+
+If this succeeds:
+- Count the number of GPU lines → `available_gpus`
+- Extract the GPU name from the first line → `device_type` (e.g., `NVIDIA B200`, `NVIDIA H100 80GB HBM3`)
+- Normalize `device_type`: strip `NVIDIA ` prefix and trailing memory info to get the short name (e.g., `B200`, `H100`, `A100`, `L40S`, `RTX 6000`)
+
+If `nvidia-smi` fails or returns no GPUs → `available_gpus = 0`, `device_type = null`.
+
+### 2. Check if local GPUs are sufficient
+
+If `available_gpus >= required_devices`, continue to step 2a to determine whether Docker is available on this host.
+
+### 2a. Check Docker availability
+
+```bash
+command -v docker >/dev/null 2>&1 && timeout 3 docker info >/dev/null 2>&1
+```
+
+- If the command succeeds (Docker CLI exists **and** the daemon responds) → **Result**: `satisfied, local, docker`
+- Otherwise (Docker CLI missing, daemon not running, or permission denied) → **Result**: `satisfied, local, direct`
+
+Include `available_gpus` in the result. Do NOT include `device_type` — local execution does not need it.
+
+### 3. Check Slurm availability
+
+If local GPUs are insufficient (or as additional detection), check for Slurm:
+
+```bash
+which squeue 2>/dev/null && squeue --version 2>/dev/null
+```
+
+If Slurm is NOT available → go to step 5.
+
+### 4. Resolve GPU type
+
+**Optional dependency.** Before doing anything else in this step, check whether `skills/internal-env-info/` exists in the toolkit. If it does **not**, skip this step entirely and return `scenario: not_satisfied` with `available_gpus = 0`, `device_type = null`, `gpus_per_node = null`, `cluster_name = null`, `default_models_repo = null`, and `default_user_root_dir = null`. Do **not** report this as an error — the skill is an internal-only dependency.
+
+When Slurm is available but local `nvidia-smi` returned no GPUs or `device_type` is null (login nodes typically have no GPUs), capture the hostname and delegate cluster identification to `internal-env-info`:
+
+```bash
+hostname -f 2>/dev/null || hostname
+```
+
+- Pass the captured hostname to `internal-env-info` in **hostname-based mode**. That skill owns the NVIDIA-internal login-host patterns and the hostname → `<cluster>` mapping; do **not** parse the hostname or hard-code any cluster identifier here.
+- It returns the standard output template (`cluster_name`, `device_type`, `gpus_per_node`, `default_models_repo`, `default_user_root_dir`) plus supplementary field (`mfa_style`).
+- Set `available_gpus` = `gpus_per_node` (if resolved).
+- Set `cluster_name` = the value returned (a placeholder `<cluster>` token in this skill). The orchestrator uses this to fetch per-cluster info (`mfa_style`, `default_models_repo`, `default_user_root_dir`, `gpus_per_node`) from `internal-env-info`; connection fields (`mounts`, `ssh_host`, `partition`, etc.) come from caller-supplied inputs in `job_spec.json` (with `internal-env-info` default values / ask-the-user fallbacks).
+- Set `default_user_root_dir` = the user root directory returned (with `<user_name>` substituted with the actual SLURM username). Set to `null` if not found.
+- If `internal-env-info` returns `null` for `cluster_name` (no pattern matched), set `device_type = null`, `cluster_name = null`, `default_user_root_dir = null` and fall through to Step 5 (`not_satisfied`).
+
+**Result**: `satisfied, slurm, local` with `device_type`, `gpus_per_node`, `cluster_name`, and `default_user_root_dir`.
+
+### 5. Not satisfied
+
+If neither local GPUs nor Slurm is available:
+- **Result**: `not_satisfied`
+
+## Output
+
+Return a single structured result:
+
+```
+scenario: <satisfied, local, docker | satisfied, local, direct | satisfied, slurm, local | not_satisfied>
+available_gpus: <N>
+device_type: <short GPU/device name or null>
+gpus_per_node: <N or null>
+cluster_name: <cluster name or null>
+default_models_repo: <host path to llm-models directory, or null>
+default_user_root_dir: <per-user root directory on cluster storage, or null>
+```
+
+**Examples:**
+
+```
+scenario: satisfied, local, docker
+available_gpus: 4
+```
+
+```
+scenario: satisfied, local, direct
+available_gpus: 4
+```
+
+```
+scenario: satisfied, slurm, local
+available_gpus: 4
+device_type: B200
+gpus_per_node: 4
+cluster_name: <cluster>
+default_user_root_dir: <default_user_root_dir>/<user_name>
+```
+
+```
+scenario: not_satisfied
+available_gpus: 0
+device_type: null
+cluster_name: null
+```
+
+## Rules
+
+- Never install drivers or modify the system
+- If `nvidia-smi` hangs, use a 5-second timeout: `timeout 5 nvidia-smi ...`
+- GPU type is derived from the hostname by matching the cluster name via the `internal-env-info` skill when it is installed — no `srun` allocation needed. If that skill is absent, the scenario falls back to `not_satisfied` (see Step 4); do not error.
+- Report the GPU type exactly as found in the mapping (do not guess or fabricate)
diff --git a/.claude/skills/exec-local-docker/SKILL.md b/.claude/skills/exec-local-docker/SKILL.md
new file mode 100644
index 000000000000..cd73eed6e043
--- /dev/null
+++ b/.claude/skills/exec-local-docker/SKILL.md
@@ -0,0 +1,121 @@
+---
+name: exec-local-docker
+description: >-
+  Execute a TensorRT-LLM workload locally in Docker. Runs a fully-resolved
+  Docker command in background, monitors completion, reads logs, and reports
+  results. Workflow-agnostic — does not need to know if the workload is pytest,
+  eval, benchmark, or a custom script.
+tags: [docker, execution, infrastructure]
+license: Apache-2.0
+metadata:
+  author: NVIDIA Corporation
+---
+
+# Local Docker Executor
+
+Run a Docker command locally, monitor it, and report results.
+
+## Input (from orchestrator prompt)
+
+The orchestrator passes these fields in the skill prompt:
+
+| Field | Description |
+|-------|-------------|
+| `docker_cmd` | Complete `docker run` command string, ready to execute |
+| `work_dir` | Local work directory for logs and artifacts |
+| `log_file` | Full path to the log file (output redirected here) |
+| `model_name` | Short model name for reporting |
+| `workflow_type` | `pytest`, `eval`, `custom`, or `benchmark` — for output parsing hints |
+| `success_patterns` | Comma-separated patterns indicating success (e.g., `passed,accuracy:`) |
+| `failure_patterns` | Comma-separated patterns indicating failure (e.g., `FAILED,Error,AssertionError`) |
+
+## Procedure
+
+### Step 0: Resolve Image and Build (when `build_project=true`)
+
+This executor owns image selection and the build for the local Docker target. Skip the entire step when `build_project=false`.
+
+1. **Detect the target GPU type.** Use `gpu_type` from `job_spec.json` if upstream env-check resolved it; otherwise probe locally with `nvidia-smi --query-gpu=name --format=csv,noheader | head -1`.
+2. **Detect the host CPU arch** with `uname -m` (`x86_64` or `aarch64`).
+3. **Resolve the container image.** Read `<repo_root>/jenkins/current_image_tags.properties` and pick the tag whose CPU-arch flavor matches the host. If the orchestrator already passed a `container_image` field in the job spec, use that and skip the lookup.
+4. **Map GPU → build arch (`-a` flag):** `H100`/`H200` → `90-real`; `B200`/`GB200`/`B300`/`GB300` → `100-real`; `A100` → `80-real`; `L40S` → `89-real`. Default `100-real` when the GPU is unknown.
+5. **Compile.** Invoke the `exec-local-compile` skill with `repo_dir=<repo_root>`, `image=<resolved tag>`, `arch=<arch>`. Wait for completion.
+6. **On failure**, do not launch the workload. Report `BUILD_FAILED` with the last 100 lines of the compile log.
+
+`build_project`, `gpu_type`, `repo_root`, and (optionally) `container_image` come from `job_spec.json`.
+
+### Step 1: Launch
+
+Run the Docker command in background using `run_in_background`:
+
+```bash
+<docker_cmd> 2>&1 | tee <log_file>
+```
+
+Report to the orchestrator: "Launched locally, log at `<log_file>`"
+
+### Step 2: Monitor for Hangs
+
+While waiting for the background process to complete, actively monitor the log
+file for hang indicators. Launch a monitoring loop using `run_in_background`:
+
+```bash
+while true; do
+  sleep 60
+  if [ -f "<log_file>" ] && grep -qi "hang detected" "<log_file>"; then
+    echo "HANG_DETECTED: Found 'hang detected' in log file"
+    docker ps --filter "ancestor=<image>" -q | xargs -r docker kill 2>/dev/null
+    exit 1
+  fi
+done
+```
+
+- If the monitor detects a hang, it kills the Docker container and exits with
+  code 1. The main background process will also terminate.
+- When the main process completes normally (background notification received),
+  kill the monitoring loop (it is no longer needed).
+- If a hang is detected, skip to Step 4 and report `HANG_DETECTED` status
+  instead of proceeding to normal result collection.
+
+### Step 3: Wait for Completion
+
+The Bash tool's `run_in_background` will notify when the process finishes
+(either normally or because the container was killed by the hang monitor).
+
+### Step 4: Read Results
+
+On completion:
+
+1. **Read exit code** from the background command result.
+2. **Read the last 100 lines** of `<log_file>` using the Read tool.
+3. **If exit code != 0**, also read the first 50 lines to catch early errors (import failures, setup crashes).
+4. **Search for patterns**:
+   - Grep `<log_file>` for each `success_patterns` entry
+   - Grep `<log_file>` for each `failure_patterns` entry
+
+### Step 5: Report
+
+Return a structured result:
+
+```
+Status: PASSED | FAILED | ERROR | HANG_DETECTED
+Exit code: <N>
+Log file: <log_file>
+Work directory: <work_dir>
+Summary: <relevant output lines — pytest summary, accuracy score, or last few lines>
+Errors: <if failed, the relevant error output>
+```
+
+### Output Parsing by Workflow Type
+
+- **pytest**: Look for `X passed, Y failed in Zs` summary line
+- **eval**: Look for `accuracy:` or `score:` lines; check for `Expected accuracy >= X, but got Y` assertion
+- **custom**: No specific patterns — report last 10 lines of output
+- **benchmark**: Look for throughput/latency numbers
+
+## Rules
+
+- Never run the Docker command in foreground — always use `run_in_background`
+- Never `cat` the full log file — use Read with offset/limit or tail
+- If the Docker command fails immediately (exit code within seconds), check if the image exists locally
+- Report results even if the log file is empty (container may have failed to start)
diff --git a/.claude/skills/exec-local-slurm/SKILL.md b/.claude/skills/exec-local-slurm/SKILL.md
new file mode 100644
index 000000000000..ecfd3a860fbc
--- /dev/null
+++ b/.claude/skills/exec-local-slurm/SKILL.md
@@ -0,0 +1,443 @@
+---
+name: exec-local-slurm
+description: >-
+  Submit and monitor a Slurm job on a local cluster. Supports two modes:
+  (1) Persistent allocation (default) — allocates nodes once via nohup salloc,
+  imports the container once, installs once, and reuses across runs by setting
+  SLURM env vars and running the sbatch script via bash. (2) One-shot sbatch —
+  submits a fully-generated Slurm script via sbatch, polls job status, reads
+  logs on completion, and reports results. Workflow-agnostic — handles pytest,
+  eval, benchmark, and custom scripts identically.
+tags: [slurm, execution, infrastructure]
+license: Apache-2.0
+metadata:
+  author: NVIDIA Corporation
+---
+
+# Local Slurm Executor
+
+Submit a Slurm job locally, monitor it, and report results. Uses persistent
+allocation by default to eliminate queue wait, container import, and install
+overhead on repeated runs.
+
+## Input (from orchestrator prompt)
+
+The orchestrator passes these fields in the skill prompt:
+
+| Field | Description |
+|-------|-------------|
+| `script_path` | Full local path to the generated `.slurm` script |
+| `work_dir` | Local work directory for logs and artifacts |
+| `model_name` | Short model name for reporting |
+| `workflow_type` | `pytest`, `eval`, `custom`, or `benchmark` — for output parsing hints |
+| `success_patterns` | Comma-separated patterns indicating success |
+| `failure_patterns` | Comma-separated patterns indicating failure |
+| `log_file_pattern` | Log filename pattern with `%j` placeholder (e.g., `llama_auto_test_%j.out`) |
+| `persistent_mode` | Default: `true` for all local slurm workflows. Set to `false` to force one-shot sbatch (opt-out). |
+| `release_allocation` | `true` to release the current allocation and stop. Only set when the user explicitly says no more jobs are needed. Default: `false`. Never auto-release. |
+| `alloc_time_limit` | Walltime for the persistent allocation. Default: `04:00:00`. |
+| `docker_image` | Container image for persistent container import. From `job_spec.json`. |
+| `container_name` | Container name used in the `.slurm` script (e.g., `llama_auto_test`). Must match exactly. From `job_spec.json`. |
+| `container_mounts` | Comma-separated mount mappings. From `job_spec.json`. |
+| `repo_root` | Repo root path (for locating state file and project path). |
+| `slurm_params` | Slurm parameters object: `partition`, `account`, `nodes`, `ntasks`, `ntasks_per_node`, `gpus_per_node`. From `job_spec.json`. |
+
+## Procedure
+
+### Pre-step: Resolve Image and Build (when `build_project=true`)
+
+This executor owns image selection and the build for the local SLURM cluster. Skip the entire pre-step when `build_project=false`.
+
+1. **Read the cluster GPU type** from `device_type` in `job_spec.json` (resolved upstream by env-check / internal-env-info). If absent, probe from any compute node via `srun -p <partition> --ntasks-per-node=1 nvidia-smi --query-gpu=name --format=csv,noheader | head -1`. If `skills/internal-env-info/` is not installed, skip the upstream lookup silently and rely on the srun probe — do not report the missing skill as an error.
+2. **Determine the partition's CPU arch.** Look up `job_spec.slurm_env.partitions[]` for the entry where `name == partition` and read its `arch` (`x86_64` / `aarch64`). Case-executor's Step 2.5 already detected this — do **not** run `scontrol show node` here. If `slurm_env` is absent or the entry's `arch` is null (case-executor ran without SLURM tools), fall back to inferring from `device_type`: Grace-based parts (`GB*` / `GH*`) → `aarch64`; everything else → `x86_64`.
+3. **Resolve the container image.** If `docker_image` is already set in `job_spec.json`, use it directly. Otherwise read `<repo_root>/jenkins/current_image_tags.properties` and pick the tag matching the partition's CPU arch.
+4. **Map GPU → build arch (`-a` flag):** `H100`/`H200` → `90-real`; `B200`/`GB200`/`B300`/`GB300` → `100-real`; `A100` → `80-real`; `L40S` → `89-real`. Default `100-real`.
+5. **Compile.** Invoke the `exec-slurm-compile` skill with `repo_dir=<repo_root>`, `partition`, `account`, `container_image=<resolved tag>`, `user_root_dir`, `arch`. Wait for completion.
+6. **On failure**, do not allocate or run the workload. Report `BUILD_FAILED` with the last 100 lines of the build log.
+
+`build_project`, `device_type`, `repo_root`, `partition`, `account`, `user_root_dir`, and (optionally) `docker_image` come from `job_spec.json`.
+
+### Step 0: Allocation Management
+
+This step runs before any execution. It determines whether to reuse an existing
+persistent allocation, create a new one, or fall through to one-shot sbatch.
+
+#### Step 0A — Release mode
+
+If `release_allocation=true`:
+
+1. Read `<repo_root>/work_dirs/.slurm_alloc.json`
+2. If file exists and `job_id` is present:
+   ```bash
+   scancel <job_id>
+   ```
+3. Delete the state file and `.salloc.log`
+4. Report: "Allocation <job_id> released." Stop.
+
+#### Step 0B — One-shot mode
+
+If `persistent_mode=false`, skip entirely to Step 1 (One-Shot Path).
+
+#### Step 0C — Validate existing allocation
+
+Read `<repo_root>/work_dirs/.slurm_alloc.json`.
+
+- If state file **does not exist** → go to Step 0D.
+- If state file exists, validate the allocation:
+  ```bash
+  squeue -j <job_id> -h -o "%T %L"
+  ```
+  - **RUNNING + remaining > 5 min + params compatible** → **reuse** (skip to Step 0F)
+  - **RUNNING + remaining <= 5 min** → warn user ("Allocation expiring soon"), `scancel <job_id>`, delete state file, go to Step 0D
+  - **RUNNING + params incompatible** → `scancel <job_id>`, delete state file, go to Step 0D
+  - **Empty output / error** (job gone) → stale state file, delete it, go to Step 0D
+
+**Params compatibility check:**
+- `partition` must match
+- `nodes` in state must be **>=** requested nodes (a 2-node allocation can serve 1-node jobs)
+- `docker_image` must match
+
+#### Step 0D — Justify and prepare allocation
+
+Before allocating, ensure no orphaned allocations exist and log the reason:
+
+1. Check for existing persistent jobs:
+   ```bash
+   squeue -u $(whoami) -h -o "%i %T %j" --name=<account>-trtllm.persistent
+   ```
+2. If found → orphan (state file was missing/corrupt). Cancel it:
+   ```bash
+   scancel <orphan_job_id>
+   ```
+   Log: "Released orphaned allocation <orphan_job_id> — state file was missing."
+3. Log justification: "No reusable allocation found. Allocating <nodes> node(s) on partition <partition> for <alloc_time_limit>."
+
+#### Step 0E — Allocate new
+
+Convert `alloc_time_limit` from `HH:MM:SS` to seconds (e.g., `04:00:00` → `14400`).
+
+```bash
+mkdir -p <repo_root>/work_dirs
+nohup salloc --partition=<partition> --account=<account> \
+  --nodes=<nodes> --time=<alloc_time_limit> \
+  --job-name=<account>-trtllm.persistent \
+  sleep <alloc_time_limit_seconds> \
+  > <repo_root>/work_dirs/.salloc.log 2>&1 &
+```
+
+Retrieve job ID:
+```bash
+squeue -u $(whoami) -h -o "%i" --name=<account>-trtllm.persistent
+```
+
+If no job appears after 10 seconds, read `<repo_root>/work_dirs/.salloc.log`
+for errors (bad partition, invalid account, etc.) and report to user.
+
+Otherwise, poll until RUNNING (every 10s, max 60 polls). Once RUNNING, get
+the nodelist:
+```bash
+squeue -j <JOB_ID> -h -o "%N"
+```
+
+**Import container** on all nodes (using the **same container name** as the
+`.slurm` script — critical for pyxis reuse):
+```bash
+srun --jobid=<JOB_ID> -N <nodes> --ntasks-per-node=1 \
+  --container-image=<docker_image> \
+  --container-name=<container_name> true
+```
+
+**Warm up filesystem mounts** on all nodes — `ls` each mounted path so
+Lustre/NFS metadata is cached for later use:
+```bash
+srun --jobid=<JOB_ID> -N <nodes> --ntasks-per-node=1 \
+  --container-name=<container_name> \
+  --container-mounts=<container_mounts> \
+  bash -c 'for p in <mount_target_1> <mount_target_2> ...; do ls "$p" > /dev/null 2>&1; done'
+```
+Parse mount targets from `container_mounts` — the right-hand side of each
+`host:container` pair.
+
+**Check GPU status** on all nodes (no container needed — `nvidia-smi` is on
+the host):
+```bash
+srun --jobid=<JOB_ID> -N <nodes> --ntasks-per-node=1 \
+  bash -c 'echo "=== $(hostname) ===" && nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader'
+```
+- If output shows no processes → GPUs are clean, proceed
+- If unexpected processes found → warn user: "GPU processes found on
+  <node>: <process_list>. These may interfere with the job."
+
+**Run install** on all nodes with `--container-writable` (so packages persist
+in the named container across srun calls). Install all common requirements
+upfront so that subsequent jobs of any workflow type can skip install:
+```bash
+srun --jobid=<JOB_ID> -N <nodes> --ntasks-per-node=1 \
+  --container-name=<container_name> --container-writable \
+  --container-mounts=<container_mounts> \
+  bash -c 'cd <repo_root> && pip install -e . && pip install -r requirements-dev.txt && \
+    if [ -f examples/trtllm-eval/requirements.txt ]; then pip install -r examples/trtllm-eval/requirements.txt; fi'
+```
+For custom workflow with `skip_install=true`, skip this step entirely.
+
+**Write state file** `<repo_root>/work_dirs/.slurm_alloc.json`:
+```json
+{
+  "job_id": "<JOB_ID>",
+  "container_name": "<container_name>",
+  "nodelist": "<NODELIST>",
+  "partition": "<partition>",
+  "account": "<account>",
+  "nodes": <nodes>,
+  "gpus_per_node": <gpus_per_node>,
+  "docker_image": "<docker_image>",
+  "container_mounts": "<container_mounts>",
+  "allocated_at": "<ISO timestamp>",
+  "time_limit": "<alloc_time_limit>",
+  "installed": true
+}
+```
+
+Proceed to Step 1 (Persistent Path).
+
+#### Step 0F — Reuse path validation
+
+Allocation is valid. Check if container name matches current request:
+
+- If `container_name` in state file **matches** `container_name` from
+  `job_spec.json` → proceed to Step 1 (Persistent Path). Container and install
+  are already set up.
+- If **different** (model changed) → import the new container, warm up mounts,
+  and install:
+  ```bash
+  srun --jobid=<job_id> -N <nodes> --ntasks-per-node=1 \
+    --container-image=<docker_image> \
+    --container-name=<new_container_name> true
+  srun --jobid=<job_id> -N <nodes> --ntasks-per-node=1 \
+    --container-name=<new_container_name> \
+    --container-mounts=<container_mounts> \
+    bash -c 'for p in <mount_targets>; do ls "$p" > /dev/null 2>&1; done'
+  srun --jobid=<job_id> -N <nodes> --ntasks-per-node=1 \
+    --container-name=<new_container_name> --container-writable \
+    --container-mounts=<container_mounts> \
+    bash -c 'cd <repo_root> && pip install -e . && pip install -r requirements-dev.txt'
+  ```
+  Update `container_name` in state file. Proceed to Step 1 (Persistent Path).
+
+---
+
+### Step 1: Execute
+
+Two execution paths depending on `persistent_mode`.
+
+#### Persistent Path (persistent_mode=true)
+
+**Pre-flight time check** — verify the allocation has enough remaining time:
+```bash
+squeue -j <job_id> -h -o "%L"
+```
+If remaining time < job's `time_limit` from `slurm_params`, warn user:
+"Allocation has <remaining> left but job expects <time_limit>. The job may
+be killed early."
+
+**Pre-flight GPU check** — verify GPUs are not occupied by leftover processes:
+```bash
+srun --jobid=<job_id> -N <nodes> --ntasks-per-node=1 \
+  bash -c 'procs=$(nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader 2>/dev/null); [ -n "$procs" ] && echo "WARNING: GPU processes on $(hostname): $procs"'
+```
+If unexpected processes found, warn user before proceeding.
+
+**Run the sbatch script** with SLURM env vars — the `#SBATCH` directives are
+comments when run via `bash`; the inner `srun` inherits the env vars and uses
+the persistent allocation:
+```bash
+export SLURM_JOB_ID=<job_id>
+export SLURM_JOB_NUM_NODES=<nodes>
+export SLURM_NNODES=<nodes>
+export SLURM_NTASKS=<ntasks>
+export SLURM_NTASKS_PER_NODE=<ntasks_per_node>
+export SLURM_NODELIST=<nodelist>
+export SLURM_JOB_NODELIST=<nodelist>
+bash <script_path> 2>&1 | tee <work_dir>/<model_name>_<workflow_type>.log
+```
+
+Run with `run_in_background`. Container reuse is automatic (pyxis skips
+re-import when `--container-name` already exists on the node). The install
+step (Step 1 srun) always runs to keep the container up to date.
+
+**Step cancellation:** If a job hangs or the user wants to abort, cancel just
+the srun step without killing the allocation:
+```bash
+# List active steps:
+squeue -s -j <job_id>
+# Cancel a specific step (e.g., step 0):
+scancel <job_id>.0
+```
+The allocation stays RUNNING — new jobs can run immediately. The job name
+`<account>-trtllm.persistent` uniquely identifies allocations created by this
+skill via `squeue --name=<account>-trtllm.persistent`.
+
+#### One-Shot Path (persistent_mode=false)
+
+```bash
+sbatch <script_path>
+```
+
+Parse the job ID from output: `Submitted batch job <JOB_ID>`.
+
+If sbatch fails, report the error immediately and stop.
+
+### Step 2: Report Submission
+
+**One-shot mode:**
+```
+Job ID: <JOB_ID>
+Script: <script_path>
+Work directory: <work_dir>
+Log files: <work_dir>/<log_file_pattern with %j replaced>
+Monitor: squeue -j <JOB_ID>
+```
+
+**Persistent mode:**
+```
+Allocation: <job_id> (persistent, remaining: <time>)
+Container: <container_name>
+Script: <script_path>
+Work directory: <work_dir>
+Log file: <work_dir>/<model_name>_<workflow_type>.log
+```
+
+### Step 3: Poll / Wait for Completion (with Hang Monitoring)
+
+Both polling modes must also monitor the log file for hang indicators. Every
+60 seconds, check the log file for the pattern `hang detected`
+(case-insensitive). If found, terminate the job and skip to Step 5 with
+`HANG_DETECTED` status.
+
+**One-shot mode:** Poll every 30 seconds, max 60 polls (30 minutes):
+
+```bash
+squeue -j <JOB_ID> -h -o "%T %M %R"
+```
+
+- Report state transitions: `PENDING` → `RUNNING` → `COMPLETING` → done
+- If `PENDING` for > 5 minutes, check the reason:
+  ```bash
+  squeue -j <JOB_ID> -o "%i %T %r %S"
+  ```
+- The job is done when `squeue -j <JOB_ID>` returns empty (no rows)
+- If max polls exceeded, report the job is still running and provide manual
+  check commands
+- **Hang check** (every other poll, i.e., every 60s): Once the job is RUNNING,
+  resolve the log file path (`<work_dir>/<log_file_pattern>` with `%j` replaced
+  by `<JOB_ID>`) and check:
+  ```bash
+  grep -qi "hang detected" <log_file> 2>/dev/null && echo "HANG_DETECTED"
+  ```
+  If `HANG_DETECTED` is printed, immediately cancel the job:
+  ```bash
+  scancel <JOB_ID>
+  ```
+  Then read the last 200 lines of the log file and skip to Step 5.
+
+**Persistent mode:** While waiting for `run_in_background` completion
+notification, launch a hang monitoring loop using `run_in_background`:
+
+```bash
+while true; do
+  sleep 60
+  LOG_FILE="<work_dir>/<model_name>_<workflow_type>.log"
+  if [ -f "$LOG_FILE" ] && grep -qi "hang detected" "$LOG_FILE"; then
+    echo "HANG_DETECTED: Found 'hang detected' in log file"
+    squeue -s -j <job_id> -h -o "%i" | head -1 | xargs -r scancel 2>/dev/null
+    exit 1
+  fi
+done
+```
+
+When the main process completes normally, kill the monitoring loop. If the
+monitor fires first, it cancels the active Slurm step (not the allocation
+itself — the allocation stays for future runs).
+
+### Step 4: Collect Results
+
+On completion:
+
+1. **Get exit code:**
+   - **One-shot:** `sacct -j <JOB_ID> --format=JobID,State,ExitCode,Elapsed,MaxRSS -n`
+   - **Persistent:** From the `run_in_background` result exit code
+
+2. **Resolve log file paths:**
+   - **One-shot:** Replace `%j` in `log_file_pattern` with `<JOB_ID>`:
+     - Output: `<work_dir>/<MODEL_NAME>_..._<JOB_ID>.out`
+     - Error: `<work_dir>/<MODEL_NAME>_..._<JOB_ID>.err`
+   - **Persistent:** Single log at `<work_dir>/<model_name>_<workflow_type>.log`
+
+3. **Read log files:**
+   - Read the last 100 lines of the output log
+   - If job failed, also read the first 50 lines to catch early errors
+   - **One-shot only:** If job failed, read the last 50 lines of the `.err` file
+
+4. **Search for patterns:**
+   - Grep the log file for each `success_patterns` entry
+   - Grep the log file for each `failure_patterns` entry
+
+### Step 5: Report
+
+Return a structured result:
+
+```
+Job ID: <JOB_ID> (or allocation <job_id> for persistent mode)
+Status: COMPLETED | FAILED | TIMEOUT | CANCELLED | OUT_OF_MEMORY | HANG_DETECTED
+Exit code: <exit code>
+Elapsed: <time>
+Log files:
+  stdout: <path>
+  stderr: <path> (one-shot only)
+Work directory: <work_dir>
+Summary: <relevant output — pytest summary, accuracy score, throughput, or last lines>
+Errors: <if failed, relevant error output>
+```
+
+### Output Parsing by Workflow Type
+
+- **pytest**: Look for `X passed, Y failed in Zs` summary line. Status maps:
+  `COMPLETED` + 0 exit → `PASSED`; non-zero → `FAILED`
+- **eval**: Look for `accuracy:` lines. If `--check_accuracy` was used,
+  assertion error means threshold not met
+- **benchmark**: Look for throughput/latency numbers. Check for
+  `8_done_*.txt` completion marker
+- **custom**: No specific patterns — report last 10 lines of output
+
+### Failure Diagnosis
+
+| Symptom | Likely Cause | Action |
+|---------|-------------|--------|
+| `TIMEOUT` | Job exceeded time limit | Report and suggest increasing time |
+| `OUT_OF_MEMORY` | OOM on node | Report and suggest reducing batch size |
+| Exit `1:0` + `FAILED` | Non-zero exit from srun | Read `.err` file for details |
+| `CANCELLED` | Preempted or user cancelled | Check `sacct --format=JobID,State,Reason` |
+| Empty output log | Container failed to start | Check `.err` / `.salloc.log` for mount or image errors |
+| `HANG_DETECTED` in log | Process/GPU hang (deadlock, NCCL timeout) | Job auto-cancelled; check last 200 log lines for root cause |
+| `srun: error: Unable to create step` | Allocation expired | Delete state file, re-run (will re-allocate) |
+| `srun: error: ... does not exist` | Named container missing | Delete state file, re-run (will re-import) |
+
+## Rules
+
+- Never use `squeue` with the `-uall` flag
+- Always use `sacct` (not `squeue`) for final job status in one-shot mode
+- Poll with `run_in_background` + `sleep 30 &&` to avoid blocking (one-shot)
+- Never read full log files — use Read with offset/limit
+- If the Slurm script path doesn't exist, report the error immediately
+- For cluster hardware info (GPU types, partitions, constraints, storage
+  paths), use the `trtllm-agent-toolkit:internal-env-info` skill
+  when it is installed. If the skill is absent, skip the lookup silently and
+  use the caller-supplied fields or the srun probe in Step 1 — do not report
+  the missing skill as an error.
+- Never auto-release an allocation after job completion — it should persist
+  for potential future runs until walltime expires or user explicitly releases
+- Must pass `-N <nodes>` to every direct `srun --jobid` call — without it,
+  srun defaults to 1 node even in a multi-node allocation
+- The container name used in setup steps must **exactly match** the
+  `--container-name` in the `.slurm` script, otherwise pyxis will re-import
+  the image and lose `--container-writable` state
diff --git a/.claude/skills/exec-remote-slurm/SKILL.md b/.claude/skills/exec-remote-slurm/SKILL.md
new file mode 100644
index 000000000000..b4b22a4ef3d5
--- /dev/null
+++ b/.claude/skills/exec-remote-slurm/SKILL.md
@@ -0,0 +1,970 @@
+---
+name: exec-remote-slurm
+description: >-
+  Remote SLURM cluster development via SSH. Use when running jobs,
+  profiling, or developing on a remote SLURM cluster with pyxis/enroot
+  containers. Covers SSH connection management, srun/sbatch/salloc job
+  patterns, tmux-based allocation persistence, file transfer, and safe
+  remote file access. Works with any SLURM cluster accessible via SSH.
+tags: [infrastructure, remote, slurm]
+license: Apache-2.0
+metadata:
+  author: NVIDIA Corporation
+---
+
+# Remote SLURM Development
+
+Develop and execute on remote SLURM clusters via SSH commands. No filesystem
+sync needed — all remote operations go through SSH.
+
+## Step 1: Resolve Cluster
+
+**Optional dependency — `internal-env-info`.** Before doing
+anything else in Step 1, check whether `skills/internal-env-info/`
+exists in the toolkit. If it does **not**, skip every sub-step that invokes
+the skill or reads its `references/` files and proceed using only the
+cluster fields the orchestrator wrote into `job_spec.json` (`ssh_host`,
+`slurm_user`, `remote_cwd`, `partition`, `account`, `container_image`,
+`mounts`, `gpus_per_node`, `mfa_style`). If a required field is also
+missing, stop and ask the user to supply it — do **not** emit an error
+about the missing skill.
+
+The orchestrator supplies one of two input shapes; pick the right path:
+
+- **Explicit cluster** — `slurm_cluster` name is provided. Skip
+  auto-selection and go straight to **Load the per-cluster info**.
+- **Auto-select cluster** — only hardware constraints are provided
+  (`device_type`, `total_required_devices`, `required_devices_per_node`).
+  Resolve a cluster from the constraints, then load its per-cluster info.
+
+### 1a. Auto-select (only when no explicit cluster was provided)
+
+Invoke the `internal-env-info` skill in **constraint-based mode**, passing
+`device_type`, `required_devices_per_node`, and (optionally)
+`total_required_devices`. That skill owns the filter rules and either
+returns a single resolved cluster entry or stops and asks the user. Do not
+implement the filter here.
+
+Use `node_count` from `job_spec.json` directly as `--nodes` and
+`required_devices_per_node` as `--ntasks-per-node`. **Do not recompute
+`node_count`** — `trtllm-case-executor` is the single source of truth and
+writes the resolved value into `job_spec.json` for every executor to read.
+
+The returned entry's `cluster_name`, `mfa_style`, `default_user_root_dir`,
+`default_models_repo`, and `gpus_per_node` feed Step 1b.
+
+### 1b. Load the per-cluster info
+
+Connection fields are read **only** from `job_spec.json` (written by
+`trtllm-case-executor`): `ssh_host`, `slurm_user`, `remote_cwd`, `account`,
+`partition`, `mounts`, `container_image`, `gpus_per_node`, and `packages`.
+There is no per-cluster connection-config file to parse — if a required
+field is absent from `job_spec.json`, stop and ask the user to supply it.
+Do not silently default.
+
+Invoke the `internal-env-info` skill in single-cluster mode with
+the resolved `cluster_name` to retrieve per-cluster info: partition limits,
+topology-aware scheduling guidance, `default_models_repo`,
+`default_user_root_dir`, `gpus_per_node`, `mfa_style` (used by Step 2 to
+choose the SSH path — see below), and shared constraint and tuning
+patterns. The runtime `container_image` always comes from `job_spec.json`,
+never from this skill.
+
+## Step 2: SSH Preflight
+
+The cluster's `mfa_style` (returned by `internal-env-info` in Step 1b) is
+the **single source of truth for which SSH path to use**. Switch on the
+value — do **not** probe-and-fall-back on the SSH error string.
+
+| `mfa_style` | Path | Why |
+|-------------|------|-----|
+| `false` | Direct (Step 2a) | Cluster accepts key-based / passwordless SSH; no special handling needed. |
+| `true`  | MFA (Step 2b)    | Cluster requires interactive MFA (Duo / device-code / one-time PIN). |
+| `null`  | Unknown (Step 2c)| `mfa_style` not declared in the cluster reference (or `internal-env-info` is not installed). Probe direct first, fall back to MFA. When the reference exists, also ask the user to update it so future runs skip the probe; otherwise just continue silently. |
+
+### 2.0. Build `<ssh_host>` from `slurm_user` + `cluster_name`
+
+Before any SSH attempt, derive the SSH destination. Priority order:
+
+1. **Caller-supplied `ssh_host`** (passed in through case-executor) → use
+   verbatim, skip everything below. This is the canonical escape hatch.
+2. **`internal-env-info` skill installed** → delegate for the
+   `ssh_host_pattern` key (defined in the skill's `Default Environment Values`
+   section). It returns a single-line snippet that concatenates `slurm_user`
+   and `cluster_name` into `SSH_HOST`. Do **not** re-derive or re-suffix
+   `slurm_user` here — `internal-env-info` already applied whatever its
+   `slurm_user` construction rule requires for that cluster.
+
+   Inputs to substitute:
+   - `slurm_user` — pre-resolved from `internal-env-info` (already in its
+     final, SSH-ready form for that cluster). Read it from
+     `job_spec.json`.
+   - `cluster_name` — from `internal-env-info` Step 1b.
+3. **Skill not installed** → stop and ask the user to supply `ssh_host`
+   explicitly. Do **not** hard-code a login-host pattern in this skill —
+   the NVIDIA-internal convention lives only in `internal-env-info`.
+
+The resolved value is the canonical `<ssh_host>` placeholder used everywhere
+in the rest of this skill (Steps 2a/2b/2c, Step 3, recipes, file-transfer,
+hang detection).
+
+### 2a. Direct path — `mfa_style: false`
+
+Verify connectivity using the user's existing SSH config:
+
+```bash
+ssh -o ConnectTimeout=5 <ssh_host> "echo ok"
+```
+
+If this succeeds, **skip to Step 3**. Do not create a trtllm-managed
+ControlMaster — reuse whatever the user already has.
+
+If it fails, do **not** silently switch to the MFA path. Stop and ask the
+user to verify their `ssh_host` value or their key-based SSH setup.
+A direct-path cluster failing here is a configuration issue, not an MFA
+prompt to interpret.
+
+### 2b. MFA path — `mfa_style: true`
+
+First check whether a ControlMaster is already live (the user may have set
+one up earlier in this session):
+
+```bash
+ssh -O check <ssh_host> 2>&1 || ssh -o ControlPath=/tmp/ssh-trtllm-%r@%h:%p -O check <ssh_host> 2>&1
+```
+
+If either succeeds, **skip to Step 3** — reuse it.
+
+Otherwise the user must complete an interactive MFA handshake. MFA uses
+browser-based device-code verification that needs an interactive SSH
+session — the skill cannot complete this on the user's behalf because the
+session must stay alive while the user authenticates in their browser and
+presses ENTER.
+
+Tell the user to run the ControlMaster command:
+
+> Run this in your terminal:
+>
+>     ssh -fNM -o ControlMaster=auto -o ControlPath=/tmp/ssh-trtllm-%r@%h:%p -o ControlPersist=24h <ssh_host>
+>
+> If the cluster uses a device-code flow you will see a PIN and a URL like:
+>
+>     Authenticate with PIN XXXXXX at <your-identity-provider-device-url>
+>
+> Open that URL, enter the PIN, approve, then press ENTER. The SSH session
+> will background itself as a persistent ControlMaster. Let me know when done.
+
+After the user confirms, verify with
+`ssh -o ControlPath=/tmp/ssh-trtllm-%r@%h:%p -O check <ssh_host>` and
+proceed to Step 3.
+
+### 2c. Unknown — `mfa_style: null`
+
+The cluster reference (`references/<cluster>.md`) does not declare an
+`mfa_style` value. Try the direct path (Step 2a) first. If it fails with
+`Permission denied (keyboard-interactive)` — and only that — switch to
+the MFA path (Step 2b). Any other failure stops here and is reported to
+the user.
+
+After connecting successfully, **ask the user to update the cluster
+reference file** with the discovered value (`mfa_style: true` or
+`false`, in the `## Authentication` table) so that future runs skip
+this probe and use the deterministic switch in Step 2a / 2b.
+
+## SSH Command Shorthand
+
+Every SSH and SCP invocation in the rest of this skill uses the placeholders below. Both pin `ControlPath` to the trtllm-managed socket so they reuse the ControlMaster set up in Step 2 — without this, OpenSSH starts a fresh handshake on every call, which on MFA clusters re-prompts for keyboard-interactive auth (and defeats the whole point of multiplexing):
+
+| Placeholder | Expansion |
+|-------------|-----------|
+| `<ssh_cmd>` | `ssh -o ControlPath=/tmp/ssh-trtllm-%r@%h:%p` |
+| `<scp_cmd>` | `scp -o ControlPath=/tmp/ssh-trtllm-%r@%h:%p` |
+
+When the user has their own ControlMaster configured in `~/.ssh/config` (no trtllm-managed socket needed), the explicit `-o ControlPath` is harmless: if the path has no live master, OpenSSH falls back to a fresh connection. The connectivity probe in Step 2 Check 2 (`ssh -o ConnectTimeout=5 <ssh_host> "echo ok"`) deliberately uses bare `ssh` to test un-multiplexed access; everywhere else use `<ssh_cmd>` / `<scp_cmd>`.
+
+## Step 3: Remote Environment Check
+
+Verify the remote environment is ready:
+
+```bash
+<ssh_cmd> <ssh_host> "test -d <remote_cwd> && echo ok"
+<ssh_cmd> <ssh_host> "which tmux"
+<ssh_cmd> <ssh_host> "squeue --version"
+```
+
+If `remote_cwd` does not exist, create it or ask the user. If tmux or SLURM
+commands are not available, stop and inform the user.
+
+## Step 4: Prepare Remote Environment
+
+This step sets up the user's working directory and clones the TensorRT-LLM
+repository on the remote cluster. It runs once before any job submission.
+
+> **MANDATORY: two-mechanism sync — both must run.**
+> - **Tracked changes** → temp branch (**4c–4f**): commit on a session-scoped
+>   branch, push to `origin`, fetch + checkout on remote, delete the ref both
+>   sides.
+> - **Untracked files** (`git ls-files --others --exclude-standard`) →
+>   `tar | ssh` pipe (**4e.1**), right after the remote checkout. Never touches
+>   `origin`.
+>
+> **Do not** shortcut with `git fetch origin <repo_branch> && git checkout <repo_branch>` —
+> `repo_branch` may not exist on `origin`, may be stale, or may lack uncommitted
+> edits. Always run 4c–4f **and** 4e.1.
+
+**Inputs received from `trtllm-case-executor`:**
+
+| Field | Description |
+|-------|-------------|
+| `default_user_root_dir` | Per-user root directory on the remote cluster (from `exec-env-check` / `internal-env-info`) |
+| `repo_url` | Git remote URL of the local TensorRT-LLM repo (from `git remote get-url origin`). **Used to clone the bare repository on the remote (Step 4b) and to push the temp branch to (Step 4d).** |
+| `repo_branch` | **Informational only** — the local current branch name (from `git rev-parse --abbrev-ref HEAD`). Used in Step 4c to remember which branch to restore locally after the temp branch is deleted. **Never used as the ref the remote checks out** — the remote always checks out the temp branch from Step 4c. |
+| `device_type` | Required short GPU name (e.g., `B200`, `B300`). Used by Step 1a auto-selection. Ignored when an explicit cluster was provided. |
+| `total_required_devices` | Total GPU processes across all nodes. Used by Step 1a auto-selection only. SLURM job sizing reads `node_count` from `job_spec.json` (precomputed by `trtllm-case-executor`) — do not recompute it here. |
+| `required_devices_per_node` | Minimum GPUs per node. Used by Step 1a auto-selection (cluster must have `gpus_per_node >= required_devices_per_node`) and as `--ntasks-per-node`. |
+
+**Procedure:**
+
+### 4a. Create user root directory
+
+```bash
+<ssh_cmd> <ssh_host> "mkdir -p <default_user_root_dir>"
+```
+
+If `default_user_root_dir` is null or empty, skip this step and use `remote_cwd` directly.
+
+### 4b. Clone or update the repository
+
+Determine the repo directory name from `repo_url` (strip `.git` suffix if present):
+
+```bash
+REPO_DIR_NAME=$(basename <repo_url> .git)
+REMOTE_REPO_PATH="<default_user_root_dir>/${REPO_DIR_NAME}"
+```
+
+Check if it already exists:
+
+```bash
+<ssh_cmd> <ssh_host> "test -d <remote_repo_path>/.git && echo exists"
+```
+
+If the repository does **not** exist, clone it:
+
+```bash
+<ssh_cmd> <ssh_host> "cd <default_user_root_dir> && git clone <repo_url>"
+```
+
+If the repository **already exists**, refresh remote refs only — do **not**
+checkout or reset to any branch here:
+
+```bash
+<ssh_cmd> <ssh_host> "cd <remote_repo_path> && git fetch origin --prune"
+```
+
+The actual ref the remote will check out is the temp branch created in 4c
+and pushed in 4d; the remote checkout happens in 4e. **Do not** add a
+`git checkout <repo_branch>` or `git reset --hard origin/<repo_branch>`
+here as a shortcut — that path is forbidden (see the MANDATORY callout
+at the top of Step 4).
+
+### 4c. Create a temporary local branch capturing all local changes
+
+Generate a unique temp branch name tied to this session. The branch name
+**must start with `<user_name>_<model_name>`** so it is easy to identify
+the owner and the workload it was created for. A timestamp suffix keeps
+it unique across sessions:
+
+```bash
+# Resolve the user name (prefer the slurm_user input over whoami,
+# so the temp branch is attributed to the cluster user rather than the
+# host that happens to dispatch the job).
+USER_NAME="${slurm_user:-$(whoami)}"
+
+# Resolve the workload name. job_spec.json's `model_name` is authoritative;
+# fall back to the repo basename if it is missing.
+MODEL_NAME="${model_name:-$(basename "<REPO_ROOT>")}"
+
+# Sanitize to a git-ref-safe identifier (alphanum + `_`/`-` only).
+# Any other character is replaced with `_`.
+sanitize() { printf '%s' "$1" | tr -c 'A-Za-z0-9_-' '_'; }
+USER_NAME_SAFE=$(sanitize "$USER_NAME")
+MODEL_NAME_SAFE=$(sanitize "$MODEL_NAME")
+
+TEMP_BRANCH="${USER_NAME_SAFE}_${MODEL_NAME_SAFE}_$(date +%Y%m%d%H%M%S)"
+```
+
+Example resolved value: `<user_name>_dwdp_deepseek_v3_lite_20260430145200`.
+
+Stash the current branch name so it can be restored later:
+
+```bash
+ORIGINAL_BRANCH=$(git -C "<REPO_ROOT>" rev-parse --abbrev-ref HEAD)
+```
+
+Create the temp branch at HEAD and stage tracked changes only
+(`-u`, not `-A` — untracked files ship in 4e.1):
+
+```bash
+git -C "<REPO_ROOT>" checkout -b "$TEMP_BRANCH"
+git -C "<REPO_ROOT>" add -u
+git -C "<REPO_ROOT>" commit --allow-empty -m "sync: temp snapshot for remote execution"
+```
+
+`--allow-empty` covers the no-tracked-changes case; the branch still exists
+at HEAD and 4e.1 still runs.
+
+### 4d. Push temp branch to the remote origin
+
+Push the temp branch so the remote cluster can fetch it:
+
+```bash
+git -C "<REPO_ROOT>" push origin "$TEMP_BRANCH"
+```
+
+### 4e. Fetch and check out temp branch on the remote cluster
+
+On the remote: stash prior-session leftovers, fetch, force-checkout the temp
+branch, and re-materialize `cpp/` from HEAD (stale `cpp/` produces
+missing-symbol `ImportError`s in Step 5):
+
+```bash
+<ssh_cmd> <ssh_host> "cd <remote_repo_path> && \
+    git stash push -u -m 'auto-stash before remote sync' || true && \
+    git fetch origin --prune && \
+    git checkout -B $TEMP_BRANCH origin/$TEMP_BRANCH && \
+    rm -rf cpp/ && \
+    git checkout HEAD -- cpp/"
+```
+
+- `stash -u` sweeps both tracked edits and untracked files left from a prior
+  session; `|| true` covers the no-op.
+- `checkout -B` force-resets even if a same-named branch already exists.
+- `rm -rf cpp/ && git checkout HEAD -- cpp/` always runs, even when checkout
+  reports no diffs.
+
+`remote_cwd` is updated to `<remote_repo_path>`. Untracked files for this
+session arrive next in 4e.1.
+
+### 4e.1. Transfer untracked files via tar pipe over SSH
+
+Ship untracked, non-gitignored files on top of the freshly checked-out tree:
+
+```bash
+UNTRACKED_COUNT=$(git -C "<REPO_ROOT>" ls-files --others --exclude-standard | wc -l)
+if [ "$UNTRACKED_COUNT" -gt 0 ]; then
+    git -C "<REPO_ROOT>" ls-files --others --exclude-standard -z | \
+        tar -C "<REPO_ROOT>" --null -T - -cf - | \
+        <ssh_cmd> <ssh_host> "tar -C <remote_repo_path> -xf -"
+fi
+```
+
+`-z` / `--null` keeps filenames with spaces or newlines intact; `tar -x`
+auto-creates parent dirs; the guard skips the round-trip when nothing to
+ship. Do **not** substitute `rsync` or `sshfs` (not reliably available on HPC).
+
+### 4f. Delete the temp branch (local and remote)
+
+After the remote checkout succeeds, immediately delete the temp branch from
+the local repo and the remote origin, then restore the original local branch:
+
+```bash
+git -C "<REPO_ROOT>" checkout "$ORIGINAL_BRANCH"
+git -C "<REPO_ROOT>" branch -D "$TEMP_BRANCH"
+git -C "<REPO_ROOT>" push origin --delete "$TEMP_BRANCH"
+```
+
+The remote cluster keeps the checked-out working tree from the deleted branch
+— only the ref is removed, not the files.
+
+### 4f.1 Create remote work_dir under the cloned repo
+
+The remote work_dir must mirror the local layout (`<REPO_ROOT>/work_dirs/<name>`)
+so that all per-run artefacts — slurm script, stdout/stderr logs, intermediate
+files — live inside the cloned repo and are co-located with the code that
+produced them.
+
+`remote_work_dir` is precomputed by `trtllm-case-executor` (Step 1) and read
+verbatim from `job_spec.json` — it is `<remote_repo_path>/work_dirs/<basename(local_work_dir)>`.
+**Do not** fall back to `<remote_cwd>/work_dirs/...` or any other root; that
+older layout is deprecated.
+
+```bash
+<ssh_cmd> <ssh_host> "mkdir -p <remote_work_dir>"
+```
+
+All subsequent steps (file uploads, sbatch submission, log paths) reference
+`<remote_work_dir>` from `job_spec.json`.
+
+### 4g. Copy build scripts to remote node
+
+The skill ships two build scripts in `scripts/` alongside this SKILL.md.
+Copy them to a dedicated directory on the remote so they are accessible
+from inside the SLURM container:
+
+```bash
+REMOTE_SCRIPTS_DIR="<default_user_root_dir>/trtllm-build-scripts"
+SKILL_SCRIPTS_DIR="<LOCAL_SKILL_DIR>/scripts"   # local path to this skill's scripts/
+
+<ssh_cmd> <ssh_host> "mkdir -p ${REMOTE_SCRIPTS_DIR}/log"
+
+<ssh_cmd> <ssh_host> "cat > ${REMOTE_SCRIPTS_DIR}/build.slurm" < "${SKILL_SCRIPTS_DIR}/build.slurm"
+<ssh_cmd> <ssh_host> "cat > ${REMOTE_SCRIPTS_DIR}/build.sh"    < "${SKILL_SCRIPTS_DIR}/build.sh"
+<ssh_cmd> <ssh_host> "chmod +x ${REMOTE_SCRIPTS_DIR}/build.slurm ${REMOTE_SCRIPTS_DIR}/build.sh"
+```
+
+`<LOCAL_SKILL_DIR>` is the directory containing this SKILL.md file
+(`.claude/skills/exec-remote-slurm/`).
+
+---
+
+## Step 5: Build on Remote SLURM Node
+
+> **Mandatory unless `build_project: false`.** Step 4e wipes `cpp/`, so
+> the remote has no compiled C++ extensions until Step 5 succeeds (exit
+> code `0:0`). Submitting any test/bench job before then fails on the
+> first `import tensorrt_llm` with missing-symbol errors.
+>
+> Do **not** replace this with `pip install -e .` inside the user's
+> `.slurm` script — it does not reliably recompile `cpp/`, and when it
+> does it blocks the GPU allocation behind a 30+ minute compile. Step 5
+> runs the build as a 1-node side job and installs into the persistent
+> container volume.
+>
+> If the caller sets `build_project: false`, skip Step 5 (the caller is
+> asserting the install is current). Do not silently downgrade `true` →
+> `false` because Step 5 looks expensive.
+
+Submit a SLURM batch job that compiles TensorRT-LLM on the remote cluster
+inside the official container. Wait for it to complete before proceeding to
+job submission.
+
+### 5a. Read the container image
+
+The container image is resolved by `trtllm-case-executor` (Step 3) and
+forwarded in `job_spec.json` as `container_image`. Read it from there —
+**do not grep `current_image_tags.properties` and do not re-derive from
+the cluster config**. case-executor is the single source of truth for the
+image URI; this skill only consumes it.
+
+Apply the enroot URI rewrite at use-time so the registry path is not
+mistaken for a Docker-style credential lookup. The rewrite recipe lives
+in the `internal-env-info` skill's `Default Environment Values` section
+(`enroot_uri_rewrite` key):
+
+- **Skill installed** → delegate to it, passing the canonical
+  `<container_image>` from `job_spec.json`. The skill returns the bash
+  snippet (an `sed` over the registry-host prefix); run it here to
+  produce `<IMAGE_PATH>`.
+- **Skill not installed** → use `<container_image>` verbatim as `<IMAGE_PATH>`.
+
+Use `<IMAGE_PATH>` as the value passed to `--container-image=` in the
+build sbatch below and in every `srun`/`sbatch` recipe later in this
+skill. The canonical (un-rewritten) URI in `job_spec.json` is unchanged;
+the rewrite is purely a transport adaptation for enroot.
+
+### 5b. Derive build architecture
+
+Map `device_type` (from the env-check output) to the `build_wheel.py` `-a` flag:
+
+| `device_type` | `-a` value |
+|------------|------------|
+| B200, GB200, GB300, B300 | `100-real` |
+| H100, H200 | `90-real` |
+| L40S | `89-real` |
+| A100 | `80-real` |
+
+Default to `100-real` if `device_type` is unknown.
+
+### 5c. Submit the build job
+
+```bash
+<ssh_cmd> <ssh_host> "sbatch \
+    --nodes=1 --ntasks=1 --ntasks-per-node=1 \
+    --partition=<partition> \
+    --account=<account> \
+    --job-name=<account>-trtllm.build \
+    --time=02:00:00 \
+    ${REMOTE_SCRIPTS_DIR}/build.slurm \
+    <IMAGE_PATH> <default_user_root_dir> ${REMOTE_SCRIPTS_DIR} <remote_repo_path> \
+    -a \"<arch>\" --nvtx"
+```
+
+Capture the job ID from `sbatch` stdout (format: `Submitted batch job <id>`).
+
+If `gpus_per_node` is set in the cluster config, add
+`--gpus-per-node=<gpus_per_node>` to the `sbatch` flags.
+
+### 5d. Monitor the build job
+
+Poll every 30–60 seconds until the job leaves the queue:
+
+```bash
+<ssh_cmd> <ssh_host> "squeue -j <job_id> -o '%.18i %.9P %.30j %.8u %.2t %.10M %R'"
+```
+
+Once running, tail the log to track compile progress:
+
+```bash
+<ssh_cmd> <ssh_host> "tail -30 ${REMOTE_SCRIPTS_DIR}/log/build_<job_id>.srun.log"
+```
+
+Progress indicators in the log:
+- `[XX%] Building CXX object ...` — compilation in progress
+- `Linking CXX ...` — link phase
+- `FAILED:` / `error:` / `fatal error:` — build failure
+- `Successfully built` — success
+
+If the job disappears from `squeue`, proceed to Step 5e.
+
+### 5e. Verify the build
+
+```bash
+<ssh_cmd> <ssh_host> "sacct -j <job_id> --format=JobID,State,ExitCode,Elapsed -n"
+<ssh_cmd> <ssh_host> "tail -50 ${REMOTE_SCRIPTS_DIR}/log/build_<job_id>.srun.log"
+```
+
+A successful build shows exit code `0:0` and ends with
+`Successfully built tensorrt_llm` or no error lines.
+
+If the build fails, report the last 100 lines of the log to the user and
+stop — do not proceed to job submission.
+
+## Build Scripts Reference
+
+| Script | Location | Purpose |
+|--------|----------|---------|
+| `build.slurm` | `scripts/build.slurm` | SLURM batch script — launches container, calls `build.sh` |
+| `build.sh` | `scripts/build.sh` | Runs inside container — calls `build_wheel.py` with defaults |
+
+Both scripts accept the same positional arguments:
+`<container_image> <user_root_dir> <scripts_dir> <repo_dir> [build_wheel_args...]`
+
+---
+
+## Recipes
+
+Pick the recipe that matches your workflow. Each recipe is self-contained —
+copy it, substitute the `<placeholders>` from config, and run.
+
+**Placeholders** (all from `job_spec.json` unless noted):
+
+| Placeholder | Source |
+|-------------|--------|
+| `<ssh_host>`, `<account>`, `<partition>`, `<mounts>` | `job_spec.json` |
+| `<image>` | `container_image` from `job_spec.json` (resolved by `trtllm-case-executor` Step 3). Apply enroot rewrite (`/` → `#`) at use-time per Step 5a — substitute the rewritten `<IMAGE_PATH>` here. Do not re-resolve from the cluster config or `current_image_tags.properties`. |
+| `<job_name>` | Pre-resolved by `trtllm-case-executor` and forwarded as `job_name` in `job_spec.json`. Substitute directly into `-J <job_name>`; do not reconstruct from `account` / `subproject` / `detail`. |
+| `<nodes>` | `node_count` from `job_spec.json` (computed by `trtllm-case-executor`). Do not recompute. |
+| `<gpus>` | `required_devices_per_node` from `job_spec.json` — probe only if absent (see bottom of this section) |
+| `<remote_repo_path>` | `remote_repo_path` from `job_spec.json` — `<default_user_root_dir>/<repo_dir_name>` precomputed by `trtllm-case-executor` Step 1. The cloned repo and `--container-workdir` for srun. Do not recompute. |
+| `<remote_work_dir>` | `remote_work_dir` from `job_spec.json` — `<remote_repo_path>/work_dirs/<basename(local_work_dir)>` precomputed by `trtllm-case-executor` Step 1. Used for `--output`/`--error` slurm log paths and any per-run artefact uploads. Do not recompute or fall back to `<remote_cwd>/work_dirs/...`. |
+| `<time>` | Walltime in HH:MM:SS |
+
+**Job naming:** Every recipe includes `-J <job_name>` using the
+pre-resolved `job_name` value from `job_spec.json` (built by
+`trtllm-case-executor`). This is required — without `-J`, SLURM emits
+warnings and jobs are hard to identify in `squeue`. Do not reconstruct
+the name from `account` / `subproject` / `detail` here; that convention
+lives in `trtllm-case-executor`.
+
+**GPU allocation:** Cluster behavior varies:
+
+- **Exclusive-access clusters** (most NVIDIA internal): Nodes are allocated
+  exclusively with all GPUs available. No `--gpus-per-node` needed.
+- **Shared/QOS clusters**: Require `--gpus-per-node=<gpus_per_node>`. Without
+  it, you may get a CPU-only node (no GPUs at all), a "Cannot find GPU
+  specification" error, or a "QOSMinGRES" error (below minimum).
+
+If `gpus_per_node` is set in the cluster config, add `--gpus-per-node=<gpus_per_node>`
+to every `srun`, `salloc`, and `sbatch` command. If not set, omit it.
+
+If a job fails with GPU-related errors and `gpus_per_node` is not configured,
+probe the node's GPU count, set `gpus_per_node` in the cluster config, and retry.
+
+Do not set `CUDA_VISIBLE_DEVICES` for MPI workloads — the MPI runtime handles
+GPU assignment.
+
+---
+
+### Recipe 1: Single-process command (srun)
+
+One-off commands on a single node: profiling, inspecting, lightweight scripts.
+
+```bash
+<ssh_cmd> <ssh_host> "srun -A <account> -p <partition> -N 1 --ntasks-per-node=1 \
+  -J <job_name> -t <time> \
+  --container-image=<image> --container-mounts=<mounts> \
+  <command>"
+```
+
+---
+
+### Recipe 2: MPI workload (srun)
+
+MPI applications across nodes: NCCL tests, HPCX benchmarks. The binary must
+be MPI-aware (e.g., `all_reduce_perf_mpi`, not `all_reduce_perf`).
+
+```bash
+<ssh_cmd> <ssh_host> "srun -A <account> -p <partition> -N <nodes> \
+  --ntasks-per-node=<gpus> --mpi=pmix \
+  -J <job_name> -t <time> \
+  --container-image=<image> --container-mounts=<mounts> \
+  <command>"
+```
+
+---
+
+### Recipe 3: Distributed training (srun)
+
+Framework-launched distributed training (PyTorch DDP, DeepSpeed). SLURM starts
+`<nodes> × <gpus>` processes. Pyxis/enroot sets `RANK`, `LOCAL_RANK`, and
+`WORLD_SIZE`. Each process runs the training script directly.
+
+```bash
+<ssh_cmd> <ssh_host> "srun -A <account> -p <partition> -N <nodes> \
+  --ntasks-per-node=<gpus> \
+  -J <job_name> -t <time> \
+  --container-image=<image> --container-mounts=<mounts> \
+  python train.py"
+```
+
+---
+
+### Recipe 4: Launcher-based training (srun)
+
+One process per node, an in-container launcher forks locally (torchrun,
+deepspeed launcher).
+
+```bash
+<ssh_cmd> <ssh_host> "srun -A <account> -p <partition> -N <nodes> \
+  --ntasks-per-node=1 \
+  -J <job_name> -t <time> \
+  --container-image=<image> --container-mounts=<mounts> \
+  torchrun --nproc-per-node=<gpus> train.py"
+```
+
+---
+
+### Recipe 5: Batch job (sbatch)
+
+Fire-and-forget background job:
+
+```bash
+<ssh_cmd> <ssh_host> "sbatch -A <account> -p <partition> -N <nodes> \
+  -J <job_name> -t <time> \
+  --container-image=<image> --container-mounts=<mounts> \
+  <script-path>"
+```
+
+Check status:
+
+```bash
+<ssh_cmd> <ssh_host> "squeue -j <jobid>"
+<ssh_cmd> <ssh_host> "sacct -j <jobid> --format=JobID,State,ExitCode,Elapsed -n"
+```
+
+---
+
+### Recipe 6: Multi-command session (salloc + srun)
+
+Interactive development with persistent container state. Use this when you need
+multiple commands on the same allocation — installing packages, compiling,
+running experiments iteratively.
+
+> **Relationship to `persistent_mode`:** The `persistent_mode` field in
+> `job_spec.json` is consumed only by `exec-local-slurm`. This
+> skill does **not** read `persistent_mode`, `container_name`, or
+> `container_mounts` from `job_spec.json` for allocation reuse — remote
+> persistence is driven by this recipe (tmux + named `salloc` + named
+> container) and controlled by inputs to this skill directly (e.g.,
+> `release_allocation`, `alloc_time_limit`). If a caller forwards a
+> `persistent_mode=true` job spec into a `remote_slurm` run, treat the field
+> as informational only and use this recipe's flow.
+
+**Step 1: Allocate via tmux**
+
+tmux keeps the allocation alive on the frontend — invisible plumbing.
+
+```bash
+<ssh_cmd> <ssh_host> "tmux new-session -d -s trtllm-<id> \
+  'salloc -A <account> -p <partition> -N <nodes> \
+  -J <job_name> -t <time>'"
+```
+
+**Step 2: Check allocation status**
+
+Wait ~30 seconds for the scheduler, then get job ID and status:
+
+```bash
+<ssh_cmd> <ssh_host> "squeue -u <user> -h -o '%i %T %S' --state=RUNNING,PENDING"
+```
+
+If RUNNING → proceed to Step 3.
+
+If PENDING → follow the **walltime negotiation flow** (see Walltime strategy
+section). Only proceed to Step 3 once the job is RUNNING.
+
+**Step 3: Create a persistent container**
+
+Import the image once. This creates a named container on each node that
+persists across `srun` calls (container persistence):
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> --ntasks-per-node=1 \
+  --container-image=<image> --container-name=trtllm-<jobid> true"
+```
+
+**Container packages:** If `packages` is defined in config, install them
+using the exact names listed — do not guess or substitute package names:
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> --ntasks-per-node=1 \
+  --container-name=trtllm-<jobid> --container-writable \
+  pip install <packages>"
+```
+
+Only install when the workload actually needs them — not for simple
+commands like `nvidia-smi` or `hostname`.
+
+**Step 4: Run commands**
+
+Use `--container-name` (without `--container-image`) to reuse the persistent
+container. All in-container state is preserved between commands.
+
+Single-process command:
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> \
+  --container-name=trtllm-<jobid> --container-mounts=<mounts> \
+  <command>"
+```
+
+Per-node operation (mkdir, pip install) — use `--ntasks-per-node=1`:
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> --ntasks-per-node=1 \
+  --container-name=trtllm-<jobid> --container-mounts=<mounts> \
+  <command>"
+```
+
+MPI workload within the allocation:
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> \
+  --ntasks-per-node=<gpus> --mpi=pmix \
+  --container-name=trtllm-<jobid> --container-mounts=<mounts> \
+  <command>"
+```
+
+Distributed training within the allocation:
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> \
+  --ntasks-per-node=<gpus> \
+  --container-name=trtllm-<jobid> --container-mounts=<mounts> \
+  python train.py"
+```
+
+Add `--container-writable` if the workload writes to container paths outside
+of mounted directories (e.g., `pip install` into site-packages).
+
+**Step 5: Release the allocation**
+
+Clean up the named container, then kill tmux. If the allocation has already
+expired (walltime exceeded), the `srun` will fail — that is fine, the container
+is already gone. Just kill the tmux session.
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> --ntasks-per-node=1 \
+  bash -c 'enroot remove -f trtllm-<jobid> 2>/dev/null; true'" 2>/dev/null
+<ssh_cmd> <ssh_host> "tmux kill-session -t trtllm-<id>"
+```
+
+---
+
+### Probing GPUs per node
+
+If you do not know `<gpus>`, probe once and record in cluster notes:
+
+```bash
+<ssh_cmd> <ssh_host> "srun --jobid=<jobid> --ntasks-per-node=1 \
+  --container-name=trtllm-<jobid> bash -c 'nvidia-smi -L | wc -l'"
+```
+
+## Cluster-Specific Flags
+
+After substituting recipe placeholders, apply the cluster-specific flags
+returned by the `internal-env-info` skill in Step 1:
+
+- **SHARP**: Add `--network=sharp` if the cluster supports it
+- **Topology constraints**: Add `--constraint=<feature>` for topology-aware
+  placement (e.g., `--constraint=nvlblk01` for a single rack)
+- **GPU clocks**: Run `sudo nvidia-smi -ac <mem>,<graphics>` as a separate
+  `srun --ntasks-per-node=1` pre-step in sbatch scripts
+- **Comments**: Add `--comment=metrics` for Prometheus/Grafana collection
+
+## File Transfer
+
+Use SSH pipes for file transfer — no extra tools needed, works on every cluster.
+
+### Local → Remote (writing files)
+
+Pipe file content directly via SSH heredoc. Never write locally then scp:
+
+```bash
+<ssh_cmd> <ssh_host> "cat > <remote_cwd>/script.py" << 'PYEOF'
+import torch
+# ... file content ...
+PYEOF
+```
+
+For multiple files, repeat the pattern — each file is one SSH command.
+
+### Remote → Local (reading results)
+
+Read remote files via SSH stdout:
+
+```bash
+<ssh_cmd> <ssh_host> "cat <remote_cwd>/output.json"
+```
+
+For large files, read only what you need:
+
+```bash
+<ssh_cmd> <ssh_host> "tail -n 100 <remote_cwd>/slurm-<jobid>.out"
+<ssh_cmd> <ssh_host> "head -n 50 <remote_cwd>/results.csv"
+```
+
+### When to use scp
+
+Fall back to `scp` only for binary files or large directories that cannot be
+piped through SSH:
+
+```bash
+<scp_cmd> <ssh_host>:<remote_cwd>/profile.nsys-rep .    # binary profile
+<scp_cmd> -r local_dir/ <ssh_host>:<remote_cwd>/        # whole directory
+```
+
+**Decision rule:** text files → SSH pipe. Binary files or directories → scp.
+Never use rsync or sshfs — not reliably available on HPC clusters.
+
+## Rules
+
+### Never read full remote files
+
+Remote output can be gigabytes. Always use `tail`, `head`, or `grep`:
+
+```bash
+<ssh_cmd> <ssh_host> "tail -n 100 <remote_cwd>/output.log"
+<ssh_cmd> <ssh_host> "head -n 50 <remote_cwd>/script.py"
+<ssh_cmd> <ssh_host> "grep -n 'error\|Error\|ERROR' <remote_cwd>/output.log"
+```
+
+Never use `cat` on remote files unless you are certain the file is small.
+
+### Never run heavy computation on the frontend
+
+The frontend node is shared. Only use it for SLURM commands, lightweight file
+operations, and tmux management.
+
+### Check before creating
+
+Before creating resources, check if they already exist:
+- SSH ControlMaster: `ssh -O check <ssh_host> 2>&1 || <ssh_cmd> -O check <ssh_host>` (probe both the user's default socket and the trtllm-managed socket from Step 2)
+- tmux sessions: `<ssh_cmd> <ssh_host> "tmux has-session -t trtllm-<id> 2>/dev/null && echo exists"`
+- Existing allocations: `<ssh_cmd> <ssh_host> "squeue -u <user> -h -o '%i %j' --state=RUNNING"`
+
+### Clean up when done
+
+- Remove named containers: `enroot remove -f` (see Recipe 6, Step 5)
+- Kill tmux sessions: `<ssh_cmd> <ssh_host> "tmux kill-session -t trtllm-<id>"`
+- Cancel unneeded jobs: `<ssh_cmd> <ssh_host> "scancel <jobid>"`
+
+### Walltime strategy
+
+- **Production (sbatch, srun):** Tight walltime with ~20% margin — gets
+  scheduled sooner.
+- **Dev/debug (salloc):** Use **walltime negotiation** — submit with max
+  walltime, then reduce if the queue wait is too long:
+
+**Walltime negotiation flow:**
+
+1. Submit `salloc` with **max walltime** for the partition.
+2. Wait ~30 seconds for the scheduler to process, then check:
+   ```bash
+   <ssh_cmd> <ssh_host> "squeue -j <jobid> -h -o '%i %T %S'"
+   ```
+   - `%S` = estimated start time. Parse it to calculate wait.
+3. **If RUNNING** → proceed to next step.
+4. **If PENDING and wait ≤ 10 minutes** → wait for it, proceed.
+5. **If PENDING and wait > 10 minutes** (or no start time shown):
+   a. Infer how long you actually need from context (profiling → ~30 min,
+      iterative dev → ~2h, training run → match expected duration).
+      Round up to the next hour boundary.
+   b. Reduce walltime in-place:
+      ```bash
+      <ssh_cmd> <ssh_host> "scontrol update JobId=<jobid> TimeLimit=<new_time>"
+      ```
+   c. Tell the user: "Reduced walltime from Xh to Yh for faster scheduling —
+      you'll have less time on the allocation."
+   d. Wait ~15 seconds, re-check estimated start time.
+   e. **If still > 10 minutes:** inform user with the ETA and suggest
+      alternatives — different partition (e.g., backfill), fewer nodes.
+
+## Hang Detection Monitoring
+
+When monitoring a running job (polling via `squeue`), also check the remote log
+file every 60 seconds for the pattern `hang detected` (case-insensitive). This
+catches GPU hangs, NCCL deadlocks, and other stalls that keep the job running
+indefinitely without making progress.
+
+**Check for hang (run alongside each `squeue` poll once job is RUNNING):**
+
+```bash
+<ssh_cmd> <ssh_host> "grep -qi 'hang detected' <remote_work_dir>/slurm-<jobid>.out 2>/dev/null && echo HANG_DETECTED"
+```
+
+Also check the `.err` file if the output file has no match:
+
+```bash
+<ssh_cmd> <ssh_host> "grep -qi 'hang detected' <remote_work_dir>/slurm-<jobid>.err 2>/dev/null && echo HANG_DETECTED"
+```
+
+**If `HANG_DETECTED` is printed:**
+
+1. Immediately cancel the job:
+   ```bash
+   <ssh_cmd> <ssh_host> "scancel <jobid>"
+   ```
+2. Read the last 200 lines of the log file for context:
+   ```bash
+   <ssh_cmd> <ssh_host> "tail -n 200 <remote_work_dir>/slurm-<jobid>.out"
+   ```
+3. Report status as `HANG_DETECTED` with the relevant log lines. Do not
+   continue polling — the job is terminated.
+
+## Recovering from Failures
+
+**SSH connection drops:** Re-run SSH preflight (Step 2). tmux-held allocations
+survive disconnects.
+
+**Job stuck in PENDING:**
+
+```bash
+<ssh_cmd> <ssh_host> "squeue -j <jobid> -o '%i %T %r %S'"
+```
+
+Check the reason (`%r`) and estimated start (`%S`). If wait > 10 minutes,
+follow the **walltime negotiation flow** (see Walltime strategy). If already
+negotiated and still stuck, suggest alternatives: backfill partition, fewer
+nodes, or a different cluster.
+
+**Job fails:** Read last lines, never dump full output:
+
+```bash
+<ssh_cmd> <ssh_host> "tail -n 100 <remote_cwd>/slurm-<jobid>.out"
+<ssh_cmd> <ssh_host> "tail -n 50 <remote_cwd>/slurm-<jobid>.err"
+```
+
+**Hang detected:** If the log file contains `hang detected`, the job is
+auto-cancelled (see Hang Detection Monitoring section above). Report
+`HANG_DETECTED` with the last 200 lines of log context.
diff --git a/.claude/skills/exec-remote-slurm/scripts/build.sh b/.claude/skills/exec-remote-slurm/scripts/build.sh
new file mode 100644
index 000000000000..e0ef1fb9457e
--- /dev/null
+++ b/.claude/skills/exec-remote-slurm/scripts/build.sh
@@ -0,0 +1,42 @@
+#!/bin/bash
+
+# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# build.sh — Runs inside the container on the compute node.
+# Builds TensorRT-LLM from source using build_wheel.py.
+#
+# Usage: build.sh <repo_dir> [build_wheel_args...]
+#
+# Default build_wheel.py flags (used when no extra args are provided):
+#   -a "100-real" --nvtx
+# Any extra arguments after repo_dir are forwarded verbatim to build_wheel.py.
+
+set -euo pipefail
+
+repo_dir=${1:?Usage: build.sh <repo_dir> [build_wheel_args...]}
+shift
+
+cd "${repo_dir}"
+
+if [[ $# -gt 0 ]]; then
+    echo "[build.sh] Running: python3 ./scripts/build_wheel.py $*"
+    python3 ./scripts/build_wheel.py "$@"
+else
+    echo "[build.sh] Running default build"
+    python3 ./scripts/build_wheel.py \
+        -a "100-real" \
+        --nvtx
+fi
diff --git a/.claude/skills/exec-remote-slurm/scripts/build.slurm b/.claude/skills/exec-remote-slurm/scripts/build.slurm
new file mode 100644
index 000000000000..a588fbc4e0b6
--- /dev/null
+++ b/.claude/skills/exec-remote-slurm/scripts/build.slurm
@@ -0,0 +1,41 @@
+#!/bin/bash
+
+# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --time=02:00:00
+#SBATCH -o log/build.sbatch.%j.out
+#
+# build.slurm — SLURM batch script that launches a container and runs build.sh.
+# Usage: sbatch [sbatch_flags...] build.slurm <container_image> <user_root_dir> <scripts_dir> <repo_dir> [build_wheel_args...]
+
+container_image=$1
+user_root_dir=$2
+scripts_dir=$3
+repo_dir=$4
+shift 4  # remaining args forwarded to build.sh
+
+logdir=${scripts_dir}/log
+mkdir -p "${logdir}"
+
+srun \
+    --container-image="${container_image}" \
+    --container-mounts="${user_root_dir}:${user_root_dir}" \
+    --no-container-mount-home \
+    bash "${scripts_dir}/build.sh" "${repo_dir}" "$@" \
+    2>&1 | tee "${logdir}/build_${SLURM_JOB_ID}.srun.log"
diff --git a/.claude/skills/exec-slurm-compile/SKILL.md b/.claude/skills/exec-slurm-compile/SKILL.md
index 8a8005d06d7a..64bca208bee5 100644
--- a/.claude/skills/exec-slurm-compile/SKILL.md
+++ b/.claude/skills/exec-slurm-compile/SKILL.md
@@ -72,7 +72,7 @@ The user must provide (or you must ask for) these values:
 |-----------|-------------|---------|
 | `container_image` | Path to `.sqsh` container image (see enroot import above) | `/path/to/pytorch.sqsh` |
 | `repo_dir` | Path to the TensorRT-LLM repository | `/path/to/TensorRT-LLM` |
-| `mount_dir` | Top-level directory to bind-mount into the container | `/shared/users` |
+| `user_root_dir` | Top-level directory to bind-mount into the container | `/shared/users` |
 | `partition` | SLURM partition | `batch` |
 | `account` | SLURM account | `my_account` |
 
@@ -123,18 +123,18 @@ If the user already has a `.sqsh` file, skip this step.
 Ask the user for any missing prerequisite values listed above. At minimum you need:
 - `container_image` (or the Docker image URL — then run Step 0 first)
 - `repo_dir`
-- `mount_dir`
+- `user_root_dir`
 - `partition` and `account`
 
 If the user has used this workflow before, check if previous values are stored in memory files.
 
 ### Step 2: Prepare the Scripts Directory
 
-The compile scripts must be accessible from inside the container (i.e., under `mount_dir`). Either:
+The compile scripts must be accessible from inside the container (i.e., under `user_root_dir`). Either:
 
-**Option A** — Copy companion scripts to a location under `mount_dir`:
+**Option A** — Copy companion scripts to a location under `user_root_dir`:
 ```bash
-scripts_dir=<mount_dir>/<username>/workspace/tensorrt_llm_scripts
+scripts_dir=<user_root_dir>/<username>/workspace/tensorrt_llm_scripts
 mkdir -p ${scripts_dir}/log
 cp skills/exec-slurm-compile/scripts/compile.sh ${scripts_dir}/
 cp skills/exec-slurm-compile/scripts/compile.slurm ${scripts_dir}/
@@ -156,7 +156,7 @@ sbatch \
     --job-name=<jobname> \
     --time=<time_limit> \
     <scripts_dir>/compile.slurm \
-    <container_image> <mount_dir> <scripts_dir> <repo_dir>
+    <container_image> <user_root_dir> <scripts_dir> <repo_dir>
 ```
 
 Capture and report the job ID from the `sbatch` output.
@@ -241,7 +241,7 @@ Common architecture values:
 **Agent actions**:
 1. Ask for container image path, repo path, mount dir (if not known)
 2. Confirm partition/account for OCI cluster
-3. Copy scripts to accessible location under mount_dir
+3. Copy scripts to accessible location under user_root_dir
 4. Submit with `sbatch`
 5. Report job ID
 6. Monitor with `squeue` until complete
diff --git a/.claude/skills/exec-slurm-compile/scripts/compile.slurm b/.claude/skills/exec-slurm-compile/scripts/compile.slurm
index 1ad698c87b68..5ab38627c200 100644
--- a/.claude/skills/exec-slurm-compile/scripts/compile.slurm
+++ b/.claude/skills/exec-slurm-compile/scripts/compile.slurm
@@ -22,10 +22,10 @@
 #SBATCH -o log/compile.sbatch.%j.out
 #
 # compile.slurm — SLURM batch script that launches a container and runs compile.sh.
-# Usage: sbatch [sbatch_flags...] compile.slurm <container_image> <mount_dir> <scripts_dir> <repo_dir> [build_wheel_args...]
+# Usage: sbatch [sbatch_flags...] compile.slurm <container_image> <user_root_dir> <scripts_dir> <repo_dir> [build_wheel_args...]
 
 container_image=$1
-mount_dir=$2
+user_root_dir=$2
 scripts_dir=$3
 repo_dir=$4
 shift 4  # remaining args forwarded to compile.sh
@@ -35,7 +35,7 @@ mkdir -p "${logdir}"
 
 srun \
     --container-image="${container_image}" \
-    --container-mounts="${mount_dir}:${mount_dir}" \
+    --container-mounts="${user_root_dir}:${user_root_dir}" \
     --no-container-mount-home \
     --mpi=pmix \
     bash "${scripts_dir}/compile.sh" "${repo_dir}" "$@" \
diff --git a/.claude/skills/exec-slurm-compile/scripts/submit_compile.sh b/.claude/skills/exec-slurm-compile/scripts/submit_compile.sh
index 14cdb414ede6..4aa7ef55354f 100755
--- a/.claude/skills/exec-slurm-compile/scripts/submit_compile.sh
+++ b/.claude/skills/exec-slurm-compile/scripts/submit_compile.sh
@@ -23,7 +23,7 @@
 #   partition        — SLURM partition name
 #   account          — SLURM account
 #   jobname          — descriptive job name
-#   mount_dir        — top-level directory to bind-mount into the container
+#   user_root_dir        — top-level directory to bind-mount into the container
 #   scripts_dir      — directory containing compile.slurm and compile.sh
 #   repo_dir         — path to the TensorRT-LLM repo to compile
 #
@@ -37,7 +37,7 @@ partition="<PARTITION>"                     # e.g., batch
 account="<ACCOUNT>"                        # e.g., my_account
 jobname="<JOB_NAME>"                       # e.g., trtllm-compile.username
 
-mount_dir="<MOUNT_DIR>"                    # e.g., /shared/users
+user_root_dir="<USER_ROOT_DIR>"            # e.g., /shared/users
 scripts_dir="<SCRIPTS_DIR>"               # directory containing compile.slurm & compile.sh
 repo_dir="<REPO_DIR>"                      # path to TensorRT-LLM repo
 # ─────────────────────────────────────────────────────────────────────────
@@ -57,5 +57,5 @@ sbatch \
     --account="${account}" \
     --job-name="${jobname}" \
     "${scripts_dir}/compile.slurm" \
-    "${container_image}" "${mount_dir}" "${scripts_dir}" "${repo_dir}" \
+    "${container_image}" "${user_root_dir}" "${scripts_dir}" "${repo_dir}" \
     "${extra_build_args[@]}"
diff --git a/.claude/skills/kernel-cute-writing/references/api-core.md b/.claude/skills/kernel-cute-writing/references/api-core.md
index 9b640b236f1b..4945971a6a88 100644
--- a/.claude/skills/kernel-cute-writing/references/api-core.md
+++ b/.claude/skills/kernel-cute-writing/references/api-core.md
@@ -126,8 +126,7 @@ cute.is_major(x, d)    # True if dimension d is stride-1
 ```python
 cute.make_tensor(ptr, layout)                  # From pointer + layout
 cute.make_identity_tensor(shape)               # Coordinate mapping tensor
-cute.make_rmem_tensor(layout)                  # Register memory tensor
-cute.make_fragment(shape, dtype)               # Register fragment
+cute.make_rmem_tensor(layout_or_shape, dtype)  # Register memory tensor
 ```
 
 ### Data Initialization
diff --git a/.claude/skills/kernel-cute-writing/references/concepts-tensors.md b/.claude/skills/kernel-cute-writing/references/concepts-tensors.md
index 5b7c9c4d67d7..09e3e775bdef 100644
--- a/.claude/skills/kernel-cute-writing/references/concepts-tensors.md
+++ b/.claude/skills/kernel-cute-writing/references/concepts-tensors.md
@@ -81,8 +81,7 @@ mA = cute.make_tensor(ptr, layout=layout)
 
 ### Register/Fragment Tensors (Owning)
 ```python
-rmem = cute.make_rmem_tensor(layout)
-frag = cute.make_fragment(shape, dtype)
+rmem = cute.make_rmem_tensor(layout_or_shape, dtype)
 ```
 
 ## Accessing Tensors
diff --git a/.claude/skills/perf-analysis/SKILL.md b/.claude/skills/perf-analysis/SKILL.md
index 0989f7e49336..cc2b86a0146f 100644
--- a/.claude/skills/perf-analysis/SKILL.md
+++ b/.claude/skills/perf-analysis/SKILL.md
@@ -58,6 +58,12 @@ When diagnosing performance issues, classify the primary bottleneck:
 | **Communication-bound** | Significant time in collective operations | Limited by inter-GPU or inter-node communication |
 | **Sync-bound** | Excessive CPU-GPU synchronization points | Stalls from unnecessary synchronization |
 
+After classifying, consult the **perf-optimization-casebook** skill for
+candidate optimizations (decision precedents) matching the bottleneck class,
+model, and hardware. Use it to populate the Recommendations section with
+proven, adapted approaches rather than inventing them — but report only
+measured numbers, never a case's expected effect as an achieved result.
+
 ## Delegation Guidelines
 
 When delegating to specialists, describe the **desired outcome** -- not the
diff --git a/.claude/skills/perf-optimization-casebook/SKILL.md b/.claude/skills/perf-optimization-casebook/SKILL.md
new file mode 100644
index 000000000000..2bb8697be87b
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/SKILL.md
@@ -0,0 +1,293 @@
+---
+name: perf-optimization-casebook
+description: >
+  Casebook of past successful and classic TensorRT-LLM optimizations
+  (runtime/execution and kernel level) recorded as reusable decision
+  precedents. Consult when deciding
+  which optimization to apply for a classified bottleneck or a given
+  config/model/hardware, to find prior art and adapt a proven approach
+  instead of guessing. Each case records applicability signals, mechanism,
+  how to apply, expected effect, accuracy risk, verification, and rollback.
+tags:
+  - optimization
+  - casebook
+  - precedents
+  - decision-support
+license: Apache-2.0
+metadata:
+  author: NVIDIA Corporation
+---
+
+# Performance Optimization Casebook
+
+A curated library of **classic and previously-successful TensorRT-LLM
+optimizations**, written as decision precedents. Use it to answer "for *this*
+bottleneck / config / model / hardware, what has worked before, and how do I
+adapt it?" — instead of inventing an approach from scratch.
+
+This skill is **reference material, not a coordinator.** It does not run
+profiling, edit code, or route to specialists by itself. It is consulted *by*
+the coordinators (`perf-analysis` after a bottleneck is classified,
+`perf-optimization` before prioritizing/routing) and by the `perf-sweep`
+system when choosing what to try next. Implementation always goes through the
+relevant specialist/skill named in each case.
+
+## When to Consult
+
+- A bottleneck has been classified (compute / memory / launch / communication /
+  sync) and you need candidate optimizations for that class.
+- You are about to prioritize or route an optimization and want to check for
+  prior art, known pitfalls, and the right specialist.
+- You are choosing the next config to try in a serving sweep.
+- You just landed a verified win and want to record it for future reuse.
+
+## How to Use a Case — the loop
+
+1. **Classify.** Know the current bottleneck class and the workload context
+   (model arch, hardware/SM, ISL/OSL, concurrency, current config). If you do
+   not have this, get it from `perf-analysis` / profiling first.
+2. **Match.** Search wide with free text, then rank — never let the query
+   vocabulary narrow what is reachable:
+   - **Recall (grep the whole case files, free text):** grep the case bodies
+     with your own natural query terms — model name, knob, mechanism word,
+     deployment feature — over the full files, not just structured fields:
+
+         grep -ril "lora" references/                 # any term, whole file
+         grep -ril "chunked prefill" references/
+
+     Every concept a case covers lives in its prose (**Applies when** /
+     **Generalizes to** / mechanism), so a free-text grep over the whole file
+     is recall-complete on its own. Frontmatter is **additive** structure layered
+     on top — it never removes a word grep could already find. Do **not** pre-
+     translate the query into canonical vocabulary before searching; that is
+     what would make grep miss LoRA / beam search / a reasoning model's thinking
+     phase and other concepts that have no canonical signal. Browse the family
+     **index** too if you want a curated table, but grep alone does not miss.
+   - **Rank (optional accelerator):** among the hits, optionally use the
+     frontmatter fields (canonical `signals`/pattern grep — see "Frontmatter
+     fields" below) to decide which case file(s) to open first. This only
+     *reorders* the free-text hits — it never defines or bounds them.
+   - Then open **only** the winning case file(s). Match on signals, not on the
+     optimization name.
+3. **Adapt — do not copy.** A case is a precedent, not a guaranteed config.
+   Adjust knob values to the current model/hardware; read the nearest
+   checked-in config and the model's deployment guide before setting numbers.
+4. **Verify.** Apply through the named specialist/skill, then measure. Numbers
+   come from real measurement only (see Principles). A precedent's expected
+   *direction* of effect does not excuse skipping measurement.
+5. **Record.** If the change is a verified win (or a clear, instructive
+   failure), add or update a case so the next decision is better informed (see
+   "Adding a New Case").
+
+### Frontmatter fields (optional precision — not a search step)
+
+Cases carry YAML frontmatter (vocabulary: `data/tags.yaml`; spec:
+`references/case-template.md`). It is **additive structure for ranking and
+judging** the free-text hits from Match step 2 — never a second search you must
+run, and never a filter that bounds what grep already found. Use it as needed:
+grep a canonical `signals`/pattern id (`data/patterns.yaml`) to rank hits, and
+read a candidate's `eligibility:` / `interactions:` before recommending (respect
+`incompatible-with`; follow `depends-on` to the producer case first). A synonym
+won't hit these fields (they store canonical terms only) — that is fine, because
+the free-text grep already found the case; never emit `no-match` from an empty
+field grep.
+
+Legacy (not-yet-migrated) cases have no frontmatter — for those, fall back
+to the family index tables. When adding or editing a migrated case, keep the
+contract by hand: frontmatter values must be canonical terms from
+`data/tags.yaml` (never synonyms), and case↔pattern pointers must stay
+bidirectionally consistent (`patterns:` in the case ↔ `instances:` in
+`data/patterns.yaml`).
+
+### Consultation trace (machine-parseable — always emit)
+
+Every consultation MUST end with exactly ONE trace line in your visible reply
+text, as soon as the Match step resolves:
+
+    CASEBOOK: matched=<family>/<case-slug> confidence=<high|medium|low> adapted="<one-line adaptation>"
+
+or, when no precedent fits:
+
+    CASEBOOK: no-match query="<bottleneck class / signals you searched for>"
+
+`<case-slug>` is the case file name without `.md`. Emit the line even when the
+consultation is a dead end: it is the join key for automated audit trails
+(which precedents were consulted vs cited vs hallucinated), and the no-match
+rate is the casebook's coverage metric for maintainers. One line per consulted
+question; never summarize multiple matches into one line. A `no-match` line
+asserts your free-text recall grep over the whole case files (Match step 2)
+turned up nothing relevant — it must never be emitted merely because a
+canonical-field grep returned empty.
+
+## Case Families — routing table
+
+Open the family **index** relevant to your classified situation, then drill into a
+single case file from its case table (each family is a directory of one-file-per-case
+behind an `index.md` of patterns + a case-picker table).
+
+| Family | Index | Covers |
+|--------|-------|--------|
+| Runtime / execution | [references/runtime-execution/](references/runtime-execution/index.md) | Overlap scheduler, PDL, CUDA graphs (piecewise/padding, split-op-for-capture), multi-stream overlap, MLA KV-cache reuse & activation freeing, auxiliary-cache management, online-EPLB overlap, chunked prefill (aligned auxiliary), sparse-attention path selection, host-overhead & sync-free, batching/scheduler knobs, speculative decoding / MTP |
+| Kernel & fusion | [references/kernel-and-fusion/](references/kernel-and-fusion/index.md) | Op/epilogue fusion (AR+residual+norm+quant, add+norm+quant, QK-norm+RoPE, pre-pass folding, data-movement+quant), fused attention & MLA kernels, sparse-MLA top-k attention & top-k selection kernels, MoE grouped-GEMM backends (trtllm-gen / DeepGEMM / MegaMoE), low-precision GEMM (FP8 block-scale / W4A8 / NVFP4), ranking-only precision, routing kernels, custom Triton/CuTe/TileIR, GEMM tactics |
+| Communication | [references/communication/](references/communication/index.md) | EP all-to-all (DeepEP, low-precision dispatch/combine), TP AllReduce (MNNVL two-shot kernel, UserBuffers / symmetric memory, shape-aware autotune), collective kernel & strategy selection |
+| Case schema + template | [references/case-template.md](references/case-template.md) | The field schema and a worked example for adding cases |
+
+Many runtime/kernel/communication cases were distilled from the TRT-LLM
+DeepSeek-R1/V3 serving stack (MLA attention, fine-grained MoE, MTP, attention-DP +
+expert-parallel), where these levers cluster — but each is recorded as a
+**transferable pattern** (see *Generalizes to*); match on signals, not on the
+DeepSeek origin. Each carries a `Commits` line so you can `git show` the source.
+
+### Quick index — bottleneck signal → where to look
+
+Routing shortcut once the bottleneck is classified. Match the signal, open the
+family index, then match a case's **Applies when** / **Generalizes to** in its
+case table and open that case file.
+
+| Profile signal | Bottleneck | Family → candidate patterns |
+|---|---|---|
+| GPU idle between steps; host prep on critical path | launch / host | runtime → overlap scheduler · move per-step bookkeeping into C++/fused op · piecewise CUDA graph · CUDA-graph padding |
+| Many small **dependent** kernels, SM≥90 | launch | runtime → PDL · kernel → mega-fuse MoE / fuse collective epilogue |
+| Residual + Norm (+quant) chain of small kernels | memory / launch | kernel → fuse AR epilogue · fuse local add+norm+quant · QK-Norm+RoPE |
+| AllReduce / all-to-all a large share at TP/EP>1 | communication | communication → DeepEP · low-precision dispatch/combine · MNNVL two-shot kernel · UserBuffers · shape-aware autotune |
+| Low MFU under data-parallel with ragged load | host / compute | runtime → eliminate attention-DP padding · CUDA-graph padding |
+| KV reads dominate decode (long ctx / high concurrency) | memory | kernel → FP8 MLA KV · runtime → MLA KV-cache reuse · free MLA intermediates |
+| Sparse-attention model (DeepSeek V3.2 / DSA): full-seq attention wasteful; lightning-indexer + top-k on the hot path | compute / launch / memory | kernel → sparse-MLA top-k attention · specialize top-k selection kernel · ranking-only TF32 · fuse data-movement+quantize · runtime → skip sparse when short-seq degenerate · chunked prefill (aligned auxiliary) · auxiliary cache in KV manager · split custom op for piecewise capture |
+| MoE grouped-GEMM dominates the step | compute | kernel → trtllm-gen FP4 MoE · mega-fuse MoE · HW-matched low-precision GEMM · optimize routing kernel |
+| Expert load imbalanced across EP ranks | comm / sync | runtime → online-EPLB + overlap rebalance |
+| Spec-decode accept rate low / draft forward exposed | compute / launch | runtime → relaxed acceptance · two-model MTP-Eagle overlap |
+| OOM caps batch / KV budget | memory | runtime → free MLA intermediates · kernel → low-precision GEMM/KV |
+
+## Case Schema
+
+Every case records these fields (see `references/case-template.md` for the
+canonical form):
+
+- **Commits** *(optional)* — primary `<hash> subject (#PR)` plus folded-in
+  `related:` follow-ups, when the case is distilled from specific commits, so a
+  reader can `git show` the source. Omit when there is no upstream commit.
+- **Applies when** — bottleneck class + signals (model arch, hardware/SM,
+  workload shape, current config) that make the case relevant.
+- **Counter-signals** — when NOT to apply: auto-disable regimes, degenerate
+  shapes, negligible-share workloads, unmet requirements. Check this before
+  recommending a matched case.
+- **Mechanism** — why it helps, in one or two lines.
+- **Generalizes to** — the transferable *pattern* this case instantiates and the
+  adjacent situations it carries to (other model/arch, neighbouring op-chains,
+  different knob values), plus what to change when adapting. This is what an
+  agent matches on when it does **not** hit the exact instance — the casebook's
+  main job. Match on this pattern, not the case title.
+- **Apply via** — the exact config knob and/or the specialist/skill to delegate
+  implementation to. Never a hand-written runnable config invented here.
+- **Expected effect** — qualitative direction and the metric that should move.
+  Quantitative deltas only when measured (cite the run) — see Principles.
+- **Accuracy risk** — `lossless`, `lossy`, or `mixed` (a case bundling moves
+  of differing risk — it must state which part is lossy). Lossy cases require an on-disk
+  accuracy record plus a rollback criterion before being promoted as a "best"
+  config (mirrors the perf-sweep correctness gate).
+- **Verify** — what to measure to confirm the win and (for lossy) parity.
+- **Rollback** — how to revert and the regression trigger that mandates it.
+- **Prior art** — pointers to the owning skill, TRT-LLM repo docs, or PRs.
+
+## Principles
+
+1. **Adapt, don't copy.** Cases are precedents. Knob values, hardware
+   constraints, and model arch differ — re-derive numbers from the nearest
+   checked-in config and the model guide. Treat any value carried over from a
+   case as unverified until measured here.
+2. **Never fabricate numbers.** Quantitative before/after figures must come
+   from real profiling/benchmark output, with the source cited. A case states
+   the *direction* of effect and the *metric to watch*; it does not assert a
+   speedup that was not measured. Do not present casebook expectations as
+   achieved results.
+3. **Respect the accuracy gate.** Optimizations split into:
+   - **Lossless** (output-equivalent): scheduling/batching/graph knobs, and
+     backend *selection* (`attn_backend`, `moe_config.backend` — same math,
+     different implementation). Safe to evaluate on throughput/latency alone.
+   - **Lossy** (output-affecting): KV-cache quant (`kv_cache_dtype` fp8/int8/
+     nvfp4), weight/GEMM quant (`quant_algo`), low-precision MoE combine,
+     low-precision EP dispatch/combine, **relaxed spec-decode acceptance** — and
+     **PDL** (`TRTLLM_ENABLE_PDL`), which changes inter-kernel ordering/visibility
+     and has produced real NaN/accuracy regressions per kernel (verify per
+     kernel). Spec/MTP is output-equivalent *by design* but can drift on
+     acceptance/draft bugs, so verify on first enable. Lossy wins need an
+     accuracy record + rollback criterion, never promoted as "best" without it.
+   This split matches `perf-sweep-workflow` `core/gate.py` — keep them aligned.
+4. **One change at a time, measure each.** Combined changes hide regressions
+   and confound attribution. Roll back on >5% slowdown or any accuracy
+   regression outside tolerance.
+5. **Interface facts are as-of their pinned commit.** Knob names, config
+   bounds, file paths, and function names in a case describe the code at the
+   case's `commits:` — they are historical anchors, not present-tense truths.
+   Before acting on them (especially machine-evaluable `eligibility:`
+   entries), re-verify against YOUR TRT-LLM checkout; cases point to where
+   the live value lives (e.g. a `file::symbol` guard). A mismatch means the
+   code moved since the pin — prefer the live source. Mechanism and pattern
+   statements ("why it helps") are not version-bound; only interface facts
+   are.
+
+## Adding a New Case
+
+After a verified win (or an instructive failure):
+
+1. Pick the family directory (`references/<family>/`); if none fits, propose a new
+   family dir + `index.md` and add a row to the routing table above.
+2. Create `references/<family>/<slug>.md` from the schema in
+   `references/case-template.md` — H1 title + the breadcrumb line + every field
+   bullet — then add one row to that family index's case table (Case → file,
+   Applies-when, Generalizes-to, Risk), and, if it instantiates a listed pattern,
+   link it from that pattern's `_(Instance: …)_`. Put real measured deltas with the
+   source run; otherwise write "measured Δ to be recorded from <run>".
+3. Keep it a *decision precedent* — applicability, mechanism, knob/specialist,
+   risk, verification, rollback. Fill **Generalizes to** with the pattern the
+   case instantiates and 2–4 adjacent situations it carries to: this is what
+   lets an agent reuse the case when its situation is a *variation*, not the
+   exact instance (match on the pattern, not the title). Do **not** re-explain
+   subsystem architecture (how MoE/attention/KV cache work); link to the TRT-LLM
+   repo developer guides for that (see
+   `docs/design/agent_toolkit_organization.md` §2.3).
+4. Cross-link the owning specialist/skill so implementation stays delegated.
+
+### Ingestion discipline (agent-scanned corpus)
+
+This corpus is scanned and written by agents. Selection — "is this a
+transferable optimization precedent?" — stays with the writing agent's
+judgment; the rules below target known agent failure modes (fabricated
+detail, plausible over-generalization, cross-run divergence), not judgment.
+
+- **Trace, don't recall.** Every number in a case (bounds, caps, shapes,
+  dims) must be traceable to a specific line in a cited commit's diff or
+  the pinned source; every model/arch applicability claim must be checked
+  against the modeling source (e.g. the routing-method enum), never
+  inferred from model-family resemblance. If you didn't look it up, don't
+  write it.
+- **Fold vs. new.** A follow-up PR that iterates the *same decision* (same
+  pattern, same knob/kernel surface) folds into the existing case as a
+  `related:` commit — a PR series refining one optimization is ONE case.
+  A new case requires new applicability signals or a new pattern.
+- **Stub threshold.** If the four core slots (applicability signals,
+  mechanism, transferable pattern, apply-via) can be at least roughly
+  filled, record the case with `maturity: stub` so it enters the match
+  surface. If not even roughly — skip it entirely; do not park
+  half-guesses in the corpus.
+- **Verify before landing.** A case is added or upgraded only after an
+  independent fact-check pass against its cited commits (quantitative
+  claims, knob names, file paths, model-applicability claims). The writing
+  agent's recall is not evidence.
+- **Scan cursor.** Sweeps of TRT-LLM main record the last commit scanned
+  here and resume from it. Deduplicate against the corpus by grepping the
+  hash across case `commits:`/`Commits:` fields before distilling.
+  Current cursor: *not yet recorded — the pre-pilot corpus was swept ad
+  hoc; set this at the next sweep.*
+
+## Relationship to Other Skills
+
+- **perf-analysis** classifies the bottleneck → this casebook supplies
+  candidate optimizations for that class.
+- **perf-optimization** prioritizes and routes → consult this casebook for
+  prior art and the right specialist before delegating.
+- **perf-sweep-workflow** chooses configs to try → cases inform hypotheses;
+  the deterministic core still validates and the challenger still gates.
+- **Implementation** is always delegated to the specialist/skill named in the
+  case (kernel-*, perf-torch-cuda-graphs, perf-torch-sync-free,
+  perf-host-optimization, trtllm-serve-config-guide, etc.).
diff --git a/.claude/skills/perf-optimization-casebook/data/aliases.yaml b/.claude/skills/perf-optimization-casebook/data/aliases.yaml
new file mode 100644
index 000000000000..c6702472a139
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/data/aliases.yaml
@@ -0,0 +1,53 @@
+# Canonical → synonym map (pilot).
+#
+# Purpose: QUERY-TIME normalization. Before grepping case frontmatter,
+# translate the terms in your head (or the user's words / profiler output)
+# to the canonical term. Frontmatter itself must contain canonical terms
+# only (tags.yaml); a synonym in frontmatter is a validation error.
+#
+# Namespaces mirror tags.yaml. Every key below must be a canonical term in
+# tags.yaml; no synonym may collide with a canonical term or another
+# synonym (a colliding synonym would make normalization ambiguous).
+
+architectures:
+  sm80: [ampere, a100, sm_80]
+  sm89: [ada, l40, l40s, sm_89]
+  sm90: [hopper, h100, h200, sm_90, sm90a]
+  sm100: [blackwell, b200, gb200, sm_100, sm100a]   # NOT umma: instruction-set term (tcgen05), also on sm103
+  sm103: [b300, gb300, sm_103]
+
+bottleneck:
+  compute: [compute-bound, math-bound, tensor-core-bound, low-mfu]
+  memory: [memory-bound, bandwidth-bound, hbm-bound]
+  launch: [launch-bound, launch-overhead, kernel-launch-overhead]
+  communication: [comm-bound, communication-bound, nccl-bound]
+  sync: [sync-bound, synchronization, stream-serialization, blocking-sync]
+  host-overhead: [host-bound, cpu-bound, cpu-overhead, host-prep-exposed, python-overhead]
+
+model_scope:
+  mla: [multi-head-latent-attention, latent-attention]
+  moe: [mixture-of-experts, fused-moe, expert-parallel]
+  sparse-attention: [dsa, deepseek-sparse-attention, lightning-indexer, block-sparse-attention, nsa]
+  spec-decode: [speculative-decoding, mtp, eagle, mtp-eagle, draft-model]
+  deepseek-v3r1: [deepseek-v3, deepseek-r1, dsv3, dsr1, r1]
+  deepseek-v32: [deepseek-v3.2, dsv32, v3.2]
+  qwen3-moe: [qwen3-235b]   # NOT qwen3-next: distinct arch (hybrid linear-attention MoE, Renormalize routing)
+  kimi-k2: [kimi]
+
+signals:
+  # NB: inter-STEP bubbles (host overhead → overlap scheduler) and inter-KERNEL
+  # gaps within a step (launch-bound → fusion/PDL) are distinct signals.
+  gpu-idle-between-steps: [inter-step-gap, gpu-bubble]
+  host-prep-on-critical-path: [host-prep-exposed-ratio, cpu-on-critical-path]
+  many-small-kernels: [kernel-fragmentation, tiny-kernels, launch-storm, inter-kernel-gap]
+  hbm-roundtrip-between-ops: [materialized-intermediates, extra-memory-pass]
+  kv-reads-dominate: [kv-bandwidth-bound, kv-cache-bound]
+  long-context: [long-isl, long-sequence, large-kv-history]
+  small-batch-decode: [low-concurrency, latency-bound-decode, bs1]
+  allreduce-dominates: [ar-bound, tp-collective-bound]
+  alltoall-dominates: [a2a-bound, dispatch-combine-bound, ep-collective-bound]
+  expert-load-imbalance: [moe-load-imbalance, ragged-experts, eplb]
+  recompilation-churn: [triton-recompile, torch-compile-recompile, jit-churn]
+  slow-path-fallback: [pytorch-fallback, unfused-fallback, is-fused-false]
+  draft-forward-exposed: [draft-not-overlapped, mtp-exposed]
+  memory-capacity-bound: [oom, kv-budget-bound, workspace-bound]
diff --git a/.claude/skills/perf-optimization-casebook/data/patterns.yaml b/.claude/skills/perf-optimization-casebook/data/patterns.yaml
new file mode 100644
index 000000000000..c4b9c4039649
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/data/patterns.yaml
@@ -0,0 +1,406 @@
+# Pattern registry (pilot) — transferable patterns as first-class entries.
+#
+# Why: a case's `Generalizes to` prose often carries SEVERAL reusable
+# patterns (the sparse-MLA case carries three), but a case has only one
+# title and one index row — burying patterns flattens the match surface.
+# Registering each pattern with a stable id gives every pattern its own
+# match entry point: symptom → pattern → instance case(s).
+#
+# Hybrid-registry contract (KernelWiki DEC-1 style — keep by hand):
+#   - every id listed in a case's frontmatter `patterns:` must exist here;
+#   - every `instances:` entry here must resolve to an existing case file
+#     whose frontmatter lists this pattern back.
+#
+# Fields: id (stable, prefix `pattern-`), name (imperative one-liner),
+# bottleneck / signals (canonical terms from tags.yaml — the top-down hook),
+# families (where instances live today; not a constraint on transfer),
+# instances (case ids), notes (adapt delta / scope caveat, optional).
+
+patterns:
+  - id: pattern-host-prep-pipelining
+    name: Pipeline host prep for step N+1 against device compute of step N
+    bottleneck: [host-overhead, launch]
+    signals: [gpu-idle-between-steps, host-prep-on-critical-path]
+    families: [runtime-execution]
+    instances: [case-overlap-scheduler]
+    notes: >
+      Adapt by keeping step N+1 prep state disjoint from state mutated by
+      step N results — the disjoint-state plumbing is the hard part.
+
+  - id: pattern-fuse-small-op-chain
+    name: Collapse a small launch-bound op-chain into one fused per-token kernel
+    bottleneck: [launch, memory]
+    signals: [many-small-kernels, hbm-roundtrip-between-ops]
+    families: [kernel-and-fusion]
+    instances: [case-fuse-moe-routing-kernel]
+    notes: >
+      Keep intermediates in registers/shared instead of HBM. Carries to any
+      per-token "activation → reduction → select" chain (routers, sampling,
+      RoPE/norm epilogues, KV scatter).
+
+  - id: pattern-reuse-warp-primitive
+    name: Reuse an existing purpose-built warp/reduction primitive over a generic algorithm
+    bottleneck: [compute, launch]
+    signals: [many-small-kernels, kernel-far-from-sol]
+    families: [kernel-and-fusion]
+    instances: [case-fuse-moe-routing-kernel]
+    notes: >
+      E.g. moeTopKFuncs.cuh::reduceTopK instead of a bitonic sort. Search
+      existing kernels for a primitive before writing a new one.
+
+  - id: pattern-bounded-fastpath-fallback
+    name: Guard a specialized fast path with explicit config bounds plus a generic fallback
+    bottleneck: [launch]
+    signals: [slow-path-fallback]
+    families: [kernel-and-fusion]
+    instances: [case-fuse-moe-routing-kernel]
+    notes: >
+      The bounds guard is part of the pattern, not incidental: the host-side
+      eligibility check must stay in sync with the kernel's hard-checked
+      supported bounds, or a config errors instead of falling back.
+
+  - id: pattern-attend-selected-kv-subset
+    name: Make the attention kernel operate on a selected KV subset, not the full sequence
+    bottleneck: [compute, memory]
+    signals: [attention-hot-path, kv-reads-dominate, long-context]
+    families: [kernel-and-fusion]
+    instances: [case-sparse-mla-topk-attention]
+    notes: >
+      Pass selected page/block indices; bound the KV loop AND the cross-CTA
+      reduction by the selection count. Carries to block-sparse,
+      sliding-window/streaming, landmark/NSA-style selection.
+
+  - id: pattern-write-direct-into-paged-cache
+    name: Write K/V straight into the paged cache to skip a materialization pass
+    bottleneck: [memory]
+    signals: [hbm-roundtrip-between-ops, kv-reads-dominate]
+    families: [kernel-and-fusion]
+    instances: [case-sparse-mla-topk-attention]
+    notes: >
+      Absorption/latent layout; also lets redundant scratch buffers be
+      dropped once K/V live in the cache.
+
+  - id: pattern-early-exit-ragged-bounds
+    name: Early-exit CTAs beyond the ragged per-sequence length so padded work is free
+    bottleneck: [compute]
+    signals: [attention-hot-path, high-concurrency]
+    families: [kernel-and-fusion]
+    instances: [case-sparse-mla-topk-attention]
+    notes: >
+      Applies to any kernel launched on a padded/max-shape grid over ragged
+      batches — clamp per-CTA work by the actual sequence bound.
+
+  # ---- runtime-execution (migration round) ----
+
+  - id: pattern-push-language-boundary-off-hot-path
+    name: Push Python/C++ language-boundary crossings off the host hot path
+    bottleneck: [host-overhead]
+    signals: [gpu-idle-between-steps, host-prep-on-critical-path]
+    families: [runtime-execution]
+    instances: [case-pybind-wrapper-pure-python]
+    notes: >
+      Pure-Python hot objects, py_*-mirrored C++ fields read once, and
+      bulk-serialize-once + lazy-deserialize for anything crossing a process
+      boundary.
+
+  - id: pattern-hoist-step-invariant-host-work
+    name: Hoist step-invariant work out of the per-layer loop
+    bottleneck: [host-overhead]
+    signals: [gpu-idle-between-steps, host-prep-on-critical-path]
+    families: [runtime-execution]
+    instances: [case-cache-step-invariant-per-layer]
+    notes: >
+      Cache once per step on the metadata object; invalidate at step start
+      with a cheap flag plus identity guard.
+
+  - id: pattern-stable-torch-compile-target
+    name: Keep a torch.compile target stable and module-level; never redefine it on the hot path
+    bottleneck: [host-overhead]
+    signals: [recompilation-churn, host-prep-on-critical-path]
+    families: [runtime-execution]
+    instances: [case-hoist-torch-compile-closures]
+    notes: >
+      Hoist the decorated function to a stable def passing only needed
+      tensors; drop the decorator entirely for single cheap ops.
+
+  - id: pattern-overlap-independent-work
+    name: Overlap two data-independent pieces of work on separate streams or engines
+    bottleneck: [launch, sync]
+    signals: [small-batch-decode, many-small-kernels, draft-forward-exposed, expert-load-imbalance]
+    families: [runtime-execution]
+    instances: [case-multi-stream-shared-routed-expert, case-overlap-mla-rope-uk-bgemm, case-overlap-online-eplb, case-two-model-mtp-eagle]
+    notes: >
+      The hard part is proving data-independence (disjoint output slices /
+      disjoint step state); stream-switching only pays off under CUDA-graph
+      capture.
+
+  - id: pattern-process-only-real-work
+    name: Process only the real work, not padded-to-max
+    bottleneck: [host-overhead, compute]
+    signals: [expert-load-imbalance, host-prep-on-critical-path]
+    families: [runtime-execution]
+    instances: [case-attention-dp-padding]
+    notes: >
+      Make the collective length-aware and remove downstream padded-shape
+      assumptions; carries to any DP/ragged-batch collective and any resource
+      sized off max_*.
+
+  - id: pattern-pad-to-captured-bucket
+    name: Pad the dynamic dimension to a captured bucket so the CUDA graph always matches
+    bottleneck: [launch]
+    signals: [slow-path-fallback, expert-load-imbalance]
+    families: [runtime-execution]
+    instances: [case-cuda-graph-padding]
+    notes: >
+      Pad with shape-correct dummies up to the nearest captured size; dummies
+      need reserved IDs and cleanup.
+
+  - id: pattern-reuse-computed-state
+    name: Reuse computed state instead of recomputing it; free large intermediates at last use
+    bottleneck: [compute, memory]
+    signals: [attention-hot-path, long-context, memory-capacity-bound]
+    families: [runtime-execution]
+    instances: [case-mla-kv-cache-reuse, case-free-mla-intermediates]
+    notes: >
+      Choose what to persist (full KV vs compressed latent) and gate on
+      kernel availability; for freeing, find the true last-use point of each
+      big tensor.
+
+  - id: pattern-pdl-producer-consumer-overlap
+    name: Overlap a consumer kernel's preamble with the producer's tail via PDL — verify accuracy per kernel
+    bottleneck: [launch]
+    signals: [many-small-kernels, small-batch-decode]
+    families: [runtime-execution]
+    instances: [case-pdl]
+    notes: >
+      Gate on sm90+; accuracy-check each kernel — PDL changes inter-kernel
+      ordering and has produced real NaN regressions.
+
+  - id: pattern-bounded-relaxed-acceptance
+    name: Trade exact-match draft verification for a bounded relaxed criterion where small divergence is tolerable
+    bottleneck: [compute]
+    signals: [low-accept-rate]
+    families: [runtime-execution]
+    instances: [case-relaxed-mtp-acceptance]
+    notes: >
+      Gate relaxed acceptance to a delimited low-stakes phase (e.g. thinking)
+      and tune topk/delta; lossy — needs accuracy record + rollback.
+
+  - id: pattern-fold-side-cache-into-kv-manager
+    name: Fold a model-specific auxiliary cache into the unified paged KV-cache manager as a typed pool
+    bottleneck: [memory]
+    signals: [memory-capacity-bound, long-context]
+    families: [runtime-execution]
+    instances: [case-auxiliary-cache-in-kv-manager]
+    notes: >
+      Register the pool type sharing block IDs with the main KV; sum
+      mixed-dtype pools additively by physical dtype.
+
+  - id: pattern-chunk-prefill-align-auxiliary
+    name: Chunk long prefill and align any auxiliary structure's chunking to the primary's boundaries
+    bottleneck: [memory]
+    signals: [memory-capacity-bound, long-context]
+    families: [runtime-execution]
+    instances: [case-chunked-prefill-aligned-auxiliary]
+    notes: >
+      The auxiliary gathers only the current chunk and inherits the primary
+      chunker's boundaries — never two chunking schemes.
+
+  - id: pattern-skip-degenerate-sparse-path
+    name: Skip the sparse/approximate path when it degenerates to the exact one below a threshold
+    bottleneck: [compute, launch]
+    signals: [attention-hot-path, small-batch-decode]
+    families: [runtime-execution]
+    instances: [case-skip-sparse-path-when-degenerate]
+    notes: >
+      Substitute a trivially-constructed exact result and add a
+      CUDA-graph-key dimension so the fast path stays captured.
+
+  - id: pattern-graph-the-graphable
+    name: Graph the graphable regions and run the un-capturable ops eagerly (piecewise capture)
+    bottleneck: [launch]
+    signals: [many-small-kernels, slow-path-fallback]
+    families: [runtime-execution]
+    instances: [case-piecewise-cuda-graph]
+    notes: >
+      Choose the split op, ensure safe resume after it, and enumerate
+      token-count buckets.
+
+  - id: pattern-split-op-at-graphable-seam
+    name: Split a monolithic custom op at the graphable/non-graphable seam so its shape-static half can be captured
+    bottleneck: [launch]
+    signals: [many-small-kernels, slow-path-fallback]
+    families: [runtime-execution]
+    instances: [case-split-custom-op-for-piecewise-capture]
+    notes: >
+      Route through a registered torch.ops op with constant-arity
+      register_fake and straight-line control flow under compile.
+
+  - id: pattern-fold-bookkeeping-into-consuming-op
+    name: Push per-step Python tensor wrangling and count-deriving syncs into the kernel/op that already runs
+    bottleneck: [host-overhead, sync]
+    signals: [host-prep-on-critical-path, gpu-idle-between-steps]
+    families: [runtime-execution]
+    instances: [case-move-bookkeeping-into-cpp-op]
+    notes: >
+      Widen the op signature to accept raw tensors plus a length; thread
+      known counts through instead of .item() syncs.
+
+  # ---- kernel-and-fusion (migration round) ----
+
+  - id: pattern-fuse-chain-feeding-gemm
+    name: Fuse the op chain that feeds the next GEMM (collective epilogue, add+norm+quant, QK-norm+RoPE, layout pre-pass, up to a mega-fused pipeline)
+    bottleneck: [launch, memory, communication]
+    signals: [many-small-kernels, hbm-roundtrip-between-ops, allreduce-dominates]
+    families: [kernel-and-fusion]
+    instances: [case-fuse-ar-epilogue, case-fuse-add-norm-quant, case-fuse-qk-norm-rope, case-fold-scale-swizzle-into-kernel, case-mega-fuse-moe-deepgemm]
+    notes: >
+      A quant step inside the fused chain makes it lossy (parity-check);
+      most land as torch.compile PatternMatcher rewrites or a single custom
+      op.
+
+  - id: pattern-quantize-memory-bound-side
+    name: Quantize the memory-bound side (KV cache + matching operand) while keeping the precision-sensitive output high
+    bottleneck: [memory]
+    signals: [kv-reads-dominate, long-context, high-concurrency]
+    families: [kernel-and-fusion]
+    instances: [case-fp8-mla-kv-cache]
+    notes: >
+      Choose a KV dtype the FMHA kernel supports, supply dequant/BMM scales,
+      keep the output dtype unquantized where the next op is sensitive.
+
+  - id: pattern-split-compute-from-reduction
+    name: Split a monolithic compute + serial cross-partition-reduction kernel into main + reduction kernels
+    bottleneck: [compute]
+    signals: [attention-hot-path, kernel-far-from-sol]
+    families: [kernel-and-fusion]
+    instances: [case-split-mla-reduction-kernel]
+    notes: >
+      Choose where to materialize partials (GMEM stats/partial-O) and gate on
+      the regime where the extra launch pays off.
+
+  - id: pattern-hw-matched-lowprec-gemm-backend
+    name: Pick the hardware-matched low-precision (grouped-)GEMM backend over the generic one
+    bottleneck: [compute, memory]
+    signals: [gemm-dominates-step, weight-reads-dominate]
+    families: [kernel-and-fusion]
+    instances: [case-trtllm-gen-fp4-moe-backend, case-hw-matched-lowprec-moe-gemm, case-mega-fuse-moe-deepgemm]
+    notes: >
+      Match format and scale layout to the SM's tensor-core support; keep a
+      non-quant fallback backend.
+
+  - id: pattern-specialize-hotpath-topk
+    name: Replace a hot-path generic-sort top-k with a specialized radix-select dispatched by (k, N, dtype, SM-wave-occupancy)
+    bottleneck: [launch, compute]
+    signals: [many-small-kernels, recompilation-churn]
+    families: [kernel-and-fusion]
+    instances: [case-specialize-topk-selection-kernel]
+    notes: >
+      Two exact short-circuits travel with it — a coarse low-precision first
+      pass with exact fallback, and a warm-start from the previous decode
+      step's answer.
+
+  - id: pattern-fuse-datamovement-into-quantize
+    name: Fuse a data-movement/layout prologue into the quantize that consumes it, reading strided views directly
+    bottleneck: [memory, launch]
+    signals: [many-small-kernels, hbm-roundtrip-between-ops]
+    families: [kernel-and-fusion]
+    instances: [case-fuse-datamovement-into-quantize]
+    notes: >
+      Accept strided input pointers, do the layout in registers, emit
+      quantized data plus scales in one pass; keep the unfused path as
+      reference.
+
+  - id: pattern-ranking-only-reduced-precision
+    name: Run a ranking-only computation on the fast tensor-core path and verify the accelerated kernel is actually dispatched
+    bottleneck: [compute]
+    signals: [kernel-far-from-sol, slow-path-fallback]
+    families: [kernel-and-fusion]
+    instances: [case-ranking-only-precision-tf32]
+    notes: >
+      Confirm the output only ranks, then verify dispatch with ncu — small-M
+      GEMMs silently fall back to CUDA-core SGEMM even when TF32 was
+      requested.
+
+  - id: pattern-reevaluate-fusion-boundary
+    name: Re-evaluate the GEMM fusion boundary per dtype/checkpoint — fusion is not monotonic
+    bottleneck: [compute, launch]
+    signals: [mixed-dtype-checkpoint, many-small-kernels]
+    families: [kernel-and-fusion]
+    instances: [case-reevaluate-fusion-boundary-per-dtype]
+    notes: >
+      Gate the fuse on a runtime dtype check and weigh whether keeping an op
+      separate unlocks a conditional skip.
+
+  - id: pattern-triton-to-precompiled-cpp-op
+    name: Move a host-bound Triton kernel to a precompiled C++/CUDA custom op with a register_fake
+    bottleneck: [host-overhead, launch]
+    signals: [many-small-kernels, recompilation-churn, host-prep-on-critical-path]
+    families: [kernel-and-fusion]
+    instances: [case-triton-to-cpp-op]
+    notes: >
+      Only worth it when host overhead (not GPU SOL) is the bottleneck; keep
+      the Triton version as reference/fallback.
+
+  - id: pattern-constexpr-only-when-required
+    name: Mark a Triton param tl.constexpr only when it must be compile-time; pass the rest as runtime args
+    bottleneck: [host-overhead]
+    signals: [recompilation-churn]
+    families: [kernel-and-fusion]
+    instances: [case-relax-tl-constexpr]
+    notes: >
+      Check each constexpr param for a genuine compile-time requirement
+      (tl.arange bounds, static unrolling) before relaxing it.
+
+  # ---- communication (migration round) ----
+
+  - id: pattern-swap-generic-collective-for-specialized-comm
+    name: Swap a generic collective for a workload-specialized comm library
+    bottleneck: [communication]
+    signals: [alltoall-dominates, slow-path-fallback]
+    families: [communication]
+    instances: [case-deepep]
+    notes: >
+      Match the backend's buffer-reservation and layout-adapter contract;
+      pick the CUDA-graph-safe variant for decode.
+
+  - id: pattern-send-activations-post-quant
+    name: Send activations post-quant, not pre-quant
+    bottleneck: [communication]
+    signals: [alltoall-dominates, full-precision-on-wire]
+    families: [communication]
+    instances: [case-low-precision-dispatch]
+    notes: >
+      Pack scale factors with the payload, re-swizzle on receive, honor
+      per-format hidden-size divisibility; lossy — needs an accuracy record
+      plus rollback criterion.
+
+  - id: pattern-tune-collective-kernel-not-algorithm
+    name: Make the collective kernel faster instead of changing the algorithm
+    bottleneck: [communication, sync]
+    signals: [allreduce-dominates, kernel-far-from-sol]
+    families: [communication]
+    instances: [case-mnnvl-twoshot-allreduce]
+    notes: >
+      Match the target collective's alignment/divisibility assumptions and
+      the buffer-flag lifecycle.
+
+  - id: pattern-preregister-symmetric-comm-buffers
+    name: Pre-register symmetric comm buffers to kill per-step allocation and unlock in-place fused collectives
+    bottleneck: [communication, launch, memory]
+    signals: [allreduce-dominates, hbm-roundtrip-between-ops]
+    families: [communication]
+    instances: [case-userbuffers-symmetric-memory]
+    notes: >
+      Route producers to write directly into the registered buffer; gate on
+      dtype support.
+
+  - id: pattern-autotune-collective-per-shape
+    name: Autotune or heuristically pick the collective implementation per shape and concurrency
+    bottleneck: [communication]
+    signals: [allreduce-dominates, high-concurrency]
+    families: [communication]
+    instances: [case-shape-aware-allreduce-autotune]
+    notes: >
+      Benchmark candidates offline into a LUT or register a TunableRunner
+      with the runtime AutoTuner.
diff --git a/.claude/skills/perf-optimization-casebook/data/tags.yaml b/.claude/skills/perf-optimization-casebook/data/tags.yaml
new file mode 100644
index 000000000000..c0400d827641
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/data/tags.yaml
@@ -0,0 +1,116 @@
+# Controlled vocabulary for casebook case frontmatter (pilot).
+#
+# Case frontmatter fields MUST use canonical terms from this file only.
+# Synonyms live in aliases.yaml and are for QUERY-TIME normalization —
+# never write a synonym into frontmatter. Check by hand when adding a case.
+#
+# Extend by adding a term here first (one term, one meaning); do not invent
+# terms inline in a case file.
+
+family:
+  - runtime-execution
+  - kernel-and-fusion
+  - communication
+
+# Bottleneck classes — aligned with the perf-analysis classification
+# (compute / memory / launch / communication / sync), plus host-overhead,
+# which perf-host-analysis detects as its own verdict and the casebook
+# uses pervasively.
+bottleneck:
+  - compute
+  - memory
+  - launch
+  - communication
+  - sync
+  - host-overhead
+
+# Profile signals — observable evidence that makes a case relevant.
+# A case lists 2–5. Keep this list short; a signal earns its place by
+# appearing in nsys/ncu/log evidence, not by being a paraphrase of a
+# bottleneck class.
+signals:
+  - gpu-idle-between-steps        # nsys: bubbles between forward steps
+  - host-prep-on-critical-path    # nsys: CPU prep serialized before launch
+  - many-small-kernels            # nsys: dense stream of tiny kernels
+  - hbm-roundtrip-between-ops     # intermediates materialized to HBM between ops
+  - kernel-far-from-sol           # ncu: hot kernel well below SOL%
+  - attention-hot-path            # attention kernel dominates the step
+  - routing-hot-path              # MoE router/gating dominates or fragments the step
+  - kv-reads-dominate             # KV-cache bytes dominate decode
+  - long-context                  # large ISL / long KV history
+  - high-concurrency              # many in-flight requests
+  - small-batch-decode            # low-concurrency / latency-bound decode
+  - allreduce-dominates           # TP collective a large share of step time
+  - alltoall-dominates            # EP dispatch/combine a large share of step time
+  - expert-load-imbalance         # ragged expert/DP load across ranks
+  - recompilation-churn           # Triton/torch.compile recompiles on hot path
+  - slow-path-fallback            # config silently fell back to a generic/slow path
+  - draft-forward-exposed         # spec-decode draft not hidden behind target
+  - memory-capacity-bound         # OOM / KV budget caps batch size
+  - gemm-dominates-step           # nsys: (grouped-)GEMM is the largest share of the step
+  - weight-reads-dominate         # decode: expert/weight bytes dominate HBM traffic
+  - mixed-dtype-checkpoint        # checkpoint puts fused-GEMM participants in differing dtypes/quant formats
+  - full-precision-on-wire        # activations cross the collective at bf16/fp16 despite a quantized model
+  - low-accept-rate               # spec-decode acceptance rate low, capping speedup
+
+architectures:
+  - sm80
+  - sm89
+  - sm90
+  - sm100
+  - sm103
+  - any-sm
+
+# Model/architecture applicability. Include `model-agnostic` when the
+# pattern transfers beyond the origin models (most cases — the casebook's
+# whole point); list the origin models too so exact-instance queries hit.
+model_scope:
+  - model-agnostic
+  - dense
+  - moe
+  - mla
+  - gqa
+  - sparse-attention
+  - spec-decode
+  - deepseek-v3r1
+  - deepseek-v32
+  - qwen3-moe
+  - kimi-k2
+  - llama-family
+  - gemma3
+
+phase:
+  - prefill
+  - decode
+  - any-phase
+
+accuracy_risk:
+  - lossless
+  - lossy
+  - mixed
+
+# stub  = skeleton entry, fields present but mechanism not yet distilled
+# full  = distilled from commits/PRs with mechanism + verify + rollback
+# verified-locally = additionally re-measured here, with a run recorded
+maturity:
+  - stub
+  - full
+  - verified-locally
+
+apply_via_kind:
+  - config-knob
+  - env-var
+  - default-on
+  - kernel-change
+  - code-change
+
+# Typed relations for the frontmatter `interactions:` field.
+# depends-on         this case consumes what the target produces (order matters)
+# composes-with      independently applicable; combined effect is the point
+# incompatible-with  mutually exclusive (mechanism auto-disables or breaks)
+# alternative-to     different lever for the same bottleneck; pick one
+relations:
+  - depends-on
+  - composes-with
+  - incompatible-with
+  - alternative-to
diff --git a/.claude/skills/perf-optimization-casebook/references/case-template.md b/.claude/skills/perf-optimization-casebook/references/case-template.md
new file mode 100644
index 000000000000..00ee89883548
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/case-template.md
@@ -0,0 +1,211 @@
+# Case Schema & Template
+
+Every casebook entry follows this schema. Copy the template below and fill
+**every** field. A case is a *decision precedent* — applicability, mechanism,
+how to apply, expected effect, accuracy risk, verification, rollback — not a
+re-explanation of how a subsystem works and not a guaranteed config.
+
+## Machine-readable frontmatter (pilot)
+
+New and migrated cases carry a YAML frontmatter block **above the H1**. It
+replaces the free-text `Tags:` bullet as the match surface: every value is a
+canonical term from `data/tags.yaml` (synonyms belong in `data/aliases.yaml`,
+never in frontmatter), and `patterns:` points into the registry
+`data/patterns.yaml` — split out **each** transferable pattern the case
+instantiates (a case often carries 2–3; each is its own match entry point).
+Keep three invariants by hand when adding or editing a case: vocab
+membership (canonical terms only), `id`/`family` structural identity
+(`id` = `case-` + filename stem; `family` = parent directory), and
+case↔pattern bidirectional consistency with `data/patterns.yaml`. Skeleton:
+
+```yaml
+---
+id: case-<file-slug>              # must equal "case-" + filename stem
+type: case
+family: <dir name>                # runtime-execution | kernel-and-fusion | communication
+maturity: stub | full | verified-locally
+bottleneck: [<tags.yaml::bottleneck>]     # 1–3
+signals: [<tags.yaml::signals>]           # 2–5 observable profile signals
+architectures: [<tags.yaml::architectures>]
+model_scope: [<tags.yaml::model_scope>]   # include model-agnostic if it transfers
+phase: [prefill | decode | any-phase]
+patterns: [<data/patterns.yaml ids>]      # EVERY pattern the case instantiates
+accuracy_risk: lossless | lossy | mixed
+apply_via_kind: [config-knob | env-var | default-on | kernel-change | code-change]
+knobs: []                         # exact knob/env names, if any
+specialists: []                   # owning specialist/skill names
+commits: []                       # 10–12 char hashes, quoted (YAML eats digits)
+log_markers: []                   # verbatim strings to grep in server logs
+                                  #   (e.g. a fallback warning) — log evidence
+eligibility: []                   # machine-checkable applicability constraints
+                                  #   an agent can evaluate against the config.
+                                  #   Volatile bounds: record a source pointer
+                                  #   (file::symbol of the live guard) PLUS the
+                                  #   values "as of <commit>" — never a bare
+                                  #   number with no anchor (it reads as
+                                  #   present-tense truth and rots silently)
+interactions: []                  # typed relations; each entry:
+                                  #   {case: case-<slug> | feature: <name>,
+                                  #    relation: <tags.yaml::relations>, note: "..."}
+                                  #   quote notes containing commas (YAML flow)
+measured: []                      # verified-result records, appended at Record
+                                  #   time; each entry requires ALL six fields:
+                                  #   {metric, value, gpu, workload, run_ref, date}
+                                  #   — no partial records (anti-fabrication)
+---
+```
+
+The human-readable bullets below remain the case body; frontmatter carries
+the *matching* metadata, bullets carry the *reasoning*. Legacy cases still
+use the `Tags:` bullet until migrated.
+
+## Field definitions
+
+- **Tags** — *(legacy — replaced by frontmatter on migrated cases.)*
+  At-a-glance match line: family, bottleneck class(es), and
+  model/hardware applicability, comma-separated. Terse; it aids scanning while
+  the bullets below carry the detail.
+- **Commits** — *Optional provenance.* When the case is distilled from specific
+  TRT-LLM commits, list the primary `<12-char hash> <subject> (#PR)` and any
+  `related:` follow-ups (the iterations folded into this one case). Lets a reader
+  `git show <hash>` the exact change. Omit for cases recorded from a local win
+  with no upstream commit yet; `Prior art` then carries the pointers.
+- **Applies when** — The bottleneck class (compute / memory / launch /
+  communication / sync) plus the signals that make this case relevant: model
+  architecture (dense / MoE, GQA / MLA), hardware/SM, workload shape (ISL/OSL,
+  concurrency), and the current config state. Match on these signals, not on
+  the optimization's name.
+- **Counter-signals** — When NOT to apply: regimes where the mechanism
+  auto-disables, degenerates, or actively hurts; workload shapes where the
+  targeted cost is a negligible share; requirements the case cannot meet
+  (e.g. bit-exact tie-breaking). Negative knowledge prevents misapplication —
+  an agent matching only positive signals will over-apply.
+- **Mechanism** — Why it helps, in one or two lines. When the case bundles
+  several numbered sub-moves (each with its own commit), keep **Accuracy
+  risk / Verify / Rollback decomposed by the same numbering** so a reader
+  knows which move to revert when one is implicated.
+- **Generalizes to** — The transferable *pattern* this case instantiates and the
+  adjacent situations it carries to: other model/architecture, neighbouring
+  op-chains, different knob values — plus what must change when adapting. An
+  agent rarely hits the exact instance; this field is how it recognises that its
+  situation is a *variation* of a known precedent. Name the pattern, list 2–4
+  transfer targets, and state the adapt delta. Match on this, not the title.
+- **Apply via** — The exact config knob(s) and/or the specialist/skill that
+  owns the implementation. Never an invented runnable config; delegate.
+- **Expected effect** — The qualitative direction of the effect and the metric
+  that should move. Put a quantitative delta **only** if it was measured, and
+  cite the run. Otherwise write "measured Δ to be recorded from <run>".
+- **Accuracy risk** — `lossless` (output-equivalent), `lossy`
+  (output-affecting), or `mixed` (a case bundling moves of differing risk —
+  state explicitly which part is lossy). Lossy cases/parts require an on-disk
+  accuracy record plus a rollback criterion before promotion as a "best"
+  config. This mirrors the `perf-sweep-workflow` correctness gate
+  (`core/gate.py`).
+- **Verify** — What to measure to confirm the win; for lossy cases, also the
+  parity/accuracy check. Include at least one **executable pointer** in
+  backticks — a test path, a profiling command, or the knob to A/B — not only
+  a described observation.
+- **Rollback** — How to revert and the regression trigger that mandates it
+  (default: >5% slowdown or any accuracy regression outside tolerance).
+- **Prior art** — Pointers to the owning skill, TRT-LLM repo developer guides,
+  or PRs. Keep links, not copied architecture text.
+
+## Template
+
+Each case is its **own file** at `references/<family>/<slug>.md`: an H1 title
+(short and descriptive), a one-line breadcrumb back to the family index and this
+schema, then every field as a bold bullet. Copy this raw form:
+
+```
+---
+<frontmatter — copy the skeleton from "Machine-readable frontmatter (pilot)"
+above and fill every field with canonical terms from data/tags.yaml>
+---
+
+# <short descriptive title>
+
+> Part of the [<Family> casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** *(optional)* `<12-char hash>` <subject> (#PR); related:
+  `<hash>` <subject> (#PR).
+- **Applies when:** <bottleneck class> + <signals: model arch, SM, ISL/OSL,
+  concurrency, current config>.
+- **Counter-signals:** <when NOT to apply: auto-disable regimes, degenerate
+  shapes, negligible-share workloads, unmet requirements>.
+- **Mechanism:** <why it helps>.
+- **Generalizes to:** <the pattern this instantiates>; carries to <2–4 adjacent
+  targets: other arch / op-chains / knob values>; adapt by <what changes>.
+- **Apply via:** `<config knob>` and/or delegate to **<specialist/skill>**.
+- **Expected effect:** <direction> in <metric>. Measured Δ to be recorded
+  from <run>, or: <metric> X → Y (source: <run/PR/doc>).
+- **Accuracy risk:** lossless | lossy | mixed (if lossy/mixed: requires
+  accuracy record + rollback criterion before promotion).
+- **Verify:** <metric(s) to measure>; <parity check if lossy>.
+- **Rollback:** <how to revert>; trigger: <regression condition>.
+- **Prior art:** <skill / repo doc / PR pointers>.
+```
+
+## Worked example
+
+A complete filled-in case, in the exact form it takes in its own
+`references/<family>/<slug>.md` file — H1 title, breadcrumb, then every field:
+
+```
+---
+id: case-cap-cuda-graph-batch-sizes   # = "case-" + filename stem
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch, memory]
+signals: [gpu-idle-between-steps, many-small-kernels, small-batch-decode, memory-capacity-bound]
+architectures: [any-sm]
+model_scope: [model-agnostic]
+phase: [decode]
+patterns: []          # register the pattern in data/patterns.yaml first, then point here
+accuracy_risk: lossless
+apply_via_kind: [config-knob]
+knobs: [cuda_graph_config.max_batch_size, cuda_graph_config.batch_sizes]
+specialists: [perf-torch-cuda-graphs]
+commits: []
+measured: []
+---
+
+# Cap CUDA-graph capture to reachable decode batch sizes
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Applies when:** launch-overhead-bound decode (many small kernels, high host
+  time, GPU idle between steps) **and** graph workspace memory is tight — e.g.
+  a low-concurrency deployment where the server `max_batch_size` is large but
+  only small decode batches are ever reached.
+- **Counter-signals:** concurrency that routinely reaches batch sizes above
+  the cap (those fall back to eager and regress step time); ample free memory
+  (capping buys nothing — leave the default, which follows `max_batch_size`).
+- **Mechanism:** CUDA graphs replay a captured launch sequence, removing
+  per-step launch overhead. Capturing graphs for batch sizes that never occur
+  wastes graph workspace memory; capping capture to reachable sizes keeps the
+  launch-overhead win without the memory cost.
+- **Generalizes to:** the pattern "size a captured/pre-allocated resource to the
+  *reachable* working set, not the configured maximum." Carries to other
+  resources keyed off `max_batch_size` — padded attention/decode buffers,
+  speculative draft-length capture, warmup shape sets. Adapt by re-deriving the
+  reachable set from the actual concurrency/ISL/OSL, then capping to it.
+- **Apply via:** `cuda_graph_config.max_batch_size` / `cuda_graph_config.batch_sizes`;
+  delegate capture/replay work to **perf-torch-cuda-graphs** /
+  **perf-torch-cuda-graph-specialist**. (Default `cuda_graph_config.max_batch_size`
+  to the server `max_batch_size`; lower it only when memory is the constraint.)
+- **Expected effect:** lower per-step latency / smaller inter-step GPU-idle gap
+  at the same accuracy; reclaimed graph workspace memory. Measured Δ to be
+  recorded from the trial run dir.
+- **Accuracy risk:** lossless — CUDA graph replay does not change numerics.
+- **Verify:** per-decode-step time and GPU-idle ratio from an nsys trace before
+  vs after; confirm batches above the cap fall back to eager (slower, not
+  erroring) and are not on the hot path.
+- **Rollback:** remove or raise `cuda_graph_config.max_batch_size`; trigger:
+  any reachable batch size falling out of capture and regressing step time
+  >5%, or graph capture OOM.
+- **Prior art:** `perf-torch-cuda-graphs` skill;
+  `trtllm-serve-config-guide` `references/knob-heuristics.md` (the DeepSeek-R1
+  conc=1 example uses `cuda_graph_config.max_batch_size: 1` with server
+  `max_batch_size: 512`).
+```
diff --git a/.claude/skills/perf-optimization-casebook/references/communication/deepep.md b/.claude/skills/perf-optimization-casebook/references/communication/deepep.md
new file mode 100644
index 000000000000..b1407afcdec2
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/communication/deepep.md
@@ -0,0 +1,37 @@
+---
+id: case-deepep
+type: case
+family: communication
+maturity: full
+bottleneck: [communication]
+signals: [alltoall-dominates, slow-path-fallback]
+architectures: [sm90, sm100]
+model_scope: [moe, deepseek-v3r1, qwen3-moe]
+phase: [any-phase]
+patterns: [pattern-swap-generic-collective-for-specialized-comm]
+accuracy_risk: lossless
+apply_via_kind: [env-var]
+knobs: [TRTLLM_CAN_USE_DEEP_EP]
+specialists: [perf-sweep-workflow, trtllm-moe-develop]
+commits: ['0b60da2c45ac']
+eligibility:
+  - "deep_ep_installed and dtype==bfloat16 (path auto-resolve gate, fused_moe_cutlass.py)"
+interactions:
+  - {feature: cuda-graphs, relation: incompatible-with, note: "the DeepEP method only; auto-resolves to DeepEPLowLatency under CUDA Graphs"}
+measured: []
+---
+
+# Use a dedicated expert-parallel all-to-all comm backend (DeepEP) for MoE dispatch/combine
+
+> Part of the [Communication casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `0b60da2c45ac` feat: large-scale EP (part 7: DeepEP integration) (#4792).
+- **Applies when:** communication-bound + signals: MoE model with expert parallelism (large-scale EP); token routing dispatch/combine all-to-all is a large share of step time; intranode NVLink and/or internode RDMA available; default `allgather`/`reducescatter` MoE path in use. Decode/low-latency with CUDA Graphs needs the LowLatency variant.
+- **Mechanism:** replaces generic allgather/reducescatter token exchange with DeepEP's purpose-built EP all-to-all kernels (patched NVSHMEM, `NVSHMEM_IBGDA_SUPPORT=1`) routing only each rank's selected tokens to their expert owners — cutting bytes and overlapping NVL/RDMA transfer with layout compute. Two methods via `AlltoallMethodType` (`fused_moe_cutlass.py`): `DeepEP` (intranode/internode, no CUDA-Graph, IBGDA for internode; currently faster) and `DeepEPLowLatency` (CUDA-Graph-compatible, IBGDA) — backed by `VariableLengthBuffer`/`VariableLengthLowLatencyBuffer` (`deep_ep_utils.py`).
+- **Generalizes to:** "swap a generic collective for a workload-specialized comm library"; carries to MoE EP combine, attention-DP all-to-all, PP send/recv; adapt by matching the backend's buffer-reservation + layout-adapter contract and picking the graph-safe variant for decode.
+- **Apply via:** env `TRTLLM_CAN_USE_DEEP_EP=1`; auto-resolves to `DeepEPLowLatency` under CUDA Graphs else `DeepEP` (only when `deep_ep_installed and dtype==bfloat16`); module `enable_alltoall` gates the path; requires the DeepEP/NVSHMEM build (`docker/common/install_deep_ep.sh`). Delegate config exploration to **perf-sweep-workflow**.
+- **Expected effect:** lower MoE dispatch/combine latency + higher EP throughput at scale; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless (bf16 token exchange; path asserts `dtype==bfloat16` — changes transport, not math).
+- **Verify:** end-to-end throughput/latency vs allgather/reducescatter baseline; correctness with `test_fused_moe.py`; sanity accuracy on a small eval.
+- **Rollback:** unset `TRTLLM_CAN_USE_DEEP_EP` (falls back to allgather/reducescatter). Trigger: DeepEP build/IBGDA unavailable, regression vs baseline, or CUDA-Graph capture failure on the non-LL method.
+- **Prior art:** PR #4792. Files: `_torch/modules/fused_moe/deep_ep_utils.py`, `fused_moe_cutlass.py` (`AlltoallMethodType`), `modeling_deepseekv3.py`, `docker/common/install_deep_ep.sh`. Owning skill: **trtllm-moe-develop**; sweep with **perf-sweep-workflow**.
diff --git a/.claude/skills/perf-optimization-casebook/references/communication/index.md b/.claude/skills/perf-optimization-casebook/references/communication/index.md
new file mode 100644
index 000000000000..1745884df307
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/communication/index.md
@@ -0,0 +1,74 @@
+# Casebook — Communication
+
+Optimizations for **multi-GPU collectives** — expert-parallel (EP) all-to-all for
+MoE dispatch/combine, and tensor-parallel (TP) AllReduce — plus the buffers and
+strategy-selection that feed them. These only matter at **TP>1 / EP>1** and only
+when a collective is a real share of step time; classify communication-bound
+first (`perf-analysis` / `perf-nsight-systems` — collective time as a fraction of
+the step, NCCL/all-to-all kernels on the critical path). Add entries using the
+schema in [case-template.md](../case-template.md); match on **Applies when** signals,
+adapt to the current world size / interconnect, and measure before claiming a win.
+
+> Routing note: a collective's *epilogue* fusion (AllReduce + Residual + RMSNorm
+> (+Quant)) is a `torch.compile` fusion pattern and lives in
+> the kernel-and-fusion casebook ([Fuse the TP collective epilogue](../kernel-and-fusion/fuse-ar-epilogue.md)); the [UserBuffers case](userbuffers-symmetric-memory.md) is what makes the *in-place* fused UB
+> variant of that possible. Many of these are bf16-only or quant-checkpoint-gated
+> — read the dtype gate in each case.
+
+## Recurring patterns in this family
+
+Match on these transferable patterns, not on a case title — most situations are a
+*variation* of a case, not its exact instance.
+
+- **Swap a generic collective for a workload-specialized comm library.** Generic
+  `allgather`/`reducescatter` moves padded, full-width data; a purpose-built EP
+  all-to-all (DeepEP) routes only each rank's selected tokens to their expert
+  owners and overlaps NVLink/RDMA transfer with layout compute. Carries to EP
+  combine, attention-DP all-to-all, PP send/recv. Pick the CUDA-graph-safe variant
+  for decode. _(Instance: [the DeepEP case](deepep.md).)_
+- **Send activations post-quant, not pre-quant.** When the model is already
+  quantized, quantize tokens *before* the collective and ship the packed
+  low-precision payload (~½ for FP8, ~¼ for FP4) instead of bf16 — pack scale
+  factors with the payload and re-swizzle on receive. **Lossy** (precision crosses
+  the wire) → accuracy record + rollback. Carries to EP dispatch+combine, TP
+  allgather/reducescatter of quantized activations, KV transfer. _(Instance: [the low-precision dispatch case](low-precision-dispatch.md).)_
+- **Make the collective *kernel* faster instead of changing the algorithm.**
+  Buffer-flag / Lamport sync, grid right-sizing, vectorized multicast loads tune a
+  fixed algorithm (e.g. MNNVL two-shot) rather than swapping it. Lossless (same
+  reduction math). Carries to one-shot AllReduce, allgather/reduce-scatter, custom
+  NVLS/multimem collectives — mind the alignment/divisibility assumptions.
+  _(Instance: [the MNNVL two-shot case](mnnvl-twoshot-allreduce.md).)_
+- **Pre-register symmetric comm buffers to kill per-step allocation and unlock
+  in-place fused collectives.** A warm-up-time symmetric-memory pool (UserBuffers /
+  NCCL-symmetric) removes dynamic allocation during inference and lets the producer
+  write directly into the comm buffer so AR + RMSNorm (+quant) fuse in place.
+  Carries to NCCL-symmetric AllReduce, AR+norm/quant fusion. _(Instance: [the UserBuffers case](userbuffers-symmetric-memory.md).)_
+- **Autotune / heuristically pick the collective per shape & concurrency.** No
+  single AllReduce algorithm wins across message sizes — choose by
+  `(SM, TP, fusion, hidden, tokens)` via an offline LUT or a runtime `TunableRunner`.
+  Lossless (bit-equivalent implementations). Carries to allgather/reduce-scatter,
+  GEMM tactic selection, attention-backend choice. _(Instance: [the shape-aware autotune case](shape-aware-allreduce-autotune.md).)_
+
+## Cases
+
+Match on **Applies when** / **Generalizes to**, then open that case file. These
+only matter at **TP>1 / EP>1** when a collective is a real share of step time.
+
+> Risk key: **lossless** = bit-/math-equivalent (only transport or kernel changes)
+> · **lossy** = precision crosses the wire; needs an accuracy record + rollback ·
+> **mixed** = the AllReduce+norm fusion is lossless but its quant variant is lossy.
+
+### EP all-to-all (MoE dispatch / combine)
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [DeepEP all-to-all backend](deepep.md) | comm-bound MoE EP; dispatch/combine all-to-all a large share of step time; NVLink/RDMA | swap a generic collective for a workload-specialized comm library | lossless |
+| [Low-precision dispatch/combine](low-precision-dispatch.md) | comm-bound MoE already on DeepEP + quantized model (FP8/NVFP4/W4A8); bf16 dispatch bytes dominate | send activations post-quant, not pre-quant | lossy |
+
+### TP AllReduce (kernel, buffers, strategy)
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [MNNVL two-shot AllReduce kernel](mnnvl-twoshot-allreduce.md) | comm-bound TP AllReduce on MNNVL/NVLink; mid/large messages; the AR kernel (not the network) is the bottleneck | make the collective kernel faster instead of changing the algorithm | lossless |
+| [UserBuffers / symmetric memory](userbuffers-symmetric-memory.md) | comm-bound TP AllReduce (PyTorch flow); AR→add→RMSNorm(+quant); per-step buffer alloc / unfused chain | pre-register symmetric comm buffers to enable in-place fused collectives | mixed |
+| [Shape-aware AllReduce autotune](shape-aware-allreduce-autotune.md) | comm-bound TP AllReduce; no single algo wins across shapes (one-shot regresses vs NCCL somewhere); strategy pinned | autotune/heuristically pick the collective per shape & concurrency | lossless |
diff --git a/.claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md b/.claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md
new file mode 100644
index 000000000000..8ea0ec641034
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md
@@ -0,0 +1,38 @@
+---
+id: case-low-precision-dispatch
+type: case
+family: communication
+maturity: full
+bottleneck: [communication]
+signals: [alltoall-dominates, full-precision-on-wire]
+architectures: [sm90, sm100]
+model_scope: [moe, deepseek-v3r1, qwen3-moe]
+phase: [any-phase]
+patterns: [pattern-send-activations-post-quant]
+accuracy_risk: lossy
+apply_via_kind: [env-var]
+knobs: [TRTLLM_MOE_POST_QUANT_ALLTOALLV, TRTLLM_MOE_USE_LOW_PRECISION_COMBINE, TRTLLM_FORCE_ALLTOALL_METHOD]
+specialists: [perf-sweep-challenger, trtllm-moe-develop]
+commits: ['854655f2f7b3', 'f172face98cf', '336c2ef5408d']
+eligibility:
+  - "quantized checkpoint required: has_fp8_qdq / has_nvfp4 / has_w4afp8"
+  - "hidden size divisible by 32 for nvfp4/fp8 low-latency dispatch (as of #7927)"
+interactions:
+  - {case: case-deepep, relation: depends-on, note: "rides on the DeepEP all-to-all path; requires DeepEP enabled"}
+measured: []
+---
+
+# Communicate MoE tokens in low precision (FP4/FP8 post-quant dispatch & combine)
+
+> Part of the [Communication casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `854655f2f7b3` deepEP fp4 post quant all2all dispatch (#5881); related: `f172face98cf` DeepEP LL dispatch FP4 (#6296), `336c2ef5408d` DeepEP LL fp8 dispatch/combine (#7927).
+- **Applies when:** communication-bound + signals: MoE already on a DeepEP all-to-all (see the [DeepEP case](deepep.md)) and the model is FP8-QDQ / NVFP4 / W4A8 quantized; dispatch/combine bf16 byte volume dominates; large EP world size where comm bytes scale with tokens×hidden.
+- **Mechanism:** quantize tokens *before* the all-to-all and ship the packed low-precision payload instead of bf16 — ~½ (FP8) or ~¼ (FP4) the dispatched bytes. FP8 dispatches `float8_e4m3fn` directly. FP4 (#5881) packs NVFP4 data + `uint8` scale factors into a bf16-typed buffer of `packed_hidden_size=2560` (LL kernel accepts bf16 hidden sizes 2560/4096/5120/7168), unpacks + re-`swizzle_sf` on receive; #6296 adds native `low_latency_dispatch_fp4`; #7927 generalizes LL fp8/nvfp4/w4afp8 dispatch and adds `low_latency_combine_low_precision(precision="fp8", ...)` so the combine leg is also low-precision.
+- **Generalizes to:** "send activations post-quant, not pre-quant" — any quantized collective: EP combine, TP allgather/reducescatter of quantized activations, KV transfer; carries to FP8 and FP4 dispatch+combine; adapt by packing scale factors with the payload, re-swizzling on the receiver, honoring per-format hidden-size/divisibility (`%32` for nvfp4/fp8 LL).
+- **Apply via:** env `TRTLLM_MOE_POST_QUANT_ALLTOALLV=1` (post-quant dispatch), `TRTLLM_MOE_USE_LOW_PRECISION_COMBINE=1` (low-precision combine), `TRTLLM_FORCE_ALLTOALL_METHOD` to pin method; requires DeepEP + a quantized checkpoint (`has_fp8_qdq`/`has_nvfp4`/`has_w4afp8`). Delegate accuracy gating to **perf-sweep-challenger**.
+- **Expected effect:** lower dispatch/combine latency + higher EP throughput from reduced comm bytes; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossy — tokens cross the wire in FP4/FP8 (and combine may be low-precision); partial sums/activations lose precision vs bf16 transport. Needs accuracy record + rollback criterion; FP4 combine is the highest-risk leg.
+- **Verify:** task-accuracy/perplexity parity vs the bf16-dispatch DeepEP baseline (and vs low-precision-combine off); confirm throughput gain; check FP4 pack/unpack round-trip (SF swizzle) correctness.
+- **Rollback:** `TRTLLM_MOE_POST_QUANT_ALLTOALLV=0` and `TRTLLM_MOE_USE_LOW_PRECISION_COMBINE=0` (revert to bf16 dispatch/combine). Trigger: accuracy drop beyond recorded threshold.
+- **Prior art:** PRs #5881, #6296, #7927. Files: `_torch/modules/fused_moe/fused_moe_wide_ep.py` (`alltoall_postquant_dispatch`, `low_latency_dispatch_fp4`, `low_latency_combine_low_precision`), `deep_ep_utils.py`, `thop/moeOp.cpp`. Owning skill: **trtllm-moe-develop**; gate with **perf-sweep-challenger**.
diff --git a/.claude/skills/perf-optimization-casebook/references/communication/mnnvl-twoshot-allreduce.md b/.claude/skills/perf-optimization-casebook/references/communication/mnnvl-twoshot-allreduce.md
new file mode 100644
index 000000000000..56da49ba9dd4
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/communication/mnnvl-twoshot-allreduce.md
@@ -0,0 +1,35 @@
+---
+id: case-mnnvl-twoshot-allreduce
+type: case
+family: communication
+maturity: full
+bottleneck: [communication, sync]
+signals: [allreduce-dominates, kernel-far-from-sol]
+architectures: [sm90, sm100]
+model_scope: [model-agnostic, dense, moe, deepseek-v3r1]
+phase: [any-phase]
+patterns: [pattern-tune-collective-kernel-not-algorithm]
+accuracy_risk: lossless
+apply_via_kind: [config-knob, kernel-change]
+knobs: [allreduce_strategy]
+specialists: [perf-nsight-compute-analysis, kernel-cuda-specialist]
+commits: ['6e1aee6fd68a']
+eligibility:
+  - "num_token*token_dim divisible by ELTS_PER_LOAD (4 for BF16, 2 for FP32), as of #5934"
+measured: []
+---
+
+# Optimize the collective kernel itself (two-shot / Lamport-sync MNNVL AllReduce)
+
+> Part of the [Communication casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `6e1aee6fd68a` [fix] Performance Optimization for MNNVL TwoShot Kernel (#5934).
+- **Applies when:** communication-bound + signals: TP AllReduce on a multi-node NVLink (MNNVL) domain; medium/large messages where two-shot (reduce-scatter + allgather) beats one-shot; `allreduce_strategy=MNNVL` or `TWOSHOT` already selected and the AR kernel (not the network) is the bottleneck.
+- **Mechanism:** tunes the MNNVL two-shot AllReduce CUDA kernel rather than swapping algorithms. Lamport-style flag sync with triple-buffered flags (`buffer_flags % 3`), grid reworked so only a subset of CTAs do the lamport sync; cross-CTA completion via `red.async.release.global.gpu.add.u32` atomics + spin on `offset_access_ptr < gridDim.x*gridDim.y`; data via vectorized `ld.volatile.global.v4.f32`/`v2.f32` over NVLink-multicast buffers (`mcastDeviceMemory`), assuming 8B-atomic writes and `num_token*token_dim` divisible by ELTS_PER_LOAD (4 BF16 / 2 FP32).
+- **Generalizes to:** "make the collective kernel faster instead of changing the algorithm" — buffer-flag/Lamport sync, grid-rightsizing, vectorized multicast loads carry to one-shot AllReduce, allgather/reduce-scatter, custom NVLS/multimem collectives; adapt by matching alignment/divisibility assumptions and the flag lifecycle.
+- **Apply via:** `allreduce_strategy=MNNVL` (or `TWOSHOT`) in `_torch/model_config.py` (`AllReduceStrategyType`: NCCL=0, MIN_LATENCY=1, UB=2, AUTO=3, ONESHOT=4, TWOSHOT=5, LOWPRECISION=6, MNNVL=7, NCCL_SYMMETRIC=8). Delegate kernel profiling to **perf-nsight-compute-analysis**; kernel edits to **kernel-cuda-specialist**.
+- **Expected effect:** lower MNNVL TwoShot AllReduce kernel latency at targeted message sizes; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless (same reduction math/precision; only sync, grid layout, vectorized access changed).
+- **Verify:** AllReduce kernel time (ncu/nsys) + end-to-end TP latency vs pre-patch; numeric parity of reduced output across ranks; respect alignment/divisibility preconditions.
+- **Rollback:** switch `allreduce_strategy` to `AUTO`/`NCCL` or revert the kernel commit. Trigger: correctness/sync failure or regression on the target shapes.
+- **Prior art:** PR #5934. Files: `cpp/.../kernels/communicationKernels/mnnvlTwoShotAllreduceKernels.cu`, `runtime/mcastDeviceMemory.{h,cpp}`, `modeling_deepseekv3.py`, `kernels/customAllReduceKernels.h`. Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/communication/shape-aware-allreduce-autotune.md b/.claude/skills/perf-optimization-casebook/references/communication/shape-aware-allreduce-autotune.md
new file mode 100644
index 000000000000..e9555da6c5e2
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/communication/shape-aware-allreduce-autotune.md
@@ -0,0 +1,33 @@
+---
+id: case-shape-aware-allreduce-autotune
+type: case
+family: communication
+maturity: full
+bottleneck: [communication]
+signals: [allreduce-dominates, high-concurrency]
+architectures: [sm80, sm90, sm100]
+model_scope: [model-agnostic, dense, moe, deepseek-v3r1]
+phase: [any-phase]
+patterns: [pattern-autotune-collective-per-shape]
+accuracy_risk: lossless
+apply_via_kind: [default-on, config-knob, env-var]
+knobs: [allreduce_strategy, TLLM_DISABLE_ALLREDUCE_AUTOTUNE, OVERRIDE_HEURISTIC_ALLREDUCE_STRATEGY]
+specialists: [perf-sweep-workflow]
+commits: ['22257457824a', 'd272f1a9bcbc']
+measured: []
+---
+
+# Shape-aware strategy selection / autotuning of collectives (pick AllReduce algo by message size & concurrency)
+
+> Part of the [Communication casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `22257457824a` [TRTLLM-8129][feat] Allreduce tuning and benchmark script revising (#7870); related: `d272f1a9bcbc` [TRTLLM-8821][feat] Apply AutoTuner to AllReduce Op for strategy tuning (#8531).
+- **Applies when:** communication-bound + signals: TP AllReduce where no single algorithm wins across shapes — a one-shot kernel regresses vs NCCL at some sizes (A100/H100) while one/two-shot win elsewhere; throughput varies strongly with message size (seq_len×hidden) and concurrency; `allreduce_strategy` currently pinned.
+- **Mechanism:** make the algorithm a function of shape. #7870 adds two `AUTO` heuristics in `customAllReduceUtils.h`: an LP threshold (`SelectStrategyLP`: per-`(SM_major, TP_size)` `(NCCL_num_token_threshold, TWO_SHOT_numel_threshold)` → TWOSHOT if `message_size ≥ two_shot_numel_threshold`, else ONESHOT, else NCCL) and a benchmark-derived LUT (`selectStrategyLookUpTable`, indexed `[sm][tp][fusionOp][log2(hidden)-7][log2(tokens)]`) generated offline by `allreduce_heuristic_code_gen.py`. Also fixes that ONESHOT/TWOSHOT couldn't be overridden, fixes a TWOSHOT perf bug, cleans dispatch. #8531 wires it into the runtime AutoTuner via `AllReduceRunner(TunableRunner)` + `trtllm::tunable_allreduce`, with `AutoTuner.choose_one` picking per-profile.
+- **Generalizes to:** "autotune/heuristically pick the collective implementation per shape & concurrency"; carries to allgather/reduce-scatter, GEMM tactic selection, attention backend choice; adapt by benchmarking candidates offline into a LUT or registering a `TunableRunner` for the AutoTuner.
+- **Apply via:** `allreduce_strategy=AUTO` (default) engages heuristic/LUT; AutoTuner runtime tuning runs under `AUTO` unless `TLLM_DISABLE_ALLREDUCE_AUTOTUNE=1`; `OVERRIDE_HEURISTIC_ALLREDUCE_STRATEGY` forces a strategy for benchmarking. Regenerate the LUT with `allreduce_heuristic_code_gen.py`. Delegate sweeping to **perf-sweep-workflow**.
+- **Expected effect:** recovers per-shape best AllReduce time, removes the one-shot-vs-NCCL regression. **PR #7870 reports ~3–4% end-to-end gain for DeepSeek-R1 at concurrency 256 and 512** (fixes a known one-shot regression). Re-measure on your shapes.
+- **Accuracy risk:** lossless (selects among bit-equivalent reduction implementations; precision unchanged).
+- **Verify:** AllReduce time across the shape/concurrency grid (`all_reduce.py` microbench + `allreduce_perf_viz.py`); end-to-end throughput at concurrency 256/512; confirm no regression vs NCCL at small sizes.
+- **Rollback:** pin `allreduce_strategy` to a fixed value (e.g. `NCCL`), or `TLLM_DISABLE_ALLREDUCE_AUTOTUNE=1` to keep AUTO heuristics without runtime autotuning. Trigger: tuner-induced regression or instability.
+- **Prior art:** PRs #7870, #8531. Files: `cpp/.../common/customAllReduceUtils.h` (`SelectStrategyLP`, `selectStrategyLookUpTable`), `thop/allreduceOp.cpp`, `_torch/custom_ops/torch_custom_ops.py` (`AllReduceRunner`), `tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py`, `tests/microbenchmarks/all_reduce.py`. Sweep with **perf-sweep-workflow**.
diff --git a/.claude/skills/perf-optimization-casebook/references/communication/userbuffers-symmetric-memory.md b/.claude/skills/perf-optimization-casebook/references/communication/userbuffers-symmetric-memory.md
new file mode 100644
index 000000000000..a1953227119f
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/communication/userbuffers-symmetric-memory.md
@@ -0,0 +1,37 @@
+---
+id: case-userbuffers-symmetric-memory
+type: case
+family: communication
+maturity: full
+bottleneck: [communication, launch, memory]
+signals: [allreduce-dominates, hbm-roundtrip-between-ops]
+architectures: [sm90, sm100]
+model_scope: [model-agnostic, dense, moe]
+phase: [any-phase]
+patterns: [pattern-preregister-symmetric-comm-buffers]
+accuracy_risk: mixed
+apply_via_kind: [config-knob]
+knobs: [allreduce_strategy]
+specialists: [perf-torch-cuda-graphs, perf-host-optimization]
+commits: ['dca6397d1e2d']
+eligibility:
+  - "AR input dtype FP16/BF16 (UB extra_check gate)"
+interactions:
+  - {case: case-fuse-ar-epilogue, relation: composes-with, note: "UB registration is the in-place backing that enables the fused AR+RMSNorm(+quant) variant"}
+measured: []
+---
+
+# UserBuffers / symmetric memory — pre-register comm buffers to enable AllReduce+quant fusion
+
+> Part of the [Communication casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `dca6397d1e2d` feat: Introduce UB allocator for pytorch flow (#3257).
+- **Applies when:** communication-bound + signals: TP AllReduce in the PyTorch flow where AR is followed by residual-add + RMSNorm (+ optional FP8/FP4 quant); per-step comm-buffer allocation or an unfused AR→norm→quant chain shows up; want a symmetric-memory AllReduce path. UB requires FP16/BF16 AR input.
+- **Mechanism:** a global `UserBuffersManager` singleton hands out pre-registered symmetric UserBuffers tensors (`create_userbuffers_tensor` op, `allocate_userbuffers`, `search_buffer`). Buffers reserved at warm-up via `initialize_userbuffers_manager(tp_size, ...)` so "there is no dynamic allocation during inference." This registration unlocks the fused UB AllReduce path: `ub_allreduce.py` rewrites the graph to fuse quant into the AllReduce as an in-place UB impl (`copy_to_userbuffers`), convert supported AllReduces to UB, fuse the producer to write directly into the userbuffer, and drop `userbuffers_allreduce_finalize` when chained — running `allreduce(..., strategy=UB, fusion=RESIDUAL_RMS_NORM_QUANT_FP8)` via launchers like `allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher`.
+- **Generalizes to:** "pre-register symmetric comm buffers to enable in-place fused collectives + kill per-step allocation"; carries to NCCL symmetric AllReduce, AR+RMSNorm fusion, AR+quant (FP8/FP4) fusion; adapt by routing producers to write into the registered buffer and gating on dtype support.
+- **Apply via:** `allreduce_strategy=UB` (`AllReduceStrategyType::UB=2`); UB manager auto-inits at warm-up; AR+RMSNorm(+FP8/FP4 quant) fusion applied by the torch-compile UB pattern pass (`register_ub_patterns`). Delegate to **perf-torch-cuda-graphs** / **perf-host-optimization** for capture + host-overhead validation.
+- **Expected effect:** lower AllReduce+norm(+quant) latency via in-place fusion + no per-step buffer allocation; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** mixed — UB AllReduce+RMSNorm fusion is lossless; the FP8/FP4-quant fusion variants are lossy (emit FP8/FP4) and need an accuracy record + rollback. UB AR requires FP16/BF16 input (`extra_check` gate).
+- **Verify:** AR+norm latency + step time vs unfused/NCCL; for quant-fusion variants, accuracy/parity vs unfused FP8/FP4; `multi_gpu/test_user_buffers.py`.
+- **Rollback:** `allreduce_strategy` → `AUTO`/`NCCL` (disables UB fusion). Trigger: capture/registration failure or accuracy regression on the quant-fused variant.
+- **Prior art:** PR #3257. Files: `cpp/.../kernels/userbuffers/userbuffersManager.{h,cpp}`, `thop/userbuffersTensor.cpp`, `thop/allreduceOp.cpp`, `_torch/compilation/patterns/ub_allreduce.py`, `pyexecutor/model_engine.py`. Related: the [AR-epilogue fusion case](../kernel-and-fusion/fuse-ar-epilogue.md) (UB is the in-place backing for it).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md
new file mode 100644
index 000000000000..701a892bec02
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md
@@ -0,0 +1,33 @@
+---
+id: case-fold-scale-swizzle-into-kernel
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [host-overhead, launch, memory]
+signals: [many-small-kernels, hbm-roundtrip-between-ops]
+architectures: [sm90, sm100]
+model_scope: [moe, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-fuse-chain-feeding-gemm]
+accuracy_risk: lossless
+apply_via_kind: [kernel-change, code-change]
+knobs: []
+specialists: [kernel-cuda-specialist, trtllm-moe-develop]
+commits: ['2f2f5cc72c51']
+measured: []
+---
+
+# Fold a host-side pre-pass into the consuming kernel (scale-factor swizzle)
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `2f2f5cc72c51` [TRTLLM-6744][feat] Remove input_sf swizzle for module WideEPMoE (#6231).
+- **Applies when:** NVFP4 WideEPMoE with expert-parallel comm where a separate `swizzle_sf(...)` op was launched on the host path to re-layout the activation scale-factor tensor into SWIZZLED layout before the grouped GEMM — adding a kernel launch + full SF read/write per MoE layer; signals: a standalone swizzle/transpose op before MoE GEMM in profiles, host launch overhead in MoE dispatch.
+- **Mechanism:** the MoE GEMM's activation-SF setup is taught to consume a non-swizzled (linear) SF directly: a `swizzled_input_sf` flag is threaded through `torch.ops.trtllm.fused_moe` → `moeOp.cpp` → CUTLASS `moe_kernels.cu`, which branches on the flag and reads SF in linear layout when `false`. The three `swizzle_sf(...)` calls in `WideEPMoE.forward` are deleted.
+- **Generalizes to:** "fold a standalone layout/transform pre-pass into the kernel that consumes the data"; carries to other quantized GEMMs that pre-swizzle scales, weight-SF layout passes, any element-wise relayout (transpose/pack) feeding a GEMM; adapt by adding a layout flag to the consumer kernel and handling the un-transformed layout in its load path, then deleting the producer op.
+- **Apply via:** plumb a layout flag (`swizzled_input_sf`) through the op chain; delegate to **kernel-cuda-specialist** (CUTLASS) and **trtllm-moe-develop** for module wiring.
+- **Expected effect:** one fewer kernel launch + one fewer full SF read/write per MoE layer → lower host launch overhead and memory traffic in FP4 WideEPMoE; direction only — measured Δ (MoE-layer time, launch count, throughput) to be recorded from run.
+- **Accuracy risk:** lossless when the flag is correct (kernel reads SF in its actual layout). Risk = layout/flag mismatch silently corrupting FP4 GEMM output — parity-check on first enable.
+- **Verify:** profile shows the swizzle op gone from MoE path; MoE-layer latency/launch count down; NVFP4 MoE output parity vs pre-swizzle path.
+- **Rollback:** set `swizzled_input_sf=True` and restore the `swizzle_sf(...)` calls. Trigger: FP4 MoE accuracy mismatch (layout/flag inconsistency).
+- **Prior art:** PR #6231. Files: `_torch/custom_ops/torch_custom_ops.py`, `_torch/modules/fused_moe/fused_moe_wide_ep.py`, `cpp/.../thop/moeOp.cpp`, `cutlass_kernels/moe_gemm/moe_kernels.cu`. Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fp8-mla-kv-cache.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fp8-mla-kv-cache.md
new file mode 100644
index 000000000000..272add499e5a
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fp8-mla-kv-cache.md
@@ -0,0 +1,38 @@
+---
+id: case-fp8-mla-kv-cache
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [memory, compute]
+signals: [kv-reads-dominate, long-context, high-concurrency, memory-capacity-bound]
+architectures: [sm90, sm100]
+model_scope: [mla, deepseek-v3r1, model-agnostic]
+phase: [decode]
+patterns: [pattern-quantize-memory-bound-side]
+accuracy_risk: lossy
+apply_via_kind: [config-knob]
+knobs: [kv_cache_config.dtype]
+specialists: [kernel-cuda-specialist, perf-nsight-compute-analysis]
+commits: ['515dd0d78fe8', '897c4bffd7']
+eligibility:
+  - "sm == 90 or sm == 100 (mFP8GenerationMLA gate: mSM==90||100)"
+  - "MLA model with FP8 KV-cache mode (QuantMode.has_fp8_kv_cache()); generation path only — context (prefill) MLA stays BF16"
+interactions:
+  - {feature: fp8-context-fmha, relation: incompatible-with, note: mFP8GenerationMLA and mFP8ContextFMHA are mutually exclusive}
+measured: []
+---
+
+# Quantize MLA generation to FP8 KV cache while keeping BF16 output
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `515dd0d78fe8` feat: Add support for FP8 MLA on Hopper and Blackwell (#3190); related: `897c4bffd7` Integrate FP4 indexer for DSA on Blackwell (#13340, the same idea applied to a *selector* cache one precision step lower).
+- **Applies when:** memory-bound MLA decode — KV-cache reads dominate generation; DeepSeek-style MLA (kv_lora_rank + qk_rope_head_dim latent KV, head_size 576); SM 90/100; long context / high concurrency where KV footprint and HBM bandwidth limit decode. Context (prefill) MLA stays BF16; the win is the generation path.
+- **Mechanism:** stores latent KV as `__nv_fp8_e4m3` (half the KV bytes + per-step bandwidth), runs the generation FMHA/FlashMLA with `dataTypeKv=DATA_TYPE_E4M3`. Q is quantized into `quant_q_buffer` with two BMM scales (`mla_bmm1_scale`/`mla_bmm2_scale`); **output is written BF16** (`dataTypeOut=DATA_TYPE_BF16`) so o_proj input stays high-precision and loss is confined to the KV/Q read path. `mFP8GenerationMLA` and `mFP8ContextFMHA` mutually exclusive; gated to `mSM==90||100`.
+- **Generalizes to:** "quantize the memory-bound side of attention (KV cache + matching operand) to FP8 but keep accumulation/output in BF16"; carries to GQA/MHA FP8 KV, FP8 context FMHA, FP4/INT8 KV; adapt by choosing the kv dtype the FMHA kernel supports, supplying BMM/dequant scales, and keeping output dtype unquantized where the next op (o_proj) is precision-sensitive. Also carries to a **model-specific *selector* cache**: DSA's indexer K cache goes one step lower to **FP4 (packed E2M1 + block-32 UE8M0)**, halving its per-token footprint (~132→68 B) and pairing a `fused_cat_fp4` cache-write with FP4×FP4 logits kernels (`indexer_k_dtype="fp4"`, SM≥100; #13340) — a selector cache tolerates coarser precision than the value KV because its output only ranks (cf. [ranking-only precision](ranking-only-precision-tf32.md)).
+- **Apply via:** `KvCacheConfig.dtype="fp8"` (`QuantMode.has_fp8_kv_cache()` → `mFP8GenerationMLA`); no separate gen-MLA knob (derived from FP8 KV mode on an MLA model). Delegate kernel work to **kernel-cuda-specialist**; profile with **perf-nsight-compute-analysis**.
+- **Expected effect:** lower KV HBM bytes, higher MLA decode throughput / lower latency under bandwidth pressure; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossy (FP8 e4m3 KV + FP8 Q); output BF16 but KV/Q rounding perturbs logits — needs accuracy record (GSM8K/MMLU vs BF16-KV) + rollback criterion before promotion.
+- **Verify:** decode throughput + KV bytes/step; accuracy parity vs `kv_cache_config.dtype="auto"`; confirm output stays BF16 (code asserts FP8-KV MLA requires bf16 output).
+- **Rollback:** `kv_cache_config.dtype="auto"`. Trigger: accuracy regression beyond recorded threshold, or FP8 gen-MLA kernel unavailable for the SM/head config.
+- **Prior art:** PR #3190. Files: `cpp/.../common/attentionOp.cpp` (`mFP8GenerationMLA`), `kernels/mlaKernels.cu`, `kernels/flashMLA/flash_fwd_mla_fp8_sm90.cu`, `_torch/modules/attention.py`. Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-add-norm-quant.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-add-norm-quant.md
new file mode 100644
index 000000000000..2e87d5dfc365
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-add-norm-quant.md
@@ -0,0 +1,39 @@
+---
+id: case-fuse-add-norm-quant
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [memory, launch]
+signals: [many-small-kernels, hbm-roundtrip-between-ops]
+architectures: [any-sm]
+model_scope: [llama-family, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-fuse-chain-feeding-gemm]
+accuracy_risk: lossy
+apply_via_kind: [code-change, config-knob, env-var]
+knobs: [TRTLLM_DISABLE_NVFP4_LAYERNORM_FUSION]
+specialists: [kernel-cuda-specialist, kernel-triton-specialist]
+commits: ['211c44b95199', '72cd7d824bb0']
+eligibility:
+  - "input dtype fp16/bf16; NO collective in the chain (TP epilogue with allreduce belongs to the AR-epilogue case)"
+  - "NVFP4 fused path: sm == 100 or sm == 103"
+  - "NVFP4 layernorm fusion is default-OFF: TRTLLM_DISABLE_NVFP4_LAYERNORM_FUSION defaults to '1'; set '0' to enable"
+interactions:
+  - {case: case-fuse-ar-epilogue, relation: alternative-to, note: "the collective sibling — when a TP allreduce precedes the add→norm→quant chain, fuse into the collective instead"}
+measured: []
+---
+
+# Fuse local residual-add + RMSNorm + quant (no collective) in one kernel / torch.compile pattern
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `211c44b95199` Adding torch ext API for FusedAddRMSNormQuant kernel (#9905); related: `72cd7d824bb0` Fuse add + norm + fp8 quant pattern (#12674).
+- **Applies when:** memory-bound layer epilogue with NO collective: `residual = x + residual; normed = rms_norm(residual); q = quantize(normed)` per layer. Signals: separate add, RMSNorm and `static_quantize_e4m3_per_tensor`/`fp4_quantize` kernels; fp16/bf16 → FP8/NVFP4 before the next GEMM; norm output only consumed by the quantizer.
+- **Mechanism:** (1) direct op `torch.ops.trtllm.fused_add_rms_norm_quant(...)` (`thop/fusedAddRMSNormQuant.cpp`, fp16/bf16) doing add+RMSNorm+NVFP4 quant in one kernel, returning packed FP4 + swizzled scales + updated residual — invoked from `RMSNorm(quantize_type="nvfp4")` and wired into `modeling_llama.py`. (2) `torch.compile` rewrite (`patterns/residual_add_norm.py::register_add_norm_quant`) matching `add → flashinfer_rmsnorm → static_quantize_e4m3_per_tensor` → `flashinfer_fused_add_rmsnorm_quant`. Both remove norm-output + residual round-trips and two launches/layer.
+- **Generalizes to:** "fuse a pointwise+reduction+quant chain that feeds the next GEMM into one op"; carries to add+LayerNorm+quant, gated-norm+quant, FP8 vs NVFP4 output, and the collective variant (sibling of the [AR-epilogue case](fuse-ar-epilogue.md)); adapt by choosing the op for the quant dtype, honoring fp16/bf16-in, passing the next GEMM's `input_scale`, and (in-place) verifying residual's last user is the add.
+- **Apply via:** `RMSNorm(..., quantize_type="nvfp4")` for the direct path; enable `torch.compile` for the FP8 rewrite. **NVFP4 layernorm fusion gated by env `TRTLLM_DISABLE_NVFP4_LAYERNORM_FUSION` (default "1" = OFF; set "0" to enable).** Delegate to **kernel-cuda-specialist** / **kernel-triton-specialist**.
+- **Expected effect:** fewer kernels + less HBM traffic per epilogue → lower latency in memory-bound regimes; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossy — quant fused on normed output may differ from unfused quant order/rounding; that is why NVFP4 fusion ships default-OFF. Always parity-check.
+- **Verify:** kernel count / nsys; numeric parity vs unfused add→norm→quant; accuracy eval (`test_nvfp4_with_norm_quant`, NVFP4 Llama-3.1-8B, SM100/103 only).
+- **Rollback:** keep `TRTLLM_DISABLE_NVFP4_LAYERNORM_FUSION=1` and/or disable `torch.compile`. Trigger: accuracy regression, or non-SM100/103 hardware.
+- **Prior art:** PRs #9905, #12674. Files: `cpp/.../thop/fusedAddRMSNormQuant.cpp`, `_torch/modules/rms_norm.py`, `modeling_llama.py`, `compilation/patterns/residual_add_norm.py`. Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-ar-epilogue.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-ar-epilogue.md
new file mode 100644
index 000000000000..1111a5d9eb49
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-ar-epilogue.md
@@ -0,0 +1,38 @@
+---
+id: case-fuse-ar-epilogue
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [communication, launch, memory]
+signals: [allreduce-dominates, many-small-kernels, hbm-roundtrip-between-ops]
+architectures: [any-sm]
+model_scope: [dense, moe, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-fuse-chain-feeding-gemm]
+accuracy_risk: mixed
+apply_via_kind: [config-knob]
+knobs: []
+specialists: [kernel-cuda-specialist, perf-optimization]
+commits: ['14d94a385641', '7b08677c0f15', '6d1f2d0fd700']
+eligibility:
+  - "TP > 1 with a per-layer allreduce → add → rms_norm → (quant) chain"
+  - "dtype fp16/bf16 and AllReduce strategy != UB (UB patterns register only when ub_enabled)"
+interactions:
+  - {case: case-userbuffers-symmetric-memory, relation: alternative-to, note: UB-backed in-place variant of the same epilogue fusion}
+measured: []
+---
+
+# Fuse the TP collective epilogue — AllReduce + Residual-add + RMSNorm (+ Quant)
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `14d94a385641` feat: Add non UB AR + Residual + Norm + Quant fusion (#6320); related: `7b08677c0f15` Add fused allreduce+RMSNorm op + optional residual (#12201), `6d1f2d0fd700` Finalize + Allreduce + add + rmsnorm fusion (#4756, MoE-finalize variant).
+- **Applies when:** communication-bound TP layer epilogue where `allreduce` is immediately followed by residual-add → RMSNorm → (optional) quant. Signals: TP>1; per-layer pattern `allreduce → add → rms_norm → static_quantize_e4m3_per_tensor`/`fp4_quantize`; many small kernels + collective launches per layer; UserBuffers unavailable/undesired.
+- **Mechanism:** collapses the post-AllReduce residual+norm(+quant) chain into one fused `torch.ops.trtllm.allreduce` carrying an `AllReduceFusionOp`, so residual-add, RMSNorm and quant ride inside the collective kernel's epilogue instead of launching separately — removing intermediate HBM round-trips and per-op launch overhead, emitting the quantized tensor (+scale) directly. Non-UB variant = a `torch.compile` PatternMatcher rewrite (`register_ar_fusions` in `patterns/ar_residual_norm.py`); UB patterns add only when `ub_enabled`.
+- **Generalizes to:** "fuse the collective's epilogue into the collective kernel"; carries to TP RMSNorm/LayerNorm epilogues, ReduceScatter+norm, MoE finalize+AllReduce+add+norm, any all-reduce-then-pointwise chain; adapt by choosing the right `AllReduceFusionOp` for the output set, gating on dtype (fp16/bf16) and strategy != UB, and (in-place residual) verifying the residual's last user is the fused node.
+- **Apply via:** enable `torch.compile` so the rewrite fires (`register_ar_fusions(custom_passes, ub_enabled)`); ops from `AllReduceFusionOp`: `RESIDUAL_RMS_NORM`, `RESIDUAL_RMS_NORM_QUANT_FP8`, `RESIDUAL_RMS_NORM_QUANT_NVFP4`, `*_OUT_QUANT_*`, plus `MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM`. Delegate to **kernel-cuda-specialist** (C++ epilogue) / **perf-optimization**. (UB-backed in-place variant: see the [UserBuffers case](../communication/userbuffers-symmetric-memory.md).)
+- **Expected effect:** fewer launches + less memory traffic per TP layer → lower per-token latency, most visible at TP>1 + small batch; no number stated — measured Δ to be recorded from run.
+- **Accuracy risk:** AR+residual+norm is lossless (same math, fused order); the `*_QUANT_FP8`/`*_QUANT_NVFP4` epilogues are lossy (quant in-kernel on normed output) — parity-check.
+- **Verify:** per-layer kernel count / nsys (fused AllReduce vs separate add/norm/quant); for quant variants run an accuracy eval vs unfused. Tests: `multi_gpu/test_allreduce.py`, `test_user_buffers.py`, `thop/test_moe.py`.
+- **Rollback:** disable `torch.compile` (pattern won't register) or restrict to `RESIDUAL_RMS_NORM` (drop quant fusion). Trigger: accuracy regression on quant variants, or UB-availability mismatch.
+- **Prior art:** PRs #6320, #12201, #4756. Files: `_torch/compilation/patterns/ar_residual_norm.py`, `compilation/backend.py` (`ub_enabled`), `cpp/.../thop/allreduceOp.cpp`, `kernels/.../moeAllReduceFusionKernels.cu`. Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-datamovement-into-quantize.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-datamovement-into-quantize.md
new file mode 100644
index 000000000000..72ad3aa69d4c
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-datamovement-into-quantize.md
@@ -0,0 +1,85 @@
+---
+id: case-fuse-datamovement-into-quantize
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [memory, launch]
+signals: [many-small-kernels, hbm-roundtrip-between-ops]
+architectures: [any-sm]
+model_scope: [sparse-attention, deepseek-v32, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-fuse-datamovement-into-quantize]
+accuracy_risk: lossless
+apply_via_kind: [kernel-change]
+knobs: []
+specialists: [kernel-cuda-specialist, kernel-triton-writing]
+commits: ['9a070ed709', '6601758d3a', '5f737b8dbe', '897c4bffd7']
+measured: []
+---
+
+# Fuse a concat/gather data-movement op into the quantize that follows it
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `9a070ed709` [perf] Add fused cat+fp8_quantize CUDA kernel for DSA
+  indexer (#11899); related: `6601758d3a` Kernel fusions in
+  `_gather_k_cache_for_chunk` of the DSA indexer (#12322), `5f737b8dbe` Use fp8
+  quant kernel in the DS3.2 indexer (#8701), `897c4bffd7` FP4 indexer
+  (`fused_cat_fp4`) (#13340).
+- **Applies when:** a **data-movement / layout op immediately precedes a quantize
+  (or dequantize)** on the hot path — `cat([pe, nope]) → rotate → view →
+  fp8_quantize`, or an index/`unravel` + advanced-indexing gather from a paged
+  cache followed by a reinterpret. Signals: a standalone `torch.cat` / `torch.split`
+  / `gather` / transpose kernel plus a separate quant kernel in the profile, with a
+  contiguous-copy intermediate between them; many small PyTorch ops
+  (`arange`/`unsqueeze`/broadcast/`_unravel_indices`/`.view`) building indices per
+  step; memory-bound activation prep for a quantized model.
+- **Mechanism:** collapse the layout op + the quantize into **one kernel** that
+  reads the (possibly non-contiguous) input **views directly via explicit row
+  strides** and emits the packed low-precision output + scale factors in a single
+  pass. #11899 reads the two BF16 halves straight from `torch.split` views
+  (`pe_row_stride`/`nope_row_stride`) — no contiguous concat copy — and writes FP8
+  E4M3 + one 1×128 scale/row (`fused_cat_fp8`). #12322 replaces ~8–12 PyTorch ops
+  with one Triton `triton_gather_k_cache` that gathers FP8 data + scales from flat
+  byte offsets. This removes the layout op's launch, the intermediate HBM
+  round-trip, and (for the gather) the index-tensor materialization. Independent
+  quantizes can also be **overlapped on a second stream** (#8701 runs Q and K quant
+  concurrently).
+- **Generalizes to:** the pattern "**fuse the data-movement/layout prologue
+  (concat / split / gather / transpose / RoPE-rotate) into the quantize (or
+  dequant) that consumes it, reading strided input views directly to skip the
+  contiguous-copy round-trip.**" A sibling of [folding a layout pre-pass into a
+  GEMM](fold-scale-swizzle-into-kernel.md) and [add+norm+quant fusion](fuse-add-norm-quant.md),
+  but here the consumer is the *quantizer*. Carries to: activation `cat`/`split` +
+  FP8/FP4 quant in any quantized model, gather-from-paged-cache + dequant,
+  transpose+quantize, KV write/read prep. Adapt by: writing the kernel to accept
+  strided input pointers, doing the layout in registers, and emitting quantized
+  data + scales in one pass; keep the unfused path as the reference.
+- **Apply via:** **not a server knob** — a fused custom op (CUDA or Triton).
+  Delegate the CUDA kernel + thop binding to **kernel-cuda-specialist**, a Triton
+  gather to **kernel-triton-writing**; add a `register_fake` so it traces under
+  `torch.compile`.
+- **Expected effect:** one kernel instead of two-plus + no contiguous-copy
+  intermediate → lower launch/HBM overhead in activation prep. Direction only —
+  measured Δ (op count, prep time, throughput) to be recorded from the run.
+- **Accuracy risk:** **lossless** — the fused kernel produces the *same quantized
+  result* as the unfused `cat/gather → quantize` path; it is a kernel-fusion
+  rewrite, so **verify vs the unfused reference** (a kernel bug changes outputs).
+  (The FP8/FP4 quantization itself is lossy, but that is pre-existing — this fusion
+  does not add precision loss.)
+- **Verify:** correctness — fused-op output equals the unfused `cat/gather +
+  quantize` reference within tolerance, incl. non-contiguous / edge strides
+  (`test_fused_cat_fp8.py`, `test_triton_gather_k_cache.py`). Perf — kernel/op count
+  and prep-phase time before vs after; confirm `register_fake` lets `torch.compile`
+  trace the op.
+- **Rollback:** swap the call site back to the unfused `cat/split/gather + quantize`
+  path (kept as reference). Trigger: any parity mismatch outside tolerance or a perf
+  regression.
+- **Prior art:** PRs #11899, #12322, #8701, #13340. Files:
+  `cpp/tensorrt_llm/kernels/fusedCatFp8.{cu,h}`, `fusedCatFp4.{cu,h}`,
+  `cpp/tensorrt_llm/thop/fusedCatFp8Op.cpp`,
+  `tensorrt_llm/_torch/attention_backend/sparse/{dsa.py,kernel.py}`
+  (`triton_gather_k_cache`, `_prep_q_or_k`). Owning specialists:
+  **kernel-cuda-specialist**, **kernel-triton-writing**. Related:
+  [fold-scale-swizzle](fold-scale-swizzle-into-kernel.md) and
+  [add+norm+quant](fuse-add-norm-quant.md) (adjacent fusion patterns).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md
new file mode 100644
index 000000000000..2930d01cb5a2
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md
@@ -0,0 +1,115 @@
+---
+id: case-fuse-moe-routing-kernel
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [launch, memory]
+signals: [routing-hot-path, many-small-kernels, hbm-roundtrip-between-ops, slow-path-fallback]
+architectures: [any-sm]
+model_scope: [moe, deepseek-v3r1, kimi-k2, model-agnostic]
+phase: [decode]
+patterns: [pattern-fuse-small-op-chain, pattern-reuse-warp-primitive, pattern-bounded-fastpath-fallback]
+accuracy_risk: lossless
+apply_via_kind: [kernel-change]
+knobs: []
+specialists: [kernel-cuda-specialist]
+commits: ['c8b9998acb8b', '84926bcb6f14']
+log_markers:
+  - "use the original pytorch implementation"   # is_fused=False fallback warning
+eligibility:
+  - "model uses the noaux_tc router (sigmoid gating + per-expert correction bias, grouped top-k)"
+  - "fused-path bounds live in tensorrt_llm/_torch/modules/fused_moe/routing.py::Deepseekv3RoutingImpl.noaux_tc (the is_fused guard) — read them from YOUR checkout before judging eligibility"
+  - "as of 84926bcb6f14, n_group > 1: requires top_k <= 8 and num_experts <= 256 and experts_per_group <= 32 and experts_per_group * topk_group <= 256"
+  - "as of 84926bcb6f14, n_group == 1: requires num_experts <= 1024 and top_k <= 32"
+interactions:
+  - {case: case-pdl, relation: composes-with, note: fused kernel can overlap neighbors via PDL (TRTLLM_ENABLE_PDL)}
+measured: []
+---
+
+# Fuse MoE grouped-top-k routing into one per-token CUDA kernel
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `c8b9998acb8b` Optimize the routing kernel for DeepseekV3 (#7761);
+  related: `84926bcb6f14` Update the deepseek routing (#13186) — widens fused-path
+  config eligibility so more workloads stay fused.
+- **Applies when:** the MoE **router/gating** step is on the decode hot path and
+  shows up as many tiny back-to-back kernels with GPU-idle gaps (launch-bound,
+  small tensors), **or** the config silently fell back to the slow PyTorch router
+  (an `is_fused=False` `warnings.warn(...)` "use the original pytorch
+  implementation") — classify with `perf-host-analysis` / `perf-nsight-systems`
+  first. The model uses the **noaux_tc** router (sigmoid gating + per-expert
+  correction bias, `n_group` groups, `topk_group`, top-`k` experts,
+  `routed_scaling_factor`): DeepSeek-V3/R1, Kimi K2. (Qwen3 models use
+  Renormalize-style routing, not noaux_tc — out of scope here.) The current path
+  computes scores as a chain of PyTorch elementwise ops (sigmoid, `+bias`) that
+  materializes `scores` / `scores_with_bias` to HBM, then runs a separate
+  group-score kernel (writing a `group_scores` intermediate) and a generic
+  bitonic-sort top-k. **Bounded config** (else it must fall back to PyTorch):
+  as of `84926bcb6f14`: `n_group>1` → `top_k ≤ 8`, `num_experts ≤ 256`,
+  `experts_per_group ≤ 32`, `experts_per_group·topk_group ≤ 256`; `n_group==1` →
+  `num_experts ≤ 1024`, `top_k ≤ 32`. The live guard is
+  `routing.py::Deepseekv3RoutingImpl.noaux_tc` — re-read it from your checkout.
+- **Counter-signals:** config outside the fused-path bounds (fails the
+  `is_fused` guard — no win without re-deriving bounds for the new shape,
+  and the kernel hard-checks rather than degrades); prefill-heavy / large-batch
+  regimes where routing is a negligible share of step time (measure the routing
+  share first); workloads requiring bit-exact parity with the PyTorch router's
+  tie-breaking order on equal scores.
+- **Mechanism:** collapse the whole router — sigmoid, bias add, per-group score
+  (sum of top-2), top-group selection, in-group expert top-k, and normalize ×
+  `routed_scaling_factor` — into **one block-per-token kernel**. Removes the
+  PyTorch elementwise launches and the `scores_with_bias` / `group_scores`
+  HBM round-trips (kernel now takes raw `logits` + `bias`), and replaces the
+  generic bitonic-sort top-k with a warp-shuffle packed-value top-k reduction
+  (the same `moeTopKFuncs.cuh::reduceTopK` the TRTLLM-Gen MoE routing kernels
+  use). Net: fewer launches + less HBM traffic + cheaper top-k; optionally
+  overlaps with neighbors via PDL.
+- **Generalizes to:** the pattern "collapse a small launch-bound op-chain into
+  one fused per-token kernel, reusing an existing warp primitive." Carries to
+  other grouped/top-k routers (e.g. Llama-4 MoE — similar shape, different bounds),
+  fused sampling top-k/top-p (reuse `reduceTopK`), and any per-token
+  "activation → reduction → select" chain. Adapt by: re-deriving the
+  supported-config bounds for the new shape and keeping the Python `is_fused`
+  guard in sync with them, and reusing an existing reduction primitive over a new
+  sort. The bounded-guard-plus-fallback is part of the pattern, not incidental.
+- **Apply via:** **not a server config knob** — a kernel/op change. Delegate to
+  **kernel-cuda-specialist** (raw CUDA C++):
+  `cpp/tensorrt_llm/kernels/noAuxTcKernels.cu` (+`.h`), the op binding
+  `cpp/tensorrt_llm/thop/noAuxTcOp.cpp`, and the model routing path
+  `tensorrt_llm/_torch/models/modeling_deepseekv3.py` (`noaux_tc`). Reuse the
+  existing `reduceTopK` rather than writing a new sort. PDL is gated on
+  `getEnvEnablePDL()` (`TRTLLM_ENABLE_PDL`).
+- **Expected effect:** lower MoE-routing latency, fewer kernel launches, and a
+  smaller inter-kernel GPU-idle gap around gating; reduced HBM traffic from the
+  eliminated intermediates. Largest at **decode / small token counts** where
+  launch overhead dominates. Measured Δ to be recorded from an nsys trace
+  (routing kernel count + total duration) before vs after — PR #7761 ships the
+  optimization but states no public speedup; do **not** cite a number not
+  measured locally.
+- **Accuracy risk:** lossless — numerically equivalent (same sigmoid / bias /
+  group-top-k / expert-top-k / normalize math, within fp tolerance). Still a
+  kernel rewrite, so **verify vs reference**; tie-breaking on equal scores can
+  differ from the PyTorch path.
+- **Verify:** correctness — `pytest tests/unittest/_torch/thop/parallel/test_noaux_tc.py`
+  and `test_moe.py`, plus the C++ `routingDeepSeekTest`; compare
+  `topk_values`/`topk_indices` to the PyTorch reference within tolerance,
+  including tie cases. Perf — nsys before/after: routing kernel count, total
+  duration, inter-kernel gap, and the fused kernel's SOL. Also confirm
+  out-of-bound configs (`top_k>8`, experts past the cap) take the PyTorch
+  fallback (the `is_fused=False` warning) and still produce correct results.
+- **Rollback:** force the PyTorch router (`self.is_fused = False`) or revert the
+  op to the precompute-scores signature; trigger: any accuracy mismatch vs
+  reference outside tolerance, or >5% perf regression. Note the kernel
+  **hard-checks** unsupported configs (`TLLM_CHECK`), so the Python guard bounds
+  must stay in sync with the kernel's supported bounds — a mismatch errors
+  instead of falling back.
+- **Prior art:** PRs #7761 (`[TRTLLM-8637]`), #13186 (widens fused-path
+  eligibility; `tensorrt_llm/_torch/modules/fused_moe/routing.py`
+  `Deepseekv3RoutingImpl`);
+  `cpp/tensorrt_llm/kernels/noAuxTcKernels.cu`, `moeTopKFuncs.cuh`,
+  `cpp/tensorrt_llm/thop/noAuxTcOp.cpp`,
+  `tensorrt_llm/_torch/models/modeling_deepseekv3.py`; analogous TRTLLM-Gen MoE
+  routing kernels under `kernels/trtllmGenKernels/blockScaleMoe/Routing*.cu`.
+  Owning specialist: **kernel-cuda-specialist**. Related: `perf-host-analysis`
+  (inter-kernel gap), `perf-nsight-systems` (launch count).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-qk-norm-rope.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-qk-norm-rope.md
new file mode 100644
index 000000000000..83297a1b5d43
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-qk-norm-rope.md
@@ -0,0 +1,36 @@
+---
+id: case-fuse-qk-norm-rope
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [launch, memory]
+signals: [many-small-kernels, hbm-roundtrip-between-ops]
+architectures: [any-sm]
+model_scope: [qwen3-moe, gemma3, llama-family, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-fuse-chain-feeding-gemm]
+accuracy_risk: lossless
+apply_via_kind: [default-on, code-change]
+knobs: [fuse_qk_norm_rope]
+specialists: [kernel-cuda-specialist]
+commits: ['9c4b8f66b454']
+eligibility:
+  - "model applies per-head RMSNorm to Q and K before RoPE (qk_norm_type == pre_rope)"
+  - "bf16 only; rope variant/layout must be supported by the fused kernel (fall back to unfused otherwise)"
+measured: []
+---
+
+# Fuse QK-Norm with RoPE in attention pre-processing
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `9c4b8f66b454` feat: Integration of Fused QKNorm+RoPE (#4611).
+- **Applies when:** launch/memory-bound attention pre-processing on models applying per-head RMSNorm to Q and K before RoPE (`qk_norm_type == pre_rope`). Signals: separate `q_norm`/`k_norm` RMSNorm + a separate RoPE kernel per layer; QKV laid out for an in-place op.
+- **Mechanism:** replaces (Q/K RMSNorm) → (RoPE) with one in-place op `torch.ops.trtllm.fused_qk_norm_rope(qkv, ...)` that normalizes Q,K and applies rotary in a single kernel over packed QKV — saving two RMSNorm launches + round-trips and the standalone RoPE launch per layer. Flag `fuse_qk_norm_rope` (default True on supporting models); when set, RoPE is removed from base attention and `apply_rope` dispatches to `apply_qk_norm_rope`. Registered in-place → CUDA-graph / torch.compile safe.
+- **Generalizes to:** "fold the norm that immediately precedes a positional/elementwise transform into that transform's kernel"; carries to post-rope QK-norm, DiT QK-norm+RoPE, partial-rotary/YaRN variants, split+norm+rope; adapt by matching norm placement (pre vs post rope), rotary interleave/mRoPE layout, and supported dtype (bf16); fall back to unfused when freq derivation doesn't match.
+- **Apply via:** set `fuse_qk_norm_rope=True` on the attention module (exposed in `modeling_qwen3.py` etc.); op `torch.ops.trtllm.fused_qk_norm_rope`. Delegate to **kernel-cuda-specialist** to wire a new model.
+- **Expected effect:** fewer attention pre-processing kernels per layer → lower attention latency / better launch efficiency at small batch; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless in principle (same QK-norm+RoPE math, one kernel) but bf16-only and layout-sensitive — enabling on a non-bf16/interleaved-rope model without kernel support is a correctness risk; parity-check Q/K on a new model.
+- **Verify:** kernel count / nsys (one fused op vs two norms + rope); Q/K output parity vs unfused; downstream accuracy.
+- **Rollback:** `fuse_qk_norm_rope=False` (falls back to separate `apply_qk_norm` + `rotary_emb`). Trigger: Q/K parity mismatch, non-bf16 dtype, unsupported rope variant.
+- **Prior art:** PR #4611. Files: `_torch/modules/attention.py` (`apply_qk_norm_rope`), `modeling_qwen3.py`, `_torch/compilation/utils.py` (in-place reg). Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md
new file mode 100644
index 000000000000..eb521cdb9f62
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md
@@ -0,0 +1,36 @@
+---
+id: case-hw-matched-lowprec-moe-gemm
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [compute, memory]
+signals: [gemm-dominates-step, weight-reads-dominate]
+architectures: [sm89, sm90, sm100]
+model_scope: [moe, deepseek-v3r1, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-hw-matched-lowprec-gemm-backend]
+accuracy_risk: lossy
+apply_via_kind: [config-knob]
+knobs: [moe_backend]
+specialists: [kernel-cuda-specialist, kernel-cute-specialist, perf-nsight-compute-analysis, trtllm-moe-develop, perf-sweep-challenger]
+commits: ['7bb0a78631de', '20b42912cef7']
+eligibility:
+  - "sm == 100: FP8 128x128 block-scale grouped GEMM via moe_backend=DEEPGEMM — requires the deep_gemm dependency and an FP8 block-scale checkpoint"
+  - "sm == 90 or sm == 89: W4A8/W4AFP8 — requires a ModelOpt-calibrated W4AFP8 checkpoint (quant_config.is_int4_weight_only_per_group())"
+measured: []
+---
+
+# Pick the hardware-matched low-precision MoE GEMM (FP8 block-scale on Blackwell, W4A8 on Hopper)
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `7bb0a78631de` Deepseek R1 FP8 Support on Blackwell (#6486); related: `20b42912cef7` [feat] Support DeepSeek-R1 W4A8 on Hopper (#4123).
+- **Applies when:** compute-bound (and weight-memory-bound) MoE on DeepSeek-class models, choosing expert-GEMM precision/kernel to match the GPU's tensor-core support. Blackwell sm_100: FP8 **128×128 block-scale** grouped GEMM (DeepGEMM). Hopper sm_90 / Ada sm_89: **W4A8/W4AFP8** (INT4 weights + FP8 activations) to shrink weight bytes and hit FP8 tensor cores.
+- **Mechanism:** Blackwell — `DeepGemmFusedMoE` calls `deep_gemm` `m_grouped` FP8 GEMM (`deepgemm_fp8_group_blockwise_gemm`, FP8 e4m3 → BF16) with 128-wide block scales as **UE8M0** (`per_token_cast_to_fp8_e8m0`, `per_block_cast_to_fp8_e8m0`, amax/448.0), re-laying weight scales via `resmooth_to_fp8_e8m0`/`transform_sf_into_required_layout` when `has_fp8_block_scales() and sm==100`; per-token act quant in a Triton kernel. Hopper — `has_w4afp8` packs experts as `torch.quint4x2` with `use_w4a8_group_scaling=True`, per-group 1×128 weight scales + per-tensor act scales (`FusedMoEQuantScalesW4A8`); mixed-precision CUTLASS MoE GEMM does INT4×FP8. Both halve/quarter expert weight bytes; Blackwell prefers native FP8 block-scale, Hopper goes INT4-weight to cut bandwidth and reuse FP8 math.
+- **Generalizes to:** "select the low-precision GEMM format and kernel backend by GPU tensor-core capability, not one-size-fits-all"; carries to dense linears (Blackwell FP8 block-scale vs Hopper W4A8/W4A16), NVFP4 on Blackwell, other grouped-expert models (Mixtral, Qwen-MoE); adapt by matching format to arch, wiring the right scale layout (UE8M0 vs 1×128 group), keeping a non-quant fallback.
+- **Apply via:** `moe_backend` (`model_config.moe_backend`, e.g. `"DEEPGEMM"`; `--moe_backend DEEPGEMM`) on Blackwell with an FP8 block-scale checkpoint; on Hopper load a W4AFP8 checkpoint (ModelOpt-calibrated act scales) so `quant_config` reports `is_int4_weight_only_per_group()`. Blackwell path needs the `deep_gemm` dependency. Delegate to **kernel-cuda-specialist** / **kernel-cute-specialist**; profile with **perf-nsight-compute-analysis**.
+- **Expected effect:** higher MoE GEMM throughput + lower expert-weight memory on the matched GPU (FP8 ~½ BF16 bytes; W4A8 ~¼ weight bytes); no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossy (FP8 e4m3 block-scale; W4A8 = INT4 weights + FP8 activations). Each needs an on-disk accuracy record (GSM8K/MMLU) + rollback criterion before promotion. W4A8 depends on a calibrated act-scale file; bad calibration degrades accuracy.
+- **Verify:** MoE GEMM throughput + weight footprint; accuracy parity vs a higher-precision backend (CUTLASS BF16/FP8); Blackwell confirm UE8M0 layout (sm_100 gate), Hopper confirm W4A8 group scaling (sm_90).
+- **Rollback:** switch `moe_backend` to `CUTLASS` (or BF16/FP8-per-tensor) / load higher-precision checkpoint. Trigger: accuracy regression beyond threshold, or backend unsupported on the running SM.
+- **Prior art:** PRs #6486, #4123. Files (Blackwell): `_torch/modules/fused_moe/fused_moe_deepgemm.py`, `create_moe.py`, `quantization/utils/fp8_utils.py`, `modeling_deepseekv3.py`. (Hopper): `_torch/modules/fused_moe.py` (`has_w4afp8`, `FusedMoEQuantScalesW4A8`), `thop/moeOp.cpp`, `examples/quantization/quantize_mixed_precision_moe.py`. Owning skill: **trtllm-moe-develop**; gate with **perf-sweep-challenger**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/index.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/index.md
new file mode 100644
index 000000000000..ac1bc40fdcaa
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/index.md
@@ -0,0 +1,158 @@
+# Casebook — Kernel & Fusion
+
+Optimizations at the kernel layer: fusing ops, swapping in fused/specialized
+kernels, and writing custom kernels. **Implementation is always delegated to a
+`kernel-*` specialist/skill** — this casebook only records when each precedent
+applies and how to verify it. A numerically-equivalent kernel rewrite is
+`lossless`, but equivalence is "within tolerance" and **must be verified
+against a reference**, since a kernel bug changes outputs. Add entries using
+the schema in [case-template.md](../case-template.md).
+
+> Routing note: classify the bottleneck first (`perf-analysis` / SOL% from
+> `perf-nsight-compute-analysis`). Kernel work pays off when a kernel is the
+> hot path and is far from SOL, or when many small ops can be fused to cut
+> launch/memory traffic. See `perf-optimization` for the specialist routing
+> table (Triton vs TileIR two-step vs CuTe DSL).
+
+## Recurring patterns in this family
+
+Match on these transferable patterns, not on a case title — most situations are
+a *variation* of a case, not its exact instance. Each pattern is the reusable
+idea; each links to its worked case.
+
+- **Collapse a small, launch-bound op-chain into one fused kernel.** Pre/post
+  steps around the big GEMMs operate on small per-token tensors — routing/gating,
+  sampling, RoPE, norm epilogues, KV scatter — yet run as many tiny kernels with
+  HBM round-trips between them. On the decode hot path their launch + memory
+  overhead is disproportionate. Fuse the chain into one block/warp-per-token
+  kernel, keep intermediates in registers/shared, and reuse an existing
+  purpose-built primitive instead of a generic algorithm. A fused kernel needs a
+  bounded-config guard + fallback to the generic path. _(Instance: [the noaux_tc routing case](fuse-moe-routing-kernel.md).)_
+- **When a Triton kernel is host-bound, move it to a precompiled C++/CUDA custom
+  op.** A small Triton kernel on the decode hot path can be cheap on the GPU yet
+  dominated by Triton's *host-side* cost — per-launch Python dispatch,
+  autotune-cache lookup, argument marshalling, and JIT/recompilation. Re-implement
+  it as a C++/CUDA `torch.ops.trtllm.*` custom op (precompiled, thin launch path)
+  with a `register_fake` meta impl so it still traces under `torch.compile`. This
+  is a host-overhead fix, not a SOL fix — classify host-bound first. _(Instance:
+  [the DSA indexer gather / index-convert ops](triton-to-cpp-op.md).)_
+- **Mark a Triton param `tl.constexpr` only when it must be compile-time.** Triton
+  recompiles a fresh kernel per distinct value of every `tl.constexpr` argument,
+  so a runtime scalar that varies across calls (e.g. a per-layer `layer_id`, a
+  stride, a block count) marked `constexpr` forces a recompile per value. Keep
+  `tl.constexpr` only where the value is genuinely needed at compile time (tile
+  sizes used in `tl.arange`, static loop unrolling); pass the rest as runtime
+  args. _(Instance: [the DSA `convert_req_index_to_global` kernel](relax-tl-constexpr.md).)_
+- **Fuse the chain that feeds the next GEMM.** Beyond a generic op-chain: the TP
+  collective *epilogue* (AllReduce + residual-add + RMSNorm + quant), a local
+  add + RMSNorm + quant, QK-Norm + RoPE, or folding a standalone layout pre-pass
+  (scale-factor swizzle) into the consumer kernel — and, when every stage is small,
+  mega-fusing a whole multi-stage pipeline (EP-dispatch + GEMM1 + SwiGLU + GEMM2 +
+  EP-combine) into one launch. Removes launches + HBM round-trips; **the quant step
+  makes it lossy** → parity-check. Most of these land as `torch.compile`
+  PatternMatcher rewrites or a single custom op. _(Instances: the [AR-epilogue](fuse-ar-epilogue.md), [add+norm+quant](fuse-add-norm-quant.md), [QK-Norm+RoPE](fuse-qk-norm-rope.md), [scale-swizzle](fold-scale-swizzle-into-kernel.md), and [MegaMoE](mega-fuse-moe-deepgemm.md) cases.)_
+- **Quantize the memory-bound side, keep the precision-sensitive boundary high.**
+  Store the KV cache (and the matching operand) in FP8/FP4 to cut HBM bytes, but
+  keep accumulation/output in BF16 where the next op (e.g. `o_proj`) is sensitive.
+  Lossy → accuracy record + rollback. Carries to GQA/MHA FP8 KV, FP8 context FMHA,
+  INT8/FP4 KV. _(Instance: [the FP8 MLA KV case](fp8-mla-kv-cache.md).)_
+- **Split a monolithic kernel that does compute + a serial cross-partition
+  reduction into two specialized kernels (main + reduction)** to shed the
+  reduction's register/smem footprint, raise occupancy, and parallelize the reduce.
+  Lossless but a kernel rewrite — verify parity. Carries to flash-attn split-KV/
+  split-K combine, GEMM split-K epilogue, any persistent kernel whose tail
+  reduction throttles occupancy. _(Instance: [the separate-reduction case](split-mla-reduction-kernel.md).)_
+- **Pick the hardware-matched low-precision grouped-GEMM backend over the generic
+  one** when the GEMM is the bottleneck and the dtype is natively supported
+  (trtllm-gen / DeepGEMM / MegaMoE on Blackwell, W4A8 on Hopper, vs generic
+  CUTLASS). Match the cubin/scale-layout to the SM and keep a non-quant fallback.
+  Lossy → accuracy record. _(Instances: the [trtllm-gen FP4](trtllm-gen-fp4-moe-backend.md), [MegaMoE](mega-fuse-moe-deepgemm.md), and [HW-matched-GEMM](hw-matched-lowprec-moe-gemm.md) cases.)_
+- **Attend to only a selected subset of KV in the attention kernel.** For sparse /
+  block-sparse / windowed attention (DeepSeek V3.2 DSA picks top-k KV per query),
+  pass the selected page/block indices and bound the kernel's KV loop **and** its
+  cross-CTA reduction by the selection count instead of the full sequence. Two
+  corollaries travel with it: **write K/V straight into the paged cache**
+  (absorption/latent layout) to skip a materialization pass, and **early-exit CTAs
+  beyond the ragged per-sequence length**. _(Instance: [the sparse-MLA top-k attention case](sparse-mla-topk-attention.md).)_
+- **The hot-path top-k / selection is a specialized kernel, not a generic sort.**
+  Fuse the selection with its index/mask post-processing into one radix-select,
+  keep several exact variants (single-CTA / multi-CTA / multi-pass radix) and
+  dispatch by `(k, N, dtype, SM-wave-occupancy)`; short-circuit with a coarse
+  low-precision first pass and a **warm-start from the previous decode step's
+  answer** (temporal correlation). Carries to MoE routing top-k, sampling,
+  beam-search. _(Instance: [the top-k selection kernel case](specialize-topk-selection-kernel.md); cf. the [MoE routing kernel](fuse-moe-routing-kernel.md).)_
+- **Fuse a data-movement prologue into the quantize that follows.** When a
+  `cat`/`split`/`gather`/`rotate` immediately precedes an FP8/FP4 quantize, collapse
+  both into one kernel and **read the strided input views directly** to skip the
+  contiguous-copy round-trip. The quantize-consumer sibling of "fold a layout
+  pre-pass into the consumer." _(Instance: [the fuse-data-movement-into-quantize case](fuse-datamovement-into-quantize.md).)_
+- **A ranking-only computation tolerates reduced precision — and verify the fast HW
+  path is actually taken.** A GEMM whose output only feeds an `argmax`/top-k/threshold
+  (a selector/scoring GEMM) can run on TF32/BF16 tensor cores; keep the
+  higher-precision contract only where a scale/softmax needs it, and **confirm the
+  accelerated kernel is dispatched** — small-M GEMMs silently fall back to CUDA-core
+  SGEMM even when you asked for it. _(Instance: [the ranking-only TF32 case](ranking-only-precision-tf32.md).)_
+- **Fusion is not monotonic — re-evaluate the fuse boundary per dtype/checkpoint.**
+  Concatenating co-resident projections into one wide GEMM wins **only when the fused
+  weights share a dtype**; a mixed-precision checkpoint can force *unfusing*, and
+  unfusing can unlock a conditional skip. The caveat that keeps every "fuse!" pattern
+  above from being applied blindly. _(Instance: [the fusion-boundary case](reevaluate-fusion-boundary-per-dtype.md).)_
+
+## Cases
+
+Match on **Applies when** / **Generalizes to**, then open that case file.
+Implementation is always delegated to a `kernel-*` specialist/skill.
+
+> Risk key: **lossless** = numerically equivalent (still verify vs a reference —
+> a kernel bug changes outputs) · **lossy** = needs an accuracy record + rollback
+> before promotion · **mixed** = base fusion lossless but the quant variant lossy.
+
+### Fusion (collapse op-chains)
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [Fuse the TP AR epilogue](fuse-ar-epilogue.md) | comm-bound TP>1; `allreduce → add → RMSNorm → (quant)` chain per layer | fuse the collective's epilogue into the collective kernel | mixed |
+| [Fuse add + norm + quant](fuse-add-norm-quant.md) | memory-bound epilogue, NO collective; separate `add → RMSNorm → quant` kernels | fuse a pointwise+reduction+quant chain that feeds the next GEMM | lossy |
+| [Fuse QK-Norm + RoPE](fuse-qk-norm-rope.md) | launch/mem-bound attention pre-proc; per-head Q/K RMSNorm before RoPE (Qwen3/Gemma3), bf16 | fold the norm preceding a positional/elementwise transform into that kernel | lossless |
+| [Fold scale-swizzle into the kernel](fold-scale-swizzle-into-kernel.md) | NVFP4 WideEPMoE; a standalone `swizzle_sf` relayout runs before the grouped GEMM | fold a standalone layout pre-pass into the kernel that consumes the data | lossless |
+| [Fuse MoE grouped-top-k routing](fuse-moe-routing-kernel.md) | MoE router = many tiny kernels on decode hot path (noaux_tc: DS-V3/Kimi-K2), or fell back to PyTorch router | collapse a small launch-bound op-chain into one per-token kernel, reuse a warp primitive | lossless |
+| [Fuse data-movement into the quantize](fuse-datamovement-into-quantize.md) | a `cat`/`split`/`gather`/`rotate` immediately precedes an FP8/FP4 quantize; standalone layout kernel + copy in the profile | fuse the layout prologue into the quantize; read strided views to skip the copy | lossless |
+| [Re-evaluate the fusion boundary per dtype](reevaluate-fusion-boundary-per-dtype.md) | deciding whether to fuse GEMMs sharing an input; a mixed-precision checkpoint splits participants, or a fused op blocks a skip | fusion is not monotonic — the optimal boundary is dtype/checkpoint-dependent | lossless |
+
+### Attention / MLA kernels
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [FP8 MLA KV cache](fp8-mla-kv-cache.md) | memory-bound MLA decode; KV reads dominate; DS-MLA, SM90/100, long ctx / high conc | quantize the memory-bound side (KV+operand) to FP8, keep output BF16 | lossy |
+| [Split MLA into a reduction kernel](split-mla-reduction-kernel.md) | compute/occupancy-bound MLA decode on TRTLLM-Gen FMHA multi-CTA-KV; big reduction tiles cap occupancy | split a compute+serial-reduction kernel into main + separate reduction | lossless |
+| [Sparse-MLA top-k attention](sparse-mla-topk-attention.md) | attention over a top-k–selected KV subset (DSA / block-sparse / windowed), SM100+; full-seq FMHA wasteful | attend to only the selected KV — bound the kernel's loop/reduction by the selection count | mixed |
+| [Specialize the top-k selection kernel](specialize-topk-selection-kernel.md) | hot-path top-k (indexer / MoE router / sampling) = a `torch.topk` sort + index/mask ops, every step | specialized radix-select; dispatch by (k,N,dtype,wave); coarse-pass + prev-step warm-start | lossless |
+
+### MoE / grouped-GEMM kernels
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [trtllm-gen FP4 MoE backend](trtllm-gen-fp4-moe-backend.md) | compute-bound MoE grouped-GEMM, sm_100a, NVFP4/FP8-block-scale, on default CUTLASS | pick a HW-specialized low-precision grouped-GEMM backend over the generic one | lossy |
+| [Mega-fuse MoE into DeepGEMM](mega-fuse-moe-deepgemm.md) | launch-bound MoE (many per-stage kernels + EP a2a), sm_100, W4A8_MXFP4_MXFP8, EP-only, shape%128 | mega-fuse a multi-stage pipeline whose stages are each launch/mem-bound | lossy |
+
+### Quantization (low-precision GEMM)
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [HW-matched low-precision MoE GEMM](hw-matched-lowprec-moe-gemm.md) | compute/weight-mem-bound MoE (DeepSeek); pick precision by GPU (FP8-block Blackwell / W4A8 Hopper) | select the low-precision GEMM format + backend by GPU tensor-core capability | lossy |
+| [Ranking-only precision (TF32)](ranking-only-precision-tf32.md) | a scoring/selector GEMM feeds only an argmax/top-k; runs as CUDA-core SGEMM at small M | a ranking-only output tolerates TF32 tensor cores — and verify the path is actually taken | lossy |
+
+### Triton / kernel host-overhead hygiene
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [Triton kernel → C++ custom op](triton-to-cpp-op.md) | a Triton kernel is host-bound (tiny GPU work, heavy launch/JIT); DSA indexer gather/index-convert | move a host-bound Triton kernel to a precompiled C++/CUDA op (+`register_fake`) | lossless |
+| [Relax unnecessary tl.constexpr](relax-tl-constexpr.md) | a Triton kernel marks runtime-varying scalars (`layer_id`, stride) `tl.constexpr` → recompiles per value | mark `tl.constexpr` only when compile-time-required; pass the rest as runtime args | lossless |
+
+## Suggested slots (optional — replace or delete)
+
+Still open in this family: elementwise/activation fusion (delegate to
+**kernel-triton-writing** / **kernel-cute-writing**), TileIR optimization of an
+existing Triton kernel on Blackwell (the two-step pipeline via
+**kernel-tileir-optimization**), CuTe DSL generation (**kernel-cute-writing**),
+and GEMM tactic / autotune selection.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md
new file mode 100644
index 000000000000..0cf219c88911
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md
@@ -0,0 +1,41 @@
+---
+id: case-mega-fuse-moe-deepgemm
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [launch, memory]
+signals: [many-small-kernels, hbm-roundtrip-between-ops]
+architectures: [sm100]
+model_scope: [moe, deepseek-v3r1, model-agnostic]
+phase: [decode]
+patterns: [pattern-fuse-chain-feeding-gemm, pattern-hw-matched-lowprec-gemm-backend]
+accuracy_risk: lossy
+apply_via_kind: [config-knob]
+knobs: [moe_backend]
+specialists: [kernel-cute-specialist]
+commits: ['6e069b69ef3c']
+eligibility:
+  - "sm == 100"
+  - "quant exactly W4A8_MXFP4_MXFP8; SwiGLU activation"
+  - "expert-parallel only: moe_tp_size == 1 and cluster_size == 1"
+  - "hidden_size % 128 == 0 and intermediate_size % 128 == 0"
+  - "live gate: mega_moe/backend.py::can_implement — bounds as of 6e069b69ef3c"
+interactions:
+  - {feature: TRTLLMGenFusedMoE, relation: alternative-to, note: "same W4A8_MXFP4_MXFP8 math, different kernel — the parity reference and rollback target"}
+measured: []
+---
+
+# Mega-fuse the entire MoE pipeline into one DeepGEMM launch (EP-dispatch + GEMM1 + SwiGLU + GEMM2 + EP-combine)
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `6e069b69ef3c` [feat] Add MegaMoEDeepGemmFusedMoE backend wrapping DeepGEMM fp8_fp4_mega_moe (#13384).
+- **Applies when:** launch-bound (many small per-stage MoE kernels + EP all-to-all) on sm_100; quant exactly `W4A8_MXFP4_MXFP8`; expert-parallel only (`moe_tp_size=1`, `cluster_size=1`); `hidden_size%128==0` and `intermediate_size%128==0`; SwiGLU activation.
+- **Mechanism:** replaces the multi-kernel MoE sequence (separate EP dispatch, GEMM1, SwiGLU, GEMM2, EP combine) with one fused kernel `deep_gemm.fp8_fp4_mega_moe`. Math matches `TRTLLMGenFusedMoE`'s W4A8_MXFP4_MXFP8 path; win is one launch vs 5+ stages, collapsing launch overhead + intermediate HBM round-trips. Inputs MXFP8 per-token quantized; weights pre-transformed via `transform_weights_for_mega_moe`.
+- **Generalizes to:** "mega-fuse a multi-stage pipeline whose stages are individually launch/memory-bound into one kernel"; carries to other fused-MoE megakernels, fused attention pipelines, any EP-dispatch+GEMM+combine chain; adapt by matching strict preconditions (arch, quant, EP-only, shape%128) and providing a clean fallback when unmet.
+- **Apply via:** `moe_backend='MEGAMOE_DEEPGEMM'` (resolved in `create_moe.py::get_moe_cls`); self-gates via `can_implement(...)`, raises `_MegaMoEUnavailable` (graceful fallback) if bundled `tensorrt_llm.deep_gemm` lacks `fp8_fp4_mega_moe`. Delegate to **kernel-cute-specialist** / DeepGEMM owners.
+- **Expected effect:** fewer launches + less intermediate memory traffic → lower MoE decode latency; no number (code notes fused MXFP8 quant ~11 us/launch vs per-seq Triton fallback) — measured Δ to be recorded from run.
+- **Accuracy risk:** lossy (W4A8_MXFP4_MXFP8 MoE GEMM) but output-equivalent to the existing `TRTLLMGenFusedMoE` W4A8 backend by design (same math, different kernel) — parity is against that backend, not bf16.
+- **Verify:** numeric parity vs `TRTLLMGenFusedMoE` same model/quant; task accuracy vs that backend; MoE collapses to a single kernel in nsys.
+- **Rollback:** `moe_backend` → `TRTLLM` or `CUTLASS`. Trigger: `can_implement` rejects env (non-SM100, wrong quant, tp>1), missing DeepGEMM symbols, or parity failure.
+- **Prior art:** PR #13384. Files: `_torch/modules/fused_moe/mega_moe/backend.py` (`can_implement`), `create_moe.py`, DeepGEMM `fp8_fp4_mega_moe`/`transform_weights_for_mega_moe`. Owning specialist: **kernel-cute-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/ranking-only-precision-tf32.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/ranking-only-precision-tf32.md
new file mode 100644
index 000000000000..ed76afc4e6cc
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/ranking-only-precision-tf32.md
@@ -0,0 +1,87 @@
+---
+id: case-ranking-only-precision-tf32
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [compute]
+signals: [kernel-far-from-sol, slow-path-fallback, small-batch-decode]
+architectures: [any-sm]
+model_scope: [sparse-attention, deepseek-v32, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-ranking-only-reduced-precision]
+accuracy_risk: lossy
+apply_via_kind: [code-change]
+knobs: []
+specialists: [kernel-cuda-specialist, perf-nsight-compute-analysis]
+commits: ['24bd7a2c20', '4b915d4f9b', '79a6c9742b']
+interactions:
+  - {case: case-reevaluate-fusion-boundary-per-dtype, relation: composes-with, note: "the same indexer GEMM — precision lever here, fuse/unfuse boundary there"}
+measured: []
+---
+
+# Run a ranking-only GEMM in TF32 — and verify the tensor-core path is actually taken
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `24bd7a2c20` [perf] Force enable TF32 tensor cores for DSA indexer
+  fused GEMM (#13452); related: `4b915d4f9b` Fuse indexer wk + weights_proj into a
+  single TF32 GEMM (#12055), `79a6c9742b` [fix] Use fp32 for indexer weight_proj
+  GEMM (#9243, the precision contract these optimize around).
+- **Applies when:** a GEMM (or reduction) produces a value that is used **only to
+  rank or select** downstream — not as a precise numeric operand. The archetype:
+  DSA's lightning indexer projects hidden states to per-token **index scores**
+  whose sole use is a top-k KV selection; the score magnitudes never feed a
+  precise computation. Signals: a scoring/gating/routing GEMM feeding an
+  `argmax` / `top-k` / threshold; the GEMM runs in FP32 "for safety" and shows up
+  as a CUDA-core SIMT SGEMM (not a tensor-core GEMM) in ncu; small-M (decode) so
+  it's launch/latency-sensitive.
+- **Mechanism:** two coupled moves. (1) **The output only ranks, so drop the
+  compute to TF32** — keep FP32 as the *storage/accumulation contract* (a
+  softmax-scale factor still deserves FP32, #9243) but let the multiply run on
+  TF32 tensor cores (~10-bit mantissa), which is far faster than FP32 CUDA-core
+  SGEMM and precise enough to preserve the ranking. (2) **Verify the tensor-core
+  path is actually taken** — asking for FP32 does *not* guarantee TF32: a custom
+  `cublas_mm` pinned to `CUBLAS_COMPUTE_32F` (and cuBLASLt heuristics for small M)
+  silently fall back to CUDA-core SGEMM (#12055 asked for TF32 but didn't get it).
+  Force it with `torch.backends.cuda.matmul.allow_tf32 = True` **and** an entry
+  point that honors the flag — `F.linear` routes through PyTorch's cuBLAS handle
+  and dispatches `CUBLAS_COMPUTE_32F_FAST_TF32` (#13452 wraps the call in a
+  `_tf32_matmul_enabled()` context manager and swaps `cublas_mm → F.linear`).
+- **Generalizes to:** the pattern "**a computation whose result is consumed only
+  for ranking/selection tolerates reduced precision — run it on the fast
+  tensor-core path (TF32 / BF16), keep the higher-precision contract only where a
+  scale/softmax needs it, and confirm the accelerated path is dispatched (small-M
+  GEMMs silently fall back to CUDA cores).**" Carries to: MoE router logits before
+  `argmax`/top-k, retrieval/similarity scoring for nearest-neighbour selection,
+  attention-sink or landmark scoring, any gate whose output only picks. Adapt by:
+  (a) confirming the output is ranking-only (a downstream select, not a numeric
+  operand); (b) choosing the precision the hardware accelerates; (c) **verifying
+  dispatch** (ncu: tensor-core GEMM, not SIMT) rather than trusting the requested
+  compute type.
+- **Apply via:** **not a server knob** — a code change at the GEMM call site
+  (`allow_tf32` context + `F.linear` in place of a fixed-compute-type custom op).
+  Delegate to **kernel-cuda-specialist**; confirm the dispatched kernel with
+  **perf-nsight-compute-analysis** (the whole point is that the requested
+  precision ≠ the executed kernel).
+- **Expected effect:** the scoring GEMM moves off CUDA cores onto tensor cores →
+  lower indexer/scoring latency at small M, with the selection unchanged.
+  Direction only — measured Δ (GEMM time, kernel name/tensor-core utilization,
+  decode throughput) to be recorded from the run.
+- **Accuracy risk:** **lossy but bounded** — TF32 rounding perturbs the score
+  values, but they are used only to rank, so the selected set is (almost always)
+  identical. Lower-risk than KV/weight quant, but still verify: it can in principle
+  flip a near-tie selection. Confirm top-k selection overlap and e2e accuracy are
+  unchanged on enable; keep an accuracy note if promoting as a default.
+- **Verify:** correctness — top-k index overlap vs the FP32 path on representative
+  logits; DSA e2e accuracy (GSM8K/MMLU) within tolerance. Perf — ncu confirms a
+  tensor-core GEMM (e.g. `…_tf32…`), not a SIMT SGEMM; GEMM time and throughput
+  before vs after.
+- **Rollback:** remove the `allow_tf32` context (revert to the FP32 path). Trigger:
+  any selection/accuracy drift beyond tolerance, or no measured speedup (path was
+  already tensor-core).
+- **Prior art:** PRs #13452, #12055, #9243. Files:
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` (`_tf32_matmul_enabled`,
+  `_fused_wk_wp_weight`, `F.linear` vs `torch.ops.trtllm.cublas_mm`). Owning
+  specialist: **kernel-cuda-specialist**; dispatch check:
+  **perf-nsight-compute-analysis**. Related: the [fusion-boundary case](reevaluate-fusion-boundary-per-dtype.md)
+  (the same indexer GEMM, the orthogonal fuse/unfuse lesson).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/reevaluate-fusion-boundary-per-dtype.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/reevaluate-fusion-boundary-per-dtype.md
new file mode 100644
index 000000000000..0b9d77c4abd0
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/reevaluate-fusion-boundary-per-dtype.md
@@ -0,0 +1,87 @@
+---
+id: case-reevaluate-fusion-boundary-per-dtype
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [compute, launch]
+signals: [many-small-kernels, mixed-dtype-checkpoint]
+architectures: [any-sm]
+model_scope: [sparse-attention, deepseek-v32, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-reevaluate-fusion-boundary]
+accuracy_risk: lossless
+apply_via_kind: [code-change]
+knobs: []
+specialists: [kernel-cuda-specialist, trtllm-moe-develop]
+commits: ['2f45640c19', 'e57d83c5dc', '4b915d4f9b']
+interactions:
+  - {case: case-ranking-only-precision-tf32, relation: composes-with, note: "the same indexer GEMM — fuse/unfuse boundary here, precision lever there"}
+measured: []
+---
+
+# Fusion is not monotonic — re-evaluate the GEMM fuse boundary per dtype/checkpoint
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `2f45640c19` [feat] Unfuse indexer.wk from attention GEMM for
+  DS-V3.2 NVFP4 (#11989); related: `e57d83c5dc` Fuse QK down_proj with indexer K +
+  weight_proj for FP4 ckpt (#8771, the original fuse), `4b915d4f9b` Fuse indexer
+  wk + weights_proj into a single GEMM (#12055, the narrower re-fuse).
+- **Applies when:** you are deciding whether to fuse several small GEMMs that read
+  the **same activation** into one wide GEMM (QKV projections, gate+up, LoRA A/B, a
+  down-projection feeding multiple heads). The trap: a fusion that was a clear win
+  for one checkpoint becomes **illegal or suboptimal** for another. Signals: a
+  **mixed-precision checkpoint** puts one fused participant in a different dtype
+  (e.g. attention weights NVFP4 but the indexer's `wk` still BF16); a fused op
+  forces a computation that a cheaper path could otherwise **skip**; the fused GEMM
+  asserts all participants share a dtype.
+- **Mechanism:** concatenating co-resident projections into one wide GEMM (weights
+  cat along the out-dim, materialized once in `post_load_weights`) cuts launches
+  and shares the activation read — **but only when the fused weights share a
+  dtype**. #8771 fused QK down_proj + `indexer.wk` + `weights_proj` into the
+  oversized `kv_a_proj_with_mqa` for an FP4 checkpoint (all same dtype). A later
+  FP4-*attention* checkpoint left `indexer.wk` in BF16 while the attention
+  projections went NVFP4, so the shared-dtype fuse was no longer valid → #11989
+  **unfused** `indexer.wk` back to a standalone GEMM, and **sited it so the
+  short-MHA path can elide it** (when the indexer isn't needed the extra GEMM is
+  skipped). #12055 then **re-fused at a narrower, same-(fp32)-dtype boundary**
+  (`wk`+`weights_proj` only). The optimum moved as the checkpoint dtypes changed.
+- **Generalizes to:** the pattern "**the optimal fusion boundary is
+  dtype/checkpoint/backend-dependent — fusion is not unconditionally a win.**"
+  Before fusing GEMMs that share an operand, check they share a dtype/quant format;
+  when a mixed-precision checkpoint splits them, **unfuse** the odd one out; and
+  weigh whether keeping an op **separate unlocks a conditional skip** (elide it when
+  its output isn't consumed). Carries to: quantized QKV/gate-up fusion across
+  FP8/NVFP4/BF16 checkpoints, LoRA-adapter fusion, any "fuse the projections"
+  refactor applied blindly across checkpoints. Adapt by: gating the fuse on a
+  runtime dtype check, keeping the boundary configurable per checkpoint, and
+  materializing the fused weight once at load — not per forward.
+- **Apply via:** **not a server knob** — model-definition code (Linear sizing +
+  `post_load_weights` weight materialization, with a dtype guard). Delegate to
+  **kernel-cuda-specialist** (GEMM) / **trtllm-moe-develop** or the model author for
+  module wiring.
+- **Expected effect:** correct + fastest fusion *for the checkpoint at hand* — fused
+  where dtypes align (fewer launches, one activation read), unfused where a
+  mixed-precision checkpoint or a skip-path makes separate better. Direction only —
+  measured Δ (projection GEMM time, launch count, throughput) to be recorded per
+  checkpoint.
+- **Accuracy risk:** **lossless** — fusing or unfusing GEMMs is numerically
+  equivalent (same math, combined or split); the risk is a *correctness* bug (a
+  dtype-mismatched fuse), not precision loss. The dtype-guard assert is what makes
+  an illegal fuse fail loudly instead of silently corrupting.
+- **Verify:** the fused/unfused path produces identical outputs vs the reference on
+  each target checkpoint; projection GEMM time + launch count before vs after **on
+  the specific checkpoint** (a win on one checkpoint does not transfer). Confirm the
+  skip-path (if relied on) actually elides the unfused op.
+- **Rollback:** flip the fuse flag / restore the separate (or fused) GEMM. Trigger:
+  a dtype-mismatch assert, an accuracy mismatch, or no measured benefit on the
+  target checkpoint.
+- **Prior art:** PRs #11989, #8771, #12055. Files:
+  `tensorrt_llm/_torch/models/modeling_deepseekv3.py` (`DeepseekV32Attention`,
+  `fuse_a_indexer_k_weight`, `post_load_weights`),
+  `tensorrt_llm/_torch/modules/attention.py`,
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`. Owning specialists:
+  **kernel-cuda-specialist**, **trtllm-moe-develop**. Related: the
+  [ranking-only TF32 case](ranking-only-precision-tf32.md) (same indexer GEMM, the
+  precision lever) and the many "fuse the chain" cases in this family — this is the
+  caveat that keeps them from being applied blindly.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/relax-tl-constexpr.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/relax-tl-constexpr.md
new file mode 100644
index 000000000000..001fca0eca15
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/relax-tl-constexpr.md
@@ -0,0 +1,69 @@
+---
+id: case-relax-tl-constexpr
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [host-overhead]
+signals: [recompilation-churn, host-prep-on-critical-path]
+architectures: [any-sm]
+model_scope: [sparse-attention, deepseek-v32, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-constexpr-only-when-required]
+accuracy_risk: lossless
+apply_via_kind: [kernel-change]
+knobs: []
+specialists: [kernel-triton-writing]
+commits: []
+interactions:
+  - {case: case-triton-to-cpp-op, relation: alternative-to, note: the heavier fix when host overhead persists after relaxing constexpr}
+measured: []
+---
+
+# Relax unnecessary tl.constexpr so a Triton kernel isn't recompiled per layer
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Applies when:** a Triton kernel marks parameters `tl.constexpr` that don't
+  need to be compile-time constants, and at least one such param takes many
+  distinct runtime values across calls — e.g. a per-layer `layer_id`, or a stride
+  / block-count that varies — so Triton recompiles a fresh kernel per value.
+  Symptoms: repeated Triton autotune/compile, a growing compile cache, host time
+  across the first N layers/shapes. Instance: the DSA
+  `_convert_req_index_to_global_index_kernel_with_stride_factor` (had `layer_id`,
+  `stride_factor`, `max_num_blocks_per_req`, `BLOCK_SIZE` all `tl.constexpr`).
+- **Mechanism:** Triton specializes (recompiles) per distinct value of every
+  `tl.constexpr` argument. Params used only as plain runtime scalars — passed to
+  arithmetic, indexing, or masks — don't need to be `constexpr`; marking them so
+  forces a recompile per value (`layer_id` → one compile per layer). Dropping
+  `tl.constexpr` from those, keeping it only where the value must be compile-time
+  (tile sizes used in `tl.arange`, static unrolling — here `BLOCK_N`), yields a
+  single kernel reused across all values.
+- **Generalizes to:** the pattern "**mark a Triton param `tl.constexpr` only when
+  it must be a compile-time constant (tiling/shape/`tl.arange`/unroll); otherwise
+  pass it as a runtime arg to avoid per-value recompilation.**" Carries to: any
+  Triton kernel with a per-layer/per-step/per-shape id or scalar marked
+  `constexpr`; block-count or stride params not used for tiling; "`constexpr` by
+  habit" annotations. Adapt by: for each `constexpr` param, check whether it is
+  used where a compile-time constant is *required* (`tl.arange` bounds, static
+  loop ranges); if not, drop the annotation. Trade-off: runtime args forgo a few
+  specialization-time optimizations, but on a host-bound path avoiding recompiles
+  wins.
+- **Apply via:** **not a server config knob** — a Triton-kernel edit. Delegate to
+  **kernel-triton-writing**. Touch point (prior art, TRT-LLM repo):
+  `tensorrt_llm/_torch/attention_backend/sparse/kernel.py` — the kernel signature
+  (keep `BLOCK_N: tl.constexpr`, relax the rest).
+- **Expected effect:** fewer Triton recompilations across layers/shapes → lower
+  first-iteration and steady host time, smaller GPU-idle; no change to GPU
+  compute. Record the local Δ (distinct-compile count, per-iteration host time) —
+  PR #12581 ships it without a public number; do **not** cite one not measured.
+- **Accuracy risk:** lossless — identical kernel math; only the compile-time-vs-
+  runtime treatment of scalar args changes.
+- **Verify:** correctness — kernel output unchanged vs reference (same DSA tests).
+  Perf — count distinct Triton compiles across a multi-layer run before vs after
+  (e.g. Triton compile logs / cache size); per-iteration host time. Confirm the
+  kept-`constexpr` param (`BLOCK_N`) still drives the `tl.arange` tiling.
+- **Rollback:** re-add `tl.constexpr`; trigger: a kernel-compute regression (if a
+  relaxed param turned out to enable a needed specialization) or any correctness
+  mismatch.
+- **Prior art:** PR #12581; `kernel.py` path above.
+  Owning specialist: **kernel-triton-writing**. Related: the [Triton→C++-op case](triton-to-cpp-op.md) (the heavier fix when host overhead persists after relaxing `constexpr`).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/sparse-mla-topk-attention.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/sparse-mla-topk-attention.md
new file mode 100644
index 000000000000..e44bcaf42198
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/sparse-mla-topk-attention.md
@@ -0,0 +1,117 @@
+---
+id: case-sparse-mla-topk-attention
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [compute, memory]
+signals: [attention-hot-path, kv-reads-dominate, long-context, high-concurrency]
+architectures: [sm100]
+model_scope: [sparse-attention, mla, deepseek-v32, model-agnostic]
+phase: [any-phase]
+patterns: [pattern-attend-selected-kv-subset, pattern-write-direct-into-paged-cache, pattern-early-exit-ragged-bounds]
+accuracy_risk: mixed
+apply_via_kind: [kernel-change]
+knobs: [kv_cache_config.dtype]
+specialists: [kernel-cuda-specialist, perf-nsight-compute-analysis]
+commits: ['f0dc746738', '497a07021d', '389b73c349']
+eligibility:
+  - "model is trained with sparse attention (DSA-style selection is the intended computation)"
+  - "sm >= 100 (trtllm-gen sparse MLA kernels)"
+  - "sparse_mla_topk < typical seqLenKv (else selection degenerates to dense plus overhead)"
+interactions:
+  - {case: case-specialize-topk-selection-kernel, relation: depends-on, note: producer of the top-k indices this kernel consumes}
+  - {case: case-fp8-mla-kv-cache, relation: composes-with, note: the quantized-KV variant rides on it (lossy part)}
+  - {case: case-split-mla-reduction-kernel, relation: composes-with, note: "independent levers on the same FMHA reduction; they co-apply — verify the combined effect"}
+  - {case: case-skip-sparse-path-when-degenerate, relation: composes-with, note: guards the short-seq degenerate regime}
+measured: []
+---
+
+# Attend only to the top-k–selected KV in the attention kernel (sparse MLA)
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `f0dc746738` [feat] Add trtllm-gen sparse MLA kernels to support
+  per-Tensor FP8 KV Cache (#8692); related: `497a07021d` optimized sparse mla
+  kernels && fix unspecified cuda launch (#8866), `389b73c349` [fix] Remove FP8
+  K/V buffer from TRTLLM sparse MLA attention kernel (#9529).
+- **Applies when:** an attention layer only needs to attend to a *selected subset*
+  of the KV history — a natively-sparse-attention model (DeepSeek V3.2's DSA,
+  where a lightning indexer picks the top-k KV per query), or block-sparse /
+  sliding-window / landmark attention. Signals: full-sequence FMHA is the hot
+  path yet most KV positions are masked out or unselected; long context / high
+  concurrency where reading all KV is wasteful; MLA on SM100+ trtllm-gen FMHA.
+  The selection indices are produced upstream (see the [top-k selection kernel
+  case](specialize-topk-selection-kernel.md)); this case is the *consumer* — the
+  attention kernel that uses them.
+- **Counter-signals:** models **not** trained with sparse attention — feeding a
+  dense-trained model selected-KV attention changes the computation itself (a
+  real accuracy change, out of scope for this precedent); short sequences where
+  `seqLenKv <= sparse_mla_topk` (selection degenerates to dense plus selection
+  overhead — see [skip-sparse-path](../runtime-execution/skip-sparse-path-when-degenerate.md));
+  non-MLA / non-trtllm-gen attention backends where the sparse dispatch does
+  not exist.
+- **Mechanism:** feed the FMHA generation kernel the selected KV page indices
+  (`kvPageIdxPtr = sparse_attn_indices`) and bound its KV loop by
+  `maxAttentionWindow = min(maxSeqLenKv, sparse_mla_topk)`, so it does O(top-k)
+  work instead of O(seqLen). Three reinforcing moves: (1) run MLA in
+  **absorption/latent space** so the RoPE kernel writes K/V straight into the
+  paged KV cache (latent dims `kv_lora_rank+qk_rope_head_dim`=576 /
+  `kv_lora_rank`=512) with per-tensor FP8, no separate materialization
+  (`applyMLARopeAndAssignQKVKernelOptContext`, `absorption_mode`); (2) in the
+  cross-CTA reduction, **early-exit CTAs whose query index exceeds the ragged
+  `seqLenQ`** and clamp `seqLenKv = min(seqLenKv, sparse_mla_topk)` so the reduce
+  only spans selected tiles; use exact page granularity (`numTokensPerPage=1`)
+  for the scattered top-k gather (#8866); (3) **drop the redundant FP8 K/V scratch
+  buffers** once K/V already live in the paged cache
+  (`fp8_k_buf_size=fp8_v_buf_size=0`, `kPtr=vPtr=nullptr`) (#9529).
+- **Generalizes to:** the pattern "**make the attention kernel operate on a
+  selected subset of KV rather than the full sequence** — pass the selected
+  page/block indices and bound the kernel's KV loop and reduction by the
+  selection count." Carries to: block-sparse attention, sliding-window / streaming
+  attention, landmark / NSA-style selection, any masked attention where the mask
+  is structured enough to skip whole tiles. Two reusable corollaries travel with
+  it: **write K/V straight into the paged cache** (absorption/latent layout) to
+  avoid a materialization pass, and **early-exit CTAs beyond the ragged
+  per-sequence length** so padded work is free. Adapt by: choosing the kernel's
+  selection-index input (page vs token indices), setting the loop/reduction bound
+  to the selection budget, and matching page granularity to how scattered the
+  access is.
+- **Apply via:** **not a server knob** — an FMHA kernel/dispatch change,
+  auto-selected for DSA/sparse-MLA on SM100+ (`useSparseMLA()`, `mSparseMla`,
+  `forward_absorption` / `forward_generation_dsa`). The selection budget
+  (`sparse_mla_topk`) is a model/config property; the per-tensor FP8 KV rides on
+  `kv_cache_config.dtype` (see [FP8 MLA KV](fp8-mla-kv-cache.md)). Delegate kernel
+  work to **kernel-cuda-specialist**; profile with **perf-nsight-compute-analysis**.
+- **Expected effect:** attention work drops from O(seqLen) to O(top-k) → lower
+  MLA decode/prefill latency and KV bandwidth at long context; reclaimed FMHA
+  workspace from the dropped scratch buffers. Direction only — measured Δ (attn
+  kernel time, KV bytes/step, workspace) to be recorded from the run.
+- **Accuracy risk:** **mixed.** The sparse *dispatch* itself is lossless relative
+  to the model's sparse-attention specification (DSA is trained sparse — attending
+  to the selected top-k **is** the intended computation, not an approximation the
+  optimization introduces). The **per-tensor FP8 (later FP4) KV cache is lossy**
+  in exactly the way [FP8 MLA KV](fp8-mla-kv-cache.md) is → accuracy record +
+  rollback when the quantized-KV variant is enabled. The dropped scratch buffer
+  (#9529) is lossless (K/V read from cache instead of a duplicate).
+- **Verify:** correctness — sparse-MLA forward parity vs the reference sparse
+  computation (`test_sparse_mla_forward.py`); DSA e2e accuracy (GSM8K / GPQA
+  references) unchanged, and for FP8/FP4 KV a parity check vs BF16 KV. Perf —
+  attn kernel time and KV bytes/step before vs after; confirm the reduction spans
+  only top-k tiles and padded CTAs early-exit (ncu).
+- **Rollback:** fall back to the dense MLA FMHA path (and restore FP8 K/V scratch
+  if the buffer removal is implicated). Trigger: sparse-kernel correctness
+  mismatch, FP8/FP4-KV accuracy drop beyond the recorded threshold, or the
+  sparse-MLA kernel unavailable for the SM/head config.
+- **Prior art:** PRs #8692, #8866, #9529. Files:
+  `cpp/tensorrt_llm/kernels/mlaKernels.cu`
+  (`applyMLARopeAndAssignQKVKernelOptContext`, `absorption_mode`),
+  `cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp` (`useSparseMLA()`, `mSparseMla`),
+  `cpp/tensorrt_llm/kernels/.../fmhaReduction.cu`,
+  `cpp/tensorrt_llm/common/attentionOp.cpp` (`mFP8ContextMLA`, K/V buffer sizing),
+  `tensorrt_llm/_torch/modules/attention.py` (`forward_absorption`,
+  `forward_generation_dsa`), `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`
+  (`sparse_mla_topk`). Owning specialist: **kernel-cuda-specialist**. Related: the
+  [top-k selection kernel](specialize-topk-selection-kernel.md) (the producer of
+  the indices), [split MLA reduction](split-mla-reduction-kernel.md) (the same
+  FMHA reduction, a different lever), and [FP8 MLA KV](fp8-mla-kv-cache.md) (the
+  quantized-KV precursor).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/specialize-topk-selection-kernel.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/specialize-topk-selection-kernel.md
new file mode 100644
index 000000000000..af8f2ff33661
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/specialize-topk-selection-kernel.md
@@ -0,0 +1,103 @@
+---
+id: case-specialize-topk-selection-kernel
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [launch, compute]
+signals: [many-small-kernels, recompilation-churn]
+architectures: [any-sm]
+model_scope: [sparse-attention, deepseek-v32, model-agnostic]
+phase: [decode]
+patterns: [pattern-specialize-hotpath-topk]
+accuracy_risk: lossless
+apply_via_kind: [kernel-change, config-knob]
+knobs: [use_cute_dsl_topk, enable_heuristic_topk]
+specialists: [kernel-cuda-specialist, kernel-cute-writing, perf-nsight-compute-analysis]
+commits: ['2e7769d1e8', '941a54c66a', 'e940e58eb9', 'b206f682f5', '29fac6b673', 'b41876f144']
+interactions:
+  - {case: case-sparse-mla-topk-attention, relation: composes-with, note: the attention kernel is the consumer of the indices this kernel produces}
+measured: []
+---
+
+# Specialize the hot-path top-k / selection kernel (don't ship a generic sort)
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `2e7769d1e8` [feat] Add customized topk and unit tests for DSA
+  (#8882); related: `941a54c66a` Update the indexer topK (#9255), `e940e58eb9`
+  [perf] Enable CuteDSL indexer_top_k in model (#12236), `b206f682f5` Indexer
+  TopK: single-block / multi-pass radix (#14268), `29fac6b673` Temporally-
+  Correlated Heuristic-guided Indexer TopK (#12385), `b41876f144` DSv4 prep:
+  IndexerTopK primitives (#15381).
+- **Applies when:** a **top-k / selection** runs on the hot path — DSA's indexer
+  picks the top-k KV per query every step; also MoE grouped-top-k routing, sampling
+  top-k/top-p, beam search. Signals: a `torch.topk(...)` (a full sort) plus a chain
+  of `arange` / `masked_fill` / cast / index-arithmetic ops in the profile;
+  Triton/PyTorch selection recompiling or launching many small kernels; the op
+  runs at every decode step so its launch + sort cost is disproportionate. You
+  only need the top-k **indices** (and then post-process them), not a full sort.
+- **Mechanism:** replace the general sort + its downstream index/mask arithmetic
+  with **one custom radix-select kernel** that emits exactly the top-k indices
+  (and does the causal mask / local-index conversion inline) — a histogram/radix
+  select over the float bit-pattern (e.g. 11/11/10-bit passes) avoids the
+  O(N log N) sort and the mask-tensor round-trips (#8882). Then **keep several
+  exact implementations and dispatch per regime**: single-CTA vs single-pass
+  multi-CTA (CuteDSL, chosen by whether the launch fits **one SM wave**
+  `num_rows·ctas_per_group ≤ num_sms`, folding dtype-dependent sync cost — fp32 =
+  4 radix rounds vs fp16/bf16 = 2 — into the crossover, #12236); single-block vs
+  multi-pass radix with per-row DRAM ping-pong scratch for large N (#14268). Two
+  **short-circuits** cut passes without changing the result: a **coarse
+  low-precision first pass** (fp16 histogram) that resolves the common case before
+  the exact fp32 passes and falls back when it can't disambiguate (#9255), and a
+  **warm-start from the previous decode step's selection** — seed the k-th-value
+  threshold search from the prior step's top-k, since consecutive autoregressive
+  steps select nearly the same tokens (#12385).
+- **Generalizes to:** the pattern "**the hot-path top-k/selection is a specialized
+  kernel, not a generic sort — fuse the selection with its index/mask
+  post-processing into one radix-select, keep multiple exact variants, and dispatch
+  by `(k, N, dtype, SM-wave-occupancy)`.**" Two reusable short-circuit sub-patterns
+  travel with it: **cheap-check-first** (a coarse low-precision pass that resolves
+  the common case, exact fallback otherwise) and **warm-start from the previous
+  iteration** when the answer drifts slowly across steps (autoregressive decode —
+  also applies to per-step threshold/quantile/routing searches). Carries to: MoE
+  grouped-top-k routing (cf. the [routing-kernel case](fuse-moe-routing-kernel.md)),
+  sampling top-k/top-p, beam-search top-k, argmax/threshold selection. Adapt by:
+  writing the radix-select for your `k`; choosing the parallelization tier by
+  N/occupancy; and pre-allocating any scratch so the op is CUDA-graph-capturable
+  (#14268 exposes `scratch` for graph-stable addresses).
+- **Apply via:** **not a server knob** for the core kernel (auto-dispatched);
+  opt-in variants are config-gated — `use_cute_dsl_topk`, `enable_heuristic_topk`
+  (`DeepSeekSparseAttentionConfig` in `llm_args.py`). Delegate the CUDA-C++
+  radix-select to **kernel-cuda-specialist**, the CuteDSL variant to
+  **kernel-cute-writing**; profile with **perf-nsight-compute-analysis**.
+- **Expected effect:** lower per-step selection time (no full sort, fewer launches,
+  fewer passes) → smaller GPU-idle and higher decode throughput where selection is
+  hot. PR #12236 reports ~1.14× over always-single-CTA for the CuteDSL path; treat
+  as the author's figure and **record the local Δ** (selection kernel time,
+  distinct-compile count, throughput) from the run — do not present a PR number as
+  a locally-achieved result.
+- **Accuracy risk:** **lossless** — every variant returns the exact top-k (the
+  coarse pass only brackets and falls back to exact; radix-select is exact). One
+  flag: the temporal warm-start (#12385) has a bounded refinement loop that
+  *targets* exactness (unit test asserts parity with `torch.topk` even with poor
+  seeds) but could deviate under pathological non-convergence — **verify accuracy
+  on first enable** of `enable_heuristic_topk`.
+- **Verify:** correctness — kernel output equals `torch.topk` on the same logits,
+  including edge cases (ties, all-equal, N<k); DSA e2e accuracy unchanged; re-check
+  when enabling the heuristic/CuteDSL variants. Perf — selection kernel time and
+  launch/compile count before vs after; confirm the chosen tier matches the
+  `(k,N,dtype)` regime (ncu / dispatch logs).
+- **Rollback:** fall back to the reference `torch.topk` path (kept as reference/
+  fallback), or disable the opt-in variant (`use_cute_dsl_topk=False`,
+  `enable_heuristic_topk=False`). Trigger: any index mismatch vs reference or a
+  perf regression in an untuned `(k,N)` regime.
+- **Prior art:** PRs #8882, #9255, #12236, #14268, #12385, #15381. Files:
+  `cpp/tensorrt_llm/kernels/indexerTopK.cu`, `heuristic_topk.cuh`, `IndexerTopK.h`,
+  `cpp/tensorrt_llm/thop/IndexerTopKOp.cpp`,
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`,
+  `.../custom_ops/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk.py`.
+  Owning specialists: **kernel-cuda-specialist**, **kernel-cute-writing**. Related:
+  the [MoE routing-kernel case](fuse-moe-routing-kernel.md) (same "collapse a
+  selection op-chain into one per-token kernel" idea) and the
+  [sparse-MLA top-k attention case](sparse-mla-topk-attention.md) (the consumer of
+  these indices).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/split-mla-reduction-kernel.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/split-mla-reduction-kernel.md
new file mode 100644
index 000000000000..6457117a09ac
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/split-mla-reduction-kernel.md
@@ -0,0 +1,35 @@
+---
+id: case-split-mla-reduction-kernel
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [compute]
+signals: [attention-hot-path, kernel-far-from-sol, small-batch-decode]
+architectures: [any-sm]
+model_scope: [mla, deepseek-v3r1, model-agnostic]
+phase: [decode]
+patterns: [pattern-split-compute-from-reduction]
+accuracy_risk: lossless
+apply_via_kind: [kernel-change]
+knobs: []
+specialists: [kernel-cuda-specialist, perf-nsight-compute-analysis]
+commits: ['da6cb541a286']
+eligibility:
+  - "TRTLLM-Gen FMHA generation with multi-CTA-along-KV enabled (isMultiCtasKvEnabled); 1-/2-CTA keepsMmaAbForGeneration MLA kernels"
+measured: []
+---
+
+# Split MLA generation attention into a separate GMEM reduction kernel
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `da6cb541a286` [feat] Optimize MLA kernels with separate reduction kernels (#7597).
+- **Applies when:** compute/occupancy-bound MLA decode using TRTLLM-Gen FMHA generation kernels with multi-CTA-along-KV splitting (`isMultiCtasKvEnabled`), specifically 1-/2-CTA `keepsMmaAbForGeneration` MLA kernels with large reduction tiles. Symptom: the monolithic kernel does cross-CTA softmax/output reduction inline, inflating register/smem pressure and capping occupancy, reduction latency exposed at low/medium batch.
+- **Mechanism:** moves the cross-CTA-KV combine (rescale partials by `exp2f(softmaxScaleLog2·(localMax−maxVal))`, accumulate, finalize) into a dedicated `fmhaReductionKernel` (new `MultiCtasKvMode::GmemReductionWithSeparateKernel` replacing in-kernel `GmemReduction`), 512-thread block that splits the reduction across CtasKv to cut latency; the specialized main kernel sheds reduction register/smem footprint → higher occupancy. `runFmhaReduction(...)` dispatched after the main kernel on the same stream.
+- **Generalizes to:** "split a monolithic kernel doing compute + a serial cross-partition reduction into two specialized kernels (main + reduction) to raise occupancy and parallelize the reduce"; carries to flash-attn split-KV/split-K combine, GEMM split-K epilogue, any persistent kernel whose tail reduction throttles occupancy; adapt by choosing where to materialize partials (GMEM via `ptrSoftmaxStats`/partial-O) and gating on the regime where the extra launch pays off.
+- **Apply via:** internal kernel/dispatch change (auto-selected when multi-CTA-KV enabled for qualifying MLA kernels) — not a server knob. Delegate to **kernel-cuda-specialist**; profile occupancy/SOL with **perf-nsight-compute-analysis**.
+- **Expected effect:** higher generation-kernel occupancy, lower MLA decode latency in the multi-CTA-KV regime; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless (numerically equivalent online-softmax reduction, just relocated). Accuracy tests updated in `test_llm_api_pytorch.py` as guard.
+- **Verify:** MLA decode latency/throughput + ncu occupancy of the generation kernel before/after; numeric parity vs in-kernel-reduction path.
+- **Rollback:** fall back to `MultiCtasKvMode::GmemReduction`. Trigger: separate-kernel launch adds net latency at small batch, or correctness mismatch.
+- **Prior art:** PR #7597. Files: `cpp/.../trtllmGenKernels/fmha/fmhaReduction.cu` (`runFmhaReduction`), `fmhaKernels.h`, `fmhaRunnerParams.h` (`MultiCtasKvMode`). Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/triton-to-cpp-op.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/triton-to-cpp-op.md
new file mode 100644
index 000000000000..4ae09efbd942
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/triton-to-cpp-op.md
@@ -0,0 +1,81 @@
+---
+id: case-triton-to-cpp-op
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [host-overhead, launch]
+signals: [many-small-kernels, recompilation-churn, host-prep-on-critical-path]
+architectures: [any-sm]
+model_scope: [sparse-attention, deepseek-v32, model-agnostic]
+phase: [decode]
+patterns: [pattern-triton-to-precompiled-cpp-op]
+accuracy_risk: lossless
+apply_via_kind: [kernel-change]
+knobs: []
+specialists: [kernel-cuda-specialist, perf-host-analysis]
+commits: []
+interactions:
+  - {case: case-relax-tl-constexpr, relation: alternative-to, note: the cheaper Triton-side fix when you keep Triton}
+measured: []
+---
+
+# Replace a Triton hot-path kernel with a precompiled C++ custom op to cut host dispatch overhead
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Applies when:** a Triton kernel sits on the **host-bound** hot path and is
+  small/cheap on the GPU, so its host-side cost dominates — Triton's per-launch
+  Python dispatch, autotune-cache lookup, argument marshalling, and JIT /
+  recompilation overhead. Symptoms: host time in the Triton launch wrapper, GPU
+  idle around the op, recompiles across layers/shapes. Classify host-bound first
+  with `perf-host-analysis` / `perf-nsight-systems` — a *compute-bound* Triton
+  kernel (near SOL) will **not** benefit from this. Instance: the DSA indexer's
+  `triton_gather_k_cache` and `triton_convert_req_index_to_global_index`.
+- **Mechanism:** Triton pays a fixed Python-side cost per launch (and recompiles
+  per new `constexpr`/shape). Re-implementing the kernel in C++/CUDA and exposing
+  it as a Torch library op (`torch.ops.trtllm.*`) makes it precompiled with a
+  thin C++ launch path, removing the JIT and most of the per-launch host
+  overhead. A `register_fake` (meta) implementation keeps it traceable under
+  `torch.compile` / fake tensors.
+- **Generalizes to:** the pattern "**when a Triton kernel is host-bound (tiny GPU
+  work, heavy launch/JIT overhead), move it to a precompiled C++/CUDA custom op
+  with a `register_fake`.**" Carries to: other small index/gather/scatter/copy
+  kernels on the decode hot path; any Triton op recompiling across layers or
+  shapes; ops where a CUDA implementation already exists or is cheap to write.
+  Adapt by: writing the `.cu` kernel + thop binding, registering the op schema +
+  `register_fake`, and swapping the call site to `torch.ops.trtllm.*`; keep the
+  Triton version as reference/fallback and port its tests. Only worth it when
+  **host** overhead, not GPU SOL, is the bottleneck.
+- **Apply via:** **not a server config knob** — a kernel/op change. Delegate to
+  **kernel-cuda-specialist** (raw CUDA C++ + thop/
+  pybind). Touch points (prior art, TRT-LLM repo): kernels
+  `cpp/tensorrt_llm/kernels/{indexerKCacheGather.cu, convertReqIndexToGlobal.cu}`
+  (+ `.h`), Torch op bindings
+  `cpp/tensorrt_llm/thop/{IndexerKCacheGatherOp.cpp, convertReqIndexToGlobalOp.cpp}`
+  (+ `CMakeLists.txt`), the `register_fake` meta impls in
+  `tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py`, and the call sites in
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`
+  (`torch.ops.trtllm.indexer_k_cache_gather_op`,
+  `torch.ops.trtllm.convert_req_index_to_global`).
+- **Expected effect:** lower host launch overhead and fewer recompiles → smaller
+  GPU-idle gaps and higher throughput when host-bound; no change to GPU compute.
+  Largest at decode / small token counts where launch overhead dominates. Record
+  the local Δ (per-iteration host time, launch overhead, throughput) — PR #12581
+  ships it without a public number; do **not** cite one not measured locally.
+- **Accuracy risk:** lossless **by intent**, but it is a kernel rewrite, so
+  **verify vs the Triton/PyTorch reference within tolerance** — a C++ kernel bug
+  changes outputs. Watch dtype/stride and edge cases (empty gather, non-contiguous
+  `k_cache`).
+- **Verify:** correctness — the new op tests
+  `tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py` compare both
+  ops to a Python reference incl. edge cases; DSA e2e still passes. Perf —
+  `perf-host-analysis`: host launch time + GPU-idle around the op, throughput
+  before vs after; confirm `register_fake` lets `torch.compile` trace the op.
+- **Rollback:** swap the call sites back to the Triton wrappers (kept in
+  `kernel.py`); trigger: any accuracy mismatch outside tolerance or >5% perf
+  regression.
+- **Prior art:** PR #12581; paths above; new op test
+  `test_cpp_custom_ops.py` (replaces the removed `test_triton_gather_k_cache.py`).
+  Owning specialist: **kernel-cuda-specialist**; detection: **perf-host-analysis**.
+  Related: the [relax-`tl.constexpr` case](relax-tl-constexpr.md) (the cheaper Triton-side fix when
+  you keep Triton).
diff --git a/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/trtllm-gen-fp4-moe-backend.md b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/trtllm-gen-fp4-moe-backend.md
new file mode 100644
index 000000000000..0f507e59d26c
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/kernel-and-fusion/trtllm-gen-fp4-moe-backend.md
@@ -0,0 +1,36 @@
+---
+id: case-trtllm-gen-fp4-moe-backend
+type: case
+family: kernel-and-fusion
+maturity: full
+bottleneck: [compute]
+signals: [gemm-dominates-step, hbm-roundtrip-between-ops]
+architectures: [sm100]
+model_scope: [moe, deepseek-v3r1, model-agnostic]
+phase: [decode]
+patterns: [pattern-hw-matched-lowprec-gemm-backend]
+accuracy_risk: lossy
+apply_via_kind: [config-knob]
+knobs: [moe_backend]
+specialists: [kernel-cute-specialist, trtllm-moe-develop]
+commits: ['31624b079a12']
+eligibility:
+  - "sm == 100 (sm_100a prebuilt cubins)"
+  - "quant is NVFP4 or FP8_BLOCK_SCALES; the cubin gemmList must cover the (dtype, tileN, transposeMmaOutput) combination"
+measured: []
+---
+
+# Select the trtllm-gen FP4 grouped-GEMM MoE backend for Blackwell DeepSeek
+
+> Part of the [Kernel & Fusion casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `31624b079a12` feat: [Deepseek] Add trtllm-gen MOE FP4 MOE backend (#3387).
+- **Applies when:** compute-bound MoE grouped-GEMM on sm_100a; quant is `NVFP4`/`FP8_BLOCK_SCALES`; DeepSeek-V3/R1 (or other block-scale MoE); currently on default `CUTLASS` MoE backend and grouped-GEMM dominates the decode step.
+- **Mechanism:** swaps the per-expert grouped GEMM to trtllm-gen prebuilt SM100a cubins (`blockScaleMoe`) that natively consume FP4/FP8 block scales, with routing fused into the moe_runner — avoiding the CUTLASS path's extra type conversions and a redundant FP4 temp output. Deliberately runs routing reduction in fp32 with tanh, and disables PDL for fc1 to stabilize numerics.
+- **Generalizes to:** "pick a hardware-specialized low-precision grouped-GEMM backend over the generic one when the GEMM is the bottleneck and the dtype is natively supported"; carries to FP8 block-scale MoE, MXFP4, other Blackwell MoE models (KimiK2, Qwen-next); adapt by checking the cubin/gemmList covers your (dtype, tileN, transposeMmaOutput) and quant exclude_modules match (trtllm-gen excludes `*kv_b_proj*`, `*k_b_proj*`, `*eh_proj` for FP8 block scales).
+- **Apply via:** `moe_backend='TRTLLM'` (on `ModelConfig`/`PyTorchConfig`, default `'CUTLASS'`; `--moe_backend {CUTLASS,TRTLLM}`). Gate `FusedMoE.is_trtllm()`; dispatches `forward_trtllmgen` → `torch.ops.trtllm.fp8_block_scale_moe_runner`/`fp4_block_scale_moe_runner`. Delegate to **kernel-cute-specialist** only if a cubin variant is missing.
+- **Expected effect:** higher MoE grouped-GEMM throughput / lower decode latency on Blackwell; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossy (FP4/FP8 block-scale GEMM). Mitigations: fp32 routing reduction + tanh, PDL off for fc1. NVFP4 weights use a distinct loader (`is_trtllm_nvfp4`) so weight layout must match.
+- **Verify:** MMLU / task accuracy vs CUTLASS baseline (commit adds `mmlu_llmapi.py` hooks, extends `test_moe.py`); grouped-GEMM kernel time drops in nsys.
+- **Rollback:** `moe_backend='CUTLASS'`. Trigger: accuracy regression beyond tolerance, or a dtype/shape lacking a trtllm-gen cubin.
+- **Prior art:** PR #3387. Files: `_torch/modules/fused_moe.py` (`is_trtllm`/`forward_trtllmgen`), `_torch/model_config.py` (`moe_backend`), `cpp/.../trtllmGenKernels/blockScaleMoe/{runner.cu,gemmList.h,RoutingKernel.cu}`, `thop/fp4BlockScaleMoe.cpp`. Owning skill: **trtllm-moe-develop**; cubins: **kernel-cute-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md
new file mode 100644
index 000000000000..ec4f11af7021
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md
@@ -0,0 +1,33 @@
+---
+id: case-attention-dp-padding
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [host-overhead, compute, memory]
+signals: [expert-load-imbalance, host-prep-on-critical-path]
+architectures: [sm90, sm100]
+model_scope: [model-agnostic, moe, deepseek-v3r1, qwen3-moe, llama-family]
+phase: [any-phase]
+patterns: [pattern-process-only-real-work]
+accuracy_risk: lossless
+apply_via_kind: [default-on]
+knobs: []
+specialists: [perf-host-optimization, perf-host-analysis]
+commits: ['07e8813984cd', 'b618e1f55b88', '15823614000b']
+measured: []
+---
+
+# Eliminate padding / wasted compute in data-parallel attention (attention-DP)
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `07e8813984cd` feat: Remove padding in attention DP (#6064); related: `b618e1f55b88` Eliminate the need for attention DP padding when possible (#3439, variable-length allgather/reducescatter), `15823614000b` only pad one dummy request for attention dp (#4664).
+- **Applies when:** attention-DP on and per-rank token counts are imbalanced/ragged (mixed ISL, uneven request assignment) so the framework padded every rank to the max — wasting attention/MoE compute and host bookkeeping on dummy tokens; signals: low MFU with attention-DP, throughput scaling poorly with DP degree, dummy work in profiles.
+- **Mechanism:** (a) #3439 replaces fixed all-padded collectives with variable-length allgather/reducescatter so ranks exchange only real tokens. (b) #4664 collapses per-step "pad to equal length" to a single dummy request (`_pad_attention_dp_dummy_request`) and drops the cleanup across executor loops, cutting per-step host overhead. (c) #6064 removes residual padding/swizzle bookkeeping in the MoE path so the un-padded path is plumbed end to end.
+- **Generalizes to:** "process only real tokens, not padded-to-max" for any DP/ragged-batch collective; carries to MoE all-to-all dispatch, other DP'd modules, any allgather/reducescatter on variable per-rank counts; adapt by making the collective length-aware and removing downstream padded-shape assumptions.
+- **Apply via:** keep attention-DP on with the variable-length collective path (default after these PRs); no extra knob beyond a recent build. Delegate residual host cleanup to **perf-host-optimization**.
+- **Expected effect:** higher throughput and MFU under ragged DP load by removing dummy-token compute + per-step host padding; direction only — measured Δ (tok/s, MFU, host prep) to be recorded from run.
+- **Accuracy risk:** lossless — removes computation on padding/dummy tokens that never contributed to outputs.
+- **Verify:** throughput/MFU up with imbalanced ISL; profile shows reduced dummy/padding work + lower host prep; output parity vs padded path.
+- **Rollback:** revert to padded collective / full per-step padding. Trigger: correctness mismatch on ragged batches or collective hang.
+- **Prior art:** PRs #6064, #3439, #4664. Files: `cpp/.../thop/{allgatherOp,reducescatterOp}.cpp`, `_torch/distributed/ops.py`, `_torch/pyexecutor/py_executor.py`, `_torch/modules/fused_moe/{fused_moe_cutlass,fused_moe_wide_ep}.py`. Detection: **perf-host-analysis**.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/auxiliary-cache-in-kv-manager.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/auxiliary-cache-in-kv-manager.md
new file mode 100644
index 000000000000..f52ba79079f1
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/auxiliary-cache-in-kv-manager.md
@@ -0,0 +1,94 @@
+---
+id: case-auxiliary-cache-in-kv-manager
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [memory]
+signals: [memory-capacity-bound, long-context]
+architectures: [any-sm]
+model_scope: [model-agnostic, sparse-attention, mla, deepseek-v32]
+phase: [any-phase]
+patterns: [pattern-fold-side-cache-into-kv-manager]
+accuracy_risk: lossless
+apply_via_kind: [code-change, config-knob]
+knobs: [kv_cache_config.enable_block_reuse, host_cache_size]
+specialists: [kernel-cuda-specialist, trtllm-serve-config-guide]
+commits: ['ae6875fe10', '356a52edf5', '402a056ac6', '6254f3a161', 'c37924f37b']
+interactions:
+  - {case: case-mla-kv-cache-reuse, relation: composes-with, note: "#9383 extends that reuse pattern to the side cache"}
+  - {case: case-move-bookkeeping-into-cpp-op, relation: composes-with, note: the scatter op there writes into this pool}
+measured: []
+---
+
+# Fold a model-specific auxiliary cache into the unified paged KV-cache manager
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `ae6875fe10` [feat] Move indexer-k-cache to KVCacheManager (#8699);
+  related: `356a52edf5` KV-cache reuse for DSv32 (#9383), `402a056ac6` host cache
+  offload for DSA (#12010), `6254f3a161` count indexer K-cache as UINT8 in the size
+  estimate (#15088), `c37924f37b` clear indexer-k-cache refs before release (#9110).
+- **Applies when:** a model needs an **auxiliary/side cache** alongside the main KV
+  cache — DSA keeps a separate small FP8/UINT8 **indexer K cache** for the lightning
+  indexer; also block-scale/scale-factor pools, sparse-attention landmark caches, any
+  per-layer compressed side state. Signals: a bespoke second allocator
+  (`indexer_k_cache_manager = BlockManager(...)`) with hand-rolled
+  `add_dummy_requests` / `prepare_resources` / `update_resources` / `free_resources`
+  / `rewind` overrides; the side cache **can't** be reused across prefixes or
+  offloaded to host because it isn't paged; a "reuse not supported" guard; its memory
+  isn't counted correctly in the auto-sizing estimate.
+- **Mechanism:** make the auxiliary cache a **first-class typed pool inside the
+  unified paged `KVCacheManager`** (a `containsIndexerKCache` pool in
+  `WindowBlockManager`, `createIndexerKCachePools()`, allocated `kUINT8`) that
+  **shares block IDs** with the main KV pool. The parallel shadow allocator and its
+  overrides are deleted; the manager pages, reuses, offloads, and accounts for the
+  side pool alongside the main blocks. That single relocation then unlocks, for free:
+  **prefix/block reuse** (drop the "unsupported" guard, teach `BlockRange` to carry
+  the side pool via `isEnableIndexerKCache()`, #9383) and **host/secondary offload**
+  (#12010). Three corollaries the relocation exposes: (1) a **mixed-dtype** pool set
+  (UINT8 side cache beside a BF16/FP8 KV) must be summed **additively per physical
+  pool dtype**, never via one multiplicative dtype factor (#15088); (2) once blocks
+  are onboarded from a secondary/host pool, **index them by the decoded memory-pool
+  index**, not the logical block ID (they diverge — `_get_pool_block_indices()` from
+  `host_kv_cache_block_offsets`, #12010); (3) **null the Python tensor views before
+  the C++ owner frees** the device memory at teardown (#9110).
+- **Generalizes to:** the pattern "**manage an auxiliary/compressed side cache as a
+  typed pool inside the unified paged KV-cache manager instead of a bolt-on shadow
+  allocator — it inherits paging, block reuse, host/secondary offload, eviction, and
+  correct per-pool memory accounting, and deletes duplicated resource-tracking
+  code.**" Carries to: scale-factor / block-scale pools, sparse-attention or
+  landmark caches, adapter/state caches, any model-specific representation that lives
+  as long as the KV blocks. Adapt by: registering the pool type + `createXPools()`,
+  sharing block IDs with the main pool, wiring size accounting to the pool's
+  *physical* dtype, and carrying the pool through the block-range/transfer machinery
+  so reuse/offload include it.
+- **Apply via:** **not a single server knob** — a KV-cache-manager code change
+  (the pool becomes first-class). The capabilities it unlocks ride existing knobs:
+  `kv_cache_config.enable_block_reuse` (reuse), the host-offload config
+  (`host_cache_size` / free-fraction). Delegate to **kernel-cuda-specialist** /
+  the KVCacheManager owner; expose YAML knobs via **trtllm-serve-config-guide**.
+- **Expected effect:** the side cache gains prefix reuse (higher prefill throughput /
+  lower TTFT on shared prefixes) and host offload (more/longer sequences fit), plus
+  correct auto-sizing (no OOM / no wasted capacity) and less duplicated code.
+  Direction only — measured Δ (KV-hit rate, TTFT, max concurrency/seq-len, memory
+  headroom) to be recorded from a prefix-heavy / long-context trace.
+- **Accuracy risk:** **lossless** — pure storage/bookkeeping relocation. The one
+  correctness-sensitive path is **reuse/offload lockstep**: the side pool must be
+  reused/offloaded in lockstep with the MLA KV blocks (shared block IDs) and indexed
+  by the correct pool index; a mismatch would silently feed wrong indexer state →
+  verify with an accuracy eval on a reuse/offload trace.
+- **Verify:** correctness — DSA e2e accuracy unchanged with `enable_block_reuse=True`
+  and with host offload on (vs off); size-estimate matches measured pool bytes. Perf
+  — KV-cache-hit rate + TTFT on a repeated-prefix trace, max concurrency/seq-len with
+  offload, memory headroom.
+- **Rollback:** `kv_cache_config.enable_block_reuse=False` / disable host offload to
+  isolate; the pool-relocation itself is structural. Trigger: accuracy drift on a
+  reuse/offload trace, or a size-estimate error causing OOM.
+- **Prior art:** PRs #8699, #9383, #12010, #15088, #9110. Files:
+  `cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp` (`createIndexerKCachePools`,
+  `containsIndexerKCache`, `isEnableIndexerKCache`), `kvCacheUtils.h` (`BlockRange`),
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` (`DSACacheManager`,
+  `_get_pool_block_indices`). Owning specialist: **kernel-cuda-specialist**. Related:
+  the [MLA KV-cache reuse case](mla-kv-cache-reuse.md) (#9383 is the side-cache
+  instance of that pattern) and [move bookkeeping into a C++ op](move-bookkeeping-into-cpp-op.md)
+  (the scatter op that writes into this pool).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/cache-step-invariant-per-layer.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/cache-step-invariant-per-layer.md
new file mode 100644
index 000000000000..afc330a25689
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/cache-step-invariant-per-layer.md
@@ -0,0 +1,85 @@
+---
+id: case-cache-step-invariant-per-layer
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [host-overhead]
+signals: [gpu-idle-between-steps, host-prep-on-critical-path]
+architectures: [any-sm]
+model_scope: [model-agnostic, sparse-attention, mla]
+phase: [any-phase]
+patterns: [pattern-hoist-step-invariant-host-work]
+accuracy_risk: lossless
+apply_via_kind: [code-change]
+knobs: []
+specialists: [perf-host-optimization, perf-host-analysis]
+commits: []
+interactions:
+  - {case: case-hoist-torch-compile-closures, relation: composes-with, note: "same PR (#12581), same host-overhead theme"}
+measured: []
+---
+
+# Cache step-invariant values across the per-layer attention loop
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Applies when:** host-overhead-bound serving — the PyExecutor loop is
+  CPU-bound, the GPU shows idle gaps between steps, and a function called **once
+  per decoder layer** recomputes values that are actually invariant across all
+  layers within a forward step. The signal it targets: per-layer attention-prep
+  code that re-derives KV-pool views/reshapes, stride factors, and block-table /
+  request-index slices that depend only on the pool shape and batch dimensions,
+  not the layer index. Confirm host-bound first with `perf-host-analysis`
+  (GPU-idle ratio, per-iteration host breakdown) — this only pays off when the
+  loop is host-bound. Most acute with many layers × non-trivial per-layer Python/
+  CUDA prep. Instance: the DSA indexer's
+  `transform_local_topk_and_prepare_pool_view()`.
+- **Mechanism:** the per-layer call recomputes step-invariant derived tensors
+  (pool `squeeze/view` reshape, `stride_factor = num_layers * tokens_per_block`,
+  context/gen block-table and request-index slices) on *every* layer, paying
+  redundant Python + CUDA overhead N times. These depend only on the KV-pool
+  shape and batch dims, so compute them once per step, cache them on the metadata
+  object (`_cached_*`), and reuse across layers; invalidate once per step.
+- **Generalizes to:** the pattern "**loop-invariant code motion on the host hot
+  path** — hoist any per-layer/per-iteration work whose inputs are constant
+  within the step out of the loop and cache it once." Carries to: other per-layer
+  metadata prep (RoPE/positional tables, slot-mapping slices, mask precompute),
+  per-step-constant reshapes/views, anything keyed off pool shape or batch dims;
+  and the same trick in any framework loop over homogeneous layers. Adapt by:
+  identifying which derived values are *truly* step-invariant (not
+  layer-dependent), caching them on the per-step metadata object, invalidating
+  once per step **before** the first per-layer use, and guarding validity with a
+  cheap key (here a boolean flag + `id(kv_cache_manager)`). Keep cache fields as
+  plain instance attributes (init in `__init__`) so they stay invisible to
+  dataclass/`torch.compile` introspection.
+- **Apply via:** **not a server config knob** — a host-side code change in the
+  attention-metadata class. Delegate optimization to **perf-host-optimization**
+  (line_profiler / nsys host rounds) and detection to **perf-host-analysis**.
+  Touch points (prior art, TRT-LLM repo):
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` —
+  `transform_local_topk_and_prepare_pool_view()` (now reads `_cached_*`),
+  `_ensure_pool_view_cached()` / `_invalidate_pool_view_cache()`, and the
+  invalidation calls in `prepare()` and `on_update_kv_lens()`.
+- **Expected effect:** lower per-layer host time in attention prep → smaller
+  GPU-idle gaps between steps and higher throughput when host-bound; no change
+  when GPU-bound. PR #12581 reports ~50µs×2 per layer (~4–5ms per forward step)
+  of host overhead removed for DSA — treat that as the author's estimate and
+  **record the local Δ** (per-iteration host time, GPU-idle ratio, throughput)
+  from a host-profiling run; do not present the PR figure as a locally-achieved
+  result.
+- **Accuracy risk:** lossless — identical values, computed once instead of N
+  times. Correctness hinges entirely on the cache being invalidated every step
+  before the first per-layer read; a stale cache would silently feed wrong
+  slices.
+- **Verify:** correctness — DSA e2e / sparse-attention tests still pass; for a
+  multi-layer step assert the cached pool view / slices equal the per-layer
+  recompute, and re-check after a batch-shape or `kv_cache_manager` change. Perf
+  — `perf-host-analysis`: per-iteration host time in attention prep, GPU-idle
+  ratio, and throughput before vs after.
+- **Rollback:** revert to per-layer recompute; trigger: any correctness
+  regression or >5% host-time / throughput regression. Likely failure mode to
+  watch: a step path that mutates batch dims or swaps `kv_cache_manager` without
+  calling the invalidation → stale cache.
+- **Prior art:** PR #12581 ("Multiple host perf
+  optimizations for DSA part"); `dsa.py` paths above. Owning skill:
+  **perf-host-optimization**; detection: **perf-host-analysis**. Related: the [torch.compile-closure-hoist case](hoist-torch-compile-closures.md) (same PR, same host-overhead theme).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/chunked-prefill-aligned-auxiliary.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/chunked-prefill-aligned-auxiliary.md
new file mode 100644
index 000000000000..7c9759a48c73
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/chunked-prefill-aligned-auxiliary.md
@@ -0,0 +1,86 @@
+---
+id: case-chunked-prefill-aligned-auxiliary
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [memory]
+signals: [memory-capacity-bound, long-context, high-concurrency]
+architectures: [any-sm]
+model_scope: [model-agnostic, mla, sparse-attention, deepseek-v32]
+phase: [prefill]
+patterns: [pattern-chunk-prefill-align-auxiliary]
+accuracy_risk: lossless
+apply_via_kind: [config-knob, code-change]
+knobs: [enable_chunked_prefill]
+specialists: [trtllm-serve-config-guide, kernel-cuda-specialist]
+commits: ['b10137fdd5', '78bb245554']
+interactions:
+  - {case: case-free-mla-intermediates, relation: alternative-to, note: the other prefill-memory lever}
+  - {case: case-auxiliary-cache-in-kv-manager, relation: depends-on, note: the chunk gather reads from that cache}
+measured: []
+---
+
+# Chunked prefill to bound memory — and align an auxiliary structure's chunking to it
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `b10137fdd5` [feat] Support MLA chunked prefill for DeepSeek V3.2
+  (#9376); related: `78bb245554` [fix] Better align MLA chunking with indexer
+  chunking for DSV32 (#10552).
+- **Applies when:** long-context **prefill** where materializing the full sequence at
+  once blows up activation/logit memory — MLA context attention over a long ISL, and
+  especially a **quadratic** intermediate (the DSA indexer's full-sequence MQA logits
+  are O(L²)). Signals: prefill OOM or capped ISL at high concurrency; a large
+  full-sequence intermediate on the prefill path; you want to overlap prefill with
+  decode. Most acute when an **auxiliary structure** (sparse indexer, landmark
+  selector) also materializes something per-sequence.
+- **Mechanism:** split a long prefill into fixed-size **KV chunks** so peak
+  activation/logit memory is bounded by the chunk size, not the sequence length. MLA
+  context attention runs chunk-by-chunk against cached KV
+  (`enable_context_mla_with_cached_kv`); the RoPE kernel recomputes a **cache-aware
+  position offset per chunk** (`cached_offset = cache_seq_len − current_seq_len`) so
+  positions are correct when earlier chunks are already cached
+  (`applyMLARopeAndAssignQKVKernelOptContext`). The reusable subtlety: the **auxiliary
+  indexer must gather only the *current chunk's* K** from cache
+  (`_gather_k_cache_for_chunk`, `slot_mapping_*_fullkv`) instead of the full sequence,
+  and it must **inherit MLA's chunk boundaries** rather than running its own second
+  chunking scheme — #10552 makes the indexer treat each MLA chunk as one indexer chunk
+  (dropping a redundant `host_cached_tokens.sum().item()` host sync and a
+  boundary-mismatch class).
+- **Generalizes to:** the pattern "**chunk long prefill into fixed windows to bound
+  the activation/logit footprint to chunk size; when an auxiliary structure exists, it
+  must (a) gather only the current chunk's data, (b) recompute cache-aware position
+  offsets per chunk, and (c) inherit the primary chunker's boundaries — never run two
+  chunking schemes over the same data.**" Carries to: chunked prefill for standard
+  MHA/GQA, chunked context with sparse/landmark/retrieval selection, any per-chunk
+  auxiliary state; and the general rule "when two subsystems chunk the same sequence,
+  the secondary defers to the primary's boundaries." Adapt by: choosing the chunk size
+  (memory vs launch trade-off), making every auxiliary gather chunk-local, and aligning
+  boundaries.
+- **Apply via:** config `enable_chunked_prefill` (+ the chunk-size knob); the
+  auxiliary-alignment is a code change on the auxiliary path. Delegate the YAML knob to
+  **trtllm-serve-config-guide**; the chunk-local auxiliary gather + boundary alignment
+  to **kernel-cuda-specialist** / the model author.
+- **Expected effect:** bounded prefill memory → longer ISL / higher concurrency fits,
+  and prefill can overlap decode; the quadratic indexer intermediate becomes per-chunk.
+  Direction only — measured Δ (peak prefill memory, max ISL/concurrency, TTFT) to be
+  recorded from a long-context run.
+- **Accuracy risk:** **lossless** — chunking is a tiling/scheduling of an exact
+  computation; the only risk is a boundary/offset bug (misaligned chunks or a wrong
+  cache offset), which is a *correctness* failure, not precision loss. Alignment
+  (#10552) exists precisely to keep the two chunkers consistent.
+- **Verify:** correctness — DSA/MLA prefill output parity vs non-chunked (GSM8K or a
+  logit check), especially at chunk boundaries and when part of the sequence is cached;
+  confirm the indexer reads only the current chunk. Perf — peak prefill memory and max
+  ISL/concurrency before vs after; TTFT.
+- **Rollback:** `enable_chunked_prefill=False`. Trigger: accuracy mismatch at chunk
+  boundaries, or no memory benefit at the target ISL.
+- **Prior art:** PRs #9376, #10552. Files:
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` (`_gather_k_cache_for_chunk`,
+  `enable_context_mla_with_cached_kv`, `split_prefill_chunks`, chunk-spec alignment),
+  `cpp/tensorrt_llm/kernels/mlaKernels.cu` (`applyMLARopeAndAssignQKVKernelOptContext`,
+  cache offset). Owning skills: **trtllm-serve-config-guide** (knob),
+  **kernel-cuda-specialist** (auxiliary gather). Related:
+  [free MLA intermediates](free-mla-intermediates.md) (the other prefill-memory lever)
+  and the [auxiliary-cache-in-manager case](auxiliary-cache-in-kv-manager.md) (the
+  cache the chunk gather reads from).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/cuda-graph-padding.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/cuda-graph-padding.md
new file mode 100644
index 000000000000..00b2f6eeb3ce
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/cuda-graph-padding.md
@@ -0,0 +1,33 @@
+---
+id: case-cuda-graph-padding
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch]
+signals: [slow-path-fallback, expert-load-imbalance]
+architectures: [any-sm]
+model_scope: [model-agnostic, spec-decode]
+phase: [decode]
+patterns: [pattern-pad-to-captured-bucket]
+accuracy_risk: lossless
+apply_via_kind: [config-knob, default-on]
+knobs: []
+specialists: [perf-torch-cuda-graphs, perf-torch-cuda-graph-specialist]
+commits: ['097636020451', '6b583f6f83cb']
+measured: []
+---
+
+# CUDA-graph padding for variable batch / speculative-decode token counts
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `097636020451` add support for MTP+cuda_graph_padding (#3096); related: `6b583f6f83cb` Enable CUDA graphs when attention DP is used and active requests are uneven (#3010).
+- **Applies when:** launch-bound generation where CUDA graphs need a fixed shape but real per-step token/batch count varies, so graphs miss. (a) MTP makes each gen request carry `1 + max_draft_tokens` tokens — padded dummies must match that width; (b) attention-DP across GPUs with uneven active-request counts → per-rank shape differs, graphs can't replay.
+- **Mechanism:** make padded/dummy work match the captured shape. (a) MTP: dummy requests created with `max_num_draft_tokens = spec_config.max_draft_tokens` so each padded request has the same `1 + max_draft_tokens` width as real ones. (b) Attention-DP: compute `expected_num_active_requests` = max active count across ranks and inject `_get_num_dummy_request()` dummies via `_merge_dummy_request(n)` so every rank runs the same number of (padded) requests → uniform graph-able shape.
+- **Generalizes to:** "pad the dynamic dimension to a captured bucket so the graph always matches"; carries to variable batch size, speculative/medusa/eagle draft widths, attention-DP/EP rank imbalance, any fixed-shape graph fed by variable work; adapt by identifying which dim varies and padding it with shape-correct dummies up to the nearest captured size.
+- **Apply via:** enable CUDA-graph padding with the spec config (MTP path in `pyexecutor/model_engine.py`); attention-DP equalization in `py_executor.py` (`expected_num_active_requests`). Delegate to **perf-torch-cuda-graphs**.
+- **Expected effect:** higher CUDA-graph hit rate (fewer eager fallbacks) under MTP and uneven DP → lower per-step launch overhead; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless — padding adds dummy requests/tokens that are discarded. Care: dummies must be cleaned up (`_finish_dummy_request`) with reserved non-conflicting IDs.
+- **Verify:** confirm CUDA graph used instead of eager on MTP / uneven-DP steps (nsys, ad-conf-check); padded dummy width == `1 + max_draft_tokens`; outputs/accuracy unchanged.
+- **Rollback:** disable CUDA-graph padding (eager on mismatched steps). Trigger: dummy-request lifecycle bugs or DP rank-balancing regressions.
+- **Prior art:** PRs #3096, #3010. Files: `_torch/pyexecutor/model_engine.py`, `pyexecutor/py_executor.py`, `pyexecutor/resource_manager.py`. Owning specialist: **perf-torch-cuda-graph-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/free-mla-intermediates.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/free-mla-intermediates.md
new file mode 100644
index 000000000000..884cbdf14a48
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/free-mla-intermediates.md
@@ -0,0 +1,33 @@
+---
+id: case-free-mla-intermediates
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [memory]
+signals: [memory-capacity-bound, high-concurrency, long-context]
+architectures: [any-sm]
+model_scope: [model-agnostic, mla, deepseek-v3r1]
+phase: [any-phase]
+patterns: [pattern-reuse-computed-state]
+accuracy_risk: lossless
+apply_via_kind: [code-change]
+knobs: []
+specialists: [perf-host-optimization, perf-analysis, perf-sweep-workflow]
+commits: ['fbcf954d9c3a']
+measured: []
+---
+
+# Free large MLA intermediate tensors right after use to cut peak activation memory
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `fbcf954d9c3a` [MLA] Deallocate tensors after use (#3286).
+- **Applies when:** memory-bound — peak GPU memory limits max batch size or KV-cache token budget (OOM at higher concurrency / longer sequences), on an MLA path that materializes several large transient activations per forward (down-proj outputs `q`/`compressed_kv`/`k_pe`, `fused_q`, context/gen attention outputs, FP8 scale tensors).
+- **Mechanism:** MLA forward holds several big intermediates simultaneously; the commit sets each to `None` immediately after its last consumer (`q = compressed_kv = k_pe = None` after `q_b_proj`; `fused_q = None`; `attn_output_context = attn_output_gen = q_nope_scales = None`; `attn_out_latent_scales = None`), dropping the last Python reference so the caching allocator can reuse that memory within the same forward → lower peak activation memory → headroom the KV estimator converts into more KV tokens or a larger max batch.
+- **Generalizes to:** "drop references to large intermediates at their last use to lower peak memory and trade it for batch/KV capacity"; carries to any attention/MLP/MoE forward with chained large temporaries, and FP8 paths where extra scale tensors add up; adapt by finding the true last-use point of each big tensor and nulling it (or reusing buffers) without breaking autograd (inference forward has no backward dependence).
+- **Apply via:** code-level change in `_torch/modules/attention.py` (MLA forward) — set spent intermediates to `None`. Not a config knob; payoff realized through the existing KV-cache auto-estimation (more `max_num_tokens`/KV tokens at the same memory). Delegate memory profiling to **perf-host-optimization** / **perf-analysis**; confirm the larger feasible batch with **perf-sweep-workflow**.
+- **Expected effect:** lower peak GPU memory in MLA forward → enables larger max batch or more KV-cache tokens at fixed memory (indirect throughput gain at higher concurrency); measured Δ (peak mem, achievable batch/KV) to be recorded from run.
+- **Accuracy risk:** lossless (frees only no-longer-referenced tensors; computation and outputs unchanged).
+- **Verify:** peak memory (`torch.cuda.max_memory_allocated`) for the MLA forward before/after; confirm a higher max batch / KV-token count is feasible without OOM; outputs identical.
+- **Rollback:** remove the `= None` deallocations. Trigger: only if a freed tensor is later needed (surfaces as an error, not silent wrongness) — none expected for the inference forward.
+- **Prior art:** PR #3286. Files: `_torch/modules/attention.py` (MLA `forward`: nulling of `q`/`compressed_kv`/`k_pe`, `fused_q`, attention outputs, scale tensors). Detection: **perf-analysis** (memory).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/hoist-torch-compile-closures.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/hoist-torch-compile-closures.md
new file mode 100644
index 000000000000..aed5e60f9362
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/hoist-torch-compile-closures.md
@@ -0,0 +1,77 @@
+---
+id: case-hoist-torch-compile-closures
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [host-overhead]
+signals: [recompilation-churn, gpu-idle-between-steps, host-prep-on-critical-path]
+architectures: [any-sm]
+model_scope: [model-agnostic, spec-decode, sparse-attention]
+phase: [any-phase]
+patterns: [pattern-stable-torch-compile-target]
+accuracy_risk: lossless
+apply_via_kind: [code-change]
+knobs: []
+specialists: [perf-host-optimization, perf-host-analysis, perf-torch-cuda-graphs]
+commits: []
+interactions:
+  - {case: case-cache-step-invariant-per-layer, relation: composes-with, note: "same PR (#12581), same host-overhead theme"}
+measured: []
+---
+
+# Hoist torch.compile closures out of hot methods (and drop torch.compile from trivial ops)
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Applies when:** host-overhead-bound serving where a hot per-call method
+  **defines a `@torch.compile` (or `@maybe_compile`)-decorated function inside
+  its body**, so a fresh compiled callable is created on every call; symptom is
+  repeated graph tracing / recompilation (visible via `TORCH_LOGS=recompiles` or
+  host time spent in Dynamo) and GPU-idle gaps. Also applies when a *trivial*
+  single op (e.g. an in-place `x[:n] += 1`) is wrapped in `torch.compile` for no
+  real benefit. Confirm host-bound first (`perf-host-analysis`). Instances:
+  `mtp.py::forward()` and `dsa.py::prepare_dense_topk_indices()`.
+- **Mechanism:** a `@torch.compile` decorator applied to a nested function
+  rebinds a new optimized callable each call; the compiled-artifact cache keys
+  off the function object, so per-call redefinition causes re-tracing/re-guarding
+  (and can defeat caching entirely), adding host overhead every iteration.
+  Hoisting the decorated function to a stable class method / module-level def
+  lets the compiled artifact persist and be reused. For a one-line op the compile
+  machinery costs more than it saves — run it eager.
+- **Generalizes to:** the pattern "**a `torch.compile` target must be a stable,
+  long-lived callable; never (re)define it on the hot path, and don't compile
+  trivial work.**" Carries to: any `@torch.compile` / `torch.compile(...)`-wrapped
+  closure created inside `forward`/`step`/`prepare`; lambdas passed to
+  `torch.compile` per call; tiny elementwise/in-place ops wrapped "just in case."
+  Adapt by: moving the decorated function to method/module scope and passing only
+  the few tensors it needs (e.g. `seq_lens_cuda`, not the whole `attn_metadata`)
+  so its guards stay small; or removing the decorator when the body is a single
+  cheap op. Confirm host-bound and inspect recompiles first.
+- **Apply via:** **not a server config knob** — a host-side code change. Detect
+  with **perf-host-analysis** (+ `TORCH_LOGS=recompiles` / Dynamo timing);
+  optimize via **perf-host-optimization**. Touch points (prior art, TRT-LLM
+  repo): `tensorrt_llm/_torch/speculative/mtp.py`
+  (`prepare_position_ids_and_last_tokens` hoisted from a `forward()` closure to a
+  method; `update_kv_lens` decorator dropped and inlined as
+  `kv_lens_cuda[:batch_size] += 1`), and
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` (`_get_dense_topk_indices`
+  hoisted out of `prepare_dense_topk_indices` into a `@maybe_compile` method).
+- **Expected effect:** fewer `torch.compile` recompilations and lower Dynamo/host
+  time per call → smaller GPU-idle gaps and higher throughput when host-bound;
+  largest where the method runs every iteration (decode / spec-dec). Record the
+  local Δ (recompile count, per-iteration host time, throughput) — PR #12581
+  ships it without a standalone number; do **not** cite one not measured locally.
+- **Accuracy risk:** lossless — same computation; only *where/when* it is compiled
+  changes (`torch.compile` is output-equivalent within fp tolerance, and the
+  dropped-compile op runs identical eager math).
+- **Verify:** correctness — spec-dec / MTP and DSA tests still pass; outputs match
+  pre-change within tolerance. Perf — `TORCH_LOGS=recompiles` shows the function
+  compiled once, not per-call; per-iteration host time and throughput before vs
+  after.
+- **Rollback:** re-inline the closure / re-add the decorator; trigger: any
+  correctness regression or >5% host-time / throughput regression.
+- **Prior art:** PR #12581; `mtp.py` + `dsa.py`
+  paths above. Owning skill: **perf-host-optimization**; detection:
+  **perf-host-analysis**; `torch.compile` background:
+  **perf-torch-cuda-graphs** (it covers `torch.compile` as a CUDA-graph API).
+  Related: the [loop-invariant-hoist case](cache-step-invariant-per-layer.md) (same PR).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/index.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/index.md
new file mode 100644
index 000000000000..af0b8dce2667
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/index.md
@@ -0,0 +1,165 @@
+# Casebook — Runtime / Execution
+
+Classic optimizations *above the kernel layer* that change how work is
+scheduled, launched, and overlapped — not what is computed. Almost all are
+**lossless** (output-equivalent). Add entries using the schema in
+[case-template.md](../case-template.md); match on **Applies when** signals, adapt
+knob values to the current model/hardware, and measure before claiming a win.
+
+> Routing note: classify host-bound vs GPU-bound **first** (`perf-host-analysis`
+> — GPU-idle ratio, host-prep-exposed ratio, per-iteration host breakdown).
+> These levers only help when the PyExecutor loop is the bottleneck; on a
+> GPU-bound loop they move nothing. Scheduling/graph/overlap knobs are config;
+> host-overhead refactors are code changes delegated to `perf-host-optimization`.
+
+## Recurring patterns in this family
+
+Match on these transferable patterns, not on a case title — most situations are
+a *variation* of a case, not its exact instance.
+
+- **Push language-boundary crossings off the host hot path.** In the PyTorch
+  executor loop, per-iteration objects that wrap a C++ binding type pay a fixed
+  cost on *every* attribute read (each `__getattr__` is a Python↔C++ crossing),
+  and an object that must cross a *process* boundary gets marshalled twice.
+  Three moves remove that cost without touching GPU work: (1) represent hot
+  per-iteration objects as pure-Python (`@dataclass`) instead of binding
+  proxies; (2) mirror the few C++-derived values the loop needs onto plain
+  Python attributes once at construction (the `py_*` field convention) so the
+  loop reads Python, not the binding; (3) when an object must cross a process
+  boundary anyway, bulk-serialize it in one native call and deserialize lazily
+  on the consumer, hoisting just the flag the dispatcher branches on. Host-bound
+  win — confirm the loop is host-bound first. _(Instance: [the LlmResponse case](pybind-wrapper-pure-python.md).)_
+- **Hoist step-invariant work out of the per-layer loop (loop-invariant code
+  motion on the host).** A function called once per decoder layer often
+  recomputes values that are constant across *all* layers within a forward step
+  — KV-pool reshapes/views, stride factors, block-table and request-index
+  slices. Compute them once per step, cache them on the per-step metadata object,
+  and reuse across layers; invalidate once at step start with a cheap boolean
+  flag plus an identity guard. Keep the cache fields as plain instance attributes
+  (init in `__init__`, not class-level annotations) so they stay invisible to
+  dataclass/`torch.compile` introspection. Host-bound win — confirm the loop is
+  host-bound first. _(Instance: [the DSA pool-view cache case](cache-step-invariant-per-layer.md).)_
+- **Keep a `torch.compile` target stable; never (re)define it on the hot path,
+  and don't compile trivial work.** A `@torch.compile` decorator on a function
+  defined *inside* a hot method rebinds a fresh compiled callable every call, so
+  the compiled artifact never sticks and Dynamo re-traces/re-guards each time.
+  Hoist the decorated function to a stable method/module-level def (passing only
+  the tensors it needs, not a whole metadata object so guards stay small); drop
+  the decorator entirely when the body is a single cheap op (the compile
+  machinery costs more than it saves). _(Instance: [the MTP/DSA torch.compile-closure case](hoist-torch-compile-closures.md).)_
+- **Overlap two independent pieces of work.** Host-prep for step N+1 vs device
+  compute for step N (overlap scheduler); two data-independent kernels on a side
+  CUDA stream (shared-vs-routed expert, MLA RoPE-vs-uk-BGEMM); periodic background
+  bookkeeping vs steady-state compute (online-EPLB rebalance); draft forward vs
+  target work (two-model MTP). The hard part is proving data-independence (disjoint
+  output slices / disjoint step state); stream-switching only pays off under
+  CUDA-graph capture. _(Instances: [overlap scheduler](overlap-scheduler.md), [multi-stream](multi-stream-shared-routed-expert.md), [EPLB](overlap-online-eplb.md), and [two-model MTP](two-model-mtp-eagle.md) cases.)_
+- **Process only the real work, not padded-to-max.** Variable-length DP
+  collectives, a single dummy request instead of pad-every-rank, capping
+  CUDA-graph capture / KV to the reachable set. Carries to any DP/ragged-batch
+  collective and any resource sized off `max_*`. _(Instance: [the attention-DP padding case](attention-dp-padding.md); cf. the CUDA-graph capping example in [case-template](../case-template.md).)_
+- **Pad the dynamic dimension to a captured bucket so the graph always matches.**
+  When CUDA graphs need a fixed shape but per-step token/batch counts vary
+  (spec-decode draft width, uneven DP ranks), inject shape-correct dummies up to
+  the nearest captured size instead of falling back to eager. _(Instance: [the CUDA-graph padding case](cuda-graph-padding.md).)_
+- **Reuse computed state instead of recomputing it.** Cache and reuse the
+  compressed/intermediate KV representation across requests on a prefix hit (MLA
+  KV-cache reuse); free large intermediates at their last use to trade peak memory
+  for batch/KV capacity. Carries to MHA/GQA block reuse, chunked-context reuse, any
+  forward with chained large temporaries. _(Instances: [the MLA KV-reuse](mla-kv-cache-reuse.md) and [free-MLA-intermediates](free-mla-intermediates.md) cases.)_
+- **Cut launch/dependency latency on SM≥90 with PDL — but verify per kernel.**
+  Programmatic Dependent Launch overlaps a consumer's preamble with the producer's
+  tail; broad and enabled by default, but it changes inter-kernel ordering and has
+  produced real NaN/accuracy regressions, so some kernels must keep it off.
+  _(Instance: [the PDL case](pdl.md).)_
+- **Trade exact-match draft verification for a bounded relaxed criterion where a
+  little divergence is tolerable** (spec-decode acceptance inside a delimited
+  low-stakes phase). Lossy — gate to where it is safe and keep an accuracy record.
+  _(Instance: [the relaxed-acceptance case](relaxed-mtp-acceptance.md).)_
+- **Fold a model-specific auxiliary cache into the unified paged KV-cache manager.**
+  Instead of a bespoke shadow allocator, register the side cache (DSA indexer K
+  cache, scale-factor pool, landmark cache) as a typed pool sharing block IDs with
+  the main KV — it then inherits paging, prefix reuse, host offload, eviction, and
+  correct per-pool size accounting. Corollaries: sum mixed-dtype pools **additively
+  by physical dtype**; index offloaded blocks by **decoded pool index** (not logical
+  block ID); release Python views before the native free. _(Instance: [the auxiliary-cache-in-manager case](auxiliary-cache-in-kv-manager.md).)_
+- **Chunk long prefill to bound its memory — and align any auxiliary structure's
+  chunking to the primary's.** Fixed-size KV chunks cap prefill activation/logit
+  memory (turning a quadratic full-sequence intermediate into a per-chunk one); an
+  auxiliary structure (sparse indexer) must gather only the current chunk, recompute
+  cache-aware position offsets per chunk, and **inherit the primary chunker's
+  boundaries** rather than running a second scheme. _(Instance: [the chunked-prefill case](chunked-prefill-aligned-auxiliary.md).)_
+- **Skip the sparse/approximate path when it degenerates to the exact one.** Below a
+  size threshold a sparse selection may pick *everything* (selectivity → 1) — detect
+  it, skip the selection machinery entirely, and substitute a trivially-constructed
+  exact result; when the fast path changes control flow, add a CUDA-graph-key
+  dimension so it stays captured. A cousin of "process only the real work." The
+  opposite-regime move: shard replicated per-token work across TP ranks at the long
+  end. _(Instance: [the degenerate-to-dense case](skip-sparse-path-when-degenerate.md).)_
+- **Author a custom op so CUDA-graph capture can include its graphable half.** Split a
+  monolithic op at the graphable/non-graphable seam (shape-static token-wise compute
+  vs batch/data-dependent work), route it through its registered `torch.ops` op so
+  `torch.compile` can trace it, and force straight-line control flow under compile
+  (constant-arity `register_fake`, disable length-adaptive branches). The producer-side
+  complement to turning piecewise capture on. _(Instance: [the split-custom-op case](split-custom-op-for-piecewise-capture.md).)_
+
+## Cases
+
+Match on **Applies when** / **Generalizes to**, then open that case file. Almost
+all runtime/execution cases are **lossless** (output-equivalent).
+
+> Risk key: **lossless** = safe to judge on perf alone · **lossy** = needs an
+> accuracy record + rollback before promotion · **lossless\*** = lossless unless a
+> quant variant (e.g. FP8 KV) is enabled.
+
+### Overlap, launch & scheduling knobs
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [Overlap scheduler](overlap-scheduler.md) | GPU idle between steps; host prep (sched/bookkeeping/input build) on the critical path | pipeline host-prep(N+1) against device compute(N) | lossless |
+| [PDL (dependent launch)](pdl.md) | launch-bound SM≥90; back-to-back dependent kernels with launch/grid-tail exposed | producer-tail / consumer-preamble overlap per dependent pair | lossy |
+| [Remove attention-DP padding](attention-dp-padding.md) | attention-DP with ragged per-rank tokens padded to max; low MFU | process only real tokens, not padded-to-max (length-aware collectives) | lossless |
+| [Skip sparse path when degenerate](skip-sparse-path-when-degenerate.md) | a sparse/approximate path equals dense below a length (DSA short seq) yet still pays selection overhead | skip the selection when it degenerates to exact; keep both paths graph-captured | lossless |
+
+### Host overhead — executor & per-step code
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [Pure-Python hot wrapper](pybind-wrapper-pure-python.md) | PyExecutor loop CPU-bound; per-iter handling on a pybind wrapper (per-field `__getattr__`) / pickled for IPC | push language-boundary crossings off the host hot path | lossless |
+| [Cache step-invariant per-layer values](cache-step-invariant-per-layer.md) | host-bound; a per-decoder-layer fn recomputes step-invariant views/strides/slices | loop-invariant code motion on the host — hoist & cache once per step | lossless |
+| [Hoist torch.compile closures](hoist-torch-compile-closures.md) | host-bound; a hot method defines a `@torch.compile` fn inside its body (recompiles/call), or compiles a trivial op | keep a `torch.compile` target stable & module-level; don't compile trivial work | lossless |
+| [Move bookkeeping into a C++ op](move-bookkeeping-into-cpp-op.md) | host-bound short GPU step; Python per-step reshape/stride + `.item()` count syncs | push Python tensor-wrangling + count syncs into the op that already runs | lossless |
+
+### System scheduling & parallelism
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [Multi-stream shared/routed expert](multi-stream-shared-routed-expert.md) | launch/latency-bound small batch; two independent per-layer sub-computations; CUDA graphs on | issue two data-independent kernels on a side stream to overlap | lossless |
+| [Overlap MLA RoPE with uk-BGEMM](overlap-mla-rope-uk-bgemm.md) | MLA decode, low-latency + CUDA graphs; up-proj BGEMM and RoPE+cache-write run serially | overlap two independent intra-op steps on an aux stream (disjoint output slices) | lossless |
+| [Overlap online-EPLB rebalance](overlap-online-eplb.md) | MoE EP>1 + ADP; expert-load imbalance; online EPLB (`layer_updates_per_iter>0`) rebalance bubble | move periodic background bookkeeping/weight-movement off the critical stream | lossless |
+| [MLA KV-cache reuse](mla-kv-cache-reuse.md) | compute-bound MLA prefill with repeated prefixes (shared prompts/multi-turn), SM90/100 | cache & reuse the compressed/intermediate representation across requests | lossless\* |
+| [Free MLA intermediates](free-mla-intermediates.md) | memory-bound MLA; peak activation memory caps batch/KV (OOM at higher concurrency) | drop refs to large intermediates at last use; trade peak mem for batch/KV | lossless |
+| [Auxiliary cache in the KV manager](auxiliary-cache-in-kv-manager.md) | a model-specific side cache (DSA indexer K) on a bespoke shadow allocator; can't reuse/offload; wrong size accounting | fold a side cache into the unified paged KV manager as a typed pool | lossless |
+| [Chunked prefill + aligned auxiliary](chunked-prefill-aligned-auxiliary.md) | long-context prefill OOM / capped ISL; a quadratic auxiliary intermediate (indexer MQA logits) | chunk prefill to bound memory; align the auxiliary's chunking to the primary's | lossless |
+
+### CUDA graphs
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [Piecewise CUDA graph](piecewise-cuda-graph.md) | launch-bound gen; full-graph capture blocked by a non-graphable op (attention); already on torch.compile | graph the graphable, run the rest eagerly (split at the un-capturable op) | lossless |
+| [CUDA-graph padding](cuda-graph-padding.md) | launch-bound gen; graphs miss because per-step token/batch count varies (MTP draft width, uneven DP) | pad the dynamic dim to a captured bucket so the graph always matches | lossless |
+| [Split a custom op for piecewise capture](split-custom-op-for-piecewise-capture.md) | one custom op mixes graphable (projections) + non-graphable (cache scatter, top-k) work → the whole op stays eager | split at the seam; route through the registered op; straight-line control flow | lossless |
+
+### Speculative decoding / MTP
+
+| Case | Applies when (signal) | Pattern (generalizes to) | Risk |
+|------|----------------------|--------------------------|------|
+| [Relaxed MTP acceptance](relaxed-mtp-acceptance.md) | MTP spec-decode, low accept rate, reasoning model with a thinking phase (R1); not with attention_dp | trade exact-match draft verify for a bounded relaxed criterion where divergence is tolerable | lossy |
+| [Two-model MTP-Eagle overlap](two-model-mtp-eagle.md) | MTP-Eagle where the one-model path constrains and the draft forward is exposed | split a fused speculator into target+draft so the draft forward overlaps target work | lossless |
+
+## Suggested slots (optional — replace or delete)
+
+Areas still open in this family: scheduler ceilings (`max_batch_size` /
+`max_num_tokens` / `max_seq_len`) and further host/sync-free refactors (delegate to
+**perf-host-optimization** / **perf-torch-sync-free**). (Chunked prefill is now
+covered by [chunked-prefill + aligned auxiliary](chunked-prefill-aligned-auxiliary.md).)
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/mla-kv-cache-reuse.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/mla-kv-cache-reuse.md
new file mode 100644
index 000000000000..aea8456a57cc
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/mla-kv-cache-reuse.md
@@ -0,0 +1,35 @@
+---
+id: case-mla-kv-cache-reuse
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [compute, memory]
+signals: [attention-hot-path, long-context]
+architectures: [sm90, sm100]
+model_scope: [model-agnostic, mla, deepseek-v3r1]
+phase: [prefill]
+patterns: [pattern-reuse-computed-state]
+accuracy_risk: mixed
+apply_via_kind: [config-knob, default-on]
+knobs: [enable_block_reuse]
+specialists: [trtllm-serve-config-guide, perf-sweep-workflow]
+commits: ['97bc680cd8a5', '8452775db86b', 'a891013e3c75']
+eligibility:
+  - "SM90–SM100 only; KV-cache quant must be none or FP8 (py_executor_creator.py gate)"
+measured: []
+---
+
+# Reuse cached/prefix KV for MLA to skip down-projection recompute (MLA KV-cache block reuse)
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `97bc680cd8a5` feat: support kv cache reuse for MLA (#3571); related: `8452775db86b` Support FP8 KV Cache Reuse for MLA (#4535), `a891013e3c75` Optimize KV Cache Reuse for MLA (#4869).
+- **Applies when:** compute-bound prefill with repeated prefixes across requests, MLA attention, SM90/SM100. Signals: high cache-hit potential (shared prompts, chat history), large ISL relative to OSL. Enabled by default; originally excluded FP8 KV cache (#4535 lifted that); plain non-FP8 quantized KV still disables reuse.
+- **Mechanism:** MLA stores the *compressed* latent KV (`compressed_kv` + `k_pe`) in the paged KV cache. On a prefix/block hit, cached tokens skip the recompute and run only the up-projection GEMM, using a 192/128 head-size MLA context kernel; shared with chunked context (`mla_context_paged_kv`, `prepare_paged_context_mla`). #4869 fuses the write path into one `invokeMLARopeAppendPagedKVAssignQ` kernel (RoPE + append-paged-KV + assign-Q), drops redundant `.contiguous()` copies, collapses to one up-projection GEMM — cutting the memcpy/GEMM overhead the feature otherwise adds.
+- **Generalizes to:** "cache and reuse the compressed/intermediate representation across requests to skip recompute"; carries to standard MHA/GQA block reuse, chunked-context prefill reuse, prefix caching for any attention variant; adapt by choosing what to persist (full KV vs MLA's compressed latent) and gating on hardware/kernel availability.
+- **Apply via:** `KvCacheConfig(enable_block_reuse=True)` (default on for MLA on SM90/SM100). Auto-disabled outside SM90–SM100 and for KV-cache quant other than none/FP8 (`py_executor_creator.py` gate). Delegate the YAML knob to **trtllm-serve-config-guide**; measure on a prefix-heavy trace with **perf-sweep-workflow**.
+- **Expected effect:** higher prefill throughput / lower TTFT on prefix-overlapping traffic. README caveat: "GPU memory consumption may be higher and the E2E performance may have regression in some cases" (extra memcpy + GEMMs) — workload-dependent; measured Δ to be recorded from run.
+- **Accuracy risk:** lossless for bf16/non-quant reuse; with FP8 KV reuse (#4535) the stored latent is FP8 — lossy in the same way FP8 KV cache already is, not an additional reuse-specific loss.
+- **Verify:** prefill throughput / TTFT + KV-cache-hit rate on a repeated-prefix trace; GPU memory headroom. For FP8 KV reuse run an accuracy/parity check vs bf16 KV.
+- **Rollback:** `KvCacheConfig(enable_block_reuse=False)`. Trigger: E2E regression on a low-reuse workload, memory pressure shrinking batch/KV, or accuracy drift with FP8 KV.
+- **Prior art:** PRs #3571, #4535, #4869. Files: `cpp/.../kernels/mlaKernels.cu` (`invokeMLARopeAppendPagedKVAssignQ`), `thop/mlaPreprocessOp.cpp`, `_torch/attention_backend/trtllm.py` (`prepare_paged_context_mla`), `pyexecutor/py_executor_creator.py` (SM/quant gating).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/move-bookkeeping-into-cpp-op.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/move-bookkeeping-into-cpp-op.md
new file mode 100644
index 000000000000..5109c7185906
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/move-bookkeeping-into-cpp-op.md
@@ -0,0 +1,36 @@
+---
+id: case-move-bookkeeping-into-cpp-op
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [host-overhead, launch, sync]
+signals: [gpu-idle-between-steps, host-prep-on-critical-path]
+architectures: [sm90, sm100]
+model_scope: [model-agnostic, sparse-attention, mla]
+phase: [decode]
+patterns: [pattern-fold-bookkeeping-into-consuming-op]
+accuracy_risk: lossless
+apply_via_kind: [code-change, kernel-change]
+knobs: []
+specialists: [kernel-cuda-specialist, perf-torch-sync-free, perf-host-analysis]
+commits: ['ae84aaddb6f1', '7081f254cf']
+interactions:
+  - {case: case-cache-step-invariant-per-layer, relation: composes-with, note: same DSA host-overhead theme}
+  - {case: case-hoist-torch-compile-closures, relation: composes-with, note: same DSA host-overhead theme}
+measured: []
+---
+
+# Move per-step tensor bookkeeping off the Python hot path into a C++/fused op
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `ae84aaddb6f1` [perf] Reduce host overhead in DSA MLA attention (#12631); related: `7081f254cf` [perf] Add custom indexer k cache scatter op (#8960, the op's origin — replaced a PyTorch advanced-indexing scatter with the custom CUDA kernel this case then widens).
+- **Applies when:** host-bound per step with a short GPU step — the Python attention/indexer path does per-step tensor reshaping/viewing/striding (byte reinterprets, `.view(uint8)`, `as_strided`, slot-map slicing) and/or `.item()` syncs to derive counts; nsys shows GPU idle, time in Python view/stride ops, a `cudaStreamSynchronize` from `.item()` on the critical path.
+- **Mechanism:** (a) fuse bookkeeping into the consuming C++ op: `indexer_k_cache_scatter_op` now takes raw `k_fp8`/`k_scale` + full slot-mapping buffers + a `num_tokens` count and does the FP8/float32 byte-reinterpret and `[:num_tokens]` slice internally — deleting the Python view/as_strided/slice. (b) Remove a host-device sync: counts previously derived via `host_context_lengths.sum().item()` and a Python `_parse_request_types` (`host_request_types.sum().item()`) are deleted; counts are passed in precomputed.
+- **Generalizes to:** "push per-step Python tensor wrangling and count-deriving syncs into the kernel/op that already runs"; carries to other backends' cache scatter/gather, RoPE/quant pre-steps, any op fed by Python reshapes; adapt by widening the op signature to accept raw tensors + a length and threading known counts through instead of `.item()`. Upstream sibling: **replace a PyTorch advanced-indexing scatter/gather** (`arange`+broadcast+`_unravel_indices`+masked assign) **with a stride-aware custom CUDA scatter** — one block per token, vectorized 4-byte stores — especially for writes into a *non-contiguous paged pool* (#8960 introduced exactly this op before #12631 pushed the remaining bookkeeping into it).
+- **Apply via:** C++ custom-op signature change (`torch.ops.trtllm.indexer_k_cache_scatter_op(..., num_tokens)`); delegate the op to **kernel-cuda-specialist** and the `.item()` removal to **perf-torch-sync-free**.
+- **Expected effect:** lower per-step host time / fewer Python ops + one fewer host-device sync → less GPU idle in DSA attention; direction only — measured Δ (host prep, GPU idle, step latency) to be recorded from run.
+- **Accuracy risk:** lossless — same scatter result and counts; only where/when the reinterpret/counting happens changed. (Op `TORCH_CHECK`s element sizes, failing loudly on dtype mismatch.)
+- **Verify:** nsys shows reduced Python view/stride time and removal of the `.item()` sync; step latency down; KV-scatter parity / DSA accuracy unchanged.
+- **Rollback:** restore the Python byte-view/slice path and old op signature; re-add `_parse_request_types`. Trigger: cache-scatter mismatch or DSA accuracy drop.
+- **Prior art:** PR #12631. Files: `cpp/.../thop/IndexerKCacheScatterOp.cpp`, `thop/attentionOp.cpp`, `_torch/attention_backend/sparse/dsa.py`, `attention_backend/trtllm_gen.py`. Detection: **perf-host-analysis**. Related: the [DSA pool-view-cache](cache-step-invariant-per-layer.md) and [torch.compile-closure](hoist-torch-compile-closures.md) cases (same DSA host-overhead theme).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/multi-stream-shared-routed-expert.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/multi-stream-shared-routed-expert.md
new file mode 100644
index 000000000000..775ca6491daa
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/multi-stream-shared-routed-expert.md
@@ -0,0 +1,37 @@
+---
+id: case-multi-stream-shared-routed-expert
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch, compute]
+signals: [small-batch-decode, many-small-kernels]
+architectures: [any-sm]
+model_scope: [model-agnostic, moe, mla, deepseek-v3r1]
+phase: [decode]
+patterns: [pattern-overlap-independent-work]
+accuracy_risk: lossless
+apply_via_kind: [config-knob, code-change]
+knobs: [cuda_graph_config, use_cuda_graph]
+specialists: [perf-torch-cuda-graphs, perf-sweep-workflow, perf-torch-cuda-graph-specialist]
+commits: ['4855431d3d4e']
+eligibility:
+  - "CUDA graphs must be enabled (gate: do_multi_stream = is_graph_capturing() and aux_stream is not None)"
+interactions:
+  - {feature: cuda-graphs, relation: depends-on, note: stream switching only pays off under graph capture/replay}
+measured: []
+---
+
+# Overlap shared-expert and routed-expert (MoE) compute on a side CUDA stream
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `4855431d3d4e` [Deepseek] Redesign multi-stream API (#3459).
+- **Applies when:** launch/exposed-latency bound (GPU not saturated, small batch / low concurrency, min-latency mode) on a model with two independent sub-computations per layer — shared experts vs routed MoE, or the two MLA layernorms (q vs kv) / MLA `_run_bmm` vs `_concat_kv_cache`. Requires CUDA Graphs enabled, because stream switching has host overhead only hidden under graph capture/replay.
+- **Mechanism:** issues two independent ops on separate CUDA streams (default = `fn0`, `aux_stream` = `fn1`) with `cuda.Event` record/wait pairs ordering them, so the GPU runs them concurrently instead of serializing — filling SM gaps a single small kernel leaves. Gate is literally `do_multi_stream = is_graph_capturing() and aux_stream is not None`.
+- **Generalizes to:** "issue two data-independent kernels on a side stream to overlap them"; carries to MoE shared+routed split, MLA dual-layernorm, attention/MoE-chunking overlap (codebase enumerates `AuxStreamType = {Attention, MoeShared, MoeChunkingOverlap}`), and EPLB statistic/weight movement; adapt by identifying a true independent pair, allocating an aux stream + two events, and only enabling under graph capture.
+- **Apply via:** `maybe_execute_in_parallel(fn0, fn1, event0, event1, aux_stream)` (`_torch/modules/multi_stream_utils.py`); aux streams created per `AuxStreamType` in model `__init__`. Operationally, enable by turning on CUDA Graphs (`cuda_graph_config` / `use_cuda_graph`) — multi-stream then self-activates. Delegate to **perf-torch-cuda-graphs**; measure with **perf-sweep-workflow**.
+- **Expected effect:** lower per-iteration latency in the low-latency/small-batch regime (overlap of shared- and routed-expert work); no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless (pure scheduling; identical math, only ordered by events).
+- **Verify:** TPOT/ITL + end-to-end latency at low concurrency with vs without CUDA Graphs; confirm two streams concurrently active in nsys. No parity check needed.
+- **Rollback:** pass `aux_stream=None` (disables multi-stream) or disable CUDA Graphs — falls back to sequential `fn0(); fn1()`. Trigger: throughput regression at high batch (tuned for low latency), or no measured overlap.
+- **Prior art:** PR #3459. Files: `_torch/modules/multi_stream_utils.py` (`maybe_execute_in_parallel`), `modeling_deepseekv3.py` (`Deepseekv3MoE.forward`), `_torch/modules/attention.py` (MLA), `_torch/utils.py` (`AuxStreamType`, `EventType`). Owning specialist: **perf-torch-cuda-graph-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-mla-rope-uk-bgemm.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-mla-rope-uk-bgemm.md
new file mode 100644
index 000000000000..5989278a5b3e
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-mla-rope-uk-bgemm.md
@@ -0,0 +1,35 @@
+---
+id: case-overlap-mla-rope-uk-bgemm
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch, sync]
+signals: [small-batch-decode, many-small-kernels]
+architectures: [any-sm]
+model_scope: [model-agnostic, mla, deepseek-v3r1]
+phase: [decode]
+patterns: [pattern-overlap-independent-work]
+accuracy_risk: lossless
+apply_via_kind: [code-change, kernel-change]
+knobs: []
+specialists: [kernel-cuda-specialist, perf-torch-cuda-graph-specialist, perf-nsight-systems]
+commits: ['51545560da00']
+eligibility:
+  - "multi-stream gated on do_multi_stream() — CUDA graphs / low-latency path enabled"
+measured: []
+---
+
+# Overlap MLA RoPE with the up/K batched-GEMM on an aux CUDA stream
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `51545560da00` [feat] Add rope and uk-bgemm overlap for mla generation (#8495).
+- **Applies when:** sync/launch-bound MLA decode where two independent pre-attention steps run serially: the q_nope up-projection BGEMM (`q_nope × k_b_proj_trans`, "uk-bgemm") and RoPE-on-Q + latent-cache write. Best on the low-latency path with CUDA graphs (multi-stream gated on `do_multi_stream()`); small per-step tensors leave SM/issue slots idle while these run back-to-back.
+- **Mechanism:** runs the two ops concurrently via `maybe_execute_in_parallel(fn0, fn1, event0, event1, aux_stream)` — `fn0` = up-proj BGEMM (`torch.ops.trtllm.bmm_out`, or `fp8_block_scaling_bmm_out` for FP8) on the default stream, `fn1` = `mla_rope_generation(...)` on `aux_stream`, synchronized with record/wait events. BGEMM writes into `fused_q[..., :kv_lora_rank]` while RoPE fills the rope portion — disjoint writes, no extra copy. New fused C++ op `torch.ops.trtllm.mla_rope_generation` (`dsv3RopeOp.cpp`) packages RoPE+cache-write into one aux-stream callable.
+- **Generalizes to:** "overlap two independent intra-op steps on an aux stream, writing into disjoint slices of a shared output, gated to the CUDA-graph/low-latency regime"; carries to MoE shared-vs-routed overlap, RoPE-vs-KV-write in GQA, projection-vs-norm pairs; adapt by proving data-independence (disjoint output regions), bracketing with record/wait events, and disabling the second stream off the CUDA-graph path.
+- **Apply via:** code-level multi-stream via `maybe_execute_in_parallel` (`_torch/modules/multi_stream_utils.py`) + aux_stream + two events; enable CUDA graphs to activate. The fused RoPE op (`dsv3RopeOp.cpp`) is a **kernel-cuda-specialist** change. Delegate to **perf-torch-cuda-graph-specialist**; confirm gap closure with **perf-nsight-systems**.
+- **Expected effect:** reduced exposed latency between up-proj and RoPE → lower MLA decode step time under CUDA graphs; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless (pure scheduling; disjoint writes). The FP8 BGEMM path carries its own FP8 risk independent of this overlap.
+- **Verify:** decode step latency + nsys timeline showing BGEMM and RoPE overlapping on two streams; parity-equal attention output vs serial path; correctness with CUDA graphs on/off.
+- **Rollback:** pass `aux_stream=None` (or disable multi-stream / CUDA graphs) to run serially. Trigger: capture failures or no measured overlap.
+- **Prior art:** PR #8495. Files: `_torch/modules/attention.py`, `cpp/.../thop/dsv3RopeOp.cpp` (`mla_rope_generation`), `attention_backend/trtllm.py`. Related: the [shared/routed-expert overlap case](multi-stream-shared-routed-expert.md) (same `maybe_execute_in_parallel` pattern).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md
new file mode 100644
index 000000000000..95f355c3ddd0
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md
@@ -0,0 +1,35 @@
+---
+id: case-overlap-online-eplb
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [communication, launch]
+signals: [expert-load-imbalance, gpu-idle-between-steps]
+architectures: [sm90, sm100]
+model_scope: [model-agnostic, moe, deepseek-v3r1, qwen3-moe]
+phase: [decode]
+patterns: [pattern-overlap-independent-work]
+accuracy_risk: lossless
+apply_via_kind: [config-knob]
+knobs: [num_slots, layer_updates_per_iter]
+specialists: [trtllm-moe-develop, perf-sweep-workflow]
+commits: ['ead89a0e408e', '9241ccaf2770']
+eligibility:
+  - "EP>1 with attention-DP (use_dp); online EPLB active (layer_updates_per_iter > 0)"
+measured: []
+---
+
+# Overlap online-EPLB rebalance with MoE compute, and extend EPLB to more backends
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `ead89a0e408e` [perf] Improve the performance of online EPLB on Hopper by better overlapping (#6624); related: `9241ccaf2770` Enable EPLB for trtllm-gen and cutlass backend (#8886).
+- **Applies when:** communication/sync-exposed MoE with expert load imbalance across EP ranks (hot experts overloading some ranks while others idle), online/dynamic EPLB enabled (`layer_updates_per_iter > 0`). Signals: MoE-heavy model, EP>1 + `use_dp`, per-rank expert-time skew, a per-iteration bubble where statistics gather / expert-weight reshuffle stalls compute.
+- **Mechanism:** #6624 routes the load-balancer's per-iteration work — local statistic collection (`update_local_statistic`, `update_statistic_with_local_ids`/`...with_gathered_statistic`) and the GPU "wait stage" that swaps in rebalanced expert weights — onto a dedicated `aux_stream` (`AuxStreamType.MoeBalancer`) ordered against compute via `EventType.Main`/`EventType.MoeBalancer`, replacing the old `cudagraph_stream`+`statistic_stream` pair, so rebalance overhead hides behind MoE GEMMs. It also splits the statistic path (`start_wait_gpu_stage`/`done_wait_gpu_stage`) to overlap allgather of statistics with routing. #8886 lifts EPLB out of `fused_moe_wide_ep` into the shared `MoEInterface` base, so the **cutlass** and **trtllm-gen** MoE backends (and GPT-OSS) gain EPLB, not just wide-EP.
+- **Generalizes to:** "move periodic background bookkeeping/weight-movement off the critical stream and overlap it with steady-state compute"; carries to KV-cache defrag/transfer, weight prefetch/streaming, any allgather of side metadata; adapt by isolating the background work on an aux stream with event ordering and only doing it when the dynamic path is active.
+- **Apply via:** `MoeLoadBalancerConfig(num_slots=<divisible by ep_size>, layer_updates_per_iter=<N>)` (`llmapi/llm_args.py`); `layer_updates_per_iter=0` = static routing (no overlap needed), `>0` = online EPLB (this overlap engages). Requires `use_dp` (ADP) and EP>1. Delegate wiring to **trtllm-moe-develop**; sweep `num_slots`/`layer_updates_per_iter` with **perf-sweep-workflow**.
+- **Expected effect:** higher MoE/decode throughput by hiding online-EPLB rebalance cost; for #8886, EPLB becomes available on cutlass/trtllm-gen backends; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless (EPLB relocates expert replicas/slots across ranks; routing math and outputs unchanged — slots map back to the same experts).
+- **Verify:** decode throughput + per-rank MoE-time balance (straggler spread) with online EPLB on vs off; confirm MoeBalancer stream concurrent with MoE GEMMs in nsys. Sanity-check accuracy unchanged.
+- **Rollback:** `layer_updates_per_iter=0` (static routing) or omit the `moe_load_balancer` config. Trigger: no measured imbalance to recover, overlap not materializing, or instability during weight swap.
+- **Prior art:** PRs #6624, #8886. Files: `_torch/modules/fused_moe/moe_load_balancer.py` (`start_wait_gpu_stage`, `update_local_statistic`, `is_static_routing`), `fused_moe/interface.py` (`_init_load_balancer`), `_torch/utils.py` (`AuxStreamType.MoeBalancer`), `llmapi/llm_args.py` (`MoeLoadBalancerConfig`).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-scheduler.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-scheduler.md
new file mode 100644
index 000000000000..09c11394d5c6
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-scheduler.md
@@ -0,0 +1,41 @@
+---
+id: case-overlap-scheduler
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [host-overhead, launch]
+signals: [gpu-idle-between-steps, host-prep-on-critical-path, small-batch-decode, high-concurrency]
+architectures: [any-sm]
+model_scope: [model-agnostic, dense, moe]
+phase: [decode]
+patterns: [pattern-host-prep-pipelining]
+accuracy_risk: lossless
+apply_via_kind: [default-on, config-knob]
+knobs: [disable_overlap_scheduler, enableTrtOverlap]
+specialists: [perf-host-optimization, perf-host-analysis]
+commits: ['72057a0a64bf']
+eligibility:
+  - "no beam search, no speculative decoding, no encoder model (mechanism auto-disables)"
+interactions:
+  - {feature: beam-search, relation: incompatible-with, note: auto-disables}
+  - {feature: spec-decode, relation: incompatible-with, note: auto-disables in the generic path}
+  - {feature: encoder-models, relation: incompatible-with, note: auto-disables}
+  - {case: case-two-model-mtp-eagle, relation: composes-with, note: re-enables overlap under MTP-Eagle two-model spec-decode}
+measured: []
+---
+
+# Overlap host scheduling with GPU compute (overlap scheduler)
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `72057a0a64bf` [TRTLLM-3429] feat: Overlap scheduling in C++ runtime (#3625); the same mechanism is the default in the PyTorch executor (`_executor_loop_overlap`, knob `disable_overlap_scheduler` in `tensorrt_llm/llmapi/llm_args.py`).
+- **Applies when:** launch/host-prep-exposed bottleneck — nsys shows GPU idle between forward steps and host prep (scheduling, request bookkeeping, decoder/sampler prep, input tensor build) is on the critical path; steady-state generation-heavy serving with many in-flight requests; small-batch / low-arithmetic-intensity decode where per-step host work rivals GPU step time.
+- **Counter-signals:** GPU-bound steps (host prep already hidden — nothing to win, and the extra in-flight state costs memory); beam search / speculative decoding / encoder models (mechanism auto-disables; the two-model MTP-Eagle path restores overlap — see [that case](two-model-mtp-eagle.md)); debugging sessions that need strictly serialized step semantics.
+- **Mechanism:** runs CPU-side preparation for step N+1 (micro-batch scheduling, decoder input/output marshaling) concurrently with engine execution of step N, hiding host latency behind GPU compute instead of serializing before it. C++ adds `enableTrtOverlap` to `ExecutorConfig` and an `invokeGatherBatch` runtime kernel so the overlapped same-batch token gather stays on GPU; keeps the two steps' state disjoint.
+- **Generalizes to:** producer/consumer pipelining of host-prep against device compute; carries to the PyTorch `_executor_loop_overlap`, to spec-decode draft/target overlap, and to any per-iteration loop where input marshaling can be hoisted one step ahead; adapt by ensuring step N+1 prep does not read state mutated by step N's results (disjoint-state plumbing is the hard part).
+- **Apply via:** PyTorch backend — leave `disable_overlap_scheduler=False` (default on). C++ runtime — `ExecutorConfig.setEnableTrtOverlap(true)`. Auto-disabled for beam search, speculative decoding, encoder. Delegate host-overhead validation to **perf-host-optimization**.
+- **Expected effect:** higher throughput and lower per-step latency by shrinking inter-step GPU idle; direction only (no number quoted) — measured Δ (tok/s, GPU-idle ratio) to be recorded from run.
+- **Accuracy risk:** lossless — pure scheduling reorder; identical math.
+- **Verify:** A/B the knob (`disable_overlap_scheduler` on/off) at fixed concurrency: nsys inter-step GPU-idle / host-prep-exposed ratio drops (detection recipe: **perf-host-analysis**); throughput up; output parity vs overlap-off.
+- **Rollback:** `disable_overlap_scheduler=True` (PyTorch) or `enableTrtOverlap=false` (C++). Trigger: throughput regression or correctness mismatch.
+- **Prior art:** PR #3625. Files: `cpp/include/.../executor/executor.h`, `batch_manager/trtGptModelInflightBatching.cpp`, `runtime/runtimeKernels.cu` (`invokeGatherBatch`), `_torch/pyexecutor/decoder.py`. Detection: **perf-host-analysis**.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/pdl.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/pdl.md
new file mode 100644
index 000000000000..ad01d108bbc8
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/pdl.md
@@ -0,0 +1,35 @@
+---
+id: case-pdl
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch]
+signals: [many-small-kernels, small-batch-decode]
+architectures: [sm90, sm100]
+model_scope: [model-agnostic, dense, moe]
+phase: [any-phase]
+patterns: [pattern-pdl-producer-consumer-overlap]
+accuracy_risk: lossy
+apply_via_kind: [default-on, env-var, kernel-change]
+knobs: [TRTLLM_ENABLE_PDL]
+specialists: [kernel-cuda-specialist]
+commits: ['8462cf6c96f1', '84d2f1281857', '21a93fbf9d10', '9e7b50aefb', '6ee8dbfe0b', 'ba25b6afae08', 'bf7142f8d1', '1c4dacb19a52', '34e2fa5c96']
+eligibility:
+  - "SM90+ only (gated on getSMVersion() >= 90)"
+measured: []
+---
+
+# Enable Programmatic Dependent Launch (PDL) to hide kernel-launch/dependency latency
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `8462cf6c96f1` [TRTLLM-9578][feat] make PDL enabled by default (#9695); related: `84d2f1281857` add PDL support for more kernels (#7977, `launchWithPdlWhenEnabled`), `21a93fbf9d10` PDL for CuteDSL + overlap MoeOutputMemset (#10043). Per-kernel accuracy WARs (PDL disabled): `9e7b50aefb` quant kernels (#10285), `6ee8dbfe0b`/`ba25b6afae08` TinyGEMM, `bf7142f8d1` TRTLLM-GEN routing NaN (#11994), `1c4dacb19a52` dsv3 MoE (#9799), `34e2fa5c96` Qwen3-235B ATP (#9530).
+- **Applies when:** launch-bound on SM≥90 — back-to-back dependent kernels (MoE comm→GEMM, layernorm→GEMM, router/quant→GEMM) where each waits for the prior to drain and launch overhead / grid-tail is exposed; chains of many small kernels at low concurrency. Not available below SM90 (gated on `getSMVersion() >= 90`).
+- **Mechanism:** PDL lets a dependent kernel begin grid setup/preamble while the producer's tail still drains, via `cudaLaunchAttributeProgrammaticStreamSerialization`, overlapping launch/tail latency. The default-on change flips `enablePDL` to `true`, making `TRTLLM_ENABLE_PDL=0` an explicit opt-out (still SM≥90 gated).
+- **Generalizes to:** producer-tail / consumer-preamble overlap for any dependent-kernel pair; carries to new kernels via the shared helper, to CuteDSL/trtllm-gen kernels, and to MoE all-to-all comm kernels; adapt by adding the launch attribute (or `launchWithPdlWhenEnabled`), gating arch ≥ sm90, AND accuracy-checking each kernel (PDL is not universally safe).
+- **Apply via:** env `TRTLLM_ENABLE_PDL` (default `1` on SM≥90; `0` disables globally). New kernels: `tensorrt_llm::common::launchWithPdlWhenEnabled(...)` or set the launch attribute guarded by `getEnvEnablePDL()`. Delegate to **kernel-cuda-specialist**.
+- **Expected effect:** lower exposed launch/dependency latency between kernels → higher throughput in kernel-chain-bound regions; direction only — measured Δ to be recorded from run.
+- **Accuracy risk:** lossy/risky per-kernel — PDL changes inter-kernel ordering/visibility and HAS produced real accuracy/NaN regressions (TinyGEMM, quant, trtllm-gen FMHA/routing, dsv3 MoE, Qwen3-235B). Needs an on-disk accuracy record + per-kernel rollback; some kernels must keep PDL off.
+- **Verify:** kernel-chain region latency / inter-kernel gap drops in nsys; AND an accuracy/parity eval (WAR commits exist because kernels regressed) — no NaN, score within tolerance vs `TRTLLM_ENABLE_PDL=0`.
+- **Rollback:** global `TRTLLM_ENABLE_PDL=0`; per-kernel disable the attribute. Trigger: accuracy drop or NaN in any affected module.
+- **Prior art:** PRs #9695, #7977, #10043 (+WARs). Files: `cpp/.../common/envUtils.{cpp,h}` (`getEnvEnablePDL`, `launchWithPdlWhenEnabled`), `cutlass_kernels/moe_gemm/moe_kernels.cu`, `fusedLayernormKernels/*`, `mlaKernels.cu`, `fusedMoeCommKernels.cu`. Owning specialist: **kernel-cuda-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/piecewise-cuda-graph.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/piecewise-cuda-graph.md
new file mode 100644
index 000000000000..b31190e89984
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/piecewise-cuda-graph.md
@@ -0,0 +1,35 @@
+---
+id: case-piecewise-cuda-graph
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch]
+signals: [many-small-kernels, slow-path-fallback]
+architectures: [any-sm]
+model_scope: [model-agnostic]
+phase: [decode]
+patterns: [pattern-graph-the-graphable]
+accuracy_risk: lossless
+apply_via_kind: [config-knob]
+knobs: [torch_compile_piecewise_cuda_graph, cuda_graph_batch_sizes]
+specialists: [perf-torch-cuda-graphs, perf-torch-cuda-graph-specialist]
+commits: ['91bf5e6a8e73']
+eligibility:
+  - "requires torch_compile_enabled; asserts torch_compile_fullgraph=True"
+measured: []
+---
+
+# Piecewise CUDA Graph — capture stable regions, leave dynamic/unsupported ops out
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `91bf5e6a8e73` Add Piecewise CUDA Graph Support (#3804).
+- **Applies when:** launch-bound generation where full-iteration CUDA-graph capture is blocked by ops that can't/shouldn't be captured (notably the fused `attention` op). Signals: launch overhead dominates between kernels; full graph unavailable; you already use `torch.compile`.
+- **Mechanism:** splits the compiled FX graph at the attention boundary and CUDA-graph-captures only the contiguous stable submodules around it, leaving attention eager. Splits on `torch.ops.trtllm.attention` (rewritten to `attention_inplace`, no new tensor → safe resume); attention submods go in `exclude_modules_id`; a `PiecewiseInterpreter` wraps each non-excluded `submod` in a `PiecewiseRunner` capturing one graph per `cuda_graph_batch_sizes` entry. All pieces share one `graph_pool_handle`. Each piece replays as a single launch.
+- **Generalizes to:** "graph the graphable, run the rest eagerly"; carries to any region broken by a dynamic-shape op, data-dependent control flow, a non-capturable custom op, or spec-decode bookkeeping; adapt by choosing the split op, ensuring safe resume after it (in-place/no-new-tensor), and enumerating token-count buckets.
+- **Apply via:** config `torch_compile_piecewise_cuda_graph: true` (requires `torch_compile_enabled`, asserts `torch_compile_fullgraph=True`) + `cuda_graph_batch_sizes`; flags `--use_torch_compile --use_piecewise_cuda_graph`. Delegate to **perf-torch-cuda-graphs** / **perf-torch-cuda-graph-specialist**.
+- **Expected effect:** reduced launch overhead / lower per-step latency by graphing stable regions; attention stays eager; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless — graph replay reproduces the same kernels. Care: a warmup flag `set_enable_piecewise_cuda_graph_capture_flag(False)` disables capture during runs that "would produce wrong results"; token counts not captured fall back to non-graph.
+- **Verify:** confirm pieces captured/replayed (ad-conf-check); nsys graph coverage vs eager attention; outputs vs eager and full-graph.
+- **Rollback:** `torch_compile_piecewise_cuda_graph=false`. Trigger: capture errors, wrong warmup results, or token counts persistently missing from `cuda_graph_batch_sizes`.
+- **Prior art:** PR #3804. Files: `_torch/compilation/piecewise_optimizer.py`, `compilation/backend.py`, `pyexecutor/config.py`. Owning specialist: **perf-torch-cuda-graph-specialist**.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/pybind-wrapper-pure-python.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/pybind-wrapper-pure-python.md
new file mode 100644
index 000000000000..ba79d3959960
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/pybind-wrapper-pure-python.md
@@ -0,0 +1,93 @@
+---
+id: case-pybind-wrapper-pure-python
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [host-overhead, launch]
+signals: [gpu-idle-between-steps, host-prep-on-critical-path, high-concurrency]
+architectures: [any-sm]
+model_scope: [model-agnostic]
+phase: [any-phase]
+patterns: [pattern-push-language-boundary-off-hot-path]
+accuracy_risk: lossless
+apply_via_kind: [code-change]
+knobs: []
+specialists: [perf-host-optimization, perf-host-analysis]
+commits: []
+measured: []
+---
+
+# Re-implement a hot-path binding wrapper in pure Python to cut pybind overhead
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Applies when:** host-overhead-bound serving — the PyExecutor loop is
+  CPU-bound, the GPU shows idle gaps between steps / low utilization, and host
+  time sits in per-iteration response/request handling rather than kernels.
+  Confirm with `perf-host-analysis` (GPU-idle ratio, host-prep-exposed ratio,
+  per-iteration host breakdown) **before** doing this — it is an invasive
+  refactor, not a knob. The signal it targets: a Python object on the
+  per-response/per-request hot path is a thin wrapper over a C++ binding type
+  (`bindings.executor.Response` / `Result`) whose attributes are read many times
+  per iteration via `__getattr__`, and/or is pickled to cross the worker→client
+  process boundary. Most acute at high concurrency / high request rate, where
+  response handling runs every iteration.
+- **Mechanism:** every attribute read on a pybind-wrapped object is a Python↔C++
+  crossing with fixed overhead, so a wrapper that delegates via `__getattr__`
+  pays it per field per iteration, and an object pickled for IPC is marshalled
+  twice. Making the hot object pure-Python (a `@dataclass`) turns those reads
+  into plain attribute lookups; mirroring the few C++-derived values it needs
+  onto `py_*` attributes once removes the remaining crossings; and serializing
+  the C++ `Result` in a single native call (returning bytes + the `is_final`
+  flag) with lazy deserialize on the consumer replaces N per-attribute crossings
+  plus a second pickle with one bulk serialize.
+- **Generalizes to:** the pattern "**push language-boundary crossings off the
+  host hot path** — pure-Python hot objects, `py_*`-mirrored C++ fields read
+  once, and bulk-serialize-once + lazy-deserialize for anything that must cross a
+  process boundary." Carries to: other executor binding-wrappers
+  (request / sampling-params / result proxies — the `py_*` mirror convention is
+  already the idiom, extend it); any tight Python loop reading attributes off a
+  C-extension proxy (nanobind / Cython / ctypes — hoist the reads to locals, cf.
+  the `begin_compute` local replacing the `request.context_current_position`
+  property read in the same PR); and IPC paths that pickle a native object
+  (serialize once in native code, defer deserialize to the consumer, lift only
+  the flag the dispatcher branches on). Adapt by: profiling the host loop first;
+  confirming the object is genuinely hot (many per-iteration crossings) before
+  refactoring; and keeping the lazy-deserialize seam correct — the consumer must
+  deserialize before touching real fields, and any control-flow flag it needs
+  pre-deserialize (here `is_final`) must be hoisted out.
+- **Apply via:** **not a server config knob** — a host-side code change in the
+  executor/binding layer. Delegate optimization to **perf-host-optimization**
+  (line_profiler / nsys host rounds) and detection to **perf-host-analysis**.
+  Touch points (prior art, TRT-LLM repo): the Python wrapper
+  `tensorrt_llm/_torch/pyexecutor/llm_request.py` (`LlmResponse` / `LlmResult`,
+  `py_*` fields), the loop `tensorrt_llm/_torch/pyexecutor/py_executor.py`
+  (`_handle_responses`), the lazy-deserialize seam
+  `tensorrt_llm/executor/result.py` (`_handle_response`), and the C++
+  bulk-serialize entry points `createSerializedResult` / `deserialize_result`
+  (`cpp/.../batch_manager/llmRequest.cpp` + pybind bindings).
+- **Expected effect:** lower host/CPU time per executor iteration in
+  response/request handling → higher throughput and smaller GPU-idle gaps when
+  the loop is host-bound; no change when it is GPU-bound. Largest at high
+  concurrency / high request rate. Measured Δ to be recorded from a
+  host-profiling run (per-iteration host time, GPU-idle ratio, throughput before
+  vs after) — PR #5224 is "a step of #3034" and reports no standalone speedup;
+  do **not** cite a number not measured locally.
+- **Accuracy risk:** lossless — purely changes how data is marshalled across the
+  binding/IPC boundary; outputs (tokens, and logits/logprobs via `PyResult`) are
+  unchanged.
+- **Verify:** correctness — return-logits / logprobs and streaming tests still
+  pass (e.g. `tests/unittest/_torch/test_return_logits.py`); confirm final-token
+  semantics intact (`is_final` still drives `request_done` in
+  `_handle_responses`) and that every consumer deserializes before reading result
+  fields. Perf — `perf-host-analysis` / `perf-host-optimization`: per-iteration
+  host time in response handling, GPU-idle ratio, and throughput at high
+  concurrency, before vs after.
+- **Rollback:** revert to the binding-wrapper `LlmResponse` / `LlmResult`;
+  trigger: any correctness regression (logits / logprobs / streaming) or >5%
+  host-time / throughput regression. Likely failure mode to watch: a consumer
+  reading result fields without first calling `deserialize()` on the lazy seam.
+- **Prior art:** PR #5224 (a step of #3034), "Re-implement LlmResponse in Python
+  to reduce host overhead of pybind"; TRT-LLM paths above. Owning skill:
+  **perf-host-optimization**; detection: **perf-host-analysis**. Related
+  runtime/execution levers: [overlap scheduler](overlap-scheduler.md), CUDA graphs ([piecewise](piecewise-cuda-graph.md), [padding](cuda-graph-padding.md)).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/relaxed-mtp-acceptance.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/relaxed-mtp-acceptance.md
new file mode 100644
index 000000000000..1866e116fab7
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/relaxed-mtp-acceptance.md
@@ -0,0 +1,37 @@
+---
+id: case-relaxed-mtp-acceptance
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [compute]
+signals: [low-accept-rate]
+architectures: [any-sm]
+model_scope: [model-agnostic, spec-decode, deepseek-v3r1]
+phase: [decode]
+patterns: [pattern-bounded-relaxed-acceptance]
+accuracy_risk: lossy
+apply_via_kind: [config-knob]
+knobs: [use_relaxed_acceptance_for_thinking, relaxed_topk, relaxed_delta]
+specialists: [perf-sweep-challenger]
+commits: ['b1621e8d4e11', '1e5e71aa4277']
+eligibility:
+  - "MTP spec-decode enabled on a reasoning model with a thinking phase (DeepSeek-R1); incompatible with attention_dp"
+interactions:
+  - {feature: attention-dp, relation: incompatible-with, note: cannot be enabled together}
+measured: []
+---
+
+# Raise MTP spec-decode acceptance with relaxed (top-N + delta) acceptance during the thinking phase
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `b1621e8d4e11` feat: add relaxed acceptance for DS (#3865); related: `1e5e71aa4277` Mtp optimizations round1 (#5689) — torch.compile fast-path for relaxed top-k.
+- **Applies when:** MTP spec decode enabled, acceptance rate (tokens accepted/step) low and capping speedup, model is a reasoning model with a thinking phase (DeepSeek-R1). NOT compatible with attention_dp (cannot be enabled together).
+- **Mechanism:** during `<think>...</think>`, instead of strict top-1 equality, accept a draft token if it lies in a candidate set: take target top-`relaxed_topk` logits, keep tokens with prob > (top-1 prob − `relaxed_delta`), accept the draft if in that set (CUDA kernel `mtpRelaxedAcceptanceKernel` / `torch.ops.trtllm.mtp_relaxed_acceptance_op`; per-request `mtpRelaxedDelta` set at think-start, 0 at think-end → greedy). More drafts accepted/step → higher tokens/step. #5689 wraps the top-k construction in `@torch.compile(max-autotune)` to cut overhead.
+- **Generalizes to:** "trade exact-match draft verification for a bounded relaxed acceptance criterion to lift accept rate where small token divergence is tolerable"; carries to Eagle/draft-target spec-decode and other reasoning models with a delimited low-stakes phase; adapt by choosing where relaxed acceptance is safe (here: thinking phase only) and tuning (topk, delta).
+- **Apply via:** `speculative_config` (MTPDecodingConfig): `use_relaxed_acceptance_for_thinking=True`, `relaxed_topk` (default 1; README ex 15), `relaxed_delta` (default 0.0; README ex 0.5). CLI `--use_relaxed_acceptance_for_thinking --relaxed_topk 15 --relaxed_delta 0.5`. Gate the accuracy trade with **perf-sweep-challenger**.
+- **Expected effect:** higher acceptance → higher tokens/step → "positive speedup" (README, no number) — measured Δ (accept rate, tokens/s) to be recorded from run.
+- **Accuracy risk:** lossy — changes which draft tokens are accepted (accepts non-top-1 tokens), so output is NOT token-equivalent to strict acceptance; bets this is acceptable inside the thinking phase only. Restricted to DeepSeek-R1. Needs accuracy record + rollback criterion.
+- **Verify:** task accuracy on R1 (PR adds e2e tests in `test_e2e.py`) AND accept-rate / tokens-per-step delta; compare strict vs relaxed output quality.
+- **Rollback:** `use_relaxed_acceptance_for_thinking=False` (strict top-1, default). Trigger: accuracy/quality regression, or need to combine with attention_dp.
+- **Prior art:** PRs #3865, #5689. Files: `cpp/.../kernels/speculativeDecoding/mtpKernels.cu`, `thop/mtpOp.cpp`, `_torch/speculative/mtp.py`, `llmapi/llm_args.py`, `examples/models/core/deepseek_v3/README.md`.
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/skip-sparse-path-when-degenerate.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/skip-sparse-path-when-degenerate.md
new file mode 100644
index 000000000000..4ba1d1fe172c
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/skip-sparse-path-when-degenerate.md
@@ -0,0 +1,90 @@
+---
+id: case-skip-sparse-path-when-degenerate
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [compute, launch]
+signals: [attention-hot-path, small-batch-decode, long-context]
+architectures: [any-sm]
+model_scope: [model-agnostic, sparse-attention, mla, deepseek-v32]
+phase: [any-phase]
+patterns: [pattern-skip-degenerate-sparse-path]
+accuracy_risk: lossless
+apply_via_kind: [config-knob, env-var, code-change]
+knobs: [skip_indexer_for_short_seqs, seq_len_threshold, TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD, q_split_threshold]
+specialists: [trtllm-serve-config-guide, perf-torch-cuda-graph-specialist]
+commits: ['8f144d9282', '695d7a0bdd', '6f3acc0614']
+interactions:
+  - {case: case-sparse-mla-topk-attention, relation: depends-on, note: length-gates the sparse top-k machinery that case implements}
+measured: []
+---
+
+# Skip the sparse/approximate path when it degenerates to the exact one
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `8f144d9282` [feat] Skip DS-v3.2 indexer MQA and Top-K for short
+  sequences (#9524); related: `695d7a0bdd` [perf] Short-sequence MHA optimization for
+  DSA MLA prefill (#11677), `6f3acc0614` [perf] Long-sequence token-parallel indexer
+  prefill (#11871, the opposite-regime move).
+- **Applies when:** a **sparse / approximate** code path has a regime where it produces
+  the **same result as the dense / exact** path but still pays the selection overhead.
+  The archetype: DSA's sparse top-k over a sequence **shorter than the top-k budget**
+  selects *every* token — so the lightning-indexer MQA + top-k are pure overhead, and
+  dense attention is both faster and identical. Signals: a selection/approximation step
+  whose selectivity → 1 below some size (short sequences, small expert counts, low
+  concurrency); the sparse path carries extra machinery (indexer, absorption BMMs,
+  larger head_dim) that dense attention avoids at that scale.
+- **Mechanism:** length-gate the path. Below a KV-length threshold, **skip the sparse
+  selection entirely** and either run dense attention (`forward_context` dense MHA:
+  `kv_b_proj` expansion + standard fused attention, `topk_indices=None`, #11677) or fill
+  the selection buffer with the **trivially-constructed exact result** (a dense causal
+  index pattern, `prepare_dense_topk_indices` / `_get_dense_topk_indices`, #9524). This
+  drops the whole indexer MQA-logits + top-k + absorption cost for short-sequence
+  batches. Because the fast path **changes control flow**, add a **CUDA-graph-key
+  dimension** (`short_seq_len_mode` in the graph key,
+  `needs_separate_short_long_cuda_graphs()`) so short-only and long/mixed batches each
+  keep a captured graph instead of falling back to eager.
+- **Generalizes to:** the pattern "**detect when a sparse/approximate path degenerates
+  to the dense/exact answer below a threshold and skip the selection work, substituting
+  a trivially-constructed exact result — and add a graph-key dimension so the
+  control-flow fast-path stays CUDA-graph-captured.**" Carries to: sparse/block-sparse
+  or sliding-window attention that is dense below a length; MoE that effectively routes
+  to all experts when `n_experts` is small; speculative decoding disabled where accept
+  rate is ~0; any approximate reduction exact below N. The complementary **long-sequence
+  regime** move travels with it: when the per-token selection work is *replicated* across
+  TP ranks, **shard it along the token axis and all-gather** (`q_split_threshold`, #11871)
+  — turning replicated work into divided work. Adapt by: finding the regime where
+  approximate == exact (or where a different parallelization wins), constructing the exact
+  result cheaply, and keeping every path graph-captured.
+- **Apply via:** config knobs — `skip_indexer_for_short_seqs` / `seq_len_threshold`, env
+  `TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD` (short-path), `q_split_threshold` (long-path
+  token-parallel); the graph-key split is code (`cuda_graph_runner.py`). Delegate the
+  YAML knobs to **trtllm-serve-config-guide** and the graph-key/path work to
+  **perf-torch-cuda-graph-specialist** / the model author.
+- **Expected effect:** short-sequence batches skip the indexer/top-k/absorption overhead
+  → lower latency at small scale with identical output; long-sequence prefill divides the
+  indexer work across TP ranks. Direction only — measured Δ (short-batch step latency,
+  long-prefill indexer time, throughput) to be recorded per regime.
+- **Accuracy risk:** **lossless** — for sequences ≤ top-k, dense attention (or a dense
+  causal index set) *is* the exact result the sparse path would compute (it would select
+  all valid tokens); the token-parallel split is an exact data-parallel decomposition
+  reconstructed by all-gather. The risk is a threshold/graph-key bug, not precision.
+- **Verify:** correctness — short-path output parity vs a standalone dense reference
+  (`test_short_seq_mha.py`) and vs the sparse path at the boundary length; long-path parity
+  across `q_split_threshold` values; confirm both regimes stay graph-captured (no eager
+  fallback). Perf — short-batch latency and long-prefill indexer time before vs after.
+- **Rollback:** `skip_indexer_for_short_seqs=False` / raise the threshold / negative
+  `q_split_threshold` to disable. Trigger: any boundary-length mismatch, an eager fallback
+  regressing captured batches, or no benefit at the target lengths.
+- **Prior art:** PRs #9524, #11677, #11871. Files:
+  `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` (`prepare_dense_topk_indices`,
+  `skip_indexer_for_short_seqs`, `q_split_threshold`),
+  `tensorrt_llm/_torch/modules/attention.py` (`_should_use_short_mha`,
+  `forward_context_dsa`), `tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`
+  (`short_seq_len_mode`). Owning skills: **trtllm-serve-config-guide**,
+  **perf-torch-cuda-graph-specialist**. Related: the
+  [sparse-MLA top-k attention](../kernel-and-fusion/sparse-mla-topk-attention.md) and
+  [top-k selection](../kernel-and-fusion/specialize-topk-selection-kernel.md) cases (the
+  machinery this skips), and [remove attention-DP padding](attention-dp-padding.md)
+  ("process only the real work" sibling).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/split-custom-op-for-piecewise-capture.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/split-custom-op-for-piecewise-capture.md
new file mode 100644
index 000000000000..4f4ffc327297
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/split-custom-op-for-piecewise-capture.md
@@ -0,0 +1,93 @@
+---
+id: case-split-custom-op-for-piecewise-capture
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch]
+signals: [many-small-kernels, slow-path-fallback]
+architectures: [any-sm]
+model_scope: [model-agnostic, sparse-attention, mla]
+phase: [decode]
+patterns: [pattern-split-op-at-graphable-seam]
+accuracy_risk: lossless
+apply_via_kind: [code-change]
+knobs: []
+specialists: [perf-torch-cuda-graph-specialist, kernel-cuda-specialist]
+commits: ['7e477ba8bf', 'b72ee4fd89']
+eligibility:
+  - "pays off only with torch_compile_enabled + torch_compile_piecewise_cuda_graph on"
+interactions:
+  - {case: case-piecewise-cuda-graph, relation: depends-on, note: "producer side of the same mechanism; pays off only with piecewise capture on"}
+measured: []
+---
+
+# Split a monolithic custom op at the graphable/non-graphable seam for CUDA-graph capture
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `7e477ba8bf` [perf] Split MLA DSA custom op for piecewise CUDA graph
+  capture (#12503); related: `b72ee4fd89` [fix] Route DSA attention through the MLA
+  custom op for torch.compile compatibility (#12186, the precondition).
+- **Applies when:** a single large fused **custom op** mixes graph-capturable and
+  non-capturable work, so piecewise CUDA graph must leave the *whole* op eager and you
+  lose capture coverage on the graphable part. Signals: launch-bound generation on the
+  `torch.compile` + piecewise-CUDA-graph path; one custom op that does both shape-static
+  token-wise math (projections, RoPE, quant) **and** batch/data-dependent work
+  (KV-cache scatter, top-k, length-adaptive branches); the op is a partition boundary so
+  its graphable half never gets captured. Prerequisite failure mode: the op is called
+  **directly in eager Python** (bypassing its registered `torch.ops` entry), so
+  `torch.compile` can't even trace it (#12186).
+- **Mechanism:** first make the op **traceable** — route it through its registered
+  `torch.library` custom op (`mla_custom_op_inplace`) instead of calling the impl
+  directly, so `torch.compile` sees one opaque node (#12186). Then **split it at the
+  seam** (#12503): hoist the **shape-static, metadata-free, token-wise** compute into a
+  capturable op (`trtllm::mla_dsa_proj` — `kv_a_proj`, layernorms, `q_b_proj`, the
+  indexer's `pre_indexer_proj`: cublas mm + RoPE + FP8 quant + weight scale; **no
+  batch metadata, no KV-cache access**), and leave the **batch/data-dependent** work as
+  a separate op **excluded** from capture (`trtllm::mla_dsa_attn_inplace` —
+  `sparse_attn_indexer`'s K-cache scatter + top-k + attention dispatch). Register the
+  attention op as a partition boundary in `piecewise_optimizer.py`. Two enabling tricks
+  keep the capturable half **unconditional**: a **constant-arity `register_fake`** (Op 1
+  always returns the same number of tensors) and **forcing straight-line control flow
+  under compile** (`_should_use_short_mha()` returns `False` when `is_torch_compiling()`,
+  so a length-adaptive branch doesn't fork the traced graph).
+- **Generalizes to:** the pattern "**split a monolithic custom op at the
+  graphable/non-graphable boundary — put the shape-static, metadata-free compute in its
+  own capturable op and leave the batch/data-dependent part as a separate excluded op —
+  and force straight-line control flow under compile (constant-arity `register_fake`,
+  disable length-adaptive branches) so the capturable half is unconditional.**" The
+  precondition — "**route work through its registered `torch.ops` custom op, never call
+  the impl directly in eager, or `torch.compile` can't trace it**" — is a reusable gotcha
+  on its own. Carries to: attention wrappers mixing projection + cache access, MoE
+  dispatch+compute+combine, any fused op with a data-dependent tail. This is the
+  **op-authoring (producer) side** of piecewise CUDA graph, complementary to
+  [turning piecewise capture on](piecewise-cuda-graph.md) (the consumer side). Adapt by:
+  locating the seam (which sub-steps touch batch metadata / the cache / data-dependent
+  branches), splitting into two registered ops, marking the data-dependent one a
+  boundary, and making the capturable op's signature/control-flow constant under compile.
+- **Apply via:** **not a server knob** — custom-op authoring; it pays off only with
+  `torch_compile_enabled` + `torch_compile_piecewise_cuda_graph` on (see the
+  [piecewise case](piecewise-cuda-graph.md)). Delegate to
+  **perf-torch-cuda-graph-specialist** + **kernel-cuda-specialist** / the model author.
+- **Expected effect:** the projection/RoPE/quant half gets captured into the CUDA graph
+  (fewer launches / lower per-step latency); only the data-dependent attention runs
+  eager. Direction only — measured Δ (graph coverage, per-step launch count, latency) to
+  be recorded from the run.
+- **Accuracy risk:** **lossless** — an op-boundary refactor; identical math, just split
+  across two registered ops. Verify the split didn't drop/reorder a side effect (the
+  in-place output arg must be marked in `inplace_info()`).
+- **Verify:** correctness — DSA e2e accuracy unchanged with piecewise CUDA graph +
+  torch.compile on (`test_nvfp4_multi_gpus_piecewise_cuda_graph`, GSM8K); outputs match
+  the eager path. Perf — confirm Op 1 is captured and Op 2 is the boundary (ad-conf-check
+  / nsys graph coverage); per-step launch count and latency before vs after.
+- **Rollback:** disable piecewise capture (`torch_compile_piecewise_cuda_graph=false`)
+  or fall back to the single-op eager path. Trigger: capture errors, an accuracy
+  mismatch, or the split adding net overhead.
+- **Prior art:** PRs #12503, #12186. Files:
+  `tensorrt_llm/_torch/modules/attention.py` (`trtllm::mla_dsa_proj` + `register_fake`,
+  `trtllm::mla_dsa_attn_inplace`, `forward_dsa_proj` / `forward_dsa_attn`,
+  `_should_use_short_mha`), `tensorrt_llm/_torch/compilation/piecewise_optimizer.py`,
+  `tensorrt_llm/_torch/compilation/utils.py` (`inplace_info`). Owning specialist:
+  **perf-torch-cuda-graph-specialist**. Related: the
+  [piecewise CUDA graph case](piecewise-cuda-graph.md) (the consumer side of the same
+  mechanism).
diff --git a/.claude/skills/perf-optimization-casebook/references/runtime-execution/two-model-mtp-eagle.md b/.claude/skills/perf-optimization-casebook/references/runtime-execution/two-model-mtp-eagle.md
new file mode 100644
index 000000000000..2c76cbb159bc
--- /dev/null
+++ b/.claude/skills/perf-optimization-casebook/references/runtime-execution/two-model-mtp-eagle.md
@@ -0,0 +1,35 @@
+---
+id: case-two-model-mtp-eagle
+type: case
+family: runtime-execution
+maturity: full
+bottleneck: [launch, sync]
+signals: [draft-forward-exposed, gpu-idle-between-steps]
+architectures: [any-sm]
+model_scope: [model-agnostic, spec-decode, deepseek-v3r1]
+phase: [decode]
+patterns: [pattern-overlap-independent-work]
+accuracy_risk: lossless
+apply_via_kind: [config-knob]
+knobs: [mtp_eagle_one_model]
+specialists: []
+commits: ['80dd8fe19733', '6151a4c9d600']
+interactions:
+  - {case: case-overlap-scheduler, relation: composes-with, note: two-model path re-enables the overlap scheduler on the draft forward}
+measured: []
+---
+
+# Run MTP-Eagle as two-model speculative decoding with overlap scheduler on the draft forward
+
+> Part of the [Runtime / Execution casebook](index.md) · schema: [case-template](../case-template.md)
+
+- **Commits:** `80dd8fe19733` Enable two-model spec dec for MTP Eagle (#7001); related: `6151a4c9d600` Add simple optimizations for MTP 2-model (#9176) — torch.compile + re-enable overlap for mtp_eagle.
+- **Applies when:** running MTP-Eagle spec decode where the single-engine (one-model) path constrains, and the draft (MTP) forward is exposed on the critical path; you want the draft forward to overlap target-side work.
+- **Mechanism:** adds a two-model MTP-Eagle mode (`SpeculativeDecodingMode.MTP_EAGLE` vs one-model `MTP_EAGLE_ONE_MODEL`), routing the DeepSeek-V3 draft as a separate `MTPDraftModelForCausalLM` (selected in `modeling_auto.py` when `max_draft_len==0`) with its own `DeepseekV3WeightLoader(is_draft_model=...)`. The two-model path participates in the overlap scheduler (`support_overlap_scheduler()` True). #9176 removes the temporary `is_mtp_eagle()->return False` block (re-enabling overlap after an accuracy fix) and adds `@torch.compile(max-autotune)` to `prepare_for_generation` in `drafting_loops.py`.
+- **Generalizes to:** "split a fused single-engine speculator into target+draft models so the draft forward can overlap target work, and torch.compile the per-step draft-prep glue"; carries to Eagle3 two-model, draft-target spec-decode, any two-engine drafting loop; adapt by wiring a separate draft model + weight loader and ensuring the spec mode reports `support_overlap_scheduler()==True`.
+- **Apply via:** `speculative_config` knob `mtp_eagle_one_model` (default True). Set `mtp_eagle_one_model=False` (`--use_one_model` controls it: not-one-model ⇒ two-model) to get the two-model path; overlap scheduler then used automatically.
+- **Expected effect:** draft forward overlaps target work → lower exposed spec-decode latency / higher throughput; no number — measured Δ to be recorded from run.
+- **Accuracy risk:** lossless in intent — execution restructuring, not a change to the acceptance rule (output should match one-model). Note: overlap for mtp_eagle was once disabled for an accuracy issue and re-enabled in #9176 — re-verify accuracy.
+- **Verify:** accuracy parity two-model vs one-model MTP-Eagle (e2e tests in `test_e2e.py`); confirm draft/target overlap + reduced GPU idle in nsys.
+- **Rollback:** `mtp_eagle_one_model=True` (`--use_one_model`) → single-engine MTP-Eagle (no overlap). Trigger: accuracy regression or scheduler instability.
+- **Prior art:** PRs #7001, #9176. Files: `_torch/speculative/interface.py` (`support_overlap_scheduler`), `speculative/drafting_loops.py`, `models/modeling_auto.py` (`MTPDraftModelForCausalLM`), `modeling_deepseekv3.py`, `llmapi/llm_args.py`.
diff --git a/.claude/skills/perf-optimization/SKILL.md b/.claude/skills/perf-optimization/SKILL.md
index 7c74affefa98..5afdc71aea66 100644
--- a/.claude/skills/perf-optimization/SKILL.md
+++ b/.claude/skills/perf-optimization/SKILL.md
@@ -29,26 +29,25 @@ You coordinate with five specialists:
   Does NOT write kernels from scratch -- receives them from kernel-triton-specialist or the user.
 - **kernel-cute-specialist**: CuTe DSL kernels (GEMM, attention, element-wise, reduction)
 
+## Consult Prior Art First
+
+Before prioritizing or routing an optimization, consult the
+**perf-optimization-casebook** skill for matching decision precedents (past
+successful and classic optimizations) for the classified bottleneck, model,
+and hardware. Adapt a proven case — its applicability signals, the right
+specialist, accuracy risk, and rollback — instead of inventing an approach.
+The casebook is reference material consulted as a skill; it does NOT replace
+the specialists below (who still own implementation).
+
 ## Delegation Rules
 
 - For actual implementation and validation, delegate to specialists.
 - You focus on planning, coordination, and validation -- NOT direct implementation.
 - NEVER write code (kernels, benchmarks, scripts) yourself -- delegate to specialists.
 - Include benchmarking in the specialist's task scope (e.g., "Write and benchmark a TileIR kernel").
-- NEVER explore or browse skill directories directly.
 - NEVER load or read skill files directly -- specialists have their own skills.
 - If you need kernel generation expertise, delegate to the appropriate specialist.
 
-**Task-to-specialist mapping:** Double-check that each delegation targets
-the CORRECT specialist for that task's domain:
-
-- CuTe DSL tasks --> Delegate to **kernel-cute-specialist** (NOT kernel-triton-specialist)
-- Triton kernel tasks --> Delegate to **kernel-triton-specialist** (NOT kernel-cute-specialist)
-- TileIR optimization --> Delegate to **kernel-tileir-specialist**
-
-Never send a CuTe DSL task to kernel-triton-specialist or vice versa. The specialist
-in each delegation must match the task domain.
-
 ### Iterative Optimization Loops
 
 When iterating toward a performance goal (optimize → profile → repeat):
@@ -62,25 +61,9 @@ When iterating toward a performance goal (optimize → profile → repeat):
 
 You are the loop controller, not the implementer. Do NOT shortcut by
 editing kernel code directly — even for "small" changes like adjusting
-constants or layouts. The specialist owns the code, handles verification,
+constants or layouts. The specialist owns the code and handles verification
 for kernels it modifies.
 
-### Remote Execution
-
-When optimizing on a remote SLURM cluster, include the
-**Remote Execution Context** block (with the SSH+srun wrapper for the target cluster) in every
-specialist delegation. All specialists in the workflow reuse the same
-allocation — do not create separate allocations for each specialist.
-
-For multi-specialist pipelines (e.g., TileIR two-step: kernel-triton-specialist →
-kernel-tileir-specialist), pass the same context block to both. Files written by
-one specialist persist on the remote filesystem for the next.
-
-**Integration code rule:** If you must write integration code (e.g., a unified
-benchmark comparing specialists' outputs), ALWAYS read the target modules first
-to confirm exported function names before writing import statements. Never guess
-export names from file names.
-
 ## Terminology -- Do NOT Confuse
 
 - **TileIR** = NVIDIA's Triton backend (nvtriton) for Blackwell GPUs --> use kernel-tileir-specialist
@@ -106,28 +89,6 @@ Example: "Apply CUDA Graph to my model"
 - Delegate to **perf-torch-cuda-graph-specialist**: "Apply CUDA Graph capture to the training loop"
 - Delegate to **perf-profiling-specialist**: "Measure performance before and after"
 
-### Autopilot Mode (Goal-Driven)
-
-When called by the Orchestrator with analysis results:
-
-1. **Review analysis**: Parse bottleneck classification and recommendations
-2. **Prioritize**: Rank optimizations by expected impact / effort
-3. **Plan**: Determine implementation order
-4. **Implement**: One optimization at a time with validation between each
-5. **Rollback**: If regression detected, revert and try next optimization
-6. **Report**: Return optimization result with before/after metrics
-
-You receive analysis data in this format:
-
-```
-Primary bottleneck: memory-bound
-Evidence: Memory bandwidth at 89% of peak, compute at 35%
-Recommendations:
-1. [High] Enable FlashAttention for self-attention layers
-2. [Medium] Apply memory pooling for attention buffers
-3. [Low] Consider gradient checkpointing for memory reduction
-```
-
 ## Optimization Workflow
 
 ### Planning Phase
diff --git a/.claude/skills/trtllm-case-executor/SKILL.md b/.claude/skills/trtllm-case-executor/SKILL.md
new file mode 100644
index 000000000000..a8ee86e9c177
--- /dev/null
+++ b/.claude/skills/trtllm-case-executor/SKILL.md
@@ -0,0 +1,440 @@
+---
+name: trtllm-case-executor
+description: >
+  Run TensorRT-LLM test cases, benchmarks, evaluations, or custom scripts by
+  checking the environment (local GPU or Slurm), selecting the appropriate
+  Docker image, and executing either locally or via Slurm job submission.
+  Accepts pre-built command strings — command construction for trtllm-bench,
+  trtllm-eval, and test_perf_sanity is handled upstream by the caller
+  (e.g. trtllm-test-specialist using build_test_command.py).
+tags:
+  - tensorrt-llm
+  - testing
+  - execution
+  - slurm
+license: Apache-2.0
+metadata:
+  author: NVIDIA Corporation
+---
+
+This skill executes TensorRT-LLM test commands by checking the environment and
+dispatching to the appropriate executor. All commands are accepted as pre-built
+strings — this skill does **not** build or modify commands.
+
+**Four workflows** (classified in this order — first match wins):
+
+1. **Custom** — `custom_cmd` or `custom_script` provided. Runs arbitrary commands via `slurm_run_custom.sh`.
+2. **Bench** — `bench_cmd` provided. Accepts a pre-built `trtllm-bench` command chain. Parses `--tp`/`--pp` for `required_devices` only; does not modify the command.
+3. **Eval** — `eval_cmd` provided. Accepts a pre-built `trtllm-eval` command. Parses `--tp_size`/`--pp_size` for `required_devices` only; does not modify the command.
+4. **Pytest / test_cmd** — `test_cmd` provided. Handles all other commands: non-perf pytest, perf-sanity (`test_e2e[...]`), python scripts, and anything else. When `perf_config_yaml` is provided alongside a perf-sanity `test_cmd`, it is passed directly to the script builder for node/GPU sizing; no auto-derivation from the test ID is performed.
+
+---
+
+## Inputs
+
+### Core Inputs
+
+| Input | Description | Required | Example |
+|-------|-------------|----------|---------|
+| `test_cmd` | Test command to run inside the container (pytest, python, perf_sanity, etc.) | **Yes** (pytest/test_cmd) | `pytest tests/integration/defs/perf/test_perf_sanity.py -v "test_e2e[aggr-gpt_oss-gpt_oss_tp4ep4pp2_mtp1]"` |
+| `bench_cmd` | Pre-built `trtllm-bench` command string. May chain `prepare-dataset` and `throughput`/`latency` with `&&`. Mutually exclusive with `test_cmd`/`eval_cmd`/`custom_cmd`. | **Yes** (bench) | `trtllm-bench ... prepare-dataset ... && trtllm-bench ... throughput ...` |
+| `eval_cmd` | Pre-built `trtllm-eval` command string. Mutually exclusive with `test_cmd`/`bench_cmd`/`custom_cmd`. | **Yes** (eval) | `trtllm-eval --model /models/llama --tp_size 4 gsm8k` |
+| `custom_cmd` | Arbitrary command to run inside the container. Mutually exclusive with `test_cmd`. | **Yes** (custom) | `python3 examples/run_inference.py --model llama` |
+| `custom_script` | Host script path — takes precedence over `custom_cmd` (becomes `bash <script>`). Auto-mounts parent dir. | No | `/home/user/my_benchmark.sh` |
+| `custom_env` | Space-separated `KEY=VALUE` pairs for custom workflow only | No | `BATCH_SIZE=32 NUM_WORKERS=4` |
+| `perf_config_yaml` | Explicit path to a perf-sanity config YAML. When provided, passed directly to the script builder for node/GPU sizing — no auto-derivation from the test ID. | No (perf_sanity) | `/work/gpt_oss_120b_fp4.yaml` |
+| `model_name` | Short name for container/job naming. Auto-derived from command if not provided. | No | `llama_7b` |
+| `required_devices` | GPU count. Auto-parsed from `--tp`/`--pp` (bench) or `--tp_size`/`--pp_size` (eval) if not provided. Default: `1`. | No | `8` |
+| `total_required_devices` | Total GPU processes across all nodes. If not provided, derived per the rules below. | No | `8` |
+| `required_devices_per_node` | GPU processes per node. If not provided, derived per the rules below. | No | `4` |
+| `skip_install` | Skip `pip install -e .` when image already has TRT-LLM. Default: `false`. | No | `true` |
+| `build_project` | Build TRT-LLM from source before running the test. Default: `true`. Set to `false` to skip the build step when the image already contains a current build. | No | `false` |
+| `container_image` | Override container image URI. If unset, case-executor resolves from `<REPO_ROOT>/jenkins/current_image_tags.properties` using the CPU arch derived from `device_type` (or local env-check). The resolved value is forwarded to every downstream consumer in `job_spec.json` — none of them re-resolve. For an example URI shape, delegate to `internal-env-info` for the `container_registry_example` key when that skill is installed; otherwise ask the user for a sample. | No | (delegated) |
+| `monitor_timeout` | Wall-clock limit. Accepts `1h`, `30m`, `3600s`, or `HH:MM:SS`. Default: `1h`. | No | `2h` |
+
+### Environment & Slurm Inputs
+
+| Input | Description | Default |
+|-------|-------------|---------|
+| `node_count` | Slurm nodes. Computed by this skill as `ceil(total_required_devices / required_devices_per_node)` and forwarded in `job_spec.json` to every downstream consumer (`trtllm-test-script-builder`, `exec-remote-slurm`). Caller may override only when an explicit value is required. | computed |
+| `partition` | Slurm partition | *(required for Slurm)* |
+| `account` | Slurm account | resolution chain (applied in [Step 2.5](#step-25-detect-slurm-account-and-valid-partitions-slurm-scenarios-only)): caller input → `default_account` from `detect_slurm_env.sh` JSON → delegate to `internal-env-info` for the `default_account` key when that skill is installed → otherwise ask the user. No NVIDIA-internal default is hard-coded in this skill. |
+| `time_limit` | Slurm walltime | `02:00:00` |
+| `device_type` | Required GPU type (e.g., `B200`, `GB200`). Auto-routes to remote cluster if local doesn't match. | — |
+| `checkpoint_path` | Host model checkpoint path | — |
+| `llm_models_root` | Models root inside container | caller input → delegate to `internal-env-info` for the `default_llm_models_root` key when that skill is installed → otherwise ask the user. No NVIDIA-internal default is hard-coded in this skill. |
+
+### Remote Slurm Inputs
+
+| Input | Description | Default |
+|-------|-------------|---------|
+| `slurm_cluster` | Remote cluster name — skips local env-check, looks up per-cluster info (`mfa_style`, `default_models_repo`, `default_user_root_dir`, `gpus_per_node`) via `internal-env-info` skill. Connection fields (`ssh_host`, `partition`, `account`, `mounts`, etc.) come from caller-supplied inputs in `job_spec.json` (or `internal-env-info` defaults / ask-the-user fallbacks). When `internal-env-info` is not installed, the executor uses whatever explicit cluster fields the caller passed in `job_spec.json`. | — |
+| `ssh_host` | Full SSH destination override. When unset, `exec-remote-slurm` Step 2.0 delegates to `internal-env-info` for the `ssh_host_pattern` key to construct the destination from `<slurm_user>` and `cluster_name`; it delegates to the same skill for `ssh_host_example` when documentation needs a sample. When `internal-env-info` is not installed and the user did not provide `ssh_host`, **ask the user** to supply `ssh_host` directly — this skill does not hard-code an NVIDIA-internal login-host pattern. `mfa_style` itself comes from `internal-env-info` Step 1b and is used by `exec-remote-slurm` Step 2 only to pick the SSH path (direct vs MFA), not to build the username; when that skill is absent, `mfa_style` defaults to `null` and `exec-remote-slurm` uses its probe-direct-first-then-MFA fall-back path. | constructed (delegated) |
+| `user_name` | Raw login username forwarded to `internal-env-info` so it can build the SSH-ready `slurm_user` (applying whatever suffix that cluster's `mfa_style` requires). This skill never builds `slurm_user` itself — it only forwards `user_name` and reads the result back. | `$(whoami)` |
+| `slurm_user` | SSH/Slurm username override. When set, used verbatim and the cluster-info value is discarded. When unset, this skill passes `user_name` to `internal-env-info` and copies the returned per-cluster `slurm_user` into `slurm_env` as-is. | (from cluster-info) |
+| `slurm_password` | SSH password (**prompted if not provided** — never stored/logged) | *(prompted)* |
+| `remote_cwd` | Remote working directory | `<REPO_ROOT>` |
+
+### Device Count Auto-Parsing
+
+If the caller passes `total_required_devices` and `required_devices_per_node` directly, use those values verbatim — skip all derivation below.
+
+Otherwise derive all three fields (`required_devices`, `total_required_devices`, `required_devices_per_node`) using the rules in the table. Derivation order: `required_devices` first, then the two per-node/total fields.
+
+| Workflow | `required_devices` | `total_required_devices` | `required_devices_per_node` |
+|----------|--------------------|--------------------------|----------------------------|
+| `bench_cmd` | `tp * pp` parsed from last `trtllm-bench` invocation | same as `required_devices` | same as `required_devices` (single-node only) |
+| `eval_cmd` | `tp_size * pp_size` parsed from `eval_cmd` | same as `required_devices` | same as `required_devices` (single-node only) |
+| `test_cmd` (non-perf pytest) | From `skip_less_device(N)` marker in test file; otherwise `1` | same as `required_devices` | same as `required_devices` (always single-node) |
+| `test_cmd` (perf-sanity) with `perf_config_yaml` | `tensor_parallel_size * pipeline_parallel_size` for the matching `server_configs` entry | same as `required_devices` | `hardware.gpus_per_node` from the YAML |
+| `test_cmd` (perf-sanity) without `perf_config_yaml` | default `1` | `1` | `1` |
+| `custom_cmd` | `required_devices` input; default `1` | same as `required_devices` | same as `required_devices` |
+
+**Perf-sanity YAML extraction procedure** (applies when `perf_config_yaml` is provided):
+
+1. Read the YAML file and extract `hardware.gpus_per_node` → `required_devices_per_node`.
+2. Identify the target server config name from `test_cmd` — it is the bracketed test ID suffix, e.g., `aggr-gpt_oss-gpt_oss_tp4ep4pp2_mtp1` in `test_e2e[aggr-gpt_oss-gpt_oss_tp4ep4pp2_mtp1]`. Match the trailing component against `server_configs[].name` in the YAML.
+3. From the matched entry, read `tensor_parallel_size` (default `1`) and `pipeline_parallel_size` (default `1`).
+4. `required_devices = tensor_parallel_size * pipeline_parallel_size`.
+5. `total_required_devices = required_devices` (total GPU processes = world size).
+6. If no matching entry is found, fall back to `required_devices = 1`, `total_required_devices = 1`; keep `required_devices_per_node` from step 1.
+
+### Models Mount Resolution
+
+Mount to `<llm_models_root>` (the resolved input — see the `llm_models_root` row above). Host-side source priority: `checkpoint_path` parent → `env_check.default_models_repo` → delegate to `internal-env-info` for the `default_models_repo_host` key when that skill is installed → otherwise stop and ask the user. No NVIDIA-specific path is hard-coded in this skill.
+
+---
+
+## Procedure
+
+### Step 1: Create Work Directory and Classify
+
+```bash
+WORK_DIR="<REPO_ROOT>/work_dirs/<model_name>_$(date +%Y%m%d_%H%M%S)"
+mkdir -p "$WORK_DIR"
+```
+
+Classify workflow (custom → bench → eval → test_cmd). Derive `model_name`, `required_devices`, `total_required_devices`, and `required_devices_per_node` per the [Device Count Auto-Parsing](#device-count-auto-parsing) rules above. If the caller already passed `total_required_devices` and `required_devices_per_node`, skip derivation and use those values directly.
+
+Then compute `node_count = ceil(total_required_devices / required_devices_per_node)` (use the caller's value only when explicitly provided). This skill is the **single source of truth for `node_count`** — every downstream consumer reads it from `job_spec.json` and must not recompute. Floor-clamp to `1` when both totals are `1`.
+
+Build the resolved `job_name` here too:
+
+- `detail = workflow_type` (`pytest`, `bench`, `eval`, `custom`) — or a more specific token derived from `model_name` when available.
+- `job_name = f"{account}.{detail}"`.
+
+This skill is the **single source of truth for `job_name`** — `exec-remote-slurm` (and any other downstream consumer) takes the resolved string from `job_spec.json` and substitutes it directly into `-J <job_name>`. The convention `<account>.<detail>` is documented and applied here only.
+
+**Remote work_dir derivation (remote_slurm scenarios only):** Compute the remote-side equivalents of `WORK_DIR` so the remote layout mirrors the local layout (`<REPO_ROOT>/work_dirs/<basename>`):
+
+```text
+repo_dir_name      = basename(repo_url stripped of .git)         # e.g., "TensorRT-LLM"
+remote_repo_path   = <default_user_root_dir>/<repo_dir_name>
+remote_work_dir    = <remote_repo_path>/work_dirs/<basename(WORK_DIR)>
+```
+
+`default_user_root_dir` comes from the cluster config (resolved by `internal-env-info` cluster-info modes with `<user_name>` substituted). When that skill is not installed, fall back in this order: caller-supplied `remote_cwd` or explicit `remote_repo_path` in `job_spec.json` → otherwise stop and ask the user. The `default_user_root_dir_example` key under `internal-env-info`'s default-values section is documentation-only — usable only if the cluster actually matches that layout. Do not error on a missing optional skill; just route to the next source. This skill is the **single source of truth for `remote_work_dir` and `remote_repo_path`** — both are written verbatim into `job_spec.json` and every downstream consumer (`trtllm-test-script-builder`, `exec-remote-slurm`) must read them from the manifest, never recompute.
+
+### Step 2: Detect Environment
+
+This skill **never parses cluster connection YAML**. It only decides local vs remote and forwards inputs downstream. YAML resolution is owned by `exec-remote-slurm`.
+
+**Optional dependency note.** If `skills/internal-env-info/` is not installed in the toolkit, the routing below still works using only env-check output and the caller's explicit inputs. Downstream executors handle the missing skill gracefully (per their own optional-dependency rubrics). Do **not** error out here when the skill is absent.
+
+1. Run `trtllm-agent-toolkit:exec-env-check` with `required_devices`. Returns: `scenario`, `available_gpus`, `device_type`, `gpus_per_node`, `default_models_repo`, `cluster_name` (the name of the cluster this skill is running on, or `null` if not on a known cluster).
+2. If `slurm_cluster` is provided, compare it to `env_check.cluster_name`:
+   - **Match** (same cluster, case-insensitive) → `local_slurm`. Continue with the local env-check output; do not call `internal-env-info`.
+   - **Mismatch** (or `env_check.cluster_name` is `null`) → `remote_slurm` with explicit cluster. Forward `slurm_cluster` (the name only) to `exec-remote-slurm`; that skill resolves per-cluster info via `internal-env-info` and reads connection fields from `job_spec.json`.
+3. Otherwise (no `slurm_cluster`): if the `device_type` input doesn't match `env_check.device_type`, route to `remote_slurm` with **auto-select**. Forward `device_type`, `total_required_devices`, and `required_devices_per_node` to `exec-remote-slurm`, which queries `internal-env-info` in catalog mode and picks the cluster.
+4. Otherwise map env-check to local scenarios: `satisfied, local, docker` → `local_docker` | `satisfied, local, direct` → `local_direct` | `satisfied, slurm, local` → `local_slurm` | `not_satisfied` → fall through to step 3's auto-select path.
+5. Prompt for `slurm_password` if not provided (any `remote_slurm` path).
+
+### Step 2.5: Detect SLURM Account and Valid Partitions (slurm scenarios only)
+
+**MANDATORY before Step 4.** For every `local_slurm` and `remote_slurm` run, this step
+**must complete successfully** before invoking `trtllm-test-script-builder`. The script
+builder relies on `account`, `partition`, and the validated `slurm_env` already being
+present in `job_spec.json` — it must not be invoked with caller-supplied / default
+values that haven't been verified against the cluster.
+
+**Do NOT skip Step 2.5 just because case-executor is running on a host that cannot
+SSH directly**. In that case use **Option A (delegation)** below
+— dispatch a `exec-remote-slurm` preflight subagent that performs the SSH
+and runs the script over it. Skipping detection and falling back to defaults is
+forbidden — defaults can produce a script that submits to a non-existent partition
+or an account the user has no association with, which only fails after the SLURM
+queue accepts it.
+
+For `local_slurm` and `remote_slurm`, run `scripts/detect_slurm_env.sh` to populate `account`, the list of valid partitions, per-partition hardware (arch / GRES / gpus_per_node), and the cluster's PMIx plugin in **one pass**. **Skip this step entirely for `local_docker` and `local_direct`** — they don't use SLURM.
+
+The script (shipped with this skill) wraps `sacctmgr`, `sinfo`, `scontrol show node`, and `srun --mpi=list` and emits a single JSON document of the form:
+
+```json
+{
+  "user": "<whoami>",
+  "default_account": "<first-account-or-null>",
+  "accounts": ["acct1", "acct2", "..."],
+  "default_partition": "<partition-marked-default-or-null>",
+  "partitions": [
+    {
+      "name": "batch",
+      "state": "up",
+      "time_limit": "08:00:00",
+      "nodes": "32",
+      "arch": "x86_64",
+      "gres": "gpu:8",
+      "gpus_per_node": 8,
+      "requires_gres": true
+    },
+    ...
+  ],
+  "pmix": {"available": ["pmix", "pmix_v3", "pmix_v5"], "preferred": "pmix_v5"},
+  "errors": ["..."]
+}
+```
+
+`arch`, `gres`, `gpus_per_node`, `requires_gres`, and the top-level `pmix` block are consumed downstream by `trtllm-test-script-builder` (Docker image arch selection, GRES handling, `--mpi=` flag) and the slurm executors (compile-step arch lookup). They must **not** be re-detected downstream.
+
+#### `local_slurm`: run the script directly
+
+Execute on the host where case-executor is running (the login or compute node):
+
+```bash
+SLURM_ENV_JSON=$(bash "<SKILL_DIR>/scripts/detect_slurm_env.sh" --format=json)
+echo "$SLURM_ENV_JSON" > "<WORK_DIR>/slurm_env.json"
+```
+
+#### `remote_slurm`: reuse `exec-remote-slurm` to run via SSH
+
+Reuse the SSH connection logic that `exec-remote-slurm` already owns — do **not** re-implement SSH host construction or preflight here. Two equivalent options:
+
+**Option A (preferred — delegation):** Dispatch a small `exec-remote-slurm` subagent in "preflight + one-off command" mode that:
+1. Resolves the cluster (its Step 1) — produces `cluster_name`, `mfa_style`, `slurm_user`, etc.
+2. Constructs `<ssh_host>` per its Step 2.0 by delegating to `internal-env-info` for the `ssh_host_pattern` key (the single-line snippet that concatenates the already-suffixed `slurm_user` with `cluster_name` into `SSH_HOST`). When the skill is not installed and the user did not provide `ssh_host`, **ask the user** for it — this skill does not hard-code an NVIDIA-internal login-host pattern.
+3. Runs SSH preflight per its Step 2a/2b/2c keyed on `mfa_style`.
+4. Streams the local script over SSH and captures the JSON output:
+   ```bash
+   <ssh_cmd> <ssh_host> "bash -s -- --format=json" < "<SKILL_DIR>/scripts/detect_slurm_env.sh"
+   ```
+5. Returns the captured stdout to case-executor.
+
+**Option B (inline — when an SSH session is already established by an earlier exec-remote-slurm invocation in the same conversation):** Just run the script over the existing `<ssh_cmd>` socket:
+```bash
+SLURM_ENV_JSON=$(<ssh_cmd> <ssh_host> "bash -s -- --format=json" < "<SKILL_DIR>/scripts/detect_slurm_env.sh")
+```
+
+In either case, save the JSON to `<WORK_DIR>/slurm_env.json`.
+
+#### Resolve `account` and `partition` from the JSON
+
+Apply this priority order (do not silently fall back past the user input):
+
+| Field | Resolution |
+|-------|------------|
+| `account` | caller input → `default_account` from JSON → delegate to `internal-env-info` for the `default_account` key when that skill is installed → otherwise **stop and ask the user**. Do not hard-code an NVIDIA-internal account in this skill. |
+| `partition` | caller input → `default_partition` from JSON → **stop and ask the user** if both are absent |
+
+**Validation:** If the caller provided `partition` and the JSON's `partitions[].name` list does not include it, surface the mismatch (list available names) and ask the user before proceeding. Same for `account` vs `accounts[]`.
+
+The resolved `account` and `partition` are written into `job_spec.json` so `trtllm-test-script-builder` and the executors can read them without re-detecting. Also include the full parsed JSON as the `slurm_env` object in `job_spec.json` — downstream consumers look up per-partition `arch` / `gres` / `gpus_per_node` / `requires_gres` (e.g. `slurm_env.partitions[name=<partition>].arch`) and the cluster's `slurm_env.pmix.preferred` from there. The on-disk `<WORK_DIR>/slurm_env.json` is the canonical artifact; `job_spec.slurm_env` is the parsed convenience copy.
+
+If the script exits non-zero, surface the `errors[]` array to the user before continuing — don't silently fall back to defaults when SLURM detection itself failed. If a specific failure mode is "SLURM client tools not found" (exit 1), forward the empty `slurm_env` so downstream skills know to consult `internal-env-info` for arch and skip pmix.
+
+#### Merge cluster YAML policy fields into `slurm_env`
+
+After the detection JSON is on disk, enrich it with the cluster's site-policy fields so that `slurm_env` is the **single object** that downstream skills (especially `trtllm-test-script-builder`) consume.
+
+Source of the policy fields:
+
+- **Known cluster (env-check returned a non-null `cluster_name`, or the caller supplied `slurm_cluster`):** invoke `trtllm-agent-toolkit:internal-env-info` in single-cluster mode, passing the `user_name` input (default `$(whoami)`). The selected cluster entry returns `cluster_name`, `mfa_style`, and the already-suffixed `slurm_user` (the skill applies its `slurm_user` construction rule against the cluster's own `mfa_style`). **Copy `slurm_user` into `slurm_env` verbatim — this skill never builds, suffixes, or re-derives `slurm_user` itself.** If the caller passed an explicit `slurm_user` in `job_spec.json`, that override wins and the cluster-info value is discarded; otherwise the cluster-info value is the only source. For the remaining policy fields (`mounts`, `ssh_host`, `remote_cwd`), follow the same fallback chain documented above for `default_user_root_dir`: caller-supplied input in `job_spec.json` → delegate to `internal-env-info` when that skill is installed → otherwise stop and ask the user. There is no per-cluster connection-config file to parse. `container_image` is owned by Step 3.
+- **Unknown cluster (no match in `internal-env-info`, or the skill is not installed):** leave the policy fields absent / `null`. Downstream consumers must handle the absence (e.g., script builder falls back to caller-supplied `checkpoint_path` / `models_path` for mounts).
+
+Merge those fields into the same JSON document and rewrite `<WORK_DIR>/slurm_env.json`. The post-merge shape is:
+
+```json
+{
+  "user": "<whoami>",
+  "default_account": "...", "accounts": [...],
+  "default_partition": "...", "partitions": [ { "name": "...", "arch": "...", "gres": "...", "gpus_per_node": <int>, "requires_gres": <bool>, ... } ],
+  "pmix": {"available": [...], "preferred": "..."},
+  "errors": [...],
+
+  "cluster_name": "<cluster>|null",
+  "mounts": ["host:container", "..."]      | null,
+  "ssh_host": "<remote_slurm only>"        | null,
+  "slurm_user": "<remote_slurm only; copied verbatim from internal-env-info (or caller override); never built here>" | null,
+  "remote_cwd": "<remote_slurm only>"      | null,
+  "mfa_style": <true|false|null>
+}
+```
+
+Refresh `job_spec.slurm_env` to the merged object. 
+
+### Step 3: Resolve Container Image and Forward Build Inputs
+
+Case-executor does **not** build, but it **does** resolve the container image once so that the test-script-builder, the run-time executor, and (for `remote_slurm`) the remote build job all use the same image.
+
+**Resolution rules (in priority order):**
+
+1. If the caller provided `container_image` as input → use it verbatim.
+2. Else read `<REPO_ROOT>/jenkins/current_image_tags.properties` and select by CPU arch:
+   - `aarch64` (when `device_type` starts with `GB` / `GH`, or local env-check reports an `aarch64` host) → `LLM_SBSA_DOCKER_IMAGE`.
+   - otherwise → `LLM_DOCKER_IMAGE` (x86_64).
+
+Store the resolved URI verbatim in `job_spec.json` as `container_image`. Transport-specific URL rewrites happen at use-time inside the consumer — when an enroot-based executor needs the rewrite, it delegates to `internal-env-info` for the `enroot_uri_rewrite` key (which provides the canonical recipe and example). Case-executor itself always stores the canonical form (no `#`).
+
+**Forwarded fields in `job_spec.json`:**
+
+- `container_image` — resolved image URI (canonical form, no enroot rewrite).
+- `build_project` (default `true`) — executors skip their build step when `false`.
+- `device_type` — for `local_*` from env-check; for `remote_slurm` from the resolved cluster.
+- `repo_root`.
+
+Each executor's build step (and the remote build job) reads `container_image` directly from `job_spec.json` and applies any transport-specific transform on the way to `--container-image=`.
+
+If an executor reports `BUILD_FAILED`, stop and surface the failure to the user without dispatching the workload.
+
+The `local_direct` scenario has no executor SKILL — its build is inlined inside the `local_direct` block in Step 5, and it does not use a container.
+
+### Step 4: Build Script
+
+> **Precondition (slurm scenarios only): Step 2.5 must have completed.** Before invoking the
+> script builder for `local_slurm` / `remote_slurm`, verify that `<WORK_DIR>/slurm_env.json`
+> exists and that `account` + `partition` were resolved (and validated against the cluster's
+> `accounts[]` / `partitions[].name`). If `slurm_env.json` is missing or `errors[]` from the
+> detection script is non-empty and unresolved, **stop and go back to Step 2.5** — do not
+> invoke the script builder with unvalidated SLURM parameters.
+
+Invoke `trtllm-test-script-builder` skill with all inputs + classified workflow type + execution scenario. Pass `work_dir`, `repo_root`, `skill_dir`, `perf_config_yaml` (if provided), `container_image` (resolved in Step 3), `total_required_devices`, `required_devices_per_node`, `node_count` (already computed in Step 1), and — for slurm scenarios — the `account` and `partition` resolved in Step 2.5 along with `slurm_env_file=<WORK_DIR>/slurm_env.json` (the path to the canonical JSON written in Step 2.5, with cluster-policy fields merged in; the script builder loads it in its Step 1 and reads `mounts` / per-partition `gpus_per_node` from there). It still resolves mounts and generates:
+
+- `<WORK_DIR>/job_spec.json` — manifest with all resolved values (including `node_count` and `container_image`)
+- `<WORK_DIR>/<MODEL_NAME>_<type>.slurm` (Slurm scenarios) or `docker_cmd.sh` (for `local_docker`)
+
+For `local_direct`, **skip the script builder entirely**. Write a minimal `<WORK_DIR>/job_spec.json`:
+
+- `scenario: local_direct`
+- `direct_cmd`: the raw `test_cmd` / `eval_cmd` / `bench_cmd` / `custom_cmd` (use `bash <custom_script>` if `custom_script` was given)
+- `custom_env`: the `KEY=VALUE` pairs to export (if provided)
+- `monitor_timeout_seconds`: normalized wall-clock limit (default `3600`)
+- `work_dir`, `log_file`, `model_name`, `workflow_type`, `success_patterns`, `failure_patterns`
+
+If the skill reports an error, stop and report to the user.
+
+### Step 5: Dispatch to Executor Subagent
+
+Read `<WORK_DIR>/job_spec.json` and dispatch by `job_spec.scenario`. The scenario was decided in Step 2 — **do not re-derive or re-map it here**; the per-scenario blocks below are dispatch *implementations* for each value of `job_spec.scenario`, not a second copy of the routing rule.
+
+**Hang detection + task timeout:** Every executor enforces hang detection (poll log for `hang detected`, case-insensitive → `HANG_DETECTED`) and `monitor_timeout_seconds` (overall wall-clock limit → `TIMEOUT`). The policy and implementation live in each executor's SKILL; case-executor does not restate them per dispatch. Before spawning, normalize `monitor_timeout` to seconds (`1h` → `3600`, `30m` → `1800`, `HH:MM:SS` → seconds) and include the value in `job_spec.json` as `monitor_timeout_seconds`. Default `3600` (1h).
+
+**For `local_docker`:**
+
+Docker is available; the workload runs locally and is short-lived enough that a separate subagent context buys nothing. Do **not** spawn a subagent — load the `exec-local-docker` skill into this context and follow its procedure with the fields from `<WORK_DIR>/job_spec.json`:
+
+```
+Skill(skill="trtllm-agent-toolkit:exec-local-docker", args="""
+Job spec (from <WORK_DIR>/job_spec.json):
+- docker_cmd, work_dir, log_file, model_name, workflow_type
+- success_patterns, failure_patterns, monitor_timeout_seconds
+- build_project, gpu_type, repo_root, container_image (consumed by the skill's optional Step 0 build pre-step)
+""")
+```
+
+When the skill returns, treat its result (status, exit code, log file, summary, errors) as this dispatch's outcome and proceed to Step 6.
+
+**For `local_direct`:**
+
+Docker is not available on this host. Do **not** spawn an executor subagent. Run the command in place:
+
+```bash
+cd "<REPO_ROOT>"
+[ -n "$CUSTOM_ENV" ] && export $CUSTOM_ENV
+
+timeout --kill-after=30s "${MONITOR_TIMEOUT_SECONDS}s" bash -c '<direct_cmd>' 2>&1 | tee "$LOG_FILE"
+EXIT_CODE=${PIPESTATUS[0]}
+```
+
+Hang detection: monitor `$LOG_FILE` in the background (every 10s, case-insensitive grep for `hang detected`). If found, send `SIGTERM` to the command's process group and report `HANG_DETECTED`.
+
+Task timeout: `timeout` enforces `monitor_timeout_seconds` (default `3600`). If `EXIT_CODE` is `124` (or `137` after the `--kill-after` grace period), report `TIMEOUT`.
+
+Map the result: `EXIT_CODE=0` → `PASSED`; otherwise apply the workflow's `failure_patterns` / `success_patterns` to `$LOG_FILE` to pick `FAILED` / `ERROR` / `TIMEOUT` / `OUT_OF_MEMORY`, etc.
+
+**For `local_slurm`:**
+```
+Agent(subagent_type="exec-local-slurm", prompt="""
+Run the job described in <WORK_DIR>/job_spec.json.
+
+Job spec (from <WORK_DIR>/job_spec.json):
+- script_path, work_dir, model_name, workflow_type
+- success_patterns, failure_patterns, log_file_pattern, monitor_timeout_seconds
+""")
+```
+
+**For `remote_slurm`:**
+7/600
+Before spawning the executor, resolve the local repo URL and branch:
+
+```bash
+REPO_URL=$(git -C "<REPO_ROOT>" remote get-url origin 2>/dev/null)
+REPO_BRANCH=$(git -C "<REPO_ROOT>" rev-parse --abbrev-ref HEAD 2>/dev/null)
+```
+
+Include these in `job_spec.json` as `repo_url` and `repo_branch`.
+
+```
+Agent(subagent_type="exec-remote-slurm", prompt="""
+Run the job described in <WORK_DIR>/job_spec.json.
+
+Job spec (from <WORK_DIR>/job_spec.json):
+- script_path, script_name, work_dir, model_name, workflow_type
+- success_patterns, failure_patterns, log_file_pattern
+- slurm_cluster (cluster name only, when explicit; absent for auto-select). The remote executor resolves per-cluster info via internal-env-info; connection fields come from job_spec.json.
+- ssh_host, slurm_user, remote_cwd, remote_work_dir (input overrides; absent required fields cause the remote executor to stop and ask the user)
+- slurm_password, extra_files
+- repo_url: Git remote URL of the local TensorRT-LLM repo (for remote environment setup)
+- repo_branch: Git branch to check out on the remote cluster
+- device_type, total_required_devices, required_devices_per_node: hardware constraints. When slurm_cluster is absent, the executor uses device_type and required_devices_per_node to auto-select a cluster from internal-env-info catalog mode.
+- container_image: resolved by case-executor (Step 3) — the executor uses this directly for both the build job and the run job. Apply transport-specific URL rewrites (e.g., enroot's `/` → `#`) at use-time. Never re-grep `current_image_tags.properties`.
+- node_count: precomputed by case-executor (Step 1) — the executor uses this directly for `--nodes`, never recomputes from totals.
+- job_name: precomputed by case-executor (Step 1) — the executor substitutes this into `-J <job_name>` directly, never reconstructs from `account` / `detail`.
+- monitor_timeout_seconds
+""")
+```
+
+### Step 6: Report Results
+
+Relay the subagent result to the user with: task type, Docker image, command/script path, status (`PASSED`/`FAILED`/`TIMEOUT`/`CANCELLED`/`HANG_DETECTED`/etc.), exit code, log file paths, summary, errors (if any), work directory, and Slurm details (if applicable).
+
+---
+
+## Failure Diagnosis
+
+| Symptom | Cause | Action |
+|---------|-------|--------|
+| Exit 1 + `FAILED` | Pytest assertion failure | Show failing tests and messages |
+| Exit 2 + `ERROR` | Collection/import failure | Check first 50 lines for ImportError |
+| `TIMEOUT` | Exceeded time limit | Increase `time_limit` |
+| `OUT_OF_MEMORY` | OOM | Reduce batch size or increase nodes |
+| `CANCELLED` | Preempted or user-cancelled | `sacct --format=JobID,State,Reason` |
+| Empty output file | Container failed to start | Check `.err` for mount/image errors |
+| `hang detected` in log | GPU/NCCL deadlock | Auto-terminated; check log context for root cause |
+| `Invalid generic resource` | Cluster doesn't support `--gres` | Omit `--gres` |
+| SSH `Broken pipe` | ControlMaster died | Subagent re-establishes |
+
+## Reference Files
+
+| File | Purpose |
+|------|---------|
+| `scripts/detect_slurm_env.sh` | One-pass SLURM env probe: accounts, partitions (with per-partition `arch` / `gres` / `gpus_per_node` / `requires_gres`), and the cluster's PMIx plugin (`pmix.available` / `pmix.preferred`). JSON or text output. Used in Step 2.5 (run locally for `local_slurm`; over SSH via `exec-remote-slurm` for `remote_slurm`). The emitted JSON is the single source of truth — downstream skills (`trtllm-test-script-builder`, the slurm executors) read it via `job_spec.slurm_env` and do **not** re-detect. |
+| `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` (in the TRT-LLM repo) | Generates perf-sanity launch scripts from config YAMLs |
+| `scripts/slurm_run_custom.sh` | Runs custom commands inside Slurm container |
+
+### Skills & Subagents
+
+| Skill | Role | Step |
+|-------|------|------|
+| `exec-env-check` | Detect GPUs + Slurm availability | 2 |
+| `trtllm-test-script-builder` | Resolve params + generate script + `job_spec.json` | 4 |
+| `exec-local-docker` | Local Docker execution | 5 |
+| `exec-local-slurm` | Local Slurm submission + monitoring | 5 |
+| `exec-remote-slurm` | Remote Slurm: SSH + submit + monitor | 5 |
diff --git a/.claude/skills/trtllm-case-executor/scripts/detect_slurm_env.sh b/.claude/skills/trtllm-case-executor/scripts/detect_slurm_env.sh
new file mode 100755
index 000000000000..8b6e119736ee
--- /dev/null
+++ b/.claude/skills/trtllm-case-executor/scripts/detect_slurm_env.sh
@@ -0,0 +1,449 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Detect a SLURM cluster's environment for the current user in a single pass:
+#   - account associations and a default account
+#   - visible partitions (sinfo) and a default partition validated via
+#     `sbatch --test-only` for the user's account
+#   - per-partition hardware: Arch (scontrol show node), GRES (sinfo %G),
+#     gpus_per_node, requires_gres
+#   - the cluster's compatible PMIx plugin (srun --mpi=list)
+#
+# Run locally for local_slurm; stream over SSH for remote_slurm.
+#
+# Usage:
+#   bash detect_slurm_env.sh                          # JSON, probe all up partitions
+#   bash detect_slurm_env.sh --format=text            # human-readable
+#   bash detect_slurm_env.sh --partitions=batch,gpu   # restrict hw probe to these
+#
+# Output (JSON, default):
+#   {
+#     "user": "<whoami>",
+#     "default_account": "<first-account-or-null>",
+#     "accounts": ["<acct1>", ...],
+#     "default_partition": "<partition-or-null>",
+#     "partitions": [
+#       {
+#         "name": "<part>",
+#         "state": "<up|down|...>",
+#         "time_limit": "<HH:MM:SS|infinite>",
+#         "nodes": "<count>",
+#         "arch": "x86_64|aarch64|null",
+#         "gres": "gpu:8|(null)",
+#         "gpus_per_node": <int|null>,
+#         "requires_gres": true|false
+#       },
+#       ...
+#     ],
+#     "pmix": {
+#       "available": ["pmix", "pmix_v3", "pmix_v4", "pmix_v5"],
+#       "preferred": "pmix_v5|null"
+#     },
+#     "errors": ["<diagnostic>", ...]
+#   }
+#
+# Exit codes:
+#   0  success (account and at least one partition resolved, no errors)
+#   1  SLURM tools not available
+#   2  partial success (one or more errors collected; see "errors")
+
+set -u
+set -o pipefail
+
+FORMAT="json"
+PARTITIONS_ARG=""
+for arg in "$@"; do
+    case "$arg" in
+        --format=json) FORMAT="json" ;;
+        --format=text) FORMAT="text" ;;
+        --partitions=*) PARTITIONS_ARG="${arg#--partitions=}" ;;
+        --partitions)
+            echo "--partitions requires a value (use --partitions=p1,p2)" >&2
+            exit 1
+            ;;
+        -h|--help)
+            sed -n '2,/^$/p' "$0" | sed 's/^# \{0,1\}//'
+            exit 0
+            ;;
+        *)
+            echo "unknown argument: $arg" >&2
+            exit 1
+            ;;
+    esac
+done
+
+ERRORS=()
+
+if ! command -v sacctmgr >/dev/null 2>&1 || ! command -v sinfo >/dev/null 2>&1; then
+    ERRORS+=("SLURM client tools (sacctmgr/sinfo) not found in PATH")
+    if [ "$FORMAT" = "json" ]; then
+        printf '{"user":"%s","default_account":null,"accounts":[],"default_partition":null,"partitions":[],"pmix":{"available":[],"preferred":null},"errors":["SLURM client tools not found"]}\n' "$(whoami)"
+    else
+        echo "ERROR: SLURM client tools (sacctmgr / sinfo) not found in PATH" >&2
+    fi
+    exit 1
+fi
+
+USER_NAME="$(whoami)"
+
+# --- Accounts -----------------------------------------------------------------
+# `sacctmgr -nP show assoc where user=$USER format=account` returns one
+# account per line. Some sites duplicate entries (per-cluster, per-partition);
+# de-duplicate while preserving order.
+ACCOUNTS_RAW="$(sacctmgr -nP show assoc where user="$USER_NAME" format=account 2>/dev/null || true)"
+ACCOUNTS=()
+if [ -n "$ACCOUNTS_RAW" ]; then
+    while IFS= read -r line; do
+        line="${line//[[:space:]]/}"
+        [ -z "$line" ] && continue
+        # de-dupe
+        local_seen=0
+        for a in "${ACCOUNTS[@]:-}"; do
+            [ "$a" = "$line" ] && local_seen=1 && break
+        done
+        [ "$local_seen" -eq 0 ] && ACCOUNTS+=("$line")
+    done <<< "$ACCOUNTS_RAW"
+fi
+DEFAULT_ACCOUNT="${ACCOUNTS[0]:-}"
+if [ -z "$DEFAULT_ACCOUNT" ]; then
+    ERRORS+=("no SLURM account associations for user $USER_NAME")
+fi
+
+# --- Partitions (basic sinfo fields) -----------------------------------------
+# `sinfo -h -o "%P|%a|%l|%D"` — partition (default has a trailing '*'),
+# avail (up/down), time limit, node count.
+
+PART_LINES=()   # "name|state|time_limit|nodes" in sinfo order, de-duped by name
+UP_PARTS=()     # ordered list of partition names with state=up
+PART_RAW="$(sinfo -h -o '%P|%a|%l|%D' 2>/dev/null || true)"
+if [ -n "$PART_RAW" ]; then
+    declare -A SEEN_PART
+    while IFS='|' read -r pname pavail ptime pnodes; do
+        pname_trim="${pname%[*]}"   # strip default '*'
+        [ -z "$pname_trim" ] && continue
+        if [ -z "${SEEN_PART[$pname_trim]:-}" ]; then
+            SEEN_PART["$pname_trim"]=1
+            PART_LINES+=("$pname_trim|$pavail|$ptime|$pnodes")
+            [ "$pavail" = "up" ] && UP_PARTS+=("$pname_trim")
+        fi
+    done <<< "$PART_RAW"
+fi
+if [ "${#PART_LINES[@]}" -eq 0 ]; then
+    ERRORS+=("sinfo returned no partitions")
+fi
+
+# --- Resolve the hardware-probe target list ----------------------------------
+# Default: every up partition. If --partitions=... was provided, restrict to
+# those (intersection with up partitions only; unknown names are warned).
+HW_TARGETS=()
+if [ -n "$PARTITIONS_ARG" ]; then
+    IFS=',' read -r -a HW_TARGETS_REQUESTED <<< "$PARTITIONS_ARG"
+    declare -A UP_SET
+    for u in "${UP_PARTS[@]:-}"; do UP_SET["$u"]=1; done
+    for r in "${HW_TARGETS_REQUESTED[@]:-}"; do
+        [ -z "$r" ] && continue
+        if [ -n "${UP_SET[$r]:-}" ]; then
+            HW_TARGETS+=("$r")
+        else
+            ERRORS+=("--partitions=$r: not an up partition (skipping hw probe)")
+        fi
+    done
+else
+    HW_TARGETS=("${UP_PARTS[@]:-}")
+fi
+
+# --- Per-partition arch + gres -----------------------------------------------
+# PART_HW[name] = "arch|gres|gpus_per_node|requires_gres"; absent if not probed.
+declare -A PART_HW
+for part in "${HW_TARGETS[@]:-}"; do
+    [ -z "$part" ] && continue
+
+    NODE="$(sinfo -p "$part" -h -o '%n' 2>/dev/null | head -1 | awk '{print $1}')"
+    ARCH=""
+    if [ -n "$NODE" ]; then
+        ARCH="$(scontrol show node "$NODE" 2>/dev/null | awk -F= '/^ *Arch=/{print $2; exit}' | tr -d '[:space:]')"
+    fi
+    if [ -z "$ARCH" ]; then
+        ERRORS+=("partition $part: could not determine Arch (node=${NODE:-unknown})")
+    fi
+
+    GRES_RAW="$(sinfo -p "$part" -h -o '%G' 2>/dev/null | head -1 | sed 's/[[:space:]]*$//')"
+    if [ -z "$GRES_RAW" ]; then
+        GRES_RAW="(null)"
+    fi
+
+    REQUIRES_GRES=false
+    GPUS_PER_NODE=""
+    case "$GRES_RAW" in
+        gpu:*)
+            REQUIRES_GRES=true
+            # gpu:4, gpu:4(IDX:0-3), gpu:a100:8, gpu:a100:8(...)
+            tmp="${GRES_RAW%%(*}"        # strip "(...)" suffix
+            tmp="${tmp##*:}"             # last colon-separated token = count
+            if [[ "$tmp" =~ ^[0-9]+$ ]]; then
+                GPUS_PER_NODE="$tmp"
+            fi
+            ;;
+    esac
+
+    PART_HW["$part"]="$ARCH|$GRES_RAW|$GPUS_PER_NODE|$REQUIRES_GRES"
+done
+
+# --- pmix plugin detection ---------------------------------------------------
+# `srun --mpi=list` is a client-side query (no allocation required). Output
+# format varies across SLURM versions; harvest every token that starts with
+# 'pmix' and dedupe.
+PMIX_AVAILABLE=()
+PMIX_PREFERRED=""
+if command -v srun >/dev/null 2>&1; then
+    MPI_LIST_RAW="$(srun --mpi=list 2>&1 || true)"
+    if [ -n "$MPI_LIST_RAW" ]; then
+        declare -A PMIX_SEEN
+        while IFS= read -r tok; do
+            tok="${tok//[[:space:]]/}"
+            tok="${tok#-}"
+            tok="${tok#*}"
+            case "$tok" in
+                pmix|pmix_v*)
+                    if [ -z "${PMIX_SEEN[$tok]:-}" ]; then
+                        PMIX_SEEN["$tok"]=1
+                        PMIX_AVAILABLE+=("$tok")
+                    fi
+                    ;;
+            esac
+        done < <(printf '%s\n' "$MPI_LIST_RAW" | tr ',' '\n')
+    else
+        ERRORS+=("srun --mpi=list returned no output")
+    fi
+
+    # Preferred: highest pmix_vN, else base 'pmix'.
+    BEST_VER=-1
+    for p in "${PMIX_AVAILABLE[@]:-}"; do
+        [ -z "$p" ] && continue
+        if [[ "$p" =~ ^pmix_v([0-9]+)$ ]]; then
+            v="${BASH_REMATCH[1]}"
+            if [ "$v" -gt "$BEST_VER" ]; then
+                BEST_VER="$v"
+                PMIX_PREFERRED="$p"
+            fi
+        fi
+    done
+    if [ -z "$PMIX_PREFERRED" ]; then
+        for p in "${PMIX_AVAILABLE[@]:-}"; do
+            if [ "$p" = "pmix" ]; then
+                PMIX_PREFERRED="pmix"
+                break
+            fi
+        done
+    fi
+    if [ -z "$PMIX_PREFERRED" ] && [ "${#PMIX_AVAILABLE[@]}" -eq 0 ]; then
+        ERRORS+=("no pmix MPI plugin reported by 'srun --mpi=list'")
+    fi
+else
+    ERRORS+=("srun not found; cannot enumerate MPI plugins")
+fi
+
+# --- Default partition selection ---------------------------------------------
+# Verify actual submission access for each "up" partition using sbatch
+# --test-only. A submission error indicates the account/QOS combination is
+# not permitted. Among accessible partitions pick by priority:
+#   1. Standard partitions  (no recognised suffix)
+#   2. Backfill partitions  (*-backfill)
+#   3. High-priority / admin partitions (*-hp, *-admin)
+_part_priority() {
+    local name="$1"
+    case "$name" in
+        *-hp|*-admin) echo 3 ;;
+        *-backfill)   echo 2 ;;
+        *)            echo 1 ;;
+    esac
+}
+
+DEFAULT_PARTITION=""
+ACCESSIBLE_PARTS=()
+if [ -n "$DEFAULT_ACCOUNT" ] && [ "${#UP_PARTS[@]}" -gt 0 ]; then
+    for pname in "${UP_PARTS[@]}"; do
+        sbatch --test-only \
+            --partition="$pname" \
+            --account="$DEFAULT_ACCOUNT" \
+            --nodes=1 --time=00:01:00 --ntasks=1 \
+            --wrap="true" >/dev/null 2>&1
+        if [ $? -eq 0 ]; then
+            ACCESSIBLE_PARTS+=("$pname")
+        fi
+    done
+fi
+
+if [ "${#ACCESSIBLE_PARTS[@]}" -gt 0 ]; then
+    best_name=""
+    best_score=99
+    for pname in "${ACCESSIBLE_PARTS[@]}"; do
+        score=$(_part_priority "$pname")
+        if [ "$score" -lt "$best_score" ]; then
+            best_score=$score
+            best_name="$pname"
+        fi
+    done
+    DEFAULT_PARTITION="$best_name"
+fi
+
+# Fallback: if sbatch --test-only is unavailable or all checks failed,
+# use the first listed partition (original behaviour).
+if [ -z "$DEFAULT_PARTITION" ] && [ "${#PART_LINES[@]}" -gt 0 ]; then
+    ERRORS+=("partition access verification failed; falling back to first available partition")
+    first="${PART_LINES[0]}"
+    DEFAULT_PARTITION="${first%%|*}"
+fi
+
+# --- Emit --------------------------------------------------------------------
+json_escape() {
+    local s="$1"
+    s="${s//\\/\\\\}"
+    s="${s//\"/\\\"}"
+    printf '%s' "$s"
+}
+
+emit_json() {
+    local acc_json="[]"
+    if [ "${#ACCOUNTS[@]}" -gt 0 ]; then
+        acc_json="["
+        local first=1
+        for a in "${ACCOUNTS[@]}"; do
+            if [ "$first" -eq 1 ]; then first=0; else acc_json+=","; fi
+            acc_json+="\"$(json_escape "$a")\""
+        done
+        acc_json+="]"
+    fi
+
+    local part_json="[]"
+    if [ "${#PART_LINES[@]}" -gt 0 ]; then
+        part_json="["
+        local first=1
+        for line in "${PART_LINES[@]}"; do
+            IFS='|' read -r p_name p_avail p_time p_nodes <<< "$line"
+            if [ "$first" -eq 1 ]; then first=0; else part_json+=","; fi
+
+            local arch_json="null"
+            local gres_json="null"
+            local gpn_json="null"
+            local req_json="false"
+            local hw="${PART_HW[$p_name]:-}"
+            if [ -n "$hw" ]; then
+                IFS='|' read -r h_arch h_gres h_gpn h_req <<< "$hw"
+                [ -n "$h_arch" ] && arch_json="\"$(json_escape "$h_arch")\""
+                gres_json="\"$(json_escape "$h_gres")\""
+                [[ "$h_gpn" =~ ^[0-9]+$ ]] && gpn_json="$h_gpn"
+                req_json="$h_req"
+            fi
+
+            part_json+="{\"name\":\"$(json_escape "$p_name")\","
+            part_json+="\"state\":\"$(json_escape "$p_avail")\","
+            part_json+="\"time_limit\":\"$(json_escape "$p_time")\","
+            part_json+="\"nodes\":\"$(json_escape "$p_nodes")\","
+            part_json+="\"arch\":$arch_json,"
+            part_json+="\"gres\":$gres_json,"
+            part_json+="\"gpus_per_node\":$gpn_json,"
+            part_json+="\"requires_gres\":$req_json}"
+        done
+        part_json+="]"
+    fi
+
+    local pmix_avail_json="[]"
+    if [ "${#PMIX_AVAILABLE[@]}" -gt 0 ]; then
+        pmix_avail_json="["
+        local first=1
+        for p in "${PMIX_AVAILABLE[@]}"; do
+            if [ "$first" -eq 1 ]; then first=0; else pmix_avail_json+=","; fi
+            pmix_avail_json+="\"$(json_escape "$p")\""
+        done
+        pmix_avail_json+="]"
+    fi
+    local pmix_pref_json="null"
+    [ -n "$PMIX_PREFERRED" ] && pmix_pref_json="\"$(json_escape "$PMIX_PREFERRED")\""
+
+    local err_json="[]"
+    if [ "${#ERRORS[@]}" -gt 0 ]; then
+        err_json="["
+        local first=1
+        for e in "${ERRORS[@]}"; do
+            if [ "$first" -eq 1 ]; then first=0; else err_json+=","; fi
+            err_json+="\"$(json_escape "$e")\""
+        done
+        err_json+="]"
+    fi
+
+    local def_acc_json="null"
+    [ -n "$DEFAULT_ACCOUNT" ] && def_acc_json="\"$(json_escape "$DEFAULT_ACCOUNT")\""
+    local def_part_json="null"
+    [ -n "$DEFAULT_PARTITION" ] && def_part_json="\"$(json_escape "$DEFAULT_PARTITION")\""
+
+    printf '{"user":"%s","default_account":%s,"accounts":%s,"default_partition":%s,"partitions":%s,"pmix":{"available":%s,"preferred":%s},"errors":%s}\n' \
+        "$(json_escape "$USER_NAME")" \
+        "$def_acc_json" "$acc_json" \
+        "$def_part_json" "$part_json" \
+        "$pmix_avail_json" "$pmix_pref_json" \
+        "$err_json"
+}
+
+emit_text() {
+    echo "user: $USER_NAME"
+    echo "default_account: ${DEFAULT_ACCOUNT:-<none>}"
+    echo "accounts: ${ACCOUNTS[*]:-<none>}"
+    echo "default_partition: ${DEFAULT_PARTITION:-<none>}"
+    echo "partitions:"
+    if [ "${#PART_LINES[@]}" -gt 0 ]; then
+        printf '  %-20s %-6s %-12s %-6s %-10s %-20s %-12s %s\n' \
+            NAME STATE TIME_LIMIT NODES ARCH GRES GPUS_PER_NODE REQUIRES_GRES
+        for line in "${PART_LINES[@]}"; do
+            IFS='|' read -r p_name p_avail p_time p_nodes <<< "$line"
+            arch="-"; gres="-"; gpn="-"; req="-"
+            hw="${PART_HW[$p_name]:-}"
+            if [ -n "$hw" ]; then
+                IFS='|' read -r h_arch h_gres h_gpn h_req <<< "$hw"
+                arch="${h_arch:-<unknown>}"
+                gres="$h_gres"
+                gpn="${h_gpn:-<n/a>}"
+                req="$h_req"
+            fi
+            printf '  %-20s %-6s %-12s %-6s %-10s %-20s %-12s %s\n' \
+                "$p_name" "$p_avail" "$p_time" "$p_nodes" "$arch" "$gres" "$gpn" "$req"
+        done
+    else
+        echo "  <none>"
+    fi
+    echo "pmix:"
+    echo "  available: ${PMIX_AVAILABLE[*]:-<none>}"
+    echo "  preferred: ${PMIX_PREFERRED:-<none>}"
+    if [ "${#ERRORS[@]}" -gt 0 ]; then
+        echo "errors:"
+        for e in "${ERRORS[@]}"; do echo "  - $e"; done
+    fi
+}
+
+if [ "$FORMAT" = "json" ]; then
+    emit_json
+else
+    emit_text
+fi
+
+# Exit code
+if [ -z "$DEFAULT_ACCOUNT" ] || [ -z "$DEFAULT_PARTITION" ]; then
+    exit 2
+fi
+if [ "${#ERRORS[@]}" -gt 0 ]; then
+    exit 2
+fi
+exit 0
diff --git a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md
index cc96eaaefa02..06e2796893e2 100644
--- a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md
+++ b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md
@@ -306,16 +306,16 @@ class {Name}Model(PreTrainedModel):
 
 Subclass **both** `BaseMultimodalInputProcessor` (drives every real request) and `BaseMultimodalDummyInputsBuilder` (drives engine warmup / KV-cache profiling). Colocate in the modeling file. References: `Qwen3VLInputProcessorBase` (image+video), `Mistral3InputProcessor` (Pixtral).
 
-**Encoder KV-cache memory profiling (deterministic dummy sizing).** The KV-cache profiler sizes the encoder's memory contribution by running the encoder **once** on a worst-case dummy (held resident through the peak measurement), **decoupled** from the text-only LLM dummy. To opt in, the model exposes `encode_multimodal_inputs` (via `MultimodalModelMixin`) and the input processor implements the modality-agnostic dummy contract on `BaseMultimodalDummyInputsBuilder`:
+**Encoder KV-cache memory profiling (deterministic dummy sizing).** The KV-cache profiler sizes the encoder's memory contribution by running the encoder **once** on a worst-case dummy (held resident through the peak measurement), **decoupled** from the text-only LLM dummy. To opt in, the model exposes `encode_multimodal_inputs` (via `MultimodalModelMixin`) and the input processor implements the existing `BaseMultimodalDummyInputsBuilder` contract:
 
-- `get_mm_max_tokens_per_item() -> {modality: tokens}` — per-modality worst-case single-item encoder-attention tokens. The keys enumerate the modalities the model encodes; the profiler splits the shared `encoder_max_num_tokens` across them in proportion to these (so they share one microbatch cap, not each the whole budget). Default `{}` → no direct encoder profiling.
-- `get_dummy_mm_data_for_tokens(*, max_tokens_per_modality, dtype) -> multimodal_data` — materialize the processed encoder tensors **directly** (zeros of the exact shape the processor would emit; no PIL image + HF-processor round-trip), merged into one `multimodal_data` dict so a single `encode_multimodal_inputs` profiles the combined peak. Default raises `NotImplementedError`.
+- `get_mm_max_tokens_per_item(max_num_encoder_tokens=None) -> {modality: tokens}` — report bounded startup maxima when the argument is `None`, or the largest legal item per modality under a concrete runtime token budget.
+- `get_dummy_mm_data(*, max_num_encoder_tokens, mm_counts, dtype) -> multimodal_data` — build the profiler-selected per-modality item counts **directly** as processed encoder tensors (no PIL image + HF-processor round-trip).
 
-Vision models implement these via the size trio: `get_num_mm_tokens(*, width, height, num_frames)` (size → **pre-merger encoder-attention tokens**; the single source of truth shared with the hashing path `get_num_tokens_per_image`/`_video`), its inverse `get_size_for_max_tokens(max_tokens)` (largest aspect-bounded size whose token count ≤ budget, capped at `max_pixels`), and `get_dummy_mm_data_for_size(...)`. Qwen builds `pixel_values`/`image_grid_thw`; Mistral builds `pixel_values`/`image_sizes` and keeps its ViT patch count off `get_num_mm_tokens` (a private `_vit_tokens` helper) so the LLM-side Pixtral hashing count is unchanged. A model with neither contract falls back to a text-only dummy (encoder memory unaccounted). Don't hardcode the encoder attention workspace (`max_num_*=8192`): inherit `MultimodalEncoderMixin` and let the engine size it via `setup_attn_metadata` at load.
+Following vLLM's separation of common budgeting from model processing geometry, the profiler selects the largest legal runtime workload and computes its repetition count under `encoder_max_num_tokens` and `encoder_max_batch_size`. It passes the resulting plan (for example, `{"image": 8}`) to the concrete processor's `get_dummy_mm_data`; the processor owns only geometry and tensor layout. Qwen builds `pixel_values`/`*_grid_thw`; Mistral builds `pixel_values`/`image_sizes` and uses its private `_vit_tokens` helper so the LLM-side Pixtral hashing count is unchanged. A model without the contract falls back to a text-only dummy (encoder memory unaccounted). Don't hardcode the encoder attention workspace (`max_num_*=8192`): inherit `MultimodalEncoderMixin` and let the engine size it via `setup_attn_metadata` at load.
 
-The workspace dimensions come from `TorchLlmArgs.get_encoder_runtime_sizes()` → `(encoder_max_batch_size, encoder_max_num_tokens)` — two prototype knobs that size the encoder's `AttentionMetadata` independently of the LLM batch and fall back to the LLM-side `max_batch_size` / `max_num_tokens` when unset. `encoder_max_num_tokens` is exactly the per-iteration encoder microbatch token cap that `get_dummy_mm_data_for_tokens` saturates (and that the profiler splits across modalities), so an encoder microbatch can be sized larger than the LLM `max_num_tokens` without inflating the KV-cache budget. Read them via `get_encoder_runtime_sizes()` rather than the raw fields so the fallback is applied.
+The workspace dimension comes from `encoder_max_num_tokens`, falling back to the LLM-side `max_num_tokens` when unset. `encoder_max_batch_size` independently bounds the number of atomic items in a multimodal encoder iteration and falls back to `max_batch_size`. The profiler resolves both limits into `mm_counts` before calling `get_dummy_mm_data`. Model-internal attention sequence counts are not the item batch size; they are derived separately from the token budget and model geometry.
 
-> Mixed image+video+audio models (nemotron-nano / phi4mm) compose multiple modality dummies through the same `get_dummy_mm_data_for_tokens` (return `{"image": ..., "audio": ...}`); a per-modality `ModalityDummySizer` composition is the planned home for the shared orchestration.
+> Mixed image+video+audio models profile the supported modality with the largest legal per-item workload under the configured limits. Runtime mixed-modality requests still share the same aggregate token limit.
 
 Implement `call_with_text_prompt(inputs, sampling_params)` — the per-model text-prompt path. **Don't override `__call__`**: the base class's concrete `__call__` dispatches here for text prompts, and also detokenizes `prompt_token_ids → prompt` and falls through to here for non-fast-path VLMs. `call_with_text_prompt` does:
 
@@ -444,7 +444,7 @@ Follow `CONTRIBUTING.md`. Title `[JIRA/NVBUG/None][type] description`, `git comm
 
 **Input processor**
 - [ ] Subclasses both `BaseMultimodalInputProcessor` and `BaseMultimodalDummyInputsBuilder`.
-- [ ] Encoder KV-cache profiling: implements the deterministic dummy contract (`get_mm_max_tokens_per_item` + `get_dummy_mm_data_for_tokens`, vision via the `get_num_mm_tokens` / `get_size_for_max_tokens` / `get_dummy_mm_data_for_size` trio) and the model exposes `encode_multimodal_inputs`; encoder inherits `MultimodalEncoderMixin` (no hardcoded `max_num_*=8192` — sized by `setup_attn_metadata`). Skipping these = text-only dummy, encoder memory unaccounted.
+- [ ] Encoder KV-cache profiling: implements the deterministic dummy contract (`get_mm_max_tokens_per_item` + `get_dummy_mm_data`) and the model exposes `encode_multimodal_inputs`; encoder inherits `MultimodalEncoderMixin` (no hardcoded `max_num_*=8192` — sized by `setup_attn_metadata`). Skipping these = text-only dummy, encoder memory unaccounted.
 - [ ] `call_with_text_prompt` (not `__call__` — that's the base-class dispatcher) runs HF AutoProcessor + tokenizer, builds `multimodal_data` by modality, computes `mrope_config` on CPU, `_postprocess`-rewrites mm token ids to the OOV sentinel.
 - [ ] `mm_processor_kwargs` flow-through preserved. (Tokenized fast path is optional: set `supports_token_id_mm_expansion = True` + implement `get_text_with_mm_placeholders` / `expand_prompt_token_ids_for_mm`; otherwise the base class detokenizes token-ID inputs automatically.)
 - [ ] `_attach_multimodal_embeddings_impl` implemented (not the `attach_multimodal_embeddings` wrapper) if `@support_multimodal_disaggregated`.
diff --git a/.claude/skills/trtllm-test-script-builder/SKILL.md b/.claude/skills/trtllm-test-script-builder/SKILL.md
new file mode 100644
index 000000000000..0812232922b9
--- /dev/null
+++ b/.claude/skills/trtllm-test-script-builder/SKILL.md
@@ -0,0 +1,385 @@
+---
+name: trtllm-test-script-builder
+description: >-
+  Build Slurm scripts or Docker commands for TensorRT-LLM workloads. Resolves
+  all parameters (docker image, mounts, parallelism,
+  MPI mode), generates the complete script from Category templates, and writes
+  both the script and a job_spec.json manifest to the work directory.
+license: Apache-2.0
+---
+
+# TensorRT-LLM Script Builder
+
+Resolve parameters and generate execution scripts for all workflow types.
+
+## Input (from orchestrator prompt)
+
+The orchestrator passes these fields. Not all fields are present for every workflow.
+
+| Field | Description | Required |
+|-------|-------------|----------|
+| `workflow_type` | `pytest`, `eval`, `bench`, `custom`, or `perf_sanity` | Always |
+| `work_dir` | Absolute path to local work directory (already created) | Always |
+| `skill_dir` | Absolute path to the test-script-builder skill directory (`.claude/skills/trtllm-test-script-builder`) | Always |
+| `repo_root` | Absolute path to the TensorRT-LLM repo root | Always |
+| `test_cmd` | Pytest command string | pytest |
+| `eval_cmd` | `trtllm-eval` command string, or `evaluation`, `accuracy`, `score`, `dataset` | eval |
+| `bench_cmd` | `trtllm-bench` command (may chain `prepare-dataset` and `throughput`/`latency` with `&&`). Auto-parses `--tp`/`--pp`/`--ep` from the last `trtllm-bench` invocation. | bench |
+| `custom_cmd` | Custom command string, `test command`, non-pytest command | custom |
+| `custom_script` | Path to custom script file | custom (optional) |
+| `custom_env` | Space-separated `KEY=VALUE` pairs | custom (optional) |
+| `perf_config_yaml` | Path to pre-built perf-sanity or benchmark config YAML | perf_sanity (required) |
+| `perf_test_name` | Server config name for aggregated mode | pytest (optional) |
+| `model_name` | Short model name (auto-derived if not provided) | Optional |
+| `required_devices` | Number of GPUs needed (auto-derived if not provided) | Optional |
+| `node_count` | Number of Slurm nodes (auto-derived if not provided) | Optional |
+| `partition` | Slurm partition. **For slurm scenarios, this MUST be the value resolved by `trtllm-case-executor` Step 2.5 against the cluster (validated to exist in `slurm_env.partitions[].name`).** Do not detect or default it here. | Required (slurm) |
+| `account` | Slurm account. **For slurm scenarios, this MUST be the value resolved by `trtllm-case-executor` Step 2.5 against the cluster (validated to exist in `slurm_env.accounts[]`).** Do not detect or default it here. | Required (slurm) |
+| `slurm_env_file` | Absolute path to the `slurm_env.json` file written by `trtllm-case-executor` Step 2.5 (canonically `<work_dir>/slurm_env.json`). Schema: `{user, default_account, accounts[], default_partition, partitions[{name,state,time_limit,nodes,arch,gres,gpus_per_node,requires_gres}], pmix:{available[],preferred}, errors[]}`. The script builder loads it once in Step 1, uses it to select the Docker image (arch), drive `--gres` / `--mpi=` flags, and embeds the parsed object into `job_spec.json` for downstream consumers. | Required (slurm) |
+| `time_limit` | Slurm time limit (default: `02:00:00`) | Optional |
+| `llm_models_root` | Models root **inside the container**. Resolution order: caller input → value already resolved by `trtllm-case-executor` and forwarded in `job_spec.json` → delegate to `internal-env-info` for the `default_llm_models_root` key when that skill is installed → otherwise **ask the user**. The script builder does not invent a default. | Optional |
+| `checkpoint_path` | Host path to model checkpoint | Optional |
+| `models_path` | Host path to LLM models directory | Optional |
+| `device_type` | Required GPU/device hardware type | Optional |
+| `execution_scenario` | `local_docker`, `local_slurm`, or `remote_slurm` | Always |
+| `persistent_mode` | `true` or `false`. **Only meaningful for `execution_scenario=local_slurm`** — when true, add persistent-mode fields to `job_spec.json` so `exec-local-slurm` can reuse a persistent allocation. Default: `true` for `local_slurm`. Ignored (no-op) for `local_docker` and `local_direct`. For `remote_slurm`, this field is **not used** — `exec-remote-slurm` owns its own persistence model (tmux-held `salloc` + named container; see remote-executor Recipe 6) and does not consume `persistent_mode` from `job_spec.json`. | Optional |
+
+> Per-partition hardware (`arch`, `gres`, `gpus_per_node`, `requires_gres`) lives under `slurm_env.partitions[]`. The chosen `partition` / `account` / `container_image` are passed as top-level inputs.
+
+## Procedure
+
+### Step 1: Resolve Docker Image
+
+1. If the top-level `container_image` input is set (case-executor resolved it as the single source of truth in its Step 3), use it directly — skip the remaining steps.
+2. Search `<repo_root>` for an image-tags properties file (e.g., under `jenkins/` or a similar CI config directory). Extract the relevant key based on arch (see step 3).
+3. Determine compute node architecture:
+
+   **Local Docker** (`execution_scenario == local_docker`):
+   - Run `uname -m` via Bash tool on the local machine.
+
+   **Local Slurm / Remote Slurm**:
+   - Load the SLURM environment from the file at `slurm_env_file` (e.g. `python3 -c 'import json,sys; print(json.load(open(sys.argv[1])))' "$slurm_env_file"` or `jq` — pick whatever's available). The path is canonically `<work_dir>/slurm_env.json`, written by `trtllm-case-executor` Step 2.5. Do **not** invoke any SLURM detection command (`sinfo`, `scontrol`, `srun --mpi=list`) — case-executor already ran the one-pass probe. If `slurm_env_file` is missing or unreadable for a slurm scenario, stop and report that case-executor did not run Step 2.5 before invoking the script builder. The loaded shape is:
+     ```json
+     {
+       "partitions": [
+         {"name": "<part>", "state": "up|down",
+          "time_limit": "...", "nodes": "...",
+          "arch": "x86_64|aarch64|null",
+          "gres": "gpu:8|(null)", "gpus_per_node": 8|null,
+          "requires_gres": true|false}
+       ],
+       "pmix": {"available": ["pmix","pmix_v3","pmix_v4","pmix_v5"], "preferred": "pmix_v5"},
+       "errors": ["..."]
+     }
+     ```
+   - Look up the entry where `name == partition` (the top-level `partition` input, resolved by case-executor Step 2.5); use its `arch` for the image selection below and pass `gres`, `gpus_per_node`, `requires_gres`, and the top-level `pmix.preferred` through to Step 3 / `job_spec.json` (consumed by the executors for the `--mpi=` flag).
+   - If `slurm_env.errors[]` indicates `"SLURM client tools not found"` (the case-executor probe ran on a host without `sinfo`/`scontrol`) or the entry's `arch` is `null`, fall back to the `internal-env-info` skill for arch and leave `pmix.preferred` unset so Step 3 falls back to the bare `pmix` plugin name. If `internal-env-info` is also not installed, skip the lookup silently and use `uname -m` for `local_slurm` or default to `x86_64` for `remote_slurm` when otherwise unknowable. Do not report the missing skill as an error.
+
+   **Arch → properties key mapping** (apply to the lookup result):
+
+   | `Arch=` value | Properties key |
+   |---------------|----------------|
+   | `aarch64` | `LLM_SBSA_DOCKER_IMAGE` |
+   | `x86_64` | `LLM_DOCKER_IMAGE` |
+
+   Use the arch for the chosen `partition` (the target partition for this job) to select the Docker image. If `slurm_env` is missing or its entry's `arch` is null, fall back as described above.
+
+   **GRES interpretation** — applied in Step 3:
+   - `gres` starts with `gpu:` (e.g., `gpu:4`, `gpu:4(IDX:0-3)`) → partition requires `--gres`
+   - `gres` is `(null)` → partition does not use GRES scheduling; omit `--gres` from the script
+
+### Step 2: Resolve Container Mounts
+
+Build comma-separated `<CONTAINER_MOUNTS>`:
+
+**Resolve models host path** — evaluate rules in priority order and stop at the first match. All mount entries below use `<llm_models_root>` as the container-side target — the resolved input from Step 2.5 of `trtllm-case-executor` (see the `llm_models_root` row in [Input](#input-from-orchestrator-prompt)).
+
+1. **`checkpoint_path` provided** → use its **parent directory** as the models host path. Construct one mount entry: `<parent_of_checkpoint_path>:<llm_models_root>`. This takes priority over `slurm_env.mounts`.
+2. **`models_path` provided** → use that as the models host path. Construct: `<models_path>:<llm_models_root>`.
+3. **`slurm_env.mounts` available** → use all entries verbatim (pre-configured for the cluster by `trtllm-case-executor` Step 2.5; includes the models directory and any symlink targets). Do **not** construct a separate models mount entry.
+4. **Fallback** → resolve `<default_models_repo_host>` in this order: `env_check.default_models_repo` (forwarded by `trtllm-case-executor`) → `internal-env-info`'s `default_models_repo_host` key if the skill is installed → otherwise **stop and ask the user** for the host path. Do not hard-code a fallback in this skill. Once resolved, construct: `<default_models_repo_host>:<llm_models_root>`.
+
+**Home directory**: For local, `$HOME`. For remote, resolve from `slurm_env.remote_cwd`'s parent (the user's home on the cluster) or `ssh <host> 'echo $HOME'`.
+
+**Project path**: For local, `<repo_root>`. For remote, `remote_repo_path` (read from `job_spec.json`; constructed by `trtllm-case-executor` as `<default_user_root_dir>/<repo_dir_name>`, where `<default_user_root_dir>` is resolved by case-executor in this order: `internal-env-info` per cluster → fall back to `internal-env-info`'s `default_user_root_dir_example` key when only documentation-level guidance is needed → otherwise case-executor asks the user. The script builder consumes the resolved `remote_repo_path` verbatim and is never hard-coded in this skill). Do **not** use `slurm_env.remote_cwd` — that is the user's root, not the cloned repo path.
+
+**Remote work_dir**: For remote scenarios, `remote_work_dir` is read directly from `job_spec.json` — the case-executor is the single source of truth for that path (`<remote_repo_path>/work_dirs/<basename(local work_dir)>`). The script builder consumes it verbatim and never recomputes.
+
+**Additional mounts** (custom workflow only): If `custom_script` parent directory is not covered by existing mounts, add it.
+
+Assemble: `<home>:<home>,<project>:<project>,<models_mount or slurm_env.mounts entries>[,<extra>:<extra>]`
+
+### Step 3: Resolve Parallelism & Slurm Parameters
+
+**By workflow type:**
+
+**Pytest**:
+- Detect sub-type: perf-sanity vs non-perf-sanity (check if `test_cmd` contains `perf-sanity` or a perf sanity test file path, or if `perf_config_yaml` is provided)
+- Non-perf: inspect test file for `skip_less_device(N)`, `@pytest.mark.parametrize` GPU markers → `required_devices`. If not derivable, use `required_devices` input (default: 1).
+- Non-perf: pytest is always single-node. If `required_devices` > `slurm_env.partitions[name=<partition>].gpus_per_node` (or `available_gpus` from env-check for non-slurm scenarios), **stop and report an error** — pytest does not support multi-node execution.
+- Non-perf: **Hard limit — `required_devices` must be `<= 8`.** If `required_devices > 8`, stop and report an error: `required_devices = N exceeds the maximum of 8 GPUs for pytest.`
+- Non-perf: **Hard limit — `node_count` must be `<= 1`.** Pytest is always single-node; `node_count` is always set to `1`. This is a redundant safety check — if any derivation produces `node_count > 1`, stop and report an error.
+- Perf-sanity: all Slurm parameters (nodes, ntasks-per-node, gpus_per_node) are derived by `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` exclusively from the config YAML hardware section — do not use `slurm_env.partitions[].gpus_per_node`, user input, or any default value
+
+**Eval**:
+- Parse from `eval_cmd`: `--tp_size N` (default: 1), `--pp_size N` (default: 1), `--gpus_per_node N` (overrides cluster config)
+- `world_size` = `tp_size * pp_size`; `required_devices` = `world_size`
+- **Single-node only**: `trtllm-eval` does not support multi-node execution. If `world_size > gpus_per_node` (i.e., the job would need more than one node), **do not generate a Category 5 script** — fall back to the perf-sanity path (Category 4) by generating a benchmark config with `accuracy.enable_accuracy_test: true` and submitting via `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py`. Same fallback applies to disaggregated accuracy testing.
+- Auto-derive `model_name` from last component of `--model` arg (e.g., `Llama-3.1-8B-Instruct` → `llama_3_1_8b_instruct`)
+- Eval task subcommand (e.g., `gsm8k`, `mmlu`) passes through as-is in `eval_cmd`
+- **Model path resolution**: If `--model` is a HuggingFace model ID (contains `/` but does NOT start with `/`), resolve it to a local path under `<llm_models_root>` (the resolved input — see the `llm_models_root` row in [Input](#input-from-orchestrator-prompt) for the priority chain). Search subdirectories (max depth 3) for a directory matching the model name component (case-insensitive). If found, rewrite `--model` in `eval_cmd` to the local absolute path. If not found, keep the original and warn that HF download will be attempted.
+
+**Bench** (trtllm-bench):
+- Parse from the last `trtllm-bench` invocation in `bench_cmd`: `--tp N` (default: 1), `--pp N` (default: 1), `--ep N` (default: 0)
+- `world_size = tp * pp`; `required_devices = world_size`
+- Split `bench_cmd` on `&&` into `PREP_CMD` (the `prepare-dataset` step) and `BENCH_RUN_CMD` (the `throughput`/`latency` step). If no `&&` is present (single command only), set `PREP_CMD` to empty and `BENCH_RUN_CMD` to the full command.
+- Single-node only: if `world_size > gpus_per_node`, stop and report an error — multi-node `trtllm-bench` is not supported.
+- Auto-derive `model_name` from the `--model` or `--model_path` argument
+- `NTASKS = world_size` (used for the benchmark run srun; the `prepare-dataset` srun always overrides to `--ntasks=1`)
+
+**Benchmark**:
+- Derived from config YAML hardware section by `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` — skip manual derivation
+
+**Custom**:
+- Use `required_devices` input (default: 1)
+- **Hard limit — `required_devices` must be `<= 8`.** If `required_devices > 8`, stop and report an error: `required_devices = N exceeds the maximum of 8 GPUs for custom commands.`
+- **Hard limit — `node_count` must be `<= 1`.** Compute `node_count = ceil(required_devices / gpus_per_node)`. If `node_count > 1`, stop and report an error: `required_devices = N requires N nodes (gpus_per_node = M), but custom commands are limited to 1 node. Reduce required_devices to at most M.`
+
+**Common Slurm parameters:**
+- `node_count` = 1 for pytest (non-perf only, always single-node); derived from config YAML by `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` for perf-sanity and benchmark; 1 for eval and bench (single-node only); `ceil(required_devices / gpus_per_node)` for custom, where `gpus_per_node` = `slurm_env.partitions[name=<partition>].gpus_per_node` for slurm scenarios, or `env_check.gpus_per_node` for local_docker; if neither resolves, stop and ask the user.
+- `NTASKS` = 1 for pytest; `world_size` for eval and bench; `required_devices` for custom
+- `NTASKS_PER_NODE` = 1 for pytest; `NTASKS / node_count` for other workflows
+- `PARTITION` = the `partition` input — the value resolved and validated by `trtllm-case-executor` Step 2.5 against the cluster's `slurm_env.partitions[].name`. **Do not fall back to any default** for slurm scenarios; if `partition` is not provided, stop and report that case-executor did not run Step 2.5 before invoking the script builder.
+- `ACCOUNT` = the `account` input — the value resolved and validated by `trtllm-case-executor` Step 2.5 against the cluster's `slurm_env.accounts[]`. **Do not fall back to any hard-coded default or to the `internal-env-info` skill's `default_account`** for slurm scenarios; if `account` is not provided, stop and report that case-executor did not run Step 2.5 before invoking the script builder.
+- `TIME_LIMIT` = `time_limit` input or `02:00:00`
+- `GRES_FLAG`: look up the `gres` value for `<PARTITION>` in the partition → {arch, gres} mapping from Step 1:
+  - `gres` starts with `gpu:` → `GRES_FLAG = --gres=gpu:<GPUS_PER_NODE>`
+  - `gres` is `(null)` or mapping unavailable → `GRES_FLAG = ` (empty string; omit the `#SBATCH --gres` line entirely from the generated script)
+
+**MPI mode:**
+
+Use `pmix.preferred` from `slurm_env` (case-executor Step 2.5's output, also embedded in `job_spec.json`) as the resolved `<MPI_PLUGIN>` (e.g., `pmix_v5`). If the field is unset (detection unavailable or fallback path taken), use the bare plugin name `pmix`. Never hardcode `pmix` — always read the resolved value.
+
+- Pytest: empty (no `--mpi`), clear MPI env vars in script (always single-node)
+- Eval: always `--mpi=pmix` on the eval run srun (including single-task). `trtllm-eval` imports `mpi4py` at module load, which auto-inits MPI; OpenMPI in the container requires srun to provide a PMI/PMIx runtime, otherwise `MPI_Init` aborts on a NULL communicator.
+- Bench (trtllm-bench): `world_size > 1` → `--mpi=<MPI_PLUGIN>` on the benchmark run srun only; `world_size == 1` → empty; `prepare-dataset` srun never uses `--mpi`
+- Custom single-node: empty; multi-node: `--mpi=<MPI_PLUGIN>`
+- Benchmark/perf-sanity: handled by `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py`
+
+### Step 4: Generate Script
+
+Select the Category template from `<skill_dir>/references/trtllm_test_template.md` and generate:
+
+> **Note on `local_direct`:** There is intentionally no Category for `execution_scenario == local_direct`. That scenario runs the command directly on the host without a container or Slurm wrapper, and `trtllm-case-executor` handles it inline (skipping this skill entirely). The Categories below cover only the three scenarios this skill supports: `local_docker`, `local_slurm`, and `remote_slurm`.
+
+**Category 1 (Local Docker)** — `execution_scenario == local_docker`:
+- Read the Category 1 template
+- Substitute placeholders
+- Write Docker command to `<work_dir>/docker_cmd.sh`
+
+**Category 2 (Non-perf pytest)** — `workflow_type == pytest` and non-perf-sanity:
+
+Run `build_slurm_script.py` to generate the script:
+```bash
+python3 <skill_dir>/scripts/build_slurm_script.py pytest \
+  --model-name <MODEL_NAME> \
+  --log-dir <LOG_DIR> \
+  --partition <PARTITION> --account <ACCOUNT> --time-limit <TIME_LIMIT> \
+  --nodes <NODES_NUM> --ntasks <NTASKS> --ntasks-per-node <NTASKS_PER_NODE> \
+  [--gres <GRES>] \
+  --docker-image <DOCKER_IMAGE> \
+  --container-mounts <CONTAINER_MOUNTS> \
+  --container-name <CONTAINER_NAME> \
+  --project-path <PROJECT_PATH> \
+  --llm-models-root <LLM_MODELS_ROOT> \
+  --test-cmd "<TEST_CMD>" \
+  --output <work_dir>/<model_name>_auto_test.slurm
+```
+- Omit `--gres` entirely when `GRES_FLAG` is empty (script omits the `#SBATCH --gres` line when `--gres` is not provided or empty)
+- For remote: use remote paths for `--project-path`, `--log-dir`, `--container-mounts`. Also pass `--local-log-dir <work_dir>` so the script's local `mkdir -p` runs against a path on the orchestrating host (`--log-dir` is cluster-only and not visible locally).
+
+**Category 3 (Custom script)** — `workflow_type == custom`:
+
+Resolve `CUSTOM_CMD` first: `bash <custom_script>` if `custom_script` was given, otherwise the literal `custom_cmd`. Then run:
+```bash
+python3 <skill_dir>/scripts/build_slurm_script.py custom \
+  --model-name <MODEL_NAME> \
+  --log-dir <LOG_DIR> \
+  --partition <PARTITION> --account <ACCOUNT> --time-limit <TIME_LIMIT> \
+  --nodes <NODES_NUM> --ntasks <NTASKS> --ntasks-per-node <NTASKS_PER_NODE> \
+  [--gres <GRES>] \
+  --docker-image <DOCKER_IMAGE> \
+  --container-mounts <CONTAINER_MOUNTS> \
+  --container-name <CONTAINER_NAME> \
+  --project-path <PROJECT_PATH> \
+  --custom-cmd "<CUSTOM_CMD>" \
+  --custom-env "<CUSTOM_ENV>" \
+  --custom-workdir <CUSTOM_WORKDIR> \
+  --run-script <skill_dir>/scripts/slurm_run_custom.sh \
+  --output <work_dir>/<model_name>_custom.slurm
+```
+- For remote: also transfer `slurm_run_custom.sh` to the cluster; use remote path for `--run-script`. Also pass `--local-log-dir <work_dir>` so the local `mkdir -p` lands on the orchestrating host's filesystem.
+
+**Category 4 (Perf-sanity / Benchmark)** — `workflow_type == perf_sanity` or pytest perf-sanity:
+
+> Benchmark requires Slurm — local Docker is not supported.
+
+The caller (trtllm-case-executor) always provides a pre-built `perf_config_yaml`. Config YAML generation is not done here. Use `perf_config_yaml` directly with the repo's `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py`.
+
+Generate the launch script via the repo's `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py`. This script is self-contained — it resolves its own helper scripts (run/install/draft) internally, so no `--run-sh`/`--install-sh`/`--draft-launch-sh` flags are passed:
+```bash
+python3 <repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py \
+  <--test-list "<pytest_test_string>" | --config-file <yaml> [--test-name <name>] [--benchmark-mode <mode>]> \
+  --partition <PARTITION> --account <ACCOUNT> \
+  --job-name <model_name>_<perf|benchmark> \
+  --image <DOCKER_IMAGE> --mounts <CONTAINER_MOUNTS> \
+  --work-dir <work_dir_or_remote_work_dir> \
+  --llm-src <repo_root_or_remote_repo_path> \
+  --llm-models-root <llm_models_root> \
+  --launch-sh <work_dir>/slurm_launch.sh \
+  --time <TIME_LIMIT>
+```
+- `--test-list`: perf-sanity with bracket string; `--config-file` + optional `--test-name`/`--benchmark-mode`: explicit config file
+- Aggregated vs disaggregated execution is selected internally by `submit.py` based on the config.
+- Use `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` (it ships with the cloned repo, so nothing extra is transferred); `--work-dir`, `--llm-src` use **remote paths**, `--launch-sh` uses a **local path**
+
+**Category 6 (Bench)** — `workflow_type == bench`:
+
+```bash
+python3 <skill_dir>/scripts/build_slurm_script.py bench \
+  --model-name <MODEL_NAME> \
+  --log-dir <LOG_DIR> \
+  --partition <PARTITION> --account <ACCOUNT> --time-limit <TIME_LIMIT> \
+  --nodes 1 --ntasks <NTASKS> --ntasks-per-node <NTASKS> \
+  [--gres <GRES>] \
+  --docker-image <DOCKER_IMAGE> \
+  --container-mounts <CONTAINER_MOUNTS> \
+  --container-name <CONTAINER_NAME> \
+  --project-path <PROJECT_PATH> \
+  --bench-cmd "<bench_cmd>" \
+  --output <work_dir>/<model_name>_bench.slurm
+```
+- `<NTASKS>` = `world_size` (tp * pp); single node so `--ntasks` == `--ntasks-per-node`
+- `bench_cmd` passed verbatim — the script splits on `&&` internally; if no `&&`, prep srun runs install only
+- For remote_slurm: pass `--log-dir <remote_work_dir>` (baked into `#SBATCH --output/--error`) and `--local-log-dir <work_dir>` (used for the local `mkdir -p`).
+
+**Category 5 (Eval)** — `workflow_type == eval`:
+
+> **Single-node only.** `trtllm-eval` does not support multi-node execution. If Step 3 computes `node_count > 1` (i.e., `world_size > gpus_per_node`), **do not generate a Category 5 script** — fall back to Category 4 (perf-sanity) with `accuracy.enable_accuracy_test: true` in the benchmark config. Same fallback applies to disaggregated accuracy testing.
+
+```bash
+python3 <skill_dir>/scripts/build_slurm_script.py eval \
+  --model-name <MODEL_NAME> \
+  --log-dir <LOG_DIR> \
+  --partition <PARTITION> --account <ACCOUNT> --time-limit <TIME_LIMIT> \
+  --nodes 1 --ntasks <NTASKS> --ntasks-per-node <NTASKS> \
+  [--gres <GRES>] \
+  --docker-image <DOCKER_IMAGE> \
+  --container-mounts <CONTAINER_MOUNTS> \
+  --container-name <CONTAINER_NAME> \
+  --project-path <PROJECT_PATH> \
+  --eval-cmd "<EVAL_CMD>" \
+  --output <work_dir>/<model_name>_eval.slurm
+```
+- `<NTASKS>` = `world_size` (tp_size * pp_size); MPI mode is always `pmix` on the eval run srun (including single-task) because `trtllm-eval` auto-inits MPI at module load
+- The script installs any required eval dependencies and, when `world_size > 1`, prepends `trtllm-llmapi-launch` to the eval command
+- For remote_slurm: pass `--log-dir <remote_work_dir>` and `--local-log-dir <work_dir>` (same pattern as bench).
+
+### Step 5: Write job_spec.json
+
+Write `<work_dir>/job_spec.json` with all resolved values:
+
+```json
+{
+  "workflow_type": "pytest|eval|custom|benchmark",
+  "model_name": "<MODEL_NAME>",
+  "execution_scenario": "local_docker|local_slurm|remote_slurm",
+  "docker_image": "<resolved image>",
+  "script_path": "<work_dir>/<script_name>",
+  "script_name": "<script_name>",
+  "log_file_pattern": "<MODEL_NAME>_<type>_%j.out",
+  "success_patterns": ["passed", "accuracy:"],
+  "failure_patterns": ["FAILED", "Error", "AssertionError"],
+  "extra_files": [
+    "<skill_dir>/scripts/slurm_run_custom.sh"
+  ],
+  "slurm_params": {
+    "partition": "<PARTITION>",
+    "account": "<ACCOUNT>",
+    "time_limit": "<TIME_LIMIT>",
+    "nodes": 1,
+    "ntasks": 4,
+    "ntasks_per_node": 4,
+    "gpus_per_node": 4,
+    "gres": "gpu:4",
+    "mpi_plugin": "pmix_v5"
+  },
+  "slurm_env": { /* the full parsed JSON loaded from `slurm_env_file` in Step 1: accounts[], partitions[{name,state,time_limit,nodes,arch,gres,gpus_per_node,requires_gres}], default_account, default_partition, pmix.{available,preferred}, errors[]. Embedded verbatim so downstream slurm executors can read arch/gres/pmix from `job_spec.slurm_env` without opening the file or re-detecting. */ }
+}
+```
+
+**Remote-slurm fields (required for `execution_scenario == remote_slurm`):** the case-executor populates these and the script builder echoes them verbatim. The builder MUST NOT recompute or substitute defaults for them.
+
+```json
+{
+  "remote_repo_path": "<default_user_root_dir>/<repo_dir_name>",
+  "remote_work_dir": "<remote_repo_path>/work_dirs/<basename(work_dir)>"
+}
+```
+
+For remote scenarios, `--log-dir <LOG_DIR>` rendered into the slurm script uses `remote_work_dir`, and `--project-path <PROJECT_PATH>` (== `--container-workdir`) uses `remote_repo_path`. If either field is missing from `job_spec.json` for a `remote_slurm` scenario, stop and surface the error — do not fall back to `remote_cwd` or any other default.
+
+**Persistent mode fields:** When `persistent_mode=true` and `execution_scenario=local_slurm`, add these fields to `job_spec.json` (all already resolved during script generation):
+
+```json
+{
+  "persistent_mode": true,
+  "container_name": "<CONTAINER_NAME>",
+  "container_mounts": "<CONTAINER_MOUNTS>"
+}
+```
+
+`container_name` is the resolved `<CONTAINER_NAME>` placeholder value (e.g., `llama_auto_test`, `gpt_custom`, `llama_eval`). `container_mounts` is the resolved comma-separated mount string. `docker_image` and `slurm_params` are already present in the base schema. The `.slurm` script generation is completely unchanged — persistent mode reuses the same scripts.
+
+**Scenario scope:** These persistent-mode fields apply **only to `local_slurm`**. Do not emit them for any other scenario:
+
+- `local_docker` / `local_direct` — no Slurm allocation exists, so `persistent_mode` has no meaning. Omit the fields even if `persistent_mode=true` was passed in.
+- `remote_slurm` — `exec-remote-slurm` runs its own tmux-held `salloc` + named-container persistence flow (see that skill's Recipe 6). It does **not** read `persistent_mode`, `container_name`, or `container_mounts` from `job_spec.json` for that purpose; remote allocation reuse is controlled directly by the remote executor's inputs (`release_allocation`, `alloc_time_limit`, etc.) rather than by this field. Always omit the persistent-mode fields from `job_spec.json` when `execution_scenario=remote_slurm`.
+
+**Pattern selection by workflow type:**
+
+| Workflow | `log_file_pattern` | `success_patterns` | `failure_patterns` |
+|----------|-------------------|-------------------|-------------------|
+| pytest | `<name>_auto_test_%j.out` | `passed` | `FAILED`, `ERROR` |
+| eval | `<name>_eval_%j.out` | `accuracy:`, `score:` | `AssertionError`, `Error` |
+| bench | `<name>_bench_%j.out` | `token/s`, `throughput` | `Error`, `Traceback` |
+| perf_sanity | `slurm-%j.out` | `8_done_` | `FAILED`, `Error` |
+| custom | `<name>_custom_%j.out` | *(none)* | `Error`, `Traceback` |
+
+For **local Docker**, use `docker_cmd` field instead of `script_path`:
+```json
+{
+  "docker_cmd": "<complete docker run command>",
+  "log_file": "<work_dir>/<model_name>_local.log"
+}
+```
+
+### Step 6: Report
+
+Print a summary of what was generated:
+```
+Generated: <work_dir>/job_spec.json
+Script: <work_dir>/<script_name>
+Workflow: <workflow_type>
+Model: <model_name>
+Docker image: <image>
+Devices: <required_devices> GPUs across <node_count> nodes
+```
+
+## Rules
+
+- Always use **absolute paths** in generated scripts and job_spec.json
+- For remote execution, the script must contain **remote paths** for container-workdir, mounts, and log directories — but `job_spec.json` records the **local** script path (the executor skill transfers it)
+- Never execute the script — only generate it. Execution is the executor skill's job.
+- If `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` fails, report the error and stop
+- Validate that required files exist before referencing them (docker image properties, test files, config YAMLs)
diff --git a/.claude/skills/trtllm-test-script-builder/references/trtllm_test_template.md b/.claude/skills/trtllm-test-script-builder/references/trtllm_test_template.md
new file mode 100644
index 000000000000..2261fa9b1774
--- /dev/null
+++ b/.claude/skills/trtllm-test-script-builder/references/trtllm_test_template.md
@@ -0,0 +1,520 @@
+# TensorRT-LLM Test Executor Templates
+
+## Placeholders
+
+All templates below use these placeholders — substitute with actual values before execution.
+
+| Placeholder | Description |
+|-------------|-------------|
+| `<PWD>` | Current working directory (result of `pwd`) |
+| `<HOME_DIR>` | User's home directory (value of `$HOME`) |
+| `<PROJECT_PATH>` | Absolute path to the project/repo root directory |
+| `<MODEL_NAME>` | Short model name from the `model_name` input |
+| `<DOCKER_IMAGE>` | Docker image selected from `jenkins/current_image_tags.properties` |
+| `<TEST_CMD>` | Test command(s) to run inside the container |
+| `<CHECKPOINT_PATH>` | Host path to the model checkpoint (from the `checkpoint_path` input) |
+| `<CHECKPOINT_PARENT_PATH>` | Parent directory of `<CHECKPOINT_PATH>` (e.g., `/data/models` if checkpoint is `/data/models/llama-7b`) |
+| `<MODELS_HOST_PATH>` | Resolved host path for models mount (see Models Mount Resolution in SKILL.md) |
+| `<REQUIRED_DEVICES>` | Number of GPU devices requested |
+| `<NODES_NUM>` | Number of Slurm nodes to allocate (default: 1) |
+| `<NTASKS>` | Total number of tasks (pytest: always 1; other workflows: `<REQUIRED_DEVICES>`) |
+| `<NTASKS_PER_NODE>` | Tasks per node (pytest: always 1; other workflows: `<REQUIRED_DEVICES>` / `<NODES_NUM>`) |
+| `<GPUS_PER_NODE>` | GPUs per node (always 4, regardless of `<NTASKS_PER_NODE>`) |
+| `<CONTAINER_NAME>` | Container name (default: `<MODEL_NAME>_auto_test`) |
+| `<CONTAINER_MOUNTS>` | Comma-separated mount mappings for `--container-mounts` |
+| `<LOG_DIR>` | Directory for log output (default: current working directory) |
+| `<PARTITION>` | Slurm partition name |
+| `<ACCOUNT>` | Slurm account name |
+| `<TIME_LIMIT>` | Slurm job time limit (default: `02:00:00`) |
+| `<MPI_FLAG>` | `--mpi=pmix` for multi-node, empty for single-node |
+| `<GRES_FLAG>` | `--gres=gpu:<GPUS_PER_NODE>` when the partition reports GPU GRES (e.g., `gpu:4`); empty string (entire `#SBATCH --gres` line omitted) when partition reports `(null)` |
+| `<EVAL_CMD>` | Complete `trtllm-eval` command string (Category 5 only) |
+| `<PREP_CMD>` | The `prepare-dataset` step from `bench_cmd` (Category 6 only). Empty string if `bench_cmd` has no `&&`. |
+| `<BENCH_RUN_CMD>` | The `throughput`/`latency` step from `bench_cmd` (Category 6 only). Full `bench_cmd` if no `&&`. |
+
+---
+
+> **Scenario coverage:** Categories below cover only `execution_scenario` values of `local_docker`, `local_slurm`, and `remote_slurm`. The `local_direct` scenario has no Category here — it runs on the host without a container or Slurm wrapper and is handled inline by `trtllm-case-executor` (which skips this skill).
+
+## Category 1: Local Docker Command
+
+Use this template when the environment check returns `satisfied, local`.
+
+### Command template
+
+```bash
+docker run -it --rm \
+  --net=host \
+  --runtime=nvidia \
+  --ipc=host \
+  --cap-add=SYS_PTRACE \
+  --cap-add=SYS_ADMIN \
+  --cap-add=DAC_READ_SEARCH \
+  --security-opt seccomp=unconfined \
+  --env REPO_DIR=tekit_sync \
+  --env HOME_DIR=<PWD> \
+  --env LLM_MODELS_ROOT=/scratch.trt_llm_data/llm-models/ \
+  --mount type=bind,source=<HOME_DIR>,target=<HOME_DIR> \
+  --mount type=bind,source=<PROJECT_PATH>,target=<PROJECT_PATH> \
+  --mount type=bind,source=<MODELS_HOST_PATH>,target=/scratch.trt_llm_data/llm-models/ \
+  --name <MODEL_NAME>_auto_test \
+  <DOCKER_IMAGE> \
+  bash -c '
+    set -euo pipefail
+    cd <PROJECT_PATH>
+    pip install -e . 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -5
+    python -c "import tensorrt_llm; print(\"tensorrt_llm version:\", tensorrt_llm.__version__)" 2>&1 \
+      || { echo "ERROR: TRT-LLM installation verification failed" >&2; exit 1; }
+    <TEST_CMD>
+  '
+```
+
+---
+
+## Category 2: Non-Perf Slurm Job Script
+
+Use this template when the environment check returns `satisfied, slurm` and the test is NOT a perf-sanity test. Supports both single-node and multi-node.
+
+### Script template
+
+Write this to `<MODEL_NAME>_auto_test.slurm` in the current working directory:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=<MODEL_NAME>_auto_test
+#SBATCH --output=<LOG_DIR>/<MODEL_NAME>_auto_test_%j.out
+#SBATCH --error=<LOG_DIR>/<MODEL_NAME>_auto_test_%j.err
+#SBATCH --nodes=<NODES_NUM>
+#SBATCH --ntasks=<NTASKS>
+#SBATCH --ntasks-per-node=<NTASKS_PER_NODE>
+#SBATCH <GRES_FLAG>
+#SBATCH --time=<TIME_LIMIT>
+#SBATCH --partition=<PARTITION>
+#SBATCH --account=<ACCOUNT>
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+# Step 1: Install TRT-LLM (one task per node) — only present for local_slurm
+# (omit entirely for remote_slurm where the image already has TRT-LLM installed)
+srun -l \
+  --ntasks-per-node=1 \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<PROJECT_PATH> \
+  --kill-on-bad-exit=1 \
+  bash -c '
+    set -euo pipefail
+    echo "Node ${SLURM_NODEID:-0}: Installing TRT-LLM from source..."
+    cd <PROJECT_PATH>
+    pip install -e . 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -5
+    echo "Node ${SLURM_NODEID:-0}: Install complete"
+  '
+
+# Step 2: Run tests
+srun -l \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<PROJECT_PATH> \
+  --kill-on-bad-exit=1 \
+  <MPI_FLAG> \
+  bash -c '
+    nvidia-smi
+    # Single-node jobs: clear MPI/SLURM env vars to prevent false MPI init
+    if [ "${SLURM_JOB_NUM_NODES:-1}" -eq 1 ]; then
+      for v in ${!PMI@} ${!PMIX@} ${!MPI@} ${!OMPI@} ${!SLURM@}; do
+        if [ "$v" != "SLURM_PROCID" ]; then
+          unset "$v"
+        fi
+      done
+    fi
+    export LLM_MODELS_ROOT=/scratch.trt_llm_data/llm-models/
+    cd <PROJECT_PATH>
+    <TEST_CMD>
+  '
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+```
+
+### Submission command
+
+```bash
+sbatch <MODEL_NAME>_auto_test.slurm
+```
+
+Report the Slurm job ID. The user can monitor with `squeue -j <JOB_ID>` and view logs at `<MODEL_NAME>_auto_test_<JOB_ID>.out`.
+
+---
+
+## Category 3: Custom Script Slurm Job Script
+
+Use this template when running an arbitrary user script or command (not a pytest test or perf-sanity benchmark) on a Slurm cluster.
+
+### Additional placeholders
+
+| Placeholder | Description |
+|-------------|-------------|
+| `<CUSTOM_CMD>` | The user's command string. If `custom_script` is provided, this is `bash <custom_script>`; otherwise the literal `custom_cmd` value. |
+| `<CUSTOM_ENV>` | Space-separated `KEY=VALUE` pairs to export inside the container (from `custom_env` input). Empty string if not provided. |
+| `<CUSTOM_WORKDIR>` | Working directory for the command inside the container. Defaults to `<PROJECT_PATH>`. |
+| `<RUN_SCRIPT>` | Absolute path to `scripts/slurm_run_custom.sh`. |
+
+### Script template
+
+Write this to `<WORK_DIR>/<MODEL_NAME>_custom.slurm`:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=<MODEL_NAME>_custom
+#SBATCH --output=<LOG_DIR>/<MODEL_NAME>_custom_%j.out
+#SBATCH --error=<LOG_DIR>/<MODEL_NAME>_custom_%j.err
+#SBATCH --nodes=<NODES_NUM>
+#SBATCH --ntasks=<NTASKS>
+#SBATCH --ntasks-per-node=<NTASKS_PER_NODE>
+#SBATCH <GRES_FLAG>
+#SBATCH --time=<TIME_LIMIT>
+#SBATCH --partition=<PARTITION>
+#SBATCH --account=<ACCOUNT>
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+export CUSTOM_COMMAND="<CUSTOM_CMD>"
+export CUSTOM_ENV="<CUSTOM_ENV>"
+export CUSTOM_WORKDIR="<CUSTOM_WORKDIR>"
+export llmSrcNode="<PROJECT_PATH>"
+
+# Step 1: Install TRT-LLM (one task per node)
+srun -l \
+  --ntasks-per-node=1 \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<PROJECT_PATH> \
+  --kill-on-bad-exit=1 \
+  bash -c '
+    set -euo pipefail
+    echo "Node ${SLURM_NODEID:-0}: Installing TRT-LLM from source..."
+    cd <PROJECT_PATH>
+    pip install -e . 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -5
+    echo "Node ${SLURM_NODEID:-0}: Install complete"
+  '
+
+# Step 2: Run custom command
+srun -l \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<CUSTOM_WORKDIR> \
+  --kill-on-bad-exit=1 \
+  <MPI_FLAG> \
+  <RUN_SCRIPT>
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+```
+
+### Submission command
+
+```bash
+sbatch <WORK_DIR>/<MODEL_NAME>_custom.slurm
+```
+
+Report the Slurm job ID. The user can monitor with `squeue -j <JOB_ID>` and view logs at `<MODEL_NAME>_custom_<JOB_ID>.out`.
+
+---
+
+## Category 4: Perf-Sanity Slurm Job Script
+
+Use the repo's `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` to generate launch scripts for perf-sanity tests and benchmarks. This handles both aggregated and disaggregated modes automatically.
+
+`submit.py` is self-contained — it bundles and resolves its own helper scripts (container run, TRT-LLM install, and the aggregated/disaggregated draft templates) internally, so none of those are passed as flags.
+
+### Commands
+
+**Step 1**: Generate the launch script from the pre-built `perf_config_yaml` provided by the caller:
+
+```bash
+python3 <repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py \
+  <--test-list "<pytest_test_string>" | --config-file <yaml> [--test-name <name>] [--benchmark-mode <mode>]> \
+  --partition <PARTITION> --account <ACCOUNT> \
+  --job-name <model_name>_<perf|benchmark> \
+  --image <DOCKER_IMAGE> --mounts <CONTAINER_MOUNTS> \
+  --work-dir <work_dir_or_remote_work_dir> \
+  --llm-src <repo_root_or_remote_repo_path> \
+  --llm-models-root <llm_models_root> \
+  --launch-sh <work_dir>/slurm_launch.sh \
+  --time <TIME_LIMIT>
+```
+
+### Generated output
+
+`submit.py` produces a `slurm_launch.sh` that combines:
+1. `#SBATCH` directives (nodes, gpus, partition, time)
+2. Exported environment variables (pytest commands, hardware config, env vars per role)
+3. `srunArgs` array (container image, mounts, workdir)
+4. Draft template content (aggregated or disaggregated execution logic)
+
+Submit with `sbatch <WORK_DIR>/slurm_launch.sh`.
+
+---
+
+## Category 5: Eval Slurm Job Script
+
+Use this template for running `trtllm-eval` accuracy evaluations on a **single node**. For multi-GPU runs (`ntasks > 1`), wraps `trtllm-eval` with `trtllm-llmapi-launch` and uses `--mpi=pmix`; for single-GPU runs, neither is applied.
+
+**Single-node only.** `trtllm-eval` does not support multi-node execution. If the requested `world_size` exceeds the node's GPUs, fall back to Category 4 (perf-sanity) with `accuracy.enable_accuracy_test: true` in the benchmark config YAML. Same fallback applies to disaggregated accuracy testing.
+
+### Additional placeholders
+
+| Placeholder | Description |
+|-------------|-------------|
+| `<EVAL_CMD>` | The complete `trtllm-eval` command string (e.g., `trtllm-eval --model /path/to/model --tp_size 4 gsm8k`). |
+
+### Script template
+
+Write this to `<WORK_DIR>/<MODEL_NAME>_eval.slurm`:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=<MODEL_NAME>_eval
+#SBATCH --output=<LOG_DIR>/<MODEL_NAME>_eval_%j.out
+#SBATCH --error=<LOG_DIR>/<MODEL_NAME>_eval_%j.err
+#SBATCH --nodes=<NODES_NUM>
+#SBATCH --ntasks=<NTASKS>
+#SBATCH --ntasks-per-node=<NTASKS_PER_NODE>
+#SBATCH <GRES_FLAG>
+#SBATCH --time=<TIME_LIMIT>
+#SBATCH --partition=<PARTITION>
+#SBATCH --account=<ACCOUNT>
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+# Step 1: Install TRT-LLM and eval dependencies (one task per node)
+srun -l \
+  --ntasks-per-node=1 \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<PROJECT_PATH> \
+  --kill-on-bad-exit=1 \
+  bash -c '
+    set -euo pipefail
+    echo "Node ${SLURM_NODEID:-0}: Installing TRT-LLM and eval dependencies..."
+    cd <PROJECT_PATH>
+    pip install -e . 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -3
+    if [ -f examples/trtllm-eval/requirements.txt ]; then
+      pip install -r examples/trtllm-eval/requirements.txt 2>&1 | tail -3
+    fi
+    echo "Node ${SLURM_NODEID:-0}: Install complete"
+  '
+
+# Step 2: Run eval
+# <MPI_FLAG> is --mpi=pmix for ntasks > 1, empty for single-GPU
+# <LAUNCHER> is "trtllm-llmapi-launch " for ntasks > 1, empty for single-GPU
+srun -l \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<PROJECT_PATH> \
+  --kill-on-bad-exit=1 \
+  <MPI_FLAG> \
+  bash -c '
+    set -euo pipefail
+    nvidia-smi
+    export LLM_MODELS_ROOT=/scratch.trt_llm_data/llm-models/
+    cd <PROJECT_PATH>
+    <LAUNCHER><EVAL_CMD>
+  '
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+```
+
+### Submission command
+
+```bash
+sbatch <WORK_DIR>/<MODEL_NAME>_eval.slurm
+```
+
+Report the Slurm job ID. The user can monitor with `squeue -j <JOB_ID>` and view logs at `<MODEL_NAME>_eval_<JOB_ID>.out`.
+
+### Notes
+
+- **Single-node constraint**: `#SBATCH --nodes` is always `1`. If the eval's `world_size` (`tp_size * pp_size`) exceeds the node's GPU count, route to Category 4 instead — do not try to schedule Category 5 across nodes.
+- **Single-GPU vs multi-GPU eval**: for `ntasks > 1`, use `srun --mpi=pmix` and prefix the eval command with `trtllm-llmapi-launch`. For `ntasks == 1` (single-GPU), omit both — `trtllm-llmapi-launch` with pmix on a single task causes unnecessary MPI overhead.
+- **`<GRES_FLAG>`**: Resolved dynamically in Step 1/3 via `sinfo -p <partition> -h -o "%G"`. The entire `#SBATCH --gres` line is omitted when the partition reports `(null)`.
+- **Eval output**: `trtllm-eval` prints accuracy scores to stdout. Only rank 0 produces meaningful output; other ranks are workers. Look for lines containing `accuracy`, `score`, or assertion errors from `--check_accuracy`.
+- **`--config` / `--extra_llm_api_options`**: The `trtllm-eval` CLI supports a YAML config file for additional LLM API options. If the user provides one, ensure it's accessible inside the container (via mounts).
+
+---
+
+## Category 6: Bench (trtllm-bench) Slurm Job Script
+
+Use this template for running `trtllm-bench` throughput/latency benchmarks on a **single node**. The `prepare-dataset` step always runs as a single process (no MPI). The benchmark run step uses `srun --mpi=pmix` + `trtllm-llmapi-launch` only for multi-GPU runs (`ntasks > 1`); single-GPU runs omit both.
+
+**Single-node only.** If `world_size > gpus_per_node`, stop and tell the caller to use the perf-sanity Benchmark (Category 4) workflow instead.
+
+### Additional placeholders
+
+| Placeholder | Description |
+|-------------|-------------|
+| `<PREP_CMD>` | The `prepare-dataset` step from `bench_cmd` (everything before `&&`). Omit the first srun block entirely if empty. |
+| `<BENCH_RUN_CMD>` | The `throughput`/`latency` step from `bench_cmd` (everything after `&&`, or the full command if no `&&`). |
+
+### Script template
+
+Write this to `<WORK_DIR>/<MODEL_NAME>_bench.slurm`:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=<MODEL_NAME>_bench
+#SBATCH --output=<LOG_DIR>/<MODEL_NAME>_bench_%j.out
+#SBATCH --error=<LOG_DIR>/<MODEL_NAME>_bench_%j.err
+#SBATCH --nodes=1
+#SBATCH --ntasks=<NTASKS>
+#SBATCH --ntasks-per-node=<NTASKS_PER_NODE>
+#SBATCH <GRES_FLAG>
+#SBATCH --time=<TIME_LIMIT>
+#SBATCH --partition=<PARTITION>
+#SBATCH --account=<ACCOUNT>
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+# Step 1: Install TRT-LLM and prepare benchmark dataset (always single process)
+srun -l --ntasks=1 \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<PROJECT_PATH> \
+  --kill-on-bad-exit=1 \
+  bash -c '
+    set -euo pipefail
+    nvidia-smi
+
+    echo "Installing TRT-LLM from source..."
+    cd <PROJECT_PATH>
+    pip install -e . 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -5
+    echo "TRT-LLM installed"
+
+    <PREP_CMD>
+  '
+
+# Step 2: Run benchmark
+# <MPI_FLAG> is --mpi=pmix for ntasks > 1, empty for single-GPU
+# <LAUNCHER> is "trtllm-llmapi-launch " for ntasks > 1, empty for single-GPU
+srun -l <MPI_FLAG> \
+  --container-image=<DOCKER_IMAGE> \
+  --container-name=<CONTAINER_NAME> \
+  --container-mounts=<CONTAINER_MOUNTS> \
+  --container-workdir=<PROJECT_PATH> \
+  --kill-on-bad-exit=1 \
+  bash -c '
+    set -euo pipefail
+    cd <PROJECT_PATH>
+    <LAUNCHER><BENCH_RUN_CMD>
+  '
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+```
+
+### Submission command
+
+```bash
+sbatch <WORK_DIR>/<MODEL_NAME>_bench.slurm
+```
+
+Report the Slurm job ID. The user can monitor with `squeue -j <JOB_ID>` and view logs at `<MODEL_NAME>_bench_<JOB_ID>.out`.
+
+### Notes
+
+- **Single-node constraint**: `#SBATCH --nodes` is always `1`. If `world_size > gpus_per_node`, route to Category 4 instead.
+- **`#SBATCH --ntasks=<NTASKS>`** allocates `world_size` tasks so the second srun can use all of them. The first srun explicitly overrides with `--ntasks=1`.
+- **Single-GPU vs multi-GPU bench**: for `ntasks > 1`, use `srun --mpi=pmix` and prefix `<BENCH_RUN_CMD>` with `trtllm-llmapi-launch`. For `ntasks == 1` (single-GPU), omit both — `trtllm-llmapi-launch` with pmix on a single task causes unnecessary MPI overhead.
+- **Container reuse**: both srun calls share `--container-name=<CONTAINER_NAME>`, so TRT-LLM installed in Step 1 is available in Step 2 without reinstalling.
+- **`<GRES_FLAG>`**: Entire `#SBATCH --gres` line is omitted when the partition reports `(null)`.
+- **Bench output**: look for lines containing `token/s` or `throughput` for success; `trtllm-bench` writes a JSON summary to the `--report_json` path.
+
+---
+
+## job_spec.json Schema
+
+Write `<work_dir>/job_spec.json` after generating the script. This manifest is consumed by executor subagents.
+
+### Slurm workflows (Categories 2–5)
+
+```json
+{
+  "workflow_type": "pytest|eval|custom|benchmark",
+  "model_name": "<MODEL_NAME>",
+  "execution_scenario": "local_docker|local_slurm|remote_slurm",
+  "docker_image": "<resolved image>",
+  "script_path": "<work_dir>/<script_name>",
+  "script_name": "<script_name>",
+  "log_file_pattern": "<MODEL_NAME>_<type>_%j.out",
+  "success_patterns": ["passed", "accuracy:"],
+  "failure_patterns": ["FAILED", "Error", "AssertionError"],
+  "extra_files": [
+    "<skill_dir>/scripts/slurm_run_custom.sh"
+  ],
+  "slurm_params": {
+    "partition": "<PARTITION>",
+    "account": "<ACCOUNT>",
+    "time_limit": "<TIME_LIMIT>",
+    "nodes": 1,
+    "ntasks": 4,
+    "ntasks_per_node": 4,
+    "gpus_per_node": 4
+  }
+}
+```
+
+### Local Docker (Category 1)
+
+Use `docker_cmd` and `log_file` instead of `script_path`:
+
+```json
+{
+  "workflow_type": "pytest|eval|custom|benchmark",
+  "model_name": "<MODEL_NAME>",
+  "execution_scenario": "local_docker",
+  "docker_image": "<resolved image>",
+  "docker_cmd": "<complete docker run command>",
+  "log_file": "<work_dir>/<model_name>_local.log",
+  "success_patterns": ["passed", "accuracy:"],
+  "failure_patterns": ["FAILED", "Error", "AssertionError"]
+}
+```
+
+### Pattern selection by workflow type
+
+| Workflow | `log_file_pattern` | `success_patterns` | `failure_patterns` |
+|----------|-------------------|-------------------|-------------------|
+| pytest | `<name>_auto_test_%j.out` | `passed` | `FAILED`, `ERROR` |
+| eval | `<name>_eval_%j.out` | `accuracy:`, `score:` | `AssertionError`, `Error` |
+| bench | `<name>_bench_%j.out` | `token/s`, `throughput` | `Error`, `Traceback` |
+| benchmark | `slurm-%j.out` | `8_done_` | `FAILED`, `Error` |
+| custom | `<name>_custom_%j.out` | *(none)* | `Error`, `Traceback` |
diff --git a/.claude/skills/trtllm-test-script-builder/scripts/build_slurm_script.py b/.claude/skills/trtllm-test-script-builder/scripts/build_slurm_script.py
new file mode 100755
index 000000000000..f5a05451d9de
--- /dev/null
+++ b/.claude/skills/trtllm-test-script-builder/scripts/build_slurm_script.py
@@ -0,0 +1,643 @@
+#!/usr/bin/env python3
+"""Build Slurm job scripts for TRT-LLM test workflows.
+
+Generates scripts for:
+
+  pytest  -> Category 2 (non-perf-sanity pytest)
+  custom  -> Category 3 (arbitrary user command/script)
+  bench   -> Category 6 (trtllm-bench throughput/latency)
+  eval    -> Category 5 (trtllm-eval accuracy)
+
+Not covered:
+  - Category 1 (local Docker) is not a Slurm workflow.
+  - Category 4 (perf-sanity / benchmark) is generated by submit.py, which
+    derives Slurm parameters from the benchmark config YAML.
+
+Argument groups by functionality:
+
+  Slurm directives          -> rendered into #SBATCH lines
+  Container / srun          -> rendered into the srun invocation
+  Category-specific body    -> rendered into the script payload
+
+Usage:
+  build_slurm_script.py pytest --test-cmd "pytest ..."        [common args]
+  build_slurm_script.py custom --custom-cmd "bash foo.sh"
+                               --run-script /path/slurm_run_custom.sh
+                                                              [common args]
+  build_slurm_script.py bench  --bench-cmd "trtllm-bench ... prepare-dataset ... && trtllm-bench ... throughput ..."
+                                                              [common args]
+  build_slurm_script.py eval   --eval-cmd "trtllm-eval ..."   [common args]
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from datetime import datetime
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Templates
+#
+# Placeholders use the <PLACEHOLDER> convention.
+# Substitution is plain string replace -- no format()/Template -- to avoid
+# escaping every bash ${VAR} / ${!PMI@} / ${foo:-0}.
+# ---------------------------------------------------------------------------
+
+# Inserted into the bench-run / eval srun bash block when ntasks == 1, so the
+# process skips MPI init entirely and avoids dynamic-spawn (MPI_ERR_SPAWN)
+# failures under a PMIx-managed step. Empty string when ntasks > 1.
+MPI_CLEANUP_BLOCK = """    # Single-task job: clear MPI/SLURM env vars to prevent MPI init / dynamic spawn
+    for v in ${!PMI@} ${!PMIX@} ${!MPI@} ${!OMPI@} ${!SLURM@}; do
+      if [ "$v" != "SLURM_PROCID" ] && [ "$v" != "SLURM_LOCALID" ]; then
+        unset "$v"
+      fi
+    done
+"""
+
+INSTALL_SRUN_FRAGMENT = """\
+# Step 1: Install TRT-LLM (one task per node)
+<INSTALL_SRUN_BLOCK>
+  bash -c '
+    set -euo pipefail
+    echo "Node ${SLURM_NODEID:-0}: Installing TRT-LLM from source..."
+    cd <PROJECT_PATH>
+    pip install -e . 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -5
+    echo "Node ${SLURM_NODEID:-0}: Install complete"
+  '
+
+"""
+
+EVAL_INSTALL_SRUN_FRAGMENT = """\
+# Step 1: Install TRT-LLM and eval dependencies (one task per node)
+<INSTALL_SRUN_BLOCK>
+  bash -c '
+    set -euo pipefail
+    echo "Node ${SLURM_NODEID:-0}: Installing TRT-LLM and eval dependencies..."
+    cd <PROJECT_PATH>
+    pip install -e . 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -3
+    if [ -f examples/trtllm-eval/requirements.txt ]; then
+      pip install -r examples/trtllm-eval/requirements.txt 2>&1 | tail -3
+    fi
+    echo "Node ${SLURM_NODEID:-0}: Install complete"
+  '
+
+"""
+
+PYTEST_BODY = """\
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+<INSTALL_BLOCK><TEST_STEP_COMMENT><SRUN_BLOCK>
+  bash -c '
+    nvidia-smi
+    # Single-node jobs: clear MPI/SLURM env vars to prevent false MPI init
+    if [ "${SLURM_JOB_NUM_NODES:-1}" -eq 1 ]; then
+      for v in ${!PMI@} ${!PMIX@} ${!MPI@} ${!OMPI@} ${!SLURM@}; do
+        if [ "$v" != "SLURM_PROCID" ]; then
+          unset "$v"
+        fi
+      done
+    fi
+    export LLM_MODELS_ROOT=<LLM_MODELS_ROOT>
+    cd <PROJECT_PATH>
+    <TEST_CMD>
+  '
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+"""
+
+CUSTOM_BODY = """\
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+export CUSTOM_COMMAND="<CUSTOM_CMD>"
+export CUSTOM_ENV="<CUSTOM_ENV>"
+export CUSTOM_WORKDIR="<CUSTOM_WORKDIR>"
+export llmSrcNode="<PROJECT_PATH>"
+
+<INSTALL_BLOCK># Step 2: Run custom command
+<SRUN_BLOCK>
+  <RUN_SCRIPT>
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+"""
+
+BENCH_PREP_FRAGMENT = """\
+# Step 2: Prepare benchmark dataset (single process; MPI mode matches bench-run)
+<PREP_SRUN_BLOCK>
+  bash -c '
+    set -euo pipefail
+    nvidia-smi
+    cd <PROJECT_PATH>
+    <PREP_CMD>
+  '
+
+"""
+
+BENCH_BODY = """\
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+<INSTALL_BLOCK><PREP_BLOCK># Step 3: Run benchmark
+<BENCH_RUN_SRUN_BLOCK>
+  bash -c '
+    set -euo pipefail
+<MPI_CLEANUP>    cd <PROJECT_PATH>
+    <BENCH_RUN_CMD>
+  '
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+"""
+
+EVAL_BODY = """\
+
+set -euo pipefail
+
+echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"
+echo "Nodes: $SLURM_JOB_NUM_NODES, Tasks: $SLURM_NTASKS"
+
+<INSTALL_BLOCK># Step 2: Run eval
+<SRUN_BLOCK>
+  bash -c '
+    set -euo pipefail
+    nvidia-smi
+<MPI_CLEANUP>    export LLM_MODELS_ROOT=<LLM_MODELS_ROOT>
+    cd <PROJECT_PATH>
+    <EVAL_CMD>
+  '
+
+exit_code=$?
+echo "Job $SLURM_JOB_ID finished at $(date) with exit code $exit_code"
+exit $exit_code
+"""
+
+# ---------------------------------------------------------------------------
+# Builders
+# ---------------------------------------------------------------------------
+
+
+def build_sbatch_header(job_name: str, args: argparse.Namespace) -> str:
+    """Render the #SBATCH directives.
+
+    --gres is emitted only when a non-empty value is supplied; otherwise the
+    line is dropped entirely (matches SKILL.md GRES_FLAG handling).
+    """
+    lines = [
+        "#!/bin/bash",
+        f"#SBATCH --job-name={job_name}",
+        f"#SBATCH --output={args.log_dir}/{job_name}_%j.out",
+        f"#SBATCH --error={args.log_dir}/{job_name}_%j.err",
+        f"#SBATCH --nodes={args.nodes}",
+        f"#SBATCH --ntasks={args.ntasks}",
+        f"#SBATCH --ntasks-per-node={args.ntasks_per_node}",
+    ]
+    if args.gres:
+        lines.append(f"#SBATCH --gres={args.gres}")
+    lines += [
+        f"#SBATCH --time={args.time_limit}",
+        f"#SBATCH --partition={args.partition}",
+        f"#SBATCH --account={args.account}",
+    ]
+    return "\n".join(lines) + "\n"
+
+
+def build_srun_block(
+    container_workdir: str,
+    mpi: str,
+    args: argparse.Namespace,
+    ntasks: int | None = None,
+    nodes: int | None = None,
+    ntasks_per_node: int | None = None,
+) -> str:
+    """Render the srun invocation up to (but not including) the payload.
+
+    The caller appends either `bash -c '...'` (pytest/eval) or a script path
+    (custom) on the following line.
+
+    nodes:           when provided, inserts an explicit --nodes override (used by
+        the bench prep step to pin dataset prep to a single node).
+    ntasks:          when provided, inserts an explicit --ntasks override (used by
+        the bench prep step to force single-process dataset prep: -N 1 -n 1).
+    ntasks_per_node: when provided, inserts --ntasks-per-node (used by the
+        install step to run exactly one installer task per allocated node).
+    """
+    lines = ["srun -l \\"]
+    if nodes is not None:
+        lines.append(f"  --nodes={nodes} \\")
+    if ntasks is not None:
+        lines.append(f"  --ntasks={ntasks} \\")
+    if ntasks_per_node is not None:
+        lines.append(f"  --ntasks-per-node={ntasks_per_node} \\")
+    lines += [
+        f"  --container-image={args.docker_image} \\",
+        f"  --container-name={args.container_name} \\",
+        f"  --container-mounts={args.container_mounts} \\",
+        f"  --container-workdir={container_workdir} \\",
+        "  --kill-on-bad-exit=1 \\",
+    ]
+    if mpi and mpi != "none":
+        lines.append(f"  --mpi={mpi} \\")
+    return "\n".join(lines)
+
+
+def substitute(template: str, values: dict) -> str:
+    out = template
+    for k, v in values.items():
+        out = out.replace(f"<{k}>", v)
+    return out
+
+
+def build_pytest(args: argparse.Namespace) -> str:
+    job_name = f"{args.model_name}_auto_test"
+    install_srun = build_srun_block(args.project_path, "none", args, ntasks=1, ntasks_per_node=1)
+    install_block = substitute(
+        INSTALL_SRUN_FRAGMENT,
+        {
+            "INSTALL_SRUN_BLOCK": install_srun,
+            "PROJECT_PATH": args.project_path,
+        },
+    )
+    body = substitute(
+        PYTEST_BODY,
+        {
+            "INSTALL_BLOCK": install_block,
+            "TEST_STEP_COMMENT": "# Step 2: Run tests\n",
+            "SRUN_BLOCK": build_srun_block(
+                args.container_workdir, args.mpi, args, ntasks=1, ntasks_per_node=1
+            ),
+            "LLM_MODELS_ROOT": args.llm_models_root,
+            "PROJECT_PATH": args.project_path,
+            "TEST_CMD": args.test_cmd,
+        },
+    )
+    return build_sbatch_header(job_name, args) + body
+
+
+def build_custom(args: argparse.Namespace) -> str:
+    job_name = f"{args.model_name}_custom"
+    custom_workdir = args.custom_workdir or args.project_path
+    install_srun = build_srun_block(args.project_path, "none", args, ntasks=1, ntasks_per_node=1)
+    install_block = substitute(
+        INSTALL_SRUN_FRAGMENT,
+        {
+            "INSTALL_SRUN_BLOCK": install_srun,
+            "PROJECT_PATH": args.project_path,
+        },
+    )
+    body = substitute(
+        CUSTOM_BODY,
+        {
+            "INSTALL_BLOCK": install_block,
+            "SRUN_BLOCK": build_srun_block(
+                custom_workdir, args.mpi, args, ntasks=1, ntasks_per_node=1
+            ),
+            "CUSTOM_CMD": args.custom_cmd,
+            "CUSTOM_ENV": args.custom_env,
+            "CUSTOM_WORKDIR": custom_workdir,
+            "PROJECT_PATH": args.project_path,
+            "RUN_SCRIPT": args.run_script,
+        },
+    )
+    return build_sbatch_header(job_name, args) + body
+
+
+def _resolve_pmix(args: argparse.Namespace) -> str:
+    # bench/eval auto-promote to a pmix plugin when --ntasks > 1. Honor an
+    # explicit pmix variant passed via --mpi (from slurm_env.pmix.preferred,
+    # populated by case-executor's detect_slurm_env.sh, e.g. 'pmix_v5');
+    # otherwise fall back to the bare 'pmix' plugin name.
+    if isinstance(args.mpi, str) and args.mpi.startswith("pmix"):
+        return args.mpi
+    return "pmix"
+
+
+def _split_bench_cmd(bench_cmd: str) -> tuple[str, str]:
+    """Split a bench_cmd of the form '<prep> && <run>' into (prep, run).
+
+    When no '&&' is present, returns ("", bench_cmd) so the prep srun block
+    runs the install only (no dataset prep) and the full command runs as the
+    benchmark step. Raises SystemExit when more than one '&&' is found.
+    """
+    parts = [p.strip() for p in bench_cmd.split("&&")]
+    if len(parts) == 1:
+        return "", parts[0]
+    if len(parts) != 2:
+        raise SystemExit(
+            "bench: --bench-cmd must contain at most one '&&' separator "
+            "between the prepare-dataset step and the throughput/latency step. "
+            f"Got {len(parts)} part(s)."
+        )
+    return parts[0], parts[1]
+
+
+def build_bench(args: argparse.Namespace) -> str:
+    prep_cmd, bench_run_cmd = _split_bench_cmd(args.bench_cmd)
+    if args.ntasks > 1 and "trtllm-llmapi-launch" not in bench_run_cmd:
+        bench_run_cmd = "trtllm-llmapi-launch " + bench_run_cmd
+
+    # Step 1: Install — single-process pip install; no MPI needed.
+    install_srun = build_srun_block(args.project_path, "none", args, ntasks=1, ntasks_per_node=1)
+    install_block = substitute(
+        INSTALL_SRUN_FRAGMENT,
+        {
+            "INSTALL_SRUN_BLOCK": install_srun,
+            "PROJECT_PATH": args.project_path,
+        },
+    )
+
+    # Step 2: prepare-dataset always runs on exactly one process on one node
+    # (-N 1 -n 1) regardless of how many nodes/GPUs the bench allocates.
+    # Step 3: bench-run uses pmix only when ntasks > 1.
+    # trtllm-bench imports mpi4py at module load, which auto-inits MPI when PMI
+    # env vars are present. For multi-task runs OpenMPI requires srun to provide
+    # a PMI/PMIx runtime. For single-task runs, dropping --mpi=pmix and clearing
+    # PMI/SLURM env vars in the bash block lets the process skip MPI init
+    # entirely and avoids dynamic-spawn (MPI_ERR_SPAWN) failures under PMIx.
+    run_mpi = "pmix" if args.ntasks > 1 else "none"
+    mpi_cleanup = "" if args.ntasks > 1 else MPI_CLEANUP_BLOCK
+    prep_srun = build_srun_block(
+        args.project_path, run_mpi, args, ntasks=1, nodes=1, ntasks_per_node=1
+    )
+    bench_run_srun = build_srun_block(
+        args.project_path, run_mpi, args, ntasks=args.ntasks, ntasks_per_node=args.ntasks_per_node
+    )
+
+    prep_block = substitute(
+        BENCH_PREP_FRAGMENT,
+        {
+            "PREP_SRUN_BLOCK": prep_srun,
+            "PROJECT_PATH": args.project_path,
+            "PREP_CMD": prep_cmd,
+        },
+    )
+
+    job_name = f"{args.model_name}_bench"
+    body = substitute(
+        BENCH_BODY,
+        {
+            "INSTALL_BLOCK": install_block,
+            "PREP_BLOCK": prep_block,
+            "BENCH_RUN_SRUN_BLOCK": bench_run_srun,
+            "PROJECT_PATH": args.project_path,
+            "BENCH_RUN_CMD": bench_run_cmd,
+            "MPI_CLEANUP": mpi_cleanup,
+        },
+    )
+    return build_sbatch_header(job_name, args) + body
+
+
+def build_eval(args: argparse.Namespace) -> str:
+    # trtllm-eval is single-node only. Multi-node accuracy must go through the
+    # perf-sanity path (submit.py with accuracy.enable_accuracy_test: true).
+    if args.nodes > 1:
+        raise SystemExit(
+            "eval: trtllm-eval is single-node only (got --nodes="
+            f"{args.nodes}). For multi-node accuracy, fall back to the "
+            "perf-sanity path: generate a benchmark config with "
+            "accuracy.enable_accuracy_test: true and submit via "
+            "submit.py."
+        )
+    eval_cmd = args.eval_cmd
+    if args.ntasks > 1 and "trtllm-llmapi-launch" not in eval_cmd:
+        eval_cmd = "trtllm-llmapi-launch " + eval_cmd
+    # trtllm-eval imports mpi4py at module load, which auto-inits MPI when PMI
+    # env vars are present. For multi-task runs OpenMPI requires srun to provide
+    # a PMI/PMIx runtime. For single-task runs, dropping --mpi=pmix and clearing
+    # PMI/SLURM env vars in the bash block lets the process skip MPI init
+    # entirely and avoids dynamic-spawn (MPI_ERR_SPAWN) failures under PMIx.
+    mpi = "pmix" if args.ntasks > 1 else "none"
+    mpi_cleanup = "" if args.ntasks > 1 else MPI_CLEANUP_BLOCK
+    job_name = f"{args.model_name}_eval"
+    install_srun = build_srun_block(args.project_path, "none", args, ntasks=1, ntasks_per_node=1)
+    install_block = substitute(
+        EVAL_INSTALL_SRUN_FRAGMENT,
+        {
+            "INSTALL_SRUN_BLOCK": install_srun,
+            "PROJECT_PATH": args.project_path,
+        },
+    )
+    body = substitute(
+        EVAL_BODY,
+        {
+            "INSTALL_BLOCK": install_block,
+            "SRUN_BLOCK": build_srun_block(
+                args.project_path,
+                mpi,
+                args,
+                ntasks=args.ntasks,
+                ntasks_per_node=args.ntasks_per_node,
+            ),
+            "PROJECT_PATH": args.project_path,
+            "LLM_MODELS_ROOT": args.llm_models_root,
+            "EVAL_CMD": eval_cmd,
+            "MPI_CLEANUP": mpi_cleanup,
+        },
+    )
+    return build_sbatch_header(job_name, args) + body
+
+
+CATEGORY_BUILDERS = {
+    "pytest": build_pytest,
+    "custom": build_custom,
+    "bench": build_bench,
+    "eval": build_eval,
+}
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def _add_common_args(p: argparse.ArgumentParser) -> None:
+    sbatch = p.add_argument_group(
+        "Slurm directives (#SBATCH lines)",
+        "Values rendered literally into the #SBATCH header.",
+    )
+    sbatch.add_argument(
+        "--model-name", required=True, help="Short model name; becomes the job-name prefix."
+    )
+    sbatch.add_argument(
+        "--log-dir",
+        default=None,
+        help=(
+            "Directory rendered into #SBATCH --output/--error "
+            "in the .slurm script. For remote_slurm runs this "
+            "is the REMOTE path on the cluster filesystem; the "
+            "script does NOT create this directory locally. "
+            "Required for pytest/custom. For bench/eval, "
+            "defaults to "
+            "<project-path>/work_dirs/<model-name>_<timestamp>."
+        ),
+    )
+    sbatch.add_argument(
+        "--local-log-dir",
+        default=None,
+        help=(
+            "Path on the LOCAL (orchestrating) node to mkdir "
+            "before writing the script. Use this for "
+            "remote_slurm runs where --log-dir points to a "
+            "cluster-only path. When unset, defaults to "
+            "--log-dir (preserves local_slurm/local_docker "
+            "behavior)."
+        ),
+    )
+    sbatch.add_argument("--partition", required=True)
+    sbatch.add_argument("--account", required=True)
+    sbatch.add_argument("--time-limit", default="02:00:00")
+    sbatch.add_argument("--nodes", type=int, default=1)
+    sbatch.add_argument("--ntasks", type=int, default=1)
+    sbatch.add_argument("--ntasks-per-node", type=int, default=1)
+    sbatch.add_argument(
+        "--gres",
+        default="",
+        help="Value for --gres (e.g. 'gpu:4'). Empty omits the line entirely.",
+    )
+
+    container = p.add_argument_group(
+        "Container / srun",
+        "Values rendered into the srun invocation.",
+    )
+    container.add_argument("--docker-image", required=True)
+    container.add_argument(
+        "--container-mounts", required=True, help="Comma-separated bind mounts for pyxis."
+    )
+    container.add_argument(
+        "--container-name", default=None, help="Default: <model-name>_<category>."
+    )
+    container.add_argument("--container-workdir", default=None, help="Default: --project-path.")
+    container.add_argument(
+        "--project-path", required=True, help="Absolute path to the TRT-LLM repo root."
+    )
+    container.add_argument(
+        "--mpi",
+        default="none",
+        help=(
+            "MPI plugin for srun. 'none' disables --mpi. Otherwise the value "
+            "is passed through verbatim (e.g. 'pmix', 'pmix_v3', 'pmix_v5'); "
+            "resolve it from slurm_env.pmix.preferred (populated by "
+            "case-executor's detect_slurm_env.sh). For bench/eval, "
+            "overridden to a pmix variant when --ntasks > 1."
+        ),
+    )
+    container.add_argument(
+        "--llm-models-root",
+        default="/scratch.trt_llm_data/llm-models/",
+        help="Path exported as LLM_MODELS_ROOT inside the container.",
+    )
+
+    output = p.add_argument_group("Output")
+    output.add_argument("--output", default=None, help="Write to this path. Default: stdout.")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    sub = p.add_subparsers(dest="category", required=True, metavar="{pytest,custom,bench,eval}")
+
+    pp = sub.add_parser("pytest", help="Non-perf-sanity pytest (Category 2)")
+    _add_common_args(pp)
+    body = pp.add_argument_group("pytest body")
+    body.add_argument(
+        "--test-cmd", required=True, help="Pytest command string run inside the container."
+    )
+
+    cp = sub.add_parser("custom", help="Custom command/script (Category 3)")
+    _add_common_args(cp)
+    body = cp.add_argument_group("custom body")
+    body.add_argument(
+        "--custom-cmd", required=True, help="Command string. 'bash <script>' when wrapping a file."
+    )
+    body.add_argument(
+        "--custom-env", default="", help="Space-separated KEY=VALUE pairs exported in-container."
+    )
+    body.add_argument(
+        "--custom-workdir",
+        default=None,
+        help="In-container cwd for the command. Default: --project-path.",
+    )
+    body.add_argument(
+        "--run-script", required=True, help="In-container path to slurm_run_custom.sh."
+    )
+
+    bp = sub.add_parser("bench", help="trtllm-bench throughput/latency (Category 6)")
+    _add_common_args(bp)
+    body = bp.add_argument_group("bench body")
+    body.add_argument(
+        "--bench-cmd",
+        required=True,
+        help=(
+            "Two trtllm-bench commands joined by '&&': "
+            "'trtllm-bench ... prepare-dataset ... && trtllm-bench ... throughput ...'. "
+            "The first (prepare-dataset) always runs single-process (-N 1 -n 1). "
+            "The second (throughput/latency) is the main benchmark run."
+        ),
+    )
+
+    ep = sub.add_parser("eval", help="trtllm-eval accuracy workflow (Category 5)")
+    _add_common_args(ep)
+    body = ep.add_argument_group("eval body")
+    body.add_argument("--eval-cmd", required=True, help="Full trtllm-eval command.")
+
+    return p
+
+
+def _apply_defaults(args: argparse.Namespace) -> None:
+    if args.container_name is None:
+        args.container_name = f"{args.model_name}_{args.category}"
+    if args.container_workdir is None:
+        args.container_workdir = args.project_path
+    if args.log_dir is None:
+        if args.category in {"bench", "eval"}:
+            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+            args.log_dir = os.path.join(args.project_path, "work_dirs", f"{args.model_name}_{ts}")
+        else:
+            sys.exit(f"--log-dir is required for category '{args.category}'")
+    # For remote_slurm, --log-dir is a cluster-only path that does not exist on
+    # the orchestrating host. Callers should pass --local-log-dir pointing to
+    # the local work_dir; we mkdir that path here. The --log-dir value still
+    # gets baked into the rendered #SBATCH --output/--error lines.
+    mkdir_target = args.local_log_dir or args.log_dir
+    os.makedirs(mkdir_target, exist_ok=True)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    _apply_defaults(args)
+
+    script = CATEGORY_BUILDERS[args.category](args)
+
+    if args.output:
+        out_path = Path(args.output)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_text(script)
+        os.chmod(out_path, 0o755)
+        print(f"Wrote {out_path}", file=sys.stderr)
+    else:
+        sys.stdout.write(script)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/.claude/skills/trtllm-test-script-builder/scripts/slurm_run_custom.sh b/.claude/skills/trtllm-test-script-builder/scripts/slurm_run_custom.sh
new file mode 100755
index 000000000000..75bd3e1ff701
--- /dev/null
+++ b/.claude/skills/trtllm-test-script-builder/scripts/slurm_run_custom.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+# Set up error handling
+set -xEeuo pipefail
+trap 'rc=$?; echo "Error in file ${BASH_SOURCE[0]} on line $LINENO: $BASH_COMMAND (exit $rc)"; exit $rc' ERR
+
+# TRT-LLM is installed in a separate srun step before this script runs.
+
+# Export user-defined env vars
+if [ -n "${CUSTOM_ENV:-}" ]; then
+    eval "export $CUSTOM_ENV"
+fi
+
+# Single-node runs: strip MPI/SLURM env vars to prevent phantom MPI init
+if [ "${SLURM_JOB_NUM_NODES:-1}" -eq 1 ]; then
+    for v in ${!PMI@} ${!PMIX@} ${!MPI@} ${!OMPI@} ${!SLURM@}; do
+        if [ "$v" != "SLURM_PROCID" ]; then
+            unset "$v"
+        fi
+    done
+fi
+
+# Change to project root (not tests/integration/defs)
+cd "${CUSTOM_WORKDIR:-$llmSrcNode}"
+
+# Turn off "exit on error" so we capture the exit code
+set +e
+
+eval "$CUSTOM_COMMAND"
+exit_code=$?
+echo "Rank${SLURM_PROCID:-0} Custom command finished with exit code $exit_code"
+
+exit $exit_code
diff --git a/.claude/skills/trtllm-test-specialist/SKILL.md b/.claude/skills/trtllm-test-specialist/SKILL.md
new file mode 100644
index 000000000000..1e6abb8dea58
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/SKILL.md
@@ -0,0 +1,681 @@
+---
+name: trtllm-test-specialist
+description: >
+  Runs model-level and module-level tests for TensorRT-LLM. First classifies
+  the test scope (module test or model test), then dispatches to the appropriate
+  workflow. Model tests are further classified by type (functionality/smoke test,
+  benchmark, or evaluation). Prompts the user for parallelism parameters
+  (tp, ep, dp), dataset paths, device type, or a config file as needed. All
+  test execution is delegated to trtllm-case-executor.
+tags:
+  - tensorrt-llm
+  - testing
+  - modeling
+license: Apache-2.0
+metadata:
+  author: NVIDIA Corporation
+---
+
+# TensorRT-LLM Model Test Specialist
+
+> **Execution mode:** After the Step 0a parameter-confirmation pause (the one mandatory stop — see Step 0a), run all remaining steps end-to-end. Invoke sub-skills (`trtllm-case-executor`, `exec-env-check`, `trtllm-test-script-builder`, etc.) sequentially and proceed immediately to the next step as soon as each returns. Do not stop to summarize intermediate results or ask the user for additional input unless a required parameter is genuinely missing.
+
+This skill handles perf-sanity, module-level, and model-level testing. It classifies the test scope, builds the appropriate test commands, and delegates execution to `trtllm-case-executor`.
+
+**Three test scopes:**
+
+- **Perf-sanity benchmark test** — `test_cmd` targets `test_perf_sanity.py` (contains `test_perf_sanity` in the path or `test_e2e[` in the selector). Routed to the Benchmark Test Workflow's pre-built path. No ≤8-GPU limit. `device_type` and `model_name` auto-derived from the test ID.
+- **Module test** — run a caller-provided test command for a module. The skill does not generate or author test scripts.
+- **Model test** — end-to-end model-level tests:
+  - **Functionality test (smoke test)** — `pytest examples/pytorch/quickstart_advanced.py` with pytorch backend.
+  - **Benchmark test** — performance benchmarking. Single-node uses `trtllm-bench` by default; multi-node falls back to the perf-sanity config-YAML path.
+  - **Evaluation test** — accuracy evaluation. Single-node uses `trtllm-eval` by default; multi-node falls back to the perf-sanity config-YAML path (`accuracy.enable_accuracy_test: true`).
+
+---
+
+## Config File
+
+`config_file` is a standalone super parameter. **The file must be a valid YAML file with a `.yaml` or `.yml` extension.** When provided, it is the **sole** source of parameters — all values specified directly in the prompt are dropped.
+
+If the user provided other parameters alongside `config_file`, issue this warning before doing anything else:
+
+> ⚠️ **Config file provided — all directly specified parameters will be dropped.**
+> The following parameters you specified will be ignored: `<param1>`, `<param2>`, ...
+> Only values from `<config_file>` will be used. To override individual settings, edit the config file directly.
+
+Then proceed:
+
+1. Parse the YAML file using `scripts/parse_config.py`.
+2. Use the parsed values as the **sole** authoritative parameter set. Discard all parameters supplied directly in the prompt.
+3. All subsequent steps operate on the config-file parameter set only.
+
+Valid keys are the input names listed in the [Inputs](#inputs) section. Unrecognized keys are warned about and ignored. A full template with all supported fields and inline documentation is at `references/test_config_template.yaml`.
+
+---
+
+## Inputs
+
+### Common Inputs
+
+| Input | Description | Required | Example |
+|-------|-------------|----------|---------|
+| `model_name` | Short model name | **Yes** (not required for pre-built perf-sanity — derived from test ID) | `llama_3_1_8b` |
+| `checkpoint_path` | Path to the HuggingFace checkpoint directory or model name | **Yes** (not required for pre-built perf-sanity) | `/models/Llama-3.1-8B-Instruct` or `meta-llama/Llama-3.1-8B-Instruct` |
+| `repo_path` | Root of the TensorRT-LLM repo (default: current working directory) | No | `/home/user/TensorRT-LLM` |
+| `report_file` | Path to write the test report (markdown). Parent directories are created automatically. Default: `./<MODEL_NAME>-auto-test-report.md` | No | `bring_up/llama/test_results/single_gpu_summary.md` |
+| `config_file` | Path to a YAML super-parameter file. See [Config File](#config-file). | No | `configs/benchmark_llama_8b.yaml` |
+
+### Module Test Inputs
+
+| Input | Description | Required | Example |
+|-------|-------------|----------|---------|
+| `test_cmd` | The exact test command to run (e.g. a `pytest` invocation). The skill runs this verbatim via `trtllm-case-executor` and does not generate test scripts. | **Yes** (module test) | `pytest tests/unittest/_torch/modules/test_attention.py -v` |
+| `class_name` | Name of the class under test (used for reporting only) | No | `Attention` or `LlamaAttention` |
+| `required_devices` | Number of GPUs required to run `test_cmd` | No (default: `1`) | `1`, `4` |
+
+### Model Test Inputs
+
+| Input | Description | Required | Example |
+|-------|-------------|----------|---------|
+| `test_type` | Type of model test: `functionality`, `benchmark`, or `evaluation` | No (auto-classified) | `functionality` |
+| `mtp_layers` | Comma-separated MTP speculation depths to benchmark (one run per value). Triggered only by explicit "mtp layers" or "speculative decoding layers" mentions. | No | `1,2,3,4,5,6` |
+| `tp_size` | Tensor parallelism size | No (default: `1`) | `4` |
+| `ep_size` | Expert parallelism size (MoE models) | No (default: `0`) | `8` |
+| `dp_size` | Data parallelism size | No (default: `1`) | `2` |
+| `pp_size` | Pipeline parallelism size | No (default: `1`) | `2` |
+| `device_type` | Target GPU device type | No | `GB200`, `B200` |
+| `dataset_path` | Path to evaluation or benchmark dataset file | No (benchmark/eval) | `/data/datasets/gsm8k.json` |
+| `benchmark_config_yaml` | Path to an existing benchmark config YAML file (perf-sanity fallback) | No (benchmark) | `examples/disaggregated/slurm/benchmark/config.yaml` |
+| `perf_config_yaml` | Path to a pre-built perf-sanity config YAML. Used with pre-built perf-sanity `test_cmd` to pass node/GPU sizing to `trtllm-case-executor`. | No (perf-sanity) | `/work/deepseek_v3_disagg.yaml` |
+| `bench_subcommand` | `trtllm-bench` subcommand: `throughput` or `latency` | No (default: `throughput`) | `latency` |
+| `backend` | Backend for `trtllm-bench`: `pytorch`, `tensorrt`, or `_autodeploy` | No (default: `pytorch`) | `pytorch` |
+| `num_requests` | Synthetic dataset size when `dataset_path` is not provided | No (default: `100`) | `512` |
+| `input_mean` / `output_mean` | Mean input / output token lengths for synthetic dataset | No (defaults: `128` / `128`) | `1024` / `1024` |
+| `concurrency` | `--concurrency` flag for `trtllm-bench` | No | `32` |
+| `max_batch_size` / `max_num_tokens` | Runtime limits for `trtllm-bench` | No | `64` / `8192` |
+| `extra_llm_api_options_yaml` | Path to a YAML file passed via `--config` / `--extra_llm_api_options` to override LLM API settings. Accepted by both `trtllm-bench` (all subcommands) and `trtllm-eval` (all tasks). | No | `/tmp/pytorch_extra.yaml` |
+| `eval_tasks` | Evaluation benchmark tasks (comma-separated) | No (evaluation) | `gsm8k,gpqa` |
+
+---
+
+## Step 0: Load Config File
+
+If `config_file` is provided:
+
+1. If any other parameters were also supplied in the prompt, output the warning described in [Config File](#config-file) before proceeding. List every extra parameter the user specified.
+2. Run the parser script via Bash:
+   ```bash
+   python3 <skill_dir>/scripts/parse_config.py --config-file <config_file>
+   ```
+   The script exits non-zero on any error (file not found, wrong extension, invalid YAML). If it fails, stop and report the error to the user before proceeding.
+3. Parse the JSON output. The output has three top-level keys:
+   - `params` — recognised parameters found in the config file plus any scope-specific defaults that were filled in.
+   - `missing_required` — required parameters still absent after defaults were applied (list of strings).
+   - `defaults_applied` — parameters whose values were filled in from built-in defaults, not from the config file (list of strings; informational only).
+4. If `missing_required` is non-empty, **stop immediately**. Report to the user which required parameters are missing, and instruct them to add those parameters to the config file before retrying. Example message:
+   > The following required parameters are missing from the config file. Please add them to `<config_file>` and re-run:
+   > - `checkpoint_path` — Path to the HuggingFace checkpoint directory or model name
+   > - `tp_size` — Tensor parallelism size
+   Do **not** proceed to any subsequent step.
+5. Use the values in `params` as the **sole** parameter set. All parameters supplied directly in the prompt are discarded.
+
+All subsequent steps operate on the config-file parameter set only.
+
+---
+
+## Step 0a: Log Resolved Parameters
+
+After all parameters are resolved — from `config_file` (Step 0) or directly from user input — output a table of every resolved parameter before proceeding to Step 1. This makes the active configuration visible and auditable.
+
+**Table format:**
+
+| Parameter | Value | Source |
+|-----------|-------|--------|
+| `model_name` | `llama_3_1_8b` | `user_input` |
+| `checkpoint_path` | `/models/Llama-3.1-8B-Instruct` | `config_file` |
+| `tp_size` | `4` | `config_file` |
+| `ep_size` | `0` | `default` |
+| `pp_size` | `1` | `default` |
+
+**Source** column values:
+
+| Source | Meaning |
+|--------|---------|
+| `config_file` | Value was read from the YAML config file |
+| `user_input` | Value was provided directly in the prompt |
+| `default` | Value was auto-filled from built-in defaults (corresponds to `defaults_applied` entries when a config file was parsed) |
+| `derived` | Value was auto-derived (e.g., `model_name` inferred from `checkpoint_path`) |
+
+Include all parameters that have a resolved value. Omit parameters that are not set and have no applicable default.
+
+After printing the table, **pause and ask the user to confirm before proceeding**:
+
+> The above parameters will be used for this test run. Do you want to continue?
+
+Wait for an explicit confirmation (`yes`, `y`, `ok`, `proceed`, or equivalent) before advancing to Step 1. If the user says no or requests changes, apply the requested changes and re-display the updated table before asking again.
+
+---
+
+## Step 1: Classify Test Scope
+
+Evaluate in this order (first match wins):
+
+1. **Perf-sanity benchmark test** — `test_cmd` is provided AND (`test_cmd` contains `test_perf_sanity` in the file path OR `test_e2e[` in the test selector). → **Benchmark Test Workflow** (pre-built perf-sanity path, see Step 0 of that workflow)
+2. **Module test** — `test_cmd` is provided. → **Module Test Workflow**
+3. **Model test** — no `test_cmd`. Classify further in Step 1a.
+4. **Ambiguous** — ask the user to clarify.
+
+### Step 1a: Classify Model Test Type
+
+**Layer disambiguation:** "mtp layers N,M,…" or "speculative decoding layers N,M,…" are interpreted as MTP speculation depths (`mtp_layers`) and route to the **Benchmark/Evaluation test**.
+
+1. **Evaluation test** — user requests accuracy evaluation, or `eval_tasks` is provided.
+2. **Benchmark test** — user requests benchmarking, `mtp_layers` is provided, or `benchmark_config_yaml`/`benchmark_model_name` is provided.
+3. **Functionality test** — all other model-level test requests (default).
+
+---
+
+## Module Test Workflow
+
+Runs a caller-provided test command against a TRT-LLM module. The skill does
+**not** search for existing tests, generate test scripts, or modify source
+files — it only executes the given `test_cmd` via `trtllm-case-executor`.
+
+### Step 1: Validate Input
+
+- `test_cmd` is required. If missing, ask the user for the exact command to run.
+
+### Step 1a: Derive `required_devices` and `device_type` from Test Markers
+
+When `required_devices` or `device_type` are **not** provided by the user, run:
+
+```bash
+python3 <skill_dir>/scripts/extract_test_markers.py --test-cmd "<test_cmd>"
+```
+
+The script parses the target class and function in the test file and returns JSON:
+
+```json
+{
+  "required_devices": 4,
+  "device_type": "Blackwell (B200/GB200)",
+  "sources": {
+    "required_devices": "derived",
+    "device_type": "derived"
+  }
+}
+```
+
+- If the user already provided `required_devices` or `device_type`, **skip** running the script for those fields — keep the user values with source = `user_input`.
+- Use the `sources` field from the JSON to set the source column in the parameter table (`derived` or `default`).
+- If `device_type` is an empty string, the field is unconstrained — omit it from the parameter table.
+
+Update the parameter table in Step 0a with the values before proceeding.
+
+### Step 1b: Validate Resource Limits (pytest / module test)
+
+After `required_devices` is resolved, enforce these hard limits before proceeding:
+
+- **Max GPUs:** `required_devices` must be `<= 8`. If it exceeds 8, **stop immediately** and report:
+  > `required_devices = N` exceeds the maximum of 8 GPUs allowed for pytest / module tests. Reduce `tp_size`/`pp_size` or split the test.
+- **Max nodes:** `node_count` for pytest is always `1` (single-node). If the derived `required_devices` would require more than one node given the available `gpus_per_node`, **stop immediately** and report:
+  > `required_devices = N` exceeds the GPUs available on a single node (`gpus_per_node = M`). Pytest does not support multi-node execution.
+
+Do not delegate to `trtllm-case-executor` if either limit is violated.
+
+### Step 2: Run
+
+[Delegate to `trtllm-case-executor`](#delegating-to-trtllm-case-executor) with:
+
+- `test_cmd`: the user-provided command (run verbatim)
+- `model_name`: `<class_name>` if provided, else `<MODEL_NAME>`
+- `required_devices`: derived from markers (Step 1a), user-provided value, or `1`
+- `device_type`: derived `device_type` if set (e.g. `B200`, `GB200`) — passed through to case-executor
+
+### Step 3: Report
+
+[Parse the results](#report-generation) and [write the report file](#write-report-file).
+
+The report must include these fields:
+
+| Field | Description |
+|-------|-------------|
+| `module_test_status` | **passed**, **failed**, or **error** |
+| `module_test_results` | Per-test results with test ID, status, duration |
+| `module_test_recommendations` | Fix recommendations for failures |
+
+---
+
+## Functionality Test Workflow
+
+Smoke test using `pytest examples/pytorch/quickstart_advanced.py` with the pytorch backend.
+
+### Step 1: Gather Parallelism Parameters
+
+If tp/ep/dp values not provided, ask the user:
+- `tp_size` (default `1`), `ep_size` (default `0`), `dp_size` (default `1`), `pp_size` (default `1`)
+
+### Step 2: Run
+
+[Delegate to `trtllm-case-executor`](#delegating-to-trtllm-case-executor) with:
+
+- `test_cmd`: `pytest examples/pytorch/quickstart_advanced.py --model_dir <checkpoint_path> --backend pytorch --tp_size <N> --pp_size <N> [--ep_size <N>] [--dp_size <N>] -v` (include `--ep_size` only when > 0, `--dp_size` only when > 1)
+- `model_name`: `<MODEL_NAME>`
+- `required_devices`: `tp_size * pp_size * dp_size`
+
+### Step 3: Report
+
+[Parse the results](#report-generation). Generate a markdown report containing: overall status table (passed/failed/skipped counts, duration), test command, per-test results, failure details with tracebacks, and fix recommendations. [Write the report file](#write-report-file).
+
+---
+
+## Benchmark Test Workflow
+
+Performance benchmarking. Three paths, selected in this order:
+
+- **Pre-built perf-sanity** — `test_cmd` is already a perf-sanity pytest command (classified via Step 1 of Classify Test Scope). Skip YAML generation; delegate directly to `trtllm-case-executor` (Step 0 below).
+- **Single-node → `trtllm-bench`** (default). Used whenever the requested world size fits on one node and no `benchmark_config_yaml` / `benchmark_model_name` is provided.
+- **Multi-node → perf-sanity config YAML** (fallback). Used when the world size exceeds one node, or when the user explicitly provides `benchmark_config_yaml` / `benchmark_model_name`. `trtllm-bench` itself is treated as single-node only — multi-node invocations (via `trtllm-llmapi-launch`) are intentionally not generated here.
+
+### Step 0: Pre-Built Perf-Sanity Path (when routed from Step 1 classification)
+
+Applies only when `test_cmd` was provided and identified as a perf-sanity test in Step 1. No YAML generation or `trtllm-bench` invocation is performed.
+
+**Step 0a: Derive `device_type` from test ID**
+
+If `device_type` is not provided by the user, extract it from the bracketed test ID in `test_cmd`. Scan the test ID tokens (split on `-`) for known device strings and map them:
+
+| Token in test ID | `device_type` |
+|-----------------|---------------|
+| `gb200` | `GB200` |
+| `gb300` | `GB300` |
+| `b200` | `B200` |
+| `b300` | `B300` |
+| `h100` | `H100` |
+| `h200` | `H200` |
+| `a100` | `A100` |
+
+If no known token is found, leave `device_type` unset (unconstrained).
+
+**Step 0b: Derive `model_name` from test ID**
+
+If `model_name` is not provided, extract it from the test ID: take the token immediately after the device token (e.g., `gb200_deepseek-v32` → `deepseek_v32`). Convert hyphens to underscores.
+
+**Step 0c: Delegate to `trtllm-case-executor`**
+
+[Delegate to `trtllm-case-executor`](#delegating-to-trtllm-case-executor) with:
+
+- `test_cmd`: the pre-built pytest command (verbatim)
+- `model_name`: derived in Step 0b or user-provided
+- `device_type`: derived in Step 0a or user-provided (omit if unset)
+- `perf_config_yaml`: `perf_config_yaml` input (omit if not provided)
+- `required_devices`: `1` (the pytest orchestrator process; actual compute is allocated by the test internally)
+
+Skip Steps 1–3b entirely. Proceed to Step 4 (Report) once the executor returns.
+
+### Step 1: Gather Parameters
+
+Required: `model_name` (HuggingFace id, e.g. `meta-llama/Llama-3.1-8B-Instruct`), `checkpoint_path`, `tp_size`, `pp_size`.
+
+Optional: `ep_size`, `mtp_layers` (comma-separated MTP depths to sweep, e.g. `1,2,3`), `backend` (default `pytorch`), `bench_subcommand` (default `throughput`), `dataset_path`, `num_requests`, `input_mean`, `output_mean`, `concurrency`, `max_batch_size`, `max_num_tokens`, `extra_llm_api_options_yaml`.
+
+Parameters are already merged from `config_file` by Step 0. If any required parameter is still missing, ask the user.
+
+### Step 2: Select Path
+
+Compute `required_devices = tp_size * pp_size` (ignore `dp_size` — `trtllm-bench` has no `--dp` flag; if the user sets `dp_size > 1` treat the run as multi-node).
+
+Query the target environment's `gpus_per_node` from `exec-env-check` (or the cluster config when `slurm_cluster` is set).
+
+- If `benchmark_config_yaml` **or** `benchmark_model_name` is provided → **perf-sanity fallback** (Step 3b).
+- Else if `required_devices > gpus_per_node` (or `dp_size > 1`) → **perf-sanity fallback** (Step 3b). Tell the user: "trtllm-bench is single-node only; falling back to perf-sanity for multi-node."
+- Else → **trtllm-bench** (Step 3a).
+
+### Step 3a: Single-Node trtllm-bench
+
+Build the `bench_cmd` using `scripts/build_test_command.py`:
+
+```bash
+python3 <skill_dir>/scripts/build_test_command.py --type bench \
+  --model <model_name> \
+  --model-path <checkpoint_path> \
+  --subcommand <bench_subcommand> \
+  --backend <backend> \
+  --tp <tp_size> [--pp <pp_size>] [--ep <ep_size>] \
+  [--num-requests <N>] [--input-mean <N>] [--input-stdev <N>] \
+  [--output-mean <N>] [--output-stdev <N>] \
+  [--concurrency <N>] [--max-batch-size <N>] [--max-num-tokens <N>] \
+  [--workspace <workspace>] \
+  [--kv-cache-free-gpu-mem-fraction <kv_cache_free_gpu_mem_fraction>] \
+  [--max-seq-len <max_seq_len>] [--beam-width <beam_width>] [--warmup <warmup>] \
+  [--streaming] [--no-chunked-context] [--scheduler-policy <scheduler_policy>] \
+  [--cluster-size <cluster_size>] [--iteration-log <iteration_log>] \
+  [--config <extra_llm_api_options_yaml>] \
+  --work-dir <WORK_DIR>
+```
+
+The script prints the fully-constructed `prepare-dataset && <subcommand>` command string.
+Defaults: `--subcommand throughput`, `--backend pytorch`, `--tp 1`, `--num-requests 100`, `--input-mean 128`, `--output-mean 128`.
+Dataset output: `<WORK_DIR>/bench_dataset.txt`; report: `<WORK_DIR>/bench_report.json`.
+Note: `--streaming`, `--no-chunked-context`, and `--scheduler-policy` are throughput-subcommand-only flags and are silently ignored for latency.
+
+Note: `trtllm-bench` uses `--tp` / `--pp` / `--ep` (not `--tp_size` / `--pp_size` / `--ep_size`), and takes the subcommand positionally (`throughput` / `latency`) after the global `--model` / `--model_path`.
+
+[Delegate to `trtllm-case-executor`](#delegating-to-trtllm-case-executor) with:
+
+- `bench_cmd`: the output of `build_test_command.py --type bench` above. The case-executor's Bench workflow auto-parses `--tp`/`--pp`/`--ep` from the run step for device count and enforces the single-node check.
+- `model_name`, `checkpoint_path`, `device_type`
+- `required_devices`: `tp_size * pp_size` (must be `<= gpus_per_node`, otherwise Step 2 already routed to the perf-sanity fallback)
+
+### Step 3b: Multi-Node Perf-Sanity Fallback (Benchmark)
+
+**Step 3b-1: Generate the perf-sanity config YAML**
+
+Use `<skill_dir>/scripts/generate_benchmark_config.py`:
+
+**MTP sweep:** When `mtp_layers` is provided (e.g. `1,2,3,4,5,6`), generate one config file per MTP value and run Steps 3b-2 through 3b-4 separately for each. Name each file `<benchmark_model_name>_mtp<N>.yaml` so the configs don't overwrite each other:
+
+```bash
+# For each MTP value in mtp_layers:
+for mtp in <mtp_layers_values>; do
+  python3 <skill_dir>/scripts/generate_benchmark_config.py \
+    --config-type aggr \
+    --model-name <benchmark_model_name> \
+    [--from-config <base_config_yaml>] \
+    --tp <tp_size> --ep <ep_size> --pp <pp_size> \
+    --mtp-layers $mtp \
+    [--max-batch-size <N>] [--max-num-tokens <N>] \
+    [--concurrency <N>] [--iterations <N>] \
+    [--input-length <N>] [--output-length <N>] \
+    [--kv-cache-dtype fp8|fp16] \
+    --output <WORK_DIR>/<benchmark_model_name>_mtp${mtp}.yaml
+done
+
+# Without MTP sweep (single config):
+python3 <skill_dir>/scripts/generate_benchmark_config.py \
+  --config-type aggr \
+  --model-name <benchmark_model_name> \
+  [--from-config <base_config_yaml>] \
+  --tp <tp_size> --ep <ep_size> --pp <pp_size> \
+  [--max-batch-size <N>] [--max-num-tokens <N>] \
+  --output <WORK_DIR>/<benchmark_model_name>.yaml
+
+# Disaggregated (separate ctx/gen workers):
+python3 <skill_dir>/scripts/generate_benchmark_config.py \
+  --model-name <benchmark_model_name> \
+  [--from-config <base_config_yaml>] \
+  --gen-tp <N> --ctx-tp <N> \
+  [--benchmark-mode e2e|gen_only|ctx_only] \
+  --output <WORK_DIR>/<benchmark_model_name>.yaml
+```
+
+**Step 3b-2: Identify available server config names** (aggregated only)
+
+Run for each generated YAML:
+
+```bash
+python3 <skill_dir>/scripts/build_test_command.py --type perf_sanity \
+  --config-file <WORK_DIR>/<benchmark_model_name>[_mtp<N>].yaml \
+  --list-test-names
+```
+
+The output lists all `server_configs[].name` values in the YAML, e.g.:
+```
+gpt_oss_tp4_ep4_pp2_mtp1
+gpt_oss_tp4_ep4_pp2_mtp2
+```
+
+**Step 3b-3: Build the perf_sanity test command**
+
+Run for each YAML / server config name:
+
+```bash
+# Aggregated — one command per server config entry to test:
+python3 <skill_dir>/scripts/build_test_command.py --type perf_sanity \
+  --serving-type aggr \
+  --config-file <WORK_DIR>/<benchmark_model_name>[_mtp<N>].yaml \
+  --test-name <server_config_name> \
+  --repo-root <repo_root>
+
+# Disaggregated:
+python3 <skill_dir>/scripts/build_test_command.py --type perf_sanity \
+  --serving-type disagg \
+  --config-file <WORK_DIR>/<benchmark_model_name>.yaml \
+  --benchmark-mode <e2e|gen_only|ctx_only> \
+  --repo-root <repo_root>
+```
+
+**Step 3b-4: Delegate to `trtllm-case-executor`**
+
+For each test command (one per MTP value when `mtp_layers` is set), [delegate to `trtllm-case-executor`](#delegating-to-trtllm-case-executor) with:
+
+- `test_cmd`: the pytest command from Step 3b-3
+- `perf_config_yaml`: `<WORK_DIR>/<benchmark_model_name>[_mtp<N>].yaml` (the generated YAML, so the executor uses it instead of auto-deriving the path from the test ID)
+- `model_name`, `device_type`, `benchmark_model_name`
+
+When running a multi-value MTP sweep, submit all jobs before monitoring any of them so they run concurrently on the cluster.
+
+### Step 4: Report
+
+Parse `<WORK_DIR>/bench_report.json` (single-node path) or the perf-sanity summary (fallback). Report: subcommand, backend, parallelism (`tp`/`pp`/`ep`), throughput / latency metrics, concurrency, dataset path and generation params (`num_requests`, `input_mean`, `output_mean`), work directory, path used (`trtllm-bench` vs perf-sanity), and any errors. [Write the report file](#write-report-file).
+
+---
+
+## Evaluation Test Workflow
+
+Accuracy evaluation. Two paths, selected in this order:
+
+- **Single-node → `trtllm-eval`** (default). Used whenever the requested world size fits on one node and no `benchmark_config_yaml` / `benchmark_model_name` is provided.
+- **Multi-node → perf-sanity config YAML** (fallback). Used when the world size exceeds one node, or when the user explicitly provides `benchmark_config_yaml` / `benchmark_model_name`. The perf-sanity run enables accuracy mode via `accuracy.enable_accuracy_test: true`. `trtllm-eval` itself is treated as single-node only.
+
+### Step 1: Gather Parameters
+
+Required: `checkpoint_path`, `eval_tasks`, `tp_size` (default `1`), `pp_size` (default `1`).
+
+Optional: `ep_size`, `dataset_path`, `device_type`. Note: `trtllm-eval` has no `--dp_size` flag — `dp_size` is ignored for this workflow.
+
+For the multi-node (perf-sanity) path, additional inputs apply: `benchmark_config_yaml`, `benchmark_model_name`, and any perf-sanity overrides (see the Benchmark workflow Step 3b inputs).
+
+Parameters are already merged from `config_file` by Step 0. If any required parameter is still missing, ask the user.
+
+### Step 2: Select Path
+
+Compute `required_devices = tp_size * pp_size`.
+
+Query the target environment's `gpus_per_node` from `exec-env-check` (or the cluster config when `slurm_cluster` is set).
+
+- If `benchmark_config_yaml` **or** `benchmark_model_name` is provided → **perf-sanity fallback** (Step 3b).
+- Else if `required_devices > gpus_per_node` → **perf-sanity fallback** (Step 3b). Tell the user: "trtllm-eval is single-node only; falling back to perf-sanity for multi-node accuracy."
+- Else → **trtllm-eval** (Step 3a).
+
+### Step 3a: Single-Node trtllm-eval
+
+Build the `eval_cmd` using `scripts/build_test_command.py`:
+
+```bash
+python3 <skill_dir>/scripts/build_test_command.py --type eval \
+  --model <model_name_or_checkpoint_path> [--model-path <checkpoint_path>] \
+  --tp-size <tp_size> [--pp-size <pp_size>] [--ep-size <ep_size>] \
+  [--tokenizer <eval_tokenizer>] \
+  [--trust-remote-code] [--disable-kv-cache-reuse] \
+  [--kv-cache-free-gpu-memory-fraction <eval_kv_cache_free_gpu_memory_fraction>] \
+  [--max-batch-size <N>] [--max-num-tokens <N>] [--max-seq-len <N>] \
+  [--config <extra_llm_api_options_yaml>] \
+  [--num-samples <eval_num_samples>] \
+  [--apply-chat-template] [--chat-template-kwargs '<eval_chat_template_kwargs>'] \
+  [--system-prompt '<eval_system_prompt>'] \
+  [--max-input-length <eval_max_input_length>] [--max-output-length <eval_max_output_length>] \
+  [--log-samples] [--output-path <eval_output_path>] \
+  [--check-accuracy] [--accuracy-threshold <eval_accuracy_threshold>] [--num-fewshot <eval_num_fewshot>] \
+  --tasks <eval_tasks>
+```
+
+The script prints one `trtllm-eval` command per task, joined with `&&` — tasks are Click subcommands and cannot be chained in a single invocation. `--tasks` accepts a comma-separated list (e.g. `gsm8k,mmlu`).
+
+Model-path resolution: `trtllm-eval`'s `--model` is also used to load the tokenizer, so it must be a real local path or a valid HuggingFace ID — a short label like `gpt_oss_120b` will fail tokenizer lookup. Pass either (a) the local `<checkpoint_path>` as `--model`, or (b) the short name as `--model` and the local path as `--model-path` — when `--model-path` is set, the builder uses it as the `--model` value in the rendered `trtllm-eval` command so the tokenizer resolves locally.
+
+Option placement rules applied by the script:
+- Global options (`--tp_size`, `--config`, `--kv_cache_free_gpu_memory_fraction`, etc.) are placed before the task name.
+- Common per-task options (`--num_samples`, `--apply_chat_template`, `--max_output_length`, etc.) are placed after each task name.
+- `--check-accuracy`, `--accuracy-threshold`, and `--num-fewshot` are mmlu-specific — appended only to the `mmlu` task invocation, not to any other task.
+
+[Delegate to `trtllm-case-executor`](#delegating-to-trtllm-case-executor) with:
+
+- `eval_cmd`: the output of `build_test_command.py --type eval` above.
+- `model_name`, `checkpoint_path`, `device_type`, plus additional eval parameters.
+- `required_devices`: `tp_size * pp_size` (must be `<= gpus_per_node`, otherwise Step 2 already routed to the perf-sanity fallback).
+
+### Step 3b: Multi-Node Perf-Sanity Fallback (Evaluation)
+
+**Step 3b-1: Generate the perf-sanity config YAML with accuracy enabled**
+
+Use `<skill_dir>/scripts/generate_benchmark_config.py`:
+The key difference from the benchmark path is `--enable-accuracy-test` and `--accuracy-task`:
+
+```bash
+# Aggregated accuracy evaluation:
+python3 <skill_dir>/scripts/generate_benchmark_config.py \
+  --config-type aggr \
+  --model-name <benchmark_model_name> \
+  [--from-config <base_config_yaml>] \
+  --tp <tp_size> --ep <ep_size> --pp <pp_size> \
+  --enable-accuracy-test \
+  --accuracy-task <task_name> \
+  [--accuracy-model <lm_eval_model_name>] \
+  [--accuracy-model-args-extra <extra_args>] \
+  [--accuracy-trust-remote-code] \
+  --output <WORK_DIR>/<benchmark_model_name>.yaml
+
+# Disaggregated accuracy evaluation:
+python3 <skill_dir>/scripts/generate_benchmark_config.py \
+  --model-name <benchmark_model_name> \
+  --gen-tp <N> --ctx-tp <N> \
+  --enable-accuracy-test \
+  --accuracy-task <task_name> \
+  --output <WORK_DIR>/<benchmark_model_name>.yaml
+```
+
+**Step 3b-2 and 3b-3** are identical to the Benchmark fallback — list server config names, then build the perf_sanity test command:
+
+```bash
+# List available names:
+python3 <skill_dir>/scripts/build_test_command.py --type perf_sanity \
+  --config-file <WORK_DIR>/<benchmark_model_name>.yaml \
+  --list-test-names
+
+# Build the test command:
+python3 <skill_dir>/scripts/build_test_command.py --type perf_sanity \
+  --serving-type aggr \
+  --config-file <WORK_DIR>/<benchmark_model_name>.yaml \
+  --test-name <server_config_name> \
+  --repo-root <repo_root>
+```
+
+**Step 3b-4: Delegate to `trtllm-case-executor`**
+
+[Delegate to `trtllm-case-executor`](#delegating-to-trtllm-case-executor) with:
+
+- `test_cmd`: the pytest command from Step 3b-3
+- `perf_config_yaml`: `<WORK_DIR>/<benchmark_model_name>.yaml` (accuracy config with `enable_accuracy_test: true`)
+- `model_name`, `device_type`, `benchmark_model_name`
+
+### Step 4: Report
+
+Parse the single-node eval log (Step 3a) or the perf-sanity accuracy summary (Step 3b). Report: accuracy scores per task, pass/fail status, parallelism (`tp`/`pp`/`ep`), configuration, work directory, path used (`trtllm-eval` vs perf-sanity), and any errors. [Write the report file](#write-report-file).
+
+---
+
+## Common Procedures
+
+### MODEL_NAME Derivation
+
+If `model_name` is not provided, derive from `checkpoint_path`:
+- HF model name `meta-llama/Llama-3.1-8B-Instruct` → `llama_3_1_8b_instruct`
+- Checkpoint directory `/models/Llama-3.1-8B-Instruct/` → `llama_3_1_8b_instruct`
+
+### Delegating to `trtllm-case-executor`
+
+All test execution is delegated to the `trtllm-case-executor` skill.
+
+```
+Skill(skill="trtllm-agent-toolkit:trtllm-case-executor", args="""
+- <key>: <value>
+- <key>: <value>
+...
+""")
+```
+
+The skill runs to completion (classifying the workflow, detecting the environment, building the script, dispatching to its own executor subagents) and returns the full result — status, exit code, log file paths, and any error output.
+
+### Monitoring Test Results
+
+While `trtllm-case-executor` runs the test, monitor for completion before advancing to report generation.
+
+**Slurm jobs (`local_slurm`, `remote_slurm`):** poll the queue for the jobs `trtllm-case-executor` submitted (it returns `job_id`s in its result; for remote runs, poll over SSH). Stop monitoring as soon as **all of those jobs have left the queue** — that means every job has finished and exited (completed, failed, cancelled, timed out, or preempted). Do not wait, re-poll, or resubmit after the queue shows none of them.
+
+```bash
+# Local Slurm — check every 30s until no tracked jobs remain
+while squeue --jobs=<job_id_csv> --noheader --states=all 2>/dev/null | grep -q .; do
+    sleep 30
+done
+# All jobs are gone → stop monitoring, collect final state via sacct, proceed to Report
+sacct -j <job_id_csv> --format=JobID,State,ExitCode,Elapsed --parsable2
+```
+
+For remote Slurm, run the same loop over SSH using the executor's `ssh_host` / `slurm_user`. If `squeue` returns a transient error (e.g. SSH hiccup), retry the poll on the next tick — do not assume the jobs are gone.
+
+**Local direct / local Docker:** the executor (or this skill, for `local_direct`) already blocks until the process exits; no separate queue polling is required. Use the exit code it reports.
+
+Once monitoring stops, pass the final log/exit code to [Report Generation](#report-generation).
+
+### Report Generation
+
+For pytest-based workflows (Core Module, Modeling Module, Functionality), parse test output with:
+
+```bash
+python3 <skill_dir>/scripts/generate_report.py \
+  --test-type <test_type> \
+  --output-file <log_file> \
+  --exit-code <exit_code> \
+  --format markdown
+```
+
+Map the active workflow to `--test-type`:
+
+| Workflow | `--test-type` value |
+|----------|---------------------|
+| Module test, Functionality test | `functional` |
+| Benchmark test | `benchmark` |
+| Feature-matrix test | `feature-matrix` |
+
+For all workflows, analyze failures and produce fix recommendations. Reference `references/trtllm_test_fix_recommendations.md` for common error patterns. Each recommendation should include: error type, root cause, concrete fix actions, and files to modify (`<file_path>:<line_number>`).
+
+### Write Report File
+
+After generating the report content, write it to the output path:
+- If `report_file` is provided → `mkdir -p $(dirname <report_file>)` then write to `<report_file>`
+- Otherwise → write to `./<MODEL_NAME>-auto-test-report.md`
+
+Confirm to the user: `Test report written to: <path>`
+
+---
+
+## Resources
+
+| File | Purpose |
+|------|---------|
+| `scripts/extract_test_markers.py` | Extract `required_devices` and `device_type` from pytest markers in a test file. Pass `--test-cmd "<pytest_cmd>"` or `--test-file <path> [--class-name X] [--test-name Y]`. Outputs JSON with `required_devices`, `device_type`, and `sources`. |
+| `scripts/build_test_command.py` | Build CLI commands for `trtllm-bench`, `trtllm-eval`, and `test_perf_sanity`. Accepts `--type {bench,eval,perf_sanity}` plus all relevant flags; prints the fully-constructed command string to stdout. Non-perf pytest, python, and custom commands are passed directly to `trtllm-case-executor` without using this script. Run `python3 build_test_command.py --help` for all options. |
+| `scripts/generate_benchmark_config.py` | Generates perf-sanity config YAMLs for aggregated and disaggregated benchmark/evaluation runs. Pass `--from-config` for a base config or `--config-type aggr/disagg` to build from a built-in template. |
+| `scripts/parse_config.py` | Parses and validates a YAML config file; prints recognised parameters as JSON |
+| `scripts/generate_report.py` | Parses pytest output into structured markdown |
+| `references/test_config_template.yaml` | Full config file template with all supported fields and inline documentation |
+| `references/agg_config_template.yaml` | Base aggregated benchmark config template (used with `generate_benchmark_config.py --from-config`) |
+| `references/benchmark_config_template.yaml` | Minimal benchmark config template for simple single-server runs |
+| `references/disagg_config_template.yaml` | Base disaggregated benchmark config template |
+| `references/trtllm_test_fix_recommendations.md` | Common error patterns and fix recommendations |
+
+## Notes
+
+- Tests requiring multiple GPUs (TP > 1) are skipped when only 1 GPU is available — expected, not a failure.
+- This skill reads definition files only for root-cause analysis — it does not modify source files. All recommendations are advisory.
diff --git a/.claude/skills/trtllm-test-specialist/references/agg_config_template.yaml b/.claude/skills/trtllm-test-specialist/references/agg_config_template.yaml
new file mode 100644
index 000000000000..cc5dbe906753
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/references/agg_config_template.yaml
@@ -0,0 +1,257 @@
+metadata:
+  model_name: gpt_oss_120b_fp4
+  supported_gpus:
+  - GB200
+hardware:
+  gpus_per_node: 4
+server_configs:
+  - name: "gpt_oss_fp4_dep4_1k8k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 4
+    moe_expert_parallel_size: 4
+    pipeline_parallel_size: 1
+    max_batch_size: 640
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: true
+    attention_dp_config:
+      enable_balance: true
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 640
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.8
+    num_postprocess_workers: 4
+    stream_interval: 20
+    client_configs:
+      - name: "con2560_iter5_1k8k"
+        concurrency: 2560
+        iterations: 5
+        isl: 1024
+        osl: 8192
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-1k8k-20480-ratio-1_for_serve.json
+
+  - name: "gpt_oss_fp4_dep2_1k1k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 2
+    moe_expert_parallel_size: 2
+    pipeline_parallel_size: 1
+    max_batch_size: 1024
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: true
+    attention_dp_config:
+      enable_balance: true
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 1024
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.8
+    num_postprocess_workers: 4
+    stream_interval: 20
+    client_configs:
+      - name: "con2048_iter5_1k1k"
+        concurrency: 2048
+        iterations: 5
+        isl: 1024
+        osl: 1024
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-1k1k-20480-ratio-1_for_serve.json
+
+  - name: "gpt_oss_fp4_tep2_1k8k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 2
+    moe_expert_parallel_size: 2
+    pipeline_parallel_size: 1
+    max_batch_size: 128
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: false
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 128
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.8
+    num_postprocess_workers: 4
+    stream_interval: 20
+    client_configs:
+      - name: "con128_iter10_1k8k"
+        concurrency: 128
+        iterations: 10
+        isl: 1024
+        osl: 8192
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-1k8k-20480-ratio-1_for_serve.json
+
+  - name: "gpt_oss_fp4_tp2_1k8k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 2
+    moe_expert_parallel_size: 1
+    pipeline_parallel_size: 1
+    max_batch_size: 8
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: false
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 8
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.8
+    num_postprocess_workers: 4
+    stream_interval: 20
+    client_configs:
+      - name: "con8_iter10_1k8k"
+        concurrency: 8
+        iterations: 10
+        isl: 1024
+        osl: 8192
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-1k8k-20480-ratio-1_for_serve.json
+
+  - name: "gpt_oss_fp4_tp4_eagle3_1k1k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 4
+    moe_expert_parallel_size: 1
+    pipeline_parallel_size: 1
+    max_batch_size: 1
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: false
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 1
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.8
+    speculative_config:
+      decoding_type: 'Eagle'
+      eagle3_layers_to_capture: [-1]
+      max_draft_len: 3
+      speculative_model: "gpt_oss/gpt-oss-120b-Eagle3"
+    stream_interval: 20
+    num_postprocess_workers: 4
+    client_configs:
+      - name: "con1_iter32_1k1k"
+        concurrency: 1
+        iterations: 32
+        isl: 1024
+        osl: 1024
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-1k1k-20480-ratio-1_for_serve.json
+
+  # GPT-OSS configs from Excel - 1k1k TP2 (TP2EP1) con4
+  - name: "gpt_oss_fp4_tp2_mtp0_1k1k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 2
+    moe_expert_parallel_size: 1
+    pipeline_parallel_size: 1
+    max_batch_size: 512
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: false
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 512
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.85
+    num_postprocess_workers: 4
+    stream_interval: 20
+    client_configs:
+      - name: "con4_iter10_1k1k"
+        concurrency: 4
+        iterations: 10
+        isl: 1024
+        osl: 1024
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-1k1k-20480-ratio-1_for_serve.json
+
+  # GPT-OSS configs from Excel - 1k1k TP1 con256
+  - name: "gpt_oss_fp4_tp1_mtp0_1k1k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 1
+    moe_expert_parallel_size: 1
+    pipeline_parallel_size: 1
+    max_batch_size: 512
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: false
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 512
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.85
+    num_postprocess_workers: 4
+    stream_interval: 20
+    client_configs:
+      - name: "con256_iter10_1k1k"
+        concurrency: 256
+        iterations: 10
+        isl: 1024
+        osl: 1024
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-1k1k-20480-ratio-1_for_serve.json
+
+  # GPT-OSS configs from Excel - 8k1k TP1 con4 and con256
+  - name: "gpt_oss_fp4_tp1_mtp0_8k1k"
+    model_name: "gpt_oss_120b_fp4"
+    tensor_parallel_size: 1
+    moe_expert_parallel_size: 1
+    pipeline_parallel_size: 1
+    max_batch_size: 512
+    max_num_tokens: 20000
+    attn_backend: "TRTLLM"
+    enable_attention_dp: false
+    moe_config:
+      backend: 'TRTLLM'
+    cuda_graph_config:
+      enable_padding: true
+      max_batch_size: 512
+    kv_cache_config:
+      dtype: 'fp8'
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.85
+    num_postprocess_workers: 4
+    stream_interval: 20
+    client_configs:
+      - name: "con4_iter10_8k1k"
+        concurrency: 4
+        iterations: 10
+        isl: 8192
+        osl: 1024
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-8k1k-20480-ratio-1_for_serve.json
+      - name: "con256_iter10_8k1k"
+        concurrency: 256
+        iterations: 10
+        isl: 8192
+        osl: 1024
+        backend: "openai"
+        dataset_file: datasets/perf-ci/gpt_oss_120b-8k1k-20480-ratio-1_for_serve.json
diff --git a/.claude/skills/trtllm-test-specialist/references/benchmark_config_template.yaml b/.claude/skills/trtllm-test-specialist/references/benchmark_config_template.yaml
new file mode 100644
index 000000000000..e4bda5707d07
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/references/benchmark_config_template.yaml
@@ -0,0 +1,133 @@
+# SLURM Configuration
+slurm:
+  script_file: "disaggr_torch.slurm"
+  partition: "<partition>"
+  account: "<account>"
+  job_time: "02:00:00"
+  job_name: "<job_name>"
+  extra_args: "" # Cluster specific arguments, e.g. "--gres=gpu:4 --exclude=node1,node2"
+  set_segment: true # Optional: whether to set the segment for the job
+  numa_bind: true # Only enable for GB200/GB300 NVL72
+
+# Benchmark Mode
+benchmark:
+  mode: "e2e"  # Options: e2e, gen_only, gen_only_no_context
+  use_nv_sa_benchmark: false  # Whether to use NVIDIA SA benchmark script
+  multi_round: 8  # Number of benchmark rounds
+  benchmark_ratio: 0.8  # Benchmark ratio
+  streaming: true  # Enable streaming mode
+  concurrency_list: "16"
+  input_length: 1024  # Input sequence length
+  output_length: 1024  # Output sequence length
+  dataset_file: "<dataset_file>"
+
+# Hardware Configuration
+hardware:
+  gpus_per_node: 4  # Modify this with your hardware configuration
+  num_ctx_servers: 1  # Number of context servers
+  num_gen_servers: 1  # Number of generation servers
+
+# Environment Configuration
+environment:
+  container_mount: "<container_mount>"  # Format: path1:path1,path2:path2
+  container_image: "<container_image>"
+  model_path: "<model_path>"
+  trtllm_repo: "<trtllm_repo>"
+  build_wheel: false  # Don't build the wheel when launching multiple jobs
+  cuda_architectures: ""  # Optional CUDA architectures to build for (e.g. "90-real;100-real"). If empty, builds for all architectures
+  trtllm_wheel_path: ""  # Path to pre-built TensorRT-LLM wheel. If provided, install from this wheel instead
+  work_dir: "<full_path_to_work_dir>"
+  worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0"
+  server_env_var: "TRTLLM_SERVER_DISABLE_GC=1"
+
+# Profiling Configuration
+profiling:
+  nsys_on: false  # Set to true to enable profiling
+  ctx_profile_range: "10-30"  # Set TLLM_PROFILE_START_STOP for ctx workers
+  gen_profile_range: "200-250"  # Set TLLM_PROFILE_START_STOP for gen workers
+
+# Accuracy Configuration
+accuracy:
+  enable_accuracy_test: false  # Set to true to enable accuracy evaluation
+  tasks:
+    gsm8k:
+      model: "local-completions"  # Model type for lm_eval
+      model_args_extra: "num_concurrent=512,max_retries=3,tokenized_requests=false,timeout=7200,max_gen_toks=16384"
+      extra_kwargs:
+        trust_remote_code: true
+
+worker_config:
+  gen:
+    tensor_parallel_size: 8
+    moe_expert_parallel_size: 8
+    enable_attention_dp: true
+    enable_lm_head_tp_in_adp: true
+    pipeline_parallel_size: 1
+    context_parallel_size: 1
+    # Uncomment this section to enable context parallelism.
+    # cp_config:
+    #   cp_type: "HELIX"
+    #   tokens_per_block: 32  # must match kv_config.tokens_per_block.
+    max_batch_size: 256
+    max_num_tokens: 512
+    max_seq_len: 2251
+    cuda_graph_config:
+      enable_padding: true
+      batch_sizes:
+      - 1
+      - 2
+      - 4
+      - 8
+      - 16
+      - 32
+      - 64
+      - 128
+      - 256
+      - 512
+      - 768
+      - 1024
+      - 2048
+      - 256
+    print_iter_log: true
+    trust_remote_code: true
+    kv_cache_config:
+      enable_block_reuse: false
+      tokens_per_block: 32
+      free_gpu_memory_fraction: 0.8
+      dtype: fp8
+    moe_config:
+      backend: CUTLASS
+      use_low_precision_moe_combine: true
+      load_balancer:
+        num_slots: 0
+    cache_transceiver_config:
+      max_tokens_in_buffer: 4608
+      backend: DEFAULT
+    stream_interval: 20
+    num_postprocess_workers: 4
+    speculative_config: # when mtp_size is 0, remove this section.
+      decoding_type: MTP
+      num_nextn_predict_layers: 1
+  ctx:
+    max_batch_size: 4
+    max_num_tokens: 4608
+    max_seq_len: 1227
+    tensor_parallel_size: 4
+    context_parallel_size: 1
+    moe_expert_parallel_size: 4
+    enable_attention_dp: true
+    pipeline_parallel_size: 1
+    print_iter_log: true
+    trust_remote_code: true
+    cuda_graph_config: null
+    disable_overlap_scheduler: true
+    kv_cache_config:
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.85
+      dtype: fp8
+    cache_transceiver_config:
+      max_tokens_in_buffer: 4608
+      backend: DEFAULT
+    speculative_config: # when mtp_size is 0, remove this section.
+      decoding_type: MTP
+      num_nextn_predict_layers: 1
diff --git a/.claude/skills/trtllm-test-specialist/references/disagg_config_template.yaml b/.claude/skills/trtllm-test-specialist/references/disagg_config_template.yaml
new file mode 100644
index 000000000000..c3c0b26cd4ef
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/references/disagg_config_template.yaml
@@ -0,0 +1,63 @@
+model: <MODLE_NAME>
+hostname: localhost
+backend: pytorch
+context_servers:
+  num_instances: 1
+  tensor_parallel_size: 2
+  pipeline_parallel_size: 1
+  moe_expert_parallel_size: 2
+  enable_attention_dp: false
+  max_num_tokens: 16640
+  max_seq_len: 8232
+  max_batch_size: 128
+  trust_remote_code: true
+  enable_chunked_prefill: true
+  kv_cache_config:
+    enable_block_reuse: false
+    free_gpu_memory_fraction: 0.8
+    dtype: fp8
+  disable_overlap_scheduler: true
+  moe_config:
+    backend: TRTLLM
+  cuda_graph_config: null
+  print_iter_log: true
+  cache_transceiver_config:
+    backend: DEFAULT
+    max_tokens_in_buffer: 16384
+generation_servers:
+  num_instances: 1
+  tensor_parallel_size: 2
+  pipeline_parallel_size: 1
+  moe_expert_parallel_size: 2
+  enable_attention_dp: false
+  max_num_tokens: 10240
+  max_seq_len: 10240
+  max_batch_size: 128
+  trust_remote_code: true
+  enable_chunked_prefill: true
+  kv_cache_config:
+    enable_block_reuse: false
+    free_gpu_memory_fraction: 0.8
+    dtype: fp8
+  disable_overlap_scheduler: true
+  moe_config:
+    backend: TRTLLM
+  cuda_graph_config:
+    enable_padding: true
+    batch_sizes:
+    - 1
+    - 2
+    - 4
+    - 8
+    - 16
+    - 32
+    - 64
+    - 128
+    - 256
+    - 512
+    - 768
+    - 1024
+  print_iter_log: true
+  cache_transceiver_config:
+    backend: DEFAULT
+    max_tokens_in_buffer: 16384
diff --git a/.claude/skills/trtllm-test-specialist/references/smoke_test_config_template.yml b/.claude/skills/trtllm-test-specialist/references/smoke_test_config_template.yml
new file mode 100644
index 000000000000..74b8076dff08
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/references/smoke_test_config_template.yml
@@ -0,0 +1,19 @@
+max_batch_size: 161
+max_num_tokens: 1160
+moe_expert_parallel_size: 1
+cuda_graph_config:
+  enable_padding: true
+  batch_sizes:
+  - 1
+  - 2
+  - 4
+  - 8
+  - 16
+  - 32
+  - 64
+  - 128
+  - 256
+  - 384
+enable_attention_dp: true
+kv_cache_config:
+  free_gpu_memory_fraction: 0.8
diff --git a/.claude/skills/trtllm-test-specialist/references/test_config_template.yaml b/.claude/skills/trtllm-test-specialist/references/test_config_template.yaml
new file mode 100644
index 000000000000..52cf3dd483d2
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/references/test_config_template.yaml
@@ -0,0 +1,227 @@
+# trtllm-test-specialist config file template
+# Parsed by scripts/parse_config.py — values override prompt inputs.
+# Unrecognized keys are warned about and ignored.
+
+# ---------------------------------------------------------------------------
+# Common
+# ---------------------------------------------------------------------------
+
+# [required] Short model name used in reports and file names (e.g. llama_3_1_8b)
+model_name: llama_3_1_8b
+
+# [required] Path to the HuggingFace checkpoint directory, or a HF model ID
+checkpoint_path: /models/Llama-3.1-8B-Instruct
+
+# [optional] Root of the TensorRT-LLM repo. Default: current working directory
+repo_path: /home/user/TensorRT-LLM
+
+# [optional] Path to write the markdown test report.
+#            Default: ./<model_name>-auto-test-report.md
+report_file: results/llama_3_1_8b-report.md
+
+# ---------------------------------------------------------------------------
+# Module test
+# ---------------------------------------------------------------------------
+
+# [required for module test] Exact pytest (or other) command to run verbatim.
+#                            Presence of this field triggers the module-test workflow.
+test_cmd: pytest tests/unittest/_torch/modules/test_attention.py -v
+
+# [optional] Name of the class under test (used in the report only).
+class_name: LlamaAttention
+
+# [optional] Number of GPUs required to run test_cmd. Default: 1
+required_devices: 1
+
+# ---------------------------------------------------------------------------
+# Model test — classification
+# ---------------------------------------------------------------------------
+
+# [optional] Test type: functionality | benchmark | evaluation
+#            Auto-classified from other inputs if not set.
+test_type: benchmark
+
+# ---------------------------------------------------------------------------
+# Model test — parallelism
+# ---------------------------------------------------------------------------
+
+# [optional] Tensor parallelism size. Default: 1
+tp_size: 1
+
+# [optional] Expert parallelism size (MoE models). Default: 0
+ep_size: 0
+
+# [optional] Data parallelism size. Default: 1
+dp_size: 1
+
+# [optional] Pipeline parallelism size. Default: 1
+pp_size: 1
+
+# ---------------------------------------------------------------------------
+# Model test — hardware
+# ---------------------------------------------------------------------------
+
+# [optional] Target GPU device type (e.g. GB200, B200). Used for cluster routing.
+device_type: GB200
+
+# ---------------------------------------------------------------------------
+# Model test — dataset / paths
+# ---------------------------------------------------------------------------
+
+# [optional] Path to evaluation or benchmark dataset file.
+dataset_path: /data/datasets/gsm8k.json
+
+# [optional] Path to an existing perf-sanity benchmark config YAML.
+#            Providing this triggers the multi-node perf-sanity fallback.
+benchmark_config_yaml: examples/disaggregated/slurm/benchmark/config.yaml
+
+# [optional] Path to a YAML file passed via --config to trtllm-bench / trtllm-eval
+#            (extra LLM API options).
+extra_llm_api_options_yaml: /tmp/pytorch_extra.yaml
+
+# ---------------------------------------------------------------------------
+# Model test — trtllm-bench
+# ---------------------------------------------------------------------------
+
+# [optional] trtllm-bench subcommand: throughput | latency. Default: throughput
+bench_subcommand: throughput
+
+# [optional] Backend: pytorch | tensorrt | _autodeploy. Default: pytorch
+backend: pytorch
+
+# [optional] trtllm-bench --workspace: directory for benchmark intermediate files.
+#            Default: /tmp (set by trtllm-bench internally)
+workspace: /tmp/bench_workspace
+
+# [optional] Synthetic dataset size when dataset_path is not provided. Default: 100
+num_requests: 100
+
+# [optional] Mean input token length for synthetic dataset. Default: 128
+input_mean: 128
+
+# [optional] Standard deviation for input token length distribution. Default: 0 (fixed length)
+input_stdev: 0
+
+# [optional] Mean output token length for synthetic dataset. Default: 128
+output_mean: 128
+
+# [optional] Standard deviation for output token length distribution. Default: 0 (fixed length)
+output_stdev: 0
+
+# [optional] Fraction of GPU memory reserved for KV cache after model load.
+#            Maps to --kv_cache_free_gpu_mem_fraction. Default: 0.90
+kv_cache_free_gpu_mem_fraction: 0.90
+
+# [optional] Maximum total sequence length (input + output tokens).
+#            Maps to --max_seq_len. Default: no limit
+max_seq_len: 4096
+
+# [optional] Number of search beams for beam search decoding.
+#            Maps to --beam_width. Default: 1 (greedy)
+beam_width: 1
+
+# [optional] Number of warmup requests to run before measuring.
+#            Maps to --warmup. Default: 2
+warmup: 2
+
+# [optional] --concurrency value for trtllm-bench.
+concurrency: 32
+
+# [optional] Maximum batch size runtime limit.
+max_batch_size: 64
+
+# [optional] Maximum number of tokens runtime limit.
+max_num_tokens: 8192
+
+# [optional] Enable streaming mode (throughput subcommand only). Default: false
+streaming: false
+
+# [optional] Disable chunked prefill context (throughput subcommand only).
+#            Passes --disable_chunked_context. Default: false (chunked is enabled by default)
+no_chunked_context: false
+
+# [optional] KV cache scheduler policy (throughput subcommand only).
+#            Choices: guaranteed_no_evict | max_utilization. Default: guaranteed_no_evict
+scheduler_policy: guaranteed_no_evict
+
+# [optional] Expert cluster parallelism size (MoE models with expert cluster parallelism).
+#            Maps to --cluster_size.
+cluster_size: 4
+
+# [optional] Path to write per-iteration benchmark log.
+#            Maps to --iteration_log.
+iteration_log: /tmp/bench_iteration_log.jsonl
+
+# ---------------------------------------------------------------------------
+# Model test — trtllm-eval
+# ---------------------------------------------------------------------------
+# trtllm-eval <global-options> <task> <task-options>
+# Supported tasks: gsm8k, mmlu, gpqa_diamond, gpqa_main, gpqa_extended,
+#                  cnn_dailymail, mmmu, longbench_v1, longbench_v2, json_mode_eval
+
+# [required for evaluation] Comma-separated evaluation task names.
+eval_tasks: gsm8k,gpqa_diamond
+
+# --- Global options (passed before the task name) ---
+
+# [optional] Tokenizer path or HF name. Only needed when using a TensorRT engine
+#            as the model (--tokenizer). Defaults to the model's own tokenizer.
+eval_tokenizer: /models/Llama-3.1-8B-Instruct
+
+# [optional] Trust remote code for HF models (--trust_remote_code). Default: false
+eval_trust_remote_code: false
+
+# [optional] Disable KV cache block reuse (--disable_kv_cache_reuse). Default: false
+eval_disable_kv_cache_reuse: false
+
+# [optional] Fraction of free GPU memory reserved for KV cache after model weights
+#            are allocated. Maps to --kv_cache_free_gpu_memory_fraction. Default: 0.90
+#            Note: flag name differs from trtllm-bench (memory vs mem).
+eval_kv_cache_free_gpu_memory_fraction: 0.90
+
+# --- Per-task options (common across most tasks) ---
+
+# [optional] Number of samples to evaluate. Default: full dataset (None)
+eval_num_samples: 500
+
+# [optional] Apply the model's chat template to each prompt (--apply_chat_template).
+#            Default: false (true for longbench_v1, longbench_v2, mmmu)
+eval_apply_chat_template: false
+
+# [optional] Chat template kwargs as a JSON string (--chat_template_kwargs).
+#            Used to enable thinking mode, e.g. '{"enable_thinking": true}' for Qwen3.
+eval_chat_template_kwargs: '{"thinking_budget": 0}'
+
+# [optional] System prompt prepended to every request (--system_prompt).
+eval_system_prompt: "You are a helpful assistant."
+
+# [optional] Maximum input prompt length in tokens (--max_input_length).
+#            Prompts exceeding this are truncated. Task-specific defaults vary
+#            (e.g. gsm8k: 4096, mmlu: 4094, cnn_dailymail: 924).
+eval_max_input_length: 4096
+
+# [optional] Maximum output generation length in tokens (--max_output_length).
+#            Task-specific defaults vary (e.g. gsm8k: 256, gpqa: 32768, mmlu: 2).
+eval_max_output_length: 256
+
+# [optional] Log per-sample outputs to stdout for debugging (--log_samples).
+#            Supported by: gsm8k, gpqa_*, mmmu, longbench_v1. Default: false
+eval_log_samples: false
+
+# [optional] Path to save per-sample evaluation results JSON (--output_path).
+#            Supported by: gsm8k, gpqa_*, mmmu, longbench_v1.
+eval_output_path: /tmp/eval_results.json
+
+# --- Task-specific options ---
+
+# [optional] Assert that accuracy meets a minimum threshold (--check_accuracy).
+#            mmlu only. Default: false
+eval_check_accuracy: false
+
+# [optional] Minimum accuracy fraction to pass when eval_check_accuracy is true
+#            (--accuracy_threshold). mmlu only. Default: 30
+eval_accuracy_threshold: 30
+
+# [optional] Number of few-shot examples prepended to each prompt (--num_fewshot).
+#            mmlu only. Default: 5
+eval_num_fewshot: 5
diff --git a/.claude/skills/trtllm-test-specialist/references/trtllm_test_fix_recommendations.md b/.claude/skills/trtllm-test-specialist/references/trtllm_test_fix_recommendations.md
new file mode 100644
index 000000000000..28e85dc20676
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/references/trtllm_test_fix_recommendations.md
@@ -0,0 +1,14 @@
+# Test Cases Failure Fix Recommendations
+
+
+| Error | Root cause pattern | Recommended fix |
+|-------|--------------------|-----------------|
+| `ImportError` on a class | Class not exported | Add missing export to `__init__.py` |
+| `AssertionError` on tensor values | Wrong computation in `forward()` | Compare `forward()` output against HuggingFace reference using `torch.allclose` |
+| Shape mismatch | Wrong dimension in `__init__` or `forward()` | Verify dimensions against model config fields |
+| `RuntimeError: CUDA error` | Device placement issue | Check tensor `.device` before ops; verify CUDA availability |
+| `KeyError` on weight name | Wrong weight name mapping | Fix HF → TRT-LLM name translation in `load_weights()` |
+| `shape mismatch in QKV fusion` | Wrong concat order | Ensure Q, K, V concatenated as `[Q, K, V]` along dim 0 |
+| `AttributeError` on config field | Missing config field | Add field with appropriate default to the model config class |
+| `CUDA out of memory` | Batch/sequence too large | Reduce test batch size or sequence length; call `torch.cuda.empty_cache()` before test |
+| `TimeoutError` / test hangs | Deadlock or infinite loop | Inspect distributed setup; reduce test size |
diff --git a/.claude/skills/trtllm-test-specialist/scripts/build_test_command.py b/.claude/skills/trtllm-test-specialist/scripts/build_test_command.py
new file mode 100644
index 000000000000..e74a81ca9eb7
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/scripts/build_test_command.py
@@ -0,0 +1,768 @@
+# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+r"""Build CLI commands for trtllm test workflows.
+
+Supports three command types:
+  bench        -- trtllm-bench prepare-dataset + throughput/latency steps
+  eval         -- trtllm-eval accuracy evaluation
+  perf_sanity  -- pytest test_perf_sanity.py with aggregated or disaggregated test ID
+
+Non-perf pytest commands, python scripts, and custom commands do not require
+this script — pass them directly to trtllm-case-executor as test_cmd / custom_cmd.
+
+Prints the fully-constructed command string to stdout (suitable for passing
+to trtllm-case-executor).  Exit code is non-zero on error.
+
+Examples:
+--------
+# trtllm-bench single GPU:
+python3 build_test_command.py --type bench \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --model-path /models/Llama-3.1-8B-Instruct \
+  --tp 1 --work-dir /tmp/work
+
+# trtllm-bench multi-GPU:
+python3 build_test_command.py --type bench \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --model-path /models/Llama-3.1-8B-Instruct \
+  --tp 4 --backend pytorch \
+  --num-requests 512 --input-mean 1024 --output-mean 1024 \
+  --concurrency 32 --work-dir /tmp/work
+
+# trtllm-eval single GPU:
+python3 build_test_command.py --type eval \
+  --model /models/Llama-3.1-8B-Instruct \
+  --tp-size 1 --tasks gsm8k
+
+# trtllm-eval multi-GPU:
+python3 build_test_command.py --type eval \
+  --model /models/Llama-3.1-8B-Instruct \
+  --tp-size 4 --tasks gsm8k,mmlu
+
+# perf_sanity aggregated — config name and test name from args:
+python3 build_test_command.py --type perf_sanity \
+  --serving-type aggr \
+  --config-name gpt_oss \
+  --test-name gpt_oss_tp4ep4pp2_mtp1 \
+  --repo-root /path/to/TensorRT-LLM
+
+# perf_sanity aggregated — config name derived from a generated YAML:
+python3 build_test_command.py --type perf_sanity \
+  --serving-type aggr \
+  --config-file /work/gpt_oss.yaml \
+  --test-name gpt_oss_tp4ep4pp2_mtp1 \
+  --repo-root /path/to/TensorRT-LLM
+
+# List available test names inside a config YAML:
+python3 build_test_command.py --type perf_sanity \
+  --config-file /work/gpt_oss.yaml \
+  --list-test-names
+
+# perf_sanity disaggregated:
+python3 build_test_command.py --type perf_sanity \
+  --serving-type disagg --benchmark-mode e2e \
+  --config-name deepseek_r1 \
+  --repo-root /path/to/TensorRT-LLM
+
+Multi-node benchmark/evaluation workflow
+-----------------------------------------
+For workloads that exceed one node, generate a perf-sanity config YAML first
+(using generate_benchmark_config.py from this skill),
+then build the test command from it:
+
+  # 1. Generate the config (aggr, benchmark):
+  python3 <skill_dir>/scripts/generate_benchmark_config.py \\
+    --config-type aggr --model-name gpt_oss_120b_fp4 \\
+    --tp 4 --ep 4 --pp 2 --mtp-layers 1 \\
+    --output /work/gpt_oss_120b_fp4.yaml
+
+  # 2. List available server config names:
+  python3 build_test_command.py --type perf_sanity \\
+    --config-file /work/gpt_oss_120b_fp4.yaml --list-test-names
+
+  # 3. Build the pytest command:
+  python3 build_test_command.py --type perf_sanity \\
+    --config-file /work/gpt_oss_120b_fp4.yaml \\
+    --test-name gpt_oss_120b_fp4_tp4ep4pp2_mtp1 \\
+    --repo-root /path/to/TensorRT-LLM
+
+  # For evaluation (accuracy), add --enable-accuracy-test when generating:
+  python3 <skill_dir>/scripts/generate_benchmark_config.py \\
+    --config-type aggr --model-name gpt_oss_120b_fp4 \\
+    --tp 4 --ep 4 --pp 2 \\
+    --enable-accuracy-test --accuracy-task gsm8k \\
+    --output /work/gpt_oss_120b_fp4.yaml
+"""
+
+import argparse
+import os
+import re
+import shlex
+import sys
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _q(s):
+    """Shell-quote a string only when it contains characters that need it."""
+    if re.search(r"[^\w@%+=:,./-]", s):
+        return shlex.quote(s)
+    return s
+
+
+def _join(*parts):
+    return " ".join(_q(str(p)) for p in parts if p is not None)
+
+
+# ---------------------------------------------------------------------------
+# bench
+# ---------------------------------------------------------------------------
+
+
+def build_bench(args):
+    """Return 'prepare-dataset && <subcommand>' command string."""
+    if not args.model:
+        sys.exit("--model is required for --type bench")
+
+    work_dir = args.work_dir or "."
+    dataset_path = args.dataset or os.path.join(work_dir, "bench_dataset.txt")
+    report_json = args.report_json or os.path.join(work_dir, "bench_report.json")
+
+    num_requests = args.num_requests if args.num_requests is not None else 100
+    input_mean = args.input_mean if args.input_mean is not None else 128
+    input_stdev = args.input_stdev if args.input_stdev is not None else 0
+    output_mean = args.output_mean if args.output_mean is not None else 128
+    output_stdev = args.output_stdev if args.output_stdev is not None else 0
+
+    # --- prepare-dataset step ---
+    prep = ["trtllm-bench", "--model", args.model]
+    if args.workspace:
+        prep += ["--workspace", args.workspace]
+    if args.model_path:
+        prep += ["--model_path", args.model_path]
+    prep += [
+        "prepare-dataset",
+        "--output",
+        dataset_path,
+        "token-norm-dist",
+        "--num-requests",
+        str(num_requests),
+        "--input-mean",
+        str(input_mean),
+        "--input-stdev",
+        str(input_stdev),
+        "--output-mean",
+        str(output_mean),
+        "--output-stdev",
+        str(output_stdev),
+    ]
+
+    # --- benchmark run step ---
+    subcommand = args.subcommand or "throughput"
+    backend = args.backend or "pytorch"
+    tp = args.tp  # argparse default: 1
+    pp = args.pp  # argparse default: 1
+    ep = args.ep  # argparse default: 0
+
+    run = ["trtllm-bench", "--model", args.model]
+    if args.workspace:
+        run += ["--workspace", args.workspace]
+    if args.model_path:
+        run += ["--model_path", args.model_path]
+    run += [subcommand, "--backend", backend, "--dataset", dataset_path]
+    run += ["--tp", str(tp)]
+    if pp > 1:
+        run += ["--pp", str(pp)]
+    if ep > 0:
+        run += ["--ep", str(ep)]
+    if args.cluster_size is not None:
+        run += ["--cluster_size", str(args.cluster_size)]
+    if args.max_seq_len is not None:
+        run += ["--max_seq_len", str(args.max_seq_len)]
+    if args.kv_cache_free_gpu_mem_fraction is not None:
+        run += ["--kv_cache_free_gpu_mem_fraction", str(args.kv_cache_free_gpu_mem_fraction)]
+    if args.beam_width is not None:
+        run += ["--beam_width", str(args.beam_width)]
+    if args.warmup is not None:
+        run += ["--warmup", str(args.warmup)]
+    if args.concurrency is not None:
+        run += ["--concurrency", str(args.concurrency)]
+    if args.max_batch_size is not None:
+        run += ["--max_batch_size", str(args.max_batch_size)]
+    if args.max_num_tokens is not None:
+        run += ["--max_num_tokens", str(args.max_num_tokens)]
+    if args.config:
+        run += ["--config", args.config]
+    if subcommand == "throughput":
+        if args.streaming:
+            run += ["--streaming"]
+        if args.no_chunked_context:
+            run += ["--disable_chunked_context"]
+        if args.scheduler_policy:
+            run += ["--scheduler_policy", args.scheduler_policy]
+    run += ["--report_json", report_json]
+    if args.iteration_log:
+        run += ["--iteration_log", args.iteration_log]
+
+    world_size = tp * pp
+    prep_cmd = _join(*prep)
+    run_cmd = _join(*run) if world_size == 1 else _join("trtllm-llmapi-launch", *run)
+    return f"{prep_cmd} && {run_cmd}"
+
+
+# ---------------------------------------------------------------------------
+# eval
+# ---------------------------------------------------------------------------
+
+
+def build_eval(args):
+    """Return trtllm-eval command string(s). One invocation per task, joined with &&.
+
+    trtllm-eval is a Click group without chain=True, so multiple tasks cannot be
+    chained in a single invocation. Each task gets its own command.
+
+    Option placement:
+      - Global options go before the task subcommand name.
+      - Common per-task options (num_samples, apply_chat_template, etc.) are
+        appended after each task name.
+      - mmlu-specific options (check_accuracy, accuracy_threshold, num_fewshot)
+        are appended only to the mmlu task invocation.
+    """
+    if not args.model:
+        sys.exit("--model is required for --type eval")
+    if not args.tasks:
+        sys.exit("--tasks is required for --type eval")
+
+    tp_size = args.tp_size  # argparse default: 1
+    pp_size = args.pp_size  # argparse default: 1
+    ep_size = args.ep_size  # argparse default: 0
+
+    # --- Global options (placed before the task subcommand name) ---
+    # trtllm-eval's --model accepts an HF ID, a local checkpoint path, or a
+    # TensorRT engine path; the tokenizer is loaded from whatever is passed
+    # here (unless --tokenizer overrides it). When a local checkpoint is
+    # available via --model-path, prefer it over a short label for --model so
+    # the tokenizer resolves locally instead of being treated as an HF ID
+    # (which fails for short names like "gpt_oss_120b").
+    model_arg = args.model_path if args.model_path else args.model
+    global_parts = ["trtllm-eval", "--model", model_arg]
+    global_parts += ["--tp_size", str(tp_size)]
+    if pp_size > 1:
+        global_parts += ["--pp_size", str(pp_size)]
+    if ep_size > 0:
+        global_parts += ["--ep_size", str(ep_size)]
+    if args.tokenizer:
+        global_parts += ["--tokenizer", args.tokenizer]
+    if args.backend and args.backend != "pytorch":
+        global_parts += ["--backend", args.backend]
+    if args.trust_remote_code:
+        global_parts += ["--trust_remote_code"]
+    if args.disable_kv_cache_reuse:
+        global_parts += ["--disable_kv_cache_reuse"]
+    if args.eval_kv_cache_free_gpu_memory_fraction is not None:
+        global_parts += [
+            "--kv_cache_free_gpu_memory_fraction",
+            str(args.eval_kv_cache_free_gpu_memory_fraction),
+        ]
+    if args.max_batch_size is not None:
+        global_parts += ["--max_batch_size", str(args.max_batch_size)]
+    if args.max_num_tokens is not None:
+        global_parts += ["--max_num_tokens", str(args.max_num_tokens)]
+    if args.max_seq_len is not None:
+        global_parts += ["--max_seq_len", str(args.max_seq_len)]
+    if args.config:
+        global_parts += ["--config", args.config]
+    if args.extra_eval_flags:
+        global_parts += args.extra_eval_flags.split()
+
+    # --- Common per-task options (appended after each task name) ---
+    common_task_parts = []
+    if args.num_samples is not None:
+        common_task_parts += ["--num_samples", str(args.num_samples)]
+    if args.apply_chat_template:
+        common_task_parts += ["--apply_chat_template"]
+    if args.chat_template_kwargs:
+        common_task_parts += ["--chat_template_kwargs", args.chat_template_kwargs]
+    if args.system_prompt:
+        common_task_parts += ["--system_prompt", args.system_prompt]
+    if args.max_input_length is not None:
+        common_task_parts += ["--max_input_length", str(args.max_input_length)]
+    if args.max_output_length is not None:
+        common_task_parts += ["--max_output_length", str(args.max_output_length)]
+    if args.log_samples:
+        common_task_parts += ["--log_samples"]
+    if args.output_path:
+        common_task_parts += ["--output_path", args.output_path]
+
+    # --- mmlu-specific options (only appended to the mmlu task invocation) ---
+    mmlu_parts = []
+    if args.check_accuracy:
+        mmlu_parts += ["--check_accuracy"]
+    if args.accuracy_threshold is not None:
+        mmlu_parts += ["--accuracy_threshold", str(args.accuracy_threshold)]
+    if args.num_fewshot is not None:
+        mmlu_parts += ["--num_fewshot", str(args.num_fewshot)]
+
+    world_size = tp_size * pp_size
+    launcher = ["trtllm-llmapi-launch"] if world_size > 1 else []
+
+    task_cmds = []
+    for task in args.tasks.split(","):
+        task = task.strip()
+        if not task:
+            continue
+        task_parts = list(common_task_parts)
+        if task == "mmlu":
+            task_parts += mmlu_parts
+        full_parts = launcher + global_parts + [task] + task_parts
+        task_cmds.append(_join(*full_parts))
+
+    if not task_cmds:
+        sys.exit("--tasks produced no valid task names")
+    return " && ".join(task_cmds)
+
+
+# ---------------------------------------------------------------------------
+# perf_sanity
+# ---------------------------------------------------------------------------
+
+
+def _load_yaml_config(path):
+    """Load a YAML file, returning the parsed dict. Exits on error."""
+    try:
+        import yaml
+    except ImportError:
+        sys.exit("PyYAML is required: pip install pyyaml")
+    try:
+        with open(path) as f:
+            return yaml.safe_load(f)
+    except FileNotFoundError:
+        sys.exit(f"Config file not found: {path}")
+    except Exception as e:
+        sys.exit(f"Failed to read {path}: {e}")
+
+
+def _aggr_server_names(config):
+    """Return list of server_config names from an aggregated config dict."""
+    return [sc.get("name", "") for sc in config.get("server_configs", [])]
+
+
+def build_perf_sanity(args):
+    """Return pytest command string for test_perf_sanity.py."""
+    repo_root = args.repo_root or "."
+    test_file = os.path.join(repo_root, "tests/integration/defs/perf/test_perf_sanity.py")
+
+    # Resolve config_name: prefer explicit --config-name; fall back to YAML filename.
+    config_file_data = None
+    if args.config_file:
+        config_file_data = _load_yaml_config(args.config_file)
+        basename = os.path.basename(args.config_file)
+        derived_name = re.sub(r"\.(yaml|yml)$", "", basename)
+        config_name = args.config_name or derived_name
+    else:
+        config_name = args.config_name
+
+    if not config_name:
+        sys.exit(
+            "--config-name (or --config-file whose basename is used as the name) "
+            "is required for --type perf_sanity"
+        )
+
+    # --list-test-names: print available server_config names and exit.
+    if args.list_test_names:
+        if config_file_data is None:
+            sys.exit("--config-file is required with --list-test-names")
+        names = _aggr_server_names(config_file_data)
+        if not names:
+            print("No server_configs found in the config file.", file=sys.stderr)
+            sys.exit(1)
+        for name in names:
+            print(name)
+        return None  # signal to main() not to print a command
+
+    serving_type = args.serving_type or "aggr"
+
+    # Build the pytest test ID from serving type and config
+    if serving_type == "aggr":
+        test_name = args.test_name
+        if not test_name:
+            # If config file is available, suggest valid names.
+            hint = ""
+            if config_file_data is not None:
+                names = _aggr_server_names(config_file_data)
+                if names:
+                    hint = f"  Available names: {', '.join(names)}"
+            sys.exit(f"--test-name is required for --serving-type aggr.{hint}")
+        test_id = f"test_e2e[aggr-{config_name}-{test_name}]"
+    else:
+        # disaggregated: mode determines the prefix segment
+        mode = args.benchmark_mode or "e2e"
+        if mode == "ctx_only":
+            test_id = f"test_e2e[aggr-ctx_only-{config_name}]"
+        else:
+            test_id = f"test_e2e[disagg-{mode}-{config_name}]"
+
+    pytest_flags = args.pytest_flags or "-v"
+    parts = ["pytest", _q(test_file), pytest_flags]
+    parts.append(f'"{test_id}"')
+    if args.output_dir:
+        parts += ["--output-dir", _q(args.output_dir)]
+    if args.extra_pytest_flags:
+        parts += args.extra_pytest_flags.split()
+    return " ".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def make_parser():
+    p = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument(
+        "--type",
+        required=True,
+        choices=["bench", "eval", "perf_sanity"],
+        dest="cmd_type",
+        help="Command type to build",
+    )
+
+    # ---- Common ----
+    g_common = p.add_argument_group("Common")
+    g_common.add_argument("--model", help="HuggingFace model ID or local path (bench/eval)")
+    g_common.add_argument(
+        "--model-path",
+        dest="model_path",
+        help="Local checkpoint path. For bench, passed as --model_path. "
+        "For eval, used as --model when set so the tokenizer loads "
+        "from the local checkpoint instead of being treated as an HF ID.",
+    )
+    g_common.add_argument(
+        "--work-dir", dest="work_dir", help="Work directory for default output paths"
+    )
+    g_common.add_argument(
+        "--repo-root", dest="repo_root", help="TensorRT-LLM repo root (perf_sanity/pytest)"
+    )
+    g_common.add_argument("--output-dir", dest="output_dir", help="Pytest --output-dir")
+
+    # ---- bench ----
+    g_bench = p.add_argument_group("bench (trtllm-bench)")
+    g_bench.add_argument(
+        "--subcommand",
+        choices=["throughput", "latency"],
+        default="throughput",
+        help="trtllm-bench subcommand (default: throughput)",
+    )
+    g_bench.add_argument(
+        "--backend", default="pytorch", help="trtllm-bench --backend (default: pytorch)"
+    )
+    g_bench.add_argument("--tp", type=int, default=1, help="Tensor parallel size (default: 1)")
+    g_bench.add_argument("--pp", type=int, default=1, help="Pipeline parallel size (default: 1)")
+    g_bench.add_argument("--ep", type=int, default=0, help="Expert parallel size (default: 0)")
+    g_bench.add_argument(
+        "--dataset", help="Dataset file path (default: <work-dir>/bench_dataset.txt)"
+    )
+    g_bench.add_argument(
+        "--num-requests",
+        dest="num_requests",
+        type=int,
+        default=100,
+        help="Synthetic dataset size (default: 100)",
+    )
+    g_bench.add_argument(
+        "--input-mean",
+        dest="input_mean",
+        type=int,
+        default=128,
+        help="Mean input token length (default: 128)",
+    )
+    g_bench.add_argument(
+        "--input-stdev",
+        dest="input_stdev",
+        type=int,
+        default=0,
+        help="Input length std dev (default: 0)",
+    )
+    g_bench.add_argument(
+        "--output-mean",
+        dest="output_mean",
+        type=int,
+        default=128,
+        help="Mean output token length (default: 128)",
+    )
+    g_bench.add_argument(
+        "--output-stdev",
+        dest="output_stdev",
+        type=int,
+        default=0,
+        help="Output length std dev (default: 0)",
+    )
+    g_bench.add_argument("--concurrency", type=int, help="--concurrency for benchmark run")
+    g_bench.add_argument(
+        "--max-batch-size", dest="max_batch_size", type=int, help="--max_batch_size"
+    )
+    g_bench.add_argument(
+        "--max-num-tokens", dest="max_num_tokens", type=int, help="--max_num_tokens"
+    )
+    g_bench.add_argument("--config", help="Extra LLM API options YAML (--config flag)")
+    g_bench.add_argument(
+        "--report-json",
+        dest="report_json",
+        help="Output report path (default: <work-dir>/bench_report.json)",
+    )
+    g_bench.add_argument(
+        "--workspace", help="trtllm-bench --workspace: directory for benchmark intermediate files"
+    )
+    g_bench.add_argument(
+        "--kv-cache-free-gpu-mem-fraction",
+        dest="kv_cache_free_gpu_mem_fraction",
+        type=float,
+        help="Fraction of GPU memory reserved for KV cache after model load (default: 0.90)",
+    )
+    g_bench.add_argument(
+        "--max-seq-len",
+        dest="max_seq_len",
+        type=int,
+        help="Maximum total sequence length (input + output tokens)",
+    )
+    g_bench.add_argument(
+        "--beam-width",
+        dest="beam_width",
+        type=int,
+        help="Number of search beams for beam search decoding (default: 1)",
+    )
+    g_bench.add_argument(
+        "--warmup", type=int, help="Number of warmup requests before benchmarking (default: 2)"
+    )
+    g_bench.add_argument(
+        "--streaming",
+        action="store_true",
+        help="Enable streaming mode (throughput subcommand only)",
+    )
+    g_bench.add_argument(
+        "--no-chunked-context",
+        dest="no_chunked_context",
+        action="store_true",
+        help="Disable chunked prefill — passes --disable_chunked_context (throughput only; chunked is on by default)",
+    )
+    g_bench.add_argument(
+        "--scheduler-policy",
+        dest="scheduler_policy",
+        choices=["guaranteed_no_evict", "max_utilization"],
+        help="KV cache scheduler policy (throughput only)",
+    )
+    g_bench.add_argument(
+        "--cluster-size",
+        dest="cluster_size",
+        type=int,
+        help="Expert cluster parallelism size (MoE models)",
+    )
+    g_bench.add_argument(
+        "--iteration-log",
+        dest="iteration_log",
+        help="Path to write per-iteration benchmark log (--iteration_log)",
+    )
+
+    # ---- eval ----
+    # Note: --max-batch-size, --max-num-tokens, --max-seq-len, --config, and --backend
+    # are defined in the bench group but are read by build_eval as well — no duplicates needed.
+    g_eval = p.add_argument_group("eval (trtllm-eval)")
+    g_eval.add_argument(
+        "--tasks",
+        help="Comma-separated eval task names (e.g. gsm8k,mmlu). "
+        "Each task becomes a separate trtllm-eval invocation joined with &&.",
+    )
+    g_eval.add_argument(
+        "--tp-size", dest="tp_size", type=int, default=1, help="TP size (default: 1)"
+    )
+    g_eval.add_argument(
+        "--pp-size", dest="pp_size", type=int, default=1, help="PP size (default: 1)"
+    )
+    g_eval.add_argument(
+        "--ep-size", dest="ep_size", type=int, default=0, help="EP size (default: 0)"
+    )
+    # Global eval options
+    g_eval.add_argument(
+        "--tokenizer",
+        dest="tokenizer",
+        help="Tokenizer path or HF name (--tokenizer; only needed for TensorRT engine)",
+    )
+    g_eval.add_argument(
+        "--trust-remote-code",
+        dest="trust_remote_code",
+        action="store_true",
+        help="Pass --trust_remote_code to trtllm-eval",
+    )
+    g_eval.add_argument(
+        "--disable-kv-cache-reuse",
+        dest="disable_kv_cache_reuse",
+        action="store_true",
+        help="Pass --disable_kv_cache_reuse to trtllm-eval",
+    )
+    g_eval.add_argument(
+        "--kv-cache-free-gpu-memory-fraction",
+        dest="eval_kv_cache_free_gpu_memory_fraction",
+        type=float,
+        help="Fraction of free GPU memory for KV cache after model weights "
+        "(--kv_cache_free_gpu_memory_fraction; note: 'memory' not 'mem')",
+    )
+    # Common per-task options (appended after each task name)
+    g_eval.add_argument(
+        "--num-samples",
+        dest="num_samples",
+        type=int,
+        help="Number of evaluation samples per task (--num_samples)",
+    )
+    g_eval.add_argument(
+        "--apply-chat-template",
+        dest="apply_chat_template",
+        action="store_true",
+        help="Apply the model's chat template to each prompt (--apply_chat_template)",
+    )
+    g_eval.add_argument(
+        "--chat-template-kwargs",
+        dest="chat_template_kwargs",
+        help="Chat template kwargs as a JSON string (--chat_template_kwargs)",
+    )
+    g_eval.add_argument(
+        "--system-prompt",
+        dest="system_prompt",
+        help="System prompt prepended to every request (--system_prompt)",
+    )
+    g_eval.add_argument(
+        "--max-input-length",
+        dest="max_input_length",
+        type=int,
+        help="Maximum input prompt length in tokens (--max_input_length)",
+    )
+    g_eval.add_argument(
+        "--max-output-length",
+        dest="max_output_length",
+        type=int,
+        help="Maximum output generation length in tokens (--max_output_length)",
+    )
+    g_eval.add_argument(
+        "--log-samples",
+        dest="log_samples",
+        action="store_true",
+        help="Log per-sample outputs to stdout for debugging (--log_samples)",
+    )
+    g_eval.add_argument(
+        "--output-path",
+        dest="output_path",
+        help="Path to save per-sample evaluation results JSON (--output_path)",
+    )
+    # mmlu-specific options (only appended when the task is 'mmlu')
+    g_eval.add_argument(
+        "--check-accuracy",
+        dest="check_accuracy",
+        action="store_true",
+        help="Assert accuracy meets minimum threshold — mmlu only (--check_accuracy)",
+    )
+    g_eval.add_argument(
+        "--accuracy-threshold",
+        dest="accuracy_threshold",
+        type=float,
+        help="Minimum accuracy fraction to pass — mmlu only (--accuracy_threshold)",
+    )
+    g_eval.add_argument(
+        "--num-fewshot",
+        dest="num_fewshot",
+        type=int,
+        help="Number of few-shot examples prepended to each prompt — mmlu only (--num_fewshot)",
+    )
+    g_eval.add_argument(
+        "--extra-eval-flags",
+        dest="extra_eval_flags",
+        help="Additional global flags appended verbatim to trtllm-eval (before task name)",
+    )
+
+    # ---- perf_sanity ----
+    g_perf = p.add_argument_group("perf_sanity (test_perf_sanity.py)")
+    g_perf.add_argument(
+        "--serving-type",
+        dest="serving_type",
+        choices=["aggr", "disagg"],
+        default="aggr",
+        help="Serving topology (default: aggr)",
+    )
+    g_perf.add_argument(
+        "--config-name",
+        dest="config_name",
+        help="Base name of the perf-sanity config YAML (without .yaml). "
+        "Auto-derived from --config-file basename when omitted.",
+    )
+    g_perf.add_argument(
+        "--config-file",
+        dest="config_file",
+        help="Path to a generated or existing perf-sanity config YAML. "
+        "config_name is derived from its basename. "
+        "Use with --list-test-names to inspect available server configs.",
+    )
+    g_perf.add_argument(
+        "--list-test-names",
+        dest="list_test_names",
+        action="store_true",
+        help="Print available server_config names from --config-file and exit.",
+    )
+    g_perf.add_argument(
+        "--test-name",
+        dest="test_name",
+        help="Server config entry name in the YAML (aggr only). "
+        "Required for aggr; if omitted and --config-file is given, "
+        "available names are shown.",
+    )
+    g_perf.add_argument(
+        "--benchmark-mode",
+        dest="benchmark_mode",
+        choices=["e2e", "gen_only", "ctx_only"],
+        default="e2e",
+        help="Benchmark mode (disagg, default: e2e)",
+    )
+    g_perf.add_argument(
+        "--pytest-flags", dest="pytest_flags", default="-v", help="Pytest flags (default: -v)"
+    )
+    g_perf.add_argument(
+        "--extra-pytest-flags",
+        dest="extra_pytest_flags",
+        help="Additional pytest flags appended verbatim",
+    )
+
+    return p
+
+
+def main():
+    parser = make_parser()
+    args = parser.parse_args()
+
+    builders = {
+        "bench": build_bench,
+        "eval": build_eval,
+        "perf_sanity": build_perf_sanity,
+    }
+    cmd = builders[args.cmd_type](args)
+    # build_perf_sanity returns None when --list-test-names is used
+    # (output already printed to stdout by the function).
+    if cmd is not None:
+        print(cmd)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.claude/skills/trtllm-test-specialist/scripts/extract_test_markers.py b/.claude/skills/trtllm-test-specialist/scripts/extract_test_markers.py
new file mode 100644
index 000000000000..37df138e9a9b
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/scripts/extract_test_markers.py
@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+"""Extract required_devices and device_type from pytest markers in a test file.
+
+Parses the target class and function in the given test file and reports
+constraints implied by TensorRT-LLM pytest markers.
+
+Output (stdout): JSON with keys:
+  required_devices  int   - minimum GPUs needed (default 1)
+  device_type       str   - e.g. "Blackwell (B200/GB200)", "" if unconstrained
+  sources           dict  - {"required_devices": "derived"|"default",
+                             "device_type":       "derived"|"default"}
+"""
+
+import argparse
+import ast
+import json
+import re
+import sys
+from pathlib import Path
+
+# Maps marker name / alias → (field, value)
+_DEVICE_TYPE_MARKERS = {
+    "skip_pre_blackwell": "Blackwell (B200/GB200)",
+    "skip_pre_hopper": "Hopper+ (H100/H200)",
+}
+
+
+def _parse_test_cmd(test_cmd: str) -> tuple[str, str | None, str | None]:
+    """Return (file_path, class_name, test_name) from a pytest node-id string."""
+    # Strip leading 'pytest' / 'python -m pytest' tokens
+    parts = test_cmd.split()
+    node_id = next((p for p in parts if p.endswith(".py") or ".py::" in p), None)
+    if node_id is None:
+        raise ValueError(f"Cannot find a .py path in test_cmd: {test_cmd!r}")
+
+    file_part, *rest = node_id.split("::")
+    class_name = rest[0] if len(rest) >= 2 else None
+    test_name = rest[1] if len(rest) >= 2 else (rest[0] if rest else None)
+    return file_part, class_name, test_name
+
+
+def _decorator_names(decorator: ast.expr) -> list[str]:
+    """Flatten a decorator AST node to a list of candidate name strings."""
+    names = []
+    if isinstance(decorator, ast.Name):
+        names.append(decorator.id)
+    elif isinstance(decorator, ast.Attribute):
+        names.append(decorator.attr)
+        names.append(f"{_attr_chain(decorator)}")
+    elif isinstance(decorator, ast.Call):
+        names.extend(_decorator_names(decorator.func))
+    return names
+
+
+def _attr_chain(node: ast.expr) -> str:
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        return f"{_attr_chain(node.value)}.{node.attr}"
+    return ""
+
+
+def _get_int_arg(call: ast.Call) -> int | None:
+    """Return the first integer argument of a Call node, if any."""
+    for arg in call.args:
+        if isinstance(arg, ast.Constant) and isinstance(arg.value, int):
+            return arg.value
+    for kw in call.keywords:
+        if isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, int):
+            return kw.value.value
+    return None
+
+
+def _extract_from_decorators(
+    decorators: list[ast.expr],
+) -> tuple[int | None, str | None]:
+    """Return (required_devices, device_type) found in a decorator list."""
+    required_devices: int | None = None
+    device_type: str | None = None
+
+    for dec in decorators:
+        names = _decorator_names(dec)
+
+        # skip_less_device(N)
+        if any(n in ("skip_less_device",) for n in names):
+            if isinstance(dec, ast.Call):
+                val = _get_int_arg(dec)
+                if val is not None:
+                    required_devices = val
+
+        # skip_pre_blackwell / skip_pre_hopper
+        for marker, dtype in _DEVICE_TYPE_MARKERS.items():
+            if any(marker in n for n in names):
+                device_type = dtype
+
+        # pytest.mark.skipif(condition, reason="... requires N GPU(s)...")
+        if any("skipif" in n for n in names) and isinstance(dec, ast.Call):
+            for kw in dec.keywords:
+                if kw.arg == "reason" and isinstance(kw.value, ast.Constant):
+                    m = re.search(r"(\d+)\s+GPU", kw.value.value, re.IGNORECASE)
+                    if m:
+                        required_devices = int(m.group(1))
+
+    return required_devices, device_type
+
+
+def extract_markers(
+    test_file: str,
+    class_name: str | None,
+    test_name: str | None,
+) -> dict:
+    path = Path(test_file)
+    if not path.exists():
+        raise FileNotFoundError(f"Test file not found: {test_file}")
+
+    tree = ast.parse(path.read_text(), filename=test_file)
+
+    class_rd: int | None = None
+    class_dt: str | None = None
+    func_rd: int | None = None
+    func_dt: str | None = None
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ClassDef) and (class_name is None or node.name == class_name):
+            c_rd, c_dt = _extract_from_decorators(node.decorator_list)
+            if c_rd is not None:
+                class_rd = c_rd
+            if c_dt is not None:
+                class_dt = c_dt
+
+            for item in node.body:
+                if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and (
+                    test_name is None or item.name == test_name
+                ):
+                    f_rd, f_dt = _extract_from_decorators(item.decorator_list)
+                    if f_rd is not None:
+                        func_rd = f_rd
+                    if f_dt is not None:
+                        func_dt = f_dt
+
+        # Also handle module-level functions (no enclosing class)
+        if (
+            class_name is None
+            and isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+            and (test_name is None or node.name == test_name)
+        ):
+            f_rd, f_dt = _extract_from_decorators(node.decorator_list)
+            if f_rd is not None:
+                func_rd = f_rd
+            if f_dt is not None:
+                func_dt = f_dt
+
+    # Function-level markers take precedence over class-level
+    required_devices = func_rd if func_rd is not None else class_rd
+    device_type = func_dt if func_dt is not None else class_dt
+
+    return {
+        "required_devices": required_devices if required_devices is not None else 1,
+        "device_type": device_type or "",
+        "sources": {
+            "required_devices": "derived" if required_devices is not None else "default",
+            "device_type": "derived" if device_type is not None else "default",
+        },
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--test-cmd",
+        help="Full pytest command (file path and node-id are parsed automatically)",
+    )
+    group.add_argument("--test-file", help="Path to the test file")
+    parser.add_argument("--class-name", help="Test class name (optional)")
+    parser.add_argument("--test-name", help="Test function name (optional)")
+    args = parser.parse_args()
+
+    if args.test_cmd:
+        try:
+            test_file, class_name, test_name = _parse_test_cmd(args.test_cmd)
+        except ValueError as exc:
+            print(json.dumps({"error": str(exc)}))
+            sys.exit(1)
+    else:
+        test_file = args.test_file
+        class_name = args.class_name
+        test_name = args.test_name
+
+    try:
+        result = extract_markers(test_file, class_name, test_name)
+    except FileNotFoundError as exc:
+        print(json.dumps({"error": str(exc)}))
+        sys.exit(1)
+
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.claude/skills/trtllm-test-specialist/scripts/generate_benchmark_config.py b/.claude/skills/trtllm-test-specialist/scripts/generate_benchmark_config.py
new file mode 100644
index 000000000000..a89fef406355
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/scripts/generate_benchmark_config.py
@@ -0,0 +1,1184 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Generate a perf-sanity compatible benchmark config YAML (aggregated or disaggregated).
+
+Takes user-provided benchmark parameters and produces a YAML file that can be
+passed to submit.py via --config-file for Slurm job submission.
+
+Config types (--config-type):
+  aggr    Aggregated mode: single-server deployment with server_configs[] format.
+          Use for standard single-role benchmarks (trtllm-serve aggregated).
+  disagg  Disaggregated mode (default): separate ctx/gen workers with worker_config format.
+          Use for disaggregated prefill/decode serving benchmarks.
+
+Supports two generation modes:
+  1. Generate from scratch: provide --model-name and parameters.
+  2. Load from existing config: provide --from-config to load a base YAML and
+     optionally override specific fields with CLI arguments. Config type is
+     auto-detected from the loaded file (server_configs → aggr, worker_config → disagg).
+
+Usage examples:
+
+  # Aggregated: generate from scratch
+  python3 generate_benchmark_config.py \
+    --config-type aggr \
+    --model-name gpt_oss_120b_fp4 \
+    --tp 4 --ep 4 --pp 2 \
+    --max-batch-size 256 --max-num-tokens 20000 \
+    --mtp-layers 1 \
+    --concurrency 256 --iterations 10 \
+    --input-length 1024 --output-length 1024 \
+    --output config.yaml
+
+  # Aggregated: load template and override MTP layers
+  python3 generate_benchmark_config.py \
+    --from-config references/agg_config_template.yaml \
+    --mtp-layers 3 \
+    --output config.yaml
+
+  # Disaggregated: model + hardware essentials
+  python3 generate_benchmark_config.py \
+    --model-name deepseek_r1_0528_fp4_v2 \
+    --gen-tp 8 --ctx-tp 4 \
+    --output config.yaml
+
+  # Disaggregated: from existing config with overrides
+  python3 generate_benchmark_config.py \
+    --from-config examples/disaggregated/slurm/benchmark/config.yaml \
+    --gen-tp 4 --concurrency-list "16,32" \
+    --output config.yaml
+
+  # Disaggregated: full control
+  python3 generate_benchmark_config.py \
+    --model-name deepseek_r1_0528_fp4_v2 \
+    --benchmark-mode e2e \
+    --gpus-per-node 8 \
+    --num-ctx-servers 1 --num-gen-servers 1 \
+    --gen-tp 8 --gen-pp 1 --gen-cp 1 \
+    --gen-max-batch-size 256 --gen-max-num-tokens 512 \
+    --ctx-tp 4 --ctx-pp 1 --ctx-cp 1 \
+    --ctx-max-batch-size 4 --ctx-max-num-tokens 4608 \
+    --input-length 1024 --output-length 1024 \
+    --concurrency-list "16,32,64" \
+    --streaming \
+    --kv-cache-dtype fp8 \
+    --mtp-layers 1 \
+    --cache-backend UCX \
+    --output config.yaml
+"""
+
+import argparse
+import os
+import sys
+
+import yaml
+
+
+def _sanitize_placeholders(obj):
+    """Recursively replace '<placeholder>' strings with None in a config dict.
+
+    Template files (e.g., benchmark_config_template.yaml) contain angle-bracket
+    placeholder strings such as "<dataset_file>" or "<model_path>" for fields the
+    user is expected to fill in.  After loading the template we strip those so they
+    do not propagate as literal strings into the generated config.
+
+    A string is considered a placeholder when it matches the pattern ^<[^>]+>$
+    (entire value is a single angle-bracketed token).
+    """
+    import re
+
+    _placeholder_re = re.compile(r"^<[^>]+>$")
+
+    if isinstance(obj, dict):
+        return {k: _sanitize_placeholders(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_sanitize_placeholders(v) for v in obj]
+    if isinstance(obj, str) and _placeholder_re.match(obj):
+        return None
+    return obj
+
+
+def load_and_normalize_config(config_path):
+    """Load an existing config YAML and normalize it to perf-sanity format.
+
+    Handles both perf-sanity style configs (which have metadata/benchmark/hardware/worker_config)
+    and examples/disaggregated/slurm/benchmark/config.yaml format (which has slurm/environment
+    sections). Strips non-perf-sanity sections (slurm, environment) and keeps the rest.
+    Placeholder strings of the form "<name>" (used in template files) are replaced with None.
+
+    Returns:
+        dict: Normalized config dict compatible with perf-sanity format.
+    """
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+
+    config = _sanitize_placeholders(config)
+
+    # If it already has the perf-sanity structure, return as-is (minus extra sections)
+    normalized = {}
+
+    # Copy over the sections we care about
+    if "metadata" in config:
+        normalized["metadata"] = config["metadata"]
+    elif "model_name" in config.get("environment", {}):
+        # Derive metadata from environment.model_path if available
+        model_path = config["environment"].get("model_path", "unknown")
+        normalized["metadata"] = {
+            "model_name": os.path.basename(model_path),
+            "benchmark_type": "",
+        }
+
+    if "benchmark" in config:
+        normalized["benchmark"] = config["benchmark"]
+
+    if "hardware" in config:
+        normalized["hardware"] = config["hardware"]
+
+    if "worker_config" in config:
+        normalized["worker_config"] = config["worker_config"]
+
+    # Aggregated format: pass server_configs through unchanged
+    if "server_configs" in config:
+        normalized["server_configs"] = config["server_configs"]
+
+    if "environment" in config:
+        # Only keep worker_env_var and server_env_var
+        env = config["environment"]
+        normalized["environment"] = {
+            "worker_env_var": env.get("worker_env_var", ""),
+            "server_env_var": env.get("server_env_var", ""),
+        }
+
+    if "profiling" in config:
+        normalized["profiling"] = config["profiling"]
+
+    if "accuracy" in config:
+        normalized["accuracy"] = _normalize_accuracy_section(config["accuracy"])
+
+    return normalized
+
+
+def _normalize_accuracy_section(acc):
+    """Convert flat accuracy format (perf-sanity) to nested tasks format (examples).
+
+    Flat format (perf-sanity):
+        accuracy:
+          enable_accuracy_test: false
+          model: local-completions
+          tasks: gsm8k
+          model_args_extra: "..."
+
+    Nested format (examples / target):
+        accuracy:
+          enable_accuracy_test: false
+          tasks:
+            gsm8k:
+              model: local-completions
+              model_args_extra: "..."
+              extra_kwargs: {}
+    """
+    if not acc:
+        return acc
+    tasks = acc.get("tasks")
+    # Already nested format — tasks is a dict mapping task name → config
+    if isinstance(tasks, dict):
+        return acc
+    # Flat format — tasks is a string like "gsm8k"
+    task_name = tasks or "gsm8k"
+    return {
+        "enable_accuracy_test": acc.get("enable_accuracy_test", False),
+        "tasks": {
+            task_name: {
+                "model": acc.get("model", "local-completions"),
+                "model_args_extra": acc.get("model_args_extra", ""),
+                "extra_kwargs": {},
+            }
+        },
+    }
+
+
+def apply_overrides(config, args, explicitly_set):
+    """Apply CLI argument overrides to a loaded config.
+
+    Only overrides fields for arguments that were explicitly provided on the
+    command line (tracked via explicitly_set set).
+
+    Args:
+        config: The base config dict to modify in place.
+        args: Parsed argparse namespace.
+        explicitly_set: Set of argument dest names that were explicitly provided.
+    """
+    # --- metadata overrides ---
+    if "model_name" in explicitly_set and args.model_name:
+        config.setdefault("metadata", {})["model_name"] = args.model_name
+
+    # --- benchmark overrides ---
+    bm = config.setdefault("benchmark", {})
+    bm_map = {
+        "benchmark_mode": "mode",
+        "input_length": "input_length",
+        "output_length": "output_length",
+        "concurrency_list": "concurrency_list",
+        "multi_round": "multi_round",
+        "streaming": "streaming",
+        "dataset_file": "dataset_file",
+        "benchmark_ratio": "benchmark_ratio",
+    }
+    for arg_name, config_key in bm_map.items():
+        if arg_name in explicitly_set:
+            bm[config_key] = getattr(args, arg_name)
+
+    # Update benchmark_type in metadata if input/output lengths changed
+    if "input_length" in explicitly_set or "output_length" in explicitly_set:
+        isl = bm.get("input_length", args.input_length)
+        osl = bm.get("output_length", args.output_length)
+        config.setdefault("metadata", {})["benchmark_type"] = (
+            f"{isl // 1000}k{osl // 1000}k" if isl >= 1000 and osl >= 1000 else f"{isl}-{osl}"
+        )
+
+    # --- hardware overrides ---
+    hw = config.setdefault("hardware", {})
+    hw_map = {
+        "gpus_per_node": "gpus_per_node",
+        "num_ctx_servers": "num_ctx_servers",
+        "num_gen_servers": "num_gen_servers",
+    }
+    for arg_name, config_key in hw_map.items():
+        if arg_name in explicitly_set:
+            hw[config_key] = getattr(args, arg_name)
+
+    # --- environment overrides ---
+    env = config.setdefault("environment", {})
+    if "worker_env_var" in explicitly_set:
+        env["worker_env_var"] = args.worker_env_var
+    if "server_env_var" in explicitly_set:
+        env["server_env_var"] = args.server_env_var
+
+    # --- worker_config gen overrides ---
+    wc = config.setdefault("worker_config", {})
+    gen = wc.setdefault("gen", {})
+    gen_simple = {
+        "gen_tp": "tensor_parallel_size",
+        "gen_pp": "pipeline_parallel_size",
+        "gen_cp": "context_parallel_size",
+        "gen_max_batch_size": "max_batch_size",
+        "gen_max_num_tokens": "max_num_tokens",
+        "gen_enable_attention_dp": "enable_attention_dp",
+        "stream_interval": "stream_interval",
+        "num_postprocess_workers": "num_postprocess_workers",
+    }
+    for arg_name, config_key in gen_simple.items():
+        if arg_name in explicitly_set:
+            gen[config_key] = getattr(args, arg_name)
+
+    if "gen_ep" in explicitly_set:
+        if args.gen_ep > 0:
+            gen["moe_expert_parallel_size"] = args.gen_ep
+        else:
+            gen.pop("moe_expert_parallel_size", None)
+
+    if "gen_max_seq_len" in explicitly_set and args.gen_max_seq_len > 0:
+        gen["max_seq_len"] = args.gen_max_seq_len
+
+    if "gen_enable_lm_head_tp_in_adp" in explicitly_set:
+        if args.gen_enable_lm_head_tp_in_adp:
+            gen["enable_lm_head_tp_in_adp"] = True
+        else:
+            gen.pop("enable_lm_head_tp_in_adp", None)
+
+    if "gen_disable_overlap_scheduler" in explicitly_set:
+        if args.gen_disable_overlap_scheduler:
+            gen["disable_overlap_scheduler"] = True
+        else:
+            gen.pop("disable_overlap_scheduler", None)
+
+    if "trust_remote_code" in explicitly_set:
+        if args.trust_remote_code:
+            gen["trust_remote_code"] = True
+        else:
+            gen.pop("trust_remote_code", None)
+
+    # Gen CUDA graph config
+    if "gen_cuda_graph_padding" in explicitly_set:
+        if args.gen_cuda_graph_padding:
+            cuda_graph = gen.get("cuda_graph_config") or {}
+            cuda_graph["enable_padding"] = True
+            if "gen_cuda_graph_batch_sizes" in explicitly_set and args.gen_cuda_graph_batch_sizes:
+                cuda_graph["batch_sizes"] = [
+                    int(x.strip()) for x in args.gen_cuda_graph_batch_sizes.split(",")
+                ]
+            gen["cuda_graph_config"] = cuda_graph
+        else:
+            gen["cuda_graph_config"] = None
+
+    # Gen KV cache config
+    gen_kv = gen.setdefault("kv_cache_config", {})
+    if "kv_cache_dtype" in explicitly_set:
+        gen_kv["dtype"] = args.kv_cache_dtype
+    if "gen_kv_free_fraction" in explicitly_set:
+        gen_kv["free_gpu_memory_fraction"] = args.gen_kv_free_fraction
+    if "kv_tokens_per_block" in explicitly_set and args.kv_tokens_per_block > 0:
+        gen_kv["tokens_per_block"] = args.kv_tokens_per_block
+
+    # Gen MoE config
+    if "moe_backend" in explicitly_set:
+        if args.moe_backend:
+            moe = gen.get("moe_config") or {}
+            moe["backend"] = args.moe_backend
+            if "moe_low_precision_combine" in explicitly_set and args.moe_low_precision_combine:
+                moe["use_low_precision_moe_combine"] = True
+            if "eplb_num_slots" in explicitly_set and args.eplb_num_slots > 0:
+                moe["load_balancer"] = {"num_slots": args.eplb_num_slots}
+            gen["moe_config"] = moe
+        else:
+            gen.pop("moe_config", None)
+
+    # Gen cache transceiver
+    gen_ct = gen.setdefault("cache_transceiver_config", {})
+    if "cache_backend" in explicitly_set:
+        gen_ct["backend"] = args.cache_backend
+    if "cache_max_tokens_in_buffer" in explicitly_set:
+        gen_ct["max_tokens_in_buffer"] = args.cache_max_tokens_in_buffer
+
+    # Gen speculative config
+    if "mtp_layers" in explicitly_set:
+        if args.mtp_layers > 0:
+            gen["speculative_config"] = {
+                "decoding_type": "MTP",
+                "num_nextn_predict_layers": args.mtp_layers,
+            }
+        else:
+            gen.pop("speculative_config", None)
+
+    # --- worker_config ctx overrides ---
+    ctx = wc.setdefault("ctx", {})
+    ctx_simple = {
+        "ctx_tp": "tensor_parallel_size",
+        "ctx_pp": "pipeline_parallel_size",
+        "ctx_cp": "context_parallel_size",
+        "ctx_max_batch_size": "max_batch_size",
+        "ctx_max_num_tokens": "max_num_tokens",
+        "ctx_enable_attention_dp": "enable_attention_dp",
+    }
+    for arg_name, config_key in ctx_simple.items():
+        if arg_name in explicitly_set:
+            ctx[config_key] = getattr(args, arg_name)
+
+    if "ctx_ep" in explicitly_set:
+        if args.ctx_ep > 0:
+            ctx["moe_expert_parallel_size"] = args.ctx_ep
+        else:
+            ctx.pop("moe_expert_parallel_size", None)
+
+    if "ctx_max_seq_len" in explicitly_set and args.ctx_max_seq_len > 0:
+        ctx["max_seq_len"] = args.ctx_max_seq_len
+
+    if "ctx_enable_lm_head_tp_in_adp" in explicitly_set:
+        if args.ctx_enable_lm_head_tp_in_adp:
+            ctx["enable_lm_head_tp_in_adp"] = True
+        else:
+            ctx.pop("enable_lm_head_tp_in_adp", None)
+
+    if "trust_remote_code" in explicitly_set:
+        if args.trust_remote_code:
+            ctx["trust_remote_code"] = True
+        else:
+            ctx.pop("trust_remote_code", None)
+
+    # Ctx KV cache config
+    ctx_kv = ctx.setdefault("kv_cache_config", {})
+    if "kv_cache_dtype" in explicitly_set:
+        ctx_kv["dtype"] = args.kv_cache_dtype
+    if "ctx_kv_free_fraction" in explicitly_set:
+        ctx_kv["free_gpu_memory_fraction"] = args.ctx_kv_free_fraction
+    if "kv_tokens_per_block" in explicitly_set and args.kv_tokens_per_block > 0:
+        ctx_kv["tokens_per_block"] = args.kv_tokens_per_block
+
+    # Ctx MoE config
+    if "moe_backend" in explicitly_set:
+        if args.moe_backend:
+            moe = ctx.get("moe_config") or {}
+            moe["backend"] = args.moe_backend
+            ctx["moe_config"] = moe
+        else:
+            ctx.pop("moe_config", None)
+
+    # Ctx cache transceiver
+    ctx_ct = ctx.setdefault("cache_transceiver_config", {})
+    if "cache_backend" in explicitly_set:
+        ctx_ct["backend"] = args.cache_backend
+    if "cache_max_tokens_in_buffer" in explicitly_set:
+        ctx_ct["max_tokens_in_buffer"] = args.cache_max_tokens_in_buffer
+
+    # Ctx speculative config
+    if "mtp_layers" in explicitly_set:
+        if args.mtp_layers > 0:
+            ctx["speculative_config"] = {
+                "decoding_type": "MTP",
+                "num_nextn_predict_layers": args.mtp_layers,
+            }
+        else:
+            ctx.pop("speculative_config", None)
+
+    # --- accuracy overrides ---
+    if any(
+        k in explicitly_set
+        for k in (
+            "enable_accuracy_test",
+            "accuracy_task",
+            "accuracy_model",
+            "accuracy_model_args_extra",
+            "accuracy_trust_remote_code",
+        )
+    ):
+        acc = config.setdefault(
+            "accuracy",
+            {
+                "enable_accuracy_test": False,
+                "tasks": {},
+            },
+        )
+        # Normalize to nested format first
+        acc = _normalize_accuracy_section(acc)
+        config["accuracy"] = acc
+
+        if "enable_accuracy_test" in explicitly_set:
+            acc["enable_accuracy_test"] = args.enable_accuracy_test
+
+        # Build / update the task entry
+        task_name = getattr(args, "accuracy_task", "gsm8k") or "gsm8k"
+        tasks = acc.setdefault("tasks", {})
+        task_cfg = tasks.setdefault(
+            task_name,
+            {
+                "model": "local-completions",
+                "model_args_extra": "",
+                "extra_kwargs": {},
+            },
+        )
+        if "accuracy_model" in explicitly_set:
+            task_cfg["model"] = args.accuracy_model
+        if "accuracy_model_args_extra" in explicitly_set:
+            task_cfg["model_args_extra"] = args.accuracy_model_args_extra
+        if "accuracy_trust_remote_code" in explicitly_set and args.accuracy_trust_remote_code:
+            task_cfg.setdefault("extra_kwargs", {})["trust_remote_code"] = True
+
+    return config
+
+
+def _detect_config_type(config):
+    """Detect if a loaded config is aggregated or disaggregated."""
+    if "server_configs" in config:
+        return "aggr"
+    if "worker_config" in config:
+        return "disagg"
+    return None
+
+
+def _apply_single_server_config_overrides(sc, args, explicitly_set):
+    """Apply CLI overrides to a single server_config dict (aggregated mode)."""
+    sc_map = {
+        "tp": "tensor_parallel_size",
+        "pp": "pipeline_parallel_size",
+        "cp": "context_parallel_size",
+        "max_batch_size": "max_batch_size",
+        "max_num_tokens": "max_num_tokens",
+        "attn_backend": "attn_backend",
+        "enable_attention_dp": "enable_attention_dp",
+        "num_postprocess_workers": "num_postprocess_workers",
+        "stream_interval": "stream_interval",
+    }
+    for arg_name, config_key in sc_map.items():
+        if arg_name in explicitly_set:
+            sc[config_key] = getattr(args, arg_name)
+
+    if "ep" in explicitly_set:
+        if args.ep > 0:
+            sc["moe_expert_parallel_size"] = args.ep
+        else:
+            sc.pop("moe_expert_parallel_size", None)
+
+    if "max_seq_len" in explicitly_set and args.max_seq_len > 0:
+        sc["max_seq_len"] = args.max_seq_len
+
+    if "enable_adp_balance" in explicitly_set:
+        if args.enable_adp_balance and sc.get("enable_attention_dp", False):
+            sc["attention_dp_config"] = {"enable_balance": True}
+        elif not args.enable_adp_balance:
+            sc.pop("attention_dp_config", None)
+
+    if "moe_backend" in explicitly_set:
+        if args.moe_backend:
+            moe = sc.get("moe_config") or {}
+            moe["backend"] = args.moe_backend
+            if "moe_low_precision_combine" in explicitly_set and args.moe_low_precision_combine:
+                moe["use_low_precision_moe_combine"] = True
+            if "eplb_num_slots" in explicitly_set and args.eplb_num_slots > 0:
+                moe["load_balancer"] = {"num_slots": args.eplb_num_slots}
+            sc["moe_config"] = moe
+        else:
+            sc.pop("moe_config", None)
+
+    if "kv_cache_dtype" in explicitly_set:
+        sc.setdefault("kv_cache_config", {})["dtype"] = args.kv_cache_dtype
+    if "kv_free_fraction" in explicitly_set:
+        sc.setdefault("kv_cache_config", {})["free_gpu_memory_fraction"] = args.kv_free_fraction
+    if "kv_tokens_per_block" in explicitly_set and args.kv_tokens_per_block > 0:
+        sc.setdefault("kv_cache_config", {})["tokens_per_block"] = args.kv_tokens_per_block
+
+    if "trust_remote_code" in explicitly_set:
+        if args.trust_remote_code:
+            sc["trust_remote_code"] = True
+        else:
+            sc.pop("trust_remote_code", None)
+
+    if "mtp_layers" in explicitly_set:
+        if args.mtp_layers > 0:
+            sc["speculative_config"] = {
+                "decoding_type": "MTP",
+                "num_nextn_predict_layers": args.mtp_layers,
+            }
+        else:
+            sc.pop("speculative_config", None)
+
+    # Client config overrides (apply to all client_configs in this server)
+    for cc in sc.get("client_configs", []):
+        if "input_length" in explicitly_set:
+            cc["isl"] = args.input_length
+        if "output_length" in explicitly_set:
+            cc["osl"] = args.output_length
+        if "concurrency" in explicitly_set and args.concurrency > 0:
+            cc["concurrency"] = args.concurrency
+        if "iterations" in explicitly_set:
+            cc["iterations"] = args.iterations
+        if "dataset_file" in explicitly_set and args.dataset_file:
+            cc["dataset_file"] = args.dataset_file
+
+
+def apply_aggr_overrides(config, args, explicitly_set):
+    """Apply CLI argument overrides to a loaded aggregated config.
+
+    Applies scalar overrides to all server_configs unless --server-name
+    is specified, in which case only the matching entry is updated.
+    """
+    if "model_name" in explicitly_set and args.model_name:
+        config.setdefault("metadata", {})["model_name"] = args.model_name
+    if "supported_gpus" in explicitly_set and args.supported_gpus:
+        config.setdefault("metadata", {})["supported_gpus"] = [
+            g.strip() for g in args.supported_gpus.split(",")
+        ]
+
+    hw = config.setdefault("hardware", {})
+    if "gpus_per_node" in explicitly_set:
+        hw["gpus_per_node"] = args.gpus_per_node
+
+    server_configs = config.get("server_configs", [])
+    targets = server_configs
+    if "server_name" in explicitly_set and args.server_name:
+        targets = [sc for sc in server_configs if sc.get("name") == args.server_name]
+        if not targets:
+            print(
+                f"Warning: no server_config named '{args.server_name}' found; "
+                "applying overrides to all server_configs",
+                file=sys.stderr,
+            )
+            targets = server_configs
+
+    for sc in targets:
+        _apply_single_server_config_overrides(sc, args, explicitly_set)
+
+    return config
+
+
+def build_server_config(args):
+    """Build a server_config entry dict for aggregated mode."""
+    isl = args.input_length
+    osl = args.output_length
+
+    # Auto-derive server name if not provided
+    server_name = args.server_name
+    if not server_name:
+        parts = [args.model_name, f"tp{args.tp}"]
+        if args.ep > 0:
+            parts.append(f"ep{args.ep}")
+        if args.pp > 1:
+            parts.append(f"pp{args.pp}")
+        if args.mtp_layers > 0:
+            parts.append(f"mtp{args.mtp_layers}")
+        server_name = "_".join(parts)
+
+    sc = {
+        "name": server_name,
+        "model_name": args.model_name,
+        "tensor_parallel_size": args.tp,
+        "pipeline_parallel_size": args.pp,
+        "max_batch_size": args.max_batch_size,
+        "max_num_tokens": args.max_num_tokens,
+        "attn_backend": args.attn_backend,
+        "enable_attention_dp": args.enable_attention_dp,
+    }
+
+    if args.cp > 1:
+        sc["context_parallel_size"] = args.cp
+    if args.ep > 0:
+        sc["moe_expert_parallel_size"] = args.ep
+    if args.max_seq_len > 0:
+        sc["max_seq_len"] = args.max_seq_len
+
+    if args.enable_attention_dp and args.enable_adp_balance:
+        sc["attention_dp_config"] = {"enable_balance": True}
+
+    if args.moe_backend:
+        moe = {"backend": args.moe_backend}
+        if args.moe_low_precision_combine:
+            moe["use_low_precision_moe_combine"] = True
+        if args.eplb_num_slots > 0:
+            moe["load_balancer"] = {"num_slots": args.eplb_num_slots}
+        sc["moe_config"] = moe
+
+    if args.gen_cuda_graph_padding:
+        cuda_graph = {"enable_padding": True, "max_batch_size": args.max_batch_size}
+        if args.gen_cuda_graph_batch_sizes:
+            cuda_graph["batch_sizes"] = [
+                int(x.strip()) for x in args.gen_cuda_graph_batch_sizes.split(",")
+            ]
+        sc["cuda_graph_config"] = cuda_graph
+
+    kv_cache = {
+        "dtype": args.kv_cache_dtype,
+        "enable_block_reuse": False,
+        "free_gpu_memory_fraction": args.kv_free_fraction,
+    }
+    if args.kv_tokens_per_block > 0:
+        kv_cache["tokens_per_block"] = args.kv_tokens_per_block
+    sc["kv_cache_config"] = kv_cache
+
+    if args.mtp_layers > 0:
+        sc["speculative_config"] = {
+            "decoding_type": "MTP",
+            "num_nextn_predict_layers": args.mtp_layers,
+        }
+
+    if args.trust_remote_code:
+        sc["trust_remote_code"] = True
+
+    sc["num_postprocess_workers"] = args.num_postprocess_workers
+    sc["stream_interval"] = args.stream_interval
+
+    # Auto-derive client config name
+    client_name = args.client_name
+    if not client_name:
+        isl_str = f"{isl // 1000}k" if isl >= 1000 else str(isl)
+        osl_str = f"{osl // 1000}k" if osl >= 1000 else str(osl)
+        client_name = f"con{args.concurrency}_iter{args.iterations}_{isl_str}{osl_str}"
+
+    client = {
+        "name": client_name,
+        "concurrency": args.concurrency,
+        "iterations": args.iterations,
+        "isl": isl,
+        "osl": osl,
+        "backend": args.client_backend,
+    }
+    if args.dataset_file:
+        client["dataset_file"] = args.dataset_file
+
+    sc["client_configs"] = [client]
+    return sc
+
+
+def build_aggr_config(args):
+    """Build a complete aggregated benchmark config dict."""
+    config = {
+        "metadata": {
+            "model_name": args.model_name,
+        },
+        "hardware": {
+            "gpus_per_node": args.gpus_per_node,
+        },
+        "server_configs": [build_server_config(args)],
+    }
+
+    if args.supported_gpus:
+        config["metadata"]["supported_gpus"] = [g.strip() for g in args.supported_gpus.split(",")]
+
+    return config
+
+
+def build_gen_worker_config(args):
+    """Build generation worker configuration dict."""
+    gen = {
+        "tensor_parallel_size": args.gen_tp,
+        "pipeline_parallel_size": args.gen_pp,
+        "context_parallel_size": args.gen_cp,
+        "max_batch_size": args.gen_max_batch_size,
+        "max_num_tokens": args.gen_max_num_tokens,
+        "enable_attention_dp": args.gen_enable_attention_dp,
+        "print_iter_log": True,
+        "stream_interval": args.stream_interval,
+        "num_postprocess_workers": args.num_postprocess_workers,
+    }
+
+    if args.gen_ep > 0:
+        gen["moe_expert_parallel_size"] = args.gen_ep
+    if args.gen_max_seq_len > 0:
+        gen["max_seq_len"] = args.gen_max_seq_len
+    if args.gen_enable_lm_head_tp_in_adp:
+        gen["enable_lm_head_tp_in_adp"] = True
+    if args.gen_disable_overlap_scheduler:
+        gen["disable_overlap_scheduler"] = True
+    if args.trust_remote_code:
+        gen["trust_remote_code"] = True
+
+    # CUDA graph config
+    if args.gen_cuda_graph_padding:
+        cuda_graph = {"enable_padding": True}
+        if args.gen_cuda_graph_batch_sizes:
+            cuda_graph["batch_sizes"] = [
+                int(x.strip()) for x in args.gen_cuda_graph_batch_sizes.split(",")
+            ]
+        gen["cuda_graph_config"] = cuda_graph
+    else:
+        gen["cuda_graph_config"] = None
+
+    # KV cache config
+    kv_cache = {
+        "enable_block_reuse": False,
+        "free_gpu_memory_fraction": args.gen_kv_free_fraction,
+        "dtype": args.kv_cache_dtype,
+    }
+    if args.kv_tokens_per_block > 0:
+        kv_cache["tokens_per_block"] = args.kv_tokens_per_block
+    gen["kv_cache_config"] = kv_cache
+
+    # MoE config
+    if args.moe_backend:
+        moe = {"backend": args.moe_backend}
+        if args.moe_low_precision_combine:
+            moe["use_low_precision_moe_combine"] = True
+        if args.eplb_num_slots > 0:
+            moe["load_balancer"] = {"num_slots": args.eplb_num_slots}
+        gen["moe_config"] = moe
+
+    # Cache transceiver config
+    gen["cache_transceiver_config"] = {
+        "max_tokens_in_buffer": args.cache_max_tokens_in_buffer,
+        "backend": args.cache_backend,
+    }
+
+    # Speculative config
+    if args.mtp_layers > 0:
+        gen["speculative_config"] = {
+            "decoding_type": "MTP",
+            "num_nextn_predict_layers": args.mtp_layers,
+        }
+
+    return gen
+
+
+def build_ctx_worker_config(args):
+    """Build context worker configuration dict."""
+    ctx = {
+        "tensor_parallel_size": args.ctx_tp,
+        "pipeline_parallel_size": args.ctx_pp,
+        "context_parallel_size": args.ctx_cp,
+        "max_batch_size": args.ctx_max_batch_size,
+        "max_num_tokens": args.ctx_max_num_tokens,
+        "enable_attention_dp": args.ctx_enable_attention_dp,
+        "print_iter_log": True,
+        "disable_overlap_scheduler": True,
+        "cuda_graph_config": None,
+    }
+
+    if args.ctx_ep > 0:
+        ctx["moe_expert_parallel_size"] = args.ctx_ep
+    if args.ctx_max_seq_len > 0:
+        ctx["max_seq_len"] = args.ctx_max_seq_len
+    if args.ctx_enable_lm_head_tp_in_adp:
+        ctx["enable_lm_head_tp_in_adp"] = True
+    if args.trust_remote_code:
+        ctx["trust_remote_code"] = True
+
+    # KV cache config
+    kv_cache = {
+        "enable_block_reuse": False,
+        "free_gpu_memory_fraction": args.ctx_kv_free_fraction,
+        "dtype": args.kv_cache_dtype,
+    }
+    if args.kv_tokens_per_block > 0:
+        kv_cache["tokens_per_block"] = args.kv_tokens_per_block
+    ctx["kv_cache_config"] = kv_cache
+
+    # MoE config (ctx may use different backend)
+    if args.moe_backend:
+        moe = {"backend": args.moe_backend}
+        ctx["moe_config"] = moe
+
+    # Cache transceiver config
+    ctx["cache_transceiver_config"] = {
+        "max_tokens_in_buffer": args.cache_max_tokens_in_buffer,
+        "backend": args.cache_backend,
+    }
+
+    # Speculative config (same as gen if set)
+    if args.mtp_layers > 0:
+        ctx["speculative_config"] = {
+            "decoding_type": "MTP",
+            "num_nextn_predict_layers": args.mtp_layers,
+        }
+
+    return ctx
+
+
+def build_config(args):
+    """Build the complete config dict."""
+    isl = args.input_length
+    osl = args.output_length
+
+    config = {
+        "metadata": {
+            "model_name": args.model_name,
+            "benchmark_type": f"{isl // 1000}k{osl // 1000}k"
+            if isl >= 1000 and osl >= 1000
+            else f"{isl}-{osl}",
+        },
+        "benchmark": {
+            "mode": args.benchmark_mode,
+            "use_nv_sa_benchmark": False,
+            "multi_round": args.multi_round,
+            "benchmark_ratio": args.benchmark_ratio,
+            "streaming": args.streaming,
+            "concurrency_list": args.concurrency_list,
+            "input_length": isl,
+            "output_length": osl,
+        },
+        "hardware": {
+            "gpus_per_node": args.gpus_per_node,
+            "num_ctx_servers": args.num_ctx_servers,
+            "num_gen_servers": args.num_gen_servers,
+        },
+        "environment": {
+            "worker_env_var": args.worker_env_var,
+            "server_env_var": args.server_env_var,
+        },
+        "profiling": {
+            "nsys_on": False,
+        },
+        "accuracy": {
+            "enable_accuracy_test": args.enable_accuracy_test,
+            "tasks": {
+                args.accuracy_task: {
+                    "model": args.accuracy_model,
+                    "model_args_extra": args.accuracy_model_args_extra,
+                    "extra_kwargs": {
+                        **({"trust_remote_code": True} if args.accuracy_trust_remote_code else {}),
+                    },
+                }
+            }
+            if args.enable_accuracy_test
+            else {},
+        },
+        "worker_config": {
+            "gen": build_gen_worker_config(args),
+            "ctx": build_ctx_worker_config(args),
+        },
+    }
+
+    if args.dataset_file:
+        config["benchmark"]["dataset_file"] = args.dataset_file
+
+    return config
+
+
+def get_explicitly_set_args(parser, argv=None):
+    """Determine which arguments were explicitly provided on the command line.
+
+    Parses argv twice: once normally, once with all defaults suppressed.
+    Arguments present in the suppressed parse were explicitly provided.
+
+    Returns:
+        tuple: (args, explicitly_set) where explicitly_set is a set of dest names.
+    """
+    args = parser.parse_args(argv)
+
+    # Parse again with defaults suppressed to find what was explicitly set
+    parser_no_defaults = argparse.ArgumentParser(parents=[], add_help=False)
+    for action in parser._actions:
+        if action.dest == "help":
+            continue
+        kwargs = {
+            "dest": action.dest,
+            "default": argparse.SUPPRESS,
+        }
+        if isinstance(action, argparse._StoreTrueAction):
+            kwargs["action"] = "store_true"
+        elif isinstance(action, argparse._StoreFalseAction):
+            kwargs["action"] = "store_false"
+        elif isinstance(action, argparse._StoreAction):
+            kwargs["type"] = action.type
+            kwargs["nargs"] = action.nargs
+            if action.choices:
+                kwargs["choices"] = action.choices
+
+        # Use the same option strings
+        if action.option_strings:
+            parser_no_defaults.add_argument(*action.option_strings, **kwargs)
+        elif action.dest != "help":
+            kwargs.pop("dest", None)
+            parser_no_defaults.add_argument(action.dest, **kwargs)
+
+    ns_explicit, _ = parser_no_defaults.parse_known_args(argv)
+    explicitly_set = set(vars(ns_explicit).keys())
+
+    return args, explicitly_set
+
+
+def main():
+    p = argparse.ArgumentParser(
+        description="Generate aggregated or disaggregated benchmark config YAML for perf-sanity submission"
+    )
+
+    p.add_argument(
+        "--model-name",
+        default=None,
+        help="Model identifier (must match MODEL_PATH_DICT in test_perf_sanity.py)",
+    )
+    p.add_argument("--output", "-o", required=True, help="Output YAML file path")
+    p.add_argument(
+        "--from-config",
+        default=None,
+        help="Path to an existing benchmark config YAML to use as base. "
+        "Config type is auto-detected (server_configs → aggr, worker_config → disagg).",
+    )
+    p.add_argument(
+        "--config-type",
+        default="disagg",
+        choices=["aggr", "disagg"],
+        help="Config type to generate from scratch: 'aggr' (aggregated) or "
+        "'disagg' (disaggregated, default). Ignored when --from-config auto-detects the type.",
+    )
+
+    # ── Aggregated server config ──────────────────────────────────────────────
+    aggr = p.add_argument_group("aggregated server config (--config-type aggr)")
+    aggr.add_argument("--tp", type=int, default=8, help="Tensor parallel size (aggr server)")
+    aggr.add_argument("--pp", type=int, default=1, help="Pipeline parallel size (aggr server)")
+    aggr.add_argument(
+        "--ep", type=int, default=0, help="Expert (MoE) parallel size (aggr server; 0 = disabled)"
+    )
+    aggr.add_argument(
+        "--cp", type=int, default=1, help="Context parallel size (aggr server; omitted when 1)"
+    )
+    aggr.add_argument(
+        "--max-batch-size", type=int, default=256, help="Server max_batch_size (aggr)"
+    )
+    aggr.add_argument(
+        "--max-num-tokens", type=int, default=20000, help="Server max_num_tokens (aggr)"
+    )
+    aggr.add_argument(
+        "--max-seq-len", type=int, default=0, help="Server max_seq_len (aggr; 0 = omit)"
+    )
+    aggr.add_argument(
+        "--attn-backend", default="TRTLLM", help="Attention backend (aggr; default: TRTLLM)"
+    )
+    aggr.add_argument(
+        "--enable-attention-dp",
+        action="store_true",
+        default=False,
+        help="Enable attention data parallelism (aggr)",
+    )
+    aggr.add_argument(
+        "--enable-adp-balance",
+        action="store_true",
+        default=False,
+        help="Set attention_dp_config.enable_balance=true (aggr, requires --enable-attention-dp)",
+    )
+    aggr.add_argument(
+        "--server-name",
+        default="",
+        help="Name for the server_config entry (auto-derived if empty). "
+        "When used with --from-config, limits overrides to this named entry.",
+    )
+    aggr.add_argument(
+        "--supported-gpus",
+        default="",
+        help="Comma-separated supported GPU types, e.g. 'B200,GB200' (aggr metadata)",
+    )
+    aggr.add_argument(
+        "--kv-free-fraction",
+        type=float,
+        default=0.8,
+        help="KV cache free GPU memory fraction (aggr)",
+    )
+
+    # ── Aggregated client config ───────────────────────────────────────────────
+    client = p.add_argument_group("aggregated client config (--config-type aggr)")
+    client.add_argument(
+        "--concurrency",
+        type=int,
+        default=256,
+        help="Benchmark concurrency for client config (aggr)",
+    )
+    client.add_argument(
+        "--iterations", type=int, default=10, help="Benchmark iterations for client config (aggr)"
+    )
+    client.add_argument(
+        "--client-backend", default="openai", help="Client backend (aggr; default: openai)"
+    )
+    client.add_argument(
+        "--client-name", default="", help="Name for the client_config entry (auto-derived if empty)"
+    )
+
+    # ── Disaggregated / shared ────────────────────────────────────────────────
+    p.add_argument(
+        "--benchmark-mode",
+        default="e2e",
+        choices=["e2e"],
+        help="Benchmark mode. Only 'e2e' is supported for benchmark tests.",
+    )
+    p.add_argument("--input-length", type=int, default=1024)
+    p.add_argument("--output-length", type=int, default=1024)
+    p.add_argument("--concurrency-list", default="16")
+    p.add_argument("--multi-round", type=int, default=8)
+    p.add_argument("--streaming", action="store_true", default=True)
+    p.add_argument("--no-streaming", dest="streaming", action="store_false")
+    p.add_argument("--dataset-file", default="")
+    p.add_argument("--benchmark-ratio", type=float, default=0.8)
+    p.add_argument("--gpus-per-node", type=int, default=8)
+    p.add_argument("--num-ctx-servers", type=int, default=1)
+    p.add_argument("--num-gen-servers", type=int, default=1)
+    p.add_argument("--gen-tp", type=int, default=8)
+    p.add_argument("--gen-pp", type=int, default=1)
+    p.add_argument("--gen-cp", type=int, default=1)
+    p.add_argument("--gen-ep", type=int, default=0)
+    p.add_argument("--gen-max-batch-size", type=int, default=256)
+    p.add_argument("--gen-max-num-tokens", type=int, default=512)
+    p.add_argument("--gen-max-seq-len", type=int, default=0)
+    p.add_argument("--gen-enable-attention-dp", action="store_true", default=False)
+    p.add_argument("--gen-enable-lm-head-tp-in-adp", action="store_true", default=False)
+    p.add_argument("--gen-disable-overlap-scheduler", action="store_true", default=False)
+    p.add_argument("--ctx-tp", type=int, default=4)
+    p.add_argument("--ctx-pp", type=int, default=1)
+    p.add_argument("--ctx-cp", type=int, default=1)
+    p.add_argument("--ctx-ep", type=int, default=0)
+    p.add_argument("--ctx-max-batch-size", type=int, default=4)
+    p.add_argument("--ctx-max-num-tokens", type=int, default=4608)
+    p.add_argument("--ctx-max-seq-len", type=int, default=0)
+    p.add_argument("--ctx-enable-attention-dp", action="store_true", default=True)
+    p.add_argument("--ctx-no-attention-dp", dest="ctx_enable_attention_dp", action="store_false")
+    p.add_argument("--ctx-enable-lm-head-tp-in-adp", action="store_true", default=False)
+    p.add_argument("--kv-cache-dtype", default="fp8")
+    p.add_argument("--gen-kv-free-fraction", type=float, default=0.8)
+    p.add_argument("--ctx-kv-free-fraction", type=float, default=0.85)
+    p.add_argument("--kv-tokens-per-block", type=int, default=0)
+    p.add_argument("--cache-backend", default="DEFAULT", choices=["DEFAULT", "UCX", "NIXL"])
+    p.add_argument("--cache-max-tokens-in-buffer", type=int, default=4608)
+    p.add_argument("--moe-backend", default="")
+    p.add_argument("--moe-low-precision-combine", action="store_true", default=False)
+    p.add_argument("--eplb-num-slots", type=int, default=0)
+    p.add_argument("--mtp-layers", type=int, default=0)
+    p.add_argument("--gen-cuda-graph-padding", action="store_true", default=True)
+    p.add_argument(
+        "--gen-no-cuda-graph-padding", dest="gen_cuda_graph_padding", action="store_false"
+    )
+    p.add_argument("--gen-cuda-graph-batch-sizes", default="")
+    p.add_argument(
+        "--worker-env-var",
+        default="TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1",
+    )
+    p.add_argument("--server-env-var", default="TRTLLM_SERVER_DISABLE_GC=1")
+    p.add_argument("--trust-remote-code", action="store_true", default=False)
+    p.add_argument("--stream-interval", type=int, default=20)
+    p.add_argument("--num-postprocess-workers", type=int, default=4)
+    # Accuracy testing (mirrors examples/disaggregated/slurm/benchmark/config.yaml structure)
+    p.add_argument(
+        "--enable-accuracy-test",
+        action="store_true",
+        default=False,
+        help="Enable accuracy evaluation after benchmark",
+    )
+    p.add_argument("--accuracy-task", default="gsm8k", help="lm_eval task name (default: gsm8k)")
+    p.add_argument(
+        "--accuracy-model",
+        default="local-completions",
+        choices=["local-completions", "local-chat-completions"],
+        help="lm_eval model type (default: local-completions)",
+    )
+    p.add_argument(
+        "--accuracy-model-args-extra",
+        default=(
+            "num_concurrent=512,max_retries=3,tokenized_requests=false,"
+            "timeout=7200,max_gen_toks=16384"
+        ),
+        help="Extra model_args passed to lm_eval",
+    )
+    p.add_argument(
+        "--accuracy-trust-remote-code",
+        action="store_true",
+        default=False,
+        help="Pass --trust_remote_code to lm_eval",
+    )
+
+    args, explicitly_set = get_explicitly_set_args(p)
+
+    if args.from_config:
+        if not os.path.exists(args.from_config):
+            print(f"Error: config file not found: {args.from_config}", file=sys.stderr)
+            sys.exit(1)
+
+        config = load_and_normalize_config(args.from_config)
+        # Auto-detect type from loaded config; fall back to --config-type flag
+        detected = _detect_config_type(config)
+        config_type = detected or args.config_type
+        if detected:
+            print(f"Auto-detected config type: {config_type}")
+
+        explicitly_set.discard("output")
+        explicitly_set.discard("from_config")
+        explicitly_set.discard("config_type")
+
+        if config_type == "aggr":
+            apply_aggr_overrides(config, args, explicitly_set)
+        else:
+            apply_overrides(config, args, explicitly_set)
+
+        print(f"Loaded base config from: {args.from_config}")
+        if explicitly_set:
+            print(f"Applied overrides for: {', '.join(sorted(explicitly_set))}")
+    else:
+        if not args.model_name:
+            print(
+                "Error: --model-name is required when generating from scratch "
+                "(not using --from-config)",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+        if args.config_type == "aggr":
+            config = build_aggr_config(args)
+        else:
+            config = build_config(args)
+
+    output_dir = os.path.dirname(os.path.abspath(args.output))
+    os.makedirs(output_dir, exist_ok=True)
+
+    with open(args.output, "w") as f:
+        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
+
+    print(f"Config written to: {args.output}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.claude/skills/trtllm-test-specialist/scripts/generate_report.py b/.claude/skills/trtllm-test-specialist/scripts/generate_report.py
new file mode 100755
index 000000000000..18e59ddd4dd9
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/scripts/generate_report.py
@@ -0,0 +1,297 @@
+#!/usr/bin/env python3
+"""Test Report Generator for TensorRT-LLM test runs.
+
+Parses pytest output and generates a formatted report.
+
+Usage:
+    generate_report.py --output-file <path_to_captured_output> \
+                       --exit-code <int> \
+                       [--format <markdown|html|json>] \
+                       [--report-file <output_path>]
+"""
+
+import argparse
+import json
+import re
+import sys
+from datetime import datetime
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Parsers
+# ---------------------------------------------------------------------------
+
+
+def parse_pytest_output(text: str) -> dict:
+    """Extract structured data from raw pytest output."""
+    result = {
+        "passed": 0,
+        "failed": 0,
+        "error": 0,
+        "skipped": 0,
+        "warnings": 0,
+        "duration_s": None,
+        "failed_tests": [],
+        "error_details": [],
+        "coverage_pct": None,
+        "summary_line": "",
+    }
+
+    # Summary line: "5 passed, 2 failed, 1 error in 3.45s"
+    summary_match = re.search(
+        r"(\d+) passed(?:,\s*(\d+) failed)?(?:,\s*(\d+) error(?:s)?)?"
+        r"(?:,\s*(\d+) warning(?:s)?)?(?:,\s*(\d+) skipped)?.*?in\s+([\d.]+)s",
+        text,
+    )
+    if summary_match:
+        result["passed"] = int(summary_match.group(1) or 0)
+        result["failed"] = int(summary_match.group(2) or 0)
+        result["error"] = int(summary_match.group(3) or 0)
+        result["warnings"] = int(summary_match.group(4) or 0)
+        result["skipped"] = int(summary_match.group(5) or 0)
+        result["duration_s"] = float(summary_match.group(6))
+        result["summary_line"] = summary_match.group(0)
+
+    # Failed test names from short summary section: "FAILED tests/foo/test_bar.py::TestClass::test_method"
+    raw_matches = re.findall(r"FAILED\s+(tests/[\w/.:_-]+)", text)
+    result["failed_tests"] = list(dict.fromkeys(raw_matches))  # deduplicate, preserve order
+
+    # Error details: capture FAILED sections
+    error_sections = re.findall(r"_{5,}\s+([\w/.:_-]+)\s+_{5,}(.*?)(?=_{5,}|\Z)", text, re.DOTALL)
+    for name, detail in error_sections:
+        snippet = detail.strip()[:500]
+        result["error_details"].append({"test": name.strip(), "detail": snippet})
+
+    # Coverage: "TOTAL  1234  56  95%"
+    cov_match = re.search(r"TOTAL\s+\d+\s+\d+\s+(\d+)%", text)
+    if cov_match:
+        result["coverage_pct"] = int(cov_match.group(1))
+
+    return result
+
+
+def parse_benchmark_metrics(text: str) -> list[dict]:
+    """Extract performance metrics from benchmark output."""
+    metrics = []
+
+    # Throughput patterns: "Throughput: 1234.5 samples/sec"
+    for match in re.finditer(r"(?i)(throughput|samples/sec|tokens/s|it/s)[:\s]+([\d.]+)", text):
+        metrics.append({"metric": match.group(1), "value": float(match.group(2))})
+
+    # Latency patterns: "Latency: 15.3 ms"
+    for match in re.finditer(r"(?i)(latency)[:\s]+([\d.]+)\s*(ms|s)?", text):
+        unit = match.group(3) or "ms"
+        metrics.append({"metric": f"Latency ({unit})", "value": float(match.group(2))})
+
+    # GPU utilization: "GPU Util: 85%"
+    for match in re.finditer(r"(?i)(gpu[\s_]?util(?:ization)?)[:\s]+([\d.]+)\s*%?", text):
+        metrics.append({"metric": "GPU Utilization (%)", "value": float(match.group(2))})
+
+    # SOL / MFU
+    for match in re.finditer(r"(?i)(SOL|MFU)[:\s]+([\d.]+)\s*%?", text):
+        metrics.append({"metric": f"{match.group(1)} (%)", "value": float(match.group(2))})
+
+    return metrics
+
+
+def parse_feature_matrix(text: str) -> list[dict]:
+    """Extract feature support results from test output."""
+    features = []
+
+    # Pattern: "PASSED tests/tools/test_trtllm_modeling.py::TestClass::test_feature_name"
+    for match in re.finditer(
+        r"(PASSED|FAILED|ERROR)\s+(tests/tools/test_trtllm_modeling\.py::[:\w]+)", text
+    ):
+        status = match.group(1)
+        test_id = match.group(2)
+        # Convert test ID to readable feature name
+        parts = test_id.split("::")
+        feature = parts[-1].replace("test_", "").replace("_", " ").title()
+        features.append({"feature": feature, "test_id": test_id, "status": status})
+
+    return features
+
+
+# ---------------------------------------------------------------------------
+# Formatters
+# ---------------------------------------------------------------------------
+
+
+def format_markdown(test_type: str, parsed: dict, exit_code: int, timestamp: str) -> str:
+    total = parsed["passed"] + parsed["failed"] + parsed["error"]
+    status = "PASS" if exit_code == 0 else "FAIL"
+    status_icon = "✅" if exit_code == 0 else "❌"
+
+    lines = [
+        f"# Test Report — {test_type.replace('-', ' ').title()}",
+        "",
+        f"**Status:** {status_icon} {status}  ",
+        f"**Generated:** {timestamp}  ",
+        f"**Duration:** {parsed['duration_s']}s" if parsed["duration_s"] else "",
+        "",
+        "## Summary",
+        "",
+        "| Metric | Count |",
+        "|--------|-------|",
+        f"| Total  | {total} |",
+        f"| Passed | {parsed['passed']} |",
+        f"| Failed | {parsed['failed']} |",
+        f"| Errors | {parsed['error']} |",
+        f"| Skipped | {parsed['skipped']} |",
+    ]
+
+    if parsed.get("coverage_pct") is not None:
+        lines += [f"| Coverage | {parsed['coverage_pct']}% |"]
+
+    if parsed["failed_tests"]:
+        lines += [
+            "",
+            "## Failed Tests",
+            "",
+        ]
+        for t in parsed["failed_tests"]:
+            lines.append(f"- `{t}`")
+
+    if parsed["error_details"]:
+        lines += ["", "## Error Details", ""]
+        for err in parsed["error_details"][:5]:  # cap at 5
+            lines += [
+                f"### `{err['test']}`",
+                "",
+                "```",
+                err["detail"],
+                "```",
+                "",
+            ]
+
+    if test_type == "benchmark" and parsed.get("benchmark_metrics"):
+        lines += ["", "## Performance Metrics", "", "| Metric | Value |", "|--------|-------|"]
+        for m in parsed["benchmark_metrics"]:
+            lines.append(f"| {m['metric']} | {m['value']} |")
+
+    if test_type == "feature-matrix" and parsed.get("features"):
+        lines += [
+            "",
+            "## Feature Support Matrix",
+            "",
+            "| Feature | Status |",
+            "|---------|--------|",
+        ]
+        for f in parsed["features"]:
+            icon = "✅" if f["status"] == "PASSED" else "❌"
+            lines.append(f"| {f['feature']} | {icon} {f['status']} |")
+
+    return "\n".join(filter(lambda x: x is not None, lines))
+
+
+def format_html(test_type: str, parsed: dict, exit_code: int, timestamp: str) -> str:
+    md = format_markdown(test_type, parsed, exit_code, timestamp)
+    # Simple HTML wrap — convert markdown tables and code blocks minimally
+    md_escaped = md.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
+    # Convert ``` blocks
+    md_escaped = re.sub(
+        r"```\n(.*?)```", r"<pre><code>\1</code></pre>", md_escaped, flags=re.DOTALL
+    )
+    # Convert headers
+    for i in range(3, 0, -1):
+        md_escaped = re.sub(
+            r"^#{" + str(i) + r"}\s+(.+)$", rf"<h{i}>\1</h{i}>", md_escaped, flags=re.MULTILINE
+        )
+    md_escaped = md_escaped.replace("\n\n", "<br/><br/>").replace("\n", "<br/>")
+
+    status_color = "#2e7d32" if exit_code == 0 else "#c62828"
+    return f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <title>Test Report — {test_type}</title>
+  <style>
+    body {{ font-family: monospace; padding: 2em; max-width: 900px; margin: auto; }}
+    h1 {{ color: {status_color}; }}
+    pre {{ background: #f5f5f5; padding: 1em; overflow-x: auto; }}
+    table {{ border-collapse: collapse; width: 100%; }}
+    td, th {{ border: 1px solid #ccc; padding: 0.5em; text-align: left; }}
+    th {{ background: #eee; }}
+  </style>
+</head>
+<body>
+{md_escaped}
+</body>
+</html>"""
+
+
+def format_json(test_type: str, parsed: dict, exit_code: int, timestamp: str) -> str:
+    output = {
+        "test_type": test_type,
+        "status": "pass" if exit_code == 0 else "fail",
+        "exit_code": exit_code,
+        "timestamp": timestamp,
+        "summary": {
+            "passed": parsed["passed"],
+            "failed": parsed["failed"],
+            "error": parsed["error"],
+            "skipped": parsed["skipped"],
+            "duration_s": parsed["duration_s"],
+            "coverage_pct": parsed.get("coverage_pct"),
+        },
+        "failed_tests": parsed["failed_tests"],
+    }
+    if test_type == "benchmark":
+        output["benchmark_metrics"] = parsed.get("benchmark_metrics", [])
+    if test_type == "feature-matrix":
+        output["features"] = parsed.get("features", [])
+    return json.dumps(output, indent=2)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate test report from pytest output")
+    parser.add_argument(
+        "--test-type", required=True, choices=["functional", "benchmark", "feature-matrix"]
+    )
+    parser.add_argument("--output-file", required=True, help="Path to captured test output file")
+    parser.add_argument(
+        "--exit-code", required=True, type=int, help="Exit code from test run (0=pass)"
+    )
+    parser.add_argument("--format", default="markdown", choices=["markdown", "html", "json"])
+    parser.add_argument("--report-file", default=None, help="Optional path to write report to")
+    args = parser.parse_args()
+
+    output_path = Path(args.output_file)
+    if not output_path.exists():
+        print(f"Error: output file not found: {output_path}", file=sys.stderr)
+        sys.exit(1)
+
+    raw = output_path.read_text(errors="replace")
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    parsed = parse_pytest_output(raw)
+
+    # Enrich with type-specific parsing
+    if args.test_type == "benchmark":
+        parsed["benchmark_metrics"] = parse_benchmark_metrics(raw)
+    elif args.test_type == "feature-matrix":
+        parsed["features"] = parse_feature_matrix(raw)
+
+    # Format
+    if args.format == "markdown":
+        report = format_markdown(args.test_type, parsed, args.exit_code, timestamp)
+    elif args.format == "html":
+        report = format_html(args.test_type, parsed, args.exit_code, timestamp)
+    else:
+        report = format_json(args.test_type, parsed, args.exit_code, timestamp)
+
+    # Output
+    if args.report_file:
+        Path(args.report_file).write_text(report)
+        print(f"Report written to: {args.report_file}")
+    else:
+        print(report)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.claude/skills/trtllm-test-specialist/scripts/parse_config.py b/.claude/skills/trtllm-test-specialist/scripts/parse_config.py
new file mode 100644
index 000000000000..4822bd108539
--- /dev/null
+++ b/.claude/skills/trtllm-test-specialist/scripts/parse_config.py
@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+"""Parse a trtllm-test-specialist YAML config file.
+
+Validates that the file is well-formed YAML, filters to known input keys, warns
+about unrecognized keys, fills in scope-specific default values for absent
+optional keys, and prints the resolved parameters as JSON to stdout.
+
+Usage:
+    parse_config.py --config-file <path>
+
+Output (JSON on stdout):
+    {
+      "params": { <key>: <value>, ... },
+      "missing_required": [ <key>, ... ],
+      "defaults_applied": [ <key>, ... ]
+    }
+
+    "params"           — recognised keys found in the config file plus any
+                         defaults that were filled in for absent optional keys.
+    "missing_required" — required keys still absent after defaults were applied,
+                         determined by test scope:
+                           * test_cmd present  → module-test required keys
+                           * test_type present → type-specific required keys
+                           * otherwise         → universally required keys
+    "defaults_applied" — keys whose values were filled in from _DEFAULTS_BY_TYPE
+                         (not present in the config file).
+
+Exit codes:
+    0  success — JSON printed to stdout
+    1  file not found, bad extension, or YAML parse error
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import yaml
+
+# All recognised input keys from the trtllm-test-specialist Inputs section.
+# Keys are grouped by section for readability; the set is what matters at runtime.
+KNOWN_KEYS = {
+    # Common
+    "model_name",
+    "checkpoint_path",
+    "repo_path",
+    "report_file",
+    # Module test
+    "test_cmd",
+    "class_name",
+    "required_devices",
+    # Model test — classification
+    "test_type",
+    # Model test — benchmark
+    "mtp_layers",
+    "benchmark_model_name",
+    # Model test — parallelism
+    "tp_size",
+    "ep_size",
+    "dp_size",
+    "pp_size",
+    # Model test — hardware
+    "device_type",
+    # Model test — dataset / paths
+    "dataset_path",
+    "benchmark_config_yaml",
+    "extra_llm_api_options_yaml",
+    # Model test — trtllm-bench
+    "bench_subcommand",
+    "backend",
+    "workspace",
+    "num_requests",
+    "input_mean",
+    "input_stdev",
+    "output_mean",
+    "output_stdev",
+    "kv_cache_free_gpu_mem_fraction",
+    "max_seq_len",
+    "beam_width",
+    "warmup",
+    "streaming",
+    "no_chunked_context",
+    "scheduler_policy",
+    "cluster_size",
+    "concurrency",
+    "max_batch_size",
+    "max_num_tokens",
+    "iteration_log",
+    # Model test — trtllm-eval (global options)
+    "eval_tasks",
+    "eval_tokenizer",
+    "eval_trust_remote_code",
+    "eval_disable_kv_cache_reuse",
+    "eval_kv_cache_free_gpu_memory_fraction",
+    # Model test — trtllm-eval (per-task options, common across most tasks)
+    "eval_num_samples",
+    "eval_apply_chat_template",
+    "eval_chat_template_kwargs",
+    "eval_system_prompt",
+    "eval_max_input_length",
+    "eval_max_output_length",
+    "eval_log_samples",
+    "eval_output_path",
+    # Model test — trtllm-eval (task-specific options)
+    "eval_check_accuracy",
+    "eval_accuracy_threshold",
+    "eval_num_fewshot",
+}
+
+# Required keys for the module-test scope (test_cmd present).
+_REQUIRED_MODULE = {"test_cmd"}
+
+# Required keys when no test_cmd and no test_type — most conservative baseline.
+_REQUIRED_ALWAYS = {"model_name", "checkpoint_path"}
+
+# Required keys per explicit test_type value.
+_REQUIRED_BY_TYPE: dict[str, set[str]] = {
+    "functionality": {"model_name", "checkpoint_path"},
+    "benchmark": {"model_name", "checkpoint_path", "tp_size", "pp_size"},
+    "evaluation": {"checkpoint_path", "eval_tasks"},
+}
+
+# Default values for absent optional (and some required) keys, keyed by test_type.
+# Applied before the missing_required check so defaults can satisfy required fields
+# (e.g. tp_size=1 satisfies the benchmark requirement without forcing explicit declaration).
+_DEFAULTS_BY_TYPE: dict[str, dict] = {
+    "benchmark": {
+        "tp_size": 1,
+        "pp_size": 1,
+        "ep_size": 0,
+        "bench_subcommand": "throughput",
+        "backend": "pytorch",
+        "num_requests": 100,
+        "input_mean": 128,
+        "input_stdev": 0,
+        "output_mean": 128,
+        "output_stdev": 0,
+        "beam_width": 1,
+        "warmup": 2,
+        "kv_cache_free_gpu_mem_fraction": 0.90,
+    },
+    "evaluation": {
+        "tp_size": 1,
+        "pp_size": 1,
+        "ep_size": 0,
+        "dp_size": 1,
+    },
+    "functionality": {
+        "tp_size": 1,
+        "pp_size": 1,
+        "ep_size": 0,
+        "dp_size": 1,
+    },
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument(
+        "--config-file", required=True, metavar="PATH", help="Path to the YAML config file"
+    )
+    return parser.parse_args()
+
+
+def load_yaml(path: Path) -> dict:
+    if not path.exists():
+        sys.exit(f"ERROR: config file not found: {path}")
+    if path.suffix.lower() not in {".yaml", ".yml"}:
+        sys.exit(f"ERROR: config file must have a .yaml or .yml extension, got: {path.suffix!r}")
+    try:
+        with path.open() as fh:
+            data = yaml.safe_load(fh)
+    except yaml.YAMLError as exc:
+        sys.exit(f"ERROR: failed to parse YAML: {exc}")
+    if data is None:
+        return {}
+    if not isinstance(data, dict):
+        sys.exit(f"ERROR: config file must be a YAML mapping (got {type(data).__name__})")
+    return data
+
+
+def apply_defaults(params: dict) -> tuple[dict, list[str]]:
+    """Fill in default values for absent keys based on test_type.
+
+    Only keys absent from *params* are touched; explicit values are never overridden.
+    Returns (updated_params, sorted list of keys that received a default value).
+    """
+    scope_defaults = _DEFAULTS_BY_TYPE.get(params.get("test_type"), {})
+    result = dict(params)
+    defaulted: list[str] = []
+    for key, value in scope_defaults.items():
+        if key not in result:
+            result[key] = value
+            defaulted.append(key)
+    return result, sorted(defaulted)
+
+
+def get_missing_required(params: dict) -> list[str]:
+    """Return sorted list of required keys absent from *params*."""
+    if "test_cmd" in params:
+        required = _REQUIRED_MODULE
+    elif "test_type" in params:
+        required = _REQUIRED_BY_TYPE.get(params["test_type"], _REQUIRED_ALWAYS)
+    else:
+        required = _REQUIRED_ALWAYS
+    return sorted(required - set(params.keys()))
+
+
+def main() -> None:
+    args = parse_args()
+    path = Path(args.config_file)
+    raw = load_yaml(path)
+
+    unknown = sorted(raw.keys() - KNOWN_KEYS)
+    if unknown:
+        print(f"WARNING: unrecognized keys ignored: {', '.join(unknown)}", file=sys.stderr)
+
+    params = {k: v for k, v in raw.items() if k in KNOWN_KEYS}
+    params, defaults_applied = apply_defaults(params)
+    missing_required = get_missing_required(params)
+
+    print(
+        json.dumps(
+            {
+                "params": params,
+                "missing_required": missing_required,
+                "defaults_applied": defaults_applied,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

From 8f16b4a707cb2e7bb1e564ff5e2b67fa19ecb6f0 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 21 Aug 2026 00:33:25 -0700
Subject: [PATCH 107/117] [None][fix] emission-tier op calls must bypass the
 tiered fast path

The tiered-GVR fast path dispatches on shape alone and ignores the
emission tensors, so a consume call carrying seed_thr/cand list/
xstate/block_max was silently routed to the stock tiered kernel:
results stay exact, but the candidate list is never read, xstate is
never published, and the closed loop never warms up. Skip the fast
path whenever any emission-tier tensor is present.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index dd95e0e327bd..a8ce5e102e67 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -7606,9 +7606,10 @@ def forward(
             # Host-only guard — no device sync. The op signature and output
             # contract are unchanged (unordered int32 indices, -1 pad only
             # for degenerate rows).
-            if _is_tiered_topk_supported(logits, pre_idx, seq_lens,
-                                         output_indices, top_k, next_n,
-                                         compress_ratio, order_row, counters):
+            if (seed_thr is None and cand_vals is None and xstate is None
+                    and block_max is None and _is_tiered_topk_supported(
+                        logits, pre_idx, seq_lens, output_indices, top_k,
+                        next_n, compress_ratio, order_row, counters)):
                 _tiered_topk(logits, pre_idx, seq_lens, output_indices, top_k,
                              next_n, compress_ratio)
                 return

From 8efc753659bcb65eeeab863cf01571b4cdb14f05 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 21 Aug 2026 00:22:21 -0700
Subject: [PATCH 108/117] [None][feat] house emission-assisted GVR inside the
 TopK module

Re-ports the GVR emission wiring onto the modular TopK structure
(#17622): the module owns GvrEmissionState (riding the caller's
gvr_prior_indices instead of a duplicate prior, own_prior=False),
prepare_gvr_emission() plans the tier and returns the scoring-op
emission kwargs, the consume-side ext kwargs are injected into the
CuTe DSL GVR call while the step is armed, and
reset_gvr_emission_rows() cold-starts reused slots at the
prefill-to-decode handoff. The indexer keeps only a thin gate
(TRTLLM_GVR_EMISSION=1, CUTE_DSL_GVR decode, next_n==1, no atom
split, <=256 rows) plus the kwargs splat into the fp8/fp4 scoring
calls.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../attention_backend/sparse/dsa/indexer.py   | 33 +++++++++
 .../attention_backend/sparse/gvr_emission.py  | 15 +++-
 tensorrt_llm/_torch/modules/top_k.py          | 72 +++++++++++++++++++
 3 files changed, 117 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index e32e806edb0e..6ec267e4a17e 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -653,6 +653,14 @@ def __init__(
             decode_implementation=decode_top_k_implementation,
             compress_ratio=self.compress_ratio,
         )
+        # GVR emission-assisted decode (opt-in, experimental): the FP4/FP8
+        # indexer epilogue emits candidates the GVR Top-K consumes (see
+        # gvr_emission / gvr_routing; state lives on the TopK module)
+        self.use_gvr_emission = (
+            os.environ.get("TRTLLM_GVR_EMISSION", "0") == "1"
+            and decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR
+            and self.use_cute_dsl_paged_mqa_logits
+        )
 
         # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM
         # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+)
@@ -1529,6 +1537,13 @@ def sparse_attn_indexer(
                 gvr_prior_indices,
                 request_offset=num_generations,
             )
+            if self.use_gvr_emission:
+                # reused slots cold-start the emission closed loop; stale
+                # lines only mis-place cuts - counts are re-measured
+                # in-kernel, so exactness never rides on this reset
+                self.top_k.reset_gvr_emission_rows(
+                    slice(num_generations, num_generations + num_contexts)
+                )
 
         reuse_topk = (
             self.mtp_index_share
@@ -1621,6 +1636,22 @@ def sparse_attn_indexer(
                     metadata.dsl_expand_factor > 1
                     and next_n == metadata.dsl_expand_factor * metadata.dsl_atom
                 )
+                gvr_emit_kwargs: dict = {}
+                # emitting for a step the Top-K cannot consume only churns
+                # the closed-loop state, so gate on the consumable shape
+                if (
+                    self.use_gvr_emission
+                    and gvr_prior_indices is not None
+                    and next_n == 1
+                    and not dsl_atom_split
+                    and num_gen_tokens <= 256
+                ):
+                    gvr_emit_kwargs = self.top_k.prepare_gvr_emission(
+                        num_generations,
+                        indexer_max_seq_len,
+                        torch.cuda.get_device_properties(q_decode.device).multi_processor_count,
+                        gvr_prior_indices,
+                    )
                 if self.use_fp4:
                     # FP4 DSL signature splits DG's (q, sf_q) tuple into two
                     # separate args and requires q.dtype == uint8 (q_decode
@@ -1653,6 +1684,7 @@ def sparse_attn_indexer(
                         dsl_block_table,
                         dsl_schedule_meta,
                         indexer_max_seq_len,
+                        **gvr_emit_kwargs,
                     )
                 else:
                     # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}.
@@ -1680,6 +1712,7 @@ def sparse_attn_indexer(
                         fp8_block_table,
                         fp8_schedule_meta,
                         indexer_max_seq_len,
+                        **gvr_emit_kwargs,
                     )
             else:
                 decode_q_scale = (
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py b/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
index 1aa77f3cc1a5..107d5af5274c 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
@@ -75,7 +75,12 @@ class GvrEmissionState:
     """Per-attention-backend emission state (persistent buffers)."""
 
     def __init__(
-        self, max_rows: int, top_k: int, device: torch.device, enable_list_tier: bool = True
+        self,
+        max_rows: int,
+        top_k: int,
+        device: torch.device,
+        enable_list_tier: bool = True,
+        own_prior: bool = True,
     ):
         self.max_rows = max_rows
         self.top_k = top_k
@@ -102,8 +107,12 @@ def __init__(
             self.cand_ctl = torch.zeros((cand_rows, 4), dtype=torch.int32, device=device)
             self.cand_cur = torch.zeros((cand_rows, 4), dtype=torch.int32, device=device)
         # previous-step top-k feedback (address-stable; zero-init ->
-        # first step's pre_idx points at index 0, a benign candidate)
-        self.prev_topk = torch.zeros((max_rows, top_k), dtype=torch.int32, device=device)
+        # first step's pre_idx points at index 0, a benign candidate).
+        # own_prior=False when the caller already keeps this state (the
+        # TopK module rides metadata's gvr_prior_indices).
+        self.prev_topk = (
+            torch.zeros((max_rows, top_k), dtype=torch.int32, device=device) if own_prior else None
+        )
         # block_max prefix ([rows, nb_pad*4] fp32 warp-partials),
         # allocated lazily once max_seq_len is known
         self.block_max: Optional[torch.Tensor] = None
diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py
index 704ce11283bc..d75ae0855ccb 100644
--- a/tensorrt_llm/_torch/modules/top_k.py
+++ b/tensorrt_llm/_torch/modules/top_k.py
@@ -55,6 +55,13 @@ def __init__(
             decode_implementation or TopKImplementation.CUDA_RADIX
         )
         self.compress_ratio = compress_ratio
+        # emission-assisted GVR (opt-in via prepare_gvr_emission): the
+        # module owns the closed-loop emission state; the caller passes
+        # the returned kwargs to the scoring op, and the consume side is
+        # injected into the GVR Top-K call while the step stays armed
+        self._gvr_emission_state = None
+        self._gvr_emission_route = None
+        self._gvr_emission_armed = False
 
     def forward(
         self,
@@ -280,6 +287,16 @@ def _forward_decode_gvr(
             )
         else:
             assert max_seq_len is not None
+            emission_kwargs: dict = {}
+            if self._gvr_emission_armed:
+                state = self._gvr_emission_state
+                num_rows = scores.shape[0]
+                emission_kwargs = state.topk_ext_kwargs(
+                    self._gvr_emission_route,
+                    num_rows,
+                    state.block_max[:num_rows] if state.block_max is not None else None,
+                )
+                self._gvr_emission_armed = False
             torch.ops.trtllm.cute_dsl_gvr_topk_decode(
                 scores,
                 gvr_prior_indices,
@@ -290,9 +307,64 @@ def _forward_decode_gvr(
                 compress_ratio=self.compress_ratio,
                 max_seq_len=max_seq_len,
                 order_row=gvr_row_order,
+                **emission_kwargs,
             )
         return output_indices
 
+    def prepare_gvr_emission(
+        self,
+        batch: int,
+        n_comp: int,
+        num_sms: int,
+        gvr_prior_indices: torch.Tensor,
+    ) -> dict:
+        """Plan the emission-assisted GVR tier for this decode step.
+
+        Returns the emission kwargs for the paged-MQA scoring op (empty
+        when the planner declines this step); the matching consume-side
+        kwargs are injected into the next GVR Top-K call automatically.
+        Host arithmetic on engine-static shapes plus capturable device
+        ops, so a captured graph bakes the tier and replays refresh the
+        state buffers in place.
+
+        Args:
+            batch: Number of decode requests this step.
+            n_comp: Engine-static compressed maximum sequence length.
+            num_sms: Device SM count.
+            gvr_prior_indices: Caller-owned previous-selection state;
+                defines the emission state's row capacity and device.
+        """
+        from ..attention_backend.sparse.gvr_emission import LIST_EMIT_MIN_N, GvrEmissionState
+
+        if self._gvr_emission_state is None:
+            self._gvr_emission_state = GvrEmissionState(
+                max_rows=gvr_prior_indices.shape[0],
+                top_k=self.top_k,
+                device=gvr_prior_indices.device,
+                enable_list_tier=n_comp >= LIST_EMIT_MIN_N,
+                own_prior=False,
+            )
+        state = self._gvr_emission_state
+        emit_tier, self._gvr_emission_route = state.plan(
+            batch, n_comp, num_sms, compress_ratio=max(self.compress_ratio, 1)
+        )
+        self._gvr_emission_armed = self._gvr_emission_route.tier != "none"
+        if emit_tier in ("counts", "list", "rungs"):
+            state.update_seed_rows(batch, emit_tier)
+        kwargs: dict = {}
+        if emit_tier in ("counts", "list"):
+            kwargs = state.indexer_emit_kwargs(emit_tier, batch)
+        if self._gvr_emission_route.attach_block_max or emit_tier in ("counts", "list"):
+            kwargs["block_max_out"] = state.ensure_block_max(n_comp)[:batch]
+        return kwargs
+
+    def reset_gvr_emission_rows(self, rows: slice) -> None:
+        """Cold-start the emission closed-loop state for reused request
+        slots (prefill-to-decode handoff): a zeroed xstate reads as
+        invalid and routes those rows to the stock path in-kernel."""
+        if self._gvr_emission_state is not None:
+            self._gvr_emission_state.xstate[rows].zero_()
+
     def update_gvr_prior_from_prefill(
         self,
         output_indices: torch.Tensor,

From 40917644d8f658e9d60924706c6a0beacd1ae8a8 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Sun, 23 Aug 2026 18:13:33 -0700
Subject: [PATCH 109/117] [None][test] cover emission-state parking for reused
 decode slots

Continuous-batching turnover can hand a request another request's
positional emission state. Exactness never rides on the seed lines -
the consumer admits on the candidate counts the emitter measures for
the current query, so a stale line only changes how many candidates
are collected - but a reset row must park on non-finite lines so it
restarts on the stock path instead of inheriting thresholds.

Verified on B200 alongside this: after handing every warm slot a
different request's score distribution with no reset, top-K stays
exact and the closed loop re-converges within one step (candidate
counts spike to ~92K, back to ~6K next step).

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tests/unittest/_torch/modules/test_top_k.py | 30 +++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py
index 9682d7981ab5..1ef992746f9e 100644
--- a/tests/unittest/_torch/modules/test_top_k.py
+++ b/tests/unittest/_torch/modules/test_top_k.py
@@ -328,3 +328,33 @@ def test_unsupported_prefill_implementation_raises() -> None:
             row_starts=torch.zeros(1, dtype=torch.int32),
             row_ends=torch.ones(1, dtype=torch.int32),
         )
+
+
+def test_gvr_emission_reset_parks_reused_slots(monkeypatch) -> None:
+    """Cold-started rows must carry non-finite lines.
+
+    Slot turnover under continuous batching can hand a request another
+    request's emission state. Exactness never rides on the lines - the
+    consumer admits on the counts the emitter measures for the current
+    query - but a reset row must park onto the stock path rather than
+    inherit finite thresholds, so the closed loop restarts cleanly.
+    """
+    gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_())
+    monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr)
+    top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR)
+    prior_indices = torch.zeros(4, 2, dtype=torch.int32)
+
+    top_k.prepare_gvr_emission(4, 1 << 17, 148, prior_indices)
+    state = top_k._gvr_emission_state
+    # emulate a warmed closed loop: every slot carries finite lines
+    state.xstate[:, 0] = 1.0
+    state.seed_row[:, :3] = torch.tensor([1.0, 2.0, 3.0])
+
+    top_k.reset_gvr_emission_rows(slice(1, 3))
+    top_k.prepare_gvr_emission(4, 1 << 17, 148, prior_indices)
+
+    lines = state.seed_row[:, 0]
+    assert torch.isinf(lines[1:3]).all(), "reset rows must park on non-finite lines"
+    assert torch.isfinite(lines[0]) and torch.isfinite(lines[3]), (
+        "untouched slots must keep their closed-loop state"
+    )

From a29d92eb03453fee0fb81fe63ca3f87996be43c3 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 25 Aug 2026 01:18:47 -0700
Subject: [PATCH 110/117] [None][chore] address perf-script and test-list
 review comments

Reviewer feedback on the committed measurement harnesses:

- Translate the remaining Chinese comments and print strings to English
  (f58_an.py, perstep.py).
- Drop every non-ASCII character from the folder: em dash, multiplication
  sign, element-of and arrow were used as typography in English prose.
- Add the SPDX header the four perf scripts were missing.
- Narrow the blanket `# ruff: noqa`. It is only warranted in ab_steps.py,
  where the explicit `del` of the per-grid buffers at the end of each loop
  body makes every closure reference read as possibly-unbound; that file
  now pins `F821` and says so. f58_an.py and perstep.py were hiding an
  unused import and `l` as a variable name behind the same escape - both
  are fixed instead, and both files now pass ruff clean.
- Stop ignoring two test files from the attention sweep only to add them
  straight back; only the file that needs its own timeout is split out.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../test_lists/test-db/l0_b200.yml            |  4 +-
 .../cute_dsl_kernels/top_k/perf/ab_steps.py   |  9 ++--
 .../cute_dsl_kernels/top_k/perf/f58_an.py     | 28 ++++++-----
 .../cute_dsl_kernels/top_k/perf/f58_grid.sh   |  2 +
 .../cute_dsl_kernels/top_k/perf/perstep.py    | 15 +++---
 .../cute_dsl_kernels/top_k/run_gvr_topk.py    | 48 +++++++++----------
 6 files changed, 55 insertions(+), 51 deletions(-)

diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index 49c5d627dcbd..fa7a9ed99495 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_b200.yml
@@ -76,9 +76,7 @@ l0_b200:
   - test_e2e.py::test_ptp_quickstart_advanced_ngram[Llama-3.1-8B-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct]
   - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-False-False]
   - test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b]
-  - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
-  - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py
-  - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
+  - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
   - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120)
   - unittest/_torch/compilation
   - unittest/_torch/debugger
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py b/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
index ba1c1ec5f9c7..8bdbc568de41 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
@@ -1,7 +1,10 @@
 #!/usr/bin/env python3
-# ruff: noqa
-# Measurement harness committed verbatim for provenance; bench idioms
-# (loop-scoped buffers, del/rebind) trip static analysis.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# ruff: noqa: F821
+# The per-grid buffers are released with an explicit `del` at the end of
+# each loop body, which makes the analyzer treat every closure reference
+# to them as possibly-unbound. All F821 reports here are that one idiom.
 # Layer x step complete grid. One NVTX range per (arm, model, isl, B,
 # layer); the COLD reps cycle through ALL usable decode steps (batch
 # refilled from the step's row BEFORE the eviction, so the row itself
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
index a4b21a0f683b..c9f9dfc50244 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
-# ruff: noqa
-# Measurement harness committed verbatim for provenance; bench idioms
-# (loop-scoped buffers, del/rebind) trip static analysis.
-# f58: B×N 四张表(flash/pro × 算数均值/最小值),逐步配对 vs PR16457,
-# 我们的臂按线上路由(plan_emission)选取。
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# f58: B x N tables (flash/pro, arithmetic mean and minimum), step-by-step
+# paired against PR16457; our arm follows the shipped routing
+# (plan_emission).
 import glob
 import os
 import pickle
@@ -31,7 +31,10 @@
             for k, v in per_step(f).items():
                 agg.setdefault(k, []).extend(v)
         data[a] = agg
-        print(f"  {a}: {len(agg)} 个(格,层)  {sum(len(v) for v in agg.values())} 步", flush=True)
+        print(
+            f"  {a}: {len(agg)} (cell, layer) pairs  {sum(len(v) for v in agg.values())} steps",
+            flush=True,
+        )
     pickle.dump(data, open(PKL, "wb"))
 
 sys.path.insert(
@@ -50,15 +53,15 @@
 BS = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
 
 cell = {}
-for m, isl, b in sorted({(m, i, b) for (m, i, b, l) in data["pr"]}):
+for m, isl, b in sorted({(m, i, b) for (m, i, b, lay) in data["pr"]}):
     n = int(isl[:-1]) * 1024 // CR[m]
     tier = plan_emission(b, n, KC[m], True)
     arm = TIER2ARM[tier]
     ratios = []
-    for (mm, ii, bb, l), pr in data["pr"].items():
+    for (mm, ii, bb, lay), pr in data["pr"].items():
         if (mm, ii, bb) != (m, isl, b):
             continue
-        ou = data[arm].get((mm, ii, bb, l), [])
+        ou = data[arm].get((mm, ii, bb, lay), [])
         for j in range(min(len(pr), len(ou))):
             ratios.append(pr[j] / ou[j])
     if ratios:
@@ -68,14 +71,15 @@
 
 
 def table(m, key):
-    print(f"\n### {m} — {'算数均值' if key == 'mean' else '最小值'}(对 PR16457,全层×全步)")
+    label = "arithmetic mean" if key == "mean" else "minimum"
+    print(f"\n### {m} - {label} (vs PR16457, all layers x all steps)")
     print("| B \\ N | " + " | ".join(NS) + " |")
     print("|---" * (len(NS) + 1) + "|")
     for b in BS:
         row = [f"**{b}**"]
         for isl in NS:
             c = cell.get((m, isl, b))
-            row.append(f"{c[key]:.3f}" if c else "—")
+            row.append(f"{c[key]:.3f}" if c else "-")
         print("| " + " | ".join(row) + " |")
 
 
@@ -85,7 +89,7 @@ def table(m, key):
 
 tot = sum(c["n"] for c in cell.values())
 print(
-    f"\n共 {len(cell)} 格 / {tot} 逐步配对;各格档位:",
+    f"\n{len(cell)} cells / {tot} per-step pairs; tier per cell:",
     {
         t: sum(1 for c in cell.values() if c["tier"] == t)
         for t in ("list", "counts", "rungs", "none")
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
index 26ca3b196945..6d5780e61638 100755
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
@@ -1,4 +1,6 @@
 #!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
 # Full-grid re-run on the current tree (post P4 fold work), fresh f58 prefix.
 BR=${GVR_BENCH_OUT:-./bench_results}
 cd "$(dirname "$0")"
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py b/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py
index 6f950dc07624..2f240b0d9001 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py
@@ -1,6 +1,6 @@
-# ruff: noqa
-# Measurement harness committed verbatim for provenance; bench idioms
-# (loop-scoped buffers, del/rebind) trip static analysis.
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
 # Per-(layer, step) pairing straight out of the nsys sqlite.
 #
 # One NVTX range = one (arm, model, isl, N, B, layer); inside it the
@@ -10,10 +10,7 @@
 # summary only carries Avg/Min/Max per range, which cannot be paired.
 import glob
 import os
-import re
 import sqlite3
-import sys
-from collections import defaultdict
 
 D = os.environ.get("GVR_BENCH_OUT", "./bench_results")
 
@@ -40,8 +37,8 @@ def per_step(path):
             while i < len(ker) and ker[i][0] < re_:
                 i += 1
             continue
-        _, arm, m, isl, _n, b, l = p
-        acc = out.setdefault((m, isl, int(b[1:]), l), [])
+        _, arm, m, isl, _n, b, lay = p
+        acc = out.setdefault((m, isl, int(b[1:]), lay), [])
         while i < len(ker) and ker[i][0] < re_:
             acc.append(ker[i][1] / 1000.0)
             i += 1
@@ -67,7 +64,7 @@ def collect(pattern):
     data = {a: collect(p) for a, p in arms.items()}
     for a, d in data.items():
         n = sum(len(v) for v in d.values())
-        print(f"  {a}: {len(d)} 个(格,层)  {n} 个步", flush=True)
+        print(f"  {a}: {len(d)} (cell, layer) pairs  {n} steps", flush=True)
     import pickle
 
     with open("/home/scratch.siyid_coreai/workspace/perstep.pkl", "wb") as f:
diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
index 041706d6fab5..73791781d666 100644
--- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
+++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
@@ -8,15 +8,15 @@
 ``use_constant_hint``, ``compress_ratio``, ``max_seq_len`` hint) so bench
 scripts can override the heuristic.
 
-**Not in CI** — this file imports the DSL kernel module directly
+**Not in CI** - this file imports the DSL kernel module directly
 (no ``trtllm`` runtime dep), to enable knob-A/B development outside the
 production op.
 
 Two usage modes:
 
-* `python -m pytest run_gvr_topk.py`` — exhaustive parameterized correctness sweep
-  (dtype × K × N × seed × next_n × T × V × warp-parallel-reduce).
-* ``python run_gvr_topk.py --dtype bf16 --top_k 1024 --N 8192`` —
+* `python -m pytest run_gvr_topk.py`` - exhaustive parameterized correctness sweep
+  (dtype x K x N x seed x next_n x T x V x warp-parallel-reduce).
+* ``python run_gvr_topk.py --dtype bf16 --top_k 1024 --N 8192`` -
   single-case correctness verification on user-specified shape; knob
   overrides via ``--num_threads`` / ``--use_256bit_load`` / etc.
 """
@@ -231,7 +231,7 @@ def _compile(
         self_scan=self_scan,
         cap_c=cap_c,
         # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed;
-        # the qfrac values are unused here — only the slot count matters.
+        # the qfrac values are unused here - only the slot count matters.
         r0_qfracs=(0.85, 0.35) if (use_ext_counts or ext_rungs) else None,
     )
     return cute.compile(
@@ -284,16 +284,16 @@ def emu_block_max(
     """``[num_rows, nb_pad*4] fp32`` warp-partial upper-bound records.
 
     tail_mode:
-      "exact":   tight bound — max over valid positions only.
+      "exact":   tight bound - max over valid positions only.
       "pad_inf": a partially-valid tail unit is forced to +FLT_MAX, the
                  worst legal inflation (the indexer masks by request-level
                  ctx >= N_eff, so tail positions can inflate the bound).
     records:
-      "rotate":     fold-correctness fixture — the 128-block max lands in
+      "rotate":     fold-correctness fixture - the 128-block max lands in
                     ONE slot rotated by blk % 4, the other 3 hold
                     -FLT_MAX. Valid ONLY for grain-128 consumers (slots
                     are NOT positional).
-      "positional": production semantics — record r is the exact max of
+      "positional": production semantics - record r is the exact max of
                     positions [r*32, r*32+32) (the indexer's TMEM T2R
                     partition gives warp w of a tile the contiguous
                     positions [tile*128 + w*32, +32)). Required for
@@ -582,7 +582,7 @@ def emu_cand(
 
     Unordered (value fp32-bits, index) pairs of all valid positions
     >= t_0 = seed_thr[:, 0]; ctl = {claimed, void}. claimed may
-    over-approximate the true count (window sentinels, idx word = -1) —
+    over-approximate the true count (window sentinels, idx word = -1) -
     ``sentinel_pad`` injects that legally. void=1 when claimed > cap; on
     overflow only the first ``cap`` entries are materialized (contract v2:
     consumers scan [0, min(claimed, cap)) skipping sentinels).
@@ -624,7 +624,7 @@ def emu_cand(
 def enc_ordered_f32(t: torch.Tensor) -> torch.Tensor:
     """Order-preserving int encoding of fp32 (an involution).
 
-    Stored back in fp32 slots — matches the indexer's encoded-int atomic
+    Stored back in fp32 slots - matches the indexer's encoded-int atomic
     min/max.
     """
     bits = t.float().contiguous().view(torch.int32)
@@ -635,7 +635,7 @@ def enc_ordered_f32(t: torch.Tensor) -> torch.Tensor:
 def hit_agg_identities(num_rows: int, device) -> torch.Tensor:
     """Identity-initialized per-row hit aggregate.
 
-    {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} — the required initial state of
+    {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} - the required initial state of
     the buffer the indexer atomically merges into.
     """
     ident = torch.tensor([_FLT_MAX, -_FLT_MAX], dtype=torch.float32, device=device)
@@ -699,9 +699,9 @@ def gvr_topk_decode(
     Args:
         logits:    ``[num_rows, max_S]`` float32 / bfloat16 / float16.
         pre_idx:   ``[num_rows // next_n, pre_idx_count]`` int32.
-                   ``pre_idx[..., 0]`` must be the argmax index — indexer invariant.
+                   ``pre_idx[..., 0]`` must be the argmax index - indexer invariant.
         seq_lens:  ``[num_rows // next_n]`` int32 (uncompressed-token space).
-        top_k:     K ∈ {512, 1024, 2048} — compile-time specialized.
+        top_k:     K in {512, 1024, 2048} - compile-time specialized.
         next_n:    Temporal stride for V3.2 ``preIdxOffset = (row % next_n) + 1``.
         compress_ratio: KV-indexer compression factor (1 = DSv3.2, 4 = DSv4).
                    When != 1, logits/preIdx live in compressed-token-index space:
@@ -725,7 +725,7 @@ def gvr_topk_decode(
                    round; every warp reduces the staged warp counts and
                    replays the classify + secant update in registers.
                    False restores the leader cadence.
-        order_row: Required iff ``seqlen_sorted=True``. Request-level —
+        order_row: Required iff ``seqlen_sorted=True``. Request-level -
                    ``int32[batch_size = num_rows // next_n]`` on the same
                    device as ``logits``; ``order_row[i]`` is the original
                    request_id of the i-th-priority request. The kernel
@@ -998,7 +998,7 @@ def gvr_topk_sort_prepare(seq_lens: torch.Tensor) -> torch.Tensor:
 
     Returns ``int32[batch_size]`` (= ``seq_lens.shape[0]`` =
     ``num_rows // next_n``) whose i-th entry is the original-batch index
-    of the i-th longest request — request-level, NOT row-level. The
+    of the i-th longest request - request-level, NOT row-level. The
     kernel expands to row level via ``order_row[req] * next_n + nn``
     inside the const_expr ``seqlen_sorted`` branch. Run once per decode
     step; the same
@@ -1075,7 +1075,7 @@ def gvr_topk_lb_prepare(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Run the LB prepare kernel.
 
-    ``seq_lens`` keeps its actual shape ``(batch_size,)`` — the kernel
+    ``seq_lens`` keeps its actual shape ``(batch_size,)`` - the kernel
     is compiled to match that exact shape; ``max_batch_size`` only
     determines the prepare kernel's block size and the ``order_row``
     buffer length. ``long_threshold`` is in SCAN-LENGTH space
@@ -1090,8 +1090,8 @@ def gvr_topk_lb_prepare(
     """
     assert seq_lens.is_cuda and seq_lens.dtype == torch.int32
     # block_prefix_sum_kernel (used inside LB prepare) constraints:
-    # num_warps = max_batch_size / 32 must be > 1 and a power of 2 →
-    # max_batch_size ∈ {64, 128, 256, 512, 1024}.
+    # num_warps = max_batch_size / 32 must be > 1 and a power of 2 ->
+    # max_batch_size in {64, 128, 256, 512, 1024}.
     if not (64 <= max_batch_size <= 1024) or (max_batch_size & (max_batch_size - 1)) != 0:
         raise ValueError(
             f"max_batch_size must be a power of 2 in [64, 1024] "
@@ -1263,9 +1263,9 @@ def _make_inputs(
     """Build (logits, pre_idx, seq_lens) for a multi-row test.
 
     Shapes:
-      logits  : [num_rows, N]                 — compressed-token-index space
-      pre_idx : [num_rows // next_n, top_k]   — argmax in slot 0 (indexer invariant)
-      seq_lens: [num_rows // next_n]          — UNCOMPRESSED-token space
+      logits  : [num_rows, N]                 - compressed-token-index space
+      pre_idx : [num_rows // next_n, top_k]   - argmax in slot 0 (indexer invariant)
+      seq_lens: [num_rows // next_n]          - UNCOMPRESSED-token space
 
     Kernel divides ``seq_lens`` by ``compress_ratio`` internally. Setting
     ``seq_lens = N * cr`` makes the kernel's
@@ -1282,7 +1282,7 @@ def _make_inputs(
     logits_f32 = torch.randn(num_rows, N, dtype=torch.float32, device="cuda") * 2.0
     logits = logits_f32.to(dtype)
     num_groups = num_rows // next_n
-    # argmax must come from the effective scan range, not full N — for
+    # argmax must come from the effective scan range, not full N - for
     # next_n>1 the kernel's row-0 N_eff is only (N - next_n + 1) cols.
     effective_len = N - next_n + 1
     argmax_idx = logits[::next_n, :effective_len].argmax(dim=-1).int()
@@ -1330,7 +1330,7 @@ def _tie_aware_correct(
         actual_kv_len = int(seq_lens_host[row // next_n]) - next_n + ofs + 1
         N_eff = actual_kv_len // compress_ratio
         if N_eff < top_k:
-            # Degenerate path — skip; caller's main() guards against this.
+            # Degenerate path - skip; caller's main() guards against this.
             continue
         row_logits = logits_f32[row, :N_eff]
         topk_vals, _ = torch.topk(row_logits, k=top_k, largest=True, sorted=True)
@@ -1388,7 +1388,7 @@ def test_gvr_topk_decode(
 ) -> None:
     # Kernel scans `N_eff = seq_lens[0] - next_n + (row_idx % next_n) + 1`
     # columns. Smallest row's N_eff = N - next_n + 1. Degenerate path
-    # (N_eff <= top_k) is a separate code branch — skip here.
+    # (N_eff <= top_k) is a separate code branch - skip here.
     if N - next_n + 1 < top_k:
         pytest.skip("N_eff < top_k is degenerate; the kernel requires N_eff >= top_k")
     seed = 42

From d9fc773cd181d25ceb3a06fd56758a810861e124 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Wed, 26 Aug 2026 20:39:15 -0700
Subject: [PATCH 111/117] [None][chore] drop the perf experiment dir, move
 gvr_emission next to the kernel

Two structural review points.

Remove tests/scripts/cute_dsl_kernels/top_k/perf/. It is an experiment
directory - three reviewers landed on the same folder - and the numbers it
backs are already recorded in the commit messages and the PR description, so
the scripts belong in a local tree rather than in-repo. This also reverts the
compliance work done on those four files in a29d92eb0 (translating comments,
stripping non-ASCII, adding SPDX, narrowing the ruff escape); deleting them is
the better answer to the same question. run_gvr_topk.py stays: it sits one
level up and is a standalone driver plus a parameterised pytest sweep, not an
experiment script, and its non-ASCII cleanup still stands.

Move gvr_emission.py from attention_backend/sparse/ into
cute_dsl_kernels/blackwell/top_k/. The constants in it are the kernel's buffer
geometry contract - the segment bases at 0 / LIST_SEG_A / 2*LIST_SEG_A have to
match the CUTLASS emitter exactly - and gvr_routing.py, which picks the tier
those buffers are emitted for, already lives there and was already imported
from across the package boundary. Folding it into modules/top_k.py instead
would move 272 lines of kernel geometry into a torch module, which is the
wrong direction. The move also removes the attention_backend -> cute_dsl_kernels
reverse dependency; the import is now local to the kernel package.

Import sites updated: modules/top_k.py and three in
test_cute_dsl_gvr_topk_decode.py.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k}/gvr_emission.py          |   8 +-
 tensorrt_llm/_torch/modules/top_k.py          |   5 +-
 .../cute_dsl_kernels/top_k/perf/README.md     |  19 -
 .../cute_dsl_kernels/top_k/perf/ab_steps.py   | 406 ------------------
 .../cute_dsl_kernels/top_k/perf/f58_an.py     |  97 -----
 .../cute_dsl_kernels/top_k/perf/f58_grid.sh   |  31 --
 .../cute_dsl_kernels/top_k/perf/perstep.py    |  72 ----
 .../sparse/test_cute_dsl_gvr_topk_decode.py   |   9 +-
 8 files changed, 11 insertions(+), 636 deletions(-)
 rename tensorrt_llm/_torch/{attention_backend/sparse => cute_dsl_kernels/blackwell/top_k}/gvr_emission.py (98%)
 delete mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/README.md
 delete mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
 delete mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
 delete mode 100755 tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
 delete mode 100644 tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py
similarity index 98%
rename from tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
rename to tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py
index 107d5af5274c..fbaec11736c1 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/gvr_emission.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py
@@ -32,13 +32,7 @@
 
 import torch
 
-from ...cute_dsl_kernels.blackwell.top_k.gvr_routing import (
-    LIST_EMIT_MAX_B,
-    LIST_EMIT_MIN_N,
-    TopkRoute,
-    pick_config,
-    plan_emission,
-)
+from .gvr_routing import LIST_EMIT_MAX_B, LIST_EMIT_MIN_N, TopkRoute, pick_config, plan_emission
 
 # Bucketed candidate-list geometry: two tight segments of LIST_SEG_A
 # entries plus a LIST_CAP_C-entry loose segment.
diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py
index d75ae0855ccb..57a267282f90 100644
--- a/tensorrt_llm/_torch/modules/top_k.py
+++ b/tensorrt_llm/_torch/modules/top_k.py
@@ -334,7 +334,10 @@ def prepare_gvr_emission(
             gvr_prior_indices: Caller-owned previous-selection state;
                 defines the emission state's row capacity and device.
         """
-        from ..attention_backend.sparse.gvr_emission import LIST_EMIT_MIN_N, GvrEmissionState
+        from ..cute_dsl_kernels.blackwell.top_k.gvr_emission import (
+            LIST_EMIT_MIN_N,
+            GvrEmissionState,
+        )
 
         if self._gvr_emission_state is None:
             self._gvr_emission_state = GvrEmissionState(
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/README.md b/tests/scripts/cute_dsl_kernels/top_k/perf/README.md
deleted file mode 100644
index 02b6fcc147bc..000000000000
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/README.md
+++ /dev/null
@@ -1,19 +0,0 @@
-# GVR top-k per-step perf grid (provenance for the PR numbers)
-
-Per-step paired protocol: each (arm, model, ISL, B, layer) is one NVTX
-range; the cold reps inside cycle every usable decode step (L2 evicted,
-batch refilled before eviction), so the k-th kernel instance of a range
-IS decode step k and arms pair index-by-index.
-
-- `ab_steps.py` - the driver (env: ARM = pr|st|va|vb|wf, UNITS =
-  "pro:64k,...", BS_LIST, OUT; GVR_CAP_ROOT points at the capture data,
-  HARNESS_ROOT at the repo).
-- `f58_grid.sh <gpu> "<units>"` - nsys wrapper per arm x unit
-  (GVR_BENCH_OUT selects the output dir).
-- `perstep.py` - per-step extraction from the nsys sqlite.
-- `f58_an.py` - pairing + production routing (plan_emission) + the B x N
-  mean/min tables. `f58_regr.py`-style regression stats derive from the
-  same pickle.
-
-Requirements: B200 (SM100), cutlass-dsl 4.5.x on PYTHONPATH, exclusive
-GPU (no concurrent process - it pollutes event timing), nsys >= 2025.3.
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py b/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
deleted file mode 100644
index 8bdbc568de41..000000000000
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/ab_steps.py
+++ /dev/null
@@ -1,406 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-# ruff: noqa: F821
-# The per-grid buffers are released with an explicit `del` at the end of
-# each loop body, which makes the analyzer treat every closure reference
-# to them as possibly-unbound. All F821 reports here are that one idiom.
-# Layer x step complete grid. One NVTX range per (arm, model, isl, B,
-# layer); the COLD reps cycle through ALL usable decode steps (batch
-# refilled from the step's row BEFORE the eviction, so the row itself
-# stays cold) -> the per-range cold mean IS the all-steps mean for that
-# layer. Tables aggregate mean-over-layers; per-layer means are the
-# per-range values (saved separately).
-#   env: ARM, UNITS ("flash:4k,pro:512k,..."), OUT, BS_LIST
-import json
-import os
-import sys
-from pathlib import Path
-
-import torch
-
-CAP = Path(
-    os.environ.get(
-        "GVR_CAP_ROOT",
-        "/home/scratch.loncheng_gpu/workspace/perf/workloads/DSV4/"
-        "E2E_exp/indexer_decode_capture/data",
-    )
-)
-LAYERS = {"flash": list(range(2, 43, 2)), "pro": list(range(2, 61, 2)), "v32": list(range(61))}
-KCFG = {"flash": (512, 4), "pro": (1024, 4), "v32": (2048, 1)}
-ALIGN = 64
-FMIN = torch.finfo(torch.float32).min
-
-ARM = os.environ["ARM"]
-OUT = os.environ["OUT"]
-_LSUB = os.environ.get("LAYER_SUBSET")
-if _LSUB:
-    _ls = [int(x) for x in _LSUB.split(",")]
-    LAYERS = {k: [x for x in v if x in _ls] for k, v in LAYERS.items()}
-UNITS = [u.split(":") for u in os.environ["UNITS"].split(",")]
-BS_LIST = [
-    int(x) for x in os.environ.get("BS_LIST", "1,2,4,8,16,32,64,128,256,512,1024").split(",")
-]
-_fb = os.environ.get("FORCE_BM")
-FORCE_BM = None if _fb is None else int(_fb)
-FORCE_CS = int(os.environ.get("FORCE_CS", "0"))
-CONLY = os.environ.get("CONLY", "0") == "1"  # codespell:ignore
-CONLY_RANK = int(os.environ.get("CONLY_RANK", "0"))
-_pw = os.environ.get("FORCE_PWR")
-FORCE_PWR = None if _pw is None else (_pw == "1")
-_tf = os.environ.get("FORCE_TAILFAST")
-FORCE_TAILFAST = None if not _tf else (_tf == "1")
-_nf = os.environ.get("FORCE_NOFINE")
-FORCE_NOFINE = None if not _nf else (_nf == "1")
-_nb = os.environ.get("FORCE_BINS")
-FORCE_BINS = int(_nb) if _nb else None
-_frt = os.environ.get("FORCE_FINERT")
-FORCE_FINERT = None if not _frt else (_frt == "1")
-_srt = os.environ.get("FORCE_SCATRT")
-FORCE_SCATRT = None if not _srt else (_srt == "1")
-_tm = os.environ.get("TIGHT_MULT")
-TIGHT_MULT = float(_tm) if _tm else None
-_nt = os.environ.get("FORCE_THREADS")
-FORCE_THREADS = int(_nt) if _nt else None
-_v3 = os.environ.get("FORCE_TAILV3")
-FORCE_TAILV3 = None if not _v3 else (_v3 == "1")
-_et = os.environ.get("FORCE_EXACTTAIL")
-FORCE_EXACTTAIL = None if not _et else (_et == "1")
-XSTATE = os.environ.get("XSTATE", "0") == "1"
-_ft = os.environ.get("FORCE_TIGHTEN")
-FORCE_TIGHTEN = None if _ft is None else bool(int(_ft))
-sms = torch.cuda.get_device_properties(0).multi_processor_count
-_EVICT = torch.empty(512 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda")
-
-if ARM == "pr":
-    sys.path.insert(0, "/home/scratch.siyid_coreai/workspace")
-    from prpkg16457.top_k.gvr_topk_decode import GvrTopKKernel  # noqa: E402
-elif ARM == "bx":
-    # loncheng PR #16877 head (BSX multi-tier): guard -> bsx dispatch,
-    # guard-fail -> that branch's own in-tree kernel (its production path)
-    sys.path.insert(0, "/home/scratch.siyid_coreai/wt-bsx/tensorrt_llm/_torch/cute_dsl_kernels")
-    from blackwell.top_k.gvr_topk_decode import GvrTopKKernel  # noqa: E402
-    from blackwell.top_k.gvr_topk_decode_bsx_dispatch import (  # noqa: E402
-        bsx_topk,
-        is_bsx_supported,
-    )
-else:
-    _root = os.environ.get("HARNESS_ROOT", str(Path(__file__).resolve().parents[5]))
-    sys.path.insert(0, _root + "/tests/scripts/cute_dsl_kernels/top_k")
-    sys.path.insert(0, _root)
-    import run_gvr_topk as rg  # noqa: E402
-    from run_gvr_topk import GvrTopKKernel  # noqa: E402
-
-    if ARM != "st":
-        from run_gvr_topk import (
-            emu_block_max,  # noqa: E402
-            emu_cand_bucketed,
-            emu_seed_counts,
-            pack_seed,
-        )
-
-
-def wf_targets(model, K, N):
-    if model == "flash":
-        T = (4096, 2048, 1024)
-    elif model == "v32":
-        T = (16384, 8192, 3072)
-    else:
-        T = (4096, 3584, 1536) if N <= 40960 else (12288, 5120, 2048)
-    return [min(t, N) for t in T]
-
-
-def main():
-    f = open(OUT, "w")
-    for model, isl in UNITS:
-        K, cr = KCFG[model]
-        Ls = LAYERS[model]
-        # ---- load per-layer step dicts (logits + captured topk) ----
-        lgs, pks = {}, {}
-        for L in Ls:
-            d = CAP / model / f"ISL_{isl}" / f"layer_{L:02d}"
-            lgs[L] = torch.load(d / "decode.logits.in.pt", map_location="cpu", weights_only=False)
-            pks[L] = torch.load(d / "decode.topk.out.pt", map_location="cpu", weights_only=False)
-        steps_all = sorted(pks[Ls[0]].keys())
-        # warmup steps carry -1 sentinel topk (and even wrong-width
-        # logits buffers): valid = captured topk present; a measured
-        # step also needs the PREVIOUS step valid (its topk is the
-        # preIdx seed)
-        valid = [s for s in steps_all if int(pks[Ls[0]][s].max()) >= 0]
-        vset = set(valid)
-        usable = [s for s in valid if (s - 1) in vset]
-        assert usable, f"no usable steps for {model}:{isl}"
-        NS = {}
-        for s in valid:
-            NS[s] = max(int(pks[L][s].max()) + 1 for L in Ls)
-        N_label = NS[valid[-1]]
-        # unit-constant pad width: block_max is validated against the
-        # BATCH buffer width (logits.shape[1]), so per-step bm must be
-        # built on rows already padded to the unit max
-        Npad_u = max((NS[s] + ALIGN - 1) // ALIGN * ALIGN for s in usable)
-        # ---- per (layer, step) prep on GPU ----
-        # row (padded), preIdx, lines, ref values, arm extras (1-row)
-        prep = {}
-        for L in Ls:
-            for s in usable:
-                Ns = NS[s]
-                Npad = Npad_u
-                row = torch.full((1, Npad), FMIN, dtype=torch.float32, device="cuda")
-                src = lgs[L][s][0]
-                row[0, :Ns] = src[:Ns].float().cuda()
-                pre = pks[L][s - 1].flatten().to(torch.int32).view(1, K)
-                pre = pre.cuda().contiguous()
-                srt = torch.sort(row[0, :Ns], descending=True).values
-                ref = srt[:K].clone()
-                ks = wf_targets(model, K, Ns)
-                t = torch.empty((1, 3), dtype=torch.float32, device="cuda")
-                for j, kc in enumerate(ks):
-                    t[0, j] = srt[kc - 1]
-                if TIGHT_MULT:
-                    # the closed loop derives its tight line from the
-                    # PREVIOUS step's k-th value, i.e. it sits just above
-                    # rank K - not at the fixed rank wf_targets uses. Model
-                    # that here so the tiers are measured on the contract
-                    # they actually ship with.
-                    t[0, 2] = srt[min(int(TIGHT_MULT * K), Ns) - 1]
-                for j in (1, 2):
-                    t[0, j] = torch.maximum(t[0, j], t[0, j - 1] + 1e-6)
-                if CONLY:  # codespell:ignore
-                    # push both tight lines above the row max so every
-                    # admitted entry lands in the loosest segment (the
-                    # one that already uses the cheap claim window).
-                    # CONLY_RANK tightens the collection line so the
-                    # claimed count stays under the consumer's line-cut
-                    # limit (and less volume is emitted).
-                    if CONLY_RANK:
-                        t[0, 0] = srt[min(CONLY_RANK * K, Ns) - 1]
-                    hi = float(row[0, :Ns].max()) + 1e4
-                    t[0, 1] = hi
-                    t[0, 2] = hi + 1.0
-                e = dict(Ns=Ns, Npad=Npad, row=row, pre=pre, ref=ref, sthr=t)
-                if ARM not in ("pr", "st", "bx"):
-                    sl1 = torch.tensor([Ns * cr], dtype=torch.int32, device="cuda")
-                    scnt = emu_seed_counts(row, sl1, t, compress_ratio=cr)
-                    if ARM == "wf":
-                        cv, ci, ctl = emu_cand_bucketed(
-                            row, sl1, t, 24576, seg_cap=8192, compress_ratio=cr, sentinel_pad=64
-                        )
-                        e.update(cv=cv, ci=ci, ctl=ctl)
-                    # thresholds mirror gvr_routing (real-capture A/B)
-                    nb = (ARM == "va" and Ns >= 65536) or (
-                        ARM == "vb" and model == "flash" and Ns >= 131072
-                    )
-                    if FORCE_BM is not None:
-                        nb = bool(FORCE_BM)
-                    if nb:
-                        e["bm"] = emu_block_max(row, sl1, compress_ratio=cr, tail_mode="exact")
-                    # col 6 = adaptive-skip pass count when bm attached
-                    e["spk"] = pack_seed(t, scnt, block_max=e.get("bm"))
-                prep[(L, s)] = e
-        Npad_max = max(e["Npad"] for e in prep.values())
-        for B in BS_LIST:
-            # reusable batch buffers (filled per rep BEFORE eviction)
-            lg_b = torch.full((B, Npad_max), FMIN, dtype=torch.float32, device="cuda")
-            pre_b = torch.zeros((B, K), dtype=torch.int32, device="cuda")
-            sl_b = torch.zeros((B,), dtype=torch.int32, device="cuda")
-            thr_b = torch.zeros((B, 3), dtype=torch.float32, device="cuda")
-            xs_b = torch.zeros((B, 8), dtype=torch.float32, device="cuda") if XSTATE else None
-            spk_b = torch.zeros((B, 8), dtype=torch.float32, device="cuda")
-            bufs = {}
-            if ARM in ("pr", "st", "bx"):
-                bufs["outb"] = torch.empty(B, K, dtype=torch.int32, device="cuda")
-            if ARM == "wf":
-                # widths from the emu contract (2*seg_cap+cap for the
-                # bucketed list), NOT the nominal cap
-                e0 = next(iter(prep.values()))
-                CW = e0["cv"].shape[1]
-                bufs["cv"] = torch.zeros((B, CW), dtype=torch.float32, device="cuda")
-                bufs["ci"] = torch.zeros((B, CW), dtype=torch.int32, device="cuda")
-                bufs["ctl"] = torch.zeros((B, e0["ctl"].shape[1]), dtype=torch.int32, device="cuda")
-            bm_b = None
-            if ARM in ("va", "vb"):
-                nbp = max((e["bm"].shape[1] for e in prep.values() if "bm" in e), default=0)
-                if nbp:
-                    bm_b = torch.zeros((B, nbp), dtype=torch.float32, device="cuda")
-            state = {}
-
-            def fill(L, s):
-                e = prep[(L, s)]
-                state.update(e=e)
-                lg_b[:, : e["Npad"]].copy_(e["row"].expand(B, e["Npad"]))
-                if e["Npad"] < Npad_max:
-                    lg_b[:, e["Npad"] :].fill_(FMIN)
-                pre_b.copy_(e["pre"].expand(B, K))
-                sl_b.fill_(e["Ns"] * cr)
-                thr_b.copy_(e["sthr"].expand(B, 3))
-                if ARM not in ("pr", "st", "bx"):
-                    spk_b.copy_(e["spk"].expand(B, 8))
-                if ARM == "wf":
-                    bufs["cv"].copy_(e["cv"].expand_as(bufs["cv"]))
-                    bufs["ci"].copy_(e["ci"].expand_as(bufs["ci"]))
-                    bufs["ctl"].copy_(e["ctl"].expand_as(bufs["ctl"]))
-                if bm_b is not None and "bm" in e:
-                    # unit-constant Npad -> bm width == bm_b width
-                    bm_b.copy_(e["bm"].expand_as(bm_b))
-                    state["bmc"] = bm_b
-                torch.cuda.synchronize()
-
-            def akw():
-                e = state["e"]
-                Ns = e["Ns"]
-                kw = dict(next_n=1, compress_ratio=cr, num_sms=sms, return_output_values=False)
-                if xs_b is not None:
-                    kw["xstate"] = xs_b
-                if FORCE_TAILFAST is not None:
-                    kw["p4_tail_fast"] = FORCE_TAILFAST
-                if FORCE_EXACTTAIL is not None:
-                    kw["p4_exact_tail"] = FORCE_EXACTTAIL
-                if FORCE_TAILV3 is not None:
-                    kw["p4_tail_v3"] = FORCE_TAILV3
-                if FORCE_NOFINE is not None:
-                    kw["p4_no_fine"] = FORCE_NOFINE
-                if FORCE_BINS is not None:
-                    kw["num_bins"] = FORCE_BINS
-                if FORCE_FINERT is not None:
-                    kw["p4_fine_rangetest"] = FORCE_FINERT
-                if FORCE_SCATRT is not None:
-                    kw["p4_scat_rangetest"] = FORCE_SCATRT
-                if FORCE_THREADS is not None:
-                    kw["num_threads_per_block"] = FORCE_THREADS
-                if ARM == "sw":
-                    return dict(kw)  # wrapper stock: no seed, no xstate
-                if ARM == "wf":
-                    a = dict(
-                        seed_thr=spk_b,
-                        cand_vals=bufs["cv"],
-                        cand_idx=bufs["ci"],
-                        cand_ctl=bufs["ctl"],
-                        cluster_size=1,
-                        **kw,
-                    )
-                    if FORCE_PWR is not None:
-                        a["p4_warp_redundant"] = FORCE_PWR
-                    return a
-                if ARM == "va":
-                    a = dict(seed_thr=spk_b, cluster_size=FORCE_CS or 1, **kw)
-                    if "bm" in e:
-                        a.update(block_max=state["bmc"], skip_min_n=None)
-                    if FORCE_PWR is not None:
-                        a["p4_warp_redundant"] = FORCE_PWR
-                    return a
-                # vb
-                cs = 1
-                if Ns >= 196608:
-                    if B * 8 <= sms // 2:
-                        cs = 8
-                    elif B * 4 <= (sms * 9) // 10:
-                        cs = 4
-                    elif B * 2 <= (sms * 9) // 10:
-                        cs = 2
-                a = dict(seed_thr=thr_b, cluster_size=cs, **kw)
-                if "bm" in e:
-                    a.update(block_max=state["bmc"], skip_min_n=None)
-                if FORCE_CS:
-                    a["cluster_size"] = FORCE_CS
-                if FORCE_PWR is not None:
-                    a["p4_warp_redundant"] = FORCE_PWR
-                return a
-
-            def call():
-                if ARM == "bx":
-                    if "bx_ok" not in state:
-                        state["bx_ok"] = is_bsx_supported(
-                            lg_b, pre_b, sl_b, bufs["outb"], K, 1, cr, None, None
-                        )
-                    if state["bx_ok"]:
-                        bsx_topk(lg_b, pre_b, sl_b, bufs["outb"], K, 1, cr)
-                    else:
-                        GvrTopKKernel.launch(lg_b, pre_b, sl_b, bufs["outb"], K, compress_ratio=cr)
-                    return
-                if ARM in ("pr", "st"):
-                    ov = {}
-                    if ARM == "st" and FORCE_PWR is not None:
-                        ov["p4_warp_redundant"] = FORCE_PWR
-                    GvrTopKKernel.launch(
-                        lg_b, pre_b, sl_b, bufs["outb"], K, compress_ratio=cr, **ov
-                    )
-                else:
-                    rg.gvr_topk_decode(lg_b, pre_b, sl_b, K, **akw())
-
-            def out_idx():
-                if ARM in ("pr", "st", "bx"):
-                    call()
-                    torch.cuda.synchronize()
-                    return bufs["outb"]
-                _, o = rg.gvr_topk_decode(lg_b, pre_b, sl_b, K, **akw())
-                torch.cuda.synchronize()
-                return o
-
-            for L in Ls:
-                base = f"{ARM}|{model}|{isl}|N{N_label}|B{B}|L{L:02d}"
-                rec = dict(
-                    model=model, isl=isl, N=N_label, K=K, B=B, arm=ARM, layer=L, steps=len(usable)
-                )
-                try:
-                    # exactness: every step at B==min; spot at larger B
-                    chk = usable if B == BS_LIST[0] else usable[:1]
-                    ok = True
-                    for s in chk:
-                        fill(L, s)
-                        o = out_idx()
-                        e = prep[(L, s)]
-                        for i in (0, B - 1) if B > 1 else (0,):
-                            gid = o[i, :K].long()
-                            # gather with an out-of-range index raises a
-                            # device-side assert that poisons the context and
-                            # kills the whole worker -> screen first
-                            if int(gid.min()) < 0 or int(gid.max()) >= e["Ns"]:
-                                ok = False
-                                rec["oob"] = rec.get("oob", 0) + 1
-                                continue
-                            got = lg_b[i, gid].sort(descending=True).values
-                            if gid.unique().numel() != K or not torch.equal(got, e["ref"]):
-                                ok = False
-                        if not ok:
-                            break
-                    rec["exact"] = ok
-                    # warmup
-                    for s in usable[:4]:
-                        fill(L, s)
-                        call()
-                    torch.cuda.synchronize()
-                    # COLD: one rep per decode step -> range mean = all-
-                    # steps mean for this layer
-                    xrec = []
-                    for s in usable:
-                        fill(L, s)
-                        if xs_b is not None:
-                            xs_b.zero_()
-                        _EVICT.uniform_(0, 1)
-                        torch.cuda.synchronize()
-                        torch.cuda.nvtx.range_push(f"c|{base}")
-                        call()
-                        torch.cuda.synchronize()
-                        torch.cuda.nvtx.range_pop()
-                        if xs_b is not None:
-                            xrec.append([round(v, 1) for v in xs_b[0].tolist()])
-                    if xs_b is not None:
-                        rec["x"] = xrec
-                    torch.cuda.synchronize()
-                except Exception as ex:  # noqa: BLE001
-                    rec["error"] = f"{type(ex).__name__}: {str(ex)[:120]}"
-                f.write(json.dumps(rec) + "\n")
-                f.flush()
-            del lg_b, pre_b, sl_b, thr_b, spk_b, bufs, bm_b, xs_b
-            torch.cuda.empty_cache()
-        prep.clear()
-        lgs.clear()
-        pks.clear()
-        torch.cuda.empty_cache()
-        print(f"[{ARM}] {model} {isl} done", flush=True)
-    f.close()
-    print(f"{ARM}_STEPS_DONE")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
deleted file mode 100644
index c9f9dfc50244..000000000000
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_an.py
+++ /dev/null
@@ -1,97 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-# f58: B x N tables (flash/pro, arithmetic mean and minimum), step-by-step
-# paired against PR16457; our arm follows the shipped routing
-# (plan_emission).
-import glob
-import os
-import pickle
-import sys
-
-sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent))
-from perstep import per_step  # noqa: E402
-
-D = os.environ.get("GVR_BENCH_OUT", "./bench_results")
-ARMS = {
-    "pr": "f58_*_pr.sqlite",
-    "st": "f58_*_st.sqlite",
-    "va": "f58_*_va.sqlite",
-    "vb": "f58_*_vb.sqlite",
-    "wf": "f58w_*.sqlite",
-}
-PKL = "./f58.pkl"
-if os.path.exists(PKL) and os.environ.get("REUSE", "1") == "1":
-    data = pickle.load(open(PKL, "rb"))
-else:
-    data = {}
-    for a, pat in ARMS.items():
-        agg = {}
-        for f in sorted(glob.glob(os.path.join(D, pat))):
-            for k, v in per_step(f).items():
-                agg.setdefault(k, []).extend(v)
-        data[a] = agg
-        print(
-            f"  {a}: {len(agg)} (cell, layer) pairs  {sum(len(v) for v in agg.values())} steps",
-            flush=True,
-        )
-    pickle.dump(data, open(PKL, "wb"))
-
-sys.path.insert(
-    0,
-    str(
-        __import__("pathlib").Path(__file__).parents[5]
-        / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k"
-    ),
-)
-from gvr_routing import plan_emission  # noqa: E402
-
-KC = {"flash": 512, "pro": 1024, "v32": 2048}
-CR = {"flash": 4, "pro": 4, "v32": 1}
-TIER2ARM = {"list": "wf", "counts": "va", "rungs": "vb", "none": "st"}
-NS = ["8k", "16k", "32k", "64k", "128k", "256k"]
-BS = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
-
-cell = {}
-for m, isl, b in sorted({(m, i, b) for (m, i, b, lay) in data["pr"]}):
-    n = int(isl[:-1]) * 1024 // CR[m]
-    tier = plan_emission(b, n, KC[m], True)
-    arm = TIER2ARM[tier]
-    ratios = []
-    for (mm, ii, bb, lay), pr in data["pr"].items():
-        if (mm, ii, bb) != (m, isl, b):
-            continue
-        ou = data[arm].get((mm, ii, bb, lay), [])
-        for j in range(min(len(pr), len(ou))):
-            ratios.append(pr[j] / ou[j])
-    if ratios:
-        cell[(m, isl, b)] = dict(
-            tier=tier, n=len(ratios), mean=sum(ratios) / len(ratios), mn=min(ratios)
-        )
-
-
-def table(m, key):
-    label = "arithmetic mean" if key == "mean" else "minimum"
-    print(f"\n### {m} - {label} (vs PR16457, all layers x all steps)")
-    print("| B \\ N | " + " | ".join(NS) + " |")
-    print("|---" * (len(NS) + 1) + "|")
-    for b in BS:
-        row = [f"**{b}**"]
-        for isl in NS:
-            c = cell.get((m, isl, b))
-            row.append(f"{c[key]:.3f}" if c else "-")
-        print("| " + " | ".join(row) + " |")
-
-
-for m in ("flash", "pro", "v32"):
-    for key in ("mean", "mn"):
-        table(m, key)
-
-tot = sum(c["n"] for c in cell.values())
-print(
-    f"\n{len(cell)} cells / {tot} per-step pairs; tier per cell:",
-    {
-        t: sum(1 for c in cell.values() if c["tier"] == t)
-        for t in ("list", "counts", "rungs", "none")
-    },
-)
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh b/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
deleted file mode 100755
index 6d5780e61638..000000000000
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/f58_grid.sh
+++ /dev/null
@@ -1,31 +0,0 @@
-#!/bin/bash
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-# Full-grid re-run on the current tree (post P4 fold work), fresh f58 prefix.
-BR=${GVR_BENCH_OUT:-./bench_results}
-cd "$(dirname "$0")"
-# PYTHONPATH must carry cutlass-dsl 4.5.x and the repo root; see README.md
-export TMPDIR=/tmp/nsys_g$1; mkdir -p $TMPDIR
-for UNIT in $2; do
-  TAG=$(echo $UNIT | tr ':' '_')
-  for ARM in ${ARMS:-pr st va vb}; do
-    rm -f $BR/f58_${TAG}_${ARM}*.csv $BR/f58_${TAG}_${ARM}.sqlite
-    CUDA_VISIBLE_DEVICES=$1 env ARM=$ARM UNITS=$UNIT \
-      GVR_BSTAR=8192 GVR_KC=8192 GVR_CAPC=16384 OUT=$BR/f58_${TAG}_${ARM}.jsonl \
-      nsys profile --trace=cuda,nvtx --force-overwrite true -o $BR/f58_${TAG}_${ARM} \
-        python3 ab_steps.py >> $BR/f58_${TAG}.log 2>&1
-    nsys stats -r nvtx_kern_sum --format csv -o $BR/f58_${TAG}_${ARM} \
-      $BR/f58_${TAG}_${ARM}.nsys-rep >> $BR/f58_${TAG}.log 2>&1
-    echo "f58 done ${TAG}_${ARM}" >> $BR/f58.log
-  done
-  # list tier (wf) under the C-only contract, same dataset
-  rm -f $BR/f58w_${TAG}*.csv $BR/f58w_${TAG}.sqlite
-  CUDA_VISIBLE_DEVICES=$1 env ARM=wf UNITS=$UNIT CONLY=1 CONLY_RANK=2 \  # codespell:ignore
-    GVR_BSTAR=8192 GVR_KC=8192 GVR_CAPC=16384 OUT=$BR/f58w_${TAG}.jsonl \
-    nsys profile --trace=cuda,nvtx --force-overwrite true -o $BR/f58w_${TAG} \
-      python3 ab_steps.py >> $BR/f58_${TAG}.log 2>&1
-  nsys stats -r nvtx_kern_sum --format csv -o $BR/f58w_${TAG} \
-    $BR/f58w_${TAG}.nsys-rep >> $BR/f58_${TAG}.log 2>&1
-  echo "f58 done ${TAG}_wf" >> $BR/f58.log
-done
-echo "f58 GPU$1 ALL DONE" >> $BR/f58.log
diff --git a/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py b/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py
deleted file mode 100644
index 2f240b0d9001..000000000000
--- a/tests/scripts/cute_dsl_kernels/top_k/perf/perstep.py
+++ /dev/null
@@ -1,72 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-# Per-(layer, step) pairing straight out of the nsys sqlite.
-#
-# One NVTX range = one (arm, model, isl, N, B, layer); inside it the
-# harness runs one kernel per decode step, in step order. So the k-th
-# kernel of a range IS step k, and the arms can be paired index by
-# index. This is the only way to get a true per-step ratio - the CSV
-# summary only carries Avg/Min/Max per range, which cannot be paired.
-import glob
-import os
-import sqlite3
-
-D = os.environ.get("GVR_BENCH_OUT", "./bench_results")
-
-
-def per_step(path):
-    """{(model, isl, B, layer): [step0_us, step1_us, ...]} for one file."""
-    out = {}
-    con = sqlite3.connect(path)
-    rng = con.execute(
-        "SELECT start, end, text FROM NVTX_EVENTS WHERE text LIKE 'c|%' ORDER BY start"
-    ).fetchall()
-    ker = con.execute(
-        "SELECT start, end - start FROM CUPTI_ACTIVITY_KIND_KERNEL ORDER BY start"
-    ).fetchall()
-    con.close()
-    # ranges do not overlap (one layer at a time), so a single merge walk
-    # assigns every kernel to the range containing it
-    i = 0
-    for rs, re_, text in rng:
-        while i < len(ker) and ker[i][0] < rs:
-            i += 1
-        p = text.lstrip(":").split("|")
-        if len(p) != 7:
-            while i < len(ker) and ker[i][0] < re_:
-                i += 1
-            continue
-        _, arm, m, isl, _n, b, lay = p
-        acc = out.setdefault((m, isl, int(b[1:]), lay), [])
-        while i < len(ker) and ker[i][0] < re_:
-            acc.append(ker[i][1] / 1000.0)
-            i += 1
-    return out
-
-
-def collect(pattern):
-    agg = {}
-    for f in sorted(glob.glob(os.path.join(D, pattern))):
-        for k, v in per_step(f).items():
-            agg.setdefault(k, []).extend(v) if k in agg else agg.setdefault(k, v)
-    return agg
-
-
-if __name__ == "__main__":
-    arms = {
-        "pr": "f22_*_pr.sqlite",
-        "st": "f22_*_st.sqlite",
-        "va": "f22_*_va.sqlite",
-        "vb": "f22_*_vb.sqlite",
-        "wf": "f30_*.sqlite",
-    }
-    data = {a: collect(p) for a, p in arms.items()}
-    for a, d in data.items():
-        n = sum(len(v) for v in d.values())
-        print(f"  {a}: {len(d)} (cell, layer) pairs  {n} steps", flush=True)
-    import pickle
-
-    with open("/home/scratch.siyid_coreai/workspace/perstep.pkl", "wb") as f:
-        pickle.dump(data, f)
-    print("saved perstep.pkl")
diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
index e54ea538e2f7..333fd24b9364 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
@@ -1419,7 +1419,7 @@ def test_cute_dsl_gvr_topk_decode_ext_list(mode, tie_aware_check):
               claim into the admission band -> the line-cut copy must
               re-measure and demote (exactness regression).
     """
-    from tensorrt_llm._torch.attention_backend.sparse.gvr_emission import (
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_emission import (
         LIST_CAP_C,
         LIST_PARK_LINE,
         LIST_SEG_A,
@@ -1546,7 +1546,10 @@ def _emulate_emission(logits, n_eff, st, tier, top_k):
     """Host-side stand-in for the indexer epilogue: fill the packed-row
     counts (and the candidate list on the list tier) against the CURRENT
     seed lines, exactly as the production emitter would."""
-    from tensorrt_llm._torch.attention_backend.sparse.gvr_emission import LIST_CAP_C, LIST_SEG_A
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_emission import (
+        LIST_CAP_C,
+        LIST_SEG_A,
+    )
 
     batch, N = logits.shape
     dev = logits.device
@@ -1589,7 +1592,7 @@ def test_cute_dsl_gvr_topk_decode_ext_closed_loop(tier_shape, tie_aware_check):
     width, so line placement must come from the fitted slope; every step
     must stay exact regardless of which internal path admission picks.
     """
-    from tensorrt_llm._torch.attention_backend.sparse.gvr_emission import GvrEmissionState
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_emission import GvrEmissionState
 
     want_tier, batch, N = tier_shape
     top_k = 512

From be40f85e691df455d88ef0ff66d9c29b61738d02 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Thu, 27 Aug 2026 00:27:40 -0700
Subject: [PATCH 112/117] [None][test] cut the new GVR top-K tests' JIT cost by
 29%

The cost of this file is dominated by CuTe DSL kernel compiles, not by
execution or by the case count: of the 60 cases this PR adds, 52 together run
in under a second, and the whole cost sits in eight cases that each instantiate
a fresh kernel config.

Two of those compiles were avoidable.

`ext_list` and `ext_closed_loop`'s list shape already agree on every launch
knob - top_k 512, N 131072, fp32, cluster_size 1, num_threads 512 - but
`emit_xstate` is part of the runner's compile key and only the latter passed an
`xstate` buffer, so the same kernel was built twice. `ext_list` now passes one
too; it costs nothing and additionally covers the xstate publish.

`ext_counts` swept top_k over {512, 1024, 2048}. The four modes it exercises
are count-based admission paths, which do not depend on K, while K is a
compile-time parameter - so the sweep bought two extra compiles and no extra
coverage. One K now.

Measured on a shared B200, new tests in isolation, cold:

  before  471s (7m51s), 60 cases
  after   333s (5m33s), 52 cases

`ext_closed_loop[list]` drops from 115.1s to 0.02s, which is the shared compile
landing. `degenerate_preidx`'s four configs are left alone: 16-bit has its own
path in this kernel and compress_ratio 1 and 4 are both production values, so
that 80s is coverage rather than waste.

For context on the file as a whole, same node, same protocol: 45m02s with these
tests, 38m21s with them deselected.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../attention/sparse/test_cute_dsl_gvr_topk_decode.py  | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
index 333fd24b9364..56ded7e2ea55 100644
--- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
@@ -1357,7 +1357,10 @@ def _pack_seed_row(logits_f32, n_eff, lines, block_max=None):
 
 
 @skip_not_sm100
-@pytest.mark.parametrize("top_k", [512, 1024, 2048])
+# One K only: the four modes exercise count-based admission, which is
+# K-independent, while K is a compile-time parameter - sweeping it cost
+# two extra kernel compiles (~43s) for no extra coverage.
+@pytest.mark.parametrize("top_k", [2048])
 @pytest.mark.parametrize("mode", ["band", "fat", "miss", "inf"])
 def test_cute_dsl_gvr_topk_decode_ext_counts(top_k, mode, tie_aware_check):
     """Packed seed row ([rows, 8]: lines + exact counts) consumption.
@@ -1482,6 +1485,10 @@ def test_cute_dsl_gvr_topk_decode_ext_list(mode, tie_aware_check):
         cand_ctl[r, 0] = int(hits.numel()) + pads
         cand_ctl[r, 1] = 1 if int(in_c.sum()) > LIST_CAP_C - pads else 0
     out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda")
+    # xstate is passed even though this test does not chain: it is part of
+    # the kernel's compile key, so sharing it with ext_closed_loop's list
+    # shape lets both reuse one compiled kernel instead of paying two.
+    xstate = torch.zeros((batch, 8), dtype=torch.float32, device=dev)
     torch.ops.trtllm.cute_dsl_gvr_topk_decode(
         logits,
         pre_idx,
@@ -1495,6 +1502,7 @@ def test_cute_dsl_gvr_topk_decode_ext_list(mode, tie_aware_check):
         cand_ctl=cand_ctl,
         accept_cap=LIST_SEG_A,
         num_threads=512,
+        xstate=xstate,
     )
     torch.cuda.synchronize()
     _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, 1)

From 2a923f6fcdece64150b77f27bbf3671bdac50c3d Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 28 Aug 2026 07:23:12 -0700
Subject: [PATCH 113/117] [None][fix] complete the #18094 port: phase 3's
 post-loop plateau terminal

The earlier port took #18094's two-sided bisection but stopped at the loop.
The 48 lines after it are load-bearing and were missing:

- On collapse with count < kK, fall back to s_thr[1] and re-count; val_lo
  admits >= kK by construction.
- Re-check adjacency: if count > kCC and the bracket is already adjacent,
  take s_thr[2] and stamp done = 3 - the plateau terminal - then re-count.

That second step is what relu_sparse_plateau was failing on. Without the
done = 3 stamp the Phase-4 plateau fill never fires, so a ReLU-sparse row
shipped its n_pos sure winners and padded the rest with -1 ("3 winners,
2045 pads"). It is also why the three earlier attempts missed: they all
edited the loop, the leader's terminal code, or the active-list flag, and
the missing piece sat directly below the loop.

Whole file on B200, serial, cold: 268 passed, 1 xpassed, 0 failed (48m14s).
The gate longcheng-nv named - hostile_hint, relu_sparse_plateau,
mtp_hostile_hint - passes, as do plateau_terminal and degenerate_preidx.

Method-level diff of base -> main vs base -> this branch shows #18094 touches
exactly two methods, _run_phases and phase3_collect_candidates, plus three new
free functions; both methods are also ones this branch changed. Only phase 3 is
ported here - _run_phases stays as it is, per longcheng-nv: the P1r rescue
covers what the synthetic bracket covers.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../blackwell/top_k/gvr_topk_decode.py        | 48 +++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index 60635a1c2d22..d8c1f05cce65 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -3106,6 +3106,54 @@ def phase3_collect_candidates(
                     cute.arch.barrier()
                 rs = rs + cutlass.Int32(1)
 
+            # Undershoot at collapse: val_lo admits >= kK by construction.
+            if s_iscalars[0] < cutlass.Int32(kK):
+                if tidx == 0:
+                    s_thr[0] = s_thr[1]
+                cute.arch.barrier()
+                self.block_count_ge(
+                    input_row,
+                    slice_start,
+                    slice_end,
+                    s_thr[0],
+                    smem_ptcnt,
+                    smem_wcnt,
+                    s_iscalars,
+                    s_cluster_partial,
+                    tidx,
+                    warp_id,
+                    lane,
+                    smem_input=smem_input,
+                    do_cluster_sync=do_cluster_sync,
+                )
+            cute.arch.barrier()
+
+            # Collapsed tie plateau (> kCC at an adjacent bracket): hand the
+            # row to the existing done=3 machinery; Phase 4 fills from the
+            # tie class in [val_lo, val_hi).
+            mid_chk, adj_chk = order_key_mid_f32(s_thr[1], s_thr[2])
+            if s_iscalars[0] > cutlass.Int32(kCC) and adj_chk:
+                if tidx == 0:
+                    s_thr[0] = s_thr[2]
+                    s_iscalars[1] = cutlass.Int32(3)  # plateau terminal
+                cute.arch.barrier()
+                self.block_count_ge(
+                    input_row,
+                    slice_start,
+                    slice_end,
+                    s_thr[0],
+                    smem_ptcnt,
+                    smem_wcnt,
+                    s_iscalars,
+                    s_cluster_partial,
+                    tidx,
+                    warp_id,
+                    lane,
+                    smem_input=smem_input,
+                    do_cluster_sync=do_cluster_sync,
+                )
+                cute.arch.barrier()
+
         # ---- Warp prefix sum over smem_ptcnt ----
         # my_total_qual = per-thread count cached by last block_count_ge.
         my_total_qual = smem_ptcnt[tidx]

From 875ce230d991e3de1c3eb250dcbc6ae86505f234 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Fri, 28 Aug 2026 11:54:01 -0700
Subject: [PATCH 114/117] [None][chore] ruff-format: drop the blank lines left
 by the phase-3 tail splice

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 .../_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py  | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
index d8c1f05cce65..672fb9209fd0 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
@@ -214,8 +214,6 @@ def order_key_mid_f32(v_lo, v_hi):
     return order_key_signed_to_f32(m_mid), m_mid == m_lo
 
 
-
-
 def _fmin_f32_inline(a, b):
     """Single PTX ``min.f32`` → one SASS FMNMX.
 

From 2bf2daecd828ee7e7391f8f9f899792de2e8b4d7 Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 1 Sep 2026 00:51:31 -0700
Subject: [PATCH 115/117] [None][fix] Gate GVR emission on the FP4 indexer
 scoring path

The FP8 scoring op's registered schema takes no emission kwargs, but the
emission gate did not require the FP4 path, so enabling TRTLLM_GVR_EMISSION
with the default FP8 indexer k dtype passed seed/candidate buffers into
cute_dsl_fp8_paged_mqa_logits and failed on unknown keywords before launch.
FP8 emission support arrives with its own kernel port in a follow-up PR.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index 3f2aa1109251..1868aa8458ba 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -677,10 +677,12 @@ def __init__(
         # GVR emission-assisted decode (opt-in, experimental): the FP4/FP8
         # indexer epilogue emits candidates the GVR Top-K consumes (see
         # gvr_emission / gvr_routing; state lives on the TopK module)
+        # only the FP4 scoring op accepts emission kwargs
         self.use_gvr_emission = (
             os.environ.get("TRTLLM_GVR_EMISSION", "0") == "1"
             and decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR
             and self.use_cute_dsl_paged_mqa_logits
+            and self.use_fp4
         )
 
         # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM

From 0d527a7882ab02dab632b893bdcc95f698a7d95c Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 1 Sep 2026 01:14:41 -0700
Subject: [PATCH 116/117] [None][fix] Do not arm GVR emission when Top-K row
 reordering is active

The emission ext tiers are single-CTA/sort-path only and the custom op
asserts order_row is None when any ext tier is present. MTP sets
kv_lens_row_reorder from the configured draft width while later draft
iterations run with next_n == 1, so the emission gate could arm a tier
whose consume call then carries both order_row and emission tensors,
tripping the assertion. Gate arming on the reorder row being absent.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index 1868aa8458ba..f11795269cfa 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -1668,6 +1668,9 @@ def sparse_attn_indexer(
                     and next_n == 1
                     and not dsl_atom_split
                     and num_gen_tokens <= 256
+                    # ext tiers are single-CTA/sort-path only; row reordering
+                    # routes the Top-K through order_row, which excludes them
+                    and metadata.kv_lens_row_reorder is None
                 ):
                     gvr_emit_kwargs = self.top_k.prepare_gvr_emission(
                         num_generations,

From eb81cc78b1477f741bb2eb9e914fe67eebf8706b Mon Sep 17 00:00:00 2001
From: siyidNV <297196620+siyidNV@users.noreply.github.com>
Date: Tue, 1 Sep 2026 01:14:52 -0700
Subject: [PATCH 117/117] [None][fix] Decline the GVR emission tier under
 torch.compile tracing

The kernel writes xstate and the emission buffers, but those mutations
cannot be declared in mutates_args (torch.library rejects a
declared-mutable Optional arg that is None at call time), so under
functionalization the writes are invisible and may be reordered or
dropped. The schema note documented the eager/CUDA-graph-only contract
without enforcing it; enforce it at the single arming point.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/top_k.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py
index d38c2592c6e6..97e6cb6cc694 100644
--- a/tensorrt_llm/_torch/modules/top_k.py
+++ b/tensorrt_llm/_torch/modules/top_k.py
@@ -388,6 +388,10 @@ def prepare_gvr_emission(
             gvr_prior_indices: Caller-owned previous-selection state;
                 defines the emission state's row capacity and device.
         """
+        # the emission/xstate writes are undeclared mutations (see the op's
+        # schema note), so the tier is eager / CUDA-graph only
+        if torch.compiler.is_dynamo_compiling():
+            return {}
         from ..cute_dsl_kernels.blackwell.top_k.gvr_emission import (
             LIST_EMIT_MIN_N,
             GvrEmissionState,