Skip to content

Commit 660aede

Browse files
cascade812cascade812
authored
[TRTLLM-13490][feat] Support cross-attention with FlashInfer TRTLLM-Gen kernels on Blackwell (#15429)
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
1 parent ccabbea commit 660aede

12 files changed

Lines changed: 246 additions & 171 deletions

File tree

‎cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1420,7 +1420,7 @@ __global__ void updateKVCacheForCrossAttention(QKVPreprocessingParams<T, KVCache
14201420
int const max_seq_len = max(decoder_seq_len, encoder_seq_len);
14211421

14221422
// Only the first chunk needs to store encoder kv input to the kv cache.
1423-
bool const store_encoder_kv_cache = (decoder_seq_len == decoder_cache_seq_len);
1423+
bool const store_encoder_kv_cache = params.cross_kv_input != nullptr && (decoder_seq_len == decoder_cache_seq_len);
14241424

14251425
// Offsets and strides.
14261426
int const head_dim_vec_idx = (threadIdx.x % VECS_PER_HEAD);

‎cpp/tensorrt_llm/nanobind/thop/bindings.cpp‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ nb::tuple trtllmGenContextPreprocessBinding(torch::Tensor qkv_input, torch::Tens
5757
double rotary_embedding_scale, int64_t rotary_embedding_max_positions, int64_t position_embedding_type,
5858
double bmm1_scale, double bmm2_scale, int64_t attention_chunk_size, bool fp8_context_fmha, bool paged_context_fmha,
5959
bool is_mla_enable, int64_t multi_processor_count, int64_t total_num_blocks, int64_t kv_factor,
60-
bool need_build_kv_cache_metadata)
60+
bool need_build_kv_cache_metadata, std::optional<torch::Tensor> cross_kv, bool cross_attention)
6161
{
6262
auto result = [&]()
6363
{
@@ -70,7 +70,7 @@ nb::tuple trtllmGenContextPreprocessBinding(torch::Tensor qkv_input, torch::Tens
7070
max_past_kv_length, rotary_embedding_dim, rotary_embedding_base, rotary_embedding_scale_type,
7171
rotary_embedding_scale, rotary_embedding_max_positions, position_embedding_type, bmm1_scale, bmm2_scale,
7272
attention_chunk_size, fp8_context_fmha, paged_context_fmha, is_mla_enable, multi_processor_count,
73-
total_num_blocks, kv_factor, need_build_kv_cache_metadata);
73+
total_num_blocks, kv_factor, need_build_kv_cache_metadata, cross_kv, cross_attention);
7474
}();
7575

7676
return nb::make_tuple(std::get<0>(result), optionalToObject(std::get<1>(result)),
@@ -92,7 +92,7 @@ nb::tuple trtllmGenGenerationPreprocessBinding(torch::Tensor qkv_input, torch::T
9292
int64_t rotary_embedding_scale_type, double rotary_embedding_scale, int64_t rotary_embedding_max_positions,
9393
int64_t position_embedding_type, double bmm1_scale, double bmm2_scale, bool fp8_context_fmha,
9494
int64_t predicted_tokens_per_seq, int64_t attention_chunk_size, int64_t multi_processor_count,
95-
int64_t total_num_blocks, int64_t kv_factor, bool need_build_kv_cache_metadata)
95+
int64_t total_num_blocks, int64_t kv_factor, bool need_build_kv_cache_metadata, bool cross_attention)
9696
{
9797
auto result = [&]()
9898
{
@@ -106,7 +106,7 @@ nb::tuple trtllmGenGenerationPreprocessBinding(torch::Tensor qkv_input, torch::T
106106
rotary_embedding_dim, rotary_embedding_base, rotary_embedding_scale_type, rotary_embedding_scale,
107107
rotary_embedding_max_positions, position_embedding_type, bmm1_scale, bmm2_scale, fp8_context_fmha,
108108
predicted_tokens_per_seq, attention_chunk_size, multi_processor_count, total_num_blocks, kv_factor,
109-
need_build_kv_cache_metadata);
109+
need_build_kv_cache_metadata, cross_attention);
110110
}();
111111

112112
return nb::make_tuple(std::get<0>(result), optionalToObject(std::get<1>(result)),
@@ -273,7 +273,8 @@ void initBindings(nb::module_& m)
273273
nb::arg("position_embedding_type"), nb::arg("bmm1_scale"), nb::arg("bmm2_scale"),
274274
nb::arg("attention_chunk_size"), nb::arg("fp8_context_fmha"), nb::arg("paged_context_fmha"),
275275
nb::arg("is_mla_enable"), nb::arg("multi_processor_count"), nb::arg("total_num_blocks"), nb::arg("kv_factor"),
276-
nb::arg("need_build_kv_cache_metadata") = true, "Fused nanobind context preprocess for trtllm-gen attention.");
276+
nb::arg("need_build_kv_cache_metadata") = true, nb::arg("cross_kv").none() = nb::none(),
277+
nb::arg("cross_attention") = false, "Fused nanobind context preprocess for trtllm-gen attention.");
277278

278279
m.def("trtllm_gen_context_postprocess", &torch_ext::trtllmGenContextPostprocess, nb::arg("qkv_input"),
279280
nb::arg("workspace"), nb::arg("sequence_lengths"), nb::arg("context_lengths"),
@@ -332,6 +333,6 @@ void initBindings(nb::module_& m)
332333
nb::arg("position_embedding_type"), nb::arg("bmm1_scale"), nb::arg("bmm2_scale"), nb::arg("fp8_context_fmha"),
333334
nb::arg("predicted_tokens_per_seq"), nb::arg("attention_chunk_size"), nb::arg("multi_processor_count"),
334335
nb::arg("total_num_blocks"), nb::arg("kv_factor"), nb::arg("need_build_kv_cache_metadata") = true,
335-
"Fused nanobind generation preprocess for trtllm-gen attention.");
336+
nb::arg("cross_attention") = false, "Fused nanobind generation preprocess for trtllm-gen attention.");
336337
}
337338
} // namespace tensorrt_llm::nanobind::thop

‎cpp/tensorrt_llm/thop/trtllmGenFusedOps.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
4242
double rotary_embedding_scale, int64_t rotary_embedding_max_positions, int64_t position_embedding_type,
4343
double bmm1_scale, double bmm2_scale, int64_t attention_chunk_size, bool fp8_context_fmha, bool paged_context_fmha,
4444
bool is_mla_enable, int64_t multi_processor_count, int64_t total_num_blocks, int64_t kv_factor,
45-
bool need_build_kv_cache_metadata);
45+
bool need_build_kv_cache_metadata, std::optional<torch::Tensor> cross_kv = std::nullopt,
46+
bool cross_attention = false);
4647

4748
void trtllmGenContextPostprocess(torch::Tensor qkv_input, torch::Tensor workspace, torch::Tensor sequence_lengths,
4849
torch::Tensor context_lengths, std::optional<torch::Tensor> kv_cache_block_offsets,
@@ -72,7 +73,7 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
7273
int64_t rotary_embedding_scale_type, double rotary_embedding_scale, int64_t rotary_embedding_max_positions,
7374
int64_t position_embedding_type, double bmm1_scale, double bmm2_scale, bool fp8_context_fmha,
7475
int64_t predicted_tokens_per_seq, int64_t attention_chunk_size, int64_t multi_processor_count,
75-
int64_t total_num_blocks, int64_t kv_factor, bool need_build_kv_cache_metadata);
76+
int64_t total_num_blocks, int64_t kv_factor, bool need_build_kv_cache_metadata, bool cross_attention = false);
7677

7778
} // namespace torch_ext
7879

‎cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp‎

Lines changed: 35 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -266,17 +266,22 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
266266
int64_t const position_embedding_type, double const bmm1_scale, double const bmm2_scale,
267267
int64_t const attention_chunk_size, bool const fp8_context_fmha, bool const paged_context_fmha,
268268
bool const is_mla_enable, int64_t const multi_processor_count, int64_t const total_num_blocks,
269-
int64_t const kv_factor, bool const need_build_kv_cache_metadata)
269+
int64_t const kv_factor, bool const need_build_kv_cache_metadata, std::optional<torch::Tensor> cross_kv,
270+
bool const cross_attention)
270271
{
271272
(void) bmm2_scale;
272273
TORCH_CHECK(host_kv_cache_pool_pointers.has_value(), "host_kv_cache_pool_pointers is required.");
273274
TORCH_CHECK(host_kv_cache_pool_mapping.has_value(), "host_kv_cache_pool_mapping is required.");
274275
TORCH_CHECK(kv_cache_block_offsets.has_value(), "kv_cache_block_offsets is required.");
276+
TORCH_CHECK(!cross_attention || !is_mla_enable, "trtllm-gen cross attention does not support MLA.");
275277

276-
bool const separateQKvOutput = paged_context_fmha || fp8_context_fmha;
278+
bool const separateQKvOutput = paged_context_fmha || fp8_context_fmha || cross_attention;
277279
auto const qkvScalarType = qkv_input.scalar_type();
278280
auto const qkvElementSize = static_cast<size_t>(qkv_input.element_size());
279281
auto const quantMode = tensorrt_llm::common::QuantMode(static_cast<uint32_t>(kv_cache_quant_mode));
282+
int64_t const effectiveMaxAttentionWindowSize = cross_attention ? max_past_kv_length : max_attention_window_size;
283+
int64_t const effectiveCyclicAttentionWindowSize
284+
= cross_attention ? max_past_kv_length : cyclic_attention_window_size;
280285
auto const views = [&]
281286
{
282287
auto const layout = TrtllmAttentionWorkspaceManager::buildContextLayout(
@@ -307,8 +312,8 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
307312
decoderInfoParams.fmhaBmm2Scale = ptrs.fmhaBmm2ScalePtr;
308313
decoderInfoParams.batchSize = static_cast<int>(batch_size);
309314
decoderInfoParams.maxQSeqLength = static_cast<int>(input_seq_length);
310-
decoderInfoParams.maxEncoderQSeqLength = 0;
311-
decoderInfoParams.attentionWindowSize = static_cast<int>(cyclic_attention_window_size);
315+
decoderInfoParams.maxEncoderQSeqLength = cross_attention ? static_cast<int>(max_past_kv_length) : 0;
316+
decoderInfoParams.attentionWindowSize = static_cast<int>(effectiveCyclicAttentionWindowSize);
312317
decoderInfoParams.numTokens = static_cast<int>(num_tokens);
313318
decoderInfoParams.removePadding = true;
314319
decoderInfoParams.attentionMaskType = static_cast<AttentionMaskType>(mask_type);
@@ -333,12 +338,13 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
333338
{
334339
return buildPagedKvCacheBuffers(kv_cache_block_offsets, host_kv_cache_pool_pointers,
335340
host_kv_cache_pool_mapping, quantMode, layer_idx, batch_size, tokens_per_block, num_kv_heads, head_size,
336-
cyclic_attention_window_size, max_attention_window_size, 0, 0, is_mla_enable, qkvElementSize);
341+
effectiveCyclicAttentionWindowSize, effectiveMaxAttentionWindowSize, 0, 0, is_mla_enable,
342+
qkvElementSize);
337343
}();
338344

339345
QKVPreprocessingParams<void, KVBlockArray> qkvParams{};
340346
qkvParams.qkv_input = qkv_input.data_ptr();
341-
qkvParams.cross_kv_input = nullptr;
347+
qkvParams.cross_kv_input = optPtr<void>(cross_kv);
342348
qkvParams.quantized_qkv_output = nullptr;
343349
qkvParams.q_output = ptrs.qBufPtr;
344350
qkvParams.kv_cache_buffer = kvArrays.kvCacheBuffer;
@@ -353,8 +359,9 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
353359
qkvParams.logn_scaling = nullptr;
354360
qkvParams.tokens_info = ptrs.tokensInfoPtr;
355361
qkvParams.seq_lens = static_cast<int*>(context_lengths.data_ptr());
356-
qkvParams.cache_seq_lens = static_cast<int*>(sequence_lengths.data_ptr());
357-
qkvParams.encoder_seq_lens = nullptr;
362+
qkvParams.cache_seq_lens = cross_attention ? static_cast<int*>(context_lengths.data_ptr())
363+
: static_cast<int*>(sequence_lengths.data_ptr());
364+
qkvParams.encoder_seq_lens = cross_attention ? static_cast<int*>(sequence_lengths.data_ptr()) : nullptr;
358365
qkvParams.cu_seq_lens = ptrs.cuQSeqlensPtr;
359366
qkvParams.cu_kv_seq_lens = ptrs.cuKvSeqlensPtr;
360367
qkvParams.sparse_kv_offsets = nullptr;
@@ -367,11 +374,11 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
367374
qkvParams.batch_size = static_cast<int>(batch_size);
368375
qkvParams.max_input_seq_len = static_cast<int>(input_seq_length);
369376
qkvParams.max_kv_seq_len = static_cast<int>(max_past_kv_length);
370-
qkvParams.cyclic_kv_cache_len = static_cast<int>(cyclic_attention_window_size);
377+
qkvParams.cyclic_kv_cache_len = static_cast<int>(effectiveCyclicAttentionWindowSize);
371378
qkvParams.token_num = static_cast<int>(num_tokens);
372379
qkvParams.remove_padding = true;
373380
qkvParams.is_last_chunk = attention_chunk_size == 0 || input_seq_length == max_past_kv_length;
374-
qkvParams.cross_attention = false;
381+
qkvParams.cross_attention = cross_attention;
375382
qkvParams.head_num = static_cast<int>(num_heads);
376383
qkvParams.kv_head_num = static_cast<int>(num_kv_heads);
377384
qkvParams.qheads_per_kv_head = static_cast<int>(num_heads / num_kv_heads);
@@ -438,7 +445,9 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
438445

439446
// FlashInfer paged context launches trtllm-gen with multi-CTA-KV mode disabled, so it does not
440447
// consume the counter slab reserved at the head of the workspace.
441-
auto const windowLeft = computeWindowLeft(cyclic_attention_window_size, max_past_kv_length, attention_chunk_size);
448+
auto const windowLeft = cross_attention
449+
? int64_t{-1}
450+
: computeWindowLeft(cyclic_attention_window_size, max_past_kv_length, attention_chunk_size);
442451
return {qProcessed, kvPool, blockTables, kvScalePool, views.fmhaBmm1Scale, views.fmhaBmm2Scale,
443452
views.trtllmGenWorkspace, views.cuQSeqlens, views.cuKvSeqlens, input_seq_length, max_past_kv_length,
444453
windowLeft};
@@ -577,17 +586,22 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
577586
int64_t const rotary_embedding_max_positions, int64_t const position_embedding_type, double const bmm1_scale,
578587
double const bmm2_scale, bool const fp8_context_fmha, int64_t const predicted_tokens_per_seq,
579588
int64_t const attention_chunk_size, int64_t const multi_processor_count, int64_t const total_num_blocks,
580-
int64_t const kv_factor, bool const need_build_kv_cache_metadata)
589+
int64_t const kv_factor, bool const need_build_kv_cache_metadata, bool const cross_attention)
581590
{
582591
TORCH_CHECK(host_kv_cache_pool_pointers.has_value(), "host_kv_cache_pool_pointers is required.");
583592
TORCH_CHECK(host_kv_cache_pool_mapping.has_value(), "host_kv_cache_pool_mapping is required.");
584593
TORCH_CHECK(kv_cache_block_offsets.has_value(), "kv_cache_block_offsets is required.");
585594
(void) bmm2_scale;
586595

587596
bool const isMultiTokenGen = spec_decoding_generation_lengths.has_value() && predicted_tokens_per_seq > 1;
597+
TORCH_CHECK(
598+
!cross_attention || !isMultiTokenGen, "trtllm-gen cross attention does not support multi-token generation.");
588599
auto const qkvScalarType = qkv_input.scalar_type();
589600
auto const qkvElementSize = static_cast<size_t>(qkv_input.element_size());
590601
auto const quantMode = tensorrt_llm::common::QuantMode(static_cast<uint32_t>(kv_cache_quant_mode));
602+
int64_t const effectiveMaxAttentionWindowSize = cross_attention ? max_past_kv_length : max_attention_window_size;
603+
int64_t const effectiveCyclicAttentionWindowSize
604+
= cross_attention ? max_past_kv_length : cyclic_attention_window_size;
591605
auto const views = [&]
592606
{
593607
auto const layout = TrtllmAttentionWorkspaceManager::buildGenerationLayout(
@@ -617,7 +631,7 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
617631
decoderInfoParams.fmhaBmm2Scale = nullptr;
618632
decoderInfoParams.batchSize = static_cast<int>(batch_beam);
619633
decoderInfoParams.maxQSeqLength = static_cast<int>(input_seq_length);
620-
decoderInfoParams.maxEncoderQSeqLength = 0;
634+
decoderInfoParams.maxEncoderQSeqLength = cross_attention ? static_cast<int>(max_past_kv_length) : 0;
621635
decoderInfoParams.attentionWindowSize = 0;
622636
decoderInfoParams.sinkTokenLength = 0;
623637
decoderInfoParams.numTokens = static_cast<int>(num_tokens);
@@ -655,7 +669,8 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
655669
{
656670
return buildPagedKvCacheBuffers(kv_cache_block_offsets, host_kv_cache_pool_pointers,
657671
host_kv_cache_pool_mapping, quantMode, layer_idx, batch_beam, tokens_per_block, num_kv_heads, head_size,
658-
cyclic_attention_window_size, max_attention_window_size, 1, seq_offset, false, qkvElementSize);
672+
effectiveCyclicAttentionWindowSize, effectiveMaxAttentionWindowSize, 1, seq_offset, false,
673+
qkvElementSize);
659674
}();
660675

661676
QKVPreprocessingParams<void, KVBlockArray> qkvParams{};
@@ -676,7 +691,7 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
676691
qkvParams.tokens_info = isMultiTokenGen ? views.tokensInfoPtr : nullptr;
677692
qkvParams.seq_lens = isMultiTokenGen ? optPtr<int>(spec_decoding_generation_lengths) : nullptr;
678693
qkvParams.cache_seq_lens = static_cast<int*>(sequence_lengths.data_ptr());
679-
qkvParams.encoder_seq_lens = nullptr;
694+
qkvParams.encoder_seq_lens = cross_attention ? static_cast<int*>(sequence_lengths.data_ptr()) : nullptr;
680695
qkvParams.cu_seq_lens = buildDecoderInfoNeeded ? views.cuSeqlensPtr : nullptr;
681696
qkvParams.cu_kv_seq_lens = buildDecoderInfoNeeded ? views.cuKvSeqlensPtr : nullptr;
682697
qkvParams.sparse_kv_offsets = nullptr;
@@ -690,11 +705,11 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
690705
qkvParams.batch_size = static_cast<int>(batch_beam);
691706
qkvParams.max_input_seq_len = static_cast<int>(input_seq_length);
692707
qkvParams.max_kv_seq_len = static_cast<int>(max_past_kv_length);
693-
qkvParams.cyclic_kv_cache_len = static_cast<int>(cyclic_attention_window_size);
708+
qkvParams.cyclic_kv_cache_len = static_cast<int>(effectiveCyclicAttentionWindowSize);
694709
qkvParams.token_num = static_cast<int>(num_tokens);
695710
qkvParams.remove_padding = true;
696711
qkvParams.is_last_chunk = false;
697-
qkvParams.cross_attention = false;
712+
qkvParams.cross_attention = cross_attention;
698713
qkvParams.head_num = static_cast<int>(num_heads);
699714
qkvParams.kv_head_num = static_cast<int>(num_kv_heads);
700715
qkvParams.qheads_per_kv_head = static_cast<int>(num_heads / num_kv_heads);
@@ -751,7 +766,9 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
751766

752767
auto qProcessed = views.qBuf.view({num_tokens, num_heads, head_size});
753768

754-
auto const windowLeft = computeWindowLeft(cyclic_attention_window_size, max_past_kv_length, attention_chunk_size);
769+
auto const windowLeft = cross_attention
770+
? int64_t{-1}
771+
: computeWindowLeft(cyclic_attention_window_size, max_past_kv_length, attention_chunk_size);
755772
return {qProcessed, kvPool, blockTables, kvScalePool, views.bmm1Scale, views.bmm2Scale, views.trtllmGenWorkspace,
756773
cuSeqlens, input_seq_length, max_past_kv_length, windowLeft, isMultiTokenGen};
757774
}

‎tensorrt_llm/_torch/attention_backend/fmha/fallback.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ def forward(
9494
block_ids_per_seq=metadata.block_ids_per_seq,
9595
tokens_per_block=metadata.tokens_per_block,
9696
max_num_requests=metadata.max_num_requests,
97-
beam_width=metadata.effective_beam_width,
97+
beam_width=metadata.beam_width,
9898
use_paged_context_fmha=metadata.use_paged_context_fmha,
9999
helix_position_offsets=metadata.helix_position_offsets,
100100
helix_is_inactive_rank=metadata.helix_is_inactive_rank,

0 commit comments

Comments
 (0)