Skip to content

Commit d21f87f

Browse files
PerkzZhengPerkzZheng
committed
add chunked-attention kernels for blackwell
Signed-off-by: Perkz Zheng <[email protected]> fix Signed-off-by: Perkz Zheng <[email protected]>
1 parent f65ed52 commit d21f87f

File tree

7 files changed

+65
-20
lines changed

7 files changed

+65
-20
lines changed

‎cpp/tensorrt_llm/common/attentionOp.cpp

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -694,7 +694,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
694694
else if (mFmhaDispatcher->isSeparateQAndKvInput())
695695
{
696696
// Paged context fmha
697-
q_buf_2_size = (mFP8ContextFMHA ? 1 : size) * max_num_tokens * local_hidden_units_qo;
697+
q_buf_2_size = (isFP8ContextFmhaInput() ? 1 : size) * max_num_tokens * local_hidden_units_qo;
698698
}
699699

700700
size_t const k_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * kv_seq_length * local_hidden_units_kv;
@@ -705,16 +705,16 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
705705
size_t const qk_buf_float_size
706706
= mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_length * kv_seq_length;
707707
size_t const fp8_qkv_buffer_size
708-
= mFP8ContextFMHA && mEnableContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput()
708+
= isFP8ContextFmhaInput() && mEnableContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput()
709709
? max_num_tokens * size_t(local_hidden_units_qo + 2 * local_hidden_units_kv)
710710
: 0;
711711
size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens;
712712
size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens;
713713
// Each token holds (batch_idx, token_idx_in_seq) int2.
714714
size_t const tokens_info_size = sizeof(int2) * max_num_tokens;
715715
size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0;
716-
size_t const fmha_bmm1_scale_size = mFP8ContextFMHA ? sizeof(float) * 2 : 0;
717-
size_t const fmha_bmm2_scale_size = mFP8ContextFMHA ? sizeof(float) : 0;
716+
size_t const fmha_bmm1_scale_size = isFP8ContextFmhaInput() ? sizeof(float) * 2 : 0;
717+
size_t const fmha_bmm2_scale_size = isFP8ContextFmhaInput() ? sizeof(float) : 0;
718718

719719
// cp workspace size upper bound
720720
size_t const cpMaxPaddedSequenceLength = max_num_tokens + batch_size * (mCpSize - 1);
@@ -1254,6 +1254,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
12541254
= mEnableContextFMHA ? 0 : sizeof(T) * params.batch_size * params.input_seq_length * kv_seq_length;
12551255
size_t const cu_seqlens_size = sizeof(int) * (params.batch_size + 1);
12561256
size_t const rotary_inv_freq_size = sizeof(float) * params.batch_size * mRotaryEmbeddingDim / 2;
1257+
TLLM_CHECK_WITH_INFO(isFP8ContextFmhaInput() == mFmhaDispatcher->isFP8InputQ(), "internal error");
12571258
size_t q_buf_2_size = 0;
12581259
if (!mEnableContextFMHA)
12591260
{
@@ -1263,7 +1264,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
12631264
else if (mFmhaDispatcher->isSeparateQAndKvInput())
12641265
{
12651266
// Paged context fmha
1266-
q_buf_2_size = (mFP8ContextFMHA ? 1 : sizeof(T)) * params.num_tokens * local_hidden_units_qo;
1267+
q_buf_2_size = (isFP8ContextFmhaInput() ? 1 : sizeof(T)) * params.num_tokens * local_hidden_units_qo;
12671268
}
12681269

12691270
size_t const k_buf_2_size
@@ -1278,7 +1279,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
12781279
? 0
12791280
: sizeof(float) * params.batch_size * mNumHeads * params.input_seq_length * kv_seq_length;
12801281
size_t const fp8_qkv_buffer_size
1281-
= mEnableContextFMHA && mFP8ContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput()
1282+
= mEnableContextFMHA && isFP8ContextFmhaInput() && !mFmhaDispatcher->isSeparateQAndKvInput()
12821283
? params.num_tokens * (local_hidden_units_qo + 2 * local_hidden_units_kv)
12831284
: 0;
12841285
size_t const padding_offset_size
@@ -1288,8 +1289,8 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
12881289
// Each token holds (batch_idx, token_idx_in_seq) int2.
12891290
size_t const tokens_info_size = sizeof(int2) * params.num_tokens;
12901291
size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0;
1291-
size_t const fmha_bmm1_scale_size = mFP8ContextFMHA ? sizeof(float) * 2 : 0;
1292-
size_t const fmha_bmm2_scale_size = mFP8ContextFMHA ? sizeof(float) : 0;
1292+
size_t const fmha_bmm1_scale_size = isFP8ContextFmhaInput() ? sizeof(float) * 2 : 0;
1293+
size_t const fmha_bmm2_scale_size = isFP8ContextFmhaInput() ? sizeof(float) : 0;
12931294

12941295
// cp workspace size upper bound
12951296
size_t const cpMaxPadedSequenceLength = params.num_tokens + params.batch_size * (mCpSize - 1);
@@ -1514,7 +1515,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
15141515
preprocessingParams.position_shift_enabled = mPosShiftEnabled;
15151516
preprocessingParams.cache_type = cache_type;
15161517
preprocessingParams.separate_q_kv_output = enablePagedKVContextFMHA || isCrossAttention();
1517-
preprocessingParams.quantized_fp8_output = mFP8ContextFMHA;
1518+
preprocessingParams.quantized_fp8_output = isFP8ContextFmhaInput();
15181519
preprocessingParams.generation_phase = false;
15191520
preprocessingParams.multi_processor_count = mMultiProcessorCount;
15201521

@@ -1614,8 +1615,8 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
16141615
// TODO: set it correctly for contiguous kv buffer (cross-attention).
16151616
fmhaParams.totalKvSeqLen = isCrossAttention() ? params.num_encoder_tokens : params.num_tokens;
16161617
// Device buffer pointers.
1617-
fmhaParams.qkvPtr = mFP8ContextFMHA ? reinterpret_cast<void const*>(fp8_qkv_buffer)
1618-
: reinterpret_cast<void const*>(attention_input);
1618+
fmhaParams.qkvPtr = isFP8ContextFmhaInput() ? reinterpret_cast<void const*>(fp8_qkv_buffer)
1619+
: reinterpret_cast<void const*>(attention_input);
16191620
fmhaParams.qPtr = reinterpret_cast<void const*>(q_buf_2_);
16201621
// TODO: add contiguous kv buffer (cross-attention).
16211622
fmhaParams.kvPtr = nullptr;
@@ -2423,6 +2424,12 @@ int AttentionOp::initialize() noexcept
24232424
if (mKVCacheQuantMode.hasFp8KvCache())
24242425
{
24252426
fmhaParams.dataTypeKv = DATA_TYPE_E4M3;
2427+
// Trtllm-gen kernels forces using FP8 FMHA kernels if fp8 kv cache is used.
2428+
// FP8 Q/KV input and FP8/BF16/FP16 output are supported.
2429+
if (mUseTllmGen)
2430+
{
2431+
fmhaParams.dataType = DATA_TYPE_E4M3;
2432+
}
24262433
}
24272434
// TODO: add FP4 KV cache support.
24282435
}

‎cpp/tensorrt_llm/common/attentionOp.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,17 @@ class AttentionOp
258258
return mPositionEmbeddingType == tensorrt_llm::kernels::PositionEmbeddingType::kALIBI_WITH_SCALE;
259259
}
260260

261+
[[nodiscard]] bool isFP8ContextFmhaInput() const
262+
{
263+
if (mUseTllmGen)
264+
{
265+
// Trtllm-gen kernels forces using FP8 FMHA kernels if fp8 kv cache is used.
266+
// FP8 Q/KV input and FP8/BF16/FP16 output are supported.
267+
return mKVCacheQuantMode.hasFp8KvCache();
268+
}
269+
return mFP8ContextFMHA;
270+
}
271+
261272
[[nodiscard]] bool isRoPE() const
262273
{
263274
return mPositionEmbeddingType == tensorrt_llm::kernels::PositionEmbeddingType::kROPE_GPTJ

‎cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,8 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams)
183183
tllmRunnerParams.mMaxSeqLenQ = runnerParams.qSeqLen;
184184
tllmRunnerParams.mMaxSeqLenKv = runnerParams.kvSeqLen;
185185
tllmRunnerParams.mAttentionWindowSize = runnerParams.slidingWindowSize;
186+
// Set chunked attention size to INT_MAX to disable chunked attention for now.
187+
tllmRunnerParams.mChunkedAttentionSize = INT_MAX;
186188
tllmRunnerParams.mSumOfSeqLensQ = runnerParams.totalQSeqLen;
187189
tllmRunnerParams.mSumOfSeqLensKv = runnerParams.totalKvSeqLen;
188190
tllmRunnerParams.mMaxNumPagesPerSeqKv = maxBlocksPerSeq;

‎cpp/tensorrt_llm/kernels/fmhaDispatcher.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,12 @@ class FmhaDispatcher
3737
// Deconstructor.
3838
~FmhaDispatcher() = default;
3939

40+
// Does the fmha kernel need FP8 inputQ ?
41+
bool isFP8InputQ() const
42+
{
43+
return mFixedParams.dataType == DATA_TYPE_E4M3;
44+
}
45+
4046
// Check if any fmha kernel meets the requirements.
4147
bool isSupported();
4248

‎cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -465,10 +465,13 @@ class TllmGenFmhaKernel
465465

466466
// The mask type.
467467
TrtllmGenAttentionMaskType maskType = params.mMaskType;
468-
// Enable sliding window causal if the max kv sequence length exceeds attention window size.
469-
if (params.mAttentionWindowSize < params.mMaxSeqLenKv && maskType == TrtllmGenAttentionMaskType::Causal)
468+
// Enable sliding window or chunked causal if the max kv sequence length exceeds attention window size or
469+
// chunked attention size.
470+
if (maskType == TrtllmGenAttentionMaskType::Causal
471+
&& (params.mMaxSeqLenKv > params.mAttentionWindowSize
472+
|| params.mMaxSeqLenKv > params.mChunkedAttentionSize))
470473
{
471-
maskType = TrtllmGenAttentionMaskType::SlidingWindowCausal;
474+
maskType = TrtllmGenAttentionMaskType::SlidingOrChunkedCausal;
472475
}
473476
// NumTokensPerPage is set to 0 when not selecting pagedKv-layout kernels.
474477
int numTokensPerPage = (!isPagedKv(params.mQkvLayout)) ? 0 : params.mNumTokensPerPage;

‎cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,8 @@ enum class TrtllmGenAttentionMaskType
3232
Dense = 0,
3333
// Causal mask.
3434
Causal,
35-
// Sliding window causal mask.
36-
SlidingWindowCausal,
35+
// Sliding window or chunked causal mask.
36+
SlidingOrChunkedCausal,
3737
// Custom mask.
3838
Custom
3939
};
@@ -50,7 +50,7 @@ enum class TrtllmGenAttentionMaskType
5050

5151
ATTENTION_MASK_TYPE_FUNCTION(Dense)
5252
ATTENTION_MASK_TYPE_FUNCTION(Causal)
53-
ATTENTION_MASK_TYPE_FUNCTION(SlidingWindowCausal)
53+
ATTENTION_MASK_TYPE_FUNCTION(SlidingOrChunkedCausal)
5454
ATTENTION_MASK_TYPE_FUNCTION(Custom)
5555

5656
#undef ATTENTION_MASK_TYPE_FUNCTION
@@ -246,8 +246,11 @@ struct TllmGenFmhaRunnerParams
246246
int mMaxSeqLenQ;
247247
// The max kv sequence length.
248248
int mMaxSeqLenKv;
249-
// The attention window size for sliding window attention.
249+
// The attention window size for sliding window attention (sliding-window-attention is enabled when seqLenKv >
250+
// mAttentionWindowSize).
250251
int mAttentionWindowSize;
252+
// The chunked attention size (chunked-context is enabled when seqLenKv > mChunkedAttentionSize).
253+
int mChunkedAttentionSize;
251254
// The sum of sequence lengths for Q and K/V. (Only used when mSupportsVarSeqLens = true)
252255
int mSumOfSeqLensQ;
253256
int mSumOfSeqLensKv;
@@ -283,8 +286,8 @@ struct TllmGenFmhaRunnerParams
283286
case 1: // tensorrt_llm::kernels::ContextAttentionMaskType::CAUSAL
284287
mMaskType = TrtllmGenAttentionMaskType::Causal;
285288
break;
286-
case 2: // tensorrt_llm::kernels::ContextAttentionMaskType::SLIDING_WINDOW_CAUSAL
287-
mMaskType = TrtllmGenAttentionMaskType::SlidingWindowCausal;
289+
case 2: // tensorrt_llm::kernels::ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL
290+
mMaskType = TrtllmGenAttentionMaskType::SlidingOrChunkedCausal;
288291
break;
289292
case 3: // tensorrt_llm::kernels::ContextAttentionMaskType::CUSTOM_MASK
290293
mMaskType = TrtllmGenAttentionMaskType::Custom;

‎cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,8 @@ struct KernelParams
108108
int32_t mAttentionWindowSize;
109109
// The batch size
110110
int32_t mBatchSize;
111+
// The chunked attention size in log2.
112+
int32_t mChunkedAttentionSizeLog2;
111113
// The log of the Sage Attention block size for K.
112114
int32_t mLogNumEltsPerSageAttnBlkK;
113115
// The log of the Sage Attention block size for P.
@@ -741,6 +743,17 @@ struct KernelParams
741743
params.ptrSoftmaxStats = options.softmaxStatsPtr;
742744

743745
params.mAttentionWindowSize = options.mAttentionWindowSize;
746+
if (isSlidingOrChunkedCausalMask(options.mMaskType) && options.mMaxSeqLenKv > options.mChunkedAttentionSize)
747+
{
748+
TLLM_CHECK_WITH_INFO((options.mChunkedAttentionSize & (options.mChunkedAttentionSize - 1)) == 0,
749+
"Chunked attention size must be a power of 2");
750+
params.mChunkedAttentionSizeLog2 = std::log2(options.mChunkedAttentionSize);
751+
}
752+
else
753+
{
754+
// Default 0 means that chunked attention is disabled.
755+
params.mChunkedAttentionSizeLog2 = 0;
756+
}
744757
params.mMaxSeqLenQ = options.mMaxSeqLenQ;
745758
params.mMaxSeqLenKv = options.mMaxSeqLenKv;
746759
params.mMaxNumCtasQ = maxNumCtasQ;

0 commit comments

Comments
 (0)