Skip to content

Draft: Support context sparse attention for mqa/gqa #10788

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

Closed
heyuhhh wants to merge 3 commits into from
Closed

Draft: Support context sparse attention for mqa/gqa #10788

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
93bb982
Add attentionOp tests and naive support for context sparse attention
heyuhhh Dec 24, 2025
1b5ff93
Update code for alignment
heyuhhh Jan 8, 2026
0df4790
minor updates
heyuhhh Jan 17, 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
19 changes: 12 additions & 7 deletions cpp/tensorrt_llm/common/attentionOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,7 +295,7 @@ bool AttentionOp::convertMMHAParamsToXQAParams(tensorrt_llm::kernels::XQAParams&
xqaParams.start_token_idx_sf = generationsParams.start_token_idx_sf;
// Parameters for sparse attention
xqaParams.sparse_params = mRuntimeSparseAttentionParams;
xqaParams.use_sparse_attention = useTllmGenSparseAttention();
xqaParams.use_sparse_attention_gen_paged = useTllmGenSparseAttentionPaged();
// Skip softmax threshold.
Comment on lines +298 to 299

@coderabbitai coderabbitai Bot Jan 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Find definitions of the sparse attention helpers
rg -n "useTllmGenSparseAttention" cpp/tensorrt_llm/common/attentionOp.cpp -C 3

Repository: NVIDIA/TensorRT-LLM

Length of output: 1979

🏁 Script executed:

#!/bin/bash
# Find where these helpers are defined
rg -n "useTllmGenSparseAttention(Paged)?\s*\(" -C 3

Repository: NVIDIA/TensorRT-LLM

Length of output: 5102

🏁 Script executed:

#!/bin/bash
# Get the file size and examine the relevant lines
wc -l cpp/tensorrt_llm/common/attentionOp.cpp

Repository: NVIDIA/TensorRT-LLM

Length of output: 107

🏁 Script executed:

#!/bin/bash
# Get the implementation of useSparseMLA and related helpers
rg -n "useSparseMLA|useSparseAttention" cpp/tensorrt_llm/common/attentionOp.h -C 2

Repository: NVIDIA/TensorRT-LLM

Length of output: 706

🏁 Script executed:

#!/bin/bash
# Get broader context around line 2754 to understand the control flow
sed -n '2740,2770p' cpp/tensorrt_llm/common/attentionOp.cpp

Repository: NVIDIA/TensorRT-LLM

Length of output: 1347

Line 2754-2756 needs a guard: useTllmGenSparseAttention() can be true when mPagedKVCache == false.

The code unconditionally sets AttentionInputLayout::Q_PAGED_KV when useTllmGenSparseAttention() returns true. However, this helper does not check mPagedKVCache—it only checks sparse attention and TllmGen flags. If contiguous KV is in use, the paged layout will cause invalid memory reads.

Line 298 is already safe because useTllmGenSparseAttentionPaged() enforces the paged invariant through useSparseAttention(), which includes && mPagedKVCache. The issue is at lines 2754-2756:

🔧 Required fix 
         else if (useTllmGenSparseAttention())
+            && mPagedKVCache)
         {
             fmhaParams.attentionInputLayout = AttentionInputLayout::Q_PAGED_KV;
         }
🤖 Prompt for AI Agents 
In `@cpp/tensorrt_llm/common/attentionOp.cpp` around lines 298 - 299, The code
unconditionally sets AttentionInputLayout::Q_PAGED_KV when
useTllmGenSparseAttention() is true, but useTllmGenSparseAttention() does not
check mPagedKVCache and can cause invalid reads when KV is contiguous; change
the guard to require the paged invariant (e.g., use
useTllmGenSparseAttentionPaged() or explicitly check mPagedKVCache) before
assigning Q_PAGED_KV so the paged layout is only used when mPagedKVCache is
true; update the branch that currently calls useTllmGenSparseAttention() (the
one setting AttentionInputLayout::Q_PAGED_KV) to
useTllmGenSparseAttentionPaged() or add && mPagedKVCache to the condition.

xqaParams.skip_softmax_threshold_scale_factor = mSkipSoftmaxThresholdScaleFactorDecode;
// Cross attention parameters.
Expand Down Expand Up @@ -939,7 +939,7 @@ size_t AttentionOp::getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32
size_t const cu_kv_seqlens_size = sizeof(int) * (batch_beam + 1);
size_t const rotary_inv_freq_size = sizeof(float) * batch_beam * mRotaryEmbeddingDim / 2;
// Two workspaces for sparse attention. One for the sequence lengths, and one for kv block offsets.
size_t const sparse_attn_cache_size = useTllmGenSparseAttention()
size_t const sparse_attn_cache_size = useTllmGenSparseAttentionPaged()
? sizeof(int) * (batch_beam + batch_beam * 2 * max_blocks_per_sequence) * mNumKVHeads
: 0;
xqa_workspaces[0] = cu_seqlens_size;
Expand Down Expand Up @@ -1111,14 +1111,14 @@ int AttentionOp::mlaGeneration(
= reinterpret_cast<float const*>(params.bmm1_scale) + bmm1_scale_offset;
}

// Set the following parameters if sparseMLA is used.
// Set the following parameters if sparseAttention is used.
if (useSparseMLA())
{
tllmRunnerParams.mSparseMla = true;
tllmRunnerParams.mSparseMlaTopK = mRuntimeSparseAttentionParams.sparse_mla_topk;
tllmRunnerParams.mSparseAttention = true;
tllmRunnerParams.mSparseTopK = mRuntimeSparseAttentionParams.sparse_topk;
tllmRunnerParams.kvPageIdxPtr = reinterpret_cast<KVCacheIndex::UnderlyingType const*>(
mRuntimeSparseAttentionParams.sparse_attn_indices);
tllmRunnerParams.kvPtr = mRuntimeSparseAttentionParams.sparse_mla_kv_cache_pool;
tllmRunnerParams.kvPtr = mRuntimeSparseAttentionParams.sparse_kv_cache_pool;
}

mTllmGenFMHARunner->run(tllmRunnerParams);
Expand Down Expand Up @@ -1883,7 +1883,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
fmhaParams.softmaxStatsPtr = params.softmax_stats;

// Sparse attention parameters
if (useSparseMLA())
if (useTllmGenSparseAttention())
{
fmhaParams.sparse_params = mRuntimeSparseAttentionParams;
}
Expand Down Expand Up @@ -2751,6 +2751,10 @@ int AttentionOp::initialize() noexcept
{
fmhaParams.attentionInputLayout = AttentionInputLayout::PACKED_QKV;
}
else if (useTllmGenSparseAttention())
{
fmhaParams.attentionInputLayout = AttentionInputLayout::Q_PAGED_KV;
}
else
{
fmhaParams.attentionInputLayout = (mPagedKVCache && mPagedContextFMHA) ? AttentionInputLayout::Q_PAGED_KV
Expand Down Expand Up @@ -2799,6 +2803,7 @@ int AttentionOp::initialize() noexcept
fmhaParams.hasAlibi = isALiBi();
fmhaParams.scaleAlibi = isAliBiWithScale();
fmhaParams.useSparseMLA = useSparseMLA();
fmhaParams.useTllmGenSparseAttention = useTllmGenSparseAttention();

// Load kernels from the pre-compiled cubins.
mFmhaDispatcher.reset(new FmhaDispatcher(fmhaParams));
Expand Down
18 changes: 12 additions & 6 deletions cpp/tensorrt_llm/common/attentionOp.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,16 +367,21 @@ class AttentionOp
return mUseSparseAttention && mPagedKVCache && mEnableXQA;
}

[[nodiscard]] bool useTllmGenSparseAttention() const
[[nodiscard]] bool useTllmGenSparseAttentionPaged() const
{
return mUseTllmGenSparseAttention && useSparseAttention();
return mUseTllmGenSparseAttentionPaged && useSparseAttention();
}

[[nodiscard]] bool useSparseMLA() const
{
return mUseSparseAttention && mUseTllmGen && mIsMLAEnabled;
}

[[nodiscard]] bool useTllmGenSparseAttention() const
{
return useSparseMLA() || (mUseSparseAttention && mUseTllmGen && mUseTllmGenSparseAttention);
}

[[nodiscard]] int smVersion() const
{
return mSM;
Expand Down Expand Up @@ -457,6 +462,7 @@ class AttentionOp
bool mIsGenerationMLA = false;
bool mUseGenFlashMLA = false;
bool mUseSparseAttention = false;
bool mUseTllmGenSparseAttentionPaged = false;
bool mUseTllmGenSparseAttention = false;
tensorrt_llm::kernels::MlaMetaParams mMLAParams;
int mCpSize = 1;
Expand Down Expand Up @@ -514,10 +520,10 @@ class AttentionOp
mPosShiftEnabled, mPagedContextFMHA, mFP8ContextFMHA, mFP8AttenOutput, mFP8ContextMLA, mFP8GenerationMLA,
mChunkPrefillBufferBatchSize, mDenseContextFMHA, mHasFullAttentionMask, mIsSpecDecodingEnabled,
mUseSpecDecoding, mIsSpecDecTree, mSpecDecodingIsGenerationLengthVariable, mSpecDecodingMaxGenerationLength,
mIsMLAEnabled, mIsGenerationMLA, mUseGenFlashMLA, mUseSparseAttention, mUseTllmGenSparseAttention,
mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin,
mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast, mEnableContextFMHA,
mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant,
mIsMLAEnabled, mIsGenerationMLA, mUseGenFlashMLA, mUseSparseAttention, mUseTllmGenSparseAttentionPaged,
mUseTllmGenSparseAttention, mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads,
mNumKVHeadsOrigin, mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast,
mEnableContextFMHA, mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant,
mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1), mSkipSoftmaxThresholdScaleFactorPrefill,
mSkipSoftmaxThresholdScaleFactorDecode);
};
Expand Down
4 changes: 4 additions & 0 deletions cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,8 @@ struct MHARunnerFixedParams
int sageBlockSizeV = 0;
// Use sparse MLA ?
bool useSparseMLA = false;
// Use sparse attention in trtllm-gen ?
bool useTllmGenSparseAttention = false;

// Convert to string for debug.
std::string convertToStrOutput()
Expand Down Expand Up @@ -193,6 +195,8 @@ struct MHARunnerFixedParams
output += ", sageBlockSizeQ = " + std::to_string(sageBlockSizeQ);
output += ", sageBlockSizeK = " + std::to_string(sageBlockSizeK);
output += ", sageBlockSizeV = " + std::to_string(sageBlockSizeV);
output += ", useSparseMLA = " + std::string(useSparseMLA ? "true" : "false");
output += ", useTllmGenSparseAttention = " + std::string(useTllmGenSparseAttention ? "true" : "false");

return output;
}
Expand Down
2 changes: 1 addition & 1 deletion cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,7 @@ void buildXQALaunchParams(XQALaunchParam<KVCacheBuffer>& launchParams, void*& in
launchParams.bmm2_scale_ptr = reinterpret_cast<float*>(workspace);
workspace = tensorrt_llm::common::nextWorkspacePtrWithAlignment(workspace, bmm2_scale_size);
// Used for block sparse attention
if (params.use_sparse_attention)
if (params.use_sparse_attention_gen_paged)
{
launchParams.sparse_kv_block_offsets = reinterpret_cast<void*>(workspace);
workspace = tensorrt_llm::common::nextWorkspacePtrWithAlignment(workspace, kv_block_offsets_size);
Expand Down
4 changes: 2 additions & 2 deletions cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ struct XQAParams

// sparse attention parameters
SparseAttentionParams sparse_params;
bool use_sparse_attention = false;
bool use_sparse_attention_gen_paged = false;

// Skip softmax threshold.
float skip_softmax_threshold_scale_factor = 0.0f;
Expand Down Expand Up @@ -197,7 +197,7 @@ struct XQAParams
<< "fp8_out_scale :" << fp8_out_scale << std ::endl
<< "encoder_input_lengths: " << encoder_input_lengths << std::endl
<< "sparse_params: " << sparse_params.toString() << std::endl
<< "use_sparse_attention :" << (use_sparse_attention ? "true" : "false") << std ::endl
<< "use_sparse_attention_gen_paged :" << (use_sparse_attention_gen_paged ? "true" : "false") << std ::endl
<< "skip_softmax_threshold_scale_factor :" << skip_softmax_threshold_scale_factor << std ::endl
<< "stream :" << stream;

Expand Down
28 changes: 18 additions & 10 deletions cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,10 +120,10 @@ bool FmhaDispatcher::isSupported()
// the kernel is supported.
tllmRunnerParams.mChunkedAttentionSize = INT_MAX;
tllmRunnerParams.mAttentionWindowSize = INT_MAX;
// Set the kernel type and mask type if sparseMLA is used.
if (mFixedParams.useSparseMLA)
// Set the kernel type and mask type if sparse attention is used.
if (mFixedParams.useTllmGenSparseAttention)
{
tllmRunnerParams.mSparseMla = true;
tllmRunnerParams.mSparseAttention = true;
tllmRunnerParams.mKernelType = FmhaKernelType::Generation;
tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Dense;
}
Expand Down Expand Up @@ -231,16 +231,24 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams)
// For skip softmax
tllmRunnerParams.mSkipSoftmaxThresholdScaleFactor = runnerParams.skipSoftmaxThresholdScaleFactor;
tllmRunnerParams.stream = runnerParams.stream;
// Set the sparse attention parameters if sparseMLA is used.
if (mFixedParams.useSparseMLA)
// Set the sparse attention parameters if trtllm-gen sparse attention is used.
if (mFixedParams.useTllmGenSparseAttention)
{
tllmRunnerParams.mSparseMla = true;
tllmRunnerParams.mSparseMlaTopK = runnerParams.sparse_params.sparse_mla_topk;
tllmRunnerParams.mSparseAttention = true;
tllmRunnerParams.mSparseTopK = runnerParams.sparse_params.sparse_topk;
tllmRunnerParams.mKernelType = FmhaKernelType::Generation;
tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Dense;
tllmRunnerParams.kvPageIdxPtr
= reinterpret_cast<int const*>(runnerParams.sparse_params.sparse_attn_indices);
tllmRunnerParams.kvPtr = runnerParams.sparse_params.sparse_mla_kv_cache_pool;
if (mFixedParams.useSparseMLA)
{
tllmRunnerParams.kvPageIdxPtr
= reinterpret_cast<int const*>(runnerParams.sparse_params.sparse_attn_indices);
}
else
{
tllmRunnerParams.kvPageIdxPtr
= reinterpret_cast<int const*>(runnerParams.sparse_params.sparse_attn_ctx_indices);
}
tllmRunnerParams.kvPtr = runnerParams.sparse_params.sparse_kv_cache_pool;
}

mTllmGenFMHARunner->run(tllmRunnerParams);
Expand Down
18 changes: 10 additions & 8 deletions cpp/tensorrt_llm/kernels/sparseAttentionKernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,13 @@ namespace kernels

struct SparseAttentionParams
{
int32_t* sparse_kv_indices{nullptr}; // [num_kv_heads, num_sparse_kv_indices]
int32_t* sparse_attn_indices{nullptr}; // [num_kv_heads, num_sparse_attn_indices]
int32_t* sparse_kv_offsets{nullptr}; // [num_contexts + 1]
int32_t* sparse_attn_offsets{nullptr}; // [num_generations + 1]
int32_t sparse_mla_topk{0}; // for DSA attention
void* sparse_mla_kv_cache_pool{nullptr}; // for DSA attention
int32_t* sparse_kv_indices{nullptr}; // [num_kv_heads, num_sparse_kv_indices]
int32_t* sparse_attn_indices{nullptr}; // [num_kv_heads, num_sparse_attn_indices]
int32_t* sparse_kv_offsets{nullptr}; // [num_contexts + 1]
int32_t* sparse_attn_offsets{nullptr}; // [num_generations + 1]
int32_t* sparse_attn_ctx_indices{nullptr}; // [num_kv_heads, num_tokens, sparse_topk]
int32_t sparse_topk{0};
void* sparse_kv_cache_pool{nullptr};

int32_t sparse_attn_indices_block_size{1};
int32_t sparse_attn_indices_stride{0};
Expand All @@ -46,8 +47,9 @@ struct SparseAttentionParams
<< "sparse_attn_indices: " << this->sparse_attn_indices << std::endl
<< "sparse_kv_offsets: " << this->sparse_kv_offsets << std::endl
<< "sparse_attn_offsets: " << this->sparse_attn_offsets << std::endl
<< "sparse_mla_topk: " << this->sparse_mla_topk << std::endl
<< "sparse_mla_kv_cache_pool: " << this->sparse_mla_kv_cache_pool << std::endl
<< "sparse_attn_ctx_indices: " << this->sparse_attn_ctx_indices << std::endl
<< "sparse_topk: " << this->sparse_topk << std::endl
<< "sparse_kv_cache_pool: " << this->sparse_kv_cache_pool << std::endl
<< "sparse_attn_indices_block_size: " << this->sparse_attn_indices_block_size << std::endl
<< "sparse_attn_indices_stride: " << this->sparse_attn_indices_stride << std::endl;
return ss.str();
Expand Down
Loading