Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lfr-0531 merged 3 commits into from
May 19, 2026
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4a84ea5
[None][feat] DSv4: enable GVR Heuristic Top-K for compress_ratio=4
longcheng-nv May 17, 2026
b6291ef
[None][feat] DSv4: warm-start GVR preIdx at prefill→decode boundary
longcheng-nv May 18, 2026
4199524
[None][chore] DSv4: yapf reformat dsa.py to satisfy pre-commit
longcheng-nv May 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 48 additions & 30 deletions cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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;
Expand Down Expand Up @@ -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))
Expand All @@ -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
Expand All @@ -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}");
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down
15 changes: 11 additions & 4 deletions cpp/tensorrt_llm/kernels/heuristicTopKDecode.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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

Expand Down
15 changes: 11 additions & 4 deletions cpp/tensorrt_llm/kernels/indexerTopK.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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)
{
Expand Down
26 changes: 26 additions & 0 deletions tensorrt_llm/_torch/model_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,19 +681,45 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs):
index_topk = sparse_attention_config.index_topk or pretrained_config.index_topk
indexer_max_chunk_size = sparse_attention_config.indexer_max_chunk_size
skip_indexer_for_short_seqs = sparse_attention_config.skip_indexer_for_short_seqs
# Pass-through DSA tuning flags so user-set values survive the
# V4 sparse_attention_config rebuild below. The V3.2 path
# already threads these explicitly (lines 723-727); without
# this block the V4 rebuild silently drops any user override
# back to subclass defaults (e.g., enable_heuristic_topk=False
# even when the user set it to True in --extra_llm_api_options).
use_cute_dsl_topk = sparse_attention_config.use_cute_dsl_topk
use_cute_dsl_paged_mqa_logits = sparse_attention_config.use_cute_dsl_paged_mqa_logits
q_split_threshold = sparse_attention_config.q_split_threshold
indexer_rope_interleave = sparse_attention_config.indexer_rope_interleave
enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk
indexer_k_dtype = sparse_attention_config.indexer_k_dtype
else:
index_n_heads = pretrained_config.index_n_heads
index_head_dim = pretrained_config.index_head_dim
index_topk = pretrained_config.index_topk
indexer_max_chunk_size = None
skip_indexer_for_short_seqs = True
# Defaults match DeepSeekSparseAttentionConfig field defaults.
use_cute_dsl_topk = False
use_cute_dsl_paged_mqa_logits = False
q_split_threshold = 8192
indexer_rope_interleave = False
enable_heuristic_topk = False
indexer_k_dtype = "fp8"
indexer_config = {}
indexer_config['index_n_heads'] = index_n_heads
indexer_config['index_head_dim'] = index_head_dim
indexer_config['index_topk'] = index_topk
indexer_config['indexer_max_chunk_size'] = indexer_max_chunk_size
indexer_config[
'skip_indexer_for_short_seqs'] = skip_indexer_for_short_seqs
indexer_config['use_cute_dsl_topk'] = use_cute_dsl_topk
indexer_config[
'use_cute_dsl_paged_mqa_logits'] = use_cute_dsl_paged_mqa_logits
indexer_config['q_split_threshold'] = q_split_threshold
indexer_config['indexer_rope_interleave'] = indexer_rope_interleave
indexer_config['enable_heuristic_topk'] = enable_heuristic_topk
indexer_config['indexer_k_dtype'] = indexer_k_dtype
return indexer_config

# Use file lock to prevent race conditions when multiple processes
Expand Down
7 changes: 4 additions & 3 deletions tensorrt_llm/llmapi/llm_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,9 +352,10 @@ class DeepSeekSparseAttentionConfig(BaseSparseAttentionConfig):
description=
"Whether to enable Guess-Verify-Refine (GVR) Top-K for the DSA decode "
"indexer. GVR reuses previous-step Top-K indices as hints to reduce "
"threshold search iterations. Currently supported for index_topk=2048 "
"on Blackwell (SM100+) and falls back to the production insertion/radix "
"Top-K path when prerequisites are not met.")
"threshold search iterations. Currently supported for index_topk ∈ "
"{512, 1024, 2048} on Blackwell (SM100+), with compress_ratio ∈ {1, 4} "
"(DSv3.2 + DSv4 indexers). Falls back to the production insertion/"
"radix Top-K path when prerequisites are not met.")
indexer_k_dtype: Literal["fp8", "fp4"] = Field(
default="fp8",
description=
Expand Down
Loading
Loading