Skip to content

[None][feat] DSv4 prep: IndexerTopK and TopK primitives #15381

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
longlee0622 merged 2 commits into from
Jun 24, 2026
Merged

[None][feat] DSv4 prep: IndexerTopK and TopK primitives #15381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f8400e4
[None][feat] DSv4 prep: IndexerTopK and TopK primitives
lfr-0531 Jun 15, 2026
7fe184c
Merge branch 'main' into user/fanrongl/dsv4-indexer-topk
lfr-0531 Jun 16, 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
72 changes: 44 additions & 28 deletions cpp/tensorrt_llm/kernels/IndexerTopK.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.
* Copyright (c) 2021, NAVER Corp. Authored by CLOVA.
*
* Licensed under the Apache License, Version 2.0 (the "License");
Expand Down Expand Up @@ -27,47 +27,63 @@ TRTLLM_NAMESPACE_BEGIN

namespace kernels
{
/// Indexer TopK decode. Three tiers:
/// - GVR Heuristic (preIdx provided, K in {512,1024,2048}, numColumns in
/// [kSeqSmall, splitWorkThreshold), numRows below the
/// architecture-derived wave/L2 bound).
/// - Single-block (numColumns < split-work threshold)
/// - Multi-pass radix (numColumns >= split-work threshold; requires
/// `scratch` sized via indexerTopKDecodeScratchBytes,
/// zero-init on first call and may be reused).
///
/// `is_prefill = true` forces single-block (split-work suppressed).
void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, int const splitWorkThreshold,
int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n,
int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0,
float* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0,
bool is_prefill = false);
// Number of blocks-per-row used by the multi-block split + merge dispatch path of
// invokeIndexerTopKDecode. Returns 1 when the single-block path is preferred.
// Callers that allocate aux buffers must use this same helper to size them, and
// must pass the same splitWorkThreshold they will pass to invokeIndexerTopKDecode
// (a value <= 0 selects the internal default).
int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold = 0);

/// Size of the multi-pass radix `scratch` buffer for these shapes.
size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int topK);
/// fp32 indexer TopK decode — L2-aware BS-threshold dispatcher with four
/// fallback tiers:
/// - GVR Heuristic (preIdx provided, kSeqSmall ≤ N < splitWork, BS < kBsLarge, K ∈ {512,1024,2048})
/// - Insertion sort (N < kSortingAlgorithmThreshold)
/// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork)
/// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux)
void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux,
int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0,
int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0,
int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1,
cudaStream_t const stream = 0);

/// bf16 overload; same contract.
/// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except
/// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and
/// the split-work tier is unsupported (the bf16/fp16 entry does not expose
/// the float aux buffers required for split-work). Insertion + radix tiers
/// share topKPerRowDecode with fp32 — histogram and sort run on float keys
/// after a static_cast<float>(InputT) at HBM-read sites.
///
/// Aborts with TLLM_CHECK if numColumns ≥ splitWorkThreshold; callers in
/// that regime must use the fp32 entry.
void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices,
int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1,
int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0,
int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, cudaStream_t const stream = 0,
void* scratch = nullptr, size_t scratchBytes = 0, bool is_prefill = false);
int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1,
cudaStream_t const stream = 0);

/// fp16 overload; same contract.
/// fp16 indexer TopK decode — see bf16 overload for dispatcher contract.
void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold,
int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n,
int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0,
__half* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0,
bool is_prefill = false);
__half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0);

void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices,
int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048,
cudaStream_t const stream = 0);

/// True iff invokeIndexerTopKDecode would pick the GVR tier for this shape:
/// K in {512,1024,2048}, numColumns in [kSeqSmall, splitWorkThreshold), and
/// numRows below the architecture-derived wave/L2 bound. Lets callers
/// provision preIdx / heuristicScratch only when needed.
/// Returns true iff invokeIndexerTopKDecode would route to the GVR Heuristic
/// kernel for this (numRows, numColumns, topK) triple, assuming valid preIdx
/// is provided and stride1 == 1. Useful for callers that need to provision a
/// preIdx tensor or heuristicScratch buffer only when GVR will be selected.
///
/// Mirrors the gating logic of the dispatcher: K ∈ {512, 1024, 2048},
/// numColumns ∈ [kSeqSmall, splitWorkThreshold), numRows < kBsLarge, where
/// kBsLarge = min(kBsWave, kBsL2) and kBsL2 scales with bytesPerElem.
///
/// @param numRows logits rows (batch · next_n)
/// @param numColumns logits columns (max sequence length)
/// @param topK requested output size
/// @param bytesPerElem element size of logits (4 for fp32, 2 for bf16/fp16)
bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem = 4);

} // namespace kernels
Expand Down
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
Loading
Loading