[None][feat] DSv4: enable GVR Heuristic Top-K for compress_ratio=4

cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -47,15 +47,19 @@ using heuristic_topk::KernelSmemTplK;
 // same kernel template. Smem layout is derived from GvrParams<float, TopK>
 // at compile time.
 template <int TopK>
-__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float const* __restrict__ logits,
-    int const* __restrict__ seqLens, int const* __restrict__ preIdx, float* __restrict__ scratchValues,
-    int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount)
+__global__ void __launch_bounds__(BLOCK_SIZE)
+    heuristicTopKMultiRowKernel(float const* __restrict__ logits, int const* __restrict__ seqLens,
+        int const* __restrict__ preIdx, float* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0,
+        int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio)
 {
     using SmemT = KernelSmemTplK<float, GvrParams<float, TopK>::kC, GvrParams<float, TopK>::kNumBins>;
 
     int const rowIdx = blockIdx.x;
     int const seq_len = seqLens[rowIdx / next_n];
-    int const N = seq_len - next_n + (rowIdx % next_n) + 1;
+    // seqLens is in uncompressed token space; the logits/preIdx live in
+    // compressed-index space when compressRatio > 1 (DSv4 indexer).
+    int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1;
+    int const N = actual_kv_len / compressRatio;
 
     float const* __restrict__ input = logits + static_cast<int64_t>(rowIdx) * stride0;
     int const* __restrict__ rowPreIdx = preIdx + static_cast<int64_t>(rowIdx / next_n) * preIdxStride;
@@ -84,9 +88,19 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float
         return;
     }
 
-    // +1 accounts for the temporal shift: prev_topk indices were computed at
-    // seq_len-1, but the current step has one additional KV token appended.
-    int const preIdxOffset = (rowIdx % next_n) + 1;
+    // Temporal-shift offset to map prev-step's top-K indices into this step's
+    // KV index space.
+    //   compressRatio == 1 (DSv3.2):   +1 — KV grew by exactly 1 token per
+    //     decode step; prev indices were at seq_len-1 so a uniform +1 maps
+    //     them to the equivalent positions under the indexer's "newest-first"
+    //     layout. The (rowIdx % next_n) addend extends this to MTP windows.
+    //   compressRatio == 4 (DSv4):     0 — in compressed-index space new
+    //     compressed entries are appended at the end; prev indices in
+    //     [0, c_prev-1] remain valid as-is. Per-row Δc varies (0 or 1) with
+    //     prev kv_len mod 4 alignment, but a uniform offset of 0 stays
+    //     within-bounds for all rows and preserves the temporal-correlation
+    //     hint (vertical top-K consistency validated offline).
+    int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0;
     gvrTopKJob<TopK>(input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset);
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     cudaTriggerProgrammaticLaunchCompletion();
@@ -103,16 +117,18 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float
 // Templated on (InputT, TopK). Smem layout is derived from
 // GvrParams<InputT, TopK>.
 template <typename InputT, int TopK>
-__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits,
-    int const* __restrict__ seqLens, int const* __restrict__ preIdx, InputT* __restrict__ scratchValues,
-    int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount)
+__global__ void __launch_bounds__(BLOCK_SIZE)
+    heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, int const* __restrict__ seqLens,
+        int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0,
+        int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio)
 {
     // dtype path uses fp32 keys[] in smem (down-conversion deferred to writeback).
     using SmemT = KernelSmemTplK<float, GvrParams<InputT, TopK>::kC, GvrParams<InputT, TopK>::kNumBins>;
 
     int const rowIdx = blockIdx.x;
     int const seq_len = seqLens[rowIdx / next_n];
-    int const N = seq_len - next_n + (rowIdx % next_n) + 1;
+    int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1;
+    int const N = actual_kv_len / compressRatio;
 
     InputT const* __restrict__ input = logits + static_cast<int64_t>(rowIdx) * stride0;
     int const* __restrict__ rowPreIdx = preIdx + static_cast<int64_t>(rowIdx / next_n) * preIdxStride;
@@ -142,7 +158,8 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(I
         return;
     }
 
-    int const preIdxOffset = (rowIdx % next_n) + 1;
+    // See fp32 path: cr==1 → (rowIdx % next_n)+1; cr!=1 (DSv4) → 0.
+    int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0;
     gvrTopKJobDtype<InputT, TopK>(
         input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset);
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
@@ -152,24 +169,25 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(I
 
 // Explicit instantiations — 6 (dtype × K) combos. Launchers dispatch on
 // runtime topK via switch, so all 6 must be available at link time.
+// Trailing `int` is the compressRatio parameter (1 = V3.2, 4 = V4 indexer).
 template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 512>(
-    __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int);
+    __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 1024>(
-    __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int);
+    __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 2048>(
-    __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int);
+    __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernelDtype<__half, 512>(
-    __half const*, int const*, int const*, __half*, int*, int, int, int, int, int);
+    __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernelDtype<__half, 1024>(
-    __half const*, int const*, int const*, __half*, int*, int, int, int, int, int);
+    __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernelDtype<__half, 2048>(
-    __half const*, int const*, int const*, __half*, int*, int, int, int, int, int);
+    __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernel<512>(
-    float const*, int const*, int const*, float*, int*, int, int, int, int, int);
+    float const*, int const*, int const*, float*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernel<1024>(
-    float const*, int const*, int const*, float*, int*, int, int, int, int, int);
+    float const*, int const*, int const*, float*, int*, int, int, int, int, int, int);
 template __global__ void heuristicTopKMultiRowKernel<2048>(
-    float const*, int const*, int const*, float*, int*, int, int, int, int, int);
+    float const*, int const*, int const*, float*, int*, int, int, int, int, int, int);
 
 // Dispatch on topK at runtime — each TopK-instantiation gets its own smem
 // size (driven by GvrParams<InputT, TopK>::kC/kNumBins) and own kfn pointer
@@ -184,7 +202,7 @@ template __global__ void heuristicTopKMultiRowKernel<2048>(
 template <typename InputT>
 void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int const* preIdx, int* outIndices,
     InputT* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows,
-    cudaStream_t stream)
+    int compressRatio, cudaStream_t stream)
 {
     TLLM_CHECK_WITH_INFO(
         topK == 512 || topK == 1024 || topK == 2048, "heuristicTopKDecode requires topK ∈ {512, 1024, 2048}");
@@ -224,7 +242,7 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int
         config.attrs = attrs;
 
         cudaLaunchKernelEx(&config, kfn, logits, seqLens, preIdx, scratchValues, outIndices, stride0, next_n, topK,
-            preIdxStride, preIdxCount);
+            preIdxStride, preIdxCount, compressRatio);
     };
 
     switch (topK)
@@ -240,26 +258,26 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int
 
 void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices,
     float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows,
-    cudaStream_t stream)
+    int compressRatio, cudaStream_t stream)
 {
     launchHeuristicTopKDecodeImpl<float>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK,
-        preIdxStride, preIdxCount, numRows, stream);
+        preIdxStride, preIdxCount, numRows, compressRatio, stream);
 }
 
 void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices,
     __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows,
-    cudaStream_t stream)
+    int compressRatio, cudaStream_t stream)
 {
     launchHeuristicTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n,
-        topK, preIdxStride, preIdxCount, numRows, stream);
+        topK, preIdxStride, preIdxCount, numRows, compressRatio, stream);
 }
 
 void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices,
     __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows,
-    cudaStream_t stream)
+    int compressRatio, cudaStream_t stream)
 {
     launchHeuristicTopKDecodeImpl<__half>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK,
-        preIdxStride, preIdxCount, numRows, stream);
+        preIdxStride, preIdxCount, numRows, compressRatio, stream);
 }
 
 } // namespace kernels

cpp/tensorrt_llm/kernels/heuristicTopKDecode.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,21 +33,28 @@ inline constexpr int kHeuristicSize = 2048;
 /// Launch heuristic TopK decode kernel — fp32 input.
 /// @param scratchValues Caller-owned buffer of size [numRows * topK] floats.
 ///        Required for CUDA Graph compatibility — must have a stable device address.
+/// @param compressRatio  KV compression ratio (1 = V3.2 indexer; 4 = V4 indexer
+///        whose logits/preIdx live in compressed-token-index space). For
+///        compressRatio != 1, preIdxOffset is forced to 0 (append-at-end in
+///        compressed space → prev-step indices remain valid as-is); the
+///        existing (rowIdx % next_n)+1 shift is used only when compressRatio==1.
 void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices,
     float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows,
-    cudaStream_t stream);
+    int compressRatio, cudaStream_t stream);
 
 /// Launch heuristic TopK decode kernel — bf16 input.
 /// scratchValues is [numRows * topK] of bf16 (matches input dtype).
+/// @param compressRatio  See fp32 overload.
 void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices,
     __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows,
-    cudaStream_t stream);
+    int compressRatio, cudaStream_t stream);
 
 /// Launch heuristic TopK decode kernel — fp16 input.
 /// scratchValues is [numRows * topK] of fp16 (matches input dtype).
+/// @param compressRatio  See fp32 overload.
 void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices,
     __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows,
-    cudaStream_t stream);
+    int compressRatio, cudaStream_t stream);
 
 } // namespace kernels
 

cpp/tensorrt_llm/kernels/indexerTopK.cu

@@ -902,7 +902,12 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic
     int const kSeqSmall = bounds.kSeqSmall;
 
     bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048);
-    bool const canUseHeuristic = compressRatio == 1 && preIdx != nullptr && stride1 == 1 && isSupportedTopK
+    // compressRatio == 1: DSv3.2 indexer (no compressor).
+    // compressRatio == 4: DSv4 indexer (overlap compressor); logits/preIdx in
+    //   compressed-token-index space. Kernel handles N = actual_kv_len/cr and
+    //   forces preIdxOffset=0 internally for cr != 1.
+    bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4);
+    bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK
         && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold
         && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge;
 
@@ -930,7 +935,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic
     if (canUseHeuristic)
     {
         launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK,
-            preIdxStride, preIdxCount, numRows, stream);
+            preIdxStride, preIdxCount, numRows, compressRatio, stream);
         sync_check_cuda_error(stream);
         return;
     }
@@ -1035,14 +1040,16 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int*
     int const kSeqSmall = bounds.kSeqSmall;
 
     bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048);
-    bool const canUseHeuristic = compressRatio == 1 && preIdx != nullptr && stride1 == 1 && isSupportedTopK
+    // See fp32 path: cr==1 (V3.2) and cr==4 (V4 indexer) are both supported.
+    bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4);
+    bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK
         && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold
         && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge;
 
     if (canUseHeuristic)
     {
         launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK,
-            preIdxStride, preIdxCount, numRows, stream);
+            preIdxStride, preIdxCount, numRows, compressRatio, stream);
     }
     else if (numColumns < kSortingAlgorithmThreshold)
     {

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -2199,6 +2199,31 @@ def sparse_attn_indexer(
             topk_indices_buffer[:num_ctx_tokens, :] = \
                 metadata.topk_indices_buffer[:num_ctx_tokens, :]
 
+        # Prefill→decode GVR handoff: seed each finishing-prefill sequence's
+        # heuristic_prev_topk slot with its own last-context-token top-K, so
+        # the FIRST decode step of that sequence gets a warm-started preIdx
+        # (~60-75% set-overlap with the eventual decode top-K on this
+        # workload) instead of the all-zero / all-(-1) cold start that the
+        # default `heuristic_prev_topk.zero_()` initialization leaves behind.
+        # Without this, GVR P2 secant on decode step 0 runs from a benign
+        # but uninformative seed (kernel +1 offset on zeros → all indices
+        # point at compressed-token position 1), wasting iterations.
+        # Slot convention (mirrors the existing decode write-back at the
+        # bottom of the decode block): new gens from finishing prefill
+        # append after currently-active gens, i.e., slots
+        # [num_generations : num_generations + num_contexts].
+        if (self._enable_heuristic_topk and has_prefill
+                and not metadata.skip_indexer_for_ctx_reqs):
+            local_layer = metadata.kv_cache_manager.layer_offsets[
+                self.layer_idx]
+            ctx_seq_lens = metadata.seq_lens[:num_contexts]
+            # Per-sequence last context-token offset (exclusive cumsum minus 1).
+            last_ctx_idx = (torch.cumsum(ctx_seq_lens, dim=0) -
+                            1).to(dtype=torch.long)
+            metadata.heuristic_prev_topk[
+                local_layer, num_generations:num_generations +
+                num_contexts].copy_(topk_indices_buffer[last_ctx_idx, :])
+
         if has_decode and not metadata.skip_indexer_for_gen_reqs:
             # Get decode lengths per request (from seq_lens) for validation
             gen_seq_lens = metadata.seq_lens[num_contexts:num_contexts +