Skip to content

[None][feat] Support TRTLLM external shared-KV MTP #18051

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
yuxianq wants to merge 3 commits into from
Closed

[None][feat] Support TRTLLM external shared-KV MTP #18051

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
2898661
[None][refactor] compose phased FMHA implementations
yuxianq Aug 20, 2026
ef26d02
[None][feat] support custom masks in TRTLLM attention
yuxianq Aug 20, 2026
4505cd3
[None][feat] support TRTLLM external shared-KV MTP
yuxianq Aug 20, 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
3 changes: 3 additions & 0 deletions cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,8 @@ struct QKVPreprocessingParams
bool separate_q_kv_output{false};
bool quantized_fp8_output{false};
bool generation_phase{false};
// The input contains Q only and the existing KV cache must not be updated.
bool q_only_input{false};
int multi_processor_count{0};
int rotary_vision_start{0};
int rotary_vision_length{0};
Expand Down Expand Up @@ -308,6 +310,7 @@ struct QKVPreprocessingParams
ss << "separate_q_kv_output: " << std::boolalpha << separate_q_kv_output << std::endl;
ss << "quantized_fp8_output: " << quantized_fp8_output << std::endl;
ss << "generation_phase: " << generation_phase << std::endl;
ss << "q_only_input: " << q_only_input << std::endl;
ss << "multi_processor_count: " << multi_processor_count << std::endl;

return ss.str();
Expand Down
96 changes: 59 additions & 37 deletions cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,7 +405,7 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams<T, KVCacheBuff
int const hidden_idx = head_idx * params.size_per_head + head_dim_idx;
int const kv_head_idx = head_idx / params.qheads_per_kv_head;
int const hidden_idx_kv = kv_head_idx * params.size_per_head + head_dim_idx;
int const hidden_size = params.hidden_size;
int const input_hidden_size = params.q_only_input ? params.q_hidden_size : params.hidden_size;
int const src_k_offset = params.q_hidden_size;
int const src_v_offset = src_k_offset + params.kv_head_num * params.size_per_head;

Expand Down Expand Up @@ -456,40 +456,49 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams<T, KVCacheBuff
// src QKV: [batch, time, 3, head_num, size_per_head]
// head_num != kv_head_num:
// src QKV: [batch, time, head_num * size_per_head + 2 * kv_head_num * size_per_head]
auto const src_q_idx = static_cast<size_t>(global_token_idx) * hidden_size + hidden_idx;
auto const src_k_idx = static_cast<size_t>(global_token_idx) * hidden_size + src_k_offset + hidden_idx_kv;
auto const src_v_idx = static_cast<size_t>(global_token_idx) * hidden_size + src_v_offset + hidden_idx_kv;

VecType q, k, v, q_pair, k_pair;
auto const src_q_idx = static_cast<size_t>(global_token_idx) * input_hidden_size + hidden_idx;
auto const src_k_idx
= static_cast<size_t>(global_token_idx) * input_hidden_size + src_k_offset + hidden_idx_kv;
auto const src_v_idx
= static_cast<size_t>(global_token_idx) * input_hidden_size + src_v_offset + hidden_idx_kv;

VecType q, q_pair;
VecType k{}, v{}, k_pair{};
// key without position embedding
VecType k_wo_pos;
VecType k_wo_pos{};

// load q,k,v and add bias
if (valid_head_dim_idx)
{
q = *reinterpret_cast<VecType const*>(&params.qkv_input[src_q_idx]);
k = *reinterpret_cast<VecType const*>(&params.qkv_input[src_k_idx]);
v = *reinterpret_cast<VecType const*>(&params.qkv_input[src_v_idx]);
q_pair = *reinterpret_cast<VecType const*>(&params.qkv_input[src_q_idx + rotated_head_dim_offset]);
k_pair = *reinterpret_cast<VecType const*>(&params.qkv_input[src_k_idx + rotated_head_dim_offset]);
if (!params.q_only_input)
{
k = *reinterpret_cast<VecType const*>(&params.qkv_input[src_k_idx]);
v = *reinterpret_cast<VecType const*>(&params.qkv_input[src_v_idx]);
k_pair = *reinterpret_cast<VecType const*>(&params.qkv_input[src_k_idx + rotated_head_dim_offset]);
}

if constexpr (ADD_BIAS)
{
auto const q_bias = *reinterpret_cast<VecType const*>(&params.qkv_bias[hidden_idx]);
auto const k_bias
= *reinterpret_cast<VecType const*>(&params.qkv_bias[hidden_idx_kv + src_k_offset]);
auto const v_bias
= *reinterpret_cast<VecType const*>(&params.qkv_bias[hidden_idx_kv + src_v_offset]);
auto const q_pair_bias
= *reinterpret_cast<VecType const*>(&params.qkv_bias[hidden_idx + rotated_head_dim_offset]);
auto const k_pair_bias = *reinterpret_cast<VecType const*>(
&params.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]);

q = mmha::add(q, q_bias);
k = mmha::add(k, k_bias);
v = mmha::add(v, v_bias);
q_pair = mmha::add(q_pair, q_pair_bias);
k_pair = mmha::add(k_pair, k_pair_bias);
if (!params.q_only_input)
{
auto const k_bias
= *reinterpret_cast<VecType const*>(&params.qkv_bias[hidden_idx_kv + src_k_offset]);
auto const v_bias
= *reinterpret_cast<VecType const*>(&params.qkv_bias[hidden_idx_kv + src_v_offset]);
auto const k_pair_bias = *reinterpret_cast<VecType const*>(
&params.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]);
k = mmha::add(k, k_bias);
v = mmha::add(v, v_bias);
k_pair = mmha::add(k_pair, k_pair_bias);
}
}
k_wo_pos = k;
}
Expand Down Expand Up @@ -550,7 +559,7 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams<T, KVCacheBuff
}
auto const channelIdx{tidx};

bool const useKVCache = params.kv_cache_buffer.data != nullptr;
bool const useKVCache = !params.q_only_input && params.kv_cache_buffer.data != nullptr;
auto token_idx_in_kv_cache = token_idx_in_seq;
bool valid_kv_cache_pos = useKVCache;

Expand Down Expand Up @@ -593,7 +602,9 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams<T, KVCacheBuff
{
*q_ptr = q;
}
if ((params.head_num == params.kv_head_num) || (head_idx == (kv_head_idx * params.qheads_per_kv_head)))
if (!params.q_only_input
&& ((params.head_num == params.kv_head_num)
|| (head_idx == (kv_head_idx * params.qheads_per_kv_head))))
{
if constexpr (STORE_QKV)
{
Expand Down Expand Up @@ -807,6 +818,7 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams<T, KVCacheBu
int const hidden_idx_kv = kv_head_idx * params.size_per_head + head_dim_idx;
int const src_k_offset = params.q_hidden_size;
int const src_v_offset = src_k_offset + params.kv_hidden_size;
int const input_hidden_size = params.q_only_input ? params.q_hidden_size : params.hidden_size;

int const rotated_head_dim_offset = first_half ? params.half_rotary_dim : -params.half_rotary_dim;
// Make sure there are multiple of tokens_per_block otherwise syncthreads will lead to deadlocks.
Expand Down Expand Up @@ -869,36 +881,44 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams<T, KVCacheBu
// src QKV: [batch, time, 3, head_num, size_per_head]
// head_num != kv_head_num:
// src QKV: [batch, time, head_num * size_per_head + 2 * kv_head_num * size_per_head]
auto const src_q_idx = static_cast<size_t>(bounded_global_token_idx) * params.hidden_size + hidden_idx;
auto const src_q_idx = static_cast<size_t>(bounded_global_token_idx) * input_hidden_size + hidden_idx;
auto const src_k_idx
= static_cast<size_t>(bounded_global_token_idx) * params.hidden_size + src_k_offset + hidden_idx_kv;
= static_cast<size_t>(bounded_global_token_idx) * input_hidden_size + src_k_offset + hidden_idx_kv;
auto const src_v_idx
= static_cast<size_t>(bounded_global_token_idx) * params.hidden_size + src_v_offset + hidden_idx_kv;
= static_cast<size_t>(bounded_global_token_idx) * input_hidden_size + src_v_offset + hidden_idx_kv;

auto q = *reinterpret_cast<VecT const*>(&params.qkv_input[src_q_idx]);
auto k = *reinterpret_cast<VecT const*>(&params.qkv_input[src_k_idx]);
auto v = *reinterpret_cast<VecT const*>(&params.qkv_input[src_v_idx]);
VecT k{};
VecT v{};
[[maybe_unused]] auto q_pair
= *reinterpret_cast<VecT const*>(&params.qkv_input[src_q_idx + rotated_head_dim_offset]);
[[maybe_unused]] auto k_pair
= *reinterpret_cast<VecT const*>(&params.qkv_input[src_k_idx + rotated_head_dim_offset]);
[[maybe_unused]] VecT k_pair{};
if (!params.q_only_input)
{
k = *reinterpret_cast<VecT const*>(&params.qkv_input[src_k_idx]);
v = *reinterpret_cast<VecT const*>(&params.qkv_input[src_v_idx]);
k_pair = *reinterpret_cast<VecT const*>(&params.qkv_input[src_k_idx + rotated_head_dim_offset]);
}

// Bias should have been fused with QKV projection, but we keep the logic here for unit tests.
if constexpr (ADD_BIAS)
{
auto const q_bias = *reinterpret_cast<VecT const*>(&params.qkv_bias[hidden_idx]);
auto const k_bias = *reinterpret_cast<VecT const*>(&params.qkv_bias[hidden_idx_kv + src_k_offset]);
auto const v_bias = *reinterpret_cast<VecT const*>(&params.qkv_bias[hidden_idx_kv + src_v_offset]);
auto const q_pair_bias
= *reinterpret_cast<VecT const*>(&params.qkv_bias[hidden_idx + rotated_head_dim_offset]);
auto const k_pair_bias = *reinterpret_cast<VecT const*>(
&params.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]);

q = mmha::add(q, q_bias);
k = mmha::add(k, k_bias);
v = mmha::add(v, v_bias);
q_pair = mmha::add(q_pair, q_pair_bias);
k_pair = mmha::add(k_pair, k_pair_bias);
if (!params.q_only_input)
{
auto const k_bias = *reinterpret_cast<VecT const*>(&params.qkv_bias[hidden_idx_kv + src_k_offset]);
auto const v_bias = *reinterpret_cast<VecT const*>(&params.qkv_bias[hidden_idx_kv + src_v_offset]);
auto const k_pair_bias = *reinterpret_cast<VecT const*>(
&params.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]);
k = mmha::add(k, k_bias);
v = mmha::add(v, v_bias);
k_pair = mmha::add(k_pair, k_pair_bias);
}
}

// Cos/sin cache.
Expand Down Expand Up @@ -967,7 +987,7 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams<T, KVCacheBu
}

auto const channelIdx = head_dim_vec_idx;
bool const useKVCache = GEN_PHASE || params.kv_cache_buffer.data != nullptr;
bool const useKVCache = !params.q_only_input && (GEN_PHASE || params.kv_cache_buffer.data != nullptr);
bool valid_kv_cache_pos = useKVCache;

auto kDst = useKVCache
Expand Down Expand Up @@ -1011,7 +1031,8 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams<T, KVCacheBu
{
*q_ptr = q;
}
if ((params.head_num == params.kv_head_num) || (head_idx == (kv_head_idx * params.qheads_per_kv_head)))
if (!params.q_only_input
&& ((params.head_num == params.kv_head_num) || (head_idx == (kv_head_idx * params.qheads_per_kv_head))))
{
if constexpr (STORE_QKV)
{
Expand Down Expand Up @@ -1646,6 +1667,7 @@ void invokeApplyBiasRopeUpdateKVCacheDispatch(QKVPreprocessingParams<T, KVCacheB
case 192: kernelV2DispatchHeadSize<192, 192, T, TCache, KVCacheBuffer>(params, stream); break;
case 224: kernelV2DispatchHeadSize<224, 224, T, TCache, KVCacheBuffer>(params, stream); break;
case 256: kernelV2DispatchHeadSize<256, 256, T, TCache, KVCacheBuffer>(params, stream); break;
case 512: kernelV2DispatchHeadSize<512, 512, T, TCache, KVCacheBuffer>(params, stream); break;
case 576: kernelV2DispatchHeadSize<576, 576, T, TCache, KVCacheBuffer>(params, stream); break;
default:
// Fall back to v1 kernel.
Expand Down
16 changes: 15 additions & 1 deletion cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,12 @@ cudaStream_t currentStreamFor(at::Tensor const& tensor)
return at::cuda::getCurrentCUDAStream(tensor.get_device()).stream();
}

bool isQOnlyInput(at::Tensor const& qkvInput, int64_t const numHeads, int64_t const headSize)
{
TORCH_CHECK(qkvInput.dim() > 0, "QKV input must have at least one dimension.");
return qkvInput.size(-1) == numHeads * headSize;
}

struct WorkspaceAccessor
{
uint8_t* base{};
Expand Down Expand Up @@ -276,7 +282,8 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
TORCH_CHECK(kv_cache_block_offsets.has_value(), "kv_cache_block_offsets is required.");
TORCH_CHECK(!cross_attention || !is_mla_enable, "trtllm-gen cross attention does not support MLA.");

bool const separateQKvOutput = paged_context_fmha || fp8_context_fmha || cross_attention;
bool const qOnlyInput = !cross_attention && isQOnlyInput(qkv_input, num_heads, head_size);
bool const separateQKvOutput = paged_context_fmha || fp8_context_fmha || cross_attention || qOnlyInput;
auto const qkvScalarType = qkv_input.scalar_type();
auto const qkvElementSize = static_cast<size_t>(qkv_input.element_size());
auto const quantMode = tensorrt_llm::common::QuantMode(static_cast<uint32_t>(kv_cache_quant_mode));
Expand Down Expand Up @@ -396,6 +403,7 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
qkvParams.separate_q_kv_output = separateQKvOutput;
qkvParams.quantized_fp8_output = fp8_context_fmha;
qkvParams.generation_phase = false;
qkvParams.q_only_input = qOnlyInput;
qkvParams.multi_processor_count = static_cast<int>(multi_processor_count);
qkvParams.rotary_vision_start = 0;
qkvParams.rotary_vision_length = 0;
Expand Down Expand Up @@ -470,6 +478,10 @@ void trtllmGenContextPostprocess(torch::Tensor qkv_input, torch::Tensor workspac
int64_t const attention_chunk_size, int64_t const multi_processor_count)
{
(void) mask_type;
if (isQOnlyInput(qkv_input, num_heads, head_size))
{
return;
}
auto const qkvScalarType = qkv_input.scalar_type();
auto const qkvElementSize = static_cast<size_t>(qkv_input.element_size());
auto const quantMode = tensorrt_llm::common::QuantMode(static_cast<uint32_t>(kv_cache_quant_mode));
Expand Down Expand Up @@ -596,6 +608,7 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
(void) bmm2_scale;

bool const isMultiTokenGen = spec_decoding_generation_lengths.has_value() && predicted_tokens_per_seq > 1;
bool const qOnlyInput = !cross_attention && isQOnlyInput(qkv_input, num_heads, head_size);
TORCH_CHECK(
!cross_attention || !isMultiTokenGen, "trtllm-gen cross attention does not support multi-token generation.");
auto const qkvScalarType = qkv_input.scalar_type();
Expand Down Expand Up @@ -728,6 +741,7 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
qkvParams.separate_q_kv_output = true;
qkvParams.quantized_fp8_output = fp8_context_fmha;
qkvParams.generation_phase = true;
qkvParams.q_only_input = qOnlyInput;
qkvParams.multi_processor_count = static_cast<int>(multi_processor_count);
qkvParams.rotary_vision_start = 0;
qkvParams.rotary_vision_length = 0;
Expand Down
14 changes: 14 additions & 0 deletions tensorrt_llm/_torch/attention_backend/flashinfer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -650,6 +650,20 @@ def get_draft_metadata(
"views.")
return self._draft_metadata

def get_shared_kv_draft_metadata(
self,
num_accepted_tokens: torch.Tensor,
num_contexts: int,
) -> "FlashInferAttentionMetadata":
"""Return the shared-target-KV draft view at the accepted prefix."""
draft_metadata = self.get_draft_metadata()
draft_metadata.update_shared_kv_draft_lengths(
self,
num_accepted_tokens,
num_contexts,
)
return draft_metadata

def _sync_draft_view(self, target: "FlashInferAttentionMetadata") -> None:
"""Refresh a shared- or separate-KV draft metadata view."""
if not (self._is_shared_kv_draft_view
Expand Down
7 changes: 6 additions & 1 deletion tensorrt_llm/_torch/attention_backend/fmha/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,24 +13,29 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from .combined import CombinedFmha
from .cute_dsl_mla import CuteDslMlaFmha
from .fallback import FallbackFmha
from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha
from .interface import Fmha
from .interface import Fmha, FmhaPhase
from .msa_sparse_gqa import MsaSparseGqaFmha
from .phased import FmhaParams, PhasedFmha
from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes
from .triton_custom_mask import TritonCustomMaskFmha

__all__ = [
"DEFAULT_FMHA_LIBS",
"FMHA_LIBS",
"CombinedFmha",
"CuteDslMlaFmha",
"FallbackFmha",
"FlashInferTrtllmGenFmha",
"Fmha",
"FmhaCls",
"FmhaParams",
"FmhaPhase",
"MsaSparseGqaFmha",
"PhasedFmha",
"TritonCustomMaskFmha",
"get_enabled_fmha_lib_classes",
]
Loading
Loading