Skip to content

[TRTLLM-14620][feat] Add ragged sparse-attention row contracts #18485

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
zhaoyuanh-nvidia merged 1 commit into from
Oct 8, 2026
Merged

[TRTLLM-14620][feat] Add ragged sparse-attention row contracts #18485

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 11 additions & 3 deletions cpp/tensorrt_llm/kernels/IndexerTopK.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,10 +38,17 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW
/// - Insertion sort (N < kSortingAlgorithmThreshold)
/// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork)
/// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux)
///
/// @param rowKvLens Optional [numRows] per-row causal KV extent in uncompressed token space.
/// Non-null selects the ragged layout: rows no longer map to requests by `rowIdx / next_n`,
/// so each row's extent is supplied directly instead of reconstructed from `next_n`.
/// Null uses the uniform request-to-row mapping. A negative extent or one whose compressed
/// length exceeds `numColumns` fails closed as an empty row, so device-updated graph inputs
/// cannot read beyond the corresponding logits row.
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 compressRatio = 1,
cudaStream_t const stream = 0);
int const* rowKvLens = nullptr, cudaStream_t const stream = 0);

/// bf16 indexer TopK decode — same dispatch tiers as the fp32 entry, except
/// the split-work tier is unsupported (the bf16/fp16 entry does not expose
Expand All @@ -53,12 +60,13 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic
/// 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 compressRatio = 1, cudaStream_t const stream = 0);
int const next_n, int const topK = 2048, int const compressRatio = 1, int const* rowKvLens = nullptr,
cudaStream_t const stream = 0);

/// 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 compressRatio = 1, cudaStream_t const stream = 0);
int const topK = 2048, int const compressRatio = 1, int const* rowKvLens = nullptr, 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,
Expand Down
36 changes: 27 additions & 9 deletions cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -298,12 +298,13 @@ __device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_
}
}

template <int HEAD_DIM, int KV_SCORE_ELEM_BYTES, int STATE_ELEM_BYTES, int COMPRESS_RATIO, int NUM_RED_WARPS = 1>
template <int HEAD_DIM, int KV_SCORE_ELEM_BYTES, int STATE_ELEM_BYTES, int COMPRESS_RATIO, int NUM_RED_WARPS = 1,
bool USE_RAGGED_TOKENS = false>
__global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape,
void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv,
int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens,
int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int next_n, int page_size,
int max_blocks, int out_elem_bytes)
int max_blocks, int out_elem_bytes, int32_t const* __restrict__ new_tokens_per_seq = nullptr)
{
using KvScoreElemT = typename std::conditional<KV_SCORE_ELEM_BYTES == 2, __nv_bfloat16, float>::type;
using StateElemT = typename std::conditional<STATE_ELEM_BYTES == 2, __nv_bfloat16, float>::type;
Expand Down Expand Up @@ -344,6 +345,10 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo
int const eff_tid = head_blk * NTHRD_INNER + tid;

int const kv_len = kv_lens[batch_idx];
if constexpr (USE_RAGGED_TOKENS)
{
next_n = new_tokens_per_seq[batch_idx];
}
int const sp = kv_len - next_n;
int const in_off = cu_seq_lens[batch_idx];
int const out_off = cu_kv_comp[batch_idx];
Expand Down Expand Up @@ -687,13 +692,17 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo
// `const`. cudafe1 rewrites every kernel parameter into a reference when it emits `__wrapper__device_stub_*`, which
// turns those otherwise-ignored qualifiers into part of the function type. Clang then rejects the instantiation;
// GCC does not diagnose it.
#define INST_DECODE(HD, KV_EB, STATE_EB, CR, NRW) \
template __global__ void pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NRW>(void const* __restrict__, \
#define INST_DECODE_RAGGED(HD, KV_EB, STATE_EB, CR, NRW, RAGGED) \
template __global__ void pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NRW, RAGGED>(void const* __restrict__, \
float const* __restrict__, void* __restrict__, void* __restrict__, int32_t const* __restrict__, \
int32_t const* __restrict__, void* __restrict__, int32_t const* __restrict__, int32_t const* __restrict__, \
int32_t const* __restrict__, int, int, int, int);
int32_t const* __restrict__, int, int, int, int, int32_t const* __restrict__);
#define INST_DECODE(HD, KV_EB, STATE_EB, CR, NRW) \
INST_DECODE_RAGGED(HD, KV_EB, STATE_EB, CR, NRW, false) \
INST_DECODE_RAGGED(HD, KV_EB, STATE_EB, CR, NRW, true)
FOREACH_DECODE_CONFIG(INST_DECODE)
#undef INST_DECODE
#undef INST_DECODE_RAGGED

// ============================================================================
// Decode Launch Wrapper
Expand All @@ -711,7 +720,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k
int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens,
int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, int max_blocks, int head_dim,
int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes,
cudaStream_t stream)
int32_t const* new_tokens_per_seq, cudaStream_t stream)
{
TLLM_CHECK_WITH_INFO(
compress_ratio == 4 || compress_ratio == 128, "pagedKvCompressLaunch only supports compress_ratio 4 or 128");
Expand Down Expand Up @@ -765,9 +774,18 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k
if (head_dim == HD && kv_score_elem_bytes == KV_EB && state_elem_bytes == STATE_EB && compress_ratio == CR \
&& num_red_warps == NRW) \
{ \
pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NRW><<<grid, nthreads, smem_bytes, stream>>>(kv_score, ape, \
paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \
next_n, page_size, max_blocks, out_elem_bytes); \
if (new_tokens_per_seq == nullptr) \
{ \
pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NRW, false><<<grid, nthreads, smem_bytes, stream>>>( \
kv_score, ape, paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, \
cu_kv_comp, next_n, page_size, max_blocks, out_elem_bytes, nullptr); \
} \
else \
{ \
pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NRW, true><<<grid, nthreads, smem_bytes, stream>>>( \
kv_score, ape, paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, \
cu_kv_comp, next_n, page_size, max_blocks, out_elem_bytes, new_tokens_per_seq); \
} \
return; \
}
FOREACH_DECODE_CONFIG(TRY_LAUNCH)
Expand Down
4 changes: 4 additions & 0 deletions cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,10 @@ void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or f
int kv_score_elem_bytes, // bytes per element for kv_score (2=bf16, 4=fp32)
int state_elem_bytes, // bytes per element for paged state (2=bf16, 4=fp32)
int out_elem_bytes, // bytes per element for output
int32_t const* new_tokens_per_seq, // [bsz] or nullptr; per-request tokens appended
// this step. nullptr means every request appended
// next_n (the uniform case); non-null makes next_n
// an upper bound only for ragged verification.
cudaStream_t stream);

// Prefill kernel: bulk compression with per-token gather/scatter + state update.
Expand Down
Loading
Loading