Skip to content

[None][feat] Add DeepSeekV4 attention kernels #13652

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 5 commits into from
May 9, 2026
Merged

[None][feat] Add DeepSeekV4 attention kernels #13652

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
339dec2
[None][feat] Add DeepSeek-V4 FMHA generated kernels
heyuhhh May 6, 2026
74f8299
[None][feat] Add DeepSeek-V4 attention kernel support
heyuhhh Apr 30, 2026
6c8440c
[None][fix] Fix AutoDeploy MLA attention op interface
heyuhhh May 2, 2026
862308a
[None][fix] Restore FMHA generation mask cubin filtering
heyuhhh May 3, 2026
e89cbb6
Merge branch 'main' into user/yuhangh/add_v4_attention_kernel
heyuhhh May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
The table of contents is too big for display.
Diff view
Diff view
  •  
  •  
  •  
The diff you're trying to view is too large. We only load the first 3000 changed files.
58 changes: 45 additions & 13 deletions cpp/tensorrt_llm/common/attentionOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@
#include "tensorrt_llm/runtime/utils/mpiUtils.h"
#include <algorithm>
#include <cstdint>
#include <type_traits>

using namespace tensorrt_llm::kernels;
namespace tc = tensorrt_llm::common;
Expand Down Expand Up @@ -796,7 +795,8 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
{
dim_q_per_head = mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
dim_k_per_head = mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
dim_v_per_head = mMLAParams.kv_lora_rank;
dim_v_per_head
= mMLAParams.rope_append ? mMLAParams.kv_lora_rank : mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
}

// Total dimension per token across all heads for Q, K, and V components respectively
Expand Down Expand Up @@ -926,7 +926,9 @@ size_t AttentionOp::getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32

int num_q_heads = mNumHeads / mCpSize;
int num_kv_heads = mNumKVHeads;
int head_size_v = mMLAParams.kv_lora_rank;
int head_size_v = (mUseSparseAttention && !mMLAParams.rope_append)
? mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim
: mMLAParams.kv_lora_rank;

int num_sm_parts = getFlashMlaNumSmParts(s_q, num_q_heads, num_kv_heads, head_size_v);

Expand Down Expand Up @@ -1042,6 +1044,7 @@ int AttentionOp::mlaGeneration(

int const num_kv_heads = 1;
int const head_size = mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
int const head_size_v = (useSparseMLA() && !mMLAParams.rope_append) ? head_size : mMLAParams.kv_lora_rank;
int32_t const batch_beam = generation_params.beam_width * generation_params.num_requests;

// The element size of the KV cache.
Expand Down Expand Up @@ -1090,7 +1093,9 @@ int AttentionOp::mlaGeneration(
TllmGenFmhaRunnerParams tllmRunnerParams{};

// Parameters to select kernels.
tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Causal;
// MLA generation kernels use dense mask. For multi-token generation, TRTLLM-Gen applies causality by
// shrinking each token's effective KV length.
tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Dense;
tllmRunnerParams.mKernelType = FmhaKernelType::Generation;
tllmRunnerParams.mMultiCtasKvMode = mMultiBlockMode;
// Note that the tileScheduler and multiCtasKvMode will be automatically tuned when using multi_block mode.
Expand Down Expand Up @@ -1122,9 +1127,12 @@ int AttentionOp::mlaGeneration(
// softmax stats if needed
tllmRunnerParams.softmaxStatsPtr = generation_params.softmax_stats;

// Per-head attention sink added to the softmax denominator.
tllmRunnerParams.attentionSinksPtr = generation_params.attention_sinks;

// MLA uses different head dimensions for Qk and V.
tllmRunnerParams.mHeadDimQk = mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
tllmRunnerParams.mHeadDimV = mMLAParams.kv_lora_rank;
tllmRunnerParams.mHeadDimQk = head_size;
tllmRunnerParams.mHeadDimV = head_size_v;

auto const num_q_heads = mNumAttnHeads;
tllmRunnerParams.mNumHeadsQ = num_q_heads;
Expand Down Expand Up @@ -1174,11 +1182,29 @@ int AttentionOp::mlaGeneration(
// Set the following parameters if sparseAttention is used.
if (useSparseMLA())
{
tllmRunnerParams.mSparseAttention = SparseType::StaticTokenSparse;
tllmRunnerParams.mSparseTopK = mRuntimeSparseAttentionParams.sparse_topk;
bool const useDynamicSparseMLA = mRuntimeSparseAttentionParams.sparse_mla_topk_lens != nullptr;
tllmRunnerParams.mSparseAttention
= useDynamicSparseMLA ? SparseType::DynamicTokenSparse : SparseType::StaticTokenSparse;
tllmRunnerParams.mSparseTopK = mRuntimeSparseAttentionParams.num_sparse_topk;
tllmRunnerParams.ptrSparseMlaTopKLens = mRuntimeSparseAttentionParams.sparse_mla_topk_lens;
tllmRunnerParams.kvPageIdxPtr = reinterpret_cast<KVCacheIndex::UnderlyingType const*>(
mRuntimeSparseAttentionParams.sparse_attn_indices);
tllmRunnerParams.kvPtr = mRuntimeSparseAttentionParams.sparse_kv_cache_pool;
if (useDynamicSparseMLA)
{
TLLM_CHECK_WITH_INFO(mRuntimeSparseAttentionParams.sliding_window_kv_cache_pool != nullptr,
"SWA KV pool must be set for dynamic sparse MLA.");
// Dynamic sparse MLA always has an SWA pool. The compressed pool is optional; when it
// is absent (ratio == 1), use SWA as kvPtr only to keep TG's primary TMA descriptor valid.
tllmRunnerParams.kvPtr = mRuntimeSparseAttentionParams.sparse_kv_cache_pool != nullptr
? mRuntimeSparseAttentionParams.sparse_kv_cache_pool
: mRuntimeSparseAttentionParams.sliding_window_kv_cache_pool;
tllmRunnerParams.slidingWindowKvPoolBasePtr
= mRuntimeSparseAttentionParams.sliding_window_kv_cache_pool;
}
else
{
tllmRunnerParams.kvPtr = mRuntimeSparseAttentionParams.sparse_kv_cache_pool;
}
}

mTllmGenFMHARunner->run(tllmRunnerParams);
Expand Down Expand Up @@ -1460,7 +1486,8 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
{
dim_q_per_head = mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
dim_k_per_head = mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
dim_v_per_head = mMLAParams.kv_lora_rank;
dim_v_per_head
= mMLAParams.rope_append ? mMLAParams.kv_lora_rank : mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
}

// Total dimension per token across all heads for Q, K, and V components respectively
Expand Down Expand Up @@ -2796,8 +2823,12 @@ int AttentionOp::initialize() noexcept
TLLM_CHECK_WITH_INFO(!mCrossAttention, "MLA(Deepseek v2) do not support cross attention right now");
TLLM_CHECK_WITH_INFO(mMaskType != tensorrt_llm::kernels::AttentionMaskType::CUSTOM_MASK,
"MLA(Deepseek v2) do not support custom mask right now");
TLLM_CHECK_WITH_INFO(mMLAParams.qk_rope_head_dim == 64 && mMLAParams.kv_lora_rank == 512,
"MLA(Deepseek v2) only support fixed kv_lora_rank(512) and fixed qk_rope_head_dim(64) right now.");
bool const mla_dims_supported = mMLAParams.qk_rope_head_dim == 64
&& ((mMLAParams.rope_append && mMLAParams.kv_lora_rank == 512)
|| (!mMLAParams.rope_append && mMLAParams.kv_lora_rank == 448));
TLLM_CHECK_WITH_INFO(mla_dims_supported,
"MLA(Deepseek v2) only supports qk_rope_head_dim=64 with kv_lora_rank=512 (rope_append=true) or "
"kv_lora_rank=448 (rope_append=false).");
Comment thread
lfr-0531 marked this conversation as resolved.
}

mDriver = CUDADriverWrapper::getInstance();
Expand Down Expand Up @@ -2930,7 +2961,8 @@ int AttentionOp::initialize() noexcept
fmhaParams.attentionInputLayout = AttentionInputLayout::Q_PAGED_KV;
fmhaParams.numKvHeads = 1;
fmhaParams.headSize = mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
fmhaParams.headSizeV = mMLAParams.kv_lora_rank;
fmhaParams.headSizeV = mMLAParams.rope_append ? mMLAParams.kv_lora_rank
: mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim;
fmhaParams.headSizeQkNope = mMLAParams.qk_nope_head_dim;
// Adjust the qScaling for the absorption mode.
fmhaParams.qScaling = mQScaling
Expand Down
28 changes: 24 additions & 4 deletions cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,10 @@ bool FmhaDispatcher::isSupported()
// Sparse context attention uses a generation-style kernel with per-token sparse indices.
if (mFixedParams.useTllmGenSparseAttention)
{
tllmRunnerParams.mSparseAttention = SparseType::StaticTokenSparse;
tllmRunnerParams.mSparseAttention
= (mFixedParams.useSparseMLA && mFixedParams.headSizeV == mFixedParams.headSize)
Comment thread
lfr-0531 marked this conversation as resolved.
? SparseType::DynamicTokenSparse
: SparseType::StaticTokenSparse;
tllmRunnerParams.mKernelType = FmhaKernelType::Generation;
Comment thread
lfr-0531 marked this conversation as resolved.
tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Causal;
Comment thread
lfr-0531 marked this conversation as resolved.
// Generation-style kernels on long KV can pick MultiCtasKv cubins
Expand Down Expand Up @@ -258,13 +261,30 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams)
// The kernel iterates over tokens via numCtasPerSeqQ when maxSeqLenQ > 1.
if (mFixedParams.useTllmGenSparseAttention)
{
tllmRunnerParams.mSparseAttention = SparseType::StaticTokenSparse;
tllmRunnerParams.mSparseTopK = runnerParams.sparse_params.sparse_topk;
bool const useDynamicSparseMLA = runnerParams.sparse_params.sparse_mla_topk_lens != nullptr;
tllmRunnerParams.mSparseAttention
= useDynamicSparseMLA ? SparseType::DynamicTokenSparse : SparseType::StaticTokenSparse;
tllmRunnerParams.mSparseTopK = runnerParams.sparse_params.num_sparse_topk;
tllmRunnerParams.ptrSparseMlaTopKLens = runnerParams.sparse_params.sparse_mla_topk_lens;
tllmRunnerParams.mKernelType = FmhaKernelType::Generation;
tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Causal;
tllmRunnerParams.kvPageIdxPtr
= reinterpret_cast<int const*>(runnerParams.sparse_params.sparse_attn_indices);
tllmRunnerParams.kvPtr = runnerParams.sparse_params.sparse_kv_cache_pool;
if (useDynamicSparseMLA)
{
TLLM_CHECK_WITH_INFO(runnerParams.sparse_params.sliding_window_kv_cache_pool != nullptr,
"SWA KV pool must be set for dynamic sparse MLA.");
// Dynamic sparse MLA always has an SWA pool. The compressed pool is optional; when it
// is absent (ratio == 1), use SWA as kvPtr only to keep TG's primary TMA descriptor valid.
tllmRunnerParams.kvPtr = runnerParams.sparse_params.sparse_kv_cache_pool != nullptr
? runnerParams.sparse_params.sparse_kv_cache_pool
: runnerParams.sparse_params.sliding_window_kv_cache_pool;
tllmRunnerParams.slidingWindowKvPoolBasePtr = runnerParams.sparse_params.sliding_window_kv_cache_pool;
}
else
{
tllmRunnerParams.kvPtr = runnerParams.sparse_params.sparse_kv_cache_pool;
}
// Enable MultiCtasKv so the autotuner can select GmemReduction cubins
tllmRunnerParams.mMultiCtasKvMode = true;
tllmRunnerParams.multiCtasKvScratchPtr = runnerParams.multiCtasKvScratchPtr;
Expand Down
Loading
Loading