From 0cb4da48b90cc20838f5f1465dea6bc27812ca70 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 30 Aug 2026 13:47:56 +0000 Subject: [PATCH 1/4] [None][feat] GVR V2 hint-free mode: pre_idx=None derives the top-k bracket from the current row run_varlen(pre_idx=None) makes the self-sampling GVR decode kernel fully self-contained -- no prev-step top-K hint tensor required: - register families: bracket = min/max fold of the first k row values, which already sit in the row-load register fragments (zero extra loads); the hint prefetch and the hint-gather bracket arms (use_bm/use_img) are compiled out under the hf arm; DEG cells keep the whole-row fold (n<=3k: the row is the sample) - clustered register family: P0 samples the first k row elements (coalesced, cluster-uniform by construction, zero barrier changes) - streaming families: gather_hint sites compiled out (sentinel pass-through; a device-truth census on 886 real decode captures x BS{1,8,64,512} shows these sites never fire on the accept path) The hinted path is unchanged; pre_idx stays in the API as the eligibility gate and future extension point (e.g. disagg true-top-K seeding). k comes from indices.shape[1]; hint-free is auto-engine only. Exactness is hint-independent by contract (tie-aware value-multiset checks all-pass on real decode captures). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 89 +++++++++++++------ .../gvr_topk_decode_self_sampling_host.py | 50 +++++++---- .../parallel/test_gvr_selfsampling_topk.py | 40 +++++++-- 3 files changed, 130 insertions(+), 49 deletions(-) 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 36c2d67ff965..a54be25626c0 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 @@ -1326,6 +1326,7 @@ def __init__( next_n: int = 1, cr_shift: int = 0, r_const: int = 1, + hf: bool = False, ): assert nbs == 256, "SNB must stay 256" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) @@ -1346,6 +1347,8 @@ def __init__( self.next_n = int(next_n) self.cr_shift = int(cr_shift) self.r_const = int(r_const) + # hint-free: gather_hint sites compiled out (sentinel pass-through) + self.hf = bool(hf) 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 @@ -1973,9 +1976,10 @@ def kern( if T > cutlass.Float32(_NEG_INF): needg = cutlass.Int32(0) if needg != cutlass.Int32(0): - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT - ) # 2 barriers inside + if cutlass.const_expr(not self.hf): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT + ) # 2 barriers inside T = GMIN if sok != cutlass.Int32(0): # HIC tighten if tot0 >= TGT: @@ -2373,9 +2377,10 @@ def kern( else: # LAZY GATHER (sentinel equality flag) if GMIN == cutlass.Float32(C.SENT_LO): - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT - ) + if cutlass.const_expr(not self.hf): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT + ) floorhit = cutlass.Int32(1) if T > GMIN: floorhit = cutlass.Int32(0) @@ -2960,19 +2965,19 @@ def __call__( _COMPILE_CACHE = {} -def get_compiled(tpl, options_extra: str = ""): +def get_compiled(tpl, options_extra: str = "", hf=False): """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) + key = (tuple(tpl), options_extra, bool(hf)) hit = _COMPILE_CACHE.get(key) if hit is not None: return hit if len(tpl) == 7: blk, u, minb, nbs, kpt, split, tshg = tpl - kern = GvrMainKernel(blk, u, minb, nbs, kpt, bool(split), bool(tshg)) + kern = GvrMainKernel(blk, u, minb, nbs, kpt, bool(split), bool(tshg), hf=bool(hf)) else: blk, u, minb, nbs, kpt, split, tshg, next_n, cr_shift, r_const = tpl kern = GvrMainKernel( @@ -2987,6 +2992,7 @@ def get_compiled(tpl, options_extra: str = ""): next_n=next_n, cr_shift=cr_shift, r_const=r_const, + hf=bool(hf), ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -3309,6 +3315,7 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, + hf: bool = False, ): assert blk in (256, 512, 1024) and vpt in (1, 2, 4) assert nbh in (256, 512, 1024, 2048) @@ -3334,8 +3341,11 @@ def __init__( # derived compile-time constants self.S = vpt * 4 self.lnbh = {256: 8, 512: 9, 2048: 11}.get(nbh, 10) - self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 - self.use_img = img and vpt == 1 + # hint-free: bracket = min/max fold of the first k row values + # (already in registers); the hint-gather bracket arms are forced off + self.hf = bool(hf) + self.use_bm = ((not deg) and (not img) and kpt >= 2 and vpt == 1 and (not hf)) + self.use_img = img and vpt == 1 and (not hf) self.brl = (minb * blk <= 1024) or (vpt == 1) # ------------------------------------------------------------------ @@ -3551,7 +3561,7 @@ def kern( # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any # dependent gather; compiled out under DEG. pvs = [] - if cutlass.const_expr(not self.deg): + if cutlass.const_expr(not (self.deg or self.hf)): for t in cutlass.range_constexpr(KPT): pv = cutlass.Int32(-1) j = tid + cutlass.Int32(t * self.blk) @@ -3646,6 +3656,18 @@ def kern( lmin = fkey(lmn) lmax = fkey(lmx) # monotone cute.arch.barrier() # bm dies + elif cutlass.const_expr(self.hf and not self.deg): + lmn = cutlass.Float32(_POS_INF) + lmx = cutlass.Float32(_NEG_INF__reg) + for s in cutlass.range_constexpr(S): + pos = ((tid + cutlass.Int32((s // 4) * self.blk)) + << cutlass.Int32(2)) + cutlass.Int32(s % 4) + if pos < k: + v = _val(frags, s) + lmn = fmin_f32(lmn, v) + lmx = fmax_f32(lmx, v) + lmin = fkey(lmn) + lmax = fkey(lmx) elif cutlass.const_expr(self.deg): lmn = cutlass.Float32(_POS_INF) lmx = cutlass.Float32(_NEG_INF__reg) @@ -4207,10 +4229,10 @@ def __call__( _COMPILE_CACHE__reg: dict = {} -def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0): +def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0, hf=False): """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, MINB, KPT, CUR, DEG, IMG, NBH).""" - key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hf)) compiled = _COMPILE_CACHE__reg.get(key) if compiled is None: from cutlass.cute import runtime as _crt @@ -4229,6 +4251,7 @@ def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_ varlen=varlen, next_n=next_n, cr_shift=cr_shift, + hf=hf, ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() @@ -4377,6 +4400,7 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, + hf: bool = False, ): assert blk == 1024, "gvr_clus is always BLK=1024" assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1)" @@ -4392,6 +4416,7 @@ def __init__( self.cr_shift = int(cr_shift) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) + self.hf = bool(hf) # hint-free: gather_hint sites compiled out self.lcs = cs.bit_length() - 1 # log2(CS) for the per-row Q shift self.blk = blk self.u = u @@ -4947,9 +4972,10 @@ def kern( needg = cutlass.Int32(0) if needg != cutlass.Int32(0): # degenerate sample: identical on every rank of the cluster - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 - ) # 2 barriers + if cutlass.const_expr(not self.hf): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers T = GMIN if sok != cutlass.Int32(0): # HIC tighten if tot0 >= TGT: @@ -5178,9 +5204,10 @@ def kern( if tshtaken == cutlass.Int32(0): # LAZY GATHER — every rank computes identical GMIN if GMIN == cutlass.Float32(C.SENT_LO): - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 - ) # 2 barriers inside + if cutlass.const_expr(not self.hf): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers inside floorhit = cutlass.Int32(1) if T > GMIN: floorhit = cutlass.Int32(0) @@ -5559,17 +5586,18 @@ def get_compiled__clus( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, + hf: bool = False, ): """Compile (or fetch) the gvr_clus variant for constexpr tuple tpl = (BLK, U, MINB, NBS, CS); scap/cmp are smem-extent keys (every reachable route has 8192/2048 — asserted by run__clus()).""" - key = (tuple(tpl), scap, cmp_, options_extra, bool(varlen), int(next_n), int(cr_shift)) + key = (tuple(tpl), scap, cmp_, options_extra, bool(varlen), int(next_n), int(cr_shift), bool(hf)) hit = _COMPILE_CACHE__clus.get(key) if hit is not None: return hit blk, u, minb, nbs, cs = tpl kern = GvrClusKernel( - blk, u, minb, nbs, cs, scap=scap, cmp_=cmp_, varlen=varlen, next_n=next_n, cr_shift=cr_shift + blk, u, minb, nbs, cs, scap=scap, cmp_=cmp_, varlen=varlen, next_n=next_n, cr_shift=cr_shift, hf=hf ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -5787,6 +5815,7 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, + hf: bool = False, ): assert blk == BLKC, "all instantiations BLK=BLKC=1024" assert vpt in (1, 2, 4) and cs in (2, 4, 8) @@ -5803,6 +5832,8 @@ def __init__( self.cr_shift = int(cr_shift) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) + # hint-free: P0 samples the first k row elements (coalesced) instead of the hint + self.hf = bool(hf) self.S = vpt * 4 self.span = blk * vpt # float4 per CTA @@ -5958,8 +5989,12 @@ def kern( # ---- P0: redundant hint gather, EVERY CTA (k<=BLK by dispatch # gate). One coalesced word per thread, NO cluster barrier — # GMIN/GMAX identical everywhere by construction. - if tid < k: - pv0 = ld_g_i32(p_addr, tid) + if cutlass.const_expr(self.hf): + if tid < k: + pv0 = tid + else: + if tid < k: + pv0 = ld_g_i32(p_addr, tid) # ---- P1: row load — predicated flat float4[VPT] batch (the CUDA # has NO exact-fit peel here, guard is per-load). Issue all loads @@ -6461,16 +6496,16 @@ def __call__( _COMPILE_CACHE__regclus: dict = {} -def get_compiled__regclus(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0): +def get_compiled__regclus(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0, hf=False): """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS).""" - key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hf)) compiled = _COMPILE_CACHE__regclus.get(key) if compiled is None: from cutlass.cute import runtime as _crt blk, vpt, cs = tpl kernel = GvrRegClusKernel( - blk, vpt, cs, pdl=pdl, varlen=varlen, next_n=next_n, cr_shift=cr_shift + blk, vpt, cs, pdl=pdl, varlen=varlen, next_n=next_n, cr_shift=cr_shift, hf=hf ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() 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 165048c8dcfe..5186cdd6e358 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,12 +679,12 @@ def _main(blk_, minb_, u_, split_): _VARLEN_CACHE = {} -def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): +def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr, hf=False): """Capture-time varlen plan + compiled launcher. The gvr_main port is the universally correct fallback; specialist family tiers below. Every choice here is a function of capture-stable quantities only — mirroring the in-tree runner's pick_tuning(graph_capture=...) discipline.""" - key = (num_rows, npad, k, n_env, next_n, cr) + key = (num_rows, npad, k, n_env, next_n, cr, bool(hf)) hit = _VARLEN_CACHE.get(key) if hit is not None: return hit @@ -699,7 +699,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): plan_free = route(num_rows, n_eff, npad, k) if plan_free["kernel"] == "reg_clus": fn = dev.get_compiled__regclus( - tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift + tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift, hf=hf ) lc = ("reg_clus", fn, n_eff) _VARLEN_CACHE[key] = lc @@ -713,7 +713,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): # / short-row handling lives in-kernel. if plan_free["kernel"] in ("reg", "regimg"): fn = dev.get_compiled__reg( - tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift + tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift, hf=hf ) rt_f = plan_free["rt"] lc = ( @@ -739,6 +739,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): varlen=True, next_n=next_n, cr_shift=cr_shift, + hf=hf, ) lc = ( "clus", @@ -754,7 +755,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): # TSHG (tpl[6]) is dead under varlen (the ctor compiles the TSH # machinery in whenever SPLIT); normalize it out of the compile key so # row counts differing only in that slot share one engine - fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const)) + fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const), hf=hf) big = num_rows * r_const <= 148 aim_base = ( ((4 * k if k >= 1024 else 2 * k) if r_const == 1 else 2 * k) @@ -1234,7 +1235,7 @@ def run_ws( def run_varlen( logits: torch.Tensor, - pre_idx: torch.Tensor, + pre_idx: torch.Tensor | None, kv_lens: torch.Tensor, indices: torch.Tensor, next_n: int = 1, @@ -1257,6 +1258,12 @@ def run_varlen( compress_ratio`` valid entries (cr 1 = DSv3.2, 4 = DSv4 Flash/Pro); ``pre_idx`` ``[batch, k]`` is REQUEST-level raw prev-step top-K, shared by a request's ``next_n`` rows (offset-free hint contract); + ``pre_idx=None`` selects HINT-FREE mode: the bracket is derived from + the current row itself (register families: min/max fold of the first + k row values, in registers / a coalesced first-k sample; streaming + families never consume the hint on the accept path), ``k`` comes from + ``indices.shape[1]``, engine "auto" only — exactness is unchanged by + contract; per-row ``n_r <= k`` takes the short path (identity + ``-1`` tail). ENGINES: ``engine="auto"`` (default) launches the per-row IN-KERNEL @@ -1306,7 +1313,9 @@ def run_varlen( batch = num_rows // nn if kv_lens.shape[0] != batch: raise RuntimeError(f"kv_lens length {kv_lens.shape[0]} != num_rows/next_n = {batch}") - if len(pre_idx.shape) != 2 or pre_idx.shape[0] != batch: + if pre_idx is not None and ( + len(pre_idx.shape) != 2 or pre_idx.shape[0] != batch + ): raise RuntimeError( f"pre_idx must be [batch={batch}, k] REQUEST-level, got {tuple(pre_idx.shape)}" ) @@ -1330,18 +1339,21 @@ def run_varlen( # batch-dim check is CRITICAL: the kernel grid comes from # logits.shape[0], so a short indices/values tensor would be written # out of bounds). - if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): + hf = pre_idx is None + if not (logits.is_cuda and indices.is_cuda and (hf or pre_idx.is_cuda)): raise RuntimeError("all tensors must be CUDA") - if logits.dtype is not _F32 or pre_idx.dtype is not _I32 or indices.dtype is not _I32: + if logits.dtype is not _F32 or indices.dtype is not _I32 or ( + not hf and pre_idx.dtype is not _I32 + ): raise RuntimeError("logits must be float32; pre_idx/indices 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)}" ) - k = pre_idx.shape[1] + k = indices.shape[1] if hf else pre_idx.shape[1] if indices.shape[1] < k: raise RuntimeError(f"indices width {indices.shape[1]} < k={k}") - if not (pre_idx.is_contiguous() and indices.is_contiguous() and kv_lens.is_contiguous()): + if not ((hf or pre_idx.is_contiguous()) and indices.is_contiguous() and kv_lens.is_contiguous()): raise RuntimeError("pre_idx/indices/kv_lens must be contiguous") # logits: accept row-major views with a wider row stride (the DSL # paged-MQA logits arena is 256-aligned and column-sliced — a legal @@ -1390,7 +1402,7 @@ def run_varlen( # R increment (bounded plans, bounded _VARLEN_CACHE) n_env = 1 << max(n_env - 1, 1).bit_length() n_env = min(max(n_env, 1), npad) - key = (num_rows, npad, k, n_env, nn, cr) + key = (num_rows, npad, k, n_env, nn, cr, hf) lc = _VARLEN_CACHE.get(key) if lc is None: if _is_capturing(): @@ -1398,26 +1410,28 @@ def run_varlen( "varlen launcher not compiled for this shape — warm up " "before CUDA graph capture" ) - lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) + lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr, hf=hf) idx = indices if idx.shape[1] != k: idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) vals = values if vals is not None and vals.shape[1] != k: vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) + # hf engines never read the pre slot; alias the out tensor in + pre_arg = idx if hf else pre_idx if lc[0] == "reg_clus": # compiled ABI: (logits, pre_idx, kv_lens, out, n_envelope) - lc[1](lg, pre_idx, kv_lens, idx, lc[2]) + lc[1](lg, pre_arg, kv_lens, idx, lc[2]) elif lc[0] == "reg": # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, CMP, QC, smem) - lc[1](lg, pre_idx, kv_lens, idx, *lc[2]) + lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) elif lc[0] == "clus": # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, npad, k, # SCAP, CMP, dead DYN x5) - lc[1](lg, pre_idx, kv_lens, idx, *lc[2]) + lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) else: _, fn, pre, tail = lc - fn(lg, pre_idx, idx, ws, *pre, kv_lens, *tail) + fn(lg, pre_arg, idx, ws, *pre, kv_lens, *tail) if vals is not None: idx64 = idx.to(torch.int64) vals.copy_(lg.gather(1, idx64.clamp_min(0))) @@ -1425,6 +1439,8 @@ def run_varlen( return if engine != "reference": raise RuntimeError(f"engine must be 'auto' or 'reference', got {engine!r}") + if pre_idx is None: + raise RuntimeError("hint-free (pre_idx=None) is auto-engine only") # ---- reference engine (differential oracle): b=1 host loop -------------- if _is_capturing(): 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 44fbb6396b0c..3d6c74385c83 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -228,7 +228,7 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): _check_exact(logits, indices, n_valid, ref_vals) -def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="auto"): +def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="auto", hint_free=False): """Build a per-row-poisoned varlen batch, run run_varlen, verify every row against its own n_r (production formula) — short rows included.""" batch, rows = len(kv), len(kv) * next_n @@ -238,10 +238,14 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut logits = torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 for r in range(rows): logits[r, n_r[r] :] = 3e38 # poison beyond each row's OWN n_r - pre_idx = torch.empty((batch, top_k), dtype=torch.int32, device=_DEV) - for q in range(batch): - nmin = max(min(n_r[q * next_n : (q + 1) * next_n]), 1) - pre_idx[q] = torch.randint(0, nmin, (top_k,), generator=gen, dtype=torch.int32, device=_DEV) + pre_idx = None + if not hint_free: + pre_idx = torch.empty((batch, top_k), dtype=torch.int32, device=_DEV) + for q in range(batch): + nmin = max(min(n_r[q * next_n : (q + 1) * next_n]), 1) + pre_idx[q] = torch.randint( + 0, nmin, (top_k,), generator=gen, dtype=torch.int32, device=_DEV + ) indices = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) values = ( torch.full((rows, top_k), 7.0, dtype=torch.float32, device=_DEV) if with_values else None @@ -301,6 +305,32 @@ def test_selfsampling_topk_varlen(kv, next_n, cr, top_k, engine): _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr, engine=engine) +@pytest.mark.parametrize( + "kv,next_n,cr,top_k", + [ + ([33000, 8200, 300], 1, 1, 512), + ([131075, 32800, 2000], 1, 4, 512), + ([9000, 5001], 2, 1, 512), + ([65540], 4, 4, 1024), + ([40000, 7003], 3, 4, 512), + ], + ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4", "cr4_mtp3"], +) +def test_selfsampling_topk_varlen_hint_free(kv, next_n, cr, top_k): + """pre_idx=None (hint-free): the bracket comes from the current row + itself; exactness must be identical to the hinted contract.""" + _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr, hint_free=True) + + +def test_selfsampling_topk_varlen_hint_free_reference_rejected(): + """The b=1 reference loop needs a hint tensor; pre_idx=None must raise.""" + kv_lens = torch.tensor([9000], dtype=torch.int32, device=_DEV) + logits = torch.randn((1, 9024), dtype=torch.float32, device=_DEV) + indices = torch.full((1, 512), -7, dtype=torch.int32, device=_DEV) + with pytest.raises(RuntimeError, match="auto-engine only"): + ss_host.run_varlen(logits, None, kv_lens, indices, engine="reference") + + def test_selfsampling_topk_varlen_engine_matches_reference(): """Differential: the in-kernel engine's per-row value multisets must equal the reference loop's on a mixed batch (deep SPLIT rows, tsh band, From f6980dd198e0f3ec4f72a6dacfb32492e68c2ea0 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 31 Aug 2026 03:06:58 +0000 Subject: [PATCH 2/4] [None][feat] GVR V2 decode is hint-free by default (TRTLLM_GVR_V2_HINTED=1 opts back) - top_k.py V2 branch launches run_varlen(pre_idx=None) by default; the env restores prev-step hint consumption, so hinted-vs-hint-free e2e A/B runs on one build. 886x11-BS real decode grid: hint-free/hinted gm 0.992, exactness 19,492/19,492. - the self-sampling warmup (metadata.py -> warmup_varlen) warms the hint-free launcher keys the dispatch looks up; fix: the launcher-cache tail pass dropped the hf bit, so a hint-free CUDA-graph capture raised "not compiled" right after a hint-free warmup. - fix: six existing _VARLEN_CACHE UT asserts still used the pre-hf 6-tuple key and would KeyError against the extended launcher key. - tests: module dispatch (default None / env opt-back passes the prior through) + hint-free warmup/CUDA-graph capture; validated on B200 (hinted and hint-free keys, capture, tie-aware exactness). - ruff-format pass on the three PR files (pre-commit CI gate). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/metadata.py | 6 +- .../top_k/gvr_topk_decode_self_sampling.py | 34 ++++++++++-- .../gvr_topk_decode_self_sampling_host.py | 40 ++++++++++---- tensorrt_llm/_torch/modules/top_k.py | 24 +++++--- tests/unittest/_torch/modules/test_top_k.py | 55 +++++++++++++++++++ .../parallel/test_gvr_selfsampling_topk.py | 46 +++++++++++++--- 6 files changed, 174 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index b69d856d743b..f9b7bb528141 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -398,7 +398,10 @@ def warmup_selfsampling_topk( if row_stride % 4: return # helper takes max_seq_len in kv-token space (get_indexer_max_seq_len - # is compressed — same multiply-back as the dispatch seam) + # is compressed — same multiply-back as the dispatch seam). The + # hint-free bit must mirror the TopK module's launch mode: the + # launcher key includes it, so warming the wrong arm leaves capture + # uncompiled. try: _ss_host.warmup_varlen( int(top_k), @@ -407,6 +410,7 @@ def warmup_selfsampling_topk( next_n=int(next_n), num_rows_list=tuple(sorted(rows)), row_stride=row_stride, + hint_free=os.environ.get("TRTLLM_GVR_V2_HINTED", "0") != "1", ) except torch.cuda.OutOfMemoryError: # warmup is best-effort: the dispatch works without it (engines 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 a54be25626c0..2d803ba08442 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 @@ -3344,7 +3344,7 @@ def __init__( # hint-free: bracket = min/max fold of the first k row values # (already in registers); the hint-gather bracket arms are forced off self.hf = bool(hf) - self.use_bm = ((not deg) and (not img) and kpt >= 2 and vpt == 1 and (not hf)) + self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 and (not hf) self.use_img = img and vpt == 1 and (not hf) self.brl = (minb * blk <= 1024) or (vpt == 1) @@ -3660,8 +3660,9 @@ def kern( lmn = cutlass.Float32(_POS_INF) lmx = cutlass.Float32(_NEG_INF__reg) for s in cutlass.range_constexpr(S): - pos = ((tid + cutlass.Int32((s // 4) * self.blk)) - << cutlass.Int32(2)) + cutlass.Int32(s % 4) + pos = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) if pos < k: v = _val(frags, s) lmn = fmin_f32(lmn, v) @@ -5591,13 +5592,32 @@ def get_compiled__clus( """Compile (or fetch) the gvr_clus variant for constexpr tuple tpl = (BLK, U, MINB, NBS, CS); scap/cmp are smem-extent keys (every reachable route has 8192/2048 — asserted by run__clus()).""" - key = (tuple(tpl), scap, cmp_, options_extra, bool(varlen), int(next_n), int(cr_shift), bool(hf)) + key = ( + tuple(tpl), + scap, + cmp_, + options_extra, + bool(varlen), + int(next_n), + int(cr_shift), + bool(hf), + ) hit = _COMPILE_CACHE__clus.get(key) if hit is not None: return hit blk, u, minb, nbs, cs = tpl kern = GvrClusKernel( - blk, u, minb, nbs, cs, scap=scap, cmp_=cmp_, varlen=varlen, next_n=next_n, cr_shift=cr_shift, hf=hf + blk, + u, + minb, + nbs, + cs, + scap=scap, + cmp_=cmp_, + varlen=varlen, + next_n=next_n, + cr_shift=cr_shift, + hf=hf, ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -6496,7 +6516,9 @@ def __call__( _COMPILE_CACHE__regclus: dict = {} -def get_compiled__regclus(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0, hf=False): +def get_compiled__regclus( + tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0, hf=False +): """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS).""" key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hf)) compiled = _COMPILE_CACHE__regclus.get(key) 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 5186cdd6e358..4664d2ffaa84 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 @@ -1313,9 +1313,7 @@ def run_varlen( batch = num_rows // nn if kv_lens.shape[0] != batch: raise RuntimeError(f"kv_lens length {kv_lens.shape[0]} != num_rows/next_n = {batch}") - if pre_idx is not None and ( - len(pre_idx.shape) != 2 or pre_idx.shape[0] != batch - ): + if pre_idx is not None and (len(pre_idx.shape) != 2 or pre_idx.shape[0] != batch): raise RuntimeError( f"pre_idx must be [batch={batch}, k] REQUEST-level, got {tuple(pre_idx.shape)}" ) @@ -1342,8 +1340,10 @@ def run_varlen( hf = pre_idx is None if not (logits.is_cuda and indices.is_cuda and (hf or pre_idx.is_cuda)): raise RuntimeError("all tensors must be CUDA") - if logits.dtype is not _F32 or indices.dtype is not _I32 or ( - not hf and pre_idx.dtype is not _I32 + if ( + logits.dtype is not _F32 + or indices.dtype is not _I32 + or (not hf and pre_idx.dtype is not _I32) ): raise RuntimeError("logits must be float32; pre_idx/indices int32") if len(indices.shape) != 2 or indices.shape[0] != num_rows: @@ -1353,7 +1353,9 @@ def run_varlen( k = indices.shape[1] if hf else pre_idx.shape[1] if indices.shape[1] < k: raise RuntimeError(f"indices width {indices.shape[1]} < k={k}") - if not ((hf or pre_idx.is_contiguous()) and indices.is_contiguous() and kv_lens.is_contiguous()): + if not ( + (hf or pre_idx.is_contiguous()) and indices.is_contiguous() and kv_lens.is_contiguous() + ): raise RuntimeError("pre_idx/indices/kv_lens must be contiguous") # logits: accept row-major views with a wider row stride (the DSL # paged-MQA logits arena is 256-aligned and column-sliced — a legal @@ -1512,6 +1514,7 @@ def warmup_varlen( next_n: int = 1, num_rows_list: Sequence[int] = (1,), row_stride: int | None = None, + hint_free: bool = False, ) -> None: """TESTING/INIT ONLY — compile the varlen engine's envelope tuples. @@ -1527,6 +1530,10 @@ def warmup_varlen( producer layout (e.g. the DSL paged-MQA arena's 256-element rounding) must pass it; the 64-element default only matches producers that round the same way. + + ``hint_free`` must match the mode dispatch will launch (the launcher key + includes the hint-free bit): warm with ``hint_free=True`` when serving + calls ``run_varlen(pre_idx=None)``. """ dev = torch.cuda.current_device() nn = max(1, int(next_n)) @@ -1576,7 +1583,16 @@ def warmup_varlen( raise RuntimeError( f"row_stride must be a float4-multiple >= n_env={n_env}, got {row_stride}" ) - key = (dev, int(top_k), int(max_seq_len), int(compress_ratio), nn, tuple(rows_list), npad) + key = ( + dev, + int(top_k), + int(max_seq_len), + int(compress_ratio), + nn, + tuple(rows_list), + npad, + bool(hint_free), + ) with _VARLEN_WARMUP_LOCK: if key in _VARLEN_WARMUP_DONE: return @@ -1585,13 +1601,17 @@ def warmup_varlen( # contiguous prefix views (compile keys depend on shapes only) logits = torch.zeros((rows_max, npad), dtype=torch.float32, device=dev) kv_lens = torch.full((rows_max // nn,), int(max_seq_len), dtype=torch.int32, device=dev) - pre_idx = torch.zeros((rows_max // nn, int(top_k)), dtype=torch.int32, device=dev) + pre_idx = ( + None + if hint_free + else torch.zeros((rows_max // nn, int(top_k)), dtype=torch.int32, device=dev) + ) out = torch.empty((rows_max, int(top_k)), dtype=torch.int32, device=dev) for rows in rows_list: batch = rows // nn run_varlen( logits[:rows], - pre_idx[:batch], + None if pre_idx is None else pre_idx[:batch], kv_lens[:batch], out[:rows], next_n=nn, @@ -1606,6 +1626,6 @@ def warmup_varlen( # CUDA-graph capture at any requested geometry finds its key immediately. n_env_l = min(max(int(max_seq_len) >> (0 if int(compress_ratio) == 1 else 2), 1), npad) for r in req_rows: - _varlen_launcher(r, npad, int(top_k), n_env_l, nn, int(compress_ratio)) + _varlen_launcher(r, npad, int(top_k), n_env_l, nn, int(compress_ratio), hf=bool(hint_free)) with _VARLEN_WARMUP_LOCK: _VARLEN_WARMUP_DONE.add(key) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 950b02c30f93..6da5ed4131a2 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from enum import Enum import torch @@ -59,6 +60,9 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio + # GVR V2 decode is hint-free by default; TRTLLM_GVR_V2_HINTED=1 + # restores prev-step hint consumption. + self._gvr_v2_hinted = os.environ.get("TRTLLM_GVR_V2_HINTED", "0") == "1" def forward( self, @@ -88,7 +92,9 @@ def forward( max_seq_len: Maximum decode score width used for GVR kernel tuning. gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices`` is the required caller-owned int32 previous selection with - shape ``[num_requests, top_k]`` on ``scores.device``. + shape ``[num_requests, top_k]`` on ``scores.device`` + (``CUTE_DSL_GVR_V2`` launches hint-free by default and reads + it only under ``TRTLLM_GVR_V2_HINTED=1``). ``gvr_row_order`` is an optional int32 request ordering with shape ``[num_requests]`` on the same device. @@ -284,19 +290,23 @@ def _forward_decode_gvr( logger.info_once( "self-sampling GVR top-K engaged " f"(K={self.top_k}, cr={self.compress_ratio}, " - f"next_n={next_n}).", + f"next_n={next_n}, " + f"{'hinted' if self._gvr_v2_hinted else 'hint-free'}).", key="selfsampling_topk_engaged", ) # Self-sampling GVR varlen engine (TRTLLM_GVR_SELF_SAMPLING=1): # one launch for the batch; per-row n from device kv_lens, # capture-stable tuning from the max-seq-len engine constant - # (no host reads — CUDA-graph safe). Hints are consumed raw - # (offset-free contract). The module receives max_seq_len in - # COMPRESSED index space; run_varlen's max_seq_len is in - # kv-token space like sequence_lengths — multiply back. + # (no host reads — CUDA-graph safe). Hint-free by default + # (pre_idx=None: the kernel brackets from the current row); + # TRTLLM_GVR_V2_HINTED=1 restores raw prev-step hint + # consumption (offset-free contract). The module receives + # max_seq_len in COMPRESSED index space; run_varlen's + # max_seq_len is in kv-token space like sequence_lengths — + # multiply back. selfsampling_topk_run_varlen( scores, - gvr_prior_indices, + gvr_prior_indices if self._gvr_v2_hinted else None, sequence_lengths, output_indices, next_n=next_n, diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 9682d7981ab5..106089843932 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -2,7 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the reusable sparse index-selection Top-K module.""" +import sys from contextlib import nullcontext +from types import SimpleNamespace from unittest.mock import Mock, call import pytest @@ -201,6 +203,59 @@ def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None: assert gvr.call_args.kwargs["order_row"] is row_order +def _install_fake_selfsampling_runner(monkeypatch) -> Mock: + """Replace the lazily imported self-sampling varlen entry with a Mock.""" + runner = Mock() + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k", + SimpleNamespace(selfsampling_topk_run_varlen=runner), + ) + return runner + + +def _run_gvr_v2_decode(top_k: TopK, prior_indices: torch.Tensor) -> None: + scores = torch.randn(1, 8) # satisfies the V2 hardware-format gate + top_k( + scores, + torch.empty(1, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=torch.tensor([32], dtype=torch.int32), + scan_lengths=torch.tensor([8], dtype=torch.int32), + next_n=1, + max_seq_len=16, + gvr_ext_kwargs={"gvr_prior_indices": prior_indices}, + ) + + +def test_gvr_v2_decode_defaults_to_hint_free(monkeypatch) -> None: + runner = _install_fake_selfsampling_runner(monkeypatch) + monkeypatch.delenv("TRTLLM_GVR_V2_HINTED", raising=False) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, + compress_ratio=4, + ) + prior_indices = torch.zeros(1, 2, dtype=torch.int32) + + _run_gvr_v2_decode(top_k, prior_indices) + + args, kwargs = runner.call_args + assert args[1] is None + assert kwargs == {"next_n": 1, "compress_ratio": 4, "max_seq_len": 64} + + +def test_gvr_v2_decode_hinted_env_restores_prior(monkeypatch) -> None: + runner = _install_fake_selfsampling_runner(monkeypatch) + monkeypatch.setenv("TRTLLM_GVR_V2_HINTED", "1") + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2) + prior_indices = torch.zeros(1, 2, dtype=torch.int32) + + _run_gvr_v2_decode(top_k, prior_indices) + + assert runner.call_args.args[1] is prior_indices + + def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) 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 3d6c74385c83..cec1473181de 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -228,7 +228,9 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): _check_exact(logits, indices, n_valid, ref_vals) -def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="auto", hint_free=False): +def _run_varlen_case( + kv, next_n, cr, top_k, seed, with_values=False, engine="auto", hint_free=False +): """Build a per-row-poisoned varlen batch, run run_varlen, verify every row against its own n_r (production formula) — short rows included.""" batch, rows = len(kv), len(kv) * next_n @@ -689,6 +691,36 @@ def test_selfsampling_warmup_row_stride_matches_arena(): assert torch.equal(ref, got) +def test_selfsampling_warmup_hint_free_matches_dispatch(): + """warmup_varlen(hint_free=True) must compile the launcher key the + hint-free dispatch (pre_idx=None) looks up: the key includes the + hint-free bit, so a hinted warmup would leave capture uncompiled.""" + k, msl = 512, 8300 + stride = (msl + 255) // 256 * 256 + rows = 2 + ss_host.warmup_varlen( + k, + msl, + compress_ratio=1, + next_n=1, + num_rows_list=(rows,), + row_stride=stride, + hint_free=True, + ) + arena = torch.randn(rows, stride, dtype=torch.float32, device=_DEV) + logits = arena[:, :msl] + kv = torch.full((rows,), msl, dtype=torch.int32, device=_DEV) + out = torch.empty(rows, k, dtype=torch.int32, device=_DEV) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + ss_host.run_varlen(logits, None, kv, out, max_seq_len=msl) + g.replay() + torch.cuda.synchronize() + ref = torch.topk(arena[:, :msl], k, dim=1).values.sort(dim=1).values + got = arena.gather(1, out.long().clamp_min(0)).sort(dim=1).values + assert torch.equal(ref, got) + + def test_selfsampling_varlen_regclus_parity_and_oracle(): """The varlen launcher must admit the clustered register-resident family exactly where the free route picks it (route() parity tier 1), and the @@ -707,7 +739,7 @@ def test_selfsampling_varlen_regclus_parity_and_oracle(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) - key = (rows, npad, k, msl_c, nn, cr) + key = (rows, npad, k, msl_c, nn, cr, False) assert ss_host._VARLEN_CACHE[key][0] == "reg_clus", ss_host._VARLEN_CACHE[key][0] ss_host.run_varlen( lg, @@ -784,7 +816,7 @@ def test_selfsampling_varlen_reg_parity_and_oracle(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) - key = (rows, npad, k, msl_c, nn, cr) + key = (rows, npad, k, msl_c, nn, cr, False) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] ss_host.run_varlen( lg, @@ -835,7 +867,7 @@ def test_selfsampling_varlen_clus_parity_and_oracle(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) - key = (rows, npad, k, msl_c, nn, cr) + key = (rows, npad, k, msl_c, nn, cr, False) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] ss_host.run_varlen( lg, @@ -872,7 +904,7 @@ def test_selfsampling_varlen_clus_cuda_graph(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr) + key = (rows, msl_c, k, msl_c, 1, cr, False) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) @@ -899,7 +931,7 @@ def test_selfsampling_varlen_reg_cuda_graph(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr) + key = (rows, msl_c, k, msl_c, 1, cr, False) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) @@ -928,7 +960,7 @@ def test_selfsampling_varlen_full_row_range(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr) + key = (rows, msl_c, k, msl_c, 1, cr, False) assert key in ss_host._VARLEN_CACHE, "row count must dispatch in-engine" ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values From a410156fda469a58db3b03520a8025d1ade575d7 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 31 Aug 2026 08:39:20 +0000 Subject: [PATCH 3/4] [None][feat] GVR V2: pin decode output width to top_k (hint-free k source) Hint-free run_varlen derives k from indices.shape[1]; a wider scratch would silently become the k. Assert at the module seam (the only in-tree caller) and cover with a negative dispatch test. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- tensorrt_llm/_torch/modules/top_k.py | 2 ++ tests/unittest/_torch/modules/test_top_k.py | 22 +++++++++++++++++++-- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 6da5ed4131a2..7c270125de64 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -285,6 +285,8 @@ def _forward_decode_gvr( and scores.data_ptr() % 16 == 0 and (scores.shape[0] > 1 or scores.shape[1] % 4 == 0) ): + # hint-free k derives from the output width; pin it to the module's k + assert output_indices.shape[1] == self.top_k from ..cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_varlen logger.info_once( diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 106089843932..31b51376d56c 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -214,11 +214,11 @@ def _install_fake_selfsampling_runner(monkeypatch) -> Mock: return runner -def _run_gvr_v2_decode(top_k: TopK, prior_indices: torch.Tensor) -> None: +def _run_gvr_v2_decode(top_k: TopK, prior_indices: torch.Tensor, out_width: int = 2) -> None: scores = torch.randn(1, 8) # satisfies the V2 hardware-format gate top_k( scores, - torch.empty(1, 2, dtype=torch.int32), + torch.empty(1, out_width, dtype=torch.int32), is_prefill=False, sequence_lengths=torch.tensor([32], dtype=torch.int32), scan_lengths=torch.tensor([8], dtype=torch.int32), @@ -256,6 +256,24 @@ def test_gvr_v2_decode_hinted_env_restores_prior(monkeypatch) -> None: assert runner.call_args.args[1] is prior_indices +def test_gvr_v2_decode_rejects_output_width_mismatch(monkeypatch) -> None: + """Hint-free k derives from the output width; a scratch wider than + top_k must be rejected before launch, not silently become the k.""" + runner = _install_fake_selfsampling_runner(monkeypatch) + monkeypatch.delenv("TRTLLM_GVR_V2_HINTED", raising=False) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, + compress_ratio=4, + ) + prior_indices = torch.zeros(1, 2, dtype=torch.int32) + + with pytest.raises(AssertionError): + _run_gvr_v2_decode(top_k, prior_indices, out_width=3) + + runner.assert_not_called() + + def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) From c1aa1eac28ff8d64316c61bd555e719fab35dfc9 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:18:20 +0000 Subject: [PATCH 4/4] [None][feat] GVR V2 decode is hint-free only: retire the hinted mode and its env switch Production V2 no longer has a hinted arm: run_varlen drops the pre_idx and engine parameters, the launcher cache key loses its mode bit, and the TRTLLM_GVR_V2_HINTED escape is removed along with the reference host-loop entry. The kernel-side flag is renamed hf -> hint_free. The V2 hardware-format fall-through now lands on the exact insertion/radix op instead of the CUDA GVR path, so a selected V2 never consumes temporal state; prefill prior seeding is skipped for V2 via TopK.needs_gvr_prior. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/metadata.py | 9 +- .../top_k/gvr_topk_decode_self_sampling.py | 98 ++++-- .../gvr_topk_decode_self_sampling_host.py | 293 +++++++--------- tensorrt_llm/_torch/modules/top_k.py | 64 ++-- tests/unittest/_torch/modules/test_top_k.py | 73 +++- .../parallel/test_gvr_selfsampling_topk.py | 329 ++++-------------- 6 files changed, 353 insertions(+), 513 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index f9b7bb528141..b4abcd5ff74f 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -397,11 +397,9 @@ def warmup_selfsampling_topk( row_stride = msl_c if row_stride % 4: return - # helper takes max_seq_len in kv-token space (get_indexer_max_seq_len - # is compressed — same multiply-back as the dispatch seam). The - # hint-free bit must mirror the TopK module's launch mode: the - # launcher key includes it, so warming the wrong arm leaves capture - # uncompiled. + # The helper takes max_seq_len in KV-token space; + # get_indexer_max_seq_len is compressed, so multiply it back as at + # the dispatch seam. try: _ss_host.warmup_varlen( int(top_k), @@ -410,7 +408,6 @@ def warmup_selfsampling_topk( next_n=int(next_n), num_rows_list=tuple(sorted(rows)), row_stride=row_stride, - hint_free=os.environ.get("TRTLLM_GVR_V2_HINTED", "0") != "1", ) except torch.cuda.OutOfMemoryError: # warmup is best-effort: the dispatch works without it (engines 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 2d803ba08442..a43d39d2e59a 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 @@ -29,6 +29,7 @@ import contextlib import sys +from typing import Any import cutlass import cutlass.cute as cute @@ -1326,8 +1327,8 @@ def __init__( next_n: int = 1, cr_shift: int = 0, r_const: int = 1, - hf: bool = False, - ): + hint_free: bool = False, + ) -> None: assert nbs == 256, "SNB must stay 256" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) assert kpt in (1, 2, 4, 8) and minb in (1, 2, 4) @@ -1348,7 +1349,7 @@ def __init__( self.cr_shift = int(cr_shift) self.r_const = int(r_const) # hint-free: gather_hint sites compiled out (sentinel pass-through) - self.hf = bool(hf) + self.hint_free = bool(hint_free) 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 @@ -1976,7 +1977,7 @@ def kern( if T > cutlass.Float32(_NEG_INF): needg = cutlass.Int32(0) if needg != cutlass.Int32(0): - if cutlass.const_expr(not self.hf): + if cutlass.const_expr(not self.hint_free): GMIN, GMAX = C.gather_hint( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT ) # 2 barriers inside @@ -2377,7 +2378,7 @@ def kern( else: # LAZY GATHER (sentinel equality flag) if GMIN == cutlass.Float32(C.SENT_LO): - if cutlass.const_expr(not self.hf): + if cutlass.const_expr(not self.hint_free): GMIN, GMAX = C.gather_hint( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT ) @@ -2965,19 +2966,21 @@ def __call__( _COMPILE_CACHE = {} -def get_compiled(tpl, options_extra: str = "", hf=False): +def get_compiled(tpl: tuple, options_extra: str = "", hint_free: 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(hf)) + key = (tuple(tpl), options_extra, bool(hint_free)) hit = _COMPILE_CACHE.get(key) if hit is not None: return hit if len(tpl) == 7: blk, u, minb, nbs, kpt, split, tshg = tpl - kern = GvrMainKernel(blk, u, minb, nbs, kpt, bool(split), bool(tshg), hf=bool(hf)) + kern = GvrMainKernel( + blk, u, minb, nbs, kpt, bool(split), bool(tshg), hint_free=bool(hint_free) + ) else: blk, u, minb, nbs, kpt, split, tshg, next_n, cr_shift, r_const = tpl kern = GvrMainKernel( @@ -2992,7 +2995,7 @@ def get_compiled(tpl, options_extra: str = "", hf=False): next_n=next_n, cr_shift=cr_shift, r_const=r_const, - hf=bool(hf), + hint_free=bool(hint_free), ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -3315,8 +3318,8 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, - hf: bool = False, - ): + hint_free: bool = False, + ) -> None: assert blk in (256, 512, 1024) and vpt in (1, 2, 4) assert nbh in (256, 512, 1024, 2048) assert nbh % blk == 0 or blk % nbh == 0 @@ -3343,9 +3346,9 @@ def __init__( self.lnbh = {256: 8, 512: 9, 2048: 11}.get(nbh, 10) # hint-free: bracket = min/max fold of the first k row values # (already in registers); the hint-gather bracket arms are forced off - self.hf = bool(hf) - self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 and (not hf) - self.use_img = img and vpt == 1 and (not hf) + self.hint_free = bool(hint_free) + self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 and (not hint_free) + self.use_img = img and vpt == 1 and (not hint_free) self.brl = (minb * blk <= 1024) or (vpt == 1) # ------------------------------------------------------------------ @@ -3561,7 +3564,7 @@ def kern( # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any # dependent gather; compiled out under DEG. pvs = [] - if cutlass.const_expr(not (self.deg or self.hf)): + if cutlass.const_expr(not (self.deg or self.hint_free)): for t in cutlass.range_constexpr(KPT): pv = cutlass.Int32(-1) j = tid + cutlass.Int32(t * self.blk) @@ -3656,7 +3659,7 @@ def kern( lmin = fkey(lmn) lmax = fkey(lmx) # monotone cute.arch.barrier() # bm dies - elif cutlass.const_expr(self.hf and not self.deg): + elif cutlass.const_expr(self.hint_free and not self.deg): lmn = cutlass.Float32(_POS_INF) lmx = cutlass.Float32(_NEG_INF__reg) for s in cutlass.range_constexpr(S): @@ -4230,10 +4233,18 @@ def __call__( _COMPILE_CACHE__reg: dict = {} -def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0, hf=False): +def get_compiled__reg( + tpl: tuple, + dump_dir: str | None = None, + pdl: bool = False, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, + hint_free: bool = False, +) -> Any: """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, MINB, KPT, CUR, DEG, IMG, NBH).""" - key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hf)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hint_free)) compiled = _COMPILE_CACHE__reg.get(key) if compiled is None: from cutlass.cute import runtime as _crt @@ -4252,7 +4263,7 @@ def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_ varlen=varlen, next_n=next_n, cr_shift=cr_shift, - hf=hf, + hint_free=hint_free, ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() @@ -4401,8 +4412,8 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, - hf: bool = False, - ): + hint_free: bool = False, + ) -> None: assert blk == 1024, "gvr_clus is always BLK=1024" assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1)" assert nbs == 256, "SNB must stay 256" @@ -4417,7 +4428,7 @@ def __init__( self.cr_shift = int(cr_shift) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) - self.hf = bool(hf) # hint-free: gather_hint sites compiled out + self.hint_free = bool(hint_free) # hint-free: gather_hint sites compiled out self.lcs = cs.bit_length() - 1 # log2(CS) for the per-row Q shift self.blk = blk self.u = u @@ -4973,7 +4984,7 @@ def kern( needg = cutlass.Int32(0) if needg != cutlass.Int32(0): # degenerate sample: identical on every rank of the cluster - if cutlass.const_expr(not self.hf): + if cutlass.const_expr(not self.hint_free): GMIN, GMAX = C.gather_hint( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 ) # 2 barriers @@ -5205,7 +5216,7 @@ def kern( if tshtaken == cutlass.Int32(0): # LAZY GATHER — every rank computes identical GMIN if GMIN == cutlass.Float32(C.SENT_LO): - if cutlass.const_expr(not self.hf): + if cutlass.const_expr(not self.hint_free): GMIN, GMAX = C.gather_hint( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 ) # 2 barriers inside @@ -5580,15 +5591,15 @@ def __call__( def get_compiled__clus( - tpl, + tpl: tuple, scap: int = 8192, cmp_: int = 2048, options_extra: str = "", varlen: bool = False, next_n: int = 1, cr_shift: int = 0, - hf: bool = False, -): + hint_free: bool = False, +) -> Any: """Compile (or fetch) the gvr_clus variant for constexpr tuple tpl = (BLK, U, MINB, NBS, CS); scap/cmp are smem-extent keys (every reachable route has 8192/2048 — asserted by run__clus()).""" @@ -5600,7 +5611,7 @@ def get_compiled__clus( bool(varlen), int(next_n), int(cr_shift), - bool(hf), + bool(hint_free), ) hit = _COMPILE_CACHE__clus.get(key) if hit is not None: @@ -5617,7 +5628,7 @@ def get_compiled__clus( varlen=varlen, next_n=next_n, cr_shift=cr_shift, - hf=hf, + hint_free=hint_free, ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -5835,8 +5846,8 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, - hf: bool = False, - ): + hint_free: bool = False, + ) -> None: assert blk == BLKC, "all instantiations BLK=BLKC=1024" assert vpt in (1, 2, 4) and cs in (2, 4, 8) self.blk = blk @@ -5853,7 +5864,7 @@ def __init__( if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) # hint-free: P0 samples the first k row elements (coalesced) instead of the hint - self.hf = bool(hf) + self.hint_free = bool(hint_free) self.S = vpt * 4 self.span = blk * vpt # float4 per CTA @@ -6009,7 +6020,7 @@ def kern( # ---- P0: redundant hint gather, EVERY CTA (k<=BLK by dispatch # gate). One coalesced word per thread, NO cluster barrier — # GMIN/GMAX identical everywhere by construction. - if cutlass.const_expr(self.hf): + if cutlass.const_expr(self.hint_free): if tid < k: pv0 = tid else: @@ -6517,17 +6528,30 @@ def __call__( def get_compiled__regclus( - tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0, hf=False -): + tpl: tuple, + dump_dir: str | None = None, + pdl: bool = False, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, + hint_free: bool = False, +) -> Any: """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS).""" - key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hf)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hint_free)) compiled = _COMPILE_CACHE__regclus.get(key) if compiled is None: from cutlass.cute import runtime as _crt blk, vpt, cs = tpl kernel = GvrRegClusKernel( - blk, vpt, cs, pdl=pdl, varlen=varlen, next_n=next_n, cr_shift=cr_shift, hf=hf + blk, + vpt, + cs, + pdl=pdl, + varlen=varlen, + next_n=next_n, + cr_shift=cr_shift, + hint_free=hint_free, ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() 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 4664d2ffaa84..cea15314e1f4 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,12 +679,19 @@ def _main(blk_, minb_, u_, split_): _VARLEN_CACHE = {} -def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr, hf=False): +def _varlen_launcher( + num_rows: int, + npad: int, + k: int, + n_env: int, + next_n: int, + cr: int, +) -> tuple: """Capture-time varlen plan + compiled launcher. The gvr_main port is the universally correct fallback; specialist family tiers below. Every choice here is a function of capture-stable quantities only — mirroring the in-tree runner's pick_tuning(graph_capture=...) discipline.""" - key = (num_rows, npad, k, n_env, next_n, cr, bool(hf)) + key = (num_rows, npad, k, n_env, next_n, cr) hit = _VARLEN_CACHE.get(key) if hit is not None: return hit @@ -699,7 +706,11 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr, hf=False): plan_free = route(num_rows, n_eff, npad, k) if plan_free["kernel"] == "reg_clus": fn = dev.get_compiled__regclus( - tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift, hf=hf + tuple(plan_free["tpl"]), + varlen=True, + next_n=next_n, + cr_shift=cr_shift, + hint_free=True, ) lc = ("reg_clus", fn, n_eff) _VARLEN_CACHE[key] = lc @@ -713,7 +724,11 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr, hf=False): # / short-row handling lives in-kernel. if plan_free["kernel"] in ("reg", "regimg"): fn = dev.get_compiled__reg( - tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift, hf=hf + tuple(plan_free["tpl"]), + varlen=True, + next_n=next_n, + cr_shift=cr_shift, + hint_free=True, ) rt_f = plan_free["rt"] lc = ( @@ -739,7 +754,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr, hf=False): varlen=True, next_n=next_n, cr_shift=cr_shift, - hf=hf, + hint_free=True, ) lc = ( "clus", @@ -755,7 +770,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr, hf=False): # TSHG (tpl[6]) is dead under varlen (the ctor compiles the TSH # machinery in whenever SPLIT); normalize it out of the compile key so # row counts differing only in that slot share one engine - fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const), hf=hf) + fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const), hint_free=True) big = num_rows * r_const <= 148 aim_base = ( ((4 * k if k >= 1024 else 2 * k) if r_const == 1 else 2 * k) @@ -1235,17 +1250,15 @@ def run_ws( def run_varlen( logits: torch.Tensor, - pre_idx: torch.Tensor | None, kv_lens: torch.Tensor, indices: torch.Tensor, next_n: int = 1, compress_ratio: int = 1, values: torch.Tensor | None = None, max_seq_len: int | None = None, - engine: str = "auto", workspace: torch.Tensor | None = None, ) -> None: - """Production-contract varlen entry (per-row device kv_lens). + """Run hint-free self-sampling Top-K with per-request device KV lengths. Row semantics (mirror of ``heuristicTopKDecode.cu`` and the in-tree ``cute_dsl_gvr_topk_decode`` runner): @@ -1256,25 +1269,18 @@ def run_varlen( not new-token seq_lens); row ``r`` uses ``n_r = (kv_lens[r // next_n] - next_n + (r % next_n) + 1) // compress_ratio`` valid entries (cr 1 = DSv3.2, 4 = DSv4 Flash/Pro); - ``pre_idx`` ``[batch, k]`` is REQUEST-level raw prev-step top-K, - shared by a request's ``next_n`` rows (offset-free hint contract); - ``pre_idx=None`` selects HINT-FREE mode: the bracket is derived from - the current row itself (register families: min/max fold of the first - k row values, in registers / a coalesced first-k sample; streaming - families never consume the hint on the accept path), ``k`` comes from - ``indices.shape[1]``, engine "auto" only — exactness is unchanged by - contract; + the bracket is derived from the current row itself (register families: + min/max fold of the first k row values; streaming families do not + consume a temporal hint on the accept path); ``k`` comes from + ``indices.shape[1]``; per-row ``n_r <= k`` takes the short path (identity + ``-1`` tail). - ENGINES: ``engine="auto"`` (default) launches the per-row IN-KERNEL - gvr_main varlen port — ONE launch for the whole batch; each CTA reads its - row's kv_len on device and re-derives the sampling ladder (route_dynamic + The per-row in-kernel engine launches once for the whole batch. Each CTA + reads its row's kv_len on device and re-derives the sampling ladder (route_dynamic formula mirror), so with ``max_seq_len`` given (a capture-stable engine constant, e.g. dsa.py's ``indexer_max_seq_len``) the call performs NO host reads. Without ``max_seq_len`` the envelope comes from ONE ``kv_lens.max()`` host read (documented sync, refused under capture). - ``engine="reference"`` keeps the b=1 host-loop reference implementation — - the differential oracle the in-kernel engine is validated against. KNOWN LIMITATION: on rows containing NaN logits the selected index SET can differ from ``heuristicTopKDecode.cu`` (both kernels order NaNs @@ -1313,10 +1319,6 @@ def run_varlen( batch = num_rows // nn if kv_lens.shape[0] != batch: raise RuntimeError(f"kv_lens length {kv_lens.shape[0]} != num_rows/next_n = {batch}") - if pre_idx is not None and (len(pre_idx.shape) != 2 or pre_idx.shape[0] != batch): - raise RuntimeError( - f"pre_idx must be [batch={batch}, k] REQUEST-level, got {tuple(pre_idx.shape)}" - ) d = logits.get_device() if not 0 <= d < _GVR_MAX_DEV: raise RuntimeError(f"device index out of range: {d}") @@ -1330,153 +1332,103 @@ def run_varlen( if ws is None: ws = default_workspace(logits) - if engine == "auto": - # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- - # Full validation battery (the engine bypasses _run_impl — every - # check the batch-uniform path enforces is replayed here; the - # batch-dim check is CRITICAL: the kernel grid comes from - # logits.shape[0], so a short indices/values tensor would be written - # out of bounds). - hf = pre_idx is None - if not (logits.is_cuda and indices.is_cuda and (hf or pre_idx.is_cuda)): - raise RuntimeError("all tensors must be CUDA") + # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- + # Full validation battery (the engine bypasses _run_impl — every + # check the batch-uniform path enforces is replayed here; the + # batch-dim check is CRITICAL: the kernel grid comes from + # logits.shape[0], so a short indices/values tensor would be written + # out of bounds). + if not (logits.is_cuda and indices.is_cuda): + raise RuntimeError("all tensors must be CUDA") + if logits.dtype is not _F32 or indices.dtype is not _I32: + raise RuntimeError("logits must be float32; 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)}" + ) + k = indices.shape[1] + if not (indices.is_contiguous() and kv_lens.is_contiguous()): + raise RuntimeError("indices/kv_lens must be contiguous") + # logits: accept row-major views with a wider row stride (the DSL + # paged-MQA logits arena is 256-aligned and column-sliced — a legal + # NON-contiguous view). The kernel only needs (base, row stride): + # widen back to a compact [rows, stride] view over the same storage; + # the tail columns are never classified (per-row n gates all reads). + if logits.stride(1) != 1: + raise RuntimeError("logits inner stride must be 1") + npad = logits.stride(0) if num_rows > 1 else logits.shape[1] + lg = logits + if not logits.is_contiguous(): + 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 npad & 3: + raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") + if lg.data_ptr() & 15: + raise RuntimeError("logits base must be 16-byte aligned") + if values is not None: + if not values.is_cuda or values.dtype is not _F32: + raise RuntimeError("values must be CUDA float32") if ( - logits.dtype is not _F32 - or indices.dtype is not _I32 - or (not hf and pre_idx.dtype is not _I32) + len(values.shape) != 2 + or values.shape[0] != num_rows + or values.shape[1] < k + or not values.is_contiguous() ): - raise RuntimeError("logits must be float32; pre_idx/indices 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)}" + f"values must be contiguous [num_rows={num_rows}, >=k], got {tuple(values.shape)}" ) - k = indices.shape[1] if hf else pre_idx.shape[1] - if indices.shape[1] < k: - raise RuntimeError(f"indices width {indices.shape[1]} < k={k}") - if not ( - (hf or pre_idx.is_contiguous()) and indices.is_contiguous() and kv_lens.is_contiguous() - ): - raise RuntimeError("pre_idx/indices/kv_lens must be contiguous") - # logits: accept row-major views with a wider row stride (the DSL - # paged-MQA logits arena is 256-aligned and column-sliced — a legal - # NON-contiguous view). The kernel only needs (base, row stride): - # widen back to a compact [rows, stride] view over the same storage; - # the tail columns are never classified (per-row n gates all reads). - if logits.stride(1) != 1: - raise RuntimeError("logits inner stride must be 1") - npad = logits.stride(0) if num_rows > 1 else logits.shape[1] - lg = logits - if not logits.is_contiguous(): - 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 npad & 3: - raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") - if lg.data_ptr() & 15: - raise RuntimeError("logits base must be 16-byte aligned") - if values is not None: - if not values.is_cuda or values.dtype is not _F32: - raise RuntimeError("values must be CUDA float32") - if ( - len(values.shape) != 2 - or values.shape[0] != num_rows - or values.shape[1] < k - or not values.is_contiguous() - ): - raise RuntimeError( - f"values must be contiguous [num_rows={num_rows}, >=k], " - f"got {tuple(values.shape)}" - ) - cshift = 0 if cr == 1 else 2 - if max_seq_len is not None: - n_env = int(max_seq_len) >> cshift - else: - if _is_capturing(): - raise RuntimeError( - "run_varlen without max_seq_len reads kv_lens.max() on " - "host — pass max_seq_len (a capture-stable engine " - "constant) under CUDA graph capture" - ) - n_env = int(kv_lens.max().item()) >> cshift - # eager mode: quantize the data-dependent envelope up to the next - # power of two so a growing decode does not recompile at every - # R increment (bounded plans, bounded _VARLEN_CACHE) - n_env = 1 << max(n_env - 1, 1).bit_length() - n_env = min(max(n_env, 1), npad) - key = (num_rows, npad, k, n_env, nn, cr, hf) - lc = _VARLEN_CACHE.get(key) - if lc is None: - if _is_capturing(): - raise RuntimeError( - "varlen launcher not compiled for this shape — warm up " - "before CUDA graph capture" - ) - lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr, hf=hf) - idx = indices - if idx.shape[1] != k: - idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) - vals = values - if vals is not None and vals.shape[1] != k: - vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) - # hf engines never read the pre slot; alias the out tensor in - pre_arg = idx if hf else pre_idx - if lc[0] == "reg_clus": - # compiled ABI: (logits, pre_idx, kv_lens, out, n_envelope) - lc[1](lg, pre_arg, kv_lens, idx, lc[2]) - elif lc[0] == "reg": - # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, CMP, QC, smem) - lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) - elif lc[0] == "clus": - # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, npad, k, - # SCAP, CMP, dead DYN x5) - lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) - else: - _, fn, pre, tail = lc - fn(lg, pre_arg, idx, ws, *pre, kv_lens, *tail) - if vals is not None: - idx64 = idx.to(torch.int64) - vals.copy_(lg.gather(1, idx64.clamp_min(0))) - vals.masked_fill_(idx < 0, torch.finfo(_F32).min) - return - if engine != "reference": - raise RuntimeError(f"engine must be 'auto' or 'reference', got {engine!r}") - if pre_idx is None: - raise RuntimeError("hint-free (pre_idx=None) is auto-engine only") - - # ---- reference engine (differential oracle): b=1 host loop -------------- - if _is_capturing(): - raise RuntimeError( - "run_varlen reference engine reads kv_lens on host, illegal under CUDA graph capture" - ) - # match the engine's flat-packed output convention for wider-than-k - # buffers (pack ONCE from the tensor base, then slice per row) - k = pre_idx.shape[1] + cshift = 0 if cr == 1 else 2 + if max_seq_len is not None: + n_env = int(max_seq_len) >> cshift + else: + if _is_capturing(): + raise RuntimeError( + "run_varlen without max_seq_len reads kv_lens.max() on " + "host — pass max_seq_len (a capture-stable engine " + "constant) under CUDA graph capture" + ) + n_env = int(kv_lens.max().item()) >> cshift + # eager mode: quantize the data-dependent envelope up to the next + # power of two so a growing decode does not recompile at every + # R increment (bounded plans, bounded _VARLEN_CACHE) + n_env = 1 << max(n_env - 1, 1).bit_length() + n_env = min(max(n_env, 1), npad) + key = (num_rows, npad, k, n_env, nn, cr) + lc = _VARLEN_CACHE.get(key) + if lc is None: + if _is_capturing(): + raise RuntimeError( + "varlen launcher not compiled for this shape — warm up before CUDA graph capture" + ) + lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) idx = indices - if len(idx.shape) != 2 or idx.shape[0] != num_rows: - raise RuntimeError( - f"indices must be [num_rows={num_rows}, >=k], got {tuple(indices.shape)}" - ) if idx.shape[1] != k: idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) vals = values if vals is not None and vals.shape[1] != k: vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) - kl = kv_lens.tolist() # the ONE documented D2H sync of this engine - for r in range(num_rows): - # production graph slots can carry kv_len < next_n (padded / evicted - # requests): clamp to the empty row, emitting all -1 — the same - # contract the in-kernel engine implements - actual = max(kl[r // nn] - nn + (r % nn) + 1, 0) - req = r // nn - _run_impl( - logits[r : r + 1], - pre_idx[req : req + 1], - actual // cr, - idx[r : r + 1], - ws, - None if vals is None else vals[r : r + 1], - ) + # Hint-free engines do not read the compiled kernel's pre_idx ABI slot. + pre_arg = idx + if lc[0] == "reg_clus": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_envelope) + lc[1](lg, pre_arg, kv_lens, idx, lc[2]) + elif lc[0] == "reg": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, CMP, QC, smem) + lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) + elif lc[0] == "clus": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, npad, k, + # SCAP, CMP, dead DYN x5) + lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) + else: + _, fn, pre, tail = lc + fn(lg, pre_arg, idx, ws, *pre, kv_lens, *tail) + if vals is not None: + idx64 = idx.to(torch.int64) + vals.copy_(lg.gather(1, idx64.clamp_min(0))) + vals.masked_fill_(idx < 0, torch.finfo(_F32).min) + return __all__ = [ @@ -1514,7 +1466,6 @@ def warmup_varlen( next_n: int = 1, num_rows_list: Sequence[int] = (1,), row_stride: int | None = None, - hint_free: bool = False, ) -> None: """TESTING/INIT ONLY — compile the varlen engine's envelope tuples. @@ -1531,9 +1482,6 @@ def warmup_varlen( must pass it; the 64-element default only matches producers that round the same way. - ``hint_free`` must match the mode dispatch will launch (the launcher key - includes the hint-free bit): warm with ``hint_free=True`` when serving - calls ``run_varlen(pre_idx=None)``. """ dev = torch.cuda.current_device() nn = max(1, int(next_n)) @@ -1591,7 +1539,6 @@ def warmup_varlen( nn, tuple(rows_list), npad, - bool(hint_free), ) with _VARLEN_WARMUP_LOCK: if key in _VARLEN_WARMUP_DONE: @@ -1601,24 +1548,18 @@ def warmup_varlen( # contiguous prefix views (compile keys depend on shapes only) logits = torch.zeros((rows_max, npad), dtype=torch.float32, device=dev) kv_lens = torch.full((rows_max // nn,), int(max_seq_len), dtype=torch.int32, device=dev) - pre_idx = ( - None - if hint_free - else torch.zeros((rows_max // nn, int(top_k)), dtype=torch.int32, device=dev) - ) out = torch.empty((rows_max, int(top_k)), dtype=torch.int32, device=dev) for rows in rows_list: batch = rows // nn run_varlen( logits[:rows], - None if pre_idx is None else pre_idx[:batch], kv_lens[:batch], out[:rows], next_n=nn, compress_ratio=int(compress_ratio), max_seq_len=int(max_seq_len), ) - del logits, kv_lens, pre_idx, out + del logits, kv_lens, out torch.cuda.synchronize() # band launches compiled every ENGINE; now populate the per-row-count # LAUNCHER cache entries for the exact requested row counts (pure host @@ -1626,6 +1567,6 @@ def warmup_varlen( # CUDA-graph capture at any requested geometry finds its key immediately. n_env_l = min(max(int(max_seq_len) >> (0 if int(compress_ratio) == 1 else 2), 1), npad) for r in req_rows: - _varlen_launcher(r, npad, int(top_k), n_env_l, nn, int(compress_ratio), hf=bool(hint_free)) + _varlen_launcher(r, npad, int(top_k), n_env_l, nn, int(compress_ratio)) with _VARLEN_WARMUP_LOCK: _VARLEN_WARMUP_DONE.add(key) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 7c270125de64..4896d5977fe9 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -4,7 +4,6 @@ from __future__ import annotations -import os from enum import Enum import torch @@ -31,6 +30,10 @@ class TopKImplementation(str, Enum): TopKImplementation.CUTE_DSL_GVR, TopKImplementation.CUTE_DSL_GVR_V2, } +_TEMPORAL_GVR_IMPLEMENTATIONS = { + TopKImplementation.CUDA_GVR, + TopKImplementation.CUTE_DSL_GVR, +} _MAX_RADIX_BLOCKS_PER_ROW = 10 @@ -60,9 +63,11 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio - # GVR V2 decode is hint-free by default; TRTLLM_GVR_V2_HINTED=1 - # restores prev-step hint consumption. - self._gvr_v2_hinted = os.environ.get("TRTLLM_GVR_V2_HINTED", "0") == "1" + + @property + def needs_gvr_prior(self) -> bool: + """Return whether decode consumes previous-step Top-K indices.""" + return self.decode_implementation in _TEMPORAL_GVR_IMPLEMENTATIONS def forward( self, @@ -91,10 +96,10 @@ def forward( next_n: Number of decode rows per request. max_seq_len: Maximum decode score width used for GVR kernel tuning. gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices`` - is the required caller-owned int32 previous selection with - shape ``[num_requests, top_k]`` on ``scores.device`` - (``CUTE_DSL_GVR_V2`` launches hint-free by default and reads - it only under ``TRTLLM_GVR_V2_HINTED=1``). + is required by the temporal CUDA and CuTe DSL GVR paths. It is + caller-owned int32 previous selection with shape + ``[num_requests, top_k]`` on ``scores.device``. GVR V2 does + not consume this state. ``gvr_row_order`` is an optional int32 request ordering with shape ``[num_requests]`` on the same device. @@ -267,7 +272,6 @@ def _forward_decode_gvr( gvr_prior_indices: torch.Tensor | None = None, gvr_row_order: torch.Tensor | None = None, ) -> torch.Tensor: - assert gvr_prior_indices is not None if self.decode_implementation == TopKImplementation.CUTE_DSL_GVR_V2: assert max_seq_len is not None if ( @@ -292,23 +296,18 @@ def _forward_decode_gvr( logger.info_once( "self-sampling GVR top-K engaged " f"(K={self.top_k}, cr={self.compress_ratio}, " - f"next_n={next_n}, " - f"{'hinted' if self._gvr_v2_hinted else 'hint-free'}).", + f"next_n={next_n}, hint-free).", key="selfsampling_topk_engaged", ) # Self-sampling GVR varlen engine (TRTLLM_GVR_SELF_SAMPLING=1): # one launch for the batch; per-row n from device kv_lens, # capture-stable tuning from the max-seq-len engine constant - # (no host reads — CUDA-graph safe). Hint-free by default - # (pre_idx=None: the kernel brackets from the current row); - # TRTLLM_GVR_V2_HINTED=1 restores raw prev-step hint - # consumption (offset-free contract). The module receives - # max_seq_len in COMPRESSED index space; run_varlen's - # max_seq_len is in kv-token space like sequence_lengths — - # multiply back. + # (no host reads — CUDA-graph safe). The module receives + # max_seq_len in compressed index space; run_varlen's value + # is in KV-token space like sequence_lengths, so multiply it + # back by the compression ratio. selfsampling_topk_run_varlen( scores, - gvr_prior_indices if self._gvr_v2_hinted else None, sequence_lengths, output_indices, next_n=next_n, @@ -320,11 +319,26 @@ def _forward_decode_gvr( "TRTLLM_GVR_SELF_SAMPLING=1 but the decode scores do not " "satisfy the engine's hardware-format gate " f"(dtype={scores.dtype}, strides={tuple(scores.stride())}); " - "falling through to the CUDA GVR top-K path.", + "falling back to the CUDA insertion/radix Top-K path.", key="selfsampling_topk_fallthrough", ) - if self.decode_implementation != TopKImplementation.CUTE_DSL_GVR: - # CUDA_GVR, or the V2 hardware-format fall-through above + radix_indices, radix_values = self._get_radix_workspace(scores) + torch.ops.trtllm.indexer_topk_decode( + scores, + sequence_lengths, + output_indices, + next_n, + self.top_k, + pre_idx=None, + heuristic_scratch=None, + compress_ratio=self.compress_ratio, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, + ) + return output_indices + + assert gvr_prior_indices is not None + if self.decode_implementation == TopKImplementation.CUDA_GVR: workspace = self._get_workspace( scores, (scores.shape[0], self.top_k), @@ -344,7 +358,7 @@ def _forward_decode_gvr( radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) - else: + elif self.decode_implementation == TopKImplementation.CUTE_DSL_GVR: assert max_seq_len is not None torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, @@ -357,6 +371,8 @@ def _forward_decode_gvr( max_seq_len=max_seq_len, order_row=gvr_row_order, ) + else: + raise AssertionError(f"Unexpected GVR implementation: {self.decode_implementation}") return output_indices def update_gvr_prior_from_prefill( @@ -378,7 +394,7 @@ def update_gvr_prior_from_prefill( The slice starting at ``request_offset`` is updated in place. request_offset: First request row to update in the prior state. """ - if self.decode_implementation not in _GVR_IMPLEMENTATIONS: + if not self.needs_gvr_prior: return assert gvr_prior_indices is not None last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 31b51376d56c..bd03c476d4da 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -214,7 +214,7 @@ def _install_fake_selfsampling_runner(monkeypatch) -> Mock: return runner -def _run_gvr_v2_decode(top_k: TopK, prior_indices: torch.Tensor, out_width: int = 2) -> None: +def _run_gvr_v2_decode(top_k: TopK, out_width: int = 2) -> None: scores = torch.randn(1, 8) # satisfies the V2 hardware-format gate top_k( scores, @@ -224,52 +224,76 @@ def _run_gvr_v2_decode(top_k: TopK, prior_indices: torch.Tensor, out_width: int scan_lengths=torch.tensor([8], dtype=torch.int32), next_n=1, max_seq_len=16, - gvr_ext_kwargs={"gvr_prior_indices": prior_indices}, ) -def test_gvr_v2_decode_defaults_to_hint_free(monkeypatch) -> None: +def test_gvr_v2_decode_is_hint_free(monkeypatch) -> None: runner = _install_fake_selfsampling_runner(monkeypatch) - monkeypatch.delenv("TRTLLM_GVR_V2_HINTED", raising=False) top_k = TopK( 2, decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, compress_ratio=4, ) - prior_indices = torch.zeros(1, 2, dtype=torch.int32) - _run_gvr_v2_decode(top_k, prior_indices) + _run_gvr_v2_decode(top_k) args, kwargs = runner.call_args - assert args[1] is None + assert len(args) == 3 + assert args[0].shape == (1, 8) + assert args[1].tolist() == [32] + assert args[2].shape == (1, 2) assert kwargs == {"next_n": 1, "compress_ratio": 4, "max_seq_len": 64} + assert not top_k.needs_gvr_prior -def test_gvr_v2_decode_hinted_env_restores_prior(monkeypatch) -> None: +def test_gvr_v2_hardware_gate_falls_back_without_prior(monkeypatch) -> None: runner = _install_fake_selfsampling_runner(monkeypatch) - monkeypatch.setenv("TRTLLM_GVR_V2_HINTED", "1") - top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2) - prior_indices = torch.zeros(1, 2, dtype=torch.int32) + decode = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, + compress_ratio=4, + ) + scores = torch.randn(1, 8, dtype=torch.bfloat16) + lengths = torch.tensor([32], dtype=torch.int32) + output = torch.empty(1, 2, dtype=torch.int32) - _run_gvr_v2_decode(top_k, prior_indices) + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=torch.tensor([8], dtype=torch.int32), + max_seq_len=16, + ) - assert runner.call_args.args[1] is prior_indices + runner.assert_not_called() + decode.assert_called_once_with( + scores, + lengths, + output, + 1, + 2, + pre_idx=None, + heuristic_scratch=None, + compress_ratio=4, + radix_aux_indices=None, + radix_aux_logits=None, + ) def test_gvr_v2_decode_rejects_output_width_mismatch(monkeypatch) -> None: """Hint-free k derives from the output width; a scratch wider than top_k must be rejected before launch, not silently become the k.""" runner = _install_fake_selfsampling_runner(monkeypatch) - monkeypatch.delenv("TRTLLM_GVR_V2_HINTED", raising=False) top_k = TopK( 2, decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, compress_ratio=4, ) - prior_indices = torch.zeros(1, 2, dtype=torch.int32) - with pytest.raises(AssertionError): - _run_gvr_v2_decode(top_k, prior_indices, out_width=3) + _run_gvr_v2_decode(top_k, out_width=3) runner.assert_not_called() @@ -287,6 +311,21 @@ def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: ) assert prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] + assert top_k.needs_gvr_prior + + +def test_gvr_v2_does_not_update_prior_from_prefill() -> None: + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2) + prior_indices = torch.zeros(1, 2, dtype=torch.int32) + + top_k.update_gvr_prior_from_prefill( + torch.tensor([[4, 5]], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + prior_indices, + ) + + assert prior_indices.tolist() == [[0, 0]] + assert not top_k.needs_gvr_prior def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: 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 cec1473181de..ad3a2dadf5b5 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -228,26 +228,16 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): _check_exact(logits, indices, n_valid, ref_vals) -def _run_varlen_case( - kv, next_n, cr, top_k, seed, with_values=False, engine="auto", hint_free=False -): +def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): """Build a per-row-poisoned varlen batch, run run_varlen, verify every row against its own n_r (production formula) — short rows included.""" - batch, rows = len(kv), len(kv) * next_n + rows = len(kv) * next_n n_r = [(kv[r // next_n] - next_n + (r % next_n) + 1) // cr for r in range(rows)] npad = (max(n_r) + 63) // 64 * 64 gen = torch.Generator(device=_DEV).manual_seed(seed) logits = torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 for r in range(rows): logits[r, n_r[r] :] = 3e38 # poison beyond each row's OWN n_r - pre_idx = None - if not hint_free: - pre_idx = torch.empty((batch, top_k), dtype=torch.int32, device=_DEV) - for q in range(batch): - nmin = max(min(n_r[q * next_n : (q + 1) * next_n]), 1) - pre_idx[q] = torch.randint( - 0, nmin, (top_k,), generator=gen, dtype=torch.int32, device=_DEV - ) indices = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) values = ( torch.full((rows, top_k), 7.0, dtype=torch.float32, device=_DEV) if with_values else None @@ -255,13 +245,11 @@ def _run_varlen_case( kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) ss_host.run_varlen( logits, - pre_idx, kv_lens, indices, next_n=next_n, compress_ratio=cr, values=values, - engine=engine, ) torch.cuda.synchronize() fmin = torch.finfo(torch.float32).min @@ -287,7 +275,28 @@ def _run_varlen_case( assert torch.equal(values[r], torch.gather(logits[r], 0, idx)) -@pytest.mark.parametrize("engine", ["auto", "reference"]) +def _reference_varlen_indices(logits, kv_lens, next_n, compress_ratio, top_k): + """Build a simple torch reference for the hint-free varlen contract.""" + reference = torch.full((logits.shape[0], top_k), -1, dtype=torch.int32, device=logits.device) + lengths = kv_lens.tolist() + for row in range(logits.shape[0]): + valid = ( + max( + lengths[row // next_n] - next_n + row % next_n + 1, + 0, + ) + // compress_ratio + ) + valid = min(valid, logits.shape[1]) + if valid <= 0: + continue + if valid <= top_k: + reference[row, :valid] = torch.arange(valid, dtype=torch.int32, device=logits.device) + else: + reference[row] = torch.topk(logits[row, :valid], top_k).indices.to(torch.int32) + return reference + + @pytest.mark.parametrize( "kv,next_n,cr,top_k", [ @@ -300,76 +309,9 @@ def _run_varlen_case( ], ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4", "cr4_mtp3"], ) -def test_selfsampling_topk_varlen(kv, next_n, cr, top_k, engine): - """run_varlen production contract: per-row n from device kv_lens with the - MTP window formula, request-level hints, per-row short path — on BOTH the - per-row in-kernel engine ("auto") and the b=1 reference loop.""" - _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr, engine=engine) - - -@pytest.mark.parametrize( - "kv,next_n,cr,top_k", - [ - ([33000, 8200, 300], 1, 1, 512), - ([131075, 32800, 2000], 1, 4, 512), - ([9000, 5001], 2, 1, 512), - ([65540], 4, 4, 1024), - ([40000, 7003], 3, 4, 512), - ], - ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4", "cr4_mtp3"], -) -def test_selfsampling_topk_varlen_hint_free(kv, next_n, cr, top_k): - """pre_idx=None (hint-free): the bracket comes from the current row - itself; exactness must be identical to the hinted contract.""" - _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr, hint_free=True) - - -def test_selfsampling_topk_varlen_hint_free_reference_rejected(): - """The b=1 reference loop needs a hint tensor; pre_idx=None must raise.""" - kv_lens = torch.tensor([9000], dtype=torch.int32, device=_DEV) - logits = torch.randn((1, 9024), dtype=torch.float32, device=_DEV) - indices = torch.full((1, 512), -7, dtype=torch.int32, device=_DEV) - with pytest.raises(RuntimeError, match="auto-engine only"): - ss_host.run_varlen(logits, None, kv_lens, indices, engine="reference") - - -def test_selfsampling_topk_varlen_engine_matches_reference(): - """Differential: the in-kernel engine's per-row value multisets must - equal the reference loop's on a mixed batch (deep SPLIT rows, tsh band, - short rows, compressed space).""" - kv = [524288, 131075, 32800, 2000, 65540, 8192, 262144, 900] - top_k, cr = 1024, 4 - rows = len(kv) - n_r = [(v - 1 + 1) // cr for v in kv] - npad = (max(n_r) + 63) // 64 * 64 - gen = torch.Generator(device=_DEV).manual_seed(77) - logits = torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 - for r in range(rows): - logits[r, n_r[r] :] = 3e38 - pre_idx = torch.empty((rows, top_k), dtype=torch.int32, device=_DEV) - for q in range(rows): - pre_idx[q] = torch.randint( - 0, max(n_r[q], 1), (top_k,), generator=gen, dtype=torch.int32, device=_DEV - ) - kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) - out_a = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) - out_r = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(logits, pre_idx, kv_lens, out_a, compress_ratio=cr) - ss_host.run_varlen(logits, pre_idx, kv_lens, out_r, compress_ratio=cr, engine="reference") - torch.cuda.synchronize() - for r in range(rows): - if n_r[r] <= top_k: - assert torch.equal(out_a[r], out_r[r]) or torch.equal( - torch.sort(out_a[r]).values, torch.sort(out_r[r]).values - ) - else: - ga = torch.sort( - torch.gather(logits[r], 0, out_a[r].to(torch.int64)) + 0.0, descending=True - ).values - gr = torch.sort( - torch.gather(logits[r], 0, out_r[r].to(torch.int64)) + 0.0, descending=True - ).values - assert torch.equal(ga, gr), f"row {r}: engine != reference" +def test_selfsampling_topk_varlen(kv, next_n, cr, top_k): + """Validate the hint-free per-row KV-length and MTP window contract.""" + _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr) def test_selfsampling_topk_varlen_values(): @@ -393,90 +335,46 @@ def test_selfsampling_topk_varlen_launch_modes(rows, base, step, top_k): def test_selfsampling_topk_varlen_zero_kv_slot(): """Padded / evicted CUDA-graph request slots can carry kv_len < next_n - (even 0): both engines must emit the empty short row (all -1), not raise — + (even 0): the engine must emit the empty short row (all -1), not raise — mixed with live MTP rows of another request in the same launch.""" - for engine in ("auto", "reference"): - gen = torch.Generator(device=_DEV).manual_seed(9) - logits = torch.randn((8, 8192), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 - pre_idx = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) - kv_lens = torch.tensor([0, 8192], dtype=torch.int32, device=_DEV) - indices = torch.full((8, 512), -7, dtype=torch.int32, device=_DEV) - for r in range(4, 8): - n = 8192 - 4 + (r - 4) + 1 - logits[r, n:] = 3e38 - ss_host.run_varlen( - logits, pre_idx, kv_lens, indices, next_n=4, compress_ratio=1, engine=engine - ) - torch.cuda.synchronize() - assert bool((indices[:4] == -1).all()), f"{engine}: kv=0 rows must be all -1" - for r in range(4, 8): - n = 8192 - 4 + (r - 4) + 1 - idx = indices[r].to(torch.int64) - assert int(idx.min()) >= 0 and int(idx.max()) < n - ref = torch.topk(logits[r, :n], 512).values - got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values - assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values) - - -def test_selfsampling_topk_varlen_wide_buffers_flat_packed(): - """indices wider than k follow the CUDA flat-packed contract (rows at - stride k from the tensor base) IDENTICALLY on both engines.""" - kv = [9000, 300] - top_k, width = 512, 512 + 64 - gen = torch.Generator(device=_DEV).manual_seed(21) - logits = torch.randn((2, 9024), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 - logits[0, 9000:] = 3e38 - logits[1, 300:] = 3e38 - pre_idx = torch.randint(0, 300, (2, top_k), generator=gen, dtype=torch.int32, device=_DEV) - kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) - outs = {} - for engine in ("auto", "reference"): - wide = torch.full((2, width), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(logits, pre_idx, kv_lens, wide, compress_ratio=1, engine=engine) - torch.cuda.synchronize() - outs[engine] = wide.reshape(-1)[: 2 * top_k].view(2, top_k).clone() - packed_a, packed_r = outs["auto"], outs["reference"] - # row 0 (kernel row): same value multiset via the SAME packed convention - ga = torch.sort( - torch.gather(logits[0], 0, packed_a[0].to(torch.int64)) + 0.0, descending=True - ).values - gr = torch.sort( - torch.gather(logits[0], 0, packed_r[0].to(torch.int64)) + 0.0, descending=True - ).values - assert torch.equal(ga, gr), "wide-buffer packing convention diverged between engines" - # row 1 (short row): identity + -1 tail at the packed location, bit-equal - expect = torch.cat( - [ - torch.arange(300, dtype=torch.int32, device=_DEV), - torch.full((top_k - 300,), -1, dtype=torch.int32, device=_DEV), - ] - ) - assert torch.equal(torch.sort(packed_a[1, :300]).values, expect[:300]) - assert bool((packed_a[1, 300:] == -1).all()) - assert torch.equal(torch.sort(packed_r[1, :300]).values, expect[:300]) - assert bool((packed_r[1, 300:] == -1).all()) + gen = torch.Generator(device=_DEV).manual_seed(9) + logits = torch.randn((8, 8192), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + kv_lens = torch.tensor([0, 8192], dtype=torch.int32, device=_DEV) + indices = torch.full((8, 512), -7, dtype=torch.int32, device=_DEV) + for r in range(4, 8): + n = 8192 - 4 + (r - 4) + 1 + logits[r, n:] = 3e38 + ss_host.run_varlen(logits, kv_lens, indices, next_n=4, compress_ratio=1) + torch.cuda.synchronize() + assert bool((indices[:4] == -1).all()), "kv=0 rows must be all -1" + for r in range(4, 8): + n = 8192 - 4 + (r - 4) + 1 + idx = indices[r].to(torch.int64) + assert int(idx.min()) >= 0 and int(idx.max()) < n + ref = torch.topk(logits[r, :n], 512).values + got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values + assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values) def test_selfsampling_topk_varlen_guards(): logits = torch.randn((2, 8192), dtype=torch.float32, device=_DEV) - pre_idx = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) indices = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) kv = torch.tensor([8192, 8192], dtype=torch.int32, device=_DEV) with pytest.raises(RuntimeError, match="kv_lens length"): - ss_host.run_varlen(logits, pre_idx, kv[:1], indices) + ss_host.run_varlen(logits, kv[:1], indices) with pytest.raises(RuntimeError, match="not divisible"): - ss_host.run_varlen(logits, pre_idx, kv, indices, next_n=3) + ss_host.run_varlen(logits, kv, indices, next_n=3) with pytest.raises(RuntimeError, match="compress_ratio"): - ss_host.run_varlen(logits, pre_idx, kv, indices, compress_ratio=2) + ss_host.run_varlen(logits, kv, indices, compress_ratio=2) with pytest.raises(RuntimeError, match="CUDA tensor"): - ss_host.run_varlen(logits, pre_idx, kv.cpu(), indices) + ss_host.run_varlen(logits, kv.cpu(), indices) with pytest.raises(RuntimeError, match="num_rows"): # request-level-shaped indices under MTP: MUST be rejected (the # kernel grid comes from logits rows — silent OOB writes otherwise) - ss_host.run_varlen(logits, pre_idx[:1], kv[:1], indices[:1], next_n=2) + ss_host.run_varlen(logits, kv[:1], indices[:1], next_n=2) with pytest.raises(RuntimeError, match="contiguous"): strided = torch.zeros((2, 2), dtype=torch.int32, device=_DEV)[:, 0] - ss_host.run_varlen(logits, pre_idx, strided, indices) + ss_host.run_varlen(logits, strided, indices) def test_selfsampling_topk_varlen_cuda_graph(): @@ -488,7 +386,6 @@ def test_selfsampling_topk_varlen_cuda_graph(): rows, top_k, msl = 4, 512, 262144 npad = msl logits = torch.randn((rows, npad), dtype=torch.float32, device=_DEV) - 2.0 - pre_idx = torch.zeros((rows, top_k), dtype=torch.int32, device=_DEV) kv_lens = torch.tensor([100, 4099, 131070, 200000], dtype=torch.int32, device=_DEV) indices = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) @@ -502,17 +399,14 @@ def refresh(step): for r in range(rows): n = min(kv[r], npad) logits[r, n:] = 3e38 - pre_idx[r] = torch.randint( - 0, max(n, 1), (top_k,), generator=gen, dtype=torch.int32, device=_DEV - ) return kv refresh(0) - ss_host.run_varlen(logits, pre_idx, kv_lens, indices, compress_ratio=1, max_seq_len=msl) + ss_host.run_varlen(logits, kv_lens, indices, compress_ratio=1, max_seq_len=msl) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - ss_host.run_varlen(logits, pre_idx, kv_lens, indices, compress_ratio=1, max_seq_len=msl) + ss_host.run_varlen(logits, kv_lens, indices, compress_ratio=1, max_seq_len=msl) for step in range(1, 6): kv = refresh(step) indices.fill_(-7) @@ -635,10 +529,9 @@ def test_selfsampling_topk_varlen_rejects_non_fp32(): the loud contract message instead of a CuTe typing failure).""" logits = torch.randn(1, 8192, device=_DEV, dtype=torch.bfloat16) kv = torch.tensor([8000], dtype=torch.int32, device=_DEV) - pre = torch.zeros(1, 512, dtype=torch.int32, device=_DEV) out = torch.empty(1, 512, dtype=torch.int32, device=_DEV) with pytest.raises(RuntimeError, match="float32"): - ss_host.run_varlen(logits, pre, kv, out, next_n=1, compress_ratio=4, max_seq_len=32768) + ss_host.run_varlen(logits, kv, out, next_n=1, compress_ratio=4, max_seq_len=32768) def test_selfsampling_topk_varlen_zero_window_rows(): @@ -652,9 +545,8 @@ def test_selfsampling_topk_varlen_zero_window_rows(): rows = kv.numel() * nn npad = (msl // cr + 63) // 64 * 64 logits = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.zeros(kv.numel(), k, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(logits, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) + ss_host.run_varlen(logits, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) torch.cuda.synchronize() assert (out[nn:] == -1).all().item(), "n<=0 rows must be fully -1-padded" for r in range(nn): @@ -679,41 +571,10 @@ def test_selfsampling_warmup_row_stride_matches_arena(): arena = torch.randn(rows, stride, dtype=torch.float32, device=_DEV) logits = arena[:, :msl] # non-contiguous column slice, like serving kv = torch.full((rows,), msl, dtype=torch.int32, device=_DEV) - pre = torch.zeros(rows, k, dtype=torch.int32, device=_DEV) out = torch.empty(rows, k, dtype=torch.int32, device=_DEV) g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): - ss_host.run_varlen(logits, pre, kv, out, max_seq_len=msl) - g.replay() - torch.cuda.synchronize() - ref = torch.topk(arena[:, :msl], k, dim=1).values.sort(dim=1).values - got = arena.gather(1, out.long().clamp_min(0)).sort(dim=1).values - assert torch.equal(ref, got) - - -def test_selfsampling_warmup_hint_free_matches_dispatch(): - """warmup_varlen(hint_free=True) must compile the launcher key the - hint-free dispatch (pre_idx=None) looks up: the key includes the - hint-free bit, so a hinted warmup would leave capture uncompiled.""" - k, msl = 512, 8300 - stride = (msl + 255) // 256 * 256 - rows = 2 - ss_host.warmup_varlen( - k, - msl, - compress_ratio=1, - next_n=1, - num_rows_list=(rows,), - row_stride=stride, - hint_free=True, - ) - arena = torch.randn(rows, stride, dtype=torch.float32, device=_DEV) - logits = arena[:, :msl] - kv = torch.full((rows,), msl, dtype=torch.int32, device=_DEV) - out = torch.empty(rows, k, dtype=torch.int32, device=_DEV) - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - ss_host.run_varlen(logits, None, kv, out, max_seq_len=msl) + ss_host.run_varlen(logits, kv, out, max_seq_len=msl) g.replay() torch.cuda.synchronize() ref = torch.topk(arena[:, :msl], k, dim=1).values.sort(dim=1).values @@ -734,24 +595,13 @@ def test_selfsampling_varlen_regclus_parity_and_oracle(): rows = batch * nn torch.manual_seed(7) lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) kv = torch.tensor([msl_c * cr, 900, nn - 1], dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) - key = (rows, npad, k, msl_c, nn, cr, False) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + key = (rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg_clus", ss_host._VARLEN_CACHE[key][0] - ss_host.run_varlen( - lg, - pre, - kv, - ref, - next_n=nn, - compress_ratio=cr, - max_seq_len=msl_c * cr, - engine="reference", - ) torch.cuda.synchronize() + ref = _reference_varlen_indices(lg, kv, nn, cr, k) for r in range(rows): if (ref[r] >= 0).any(): row = lg[r].float() @@ -770,15 +620,14 @@ def test_selfsampling_varlen_regclus_cuda_graph(): rows = 8 torch.manual_seed(11) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() g = torch.cuda.CUDAGraph() out.fill_(-7) with torch.cuda.graph(g): - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values for _ in range(2): out.fill_(-7) @@ -811,24 +660,13 @@ def test_selfsampling_varlen_reg_parity_and_oracle(): assert fam == want, (fam, want, k, msl_c) msl = msl_c * cr lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) kv = torch.tensor([msl, max((k - 3) * cr, nn), nn - 1], dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) - key = (rows, npad, k, msl_c, nn, cr, False) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) + key = (rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] - ss_host.run_varlen( - lg, - pre, - kv, - ref, - next_n=nn, - compress_ratio=cr, - max_seq_len=msl, - engine="reference", - ) torch.cuda.synchronize() + ref = _reference_varlen_indices(lg, kv, nn, cr, k) for r in range(rows): if (ref[r] >= 0).any(): row = lg[r].float() @@ -859,27 +697,16 @@ def test_selfsampling_varlen_clus_parity_and_oracle(): assert plan["kernel"] == "clus" and plan["cluster"] == want_cs, plan batch = rows // nn lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) lens = [msl_c * cr, 900, nn - 1, 20000, 40000, 300000, msl_c * cr // 2, 5000] kv = torch.tensor( [lens[i % len(lens)] for i in range(batch)], dtype=torch.int32, device=_DEV ) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) - key = (rows, npad, k, msl_c, nn, cr, False) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + key = (rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] - ss_host.run_varlen( - lg, - pre, - kv, - ref, - next_n=nn, - compress_ratio=cr, - max_seq_len=msl_c * cr, - engine="reference", - ) torch.cuda.synchronize() + ref = _reference_varlen_indices(lg, kv, nn, cr, k) for r in range(rows): if (ref[r] >= 0).any(): row = lg[r].float() @@ -899,17 +726,16 @@ def test_selfsampling_varlen_clus_cuda_graph(): rows = 32 torch.manual_seed(29) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr, False) + key = (rows, msl_c, k, msl_c, 1, cr) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) with torch.cuda.graph(g): - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values for _ in range(2): out.fill_(-7) @@ -926,17 +752,16 @@ def test_selfsampling_varlen_reg_cuda_graph(): rows = 16 torch.manual_seed(17) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr, False) + key = (rows, msl_c, k, msl_c, 1, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) with torch.cuda.graph(g): - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values for _ in range(2): out.fill_(-7) @@ -955,12 +780,11 @@ def test_selfsampling_varlen_full_row_range(): torch.manual_seed(3) for rows in (304, 1024): lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr, False) + key = (rows, msl_c, k, msl_c, 1, cr) assert key in ss_host._VARLEN_CACHE, "row count must dispatch in-engine" ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values @@ -997,14 +821,13 @@ def test_selfsampling_varlen_heterogeneous_lengths_main(): rows = batch * nn # 304 -> route_streaming main family torch.manual_seed(5) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) kv = torch.randint(1, msl_c * cr, (batch,), dtype=torch.int32, device=_DEV) kv[0] = msl_c * cr # full length kv[1] = 900 # short (n <= k) kv[2] = nn - 1 # zero-window (every row of the request empty) kv[3] = k * cr + nn # just above the short path out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() kl = kv.tolist() for r in range(rows):