[None][fix] Self-sampling top-k host: physical row-width envelope and exact-row warmup population

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py

@@ -695,15 +695,20 @@ 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 --------
     # Admit reg_clus exactly where the free route picks it; its whole
     # 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,24 +899,23 @@ 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:
     """Raw-pointer view of a user workspace tensor: alias the first WS_BYTES
     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,33 +1544,39 @@ 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
     # 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))
-    with _VARLEN_WARMUP_LOCK:
-        _VARLEN_WARMUP_DONE.add(key)
+    if not bands_done:
+        with _VARLEN_WARMUP_LOCK:
+            _VARLEN_WARMUP_DONE.add(key)

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() -> 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
+    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() -> 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)."""
+    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() -> 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)."""
+    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