[None][feat] Self-sampling GVR V2 prefill indexer top-K

tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py

@@ -754,9 +754,19 @@ def __init__(
                 if self.use_cute_dsl_topk
                 else TopKImplementation.CUDA_RADIX
             )
+        # Prefill uses the self-sampling engine on exactly the layers where the
+        # decode dispatch picks it; the temporal-hint engine has no prefill form.
+        prefill_top_k_implementation = (
+            TopKImplementation.CUTE_DSL_GVR
+            if (
+                decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR
+                and self._use_self_sampling_topk
+            )
+            else TopKImplementation.CUDA_RADIX
+        )
         self.top_k = TopK(
             self.index_topk,
-            prefill_implementation=TopKImplementation.CUDA_RADIX,
+            prefill_implementation=prefill_top_k_implementation,
             decode_implementation=decode_top_k_implementation,
             compress_ratio=self.compress_ratio,
             gvr_self_sampling=self._use_self_sampling_topk,

tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py

@@ -437,6 +437,15 @@ def warmup_selfsampling_topk(
         for bs in batch_sizes or ():
             rows.add(int(bs) * nn)
         msl_c = int(self.get_indexer_max_seq_len())
+        # Prefill leg first: the decode-stride guard below may return early for
+        # an odd msl_c, but the DeepGEMM prefill stride is always 256-aligned.
+        try:
+            _ss_host.warmup_prefill(int(top_k), max(msl_c, 32768))
+        except torch.cuda.OutOfMemoryError:
+            logger.warning(
+                "self-sampling GVR prefill warmup ran out of memory; prefill "
+                "engines will JIT-compile lazily on first touch instead."
+            )
         if self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits:
             # mirror the DSL paged-MQA arena stride (rows round up to 256
             # elements). A drift here only degrades warmup to unused keys —

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

@@ -20,6 +20,8 @@
 from .gvr_topk_decode_direct import DirectTopKKernel
 from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk
 from .gvr_topk_decode_reg import GvrRegKernel
+from .gvr_topk_decode_self_sampling_host import prefill_ready as selfsampling_topk_prefill_ready
+from .gvr_topk_decode_self_sampling_host import run_prefill as selfsampling_topk_run_prefill
 from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen
 from .gvr_topk_decode_tp import GvrTpKernel
 from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel
@@ -36,4 +38,6 @@
     "tiered_topk",
     "is_tiered_topk_supported",
     "selfsampling_topk_run_varlen",
+    "selfsampling_topk_run_prefill",
+    "selfsampling_topk_prefill_ready",
 ]

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

@@ -678,6 +678,63 @@ def _main(blk_, minb_, u_, split_):
 
 _VARLEN_CACHE = {}
 
+# ---- prefill launcher cache ------------------------------------------------
+# Prefill forces R==1 (route_streaming gives R>1 only for b<=74). The compiled
+# launcher depends only on the row tier, k and the envelope bucket — never on the
+# exact row count or npad — so the cache stays bounded on a long-running server.
+_PREFILL_CACHE = {}
+_PREFILL_ROW_SLAB = 32768  # gridDim.y <= 65535; slab so keys stay bounded
+_PREFILL_TIER_ROWS = (75, 149, 297)  # (rows<=148, 149..296, >296) band reps
+
+
+def _prefill_tier(rows: int) -> int:
+    return 0 if rows <= 148 else 1 if rows <= 296 else 2
+
+
+def _prefill_bucket(n_env: int) -> int:
+    # pow2-quantize the envelope so a growing envelope reuses one plan; cap at
+    # 32768 because U=8 for every n>=32768 on the tier-0 arm.
+    return min(1 << max(int(n_env) - 1, 1).bit_length(), 32768)
+
+
+def _prefill_cache_key(tier: int, k: int, n_bucket: int):
+    # tiers 1/2 fix U, so the bucket does not change their engine — collapse it
+    # to one key so warmup covers them with a single launch.
+    return (tier, k, n_bucket if tier == 0 else 0)
+
+
+def _prefill_launcher(tier: int, k: int, n_bucket: int) -> tuple:
+    """Prefill plan + compiled launcher: ``_varlen_launcher``'s main branch with
+    r_const=1, split=False and the prefill compile flag. SCAP_/CMP_ are envelope
+    upper bounds; npad is filled per call in ``run_prefill``."""
+    key = _prefill_cache_key(tier, k, n_bucket)
+    hit = _PREFILL_CACHE.get(key)
+    if hit is not None:
+        return hit
+    b_route = _PREFILL_TIER_ROWS[tier]
+    n_route = max(n_bucket, k + 1)
+    plan = route_streaming(b_route, n_route, n_route, k, force_main=True)
+    if plan["kernel"] != "main":
+        raise RuntimeError(f"prefill route did not land on gvr_main: {plan['kernel']}")
+    rt = plan["rt"]
+    if rt["R"] != 1:
+        raise RuntimeError(f"prefill requires R==1 (got {rt['R']})")
+    tpl = tuple(plan["tpl"])
+    dev = _device()
+    fn = dev.get_compiled(tpl[:6] + (False,) + (1, 0, 1), hint_free=True, prefill=True)
+    big = tier == 0
+    # r_const==1 branch of the _varlen_launcher tuning scalars
+    aim_base = (
+        (4 * k if k >= 1024 else 2 * k) if big else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2)
+    )
+    sfac = 64 if k >= 1024 else 32
+    amin = (7 * k) // 2
+    sd_en = 1 if (k > 1024 and not big) else 0
+    tail = (aim_base, sfac, amin, sd_en, 0)  # tsh_en=0 (split=False)
+    lc = ("main", fn, (rt["SCAP_"], rt["CMP_"]), tail)
+    _PREFILL_CACHE[key] = lc
+    return lc
+
 
 def _varlen_launcher(
     num_rows: int,
@@ -1435,6 +1492,119 @@ def run_varlen(
     return
 
 
+def run_prefill(
+    logits: torch.Tensor,
+    row_starts: torch.Tensor,
+    row_ends: torch.Tensor,
+    indices: torch.Tensor,
+    max_row_len: int | None = None,
+    workspace: torch.Tensor | None = None,
+) -> None:
+    """Hint-free self-sampling Top-K for prefill: row ``r`` selects the Top-K of
+    ``logits[r, ks:ke]`` (compressed columns) into the local frame (column - ks)
+    with a -1 pad; ``nv <= k`` rows get the identity, as ``indexer_topk_prefill``.
+    No device reads, never compiles under capture; trusts 0 <= ks <= ke <= shape[1]."""
+    if logits.dtype is not _F32:
+        raise RuntimeError(
+            f"logits must be float32 (got {logits.dtype}); bf16/fp16 paths "
+            "are a follow-up — see the PR roadmap"
+        )
+    for _nm, _t in (("row_starts", row_starts), ("row_ends", row_ends)):
+        if not (isinstance(_t, _TENSOR) and _t.is_cuda):
+            raise RuntimeError(f"{_nm} must be a CUDA tensor")
+        if _t.dtype is not _I32:
+            raise RuntimeError(f"{_nm} must be int32")
+        if _t.dim() != 1:
+            raise RuntimeError(f"{_nm} must be 1-D")
+        if not _t.is_contiguous():
+            raise RuntimeError(f"{_nm} must be contiguous")
+    if len(logits.shape) != 2:
+        raise RuntimeError("logits must be 2-D")
+    num_rows = logits.shape[0]
+    if num_rows == 0:
+        return
+    if row_starts.shape[0] != num_rows or row_ends.shape[0] != num_rows:
+        raise RuntimeError(
+            f"row_starts/row_ends length must equal logits.shape[0]={num_rows}, "
+            f"got {row_starts.shape[0]}/{row_ends.shape[0]}"
+        )
+    if not (logits.is_cuda and indices.is_cuda):
+        raise RuntimeError("all tensors must be CUDA")
+    if indices.dtype is not _I32:
+        raise RuntimeError("indices must be int32")
+    if len(indices.shape) != 2 or indices.shape[0] != num_rows:
+        raise RuntimeError(f"indices must be [num_rows={num_rows}, k], got {tuple(indices.shape)}")
+    if not indices.is_contiguous():
+        raise RuntimeError("indices must be contiguous")
+    k = indices.shape[1]
+    if k < 4 or (k & 3):
+        raise RuntimeError(f"index_topk must be a multiple of 4 and >= 4, got {k}")
+    if indices.data_ptr() & 15:
+        raise RuntimeError("indices base must be 16-byte aligned")
+    if logits.stride(1) != 1:
+        raise RuntimeError("logits inner stride must be 1")
+    # key on stride(0) for every row count: DeepGEMM prefill rows are 1024B-aligned
+    # with slack, and the varlen 1-row shape[1] rule would reject odd-width tiles.
+    npad = logits.stride(0)
+    if npad & 3:
+        raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}")
+    if logits.data_ptr() & 15:
+        raise RuntimeError("logits base must be 16-byte aligned")
+    d = logits.get_device()
+    if not 0 <= d < _GVR_MAX_DEV:
+        raise RuntimeError(f"device index out of range: {d}")
+    lg = logits
+    if logits.shape[1] != npad:
+        need = logits.storage_offset() + num_rows * npad
+        if logits.untyped_storage().size() // 4 < need:
+            raise RuntimeError("logits view storage too small to widen to its row stride")
+        lg = logits.as_strided((num_rows, npad), (npad, 1), logits.storage_offset())
+    if workspace is not None:
+        validate_run_ws(workspace, logits)
+        ws = kernel_view(workspace)
+    else:
+        ws = _ws_hot.get(d)
+        if ws is None:
+            ws = default_workspace(logits)
+    n_env = _index(max_row_len) if max_row_len is not None else logits.shape[1]
+    n_env = min(max(n_env, 1), npad)
+    n_bucket = _prefill_bucket(n_env)
+    for r0 in range(0, num_rows, _PREFILL_ROW_SLAB):
+        r1 = min(r0 + _PREFILL_ROW_SLAB, num_rows)
+        tier = _prefill_tier(r1 - r0)
+        lc = _PREFILL_CACHE.get(_prefill_cache_key(tier, k, n_bucket))
+        if lc is None:
+            if _is_capturing():
+                raise RuntimeError(
+                    "prefill launcher not compiled for this shape — warm up "
+                    "before CUDA graph capture"
+                )
+            lc = _prefill_launcher(tier, k, n_bucket)
+        _, fn, (scap, cmp_), tail = lc
+        # varlen main ABI: pre_idx slot = row_ends, kv_lens slot = row_starts;
+        # only npad / k / SCAP_ / CMP_ matter (R=1), the other scalars are dead.
+        pre = (0, npad, k, scap, cmp_, 1, 0, 0, 0, 0, 0)
+        fn(lg[r0:r1], row_ends[r0:r1], indices[r0:r1], ws, *pre, row_starts[r0:r1], *tail)
+    return
+
+
+def prefill_ready(logits: torch.Tensor, indices: torch.Tensor) -> bool:
+    """True iff ``run_prefill(logits, ..., indices)`` would launch without
+    compiling — the same (tier, k, envelope bucket) keys it looks up, so a
+    caller can route around the engine under CUDA graph capture. Host-only."""
+    num_rows = logits.shape[0]
+    if num_rows == 0:
+        return True
+    k = indices.shape[1]
+    npad = logits.stride(0)
+    n_bucket = _prefill_bucket(min(max(logits.shape[1], 1), max(npad, 1)))
+    for r0 in range(0, num_rows, _PREFILL_ROW_SLAB):
+        tier = _prefill_tier(min(r0 + _PREFILL_ROW_SLAB, num_rows) - r0)
+        if _prefill_cache_key(tier, k, n_bucket) not in _PREFILL_CACHE:
+            return False
+    return True
+
+
 __all__ = [
     "route",
     "route_static",
@@ -1444,7 +1614,10 @@ def run_varlen(
     "run",
     "run_ws",
     "run_varlen",
+    "run_prefill",
+    "prefill_ready",
     "warmup_varlen",
+    "warmup_prefill",
     "workspace_bytes",
     "WS_BYTES",
     "default_workspace",
@@ -1580,3 +1753,54 @@ def warmup_varlen(
     if not bands_done:
         with _VARLEN_WARMUP_LOCK:
             _VARLEN_WARMUP_DONE.add(key)
+
+
+_PREFILL_WARMUP_DONE: set = set()
+_PREFILL_WARMUP_LOCK = threading.Lock()
+
+
+def warmup_prefill(
+    top_k: int,
+    max_cols: int,
+    num_rows_list: Sequence[int] = (1, 149, 297),
+    row_stride: int | None = None,
+) -> None:
+    """Compile the prefill engine set before serving (<=6 per k): the tier-0 arm
+    walks the pow2 envelope buckets up to 32768, tiers 1/2 need one launch each.
+    ``max_cols`` is the compressed max column count; idempotent per done-key."""
+    dev = torch.cuda.current_device()
+    k = int(top_k)
+    max_cols = int(max_cols)
+    lo = _prefill_bucket(k + 1)
+    hi = _prefill_bucket(max_cols)
+    buckets = []
+    b = lo
+    while b <= hi:
+        buckets.append(b)
+        b <<= 1
+    if not buckets:
+        buckets = [hi]
+    keys = {}  # cache_key -> (tier, bucket) representative for the launch
+    for rows in num_rows_list:
+        tier = _prefill_tier(int(rows))
+        bset = buckets if tier == 0 else buckets[:1]
+        for bk in bset:
+            keys.setdefault(_prefill_cache_key(tier, k, bk), (tier, bk))
+    done_key = (dev, k, max_cols, tuple(sorted(int(r) for r in num_rows_list)), row_stride)
+    with _PREFILL_WARMUP_LOCK:
+        if done_key in _PREFILL_WARMUP_DONE:
+            return
+    for tier, bk in keys.values():
+        rows = _PREFILL_TIER_ROWS[tier]
+        stride = row_stride if row_stride is not None else ((bk + 256 + 255) // 256 * 256)
+        if stride < bk or stride % 4:
+            stride = (max(stride, bk) + 256 + 255) // 256 * 256
+        logits = torch.zeros((rows, stride), dtype=torch.float32, device=dev)
+        ks = torch.zeros((rows,), dtype=torch.int32, device=dev)
+        ke = torch.full((rows,), bk, dtype=torch.int32, device=dev)
+        out = torch.empty((rows, k), dtype=torch.int32, device=dev)
+        run_prefill(logits[:, :bk], ks, ke, out, max_row_len=bk)
+        del logits, ks, ke, out
+    torch.cuda.synchronize()
+    with _PREFILL_WARMUP_LOCK:
+        _PREFILL_WARMUP_DONE.add(done_key)

tensorrt_llm/_torch/modules/top_k.py

@@ -143,7 +143,59 @@ def _forward_prefill(
                 row_ends,
                 output_indices,
             )
-        if self.prefill_implementation == TopKImplementation.CUTE_DSL_RADIX:
+        if self.prefill_implementation == TopKImplementation.CUTE_DSL_GVR:
+            # hint-free k derives from the output width; pin it to the module's k
+            assert output_indices.shape[1] == self.top_k
+            if not self.gvr_self_sampling:
+                # the temporal (hint) GVR engine has no prefill form
+                logger.warning_once(
+                    "temporal GVR has no prefill engine; using the CUDA radix prefill Top-K.",
+                    key="gvr_temporal_prefill_radix",
+                )
+            elif scores.shape[1] <= self.top_k:
+                # every row is short (nv <= k): the exact radix path emits the
+                # identity/-1 answer without reading logits — cheaper than a
+                # zero-work self-sampling launch. Deliberate, no warning.
+                pass
+            elif self._selfsampling_prefill_ok(scores):
+                from ..cute_dsl_kernels.blackwell.top_k import (
+                    selfsampling_topk_prefill_ready,
+                    selfsampling_topk_run_prefill,
+                )
+
+                if self._prefill_capturing(scores) and not selfsampling_topk_prefill_ready(
+                    scores, output_indices
+                ):
+                    # the engine never JIT-compiles under capture; an engine
+                    # missed by warmup takes the exact radix path in the graph
+                    logger.warning_once(
+                        "self-sampling GVR prefill engine is not compiled for this "
+                        "shape and cannot JIT under CUDA graph capture; using the "
+                        "CUDA radix prefill Top-K.",
+                        key="selfsampling_topk_prefill_capture_radix",
+                    )
+                else:
+                    logger.info_once(
+                        "self-sampling GVR prefill top-K engaged "
+                        f"(K={self.top_k}, cr={self.compress_ratio}, hint-free).",
+                        key="selfsampling_topk_prefill_engaged",
+                    )
+                    # ks/ke are already in compressed column units; run_prefill
+                    # writes the local (column - ks) frame with -1 pad and no
+                    # host reads (envelope from scores.shape[1]).
+                    selfsampling_topk_run_prefill(scores, row_starts, row_ends, output_indices)
+                    return output_indices
+            else:
+                # engine hardware-format gate missed (e.g. a non-fp4 layer with
+                # an odd DeepGEMM width, or a bf16 producer): exact radix.
+                logger.warning_once(
+                    "self-sampling GVR prefill is selected but the scores do "
+                    "not satisfy the engine's hardware-format gate "
+                    f"(dtype={scores.dtype}, strides={tuple(scores.stride())}); "
+                    "falling back to the CUDA radix prefill Top-K.",
+                    key="selfsampling_topk_prefill_fallthrough",
+                )
+        elif self.prefill_implementation == TopKImplementation.CUTE_DSL_RADIX:
             # Keep the op's reread policy default; only its copy width is tuned.
             torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell(
                 scores,
@@ -154,7 +206,7 @@ def _forward_prefill(
                 _CUTE_DSL_PREFILL_COPY_BITS,
             )
             return output_indices
-        if self.prefill_implementation != TopKImplementation.CUDA_RADIX:
+        elif self.prefill_implementation != TopKImplementation.CUDA_RADIX:
             raise NotImplementedError(
                 f"{self.prefill_implementation.value} does not support prefill Top-K"
             )
@@ -167,6 +219,20 @@ def _forward_prefill(
         )
         return output_indices
 
+    def _selfsampling_prefill_ok(self, scores: torch.Tensor) -> bool:
+        """Engine format gate: fp32 row-major scores with a float4-aligned row
+        stride and a 16B base (the DeepGEMM prefill logits arena)."""
+        return (
+            scores.dtype == torch.float32
+            and scores.stride(1) == 1
+            and scores.stride(0) % 4 == 0
+            and scores.data_ptr() % 16 == 0
+        )
+
+    @staticmethod
+    def _prefill_capturing(scores: torch.Tensor) -> bool:
+        return scores.is_cuda and torch.cuda.is_current_stream_capturing()
+
     def _forward_decode(
         self,
         scores: torch.Tensor,