From eae7920076fac986e445f804117e4df034f17b36 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 8 Sep 2026 01:29:40 +0000 Subject: [PATCH 1/4] [None][feat] Self-sampling GVR V2 prefill indexer top-K MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rebased onto merged main (#18446): drops the now-duplicate two-level decode-dispatch and CUDA-heuristic-removal changes (they landed with #18446), keeping only the varlen prefill path — run_prefill / _prefill_launcher / warmup_prefill and the per-row [ks, ke) window support. Preserves main's merged #18683 envelope/alignment fix (n_kernel/n_route split, 16-byte workspace check, bands_done warmup guard) and the kv_cache_manager_v2 import path move. Made-with: Claude Code (Opus 4.8, 1M context) Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention/backends/sparse/dsa/indexer.py | 15 +- .../attention/backends/sparse/dsa/metadata.py | 13 + .../blackwell/top_k/__init__.py | 2 + .../top_k/gvr_topk_decode_self_sampling.py | 258 ++++++++++++--- .../gvr_topk_decode_self_sampling_host.py | 244 ++++++++++++++ tensorrt_llm/_torch/modules/top_k.py | 54 ++- .../attention/sparse/dsa/test_dsa_indexer.py | 54 +++ tests/unittest/_torch/modules/test_top_k.py | 131 +++++++- .../parallel/test_gvr_selfsampling_topk.py | 313 ++++++++++++++++++ 9 files changed, 1014 insertions(+), 70 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index d6655133dbf6..42c38ecc2903 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -754,9 +754,22 @@ def __init__( if self.use_cute_dsl_topk else TopKImplementation.CUDA_RADIX ) + # The self-sampling engine has a prefill form (per-row [ks, ke) + # windows); select it for prefill on exactly the layers where the + # two-level dispatch picks self-sampling for decode, so both phases + # share one config and one warmup. The temporal-hint engine has no + # prefill form, so those layers keep the exact radix prefill. + prefill_top_k_implementation = ( + TopKImplementation.CUTE_DSL_GVR + if ( + decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR + and self._use_self_sampling_topk + ) + else TopKImplementation.CUDA_RADIX + ) self.top_k = TopK( self.index_topk, - prefill_implementation=TopKImplementation.CUDA_RADIX, + prefill_implementation=prefill_top_k_implementation, decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, gvr_self_sampling=self._use_self_sampling_topk, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index f74c6cacbbc7..b01142e42eb1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -437,6 +437,19 @@ def warmup_selfsampling_topk( for bs in batch_sizes or (): rows.add(int(bs) * nn) msl_c = int(self.get_indexer_max_seq_len()) + # Prefill leg: the self-sampling engine also serves prefill (per-row + # [ks, ke) windows). It is placed BEFORE the DeepGEMM decode-stride + # guard below (which would return early for an odd msl_c) because the + # DeepGEMM prefill stride is always a 256-multiple. Bounded to the six + # tier x U engines per k; best-effort under the same OOM guard as the + # decode leg. + try: + _ss_host.warmup_prefill(int(top_k), max(msl_c, 32768)) + except torch.cuda.OutOfMemoryError: + logger.warning( + "self-sampling GVR prefill warmup ran out of memory; prefill " + "engines will JIT-compile lazily on first touch instead." + ) if self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits: # mirror the DSL paged-MQA arena stride (rows round up to 256 # elements). A drift here only degrades warmup to unused keys — diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index d5cab389489b..25d06bff8bd0 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -20,6 +20,7 @@ from .gvr_topk_decode_direct import DirectTopKKernel from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel +from .gvr_topk_decode_self_sampling_host import run_prefill as selfsampling_topk_run_prefill from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -36,4 +37,5 @@ "tiered_topk", "is_tiered_topk_supported", "selfsampling_topk_run_varlen", + "selfsampling_topk_run_prefill", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index a07ba529eb6d..6ba2c0e5814c 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -1328,6 +1328,7 @@ def __init__( cr_shift: int = 0, r_const: int = 1, hint_free: bool = False, + prefill: bool = False, ) -> None: assert nbs == 256, "SNB must stay 256" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) @@ -1350,6 +1351,21 @@ def __init__( self.r_const = int(r_const) # hint-free: gather_hint sites compiled out (sentinel pass-through) self.hint_free = bool(hint_free) + # prefill: per-row window [ks, ke) from row_starts/row_ends (rides the + # kv_lens / pre_idx ABI slots); base rounds down to a 16B boundary and + # the <=3 lead lanes are positionally masked. Single-CTA-per-row only + # (no SPLIT/workspace/TSH); next_n==1, cr_shift==0 (ks/ke are already + # in compressed column units). All prefill edits are const_expr-gated + # so legacy/varlen codegen stays byte-identical. + self.prefill = bool(prefill) + if self.prefill: + assert ( + self.varlen + and self.hint_free + and self.next_n == 1 + and self.cr_shift == 0 + and not self.split + ) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) and self.r_const >= 1 # TSH-floor staging arm. SPLIT-only compile-time key; the CUDA form @@ -1506,17 +1522,41 @@ def kern( short = cutlass.Int32(0) n_row = cutlass.Int32(0) tsh_run = cutlass.Int32(1) + # prefill window offset: lead = ks & 3 low lanes masked, col0 = ks + # rounded down to a float4 boundary (declared before the dynamic ifs + # per the scoping rule; stay 0 in every non-prefill compile). + lead = cutlass.Int32(0) + col0 = cutlass.Int32(0) if cutlass.const_expr(self.varlen): - req = row // cutlass.Int32(self.next_n) - rr = row % cutlass.Int32(self.next_n) - kvl = kv_lens[req] - nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( - self.cr_shift - ) - if nv < cutlass.Int32(0): - nv = cutlass.Int32(0) - if nv > npad: - nv = npad + if cutlass.const_expr(self.prefill): + # per-row window [ks, ke) already in compressed column units + # (kv_lens slot = row_starts, pre_idx slot = row_ends); no + # next_n / cr_shift math. Clamp only for memory safety — the + # indexer guarantees 0 <= ks <= ke <= logits.shape[1]. + ks = kv_lens[row] + ke = pre_idx[row] + if ks < cutlass.Int32(0): + ks = cutlass.Int32(0) + if ke > npad: + ke = npad + if ks > ke: + ks = ke + nv = ke - ks + lead = ks & cutlass.Int32(3) + col0 = ks - lead + if col0 > npad - cutlass.Int32(4): + col0 = npad - cutlass.Int32(4) + else: + req = row // cutlass.Int32(self.next_n) + rr = row % cutlass.Int32(self.next_n) + kvl = kv_lens[req] + nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( + self.cr_shift + ) + if nv < cutlass.Int32(0): + nv = cutlass.Int32(0) + if nv > npad: + nv = npad n_row = nv if nv <= k: short = cutlass.Int32(1) @@ -1530,7 +1570,12 @@ def kern( TGT2 = cutlass.Int32(0x3FFFFFFF) Q = cutlass.Int32(0) if short == cutlass.Int32(0): - n = nv + if cutlass.const_expr(self.prefill): + # scan extent from the rounded-down base col0 spans the + # lead pad plus the real window: [col0, ke) = nv + lead. + n = nv + lead + else: + n = nv n4v = n >> cutlass.Int32(2) # Ladder-scalar baselines only: the real SMP/SS2/TGT/TGT2 are # derived by warp0 alone in the block below (bit-identical @@ -1595,6 +1640,13 @@ def kern( s_lad = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 ) + if cutlass.const_expr(self.prefill): + # per-row lead (0..3) broadcast slot: warp0 publishes it under the + # existing s_lad barrier; every masked-lane site reloads it from + # smem so no live register is carried on the 64-register arms. + s_lead = smem.allocate_tensor( + cutlass.Int32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 + ) blob = smem.allocate_tensor( # dynamic-equivalent region cutlass.Int8, cute.make_ordered_layout((self.dyn_bytes,), order=(0,)), byte_alignment=16 ) @@ -1640,14 +1692,27 @@ def kern( row64 = cutlass.Int64(row) # _pin_i64: keep the row base a REGISTER across the attempt/tile scf # regions (NVVM otherwise re-derives ld.param+%ctaid.y+mul per region) - x_addr = _pin_i64(logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4)) + if cutlass.const_expr(self.prefill): + # base rounded down to the col0 float4 boundary (16B aligned since + # the row base is 16B aligned and col0 is a multiple of 4). + x_addr = _pin_i64( + logits.iterator.toint() + + (row64 * cutlass.Int64(npad) + cutlass.Int64(col0)) * cutlass.Int64(4) + ) + else: + x_addr = _pin_i64( + logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4) + ) # varlen: pre_idx is REQUEST-level [num_rows/next_n, k] — a request's # next_n rows share one hint row (production contract); legacy mode - # keeps the per-row mapping (next_n == 1 makes them identical). - prow64 = row64 - if cutlass.const_expr(self.varlen): - prow64 = cutlass.Int64(row // cutlass.Int32(self.next_n)) - p_addr = pre_idx.iterator.toint() + prow64 * cutlass.Int64(k) * cutlass.Int64(4) + # keeps the per-row mapping (next_n == 1 makes them identical). In + # prefill the pre_idx slot is 1-D row_ends (already consumed in the + # prologue), so this dead hint pointer is compiled out. + if cutlass.const_expr(not self.prefill): + prow64 = row64 + if cutlass.const_expr(self.varlen): + prow64 = cutlass.Int64(row // cutlass.Int32(self.next_n)) + p_addr = pre_idx.iterator.toint() + prow64 * cutlass.Int64(k) * cutlass.Int64(4) out_row = out[row, None] ws_addr = ws.iterator.toint() gdon_addr = ws_addr # slab views @@ -1789,11 +1854,22 @@ def kern( s_lad[1] = SS2 s_lad[2] = TGT s_lad[3] = TGT2 + if cutlass.const_expr(self.prefill): + s_lead[0] = lead # Register-free L2 hints for the first U-batch of this CTA's own # P3 slice (clamped in-row): the data P3 touches first starts # flowing while warp0 walks the chain. Short rows clamp every # hint to the row's last line — harmless. - plim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) + # prefill: the base is shifted to col0, so the clamp must stay in + # the row's own window [col0, ke) — an npad-based clamp would + # over-read col0 columns past the last row's allocation. n4-1 is + # the last full in-window float4 (>=0 even for the n=0 short pass). + if cutlass.const_expr(self.prefill): + plim4 = n4 - cutlass.Int32(1) + if plim4 < cutlass.Int32(0): + plim4 = cutlass.Int32(0) + else: + plim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) for uu in cutlass.range_constexpr(U): # NOTE: names must not collide with the PRIME-LATE block's # i_/ic — the DSL kills inner-scope names at region exit and @@ -1820,6 +1896,17 @@ def kern( p4 = tidx * SS2 * cutlass.Int32(2) C.ld_g_f32x4(atom128, x_addr, p4, fsa) C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) + if cutlass.const_expr(self.prefill): + # only thread 0's fsa (float4 index 0) can hold the <=3 masked + # lead lanes; substitute the always-valid lane 3 so the sample + # min/max fold and histogram stay finite and count-invariant + # (a materialized -inf would drive f2s_rz to INT_MIN and write + # out of bounds in the sample histogram at :1930). + if tidx == cutlass.Int32(0): + ld_ = s_lead[0] + for q in cutlass.range_constexpr(3): + if cutlass.Int32(q) < ld_: + fsa[q] = fsa[3] # ============ P2: quantile rung from the sample ====================== smn = cutlass.Float32(float("inf")) @@ -1853,7 +1940,13 @@ def kern( cute.arch.barrier() # ---- barrier (sample redux publish) ---- # PRIME-LATE prefetch block: strictly after the barrier. - lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) + # prefill clamps to the last in-window float4 (see plim4 note above). + if cutlass.const_expr(self.prefill): + lim4 = n4 - cutlass.Int32(1) + if lim4 < cutlass.Int32(0): + lim4 = cutlass.Int32(0) + else: + lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) pf = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(max(PFD, 1))] if cutlass.const_expr(self.pf): fullsl = cutlass.Int32(0) @@ -2142,6 +2235,12 @@ def kern( if okq != cutlass.Int32(0): # ok-gated (+inf-pad escape) for q in cutlass.range_constexpr(4): M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + if cutlass.const_expr(self.prefill): + # the <=lead lead lanes live only in bits 0..lead-1 of the + # i0==0 tile (float4 0, thread 0, part 0); clear them so the + # reservation, survivor walk and re-reads never see them. + if i0 == cutlass.Int32(0): + M = M & (cutlass.Int32(-1) << s_lead[0]) # prefetch roll-forward BEFORE reservation/walk if cutlass.const_expr(self.pf): hasnext = cutlass.Int32(0) @@ -2450,7 +2549,12 @@ def kern( if bq >= B: p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) if p < lim1: - out_row[p] = idv + # prefill: staged idx are in the col0 frame; the + # local output frame is relative to ks = col0+lead. + if cutlass.const_expr(self.prefill): + out_row[p] = idv - s_lead[0] + else: + out_row[p] = idv else: if whole == cutlass.Int32(0): q2 = p - above @@ -2473,22 +2577,32 @@ def kern( i_ = i0_ if i0_ >= hi2: i_ = tail0 + (i0_ - hi2) - x = C.ldg_f32(x_addr, i_) - if x >= TF: - bq = C.f2s_rz((x - TF) * SC) - if bq > cutlass.Int32(NBS - 1): - bq = cutlass.Int32(NBS - 1) - if bq >= B: - p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) - if p < lim1: - out_row[p] = i_ - else: - if whole == cutlass.Int32(0): - q2 = p - above - if q2 < cutlass.Int32(CMPB): - s_ck64[q2] = ( - cutlass.Uint64(C.fkey(x)) << cutlass.Uint64(32) - ) | cutlass.Uint64(cutlass.Uint32(i_)) + masked = cutlass.Int32(0) + if cutlass.const_expr(self.prefill): + # skip the <=lead lead lanes (col0-frame positions + # 0..lead-1 hold the previous request's finite logits) + if i_ < s_lead[0]: + masked = cutlass.Int32(1) + if masked == cutlass.Int32(0): + x = C.ldg_f32(x_addr, i_) + if x >= TF: + bq = C.f2s_rz((x - TF) * SC) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + if bq >= B: + p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + if p < lim1: + if cutlass.const_expr(self.prefill): + out_row[p] = i_ - s_lead[0] + else: + out_row[p] = i_ + else: + if whole == cutlass.Int32(0): + q2 = p - above + if q2 < cutlass.Int32(CMPB): + s_ck64[q2] = ( + cutlass.Uint64(C.fkey(x)) << cutlass.Uint64(32) + ) | cutlass.Uint64(cutlass.Uint32(i_)) i0_ = i0_ + cutlass.Int32(BLK) # ---- P6 refine ---- @@ -2517,11 +2631,14 @@ def kern( cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) ) if r_ < need: - out_row[above + r_] = cutlass.Int32( + idv6 = cutlass.Int32( cutlass.Uint32( cutlass.Uint64(u64v) & cutlass.Uint64(0xFFFFFFFF) ) ) + if cutlass.const_expr(self.prefill): + idv6 = idv6 - s_lead[0] + out_row[above + r_] = idv6 i = i + cutlass.Int32(BLK) else: # key-space narrowing over ck64 @@ -2623,6 +2740,11 @@ def kern( p1 = cutlass.Int32(1) if iu == ethr: p2 = cutlass.Int32(1) + # staged idx are col0-frame; shift to the ks-relative + # local output frame (p1=p2=0 for i>=mc, so the -lead + # on the idv=0 default is never emitted). + if cutlass.const_expr(self.prefill): + idv = idv - s_lead[0] self._ballot_pair_emit( p1, p2, @@ -2786,6 +2908,10 @@ def kern( if tie_m != cutlass.Int32(0): if iu == ethr: p2 = cutlass.Int32(1) + # staged idx (already >= lead via the P3 M-mask) -> local + # frame; the x_addr re-read above stays in the col0 frame. + if cutlass.const_expr(self.prefill): + idv = idv - s_lead[0] self._ballot_pair_emit( p1, p2, idv, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane ) @@ -2796,7 +2922,13 @@ def kern( rhi = cutlass.Uint32(0xFFFFFFFF) above2 = cutlass.Int32(0) need2 = k - m2 = n + # prefill: the genuine window is [lead, n); the <=lead lead + # lanes are excluded from the histogram, the emit and the + # candidate count so they never join a tie class. + lead_db = cutlass.Int32(0) + if cutlass.const_expr(self.prefill): + lead_db = s_lead[0] + m2 = n - lead_db ethr = cutlass.Int64(0) tie_m = cutlass.Int32(1) if tidx < cutlass.Int32(NBS): @@ -2826,8 +2958,8 @@ def kern( if sh2 < cutlass.Int32(0): sh2 = cutlass.Int32(0) sh2u = cutlass.Uint32(sh2) - i = tidx - while i < n: # whole row + i = tidx + lead_db # prefill: skip lead lanes + while i < n: # whole row (window [lead, n)) uq = C.fkey(C.ldg_f32(x_addr, i)) if uq >= cutlass.Uint32(rlo): if uq <= cutlass.Uint32(rhi): @@ -2876,15 +3008,16 @@ def kern( p1 = cutlass.Int32(0) p2 = cutlass.Int32(0) if i < n: - uq = C.fkey(C.ldg_f32(x_addr, i)) - iu = cutlass.Int64(uq) - if iu > ethr: - p1 = cutlass.Int32(1) - if tie_m != cutlass.Int32(0): - if iu == ethr: - p2 = cutlass.Int32(1) + if i >= lead_db: # prefill: exclude lead lanes + uq = C.fkey(C.ldg_f32(x_addr, i)) + iu = cutlass.Int64(uq) + if iu > ethr: + p1 = cutlass.Int32(1) + if tie_m != cutlass.Int32(0): + if iu == ethr: + p2 = cutlass.Int32(1) self._ballot_pair_emit( - p1, p2, i, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane + p1, p2, i - lead_db, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane ) it = it + cutlass.Int32(1) @@ -2966,16 +3099,26 @@ def __call__( _COMPILE_CACHE = {} -def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) -> Any: +def get_compiled( + tpl: tuple, options_extra: str = "", hint_free: bool = False, prefill: bool = False +) -> Any: """Compile (or fetch) the gvr_main variant for constexpr tuple tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG) — legacy, or tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG, NEXT_N, CR_SHIFT, R_CONST) — per-row varlen mode (TSHG slot is ignored: varlen compiles the TSH - machinery in whenever SPLIT and gates it per row at runtime).""" - key = (tuple(tpl), options_extra, bool(hint_free)) + machinery in whenever SPLIT and gates it per row at runtime). + + ``prefill`` selects the per-row window mode. It shares the varlen tuple + (next_n=1, cr_shift=0) but has a distinct prologue, so it MUST be part of + the cache key — otherwise a DSv3.2 decode varlen engine and the prefill + engine collide on the same tuple. The prefill compile also retypes the + pre_idx ABI slot to a 1-D align-4 fake (it carries 4B-aligned row_ends).""" + key = (tuple(tpl), options_extra, bool(hint_free), bool(prefill)) hit = _COMPILE_CACHE.get(key) if hit is not None: return hit + if prefill: + assert len(tpl) == 10, "prefill compile requires the varlen tuple" if len(tpl) == 7: blk, u, minb, nbs, kpt, split, tshg = tpl kern = GvrMainKernel( @@ -2996,6 +3139,7 @@ def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) - cr_shift=cr_shift, r_const=r_const, hint_free=bool(hint_free), + prefill=bool(prefill), ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -3005,9 +3149,15 @@ def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) - logits_fake = _crt.make_fake_compact_tensor( cutlass.Float32, (r0, c0), stride_order=(1, 0), assumed_align=16 ) - pre_fake = _crt.make_fake_compact_tensor( - cutlass.Int32, (r1, c1), stride_order=(1, 0), assumed_align=16 - ) + if prefill: + # pre_idx slot carries row_ends [rows] int32 (4B-aligned slices). + pre_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r1,), stride_order=(0,), assumed_align=4 + ) + else: + pre_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r1, c1), stride_order=(1, 0), assumed_align=16 + ) out_fake = _crt.make_fake_compact_tensor( cutlass.Int32, (r2, c2), stride_order=(1, 0), assumed_align=16 ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 769452c366f5..9ee5a5730f7b 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -678,6 +678,69 @@ def _main(blk_, minb_, u_, split_): _VARLEN_CACHE = {} +# ---- prefill launcher cache ------------------------------------------------ +# Prefill routes always force R==1 (single CTA per row): route_streaming gives +# R>1 only for b<=74, so the representative row counts below (first row of each +# route band) pin R=1 and reduce the engine set to <=6 per k. The launcher +# compiled function depends only on the row TIER, k and the envelope bucket +# (which selects U on the tier-0 1024-thread arm; tiers 1/2 fix U), never on +# the exact row count (arbitrary q-tile / q-split remainders) or npad (a +# runtime scalar), so the cache stays bounded over a long-running server. +_PREFILL_CACHE = {} +_PREFILL_ROW_SLAB = 32768 # gridDim.y <= 65535; slab so keys stay bounded +_PREFILL_TIER_ROWS = (75, 149, 297) # (rows<=148, 149..296, >296) band reps + + +def _prefill_tier(rows: int) -> int: + return 0 if rows <= 148 else 1 if rows <= 296 else 2 + + +def _prefill_bucket(n_env: int) -> int: + # pow2-quantize the envelope so a growing envelope reuses one plan; cap at + # 32768 because U=8 for every n>=32768 on the tier-0 arm. + return min(1 << max(int(n_env) - 1, 1).bit_length(), 32768) + + +def _prefill_cache_key(tier: int, k: int, n_bucket: int): + # tiers 1/2 fix U, so the bucket does not change their engine — collapse it + # to one key so warmup covers them with a single launch. + return (tier, k, n_bucket if tier == 0 else 0) + + +def _prefill_launcher(tier: int, k: int, n_bucket: int) -> tuple: + """Capture-time prefill plan + compiled launcher (main family, R=1). + + Mirrors ``_varlen_launcher``'s main branch but with r_const=1, split=False + (so tsh_en=0) and the prefill compile flag. SCAP_/CMP_/aim are envelope + upper bounds; npad is filled per call in ``run_prefill``.""" + key = _prefill_cache_key(tier, k, n_bucket) + hit = _PREFILL_CACHE.get(key) + if hit is not None: + return hit + b_route = _PREFILL_TIER_ROWS[tier] + n_route = max(n_bucket, k + 1) + plan = route_streaming(b_route, n_route, n_route, k, force_main=True) + if plan["kernel"] != "main": + raise RuntimeError(f"prefill route did not land on gvr_main: {plan['kernel']}") + rt = plan["rt"] + if rt["R"] != 1: + raise RuntimeError(f"prefill requires R==1 (got {rt['R']})") + tpl = tuple(plan["tpl"]) + dev = _device() + fn = dev.get_compiled(tpl[:6] + (False,) + (1, 0, 1), hint_free=True, prefill=True) + big = tier == 0 + # r_const==1 branch of the _varlen_launcher tuning scalars + aim_base = ( + (4 * k if k >= 1024 else 2 * k) if big else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2) + ) + sfac = 64 if k >= 1024 else 32 + amin = (7 * k) // 2 + sd_en = 1 if (k > 1024 and not big) else 0 + tail = (aim_base, sfac, amin, sd_en, 0) # tsh_en=0 (split=False) + lc = ("main", fn, (rt["SCAP_"], rt["CMP_"]), tail) + _PREFILL_CACHE[key] = lc + return lc + def _varlen_launcher( num_rows: int, @@ -1435,6 +1498,127 @@ def run_varlen( return +def run_prefill( + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + max_row_len: int | None = None, + workspace: torch.Tensor | None = None, +) -> None: + """Hint-free self-sampling Top-K for the prefill phase, per-row windows. + + Row semantics (mirror of ``topKPerRowPrefill`` / ``indexer_topk_prefill``): + row ``r`` selects the Top-K of ``logits[r, ks:ke]`` where + ``ks = row_starts[r]``, ``ke = row_ends[r]`` (both int32, in the SAME + compressed column units the DeepGEMM prefill producer emits — no + ``next_n`` / ``compress_ratio`` math). ``k`` comes from + ``indices.shape[1]``. The output is written in the LOCAL frame (column + minus ``ks``) with a trailing ``-1`` pad; rows with ``nv = ke - ks <= k`` + get the identity ``0..nv-1`` (matching the radix short-row contract). The + engine reads exactly ``[r*npad + (ks & ~3), r*npad + ke)`` — no dependence + on any producer slack. + + Envelope: ``max_row_len`` (a capture-stable engine constant) or, when + omitted, ``logits.shape[1]`` — a host int, so the call performs NO device + reads and is CUDA-graph-replay safe (it refuses to compile a new plan + under capture). Launches in ``<=65535``-row slabs so ``gridDim.y`` never + overflows. + + KNOWN LIMITATION: rows containing NaN inside the window are out of + contract (as for the radix reference — both order NaN implementation- + specifically). DeepGEMM prefill logits are finite in-window. Trusted + invariant: ``0 <= ks <= ke <= logits.shape[1]`` (the indexer guarantees + it); the kernel clamps ``ke <= npad`` for memory safety only. + """ + if logits.dtype is not _F32: + raise RuntimeError( + f"logits must be float32 (got {logits.dtype}); bf16/fp16 paths " + "are a follow-up — see the PR roadmap" + ) + for _nm, _t in (("row_starts", row_starts), ("row_ends", row_ends)): + if not (isinstance(_t, _TENSOR) and _t.is_cuda): + raise RuntimeError(f"{_nm} must be a CUDA tensor") + if _t.dtype is not _I32: + raise RuntimeError(f"{_nm} must be int32") + if _t.dim() != 1: + raise RuntimeError(f"{_nm} must be 1-D") + if not _t.is_contiguous(): + raise RuntimeError(f"{_nm} must be contiguous") + if len(logits.shape) != 2: + raise RuntimeError("logits must be 2-D") + num_rows = logits.shape[0] + if num_rows == 0: + return + if row_starts.shape[0] != num_rows or row_ends.shape[0] != num_rows: + raise RuntimeError( + f"row_starts/row_ends length must equal logits.shape[0]={num_rows}, " + f"got {row_starts.shape[0]}/{row_ends.shape[0]}" + ) + if not (logits.is_cuda and indices.is_cuda): + raise RuntimeError("all tensors must be CUDA") + if indices.dtype is not _I32: + raise RuntimeError("indices must be int32") + if len(indices.shape) != 2 or indices.shape[0] != num_rows: + raise RuntimeError(f"indices must be [num_rows={num_rows}, k], got {tuple(indices.shape)}") + if not indices.is_contiguous(): + raise RuntimeError("indices must be contiguous") + k = indices.shape[1] + if k < 4 or (k & 3): + raise RuntimeError(f"index_topk must be a multiple of 4 and >= 4, got {k}") + if indices.data_ptr() & 15: + raise RuntimeError("indices base must be 16-byte aligned") + if logits.stride(1) != 1: + raise RuntimeError("logits inner stride must be 1") + # DeepGEMM prefill rows are 1024B-aligned with >=256 float slack, so the + # row stride is valid for EVERY row count (the varlen 1-row shape[1] rule + # is a paged-MQA-arena quirk that would reject odd-width single-token + # prefill tiles — the common fully-cached follow-up turn). + npad = logits.stride(0) + if npad & 3: + raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") + if logits.data_ptr() & 15: + raise RuntimeError("logits base must be 16-byte aligned") + d = logits.get_device() + if not 0 <= d < _GVR_MAX_DEV: + raise RuntimeError(f"device index out of range: {d}") + lg = logits + if logits.shape[1] != npad: + need = logits.storage_offset() + num_rows * npad + if logits.untyped_storage().size() // 4 < need: + raise RuntimeError("logits view storage too small to widen to its row stride") + lg = logits.as_strided((num_rows, npad), (npad, 1), logits.storage_offset()) + if workspace is not None: + validate_run_ws(workspace, logits) + ws = kernel_view(workspace) + else: + ws = _ws_hot.get(d) + if ws is None: + ws = default_workspace(logits) + n_env = _index(max_row_len) if max_row_len is not None else logits.shape[1] + n_env = min(max(n_env, 1), npad) + n_bucket = _prefill_bucket(n_env) + for r0 in range(0, num_rows, _PREFILL_ROW_SLAB): + r1 = min(r0 + _PREFILL_ROW_SLAB, num_rows) + tier = _prefill_tier(r1 - r0) + lc = _PREFILL_CACHE.get(_prefill_cache_key(tier, k, n_bucket)) + if lc is None: + if _is_capturing(): + raise RuntimeError( + "prefill launcher not compiled for this shape — warm up " + "before CUDA graph capture" + ) + lc = _prefill_launcher(tier, k, n_bucket) + _, fn, (scap, cmp_), tail = lc + # ABI parity with the varlen main call: pre_idx slot = row_ends, + # kv_lens slot = row_starts. The n / SMP / TGT / Q / SS2 / TGT2 launch + # scalars are dead (re-derived per row); only npad / k / SCAP_ / CMP_ + # matter, R=1. + pre = (0, npad, k, scap, cmp_, 1, 0, 0, 0, 0, 0) + fn(lg[r0:r1], row_ends[r0:r1], indices[r0:r1], ws, *pre, row_starts[r0:r1], *tail) + return + + __all__ = [ "route", "route_static", @@ -1444,7 +1628,9 @@ def run_varlen( "run", "run_ws", "run_varlen", + "run_prefill", "warmup_varlen", + "warmup_prefill", "workspace_bytes", "WS_BYTES", "default_workspace", @@ -1580,3 +1766,61 @@ def warmup_varlen( if not bands_done: with _VARLEN_WARMUP_LOCK: _VARLEN_WARMUP_DONE.add(key) + + +_PREFILL_WARMUP_DONE: set = set() +_PREFILL_WARMUP_LOCK = threading.Lock() + + +def warmup_prefill( + top_k: int, + max_cols: int, + num_rows_list: Sequence[int] = (1, 149, 297), + row_stride: int | None = None, +) -> None: + """TESTING/INIT ONLY — compile the prefill engine set before serving. + + Six engines per k at most: the tier-0 (1024-thread) arm walks the pow2 + envelope buckets (U = 1/2/4/8), tiers 1/2 fix U so one launch each. One + tiny real launch per distinct ``(tier, k, bucket)`` cache key; ``ks=0``, + ``ke=n_env`` (all long rows). ``max_cols`` is the compressed max column + count (``get_indexer_max_seq_len``); the bucket caps at 32768 (U=8 above), + so envelopes past it share one key. The done-key gates only the GPU + launches — the ``_PREFILL_CACHE`` population is idempotent. + """ + dev = torch.cuda.current_device() + k = int(top_k) + max_cols = int(max_cols) + lo = _prefill_bucket(k + 1) + hi = _prefill_bucket(max_cols) + buckets = [] + b = lo + while b <= hi: + buckets.append(b) + b <<= 1 + if not buckets: + buckets = [hi] + keys = {} # cache_key -> (tier, bucket) representative for the launch + for rows in num_rows_list: + tier = _prefill_tier(int(rows)) + bset = buckets if tier == 0 else buckets[:1] + for bk in bset: + keys.setdefault(_prefill_cache_key(tier, k, bk), (tier, bk)) + done_key = (dev, k, max_cols, tuple(sorted(int(r) for r in num_rows_list)), row_stride) + with _PREFILL_WARMUP_LOCK: + if done_key in _PREFILL_WARMUP_DONE: + return + for tier, bk in keys.values(): + rows = _PREFILL_TIER_ROWS[tier] + stride = row_stride if row_stride is not None else ((bk + 256 + 255) // 256 * 256) + if stride < bk or stride % 4: + stride = (max(stride, bk) + 256 + 255) // 256 * 256 + logits = torch.zeros((rows, stride), dtype=torch.float32, device=dev) + ks = torch.zeros((rows,), dtype=torch.int32, device=dev) + ke = torch.full((rows,), bk, dtype=torch.int32, device=dev) + out = torch.empty((rows, k), dtype=torch.int32, device=dev) + run_prefill(logits[:, :bk], ks, ke, out, max_row_len=bk) + del logits, ks, ke, out + torch.cuda.synchronize() + with _PREFILL_WARMUP_LOCK: + _PREFILL_WARMUP_DONE.add(done_key) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 8093a96a016d..3eec56fa6c7c 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -143,7 +143,44 @@ def _forward_prefill( row_ends, output_indices, ) - if self.prefill_implementation == TopKImplementation.CUTE_DSL_RADIX: + if self.prefill_implementation == TopKImplementation.CUTE_DSL_GVR: + # hint-free k derives from the output width; pin it to the module's k + assert output_indices.shape[1] == self.top_k + if not self.gvr_self_sampling: + # the temporal (hint) GVR engine has no prefill form + logger.warning_once( + "temporal GVR has no prefill engine; using the CUDA radix prefill Top-K.", + key="gvr_temporal_prefill_radix", + ) + elif scores.shape[1] <= self.top_k: + # every row is short (nv <= k): the exact radix path emits the + # identity/-1 answer without reading logits — cheaper than a + # zero-work self-sampling launch. Deliberate, no warning. + pass + elif self._selfsampling_prefill_ok(scores): + from ..cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_prefill + + logger.info_once( + "self-sampling GVR prefill top-K engaged " + f"(K={self.top_k}, cr={self.compress_ratio}, hint-free).", + key="selfsampling_topk_prefill_engaged", + ) + # ks/ke are already in compressed column units; run_prefill + # writes the local (column - ks) frame with -1 pad and no host + # reads (envelope from scores.shape[1]). + selfsampling_topk_run_prefill(scores, row_starts, row_ends, output_indices) + return output_indices + else: + # engine hardware-format gate missed (e.g. a non-fp4 layer with + # an odd DeepGEMM width, or a bf16 producer): exact radix. + logger.warning_once( + "self-sampling GVR prefill is selected but the scores do " + "not satisfy the engine's hardware-format gate " + f"(dtype={scores.dtype}, strides={tuple(scores.stride())}); " + "falling back to the CUDA radix prefill Top-K.", + key="selfsampling_topk_prefill_fallthrough", + ) + elif self.prefill_implementation == TopKImplementation.CUTE_DSL_RADIX: # Keep the op's reread policy default; only its copy width is tuned. torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell( scores, @@ -154,7 +191,7 @@ def _forward_prefill( _CUTE_DSL_PREFILL_COPY_BITS, ) return output_indices - if self.prefill_implementation != TopKImplementation.CUDA_RADIX: + elif self.prefill_implementation != TopKImplementation.CUDA_RADIX: raise NotImplementedError( f"{self.prefill_implementation.value} does not support prefill Top-K" ) @@ -167,6 +204,19 @@ def _forward_prefill( ) return output_indices + def _selfsampling_prefill_ok(self, scores: torch.Tensor) -> bool: + """Engine hardware-format gate for the self-sampling prefill Top-K. + + fp32 row-major scores with a float4-aligned row stride and a 16B base + (the DeepGEMM prefill logits arena, whose rows are 1024B-aligned). The + all-short tile case is handled by the caller before this check.""" + return ( + scores.dtype == torch.float32 + and scores.stride(1) == 1 + and scores.stride(0) % 4 == 0 + and scores.data_ptr() % 16 == 0 + ) + def _forward_decode( self, scores: torch.Tensor, diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 6ac79f345954..0b9f1552ff5a 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -281,6 +281,54 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( cute_dsl_radix.assert_not_called() +@pytest.mark.parametrize( + "use_self_sampling,sm_version,msl_c,should_warmup", + [ + (True, 100, 65536, True), + (True, 100, 30001, True), # odd msl_c must not skip the prefill leg + (False, 100, 65536, False), # temporal-hint layers: no prefill engine + (True, 90, 65536, False), # non-datacenter Blackwell + ], +) +def test_metadata_warmup_selfsampling_prefill_leg( + use_self_sampling, sm_version, msl_c, should_warmup +): + """The self-sampling warmup drives BOTH the decode (varlen) and the prefill + engines; the prefill leg sits before the DeepGEMM decode-stride guard so an + odd msl_c cannot skip it.""" + metadata = SimpleNamespace( + enable_gvr_topk=True, + use_self_sampling_topk=use_self_sampling, + sparse_mla_topk=512, + _indexer_compress_ratio=4, + kv_cache_manager=SimpleNamespace(), + get_indexer_max_seq_len=Mock(return_value=msl_c), + sparse_metadata_params=SimpleNamespace(use_cute_dsl_paged_mqa_logits=True), + num_sms=148, + ) + ss_host = ( + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_self_sampling_host" + ) + with ( + patch( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_sm_version", + return_value=sm_version, + ), + patch(f"{ss_host}.warmup_prefill") as warmup_prefill, + patch(f"{ss_host}.warmup_varlen"), + ): + DSAtrtllmAttentionMetadata.warmup_selfsampling_topk(metadata, next_n=1, batch_sizes=[8]) + + if should_warmup: + warmup_prefill.assert_called_once_with(512, max(msl_c, 32768)) + else: + warmup_prefill.assert_not_called() + + def test_kv_lens_row_reorder_threshold(): """Prepare row order only when CuTe DSL GVR has enough decode rows.""" num_sms = 16 @@ -601,6 +649,12 @@ def test_indexer_two_level_gvr_dispatch( assert indexer.top_k.decode_implementation == expected_decode assert indexer.top_k.gvr_self_sampling == use_self_sampling assert indexer.top_k.needs_gvr_prior == (not use_self_sampling) + # Prefill uses the self-sampling engine on exactly the self-sampling + # layers; the temporal-hint layers keep the exact radix prefill. + expected_prefill = ( + TopKImplementation.CUTE_DSL_GVR if use_self_sampling else TopKImplementation.CUDA_RADIX + ) + assert indexer.top_k.prefill_implementation == expected_prefill @skip_pre_hopper diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index a9d319eac9f5..cf69bd18523e 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -425,19 +425,6 @@ def test_cute_dsl_prefill_dispatches_to_blackwell_kernel(monkeypatch) -> None: ) -def test_unsupported_prefill_implementation_raises() -> None: - top_k = TopK(1, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) - - with pytest.raises(NotImplementedError, match="does not support prefill Top-K"): - top_k( - torch.ones(1, 1), - torch.empty(1, 1, dtype=torch.int32), - is_prefill=True, - 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. @@ -466,3 +453,121 @@ def test_gvr_emission_reset_parks_reused_slots(monkeypatch) -> None: assert torch.isfinite(lines[0]) and torch.isfinite(lines[3]), ( "untouched slots must keep their closed-loop state" ) + + +def _install_fake_prefill_runner(monkeypatch) -> Mock: + """Stub both self-sampling entries (a test may exercise decode and prefill + through the same lazily imported module); return the prefill Mock.""" + prefill = Mock() + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k", + SimpleNamespace( + selfsampling_topk_run_varlen=Mock(), + selfsampling_topk_run_prefill=prefill, + ), + ) + return prefill + + +def _prefill_call(top_k: TopK, scores: torch.Tensor, out_width: int = 2): + rows = scores.shape[0] + row_starts = torch.zeros(rows, dtype=torch.int32) + row_ends = torch.full((rows,), scores.shape[1], dtype=torch.int32) + output = torch.full((rows, out_width), -1, dtype=torch.int32) + top_k( + scores, + output, + is_prefill=True, + row_starts=row_starts, + row_ends=row_ends, + ) + return row_starts, row_ends, output + + +def test_gvr_v2_prefill_routes_to_selfsampling_runner(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + top_k = TopK( + 2, + prefill_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) + scores = torch.randn(3, 8) # fp32, stride(0)=8 %4==0, contiguous -> gate ok + row_starts, row_ends, output = _prefill_call(top_k, scores) + + runner.assert_called_once() + args, kwargs = runner.call_args + assert args[0] is scores and args[1] is row_starts and args[2] is row_ends + assert args[3] is output + assert kwargs == {} + + +def test_gvr_v2_prefill_format_gate_falls_back_to_radix(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK( + 2, + prefill_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) + scores = torch.randn(3, 8, dtype=torch.bfloat16) # dtype gate miss + row_starts, row_ends, output = _prefill_call(top_k, scores) + + runner.assert_not_called() + radix.assert_called_once_with(scores, row_starts, row_ends, output, 2) + + +def test_gvr_v2_prefill_odd_stride_falls_back_to_radix(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK(2, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + scores = torch.randn(3, 6) # stride(0)=6, 6 % 4 != 0 -> gate miss + + _prefill_call(top_k, scores) + + runner.assert_not_called() + radix.assert_called_once() + + +def test_gvr_v2_prefill_all_short_uses_radix(monkeypatch) -> None: + """scores.shape[1] <= top_k: every row is short, so the exact radix + identity/-1 path runs (no logits read) with no fallthrough warning.""" + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK(4, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + scores = torch.randn(3, 4) # shape[1] == top_k + + _prefill_call(top_k, scores, out_width=4) + + runner.assert_not_called() + radix.assert_called_once() + + +def test_gvr_v2_prefill_temporal_mode_uses_radix(monkeypatch) -> None: + """CUTE_DSL_GVR + gvr_self_sampling=False has no prefill engine -> radix.""" + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK( + 2, + prefill_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ) + scores = torch.randn(3, 8) + + _prefill_call(top_k, scores) + + runner.assert_not_called() + radix.assert_called_once() + + +def test_gvr_v2_prefill_rejects_output_width_mismatch(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + top_k = TopK(2, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + scores = torch.randn(3, 8) + with pytest.raises(AssertionError): + _prefill_call(top_k, scores, out_width=3) + runner.assert_not_called() diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index be0db369a452..8b938866dd74 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -1008,3 +1008,316 @@ def test_validate_run_ws_requires_16_byte_alignment() -> None: ss_host.validate_run_ws(base[4:], logits) # +16 B with pytest.raises(RuntimeError, match="16-byte"): ss_host.validate_run_ws(base[2:], logits) # +8 B + + +# =========================================================================== +# ==== prefill: per-row [ks, ke) windows (run_prefill) ====================== +# =========================================================================== +# Contract: row r selects the Top-K of logits[r, ks:ke] (ks=row_starts[r], +# ke=row_ends[r], compressed column units), output in the LOCAL frame +# (column - ks) with a trailing -1 pad; nv=ke-ks <= k gives identity 0..nv-1. +# The base is rounded down to a 16B boundary and the <=3 lead lanes are masked, +# so ks % 4 in {1,2,3} (2nd+ request of a multi-request chunk) is exercised +# with poison (+inf/NaN/3e38/-inf) written at [ks-3, ks) and [ke, npad). + + +def _prefill_reference(logits, row_starts, row_ends, top_k): + rows, ncols = logits.shape + out = torch.full((rows, top_k), -1, dtype=torch.int32, device=logits.device) + ks, ke = row_starts.tolist(), row_ends.tolist() + for r in range(rows): + nv = max(min(ke[r], ncols) - ks[r], 0) + if nv == 0: + continue + if nv <= top_k: + out[r, :nv] = torch.arange(nv, dtype=torch.int32, device=logits.device) + else: + out[r] = torch.topk(logits[r, ks[r] : ks[r] + nv], top_k).indices.to(torch.int32) + return out + + +def _check_prefill_exact(logits, got, row_starts, row_ends, top_k): + """Tie-aware radix-parity check: trailing -1 pad from lengths; head unique + and in [0, nv); identity for nv <= k; exact index set when the k-th value + is unique, else strictly-above set + tie-class count (signed zeros and + genuine +/-inf compare like radix). NaN-in-window rows are structure-only.""" + assert got.shape == (logits.shape[0], top_k) and got.dtype == torch.int32 + ks, ke = row_starts.tolist(), row_ends.tolist() + got64 = got.to(torch.int64) + dev = logits.device + for r in range(logits.shape[0]): + nv = max(min(ke[r], logits.shape[1]) - ks[r], 0) + m = min(nv, top_k) + row = got64[r] + assert bool((row[m:] == -1).all()), f"row {r}: pad must be trailing -1 x{top_k - m}" + head = row[:m] + if m == 0: + continue + assert bool((head != -1).all()), f"row {r}: -1 inside the valid head" + assert int(head.min()) >= 0 and int(head.max()) < nv, f"row {r}: index outside [0,{nv})" + assert int(torch.unique(head).numel()) == m, f"row {r}: duplicate indices" + win = logits[r, ks[r] : ks[r] + nv] + if bool(torch.isnan(win).any()): + continue # NaN out of contract for both kernels; structure only + if nv <= top_k: + assert torch.equal(torch.sort(head).values, torch.arange(nv, device=dev)), ( + f"row {r}: short row must be identity" + ) + continue + vals = torch.sort(win, descending=True).values + v_k, v_next = vals[top_k - 1], vals[top_k] + got_vals = win[head] + if bool(v_k != v_next): + ref = (win >= v_k).nonzero(as_tuple=True)[0] + assert ref.numel() == top_k + assert torch.equal(torch.sort(head).values, ref), f"row {r}: index set mismatch" + else: + above = (win > v_k).nonzero(as_tuple=True)[0] + got_above = head[got_vals > v_k] + assert torch.equal(torch.sort(got_above).values, above), ( + f"row {r}: strictly-above set mismatch" + ) + assert int((got_vals == v_k).sum()) == top_k - above.numel(), ( + f"row {r}: wrong number of boundary-tied picks" + ) + assert bool((got_vals >= v_k).all()), f"row {r}: value below k-th selected" + + +def _make_prefill_case(rows, ncols, ks_list, ke_list, *, top_k, seed, dist="randn"): + """DeepGEMM-like storage: stride = align(ncols + 256, 256), column slice + [:, :ncols]; outside-window columns poisoned so an over-read/frame bug is + caught (+inf at [ks-3, ks), rotating NaN/inf/3e38/-inf elsewhere).""" + gen = torch.Generator(device=_DEV).manual_seed(seed) + stride = ((ncols + 256 + 255) // 256) * 256 + if dist == "randn": + full = torch.randn((rows, stride), generator=gen, dtype=torch.float32, device=_DEV) + elif dist == "equal": + full = torch.ones((rows, stride), dtype=torch.float32, device=_DEV) + elif dist == "twoval": + full = torch.randint(0, 2, (rows, stride), generator=gen, device=_DEV).float() + else: + raise ValueError(dist) + logits = full[:, :ncols] + row_starts = torch.tensor(ks_list, dtype=torch.int32, device=_DEV) + row_ends = torch.tensor(ke_list, dtype=torch.int32, device=_DEV) + cols = torch.arange(stride, device=_DEV).unsqueeze(0) + outside = (cols < row_starts.unsqueeze(1)) | (cols >= row_ends.unsqueeze(1)) + pat = torch.tensor([float("nan"), float("inf"), 3e38, float("-inf")], device=_DEV)[ + cols % 4 + ].expand(rows, -1) + full.masked_scatter_(outside, pat[outside]) + for r, ks in enumerate(ks_list): + full[r, max(ks - 3, 0) : ks] = float("inf") + return logits, row_starts, row_ends + + +@pytest.mark.parametrize("top_k", [512, 1024, 2048], ids=lambda k: f"k{k}") +def test_prefill_causal_ramp(top_k): + """Single-request causal ramp (ks=0): a run of short rows then long rows + in one launch straddles the k boundary. Covers nv < k, == k, > k.""" + rows = 148 + ks = [0] * rows + ke = list(range(1, rows + 1)) + lg, rs, re = _make_prefill_case(rows, rows, ks, ke, top_k=top_k, seed=top_k) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("lead", [1, 2, 3], ids=lambda x: f"lead{x}") +@pytest.mark.parametrize("top_k", [512, 2048], ids=lambda k: f"k{k}") +def test_prefill_packed_misaligned_ks(top_k, lead): + """Multi-request chunk: request 2 starts at ks % 4 == lead with +inf poison + at [ks-lead, ks). A leaked lead lane would become top-1 (wrong).""" + a = 300 + ks1 = ((a + 3) // 4) * 4 + lead + n1 = 4096 + 17 + rows = a + n1 + ncols = ks1 + n1 + ks = [0] * a + [ks1] * n1 + ke = list(range(1, a + 1)) + [ks1 + n1] * n1 + lg, rs, re = _make_prefill_case(rows, ncols, ks, ke, top_k=top_k, seed=top_k * 100 + lead) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + assert int(out.min()) >= -1, "negative index leaked (missed -lead correction / guard)" + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("top_k", [512, 1024], ids=lambda k: f"k{k}") +def test_prefill_short_rows(top_k): + """nv in {0, 1, k-1, k, k+1}: identity 0..nv-1 + trailing -1 (radix short + contract); nv==0 (ks==ke) -> all -1.""" + for nv in (0, 1, top_k - 1, top_k, top_k + 1): + rows = 4 + ncols = max(nv, 1) + 8 + ks = [0] * rows + ke = [nv] * rows + lg, rs, re = _make_prefill_case(rows, ncols, ks, ke, top_k=top_k, seed=nv + top_k) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("dist", ["equal", "twoval"], ids=lambda d: d) +@pytest.mark.parametrize("top_k", [512, 1024], ids=lambda k: f"k{k}") +def test_prefill_ties_degenerate(top_k, dist): + """All-equal (whole tie class) and two-valued (massive ties) rows drive the + degenerate A/B narrowing paths; tie-aware acceptance.""" + rows = 16 + n = 4096 + ks = [0] * rows + ke = [n] * rows + lg, rs, re = _make_prefill_case(rows, n, ks, ke, top_k=top_k, seed=top_k, dist=dist) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("lead", [1, 2, 3], ids=lambda x: f"lead{x}") +@pytest.mark.parametrize("top_k", [512, 1024], ids=lambda k: f"k{k}") +def test_prefill_neginf_tie_class(top_k, lead): + """nv > k with fewer than k finite values -> the k-th boundary is in the + -inf tie class (degen B), crossed with misaligned lead. A -inf-valued mask + would be emitted here as a negative index (== -lead); assert none leaks.""" + n_finite = top_k - 100 + nv = top_k + 400 + ks1 = ((37 + 3) // 4) * 4 + lead + ncols = ks1 + nv + rows = 5 + gen = torch.Generator(device=_DEV).manual_seed(top_k * 10 + lead) + stride = ((ncols + 256 + 255) // 256) * 256 + full = torch.full((rows, stride), float("-inf"), dtype=torch.float32, device=_DEV) + for r in range(rows): + full[r, ks1 : ks1 + n_finite] = torch.randn(n_finite, generator=gen, device=_DEV) + full[r, :ks1] = float("inf") + full[r, ks1 + nv :] = 3e38 + logits = full[:, :ncols] + rs = torch.tensor([ks1] * rows, dtype=torch.int32, device=_DEV) + re = torch.tensor([ks1 + nv] * rows, dtype=torch.int32, device=_DEV) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(logits, rs, re, out) + torch.cuda.synchronize() + assert int(out.min()) >= -1, "negative index leaked in a -inf tie class" + _check_prefill_exact(logits, out, rs, re, top_k) + + +@pytest.mark.parametrize("top_k, n", [(512, 4099), (2048, 131075)], ids=lambda v: f"n{v}") +def test_prefill_deepgemm_single_row_odd_width(top_k, n): + """A 1-row tile with an odd num_k_tokens on a DeepGEMM-strided view must + use stride(0) (not shape[1]) and stay exact — the fully-cached follow-up + turn that the varlen 1-row rule would wrongly reject.""" + rows = 1 + ks = [0] + ke = [n] + lg, rs, re = _make_prefill_case(rows, n, ks, ke, top_k=top_k, seed=n) + assert lg.stride(0) % 256 == 0 and lg.shape[1] == n # DeepGEMM-like view + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +def test_prefill_slab_over_gridy_limit(): + """> 65535 rows in one call must be slabbed (gridDim.y <= 65535).""" + k = 512 + rows = 70000 + n = 2048 + stride = ((n + 256 + 255) // 256) * 256 + lg = torch.randn((rows, stride), dtype=torch.float32, device=_DEV)[:, :n] + rs = torch.zeros((rows,), dtype=torch.int32, device=_DEV) + re = torch.full((rows,), n, dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + idx = torch.tensor([0, 1, 32767, 32768, 65535, 65536, 69999], device=_DEV) + _check_prefill_exact(lg[idx], out[idx].contiguous(), rs[idx], re[idx], k) + + +def test_prefill_engine_key_distinct_from_decode(): + """The prefill compile shares the DSv3.2 decode varlen tuple (next_n=1, + cr_shift=0) but has a distinct prologue, so the compile keys must differ.""" + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_self_sampling as dev, + ) + + tpl = (256, 8, 4, 256, 2, False, False, 1, 0, 1) + a = dev.get_compiled(tpl, hint_free=True) + b = dev.get_compiled(tpl, hint_free=True, prefill=True) + assert a is not b + + +def test_prefill_guards(): + k = 512 + n = 4096 + stride = ((n + 256 + 255) // 256) * 256 + lg = torch.randn((3, stride), dtype=torch.float32, device=_DEV)[:, :n] + rs = torch.zeros((3,), dtype=torch.int32, device=_DEV) + re = torch.full((3,), n, dtype=torch.int32, device=_DEV) + out = torch.full((3, k), -7, dtype=torch.int32, device=_DEV) + with pytest.raises(RuntimeError, match="float32"): + ss_host.run_prefill(lg.to(torch.bfloat16), rs, re, out) + with pytest.raises(RuntimeError, match="row_starts"): + ss_host.run_prefill(lg, rs.to(torch.int64), re, out) + with pytest.raises(RuntimeError, match="row_starts/row_ends length"): + ss_host.run_prefill(lg, rs[:2], re, out) + with pytest.raises(RuntimeError, match="multiple of 4"): + ss_host.run_prefill(lg, rs, re, torch.full((3, k + 2), -7, dtype=torch.int32, device=_DEV)) + with pytest.raises(RuntimeError, match="16-byte aligned"): + ss_host.run_prefill(lg[:, 1:], rs, re, out) # base offset by 1 float + + +def test_prefill_warmup_idempotent_and_no_rejit(): + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_self_sampling as dev, + ) + + k = 512 + ss_host.warmup_prefill(k, 32768) + before = len(ss_host._PREFILL_WARMUP_DONE) + ss_host.warmup_prefill(k, 32768) + assert len(ss_host._PREFILL_WARMUP_DONE) == before, "warmup not idempotent" + orig = dev.get_compiled + calls = {"n": 0} + + def counting(*a, **kw): + calls["n"] += 1 + return orig(*a, **kw) + + dev.get_compiled = counting + try: + for rows in (1, 8, 37, 74, 100, 296, 297, 4096): + for nkv in (4096, 16384, 32768): + stride = ((nkv + 256 + 255) // 256) * 256 + lg = torch.zeros((rows, stride), dtype=torch.float32, device=_DEV)[:, :nkv] + rs = torch.zeros((rows,), dtype=torch.int32, device=_DEV) + re = torch.full((rows,), nkv, dtype=torch.int32, device=_DEV) + out = torch.empty((rows, k), dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out, max_row_len=nkv) + finally: + dev.get_compiled = orig + assert calls["n"] == 0, f"warmup missed keys: {calls['n']} live compiles" + + +def test_prefill_capture_no_host_sync(): + """After warmup, a run_prefill call captures under a CUDA graph (proves no + .item()/.max() host read).""" + k = 512 + n = 8192 + ss_host.warmup_prefill(k, max(n, 32768)) + stride = ((n + 256 + 255) // 256) * 256 + lg = torch.randn((64, stride), dtype=torch.float32, device=_DEV)[:, :n] + rs = torch.zeros((64,), dtype=torch.int32, device=_DEV) + re = torch.full((64,), n, dtype=torch.int32, device=_DEV) + out = torch.full((64, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out, max_row_len=n) # compile outside capture + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + ss_host.run_prefill(lg, rs, re, out, max_row_len=n) + g.replay() + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, k) From c2bb2fe705336e217651acafec09e145c6e6e42b Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 8 Sep 2026 04:31:11 +0000 Subject: [PATCH 2/4] [None][test] Label the SM cases in the prefill-warmup test; tighten the causal-ramp docstring The `sm_version=90` fall-through case was commented as non-datacenter Blackwell; it is Hopper. Relabel it and add an SM120 (consumer Blackwell) case, both expecting no prefill-engine warmup since the self-sampling engine is SM100/103-only. The causal-ramp docstring now states what the 148-row launch actually exercises (all-short identity path) and points at the sibling tests that cover the k boundary and mixed short/long rows. Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/attention/sparse/dsa/test_dsa_indexer.py | 3 ++- .../_torch/thop/parallel/test_gvr_selfsampling_topk.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 0b9f1552ff5a..07e8cf17e134 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -287,7 +287,8 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( (True, 100, 65536, True), (True, 100, 30001, True), # odd msl_c must not skip the prefill leg (False, 100, 65536, False), # temporal-hint layers: no prefill engine - (True, 90, 65536, False), # non-datacenter Blackwell + (True, 90, 65536, False), # Hopper + (True, 120, 65536, False), # consumer Blackwell (SM120): engine is SM100/103-only ], ) def test_metadata_warmup_selfsampling_prefill_leg( diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 8b938866dd74..a6286bb28376 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -1113,8 +1113,10 @@ def _make_prefill_case(rows, ncols, ks_list, ke_list, *, top_k, seed, dist="rand @pytest.mark.parametrize("top_k", [512, 1024, 2048], ids=lambda k: f"k{k}") def test_prefill_causal_ramp(top_k): - """Single-request causal ramp (ks=0): a run of short rows then long rows - in one launch straddles the k boundary. Covers nv < k, == k, > k.""" + """Single-request causal ramp (ks=0), 148 rows with nv = 1..148 < k: one + tier-0 launch where every row takes the short-row identity path. The k + boundary is test_prefill_short_rows; mixed short/long rows in one launch is + test_prefill_packed_misaligned_ks.""" rows = 148 ks = [0] * rows ke = list(range(1, rows + 1)) From 48679aad12974b0c2c2c8071c97b3ed4dfc34e76 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 8 Sep 2026 05:47:51 +0000 Subject: [PATCH 3/4] [None][fix] Prefill GVR: warm up without a forward, radix under capture, pin the DeepGEMM stride The DSA top-K pre-compile hooks read `attn_metadata`, which only a warmup forward creates; a draft engine, a guided decoder or a context-only server without general warmup skips every forward, so the engines JIT-compiled on the first live request. Warmup now builds the DSA metadata itself in that case. Under CUDA graph capture an engine that warmup missed used to raise from `run_prefill`; `TopK` now queries `prefill_ready` and captures the exact radix path instead. The engine's format gate relies on DeepGEMM returning a column-sliced view with a 256-float-aligned row stride; two GPU tests (fp8 and fp8/fp4 producers, odd widths) pin that contract so a producer change surfaces as a test failure rather than a silent radix fallback. Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/__init__.py | 2 + .../gvr_topk_decode_self_sampling_host.py | 18 +++++++ tensorrt_llm/_torch/modules/top_k.py | 41 +++++++++++----- .../_torch/pyexecutor/model_engine.py | 31 ++++++++++++ .../sparse/dsa/test_dsa_fp4_indexer.py | 36 ++++++++++++++ .../attention/sparse/dsa/test_dsa_indexer.py | 25 ++++++++++ .../executor/test_pytorch_model_engine.py | 49 +++++++++++++++++++ tests/unittest/_torch/modules/test_top_k.py | 30 ++++++++++++ .../parallel/test_gvr_selfsampling_topk.py | 4 ++ 9 files changed, 225 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index 25d06bff8bd0..7203c41fcd80 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -20,6 +20,7 @@ from .gvr_topk_decode_direct import DirectTopKKernel from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel +from .gvr_topk_decode_self_sampling_host import prefill_ready as selfsampling_topk_prefill_ready from .gvr_topk_decode_self_sampling_host import run_prefill as selfsampling_topk_run_prefill from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen from .gvr_topk_decode_tp import GvrTpKernel @@ -38,4 +39,5 @@ "is_tiered_topk_supported", "selfsampling_topk_run_varlen", "selfsampling_topk_run_prefill", + "selfsampling_topk_prefill_ready", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 9ee5a5730f7b..bf639ab72ad9 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -1619,6 +1619,23 @@ def run_prefill( return +def prefill_ready(logits: torch.Tensor, indices: torch.Tensor) -> bool: + """True iff ``run_prefill(logits, ..., indices)`` would launch without + compiling — the same (tier, k, envelope bucket) keys it looks up, so a + caller can route around the engine under CUDA graph capture. Host-only.""" + num_rows = logits.shape[0] + if num_rows == 0: + return True + k = indices.shape[1] + npad = logits.stride(0) + n_bucket = _prefill_bucket(min(max(logits.shape[1], 1), max(npad, 1))) + for r0 in range(0, num_rows, _PREFILL_ROW_SLAB): + tier = _prefill_tier(min(r0 + _PREFILL_ROW_SLAB, num_rows) - r0) + if _prefill_cache_key(tier, k, n_bucket) not in _PREFILL_CACHE: + return False + return True + + __all__ = [ "route", "route_static", @@ -1629,6 +1646,7 @@ def run_prefill( "run_ws", "run_varlen", "run_prefill", + "prefill_ready", "warmup_varlen", "warmup_prefill", "workspace_bytes", diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 3eec56fa6c7c..60810d05a2dc 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -158,18 +158,33 @@ def _forward_prefill( # zero-work self-sampling launch. Deliberate, no warning. pass elif self._selfsampling_prefill_ok(scores): - from ..cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_prefill - - logger.info_once( - "self-sampling GVR prefill top-K engaged " - f"(K={self.top_k}, cr={self.compress_ratio}, hint-free).", - key="selfsampling_topk_prefill_engaged", + from ..cute_dsl_kernels.blackwell.top_k import ( + selfsampling_topk_prefill_ready, + selfsampling_topk_run_prefill, ) - # ks/ke are already in compressed column units; run_prefill - # writes the local (column - ks) frame with -1 pad and no host - # reads (envelope from scores.shape[1]). - selfsampling_topk_run_prefill(scores, row_starts, row_ends, output_indices) - return output_indices + + if self._prefill_capturing(scores) and not selfsampling_topk_prefill_ready( + scores, output_indices + ): + # the engine never JIT-compiles under capture; an engine + # missed by warmup takes the exact radix path in the graph + logger.warning_once( + "self-sampling GVR prefill engine is not compiled for this " + "shape and cannot JIT under CUDA graph capture; using the " + "CUDA radix prefill Top-K.", + key="selfsampling_topk_prefill_capture_radix", + ) + else: + logger.info_once( + "self-sampling GVR prefill top-K engaged " + f"(K={self.top_k}, cr={self.compress_ratio}, hint-free).", + key="selfsampling_topk_prefill_engaged", + ) + # ks/ke are already in compressed column units; run_prefill + # writes the local (column - ks) frame with -1 pad and no + # host reads (envelope from scores.shape[1]). + selfsampling_topk_run_prefill(scores, row_starts, row_ends, output_indices) + return output_indices else: # engine hardware-format gate missed (e.g. a non-fp4 layer with # an odd DeepGEMM width, or a bf16 producer): exact radix. @@ -217,6 +232,10 @@ def _selfsampling_prefill_ok(self, scores: torch.Tensor) -> bool: and scores.data_ptr() % 16 == 0 ) + @staticmethod + def _prefill_capturing(scores: torch.Tensor) -> bool: + return scores.is_cuda and torch.cuda.is_current_stream_capturing() + def _forward_decode( self, scores: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 82543b67b43c..b2c44a37a749 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1494,6 +1494,9 @@ def warmup(self, resource_manager: ResourceManager) -> None: # nvcc-driven JIT compile (~3s stall inside _prepare_inputs) on # first touch. Pre-touching every bucket funnels that cost into # warmup. No-op on non-DSA models. + # Both DSA hooks read attn_metadata, which only a warmup forward + # creates; build it when every forward above was skipped. + self._ensure_dsa_attn_metadata_for_warmup(resource_manager) self._warmup_dg_paged_mqa_logits_metadata() log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata") self._warmup_cute_dsl_radix_topk() @@ -1673,6 +1676,34 @@ def _prewarm_cute_dsl_indexer_q(self) -> None: self.dist.tp_allgather(1) logger.info("indexer-Q CuTe DSL prewarm complete") + def _ensure_dsa_attn_metadata_for_warmup( + self, resource_manager: ResourceManager) -> None: + """Build the DSA attention metadata if no warmup forward created it. + + A draft engine, a guided decoder, or a context-only server without + general warmup can skip every warmup forward; the DSA pre-compile + hooks would then find no metadata and the indexer top-K engines + would JIT on the first live request. No-op unless the backend is DSA. + """ + if getattr(self, "attn_metadata", None) is not None: + return + try: + from ..attention.backends.sparse.dsa import \ + DSAtrtllmAttentionMetadata + except ImportError: + return + metadata_cls = getattr(self.attn_backend, "Metadata", None) + if metadata_cls is None or not issubclass(metadata_cls, + DSAtrtllmAttentionMetadata): + return + kv_cache_manager = resource_manager.get_resource_manager( + self.kv_cache_manager_key) + if kv_cache_manager is None: + return + self._set_up_attn_metadata( + kv_cache_manager, + self._get_draft_kv_cache_manager(resource_manager)) + def _warmup_cute_dsl_radix_topk(self) -> None: """Pre-compile the DSA radix-filter CuTe DSL decode top-k for every cluster_size band during warmup, before serving. diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py index b4a8678132a0..60580764bc3b 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py @@ -38,6 +38,8 @@ from utils.util import skip_pre_blackwell # noqa: E402 +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation # noqa: E402 + FP4_MQA_NUM_HEADS = [ pytest.param( 32, @@ -78,6 +80,40 @@ def _dense_context_bounds(seq_len: int, seq_len_kv: int, device): return cu_ks, cu_ke.to(torch.int32) +@pytest.mark.skipif(not HAS_DEEP_GEMM, reason="fp8_fp4_mqa_logits not available") +@skip_pre_blackwell +@pytest.mark.parametrize("seq_len_kv", [1027, 4099]) +def test_fp4_mqa_logits_pass_selfsampling_prefill_format_gate(seq_len_kv): + """DeepSeek-V4 prefill (cr=4) hands these logits to the self-sampling GVR + prefill engine, whose format gate needs a float4-aligned row stride on + odd compressed widths; an exact-width producer would silently fall back + to radix.""" + torch.manual_seed(0) + num_heads, head_dim, seq_len = 64, 128, 64 + q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn(seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(seq_len, num_heads, device="cuda", dtype=torch.float32) + cu_ks = torch.zeros(seq_len, dtype=torch.int32, device="cuda") + cu_ke = torch.full((seq_len,), seq_len_kv, dtype=torch.int32, device="cuda") + q_fp4, q_scale_full = _fp4_quantize_sf_transpose(q) + k_fp4, k_scale_full = _fp4_quantize_sf_transpose(k) + + logits = deep_gemm.fp8_fp4_mqa_logits( + (q_fp4, q_scale_full.view(seq_len, num_heads)), + (k_fp4, k_scale_full.reshape(-1)), + weights, + cu_ks, + cu_ke, + False, # clean_logits + 0, # max_seqlen_k + torch.float32, # logits_dtype + ) + + assert logits.shape == (seq_len, seq_len_kv) + top_k = TopK(512, prefill_implementation=TopKImplementation.CUTE_DSL_GVR, compress_ratio=4) + assert top_k._selfsampling_prefill_ok(logits), (logits.dtype, tuple(logits.stride())) + + @pytest.mark.skipif(not HAS_DEEP_GEMM, reason="fp8_fp4_mqa_logits not available") @skip_pre_blackwell @pytest.mark.parametrize("num_heads", FP4_MQA_NUM_HEADS) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 07e8cf17e134..65baad615620 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -1534,6 +1534,31 @@ def test_deepgemm_fp8_mqa_logits_basic(compress_ratio): ) # double check for per-element similarity +@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") +@skip_pre_hopper +@pytest.mark.parametrize("seq_len_kv", [1027, 4099]) +def test_deepgemm_prefill_logits_pass_selfsampling_format_gate(seq_len_kv): + """The self-sampling GVR prefill engine engages only while DeepGEMM keeps + a float4-aligned row stride on odd-width logits; an exact-width producer + would silently route every such tile back to radix.""" + torch.manual_seed(0) + num_heads, head_dim, seq_len = 64, 128, 64 + q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(seq_len, num_heads, device="cuda", dtype=torch.float32) + ks = torch.zeros(seq_len, dtype=torch.int32, device="cuda") + ke = torch.full((seq_len,), seq_len_kv, dtype=torch.int32, device="cuda") + kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False) + + logits = deep_gemm.fp8_mqa_logits( + q.to(torch.float8_e4m3fn), kv_fp8, weights, ks, ke, clean_logits=False + ) + + assert logits.shape == (seq_len, seq_len_kv) + top_k = TopK(512, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + assert top_k._selfsampling_prefill_ok(logits), (logits.dtype, tuple(logits.stride())) + + def _create_mock_metadata( request_ids, batch_size, diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index ea673f9b61bd..ac660bce291c 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1521,6 +1521,55 @@ def test_encoder_cuda_graph_stages_and_restores_fixed_sequence_slots( (fixed_slot_output[511:512], fixed_slot_output[:400])) torch.testing.assert_close(restored_output, expected_output) + def test_warmup_builds_dsa_attn_metadata_when_no_forward_ran(self) -> None: + """The DSA top-K pre-compile hooks read attn_metadata; when every + warmup forward was skipped it is built on demand so the engines still + compile before serving.""" + from tensorrt_llm._torch.attention.backends.sparse.dsa import \ + DSAtrtllmAttentionMetadata + + engine = object.__new__(PyTorchModelEngine) + engine.attn_metadata = None + engine.attn_backend = SimpleNamespace( + Metadata=DSAtrtllmAttentionMetadata) + engine.kv_cache_manager_key = "kv" + engine.original_max_draft_len = 2 + engine._cuda_graph_batch_sizes = [8] + engine._get_draft_kv_cache_manager = Mock(return_value=None) + metadata = Mock(spec=DSAtrtllmAttentionMetadata) + + def build(kv_cache_manager, draft_kv_cache_manager): + engine.attn_metadata = metadata + return metadata + + engine._set_up_attn_metadata = Mock(side_effect=build) + kv_cache_manager = Mock() + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = kv_cache_manager + + engine._ensure_dsa_attn_metadata_for_warmup(resource_manager) + engine._warmup_cute_dsl_radix_topk() + + engine._set_up_attn_metadata.assert_called_once_with( + kv_cache_manager, None) + metadata.warmup_cute_dsl_radix_topk.assert_called_once_with(3) + metadata.warmup_selfsampling_topk.assert_called_once_with( + 3, batch_sizes=[8]) + + def test_warmup_does_not_build_metadata_for_non_dsa_backend(self) -> None: + from tensorrt_llm._torch.attention.backends.trtllm import \ + TrtllmAttentionMetadata + + engine = object.__new__(PyTorchModelEngine) + engine.attn_metadata = None + engine.attn_backend = SimpleNamespace(Metadata=TrtllmAttentionMetadata) + engine._set_up_attn_metadata = Mock() + + engine._ensure_dsa_attn_metadata_for_warmup(Mock()) + engine._warmup_cute_dsl_radix_topk() + + engine._set_up_attn_metadata.assert_not_called() + def test_breakable_rejects_multimodal_models(self) -> None: engine = object.__new__(PyTorchModelEngine) engine.model = DummyLegacyMultimodalIndexModel() diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index cf69bd18523e..1e4569d35af0 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -465,6 +465,7 @@ def _install_fake_prefill_runner(monkeypatch) -> Mock: SimpleNamespace( selfsampling_topk_run_varlen=Mock(), selfsampling_topk_run_prefill=prefill, + selfsampling_topk_prefill_ready=Mock(return_value=True), ), ) return prefill @@ -571,3 +572,32 @@ def test_gvr_v2_prefill_rejects_output_width_mismatch(monkeypatch) -> None: with pytest.raises(AssertionError): _prefill_call(top_k, scores, out_width=3) runner.assert_not_called() + + +def test_gvr_v2_prefill_capture_uncompiled_uses_radix(monkeypatch) -> None: + """Under CUDA graph capture an engine missed by warmup must not JIT; the + exact radix path is captured instead.""" + runner = _install_fake_prefill_runner(monkeypatch) + fake = sys.modules["tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k"] + fake.selfsampling_topk_prefill_ready = Mock(return_value=False) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + monkeypatch.setattr(TopK, "_prefill_capturing", staticmethod(lambda scores: True)) + top_k = TopK(2, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + scores = torch.randn(3, 8) + + row_starts, row_ends, output = _prefill_call(top_k, scores) + + runner.assert_not_called() + fake.selfsampling_topk_prefill_ready.assert_called_once_with(scores, output) + radix.assert_called_once_with(scores, row_starts, row_ends, output, 2) + + +def test_gvr_v2_prefill_capture_compiled_uses_engine(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + monkeypatch.setattr(TopK, "_prefill_capturing", staticmethod(lambda scores: True)) + top_k = TopK(2, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + + _prefill_call(top_k, torch.randn(3, 8)) + + runner.assert_called_once() diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index a6286bb28376..5048dfc0c0de 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -1302,6 +1302,10 @@ def counting(*a, **kw): finally: dev.get_compiled = orig assert calls["n"] == 0, f"warmup missed keys: {calls['n']} live compiles" + # the capture-time readiness query agrees with the cache: warmed shape + # ready, an unwarmed k (4 is valid but never compiled) not ready + assert ss_host.prefill_ready(lg, out) + assert not ss_host.prefill_ready(lg[:1], torch.empty((1, 4), dtype=torch.int32, device=_DEV)) def test_prefill_capture_no_host_sync(): From 7708dc35f9407e8e38ba1d5afb66fbb7bbc747fb Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 8 Sep 2026 06:11:52 +0000 Subject: [PATCH 4/4] [None][chore] Trim the prefill top-K comments to the invariants Every comment or docstring this PR added that ran past three lines is cut to the invariant it protects (18 blocks, no code change). Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention/backends/sparse/dsa/indexer.py | 7 +- .../attention/backends/sparse/dsa/metadata.py | 8 +-- .../top_k/gvr_topk_decode_self_sampling.py | 29 +++----- .../gvr_topk_decode_self_sampling_host.py | 70 +++++-------------- tensorrt_llm/_torch/modules/top_k.py | 7 +- .../_torch/pyexecutor/model_engine.py | 10 +-- .../sparse/dsa/test_dsa_fp4_indexer.py | 6 +- .../parallel/test_gvr_selfsampling_topk.py | 23 +++--- 8 files changed, 46 insertions(+), 114 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index 42c38ecc2903..2d894442753f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -754,11 +754,8 @@ def __init__( if self.use_cute_dsl_topk else TopKImplementation.CUDA_RADIX ) - # The self-sampling engine has a prefill form (per-row [ks, ke) - # windows); select it for prefill on exactly the layers where the - # two-level dispatch picks self-sampling for decode, so both phases - # share one config and one warmup. The temporal-hint engine has no - # prefill form, so those layers keep the exact radix prefill. + # Prefill uses the self-sampling engine on exactly the layers where the + # decode dispatch picks it; the temporal-hint engine has no prefill form. prefill_top_k_implementation = ( TopKImplementation.CUTE_DSL_GVR if ( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index b01142e42eb1..43615b2a7ade 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -437,12 +437,8 @@ def warmup_selfsampling_topk( for bs in batch_sizes or (): rows.add(int(bs) * nn) msl_c = int(self.get_indexer_max_seq_len()) - # Prefill leg: the self-sampling engine also serves prefill (per-row - # [ks, ke) windows). It is placed BEFORE the DeepGEMM decode-stride - # guard below (which would return early for an odd msl_c) because the - # DeepGEMM prefill stride is always a 256-multiple. Bounded to the six - # tier x U engines per k; best-effort under the same OOM guard as the - # decode leg. + # Prefill leg first: the decode-stride guard below may return early for + # an odd msl_c, but the DeepGEMM prefill stride is always 256-aligned. try: _ss_host.warmup_prefill(int(top_k), max(msl_c, 32768)) except torch.cuda.OutOfMemoryError: diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 6ba2c0e5814c..8bbd903075f8 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -1351,12 +1351,9 @@ def __init__( self.r_const = int(r_const) # hint-free: gather_hint sites compiled out (sentinel pass-through) self.hint_free = bool(hint_free) - # prefill: per-row window [ks, ke) from row_starts/row_ends (rides the - # kv_lens / pre_idx ABI slots); base rounds down to a 16B boundary and - # the <=3 lead lanes are positionally masked. Single-CTA-per-row only - # (no SPLIT/workspace/TSH); next_n==1, cr_shift==0 (ks/ke are already - # in compressed column units). All prefill edits are const_expr-gated - # so legacy/varlen codegen stays byte-identical. + # prefill: per-row [ks, ke) window riding the kv_lens/pre_idx ABI slots, + # base rounded down to 16B with the <=3 lead lanes masked, one CTA per + # row (no SPLIT); every edit is const_expr-gated so other codegen is unchanged. self.prefill = bool(prefill) if self.prefill: assert ( @@ -1529,10 +1526,8 @@ def kern( col0 = cutlass.Int32(0) if cutlass.const_expr(self.varlen): if cutlass.const_expr(self.prefill): - # per-row window [ks, ke) already in compressed column units - # (kv_lens slot = row_starts, pre_idx slot = row_ends); no - # next_n / cr_shift math. Clamp only for memory safety — the - # indexer guarantees 0 <= ks <= ke <= logits.shape[1]. + # ks/ke are already compressed column units (kv_lens = row_starts, + # pre_idx = row_ends); the clamps are for memory safety only. ks = kv_lens[row] ke = pre_idx[row] if ks < cutlass.Int32(0): @@ -1860,10 +1855,8 @@ def kern( # P3 slice (clamped in-row): the data P3 touches first starts # flowing while warp0 walks the chain. Short rows clamp every # hint to the row's last line — harmless. - # prefill: the base is shifted to col0, so the clamp must stay in - # the row's own window [col0, ke) — an npad-based clamp would - # over-read col0 columns past the last row's allocation. n4-1 is - # the last full in-window float4 (>=0 even for the n=0 short pass). + # prefill: the base is shifted to col0, so clamp within the row's own + # window [col0, ke) — n4-1 is the last full in-window float4. if cutlass.const_expr(self.prefill): plim4 = n4 - cutlass.Int32(1) if plim4 < cutlass.Int32(0): @@ -1897,11 +1890,9 @@ def kern( C.ld_g_f32x4(atom128, x_addr, p4, fsa) C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) if cutlass.const_expr(self.prefill): - # only thread 0's fsa (float4 index 0) can hold the <=3 masked - # lead lanes; substitute the always-valid lane 3 so the sample - # min/max fold and histogram stay finite and count-invariant - # (a materialized -inf would drive f2s_rz to INT_MIN and write - # out of bounds in the sample histogram at :1930). + # only thread 0's float4 0 holds the <=3 lead lanes; substitute + # lane 3 (a -inf would drive f2s_rz to INT_MIN and index the + # sample histogram out of bounds). if tidx == cutlass.Int32(0): ld_ = s_lead[0] for q in cutlass.range_constexpr(3): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index bf639ab72ad9..a88937fef940 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -679,13 +679,9 @@ def _main(blk_, minb_, u_, split_): _VARLEN_CACHE = {} # ---- prefill launcher cache ------------------------------------------------ -# Prefill routes always force R==1 (single CTA per row): route_streaming gives -# R>1 only for b<=74, so the representative row counts below (first row of each -# route band) pin R=1 and reduce the engine set to <=6 per k. The launcher -# compiled function depends only on the row TIER, k and the envelope bucket -# (which selects U on the tier-0 1024-thread arm; tiers 1/2 fix U), never on -# the exact row count (arbitrary q-tile / q-split remainders) or npad (a -# runtime scalar), so the cache stays bounded over a long-running server. +# Prefill forces R==1 (route_streaming gives R>1 only for b<=74). The compiled +# launcher depends only on the row tier, k and the envelope bucket — never on the +# exact row count or npad — so the cache stays bounded on a long-running server. _PREFILL_CACHE = {} _PREFILL_ROW_SLAB = 32768 # gridDim.y <= 65535; slab so keys stay bounded _PREFILL_TIER_ROWS = (75, 149, 297) # (rows<=148, 149..296, >296) band reps @@ -708,10 +704,8 @@ def _prefill_cache_key(tier: int, k: int, n_bucket: int): def _prefill_launcher(tier: int, k: int, n_bucket: int) -> tuple: - """Capture-time prefill plan + compiled launcher (main family, R=1). - - Mirrors ``_varlen_launcher``'s main branch but with r_const=1, split=False - (so tsh_en=0) and the prefill compile flag. SCAP_/CMP_/aim are envelope + """Prefill plan + compiled launcher: ``_varlen_launcher``'s main branch with + r_const=1, split=False and the prefill compile flag. SCAP_/CMP_ are envelope upper bounds; npad is filled per call in ``run_prefill``.""" key = _prefill_cache_key(tier, k, n_bucket) hit = _PREFILL_CACHE.get(key) @@ -1506,31 +1500,10 @@ def run_prefill( max_row_len: int | None = None, workspace: torch.Tensor | None = None, ) -> None: - """Hint-free self-sampling Top-K for the prefill phase, per-row windows. - - Row semantics (mirror of ``topKPerRowPrefill`` / ``indexer_topk_prefill``): - row ``r`` selects the Top-K of ``logits[r, ks:ke]`` where - ``ks = row_starts[r]``, ``ke = row_ends[r]`` (both int32, in the SAME - compressed column units the DeepGEMM prefill producer emits — no - ``next_n`` / ``compress_ratio`` math). ``k`` comes from - ``indices.shape[1]``. The output is written in the LOCAL frame (column - minus ``ks``) with a trailing ``-1`` pad; rows with ``nv = ke - ks <= k`` - get the identity ``0..nv-1`` (matching the radix short-row contract). The - engine reads exactly ``[r*npad + (ks & ~3), r*npad + ke)`` — no dependence - on any producer slack. - - Envelope: ``max_row_len`` (a capture-stable engine constant) or, when - omitted, ``logits.shape[1]`` — a host int, so the call performs NO device - reads and is CUDA-graph-replay safe (it refuses to compile a new plan - under capture). Launches in ``<=65535``-row slabs so ``gridDim.y`` never - overflows. - - KNOWN LIMITATION: rows containing NaN inside the window are out of - contract (as for the radix reference — both order NaN implementation- - specifically). DeepGEMM prefill logits are finite in-window. Trusted - invariant: ``0 <= ks <= ke <= logits.shape[1]`` (the indexer guarantees - it); the kernel clamps ``ke <= npad`` for memory safety only. - """ + """Hint-free self-sampling Top-K for prefill: row ``r`` selects the Top-K of + ``logits[r, ks:ke]`` (compressed columns) into the local frame (column - ks) + with a -1 pad; ``nv <= k`` rows get the identity, as ``indexer_topk_prefill``. + No device reads, never compiles under capture; trusts 0 <= ks <= ke <= shape[1].""" if logits.dtype is not _F32: raise RuntimeError( f"logits must be float32 (got {logits.dtype}); bf16/fp16 paths " @@ -1570,10 +1543,8 @@ def run_prefill( raise RuntimeError("indices base must be 16-byte aligned") if logits.stride(1) != 1: raise RuntimeError("logits inner stride must be 1") - # DeepGEMM prefill rows are 1024B-aligned with >=256 float slack, so the - # row stride is valid for EVERY row count (the varlen 1-row shape[1] rule - # is a paged-MQA-arena quirk that would reject odd-width single-token - # prefill tiles — the common fully-cached follow-up turn). + # key on stride(0) for every row count: DeepGEMM prefill rows are 1024B-aligned + # with slack, and the varlen 1-row shape[1] rule would reject odd-width tiles. npad = logits.stride(0) if npad & 3: raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") @@ -1610,10 +1581,8 @@ def run_prefill( ) lc = _prefill_launcher(tier, k, n_bucket) _, fn, (scap, cmp_), tail = lc - # ABI parity with the varlen main call: pre_idx slot = row_ends, - # kv_lens slot = row_starts. The n / SMP / TGT / Q / SS2 / TGT2 launch - # scalars are dead (re-derived per row); only npad / k / SCAP_ / CMP_ - # matter, R=1. + # varlen main ABI: pre_idx slot = row_ends, kv_lens slot = row_starts; + # only npad / k / SCAP_ / CMP_ matter (R=1), the other scalars are dead. pre = (0, npad, k, scap, cmp_, 1, 0, 0, 0, 0, 0) fn(lg[r0:r1], row_ends[r0:r1], indices[r0:r1], ws, *pre, row_starts[r0:r1], *tail) return @@ -1796,16 +1765,9 @@ def warmup_prefill( num_rows_list: Sequence[int] = (1, 149, 297), row_stride: int | None = None, ) -> None: - """TESTING/INIT ONLY — compile the prefill engine set before serving. - - Six engines per k at most: the tier-0 (1024-thread) arm walks the pow2 - envelope buckets (U = 1/2/4/8), tiers 1/2 fix U so one launch each. One - tiny real launch per distinct ``(tier, k, bucket)`` cache key; ``ks=0``, - ``ke=n_env`` (all long rows). ``max_cols`` is the compressed max column - count (``get_indexer_max_seq_len``); the bucket caps at 32768 (U=8 above), - so envelopes past it share one key. The done-key gates only the GPU - launches — the ``_PREFILL_CACHE`` population is idempotent. - """ + """Compile the prefill engine set before serving (<=6 per k): the tier-0 arm + walks the pow2 envelope buckets up to 32768, tiers 1/2 need one launch each. + ``max_cols`` is the compressed max column count; idempotent per done-key.""" dev = torch.cuda.current_device() k = int(top_k) max_cols = int(max_cols) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 60810d05a2dc..6e3f5f4394e4 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -220,11 +220,8 @@ def _forward_prefill( return output_indices def _selfsampling_prefill_ok(self, scores: torch.Tensor) -> bool: - """Engine hardware-format gate for the self-sampling prefill Top-K. - - fp32 row-major scores with a float4-aligned row stride and a 16B base - (the DeepGEMM prefill logits arena, whose rows are 1024B-aligned). The - all-short tile case is handled by the caller before this check.""" + """Engine format gate: fp32 row-major scores with a float4-aligned row + stride and a 16B base (the DeepGEMM prefill logits arena).""" return ( scores.dtype == torch.float32 and scores.stride(1) == 1 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b2c44a37a749..4bb8bceba566 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1678,13 +1678,9 @@ def _prewarm_cute_dsl_indexer_q(self) -> None: def _ensure_dsa_attn_metadata_for_warmup( self, resource_manager: ResourceManager) -> None: - """Build the DSA attention metadata if no warmup forward created it. - - A draft engine, a guided decoder, or a context-only server without - general warmup can skip every warmup forward; the DSA pre-compile - hooks would then find no metadata and the indexer top-K engines - would JIT on the first live request. No-op unless the backend is DSA. - """ + """Build the DSA attention metadata if no warmup forward created it, so + the top-K pre-compile hooks still run (draft engine, guided decoder or + context-only server without general warmup). No-op unless DSA.""" if getattr(self, "attn_metadata", None) is not None: return try: diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py index 60580764bc3b..31fc078f3f80 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py @@ -84,10 +84,8 @@ def _dense_context_bounds(seq_len: int, seq_len_kv: int, device): @skip_pre_blackwell @pytest.mark.parametrize("seq_len_kv", [1027, 4099]) def test_fp4_mqa_logits_pass_selfsampling_prefill_format_gate(seq_len_kv): - """DeepSeek-V4 prefill (cr=4) hands these logits to the self-sampling GVR - prefill engine, whose format gate needs a float4-aligned row stride on - odd compressed widths; an exact-width producer would silently fall back - to radix.""" + """The self-sampling prefill gate needs DeepGEMM's float4-aligned row stride + on odd compressed widths; an exact-width producer would fall back to radix.""" torch.manual_seed(0) num_heads, head_dim, seq_len = 64, 128, 64 q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 5048dfc0c0de..c10934cf8b46 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -1013,12 +1013,9 @@ def test_validate_run_ws_requires_16_byte_alignment() -> None: # =========================================================================== # ==== prefill: per-row [ks, ke) windows (run_prefill) ====================== # =========================================================================== -# Contract: row r selects the Top-K of logits[r, ks:ke] (ks=row_starts[r], -# ke=row_ends[r], compressed column units), output in the LOCAL frame -# (column - ks) with a trailing -1 pad; nv=ke-ks <= k gives identity 0..nv-1. -# The base is rounded down to a 16B boundary and the <=3 lead lanes are masked, -# so ks % 4 in {1,2,3} (2nd+ request of a multi-request chunk) is exercised -# with poison (+inf/NaN/3e38/-inf) written at [ks-3, ks) and [ke, npad). +# Contract: Top-K of logits[r, ks:ke] in the local frame (column - ks), -1 pad, +# identity for nv <= k. ks % 4 in {1,2,3} is exercised with poison written at +# [ks-3, ks) and [ke, npad). def _prefill_reference(logits, row_starts, row_ends, top_k): @@ -1037,10 +1034,9 @@ def _prefill_reference(logits, row_starts, row_ends, top_k): def _check_prefill_exact(logits, got, row_starts, row_ends, top_k): - """Tie-aware radix-parity check: trailing -1 pad from lengths; head unique - and in [0, nv); identity for nv <= k; exact index set when the k-th value - is unique, else strictly-above set + tie-class count (signed zeros and - genuine +/-inf compare like radix). NaN-in-window rows are structure-only.""" + """Tie-aware radix-parity check: -1 pad from lengths, unique head in [0, nv), + identity for nv <= k; exact set when the k-th value is unique, else the + strictly-above set plus a tie-class count. NaN-in-window rows: structure only.""" assert got.shape == (logits.shape[0], top_k) and got.dtype == torch.int32 ks, ke = row_starts.tolist(), row_ends.tolist() got64 = got.to(torch.int64) @@ -1113,10 +1109,9 @@ def _make_prefill_case(rows, ncols, ks_list, ke_list, *, top_k, seed, dist="rand @pytest.mark.parametrize("top_k", [512, 1024, 2048], ids=lambda k: f"k{k}") def test_prefill_causal_ramp(top_k): - """Single-request causal ramp (ks=0), 148 rows with nv = 1..148 < k: one - tier-0 launch where every row takes the short-row identity path. The k - boundary is test_prefill_short_rows; mixed short/long rows in one launch is - test_prefill_packed_misaligned_ks.""" + """Causal ramp (ks=0), 148 rows with nv = 1..148 < k: one tier-0 launch, all + short-row identity. k boundary: test_prefill_short_rows; mixed short/long + rows in one launch: test_prefill_packed_misaligned_ks.""" rows = 148 ks = [0] * rows ke = list(range(1, rows + 1))