From 12766635237665f8ac7367800bc139cbdd2246e4 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:52:45 +0000 Subject: [PATCH 1/2] [None][fix] Self-sampling top-k host: physical row-width envelope and exact-row warmup population _varlen_launcher inflated the envelope to max(min(n_env, npad), k + 1) and passed it to the reg / reg_clus / clus launch tuples; the kernels clamp each row's valid length to that envelope before the short-path decision, so logits narrower than k + 1 with kv lengths beyond k read k + 1 elements at the physical row stride (into the next row, past the tensor for the last row). Only routing needs the k + 1 floor: the launch tuples now carry min(n_env, npad) and such rows take the short path. warmup_varlen keyed its done-set on the band representatives and returned before the exact-row launcher population, so a later call with a new row count inside an already-warmed band left that row count uncompiled and CUDA-graph capture at it raised. The done key now gates only the band launches; the exact-row population always runs. validate_run_ws now requires 16-byte base alignment, matching the DSL workspace fake. Neither defect is reachable through the TensorRT-LLM indexer path (its logits width, kv_lens and max_seq_len share one bound and graph launchers come from the eager warmup forwards); both hit direct callers of run_varlen / warmup_varlen (DKG issue #60). Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../gvr_topk_decode_self_sampling_host.py | 80 +++++++++++-------- .../parallel/test_gvr_selfsampling_topk.py | 57 +++++++++++++ 2 files changed, 102 insertions(+), 35 deletions(-) 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 cea15314e1f4..769452c366f5 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 @@ -695,7 +695,12 @@ def _varlen_launcher( hit = _VARLEN_CACHE.get(key) if hit is not None: return hit - n_eff = max(min(n_env, npad), k + 1) + # Two envelopes: the kernel gets the PHYSICAL bound (never past the row + # stride, so a row whose kv length exceeds the logits width clamps and + # takes the short path instead of reading into the next row); the router + # gets the k+1 floor it needs to pick a non-degenerate family. + n_kernel = min(n_env, npad) + n_route = max(n_kernel, k + 1) cr_shift = 0 if cr == 1 else 2 dev = _device() # ---- route() parity, family tier 1: clustered register-resident -------- @@ -703,7 +708,7 @@ def _varlen_launcher( # admission window (n4 <= 32768) fits capture-frozen envelopes. The # choice is a pure function of this cache key, so CUDA-graph replay # safety is unchanged; per-row n / short-row handling lives in-kernel. - plan_free = route(num_rows, n_eff, npad, k) + plan_free = route(num_rows, n_route, npad, k) if plan_free["kernel"] == "reg_clus": fn = dev.get_compiled__regclus( tuple(plan_free["tpl"]), @@ -712,7 +717,7 @@ def _varlen_launcher( cr_shift=cr_shift, hint_free=True, ) - lc = ("reg_clus", fn, n_eff) + lc = ("reg_clus", fn, n_kernel) _VARLEN_CACHE[key] = lc return lc # ---- route() parity, family tier 2: register-resident (+img flavor) ---- @@ -734,7 +739,7 @@ def _varlen_launcher( lc = ( "reg", fn, - (rt_f["n"], rt_f["CMP"], rt_f["QC"], dev.STATIC_BYTES + plan_free["smem"]), + (n_kernel, rt_f["CMP"], rt_f["QC"], dev.STATIC_BYTES + plan_free["smem"]), ) _VARLEN_CACHE[key] = lc return lc @@ -759,11 +764,11 @@ def _varlen_launcher( lc = ( "clus", fn, - (n_eff, npad, k, rt_f["SCAP"], rt_f["CMP"], 0, 0, 0, 0, 0), + (n_kernel, npad, k, rt_f["SCAP"], rt_f["CMP"], 0, 0, 0, 0, 0), ) _VARLEN_CACHE[key] = lc return lc - plan = route_streaming(num_rows, n_eff, npad, k, force_main=True) + plan = route_streaming(num_rows, n_route, npad, k, force_main=True) tpl = tuple(plan["tpl"]) # (BLK, U, MINB, SNB, KPT, SPLIT, TSHG) rt = plan["rt"] r_const = rt["R"] @@ -894,13 +899,13 @@ def default_workspace(ref: torch.Tensor) -> torch.Tensor: def validate_run_ws(workspace: torch.Tensor, logits: torch.Tensor) -> None: """run_ws() workspace hardening, in a fixed predicate order: CUDA + same device as logits; numel*element_size >= workspace_bytes(); - base 8-byte aligned.""" + base 16-byte aligned (the DSL workspace fake declares assumed_align=16).""" if not (workspace.is_cuda and workspace.get_device() == logits.get_device()): raise RuntimeError("workspace must be a CUDA tensor on the same device") if workspace.numel() * workspace.element_size() < WS_BYTES: raise RuntimeError(f"workspace too small: need {WS_BYTES} bytes") - if workspace.data_ptr() & 7: - raise RuntimeError("workspace must be 8-byte aligned") + if workspace.data_ptr() & 15: + raise RuntimeError("workspace must be 16-byte aligned") def kernel_view(workspace: torch.Tensor) -> torch.Tensor: @@ -908,10 +913,9 @@ def kernel_view(workspace: torch.Tensor) -> torch.Tensor: bytes at the tensor's data_ptr() as int32[WS_BYTES/4], ignoring dtype/shape. - NOTE: the DSL-side fake tensor declares assumed_align=16; a workspace at - 8-but-not-16-byte alignment passes the validate_run_ws check but is - rejected by the DSL at conversion -- surfaced as a launch failure with - shape context.""" + NOTE: the DSL-side fake tensor declares assumed_align=16, matching the + validate_run_ws base-alignment check, so misaligned workspaces fail on + the host with a clear message instead of at DSL conversion.""" if ( workspace.dtype is torch.int32 and workspace.dim() == 1 @@ -1540,27 +1544,32 @@ def warmup_varlen( tuple(rows_list), npad, ) + # The done key covers the GPU band launches only (one per engine compile + # key). The exact-row launcher population below is keyed by the requested + # row counts, which the band key does not see, so it always runs: a later + # call with a new row count inside an already-warmed band must still + # create that row count's entry, or capture at it raises not-compiled. with _VARLEN_WARMUP_LOCK: - if key in _VARLEN_WARMUP_DONE: - return - rows_max = rows_list[-1] - # one allocation at the largest geometry; smaller row counts run on - # 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) - 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], - kv_lens[:batch], - out[:rows], - next_n=nn, - compress_ratio=int(compress_ratio), - max_seq_len=int(max_seq_len), - ) - del logits, kv_lens, out - torch.cuda.synchronize() + bands_done = key in _VARLEN_WARMUP_DONE + if not bands_done: + rows_max = rows_list[-1] + # one allocation at the largest geometry; smaller row counts run on + # 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) + 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], + kv_lens[:batch], + out[:rows], + next_n=nn, + compress_ratio=int(compress_ratio), + max_seq_len=int(max_seq_len), + ) + 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 # work, zero allocation/launch — engines hit the compile cache), so a @@ -1568,5 +1577,6 @@ def warmup_varlen( 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)) - with _VARLEN_WARMUP_LOCK: - _VARLEN_WARMUP_DONE.add(key) + if not bands_done: + with _VARLEN_WARMUP_LOCK: + _VARLEN_WARMUP_DONE.add(key) 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 3408dd325fd6..9d6a3758fd67 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -951,3 +951,60 @@ def test_selfsampling_varlen_heterogeneous_lengths_main(): ref_v = torch.topk(row, k).values.sort().values got = row[out[r].long()].sort().values assert torch.equal(got, ref_v), f"row {r}: value multiset mismatch (n={n_r})" + + +def test_selfsampling_varlen_narrow_logits_clamps_envelope_to_row_width(): + """Logits narrower than k + 1 with kv lengths beyond k: the launch + envelope must be the physical row width, so every row takes the in-kernel + short path (identity + -1 tail) instead of reading k + 1 elements across + the row stride (DKG issue #60, item 1).""" + rows, n, k = 2, 256, 512 + lg = torch.randn(rows, n, dtype=torch.float32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + lc = ss_host._varlen_launcher(rows, n, k, n, 1, 1) + if lc[0] != "main": # main derives the envelope from npad in-kernel + env = lc[2] if not isinstance(lc[2], tuple) else lc[2][0] + assert env == n, (lc[0], env) + want = torch.full((rows, k), -1, dtype=torch.int32, device=_DEV) + want[:, :n] = torch.arange(n, dtype=torch.int32, device=_DEV) + for kv_len in (100000, n): # kv far beyond the row width, and exactly the row width + kv = torch.full((rows,), kv_len, dtype=torch.int32, device=_DEV) + out.fill_(-7) + ss_host.run_varlen(lg, kv, out, max_seq_len=n) + torch.cuda.synchronize() + assert torch.equal(out, want), kv_len + + +def test_selfsampling_warmup_incremental_rows_populate_exact_launchers(): + """A later warmup_varlen call whose new row count falls inside an already + warmed engine band must still create that row count's launcher entry, so a + CUDA-graph capture at it finds its key (DKG issue #60, item 2).""" + k, msl, cr, nn = 512, 8192, 1, 1 + npad = (msl + 63) // 64 * 64 + ss_host.warmup_varlen(k, msl, compress_ratio=cr, next_n=nn, num_rows_list=(64,)) + key63 = (63, npad, k, min(msl, npad), nn, cr) + ss_host._VARLEN_CACHE.pop(key63, None) # independent of earlier tests in this process + ss_host.warmup_varlen(k, msl, compress_ratio=cr, next_n=nn, num_rows_list=(63, 64)) + assert key63 in ss_host._VARLEN_CACHE + lg = torch.randn(63, npad, dtype=torch.float32, device=_DEV) + kv = torch.full((63,), msl, dtype=torch.int32, device=_DEV) + out = torch.empty(63, k, dtype=torch.int32, device=_DEV) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + ss_host.run_varlen(lg[:, :msl], kv, out, max_seq_len=msl) + g.replay() + torch.cuda.synchronize() + ref = torch.topk(lg[:, :msl], k, dim=1).values.sort(dim=1).values + got = lg.gather(1, out.long().clamp_min(0)).sort(dim=1).values + assert torch.equal(ref, got) + + +def test_validate_run_ws_requires_16_byte_alignment(): + """The workspace fake declares 16-byte alignment; the host check matches + it so an 8-but-not-16-byte workspace fails here instead of at DSL + conversion (DKG issue #60, item 3).""" + logits = torch.zeros(1, 64, dtype=torch.float32, device=_DEV) + base = torch.zeros(ss_host.WS_BYTES // 4 + 4, dtype=torch.int32, device=_DEV) + ss_host.validate_run_ws(base[4:], logits) # +16 B + with pytest.raises(RuntimeError, match="16-byte"): + ss_host.validate_run_ws(base[2:], logits) # +8 B From eed52e18dca2e093ccdb66d64eb68d0425c04037 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 21:04:39 +0000 Subject: [PATCH 2/2] [None][test] Annotate the issue #60 regression tests Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/thop/parallel/test_gvr_selfsampling_topk.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 9d6a3758fd67..be0db369a452 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -953,7 +953,7 @@ def test_selfsampling_varlen_heterogeneous_lengths_main(): assert torch.equal(got, ref_v), f"row {r}: value multiset mismatch (n={n_r})" -def test_selfsampling_varlen_narrow_logits_clamps_envelope_to_row_width(): +def test_selfsampling_varlen_narrow_logits_clamps_envelope_to_row_width() -> None: """Logits narrower than k + 1 with kv lengths beyond k: the launch envelope must be the physical row width, so every row takes the in-kernel short path (identity + -1 tail) instead of reading k + 1 elements across @@ -975,7 +975,7 @@ def test_selfsampling_varlen_narrow_logits_clamps_envelope_to_row_width(): assert torch.equal(out, want), kv_len -def test_selfsampling_warmup_incremental_rows_populate_exact_launchers(): +def test_selfsampling_warmup_incremental_rows_populate_exact_launchers() -> None: """A later warmup_varlen call whose new row count falls inside an already warmed engine band must still create that row count's launcher entry, so a CUDA-graph capture at it finds its key (DKG issue #60, item 2).""" @@ -999,7 +999,7 @@ def test_selfsampling_warmup_incremental_rows_populate_exact_launchers(): assert torch.equal(ref, got) -def test_validate_run_ws_requires_16_byte_alignment(): +def test_validate_run_ws_requires_16_byte_alignment() -> None: """The workspace fake declares 16-byte alignment; the host check matches it so an 8-but-not-16-byte workspace fails here instead of at DSL conversion (DKG issue #60, item 3)."""