diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index d6655133dbf6..2d894442753f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -754,9 +754,19 @@ def __init__( if self.use_cute_dsl_topk else TopKImplementation.CUDA_RADIX ) + # 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 ( + 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..43615b2a7ade 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -437,6 +437,15 @@ 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 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: + 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..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,8 @@ 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 from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -36,4 +38,6 @@ "tiered_topk", "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.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index a07ba529eb6d..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 @@ -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,18 @@ 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 [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 ( + 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 +1519,39 @@ 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): + # 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): + 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 +1565,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 +1635,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 +1687,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 +1849,20 @@ 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 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): + 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 +1889,15 @@ 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 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): + if cutlass.Int32(q) < ld_: + fsa[q] = fsa[3] # ============ P2: quantile rung from the sample ====================== smn = cutlass.Float32(float("inf")) @@ -1853,7 +1931,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 +2226,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 +2540,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 +2568,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 +2622,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 +2731,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 +2899,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 +2913,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 +2949,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 +2999,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 +3090,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 +3130,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 +3140,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..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 @@ -678,6 +678,63 @@ def _main(blk_, minb_, u_, split_): _VARLEN_CACHE = {} +# ---- prefill launcher cache ------------------------------------------------ +# 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 + + +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: + """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) + 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 +1492,119 @@ 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 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 " + "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") + # 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}") + 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 + # 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 + + +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", @@ -1444,7 +1614,10 @@ def run_varlen( "run", "run_ws", "run_varlen", + "run_prefill", + "prefill_ready", "warmup_varlen", + "warmup_prefill", "workspace_bytes", "WS_BYTES", "default_workspace", @@ -1580,3 +1753,54 @@ 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: + """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) + 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..6e3f5f4394e4 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -143,7 +143,59 @@ 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_prefill_ready, + selfsampling_topk_run_prefill, + ) + + 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. + 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 +206,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 +219,20 @@ def _forward_prefill( ) return output_indices + def _selfsampling_prefill_ok(self, scores: torch.Tensor) -> bool: + """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 + and scores.stride(0) % 4 == 0 + 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..4bb8bceba566 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,30 @@ 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, 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: + 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..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 @@ -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,38 @@ 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): + """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) + 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 6ac79f345954..65baad615620 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,55 @@ 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), # Hopper + (True, 120, 65536, False), # consumer Blackwell (SM120): engine is SM100/103-only + ], +) +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 +650,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 @@ -1479,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 a9d319eac9f5..1e4569d35af0 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,151 @@ 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, + selfsampling_topk_prefill_ready=Mock(return_value=True), + ), + ) + 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() + + +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 be0db369a452..c10934cf8b46 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,317 @@ 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: 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): + 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: -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) + 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): + """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)) + 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" + # 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(): + """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)