From 40bb9dafb9199046eb9a7700bf80de1c81fd650c Mon Sep 17 00:00:00 2001 From: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> Date: Mon, 30 Jun 2025 00:10:26 -0700 Subject: [PATCH 01/14] use separate qkv for context MLA Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 156 ++++++++------- cpp/tensorrt_llm/common/attentionOp.h | 6 + .../fused_multihead_attention_common.h | 10 +- cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp | 33 ++- cpp/tensorrt_llm/kernels/mlaKernels.cu | 188 ++++++++++++------ cpp/tensorrt_llm/kernels/mlaKernels.h | 20 +- cpp/tensorrt_llm/thop/attentionOp.cpp | 99 +++++---- .../_torch/attention_backend/interface.py | 4 + .../_torch/attention_backend/trtllm.py | 58 +++--- tensorrt_llm/_torch/metadata.py | 2 +- .../_torch/models/modeling_deepseekv3.py | 1 - tensorrt_llm/_torch/modules/attention.py | 174 ++++++++-------- 12 files changed, 447 insertions(+), 304 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index fba7f729c903..87daafcbff42 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -747,17 +747,19 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t = mNumAttnHeads * dim_k_per_head; // Assuming effective num_kv_heads = head_num for layout int const total_v_dim_all_heads = mNumAttnHeads * dim_v_per_head; // Assuming effective num_kv_heads = head_num for layout - - int const num_total_qkv_elements - = max_num_tokens * (total_q_dim_all_heads + total_k_dim_all_heads + total_v_dim_all_heads); - + // Packed fp8 qkv buffer size for normal fp8 context FMHA size_t fp8_qkv_buffer_size = mFP8ContextFMHA && mEnableContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput() ? max_num_tokens * size_t(local_hidden_units_qo + 2 * local_hidden_units_kv) : 0; - if (mFP8ContextMLA) + // Separate fp8 q/k/v buffer size for fp8 context MLA + size_t fp8_q_buf_size = 0; + size_t fp8_k_buf_size = 0; + size_t fp8_v_buf_size = 0; + if (mEnableContextFMHA && mFP8ContextMLA && mFmhaDispatcher->isSeparateQAndKvInput()) { - fp8_qkv_buffer_size - = mEnableContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput() ? num_total_qkv_elements : 0; + fp8_q_buf_size = max_num_tokens * total_q_dim_all_heads; + fp8_k_buf_size = max_num_tokens * total_k_dim_all_heads; + fp8_v_buf_size = max_num_tokens * total_v_dim_all_heads; } size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens; @@ -774,7 +776,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t ? 0 : (2 * size * cpMaxPaddedSequenceLength * getHeadSize() * (mNumHeads + 2 * mNumKVHeads) + cu_seqlens_size); - int const NUM_BUFFERS = 20; + int const NUM_BUFFERS = 23; size_t workspaces[NUM_BUFFERS]; workspaces[0] = CUBLAS_WORKSPACE_SIZE; workspaces[1] = attention_mask_size; @@ -789,13 +791,16 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t workspaces[10] = qkv_buf_2_size; workspaces[11] = qk_buf_float_size; workspaces[12] = fp8_qkv_buffer_size; - workspaces[13] = padding_offset_size; - workspaces[14] = encoder_padding_offset_size; - workspaces[15] = tokens_info_size; - workspaces[16] = fmha_scheduler_counter; - workspaces[17] = fmha_bmm1_scale_size; - workspaces[18] = fmha_bmm2_scale_size; - workspaces[19] = cpWorkspaceSize; + workspaces[13] = fp8_q_buf_size; + workspaces[14] = fp8_k_buf_size; + workspaces[15] = fp8_v_buf_size; + workspaces[16] = padding_offset_size; + workspaces[17] = encoder_padding_offset_size; + workspaces[18] = tokens_info_size; + workspaces[19] = fmha_scheduler_counter; + workspaces[20] = fmha_bmm1_scale_size; + workspaces[21] = fmha_bmm2_scale_size; + workspaces[22] = cpWorkspaceSize; context_workspace_size = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); return context_workspace_size; @@ -979,7 +984,7 @@ int AttentionOp::mlaGeneration( params.fmha_tile_counter = fmha_tile_counter_ptr; params.bmm1_scale = mla_bmm1_scale_ptr; params.bmm2_scale = mla_bmm2_scale_ptr; - params.quant_attention_input_buf = quant_q_buffer_ptr; + params.quant_q_buf = quant_q_buffer_ptr; params.quant_scale_o = generation_params.attention_output_orig_quant; params.quant_scale_q = generation_params.kv_scale_orig_quant; @@ -1020,8 +1025,8 @@ int AttentionOp::mlaGeneration( tllmRunnerParams.mTileScheduler = mMultiBlockMode ? TileScheduler::Static : TileScheduler::Persistent; // Q buffer. - tllmRunnerParams.qPtr = mFP8GenerationMLA ? reinterpret_cast(params.quant_attention_input_buf) - : reinterpret_cast(params.attention_input_buf); + tllmRunnerParams.qPtr = mFP8GenerationMLA ? reinterpret_cast(params.quant_q_buf) + : reinterpret_cast(params.q_buf); // KV buffer // Paged KV @@ -1147,8 +1152,8 @@ int AttentionOp::mlaGeneration( flashMlaParams.scale_softmax_log2 = float(softmax_scale * M_LOG2E); flashMlaParams.q_ptr = mFP8GenerationMLA - ? const_cast(reinterpret_cast(params.quant_attention_input_buf)) - : const_cast(reinterpret_cast(params.attention_input_buf)); + ? const_cast(reinterpret_cast(params.quant_q_buf)) + : const_cast(reinterpret_cast(params.q_buf)); flashMlaParams.k_ptr = kv_cache_buffer.mPrimaryPoolPtr; flashMlaParams.v_ptr = flashMlaParams.k_ptr; flashMlaParams.o_ptr = reinterpret_cast(params.context_buf); @@ -1253,8 +1258,8 @@ int AttentionOp::mlaGeneration( // fmhaParams.totalKvSeqLen = params.num_tokens; // Device buffer pointers. // fmhaParams.qkvPtr = reinterpret_cast(params.attention_input); - fmhaParams.qPtr = mFP8GenerationMLA ? reinterpret_cast(params.quant_attention_input_buf) - : reinterpret_cast(params.attention_input_buf); + fmhaParams.qPtr = mFP8GenerationMLA ? reinterpret_cast(params.quant_q_buf) + : reinterpret_cast(params.q_buf); // TODO: add contiguous kv buffer (cross-attention). fmhaParams.kvPtr = nullptr; @@ -1379,15 +1384,19 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea = mNumAttnHeads * dim_k_per_head; // Assuming effective num_kv_heads = head_num for layout int const total_v_dim_all_heads = mNumAttnHeads * dim_v_per_head; // Assuming effective num_kv_heads = head_num for layout - int const num_total_qkv_elements - = params.num_tokens * (total_q_dim_all_heads + total_k_dim_all_heads + total_v_dim_all_heads); + // Packed fp8 qkv buffer size for normal fp8 context FMHA size_t fp8_qkv_buffer_size = mEnableContextFMHA && mFP8ContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput() ? params.num_tokens * (local_hidden_units_qo + 2 * local_hidden_units_kv) : 0; - if (mFP8ContextMLA) + // Separate fp8 q/k/v buffer size for fp8 context MLA + size_t fp8_q_buf_size = 0; + size_t fp8_k_buf_size = 0; + size_t fp8_v_buf_size = 0; + if (mEnableContextFMHA && mFP8ContextMLA && mFmhaDispatcher->isSeparateQAndKvInput()) { - fp8_qkv_buffer_size - = mEnableContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput() ? num_total_qkv_elements : 0; + fp8_q_buf_size = params.num_tokens * total_q_dim_all_heads; + fp8_k_buf_size = params.num_tokens * total_k_dim_all_heads; + fp8_v_buf_size = params.num_tokens * total_v_dim_all_heads; } size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * params.batch_size * params.input_seq_length; @@ -1424,6 +1433,9 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea float* qk_buf_float_ = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, qk_buf_float_size)); __nv_fp8_e4m3* fp8_qkv_buffer = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, fp8_qkv_buffer_size)); + __nv_fp8_e4m3* fp8_q_buf = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, fp8_q_buf_size)); + __nv_fp8_e4m3* fp8_k_buf = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, fp8_k_buf_size)); + __nv_fp8_e4m3* fp8_v_buf = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, fp8_v_buf_size)); int* padding_offset = mEnableContextFMHA ? nullptr : reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, padding_offset_size)); @@ -1638,7 +1650,6 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea "Found invalid number (NaN or Inf) in " + beforeRopeStr); } - KVBlockArray mla_context_paged_kv_cache_buffer; if (mIsMLAEnabled) { params.mla_param->cache_type = cache_type; @@ -1647,7 +1658,9 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea // Set BMM scales for FP8 context computation params.mla_param->bmm1_scale = fmha_bmm1_scale_ptr; params.mla_param->bmm2_scale = fmha_bmm2_scale_ptr; - params.mla_param->quant_attention_input_buf = mFP8ContextMLA ? fp8_qkv_buffer : nullptr; + params.mla_param->quant_q_buf = mFP8ContextMLA ? fp8_q_buf : nullptr; + params.mla_param->quant_k_buf = mFP8ContextMLA ? fp8_k_buf : nullptr; + params.mla_param->quant_v_buf = mFP8ContextMLA ? fp8_v_buf : nullptr; // Set additional scales for context phase params.mla_param->quant_scale_o = params.attention_output_orig_quant; params.mla_param->quant_scale_q = params.kv_scale_orig_quant; @@ -1656,36 +1669,14 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea params.mla_param->dequant_scale_kv = params.kv_scale_quant_orig; params.mla_param->host_bmm1_scale = 1 / (mQScaling * sqrt((float) (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim))); - if (mPagedContextFMHA && mPagedKVCache) + if (params.mla_param->latent_cache != nullptr) { - TLLM_CHECK_WITH_INFO(params.mla_param->context_paged_kv_ptr != nullptr, - "Paged kv cache is not set for MLA context kernel"); - TLLM_CHECK_WITH_INFO(params.mla_param->context_kv_cache_block_offsets_ptr != nullptr, - "Paged kv cache block offsets is not set for MLA context kernel"); - // build another KVBlockArray for MLA context kernel to read paged kv cache, which is built by the - // PyTorch backend assume the dtype of paged kv cache is the same as the T - auto const elemSize = sizeof(T); - auto const headSize = params.mla_param->meta.qk_nope_head_dim + params.mla_param->meta.qk_rope_head_dim; - // mNumKVHeads is 1 for writing, we use mNumHeads for reading paged kv cache - auto sizePerToken = mNumHeads * headSize * elemSize; - auto maxBlocksPerSeq = params.mla_param->context_paged_kv_max_blocks_per_seq; - TLLM_LOG_DEBUG( - "AttentionOp building KVBlockArray for MLA context kernel, elemSize: %d, headSize: %d, mNumHeads: " - "%d, sizePerToken: %d, batchSize: %d, maxBlocksPerSeq: %d, tokensPerBlock: %d, maxAttentionWindow: " - "%d, " - "sinkTokenLen: %d, canUseOneMoreBlock: %d", - elemSize, headSize, mNumHeads, sizePerToken, params.batch_size, maxBlocksPerSeq, mTokensPerBlock, - params.cyclic_attention_window_size, params.sink_token_length, params.can_use_one_more_block); - mla_context_paged_kv_cache_buffer = KVBlockArray(params.batch_size, maxBlocksPerSeq, mTokensPerBlock, - sizePerToken, params.cyclic_attention_window_size, params.max_cyclic_attention_window_size, - params.sink_token_length, params.can_use_one_more_block, params.mla_param->context_paged_kv_ptr, - nullptr, - static_cast(params.mla_param->context_kv_cache_block_offsets_ptr)); + // compute RoPE and set compressed_kv + k_pe by invokeMLARopeContext if latent_cache is not nullptr + invokeMLARopeContext(*params.mla_param, kv_cache_buffer, stream); } - else + if (mFP8ContextMLA) { - // compute RoPE and set compressed_kv + k_pe by invokeMLARopeContext if not using paged context FMHA - invokeMLARopeContext(*params.mla_param, kv_cache_buffer, stream); + invokeMLAContextFp8Quantize(*params.mla_param, params.total_kv_len, stream); } } else @@ -1709,7 +1700,8 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea mFMHAForceFP32Acc = mFMHAForceFP32Acc || enable_context_fmha_fp32_acc_val == 1; } - // Unified FMHA runner interface for both packed QKV FMHA, contiguous Q_KV and paged KV FMHA. + // Unified FMHA runner interface for both packed QKV FMHA, contiguous Q_KV, paged KV FMHA, and separate QKV + // FMHA. // Page KV input layout: // - q_ptr: [B, S, H, D], which supports variable sequence length // - paged_kv_cache: paged kv buffer @@ -1721,6 +1713,13 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea // - kv_ptr: [B, S, 2, H, D], which supports variable sequence length // - cu_q_seqlens: the cumulative query sequence lengths, needed for variable sequence length. // - cu_kv_seqlens: the cumulative kv sequence lengths, needed for variable sequence length. + // + // Separate QKV input layout (only for context MLA now): + // - q_ptr: [B, S, H, D], which supports variable sequence length + // - k_ptr: [B, S, H_kv, D], which supports variable sequence length + // - v_ptr: [B, S, H_kv, D_v], which supports variable sequence length + // - cu_q_seqlens: the cumulative query sequence lengths, needed for variable sequence length. + // - cu_kv_seqlens: the cumulative kv sequence lengths, needed for variable sequence length. // Construct the fmha params for running kernels. MHARunnerParams fmhaParams{}; @@ -1732,11 +1731,34 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea = (mDenseContextFMHA || isCrossAttention()) ? max_kv_seq_len : params.cyclic_attention_window_size; fmhaParams.totalQSeqLen = params.num_tokens; // TODO: set it correctly for contiguous kv buffer (cross-attention). - fmhaParams.totalKvSeqLen = isCrossAttention() ? params.num_encoder_tokens : params.num_tokens; + fmhaParams.totalKvSeqLen = isCrossAttention() ? params.num_encoder_tokens : params.total_kv_len; // Device buffer pointers. - fmhaParams.qkvPtr = (mFP8ContextFMHA || mFP8ContextMLA) ? reinterpret_cast(fp8_qkv_buffer) - : reinterpret_cast(attention_input); - fmhaParams.qPtr = reinterpret_cast(q_buf_2_); + if (mIsMLAEnabled) + { + // separate QKV input for context MLA + if (mFP8ContextMLA) + { + TLLM_CHECK_WITH_INFO(mFmhaDispatcher->isSeparateQAndKvInput(), "Separate QKV input is required for fp8 context MLA"); + TLLM_CHECK_WITH_INFO(fp8_q_buf != nullptr, "FP8 q buffer is required for fp8 context MLA"); + TLLM_CHECK_WITH_INFO(fp8_k_buf != nullptr, "FP8 k buffer is required for fp8 context MLA"); + TLLM_CHECK_WITH_INFO(fp8_v_buf != nullptr, "FP8 v buffer is required for fp8 context MLA"); + fmhaParams.qPtr = reinterpret_cast(fp8_q_buf); + fmhaParams.kPtr = reinterpret_cast(fp8_k_buf); + fmhaParams.vPtr = reinterpret_cast(fp8_v_buf); + } + else + { + fmhaParams.qPtr = attention_input; + fmhaParams.kPtr = params.k_ptr; + fmhaParams.vPtr = params.v_ptr; + } + } + else + { + fmhaParams.qkvPtr = mFP8ContextFMHA ? reinterpret_cast(fp8_qkv_buffer) + : reinterpret_cast(attention_input); + fmhaParams.qPtr = reinterpret_cast(q_buf_2_); + } // TODO: add contiguous kv buffer (cross-attention). fmhaParams.kvPtr = nullptr; if (isCrossAttention() && !useKVCache()) @@ -1750,15 +1772,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea fmhaParams.packedMaskPtr = params.attention_packed_mask; if constexpr (std::is_same_v) { - if (mIsMLAEnabled && mPagedContextFMHA && mPagedKVCache) - { - fmhaParams.pagedKvCache = mla_context_paged_kv_cache_buffer; - fmhaParams.qPtr = reinterpret_cast(attention_input); - } - else - { - fmhaParams.pagedKvCache = kv_cache_buffer; - } + fmhaParams.pagedKvCache = kv_cache_buffer; } fmhaParams.cuQSeqLenPtr = cu_q_seqlens; fmhaParams.kvSeqLenPtr = decoder_params.seqKVLengths; @@ -2606,6 +2620,8 @@ int AttentionOp::initialize() noexcept // the wrong kernel, no matter mIsGenerationMLA is true or false if (mIsMLAEnabled) { + // context MLA always use separate_qkv layout + fmhaParams.attentionInputLayout = AttentionInputLayout::SEPARATE_QKV; // Context attention of MLA is different fmhaParams.numKvHeads = mNumHeads; fmhaParams.headSize = mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim; diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index 25d95dfea2bc..6b3b00a50dd6 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -95,6 +95,7 @@ class AttentionOp void* host_primary_pool_pointer = nullptr; void* host_secondary_pool_pointer = nullptr; int32_t num_tokens = 0; + int32_t total_kv_len = 0; int32_t max_blocks_per_sequence = 0; int32_t const* sequence_lengths = nullptr; int32_t const* context_lengths = nullptr; @@ -128,6 +129,9 @@ class AttentionOp // For MLA chunked prefill void* softmaxStatsPtr = nullptr; + // optional for separate QKV input, currently only used for context MLA + T const* k_ptr = nullptr; + T const* v_ptr = nullptr; std::string enqueueContextParamsToString() const { @@ -179,6 +183,8 @@ class AttentionOp ss << "encoder_input_lengths: " << this->encoder_input_lengths << std::endl; ss << "num_encoder_tokens: " << this->num_encoder_tokens << std::endl; ss << "softmaxStatsPtr: " << this->softmaxStatsPtr << std::endl; + ss << "k_ptr: " << this->k_ptr << std::endl; + ss << "v_ptr: " << this->v_ptr << std::endl; return ss.str(); } }; diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h index e9098866161a..d8fe8c49fc68 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h @@ -81,7 +81,10 @@ enum class AttentionInputLayout // Q has contiguous [B, S, H, D] layout, while paged KV has [B, 2, Max_blocks_per_seq] layout // that contains paged block indices. The indices indicate the block offset to the pool ptr in // global memory - Q_PAGED_KV + Q_PAGED_KV, + // Q has contiguous [B, S, H, D] layout, while K has contiguous [B, S, H_kv, D] layout, and V has + // contiguous [B, S, H_kv, D_v] layout. Only used for context MLA now. + SEPARATE_QKV, }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -166,6 +169,7 @@ struct MHARunnerFixedParams case AttentionInputLayout::PACKED_QKV: output += "packed_qkv"; break; case AttentionInputLayout::Q_CONTIGUOUS_KV: output += "q_contiguous_kv"; break; case AttentionInputLayout::Q_PAGED_KV: output += "q_paged_kv"; break; + case AttentionInputLayout::SEPARATE_QKV: output += "separate_qkv"; break; default: output += std::to_string(static_cast(attentionInputLayout)) + " (unknown)"; break; } @@ -255,6 +259,10 @@ struct MHARunnerParams void const* qPtr; // The contiguous Kv buffer ptr; void const* kvPtr; + // The K buffer ptr (for separate K input). + void const* kPtr; + // The V buffer ptr (for separate V input). + void const* vPtr; // The paged kv cache array. KVBlockArray pagedKvCache; // The output buffer ptr. diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp index af7468415a49..5e9adf69cd26 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp @@ -36,6 +36,10 @@ QkvLayout AttentionInputLayoutToQkvLayout(AttentionInputLayout layout) { return QkvLayout::PagedKv; } + else if (layout == AttentionInputLayout::SEPARATE_QKV) + { + return QkvLayout::SeparateQkv; + } TLLM_CHECK_WITH_INFO(false, "Unexpected AttentionInputLayout"); return QkvLayout::SeparateQkv; } @@ -148,6 +152,10 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) maxBlocksPerSeq = pagedKvCache.mMaxBlocksPerSeq; numTokensPerBlock = pagedKvCache.mTokensPerBlock; } + else if (mFixedParams.attentionInputLayout == AttentionInputLayout::SEPARATE_QKV) + { + qkvLayout = kernels::QkvLayout::SeparateQkv; + } TllmGenFmhaRunnerParams tllmRunnerParams; memset(&tllmRunnerParams, 0, sizeof(tllmRunnerParams)); @@ -161,8 +169,8 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) tllmRunnerParams.mMultiCtasKvMode = false; tllmRunnerParams.qPtr = runnerParams.qPtr; - tllmRunnerParams.kPtr = nullptr; - tllmRunnerParams.vPtr = nullptr; + tllmRunnerParams.kPtr = runnerParams.kPtr; + tllmRunnerParams.vPtr = runnerParams.vPtr; tllmRunnerParams.kvPtr = kvPoolPtr; tllmRunnerParams.qkvPtr = runnerParams.qkvPtr; tllmRunnerParams.attentionSinksPtr = runnerParams.attentionSinksPtr; @@ -202,6 +210,27 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) // For mla chunked prefill tllmRunnerParams.softmaxStatsPtr = reinterpret_cast(runnerParams.softmaxStatsPtr); tllmRunnerParams.stream = runnerParams.stream; + + // if (runnerParams.totalQSeqLen < 8192) { + // std::vector cumSeqLensQ(runnerParams.b + 1); + // std::vector cumSeqLensKv(runnerParams.b + 1); + // cudaMemcpy(cumSeqLensQ.data(), runnerParams.cuQSeqLenPtr, (runnerParams.b + 1) * sizeof(int), + // cudaMemcpyDeviceToHost); cudaMemcpy(cumSeqLensKv.data(), runnerParams.cuKvSeqLenPtr, (runnerParams.b + 1) + // * sizeof(int), cudaMemcpyDeviceToHost); std::stringstream ss; for (int i = 0; i < runnerParams.b + 1; + // i++) { + // ss << cumSeqLensQ[i] << " "; + // } + // TLLM_LOG_INFO("cumSeqLensQ: %s", ss.str().c_str()); + // ss.str(""); + // for (int i = 0; i < runnerParams.b + 1; i++) { + // ss << cumSeqLensKv[i] << " "; + // } + // TLLM_LOG_INFO("cumSeqLensKv: %s", ss.str().c_str()); + // TLLM_LOG_INFO("mMaxSeqLenCacheKv: %d, mMaxSeqLenQ: %d, mMaxSeqLenKv: %d", + // tllmRunnerParams.mMaxSeqLenCacheKv, tllmRunnerParams.mMaxSeqLenQ, tllmRunnerParams.mMaxSeqLenKv); + // TLLM_LOG_INFO("mSumOfSeqLensQ: %d, mSumOfSeqLensKv: %d", tllmRunnerParams.mSumOfSeqLensQ, + // tllmRunnerParams.mSumOfSeqLensKv); + // } mTllmGenFMHARunner->run(tllmRunnerParams); } else diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index cdb7abbb91bf..2ffee12b788f 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -205,7 +205,7 @@ inline __device__ void dequantCopy( } template -__global__ void applyMLARopeAndAssignQKVKernelOptContext(T* qkv_output, T const* fuse_buf, KVCacheBuffer kv_cache, +__global__ void applyMLARopeAndAssignQKVKernelOptContext(T* q_ptr, T* k_ptr, T const* fuse_buf, KVCacheBuffer kv_cache, float2 const* cos_sin_cache, size_t head_num, int head_size, int c_k, int* cu_q_seqlens, int32_t const* kv_cache_lengths, uint32_t max_input_seq_len, KvCacheDataType cache_type, float* bmm1_scale, float* bmm2_scale, float const* quant_scale_o, float const* quant_scale_kv, float const* dequant_scale_q, @@ -287,11 +287,10 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* qkv_output, T const* VecT q, k; auto const src_k_global_offset = static_cast(global_token_idx) * (c_k + ROPE_DIM) + c_k; - auto const src_q_global_offset - = static_cast(global_token_idx) * head_num * ((head_size + ROPE_DIM) * 2 + head_size) + auto const src_q_global_offset = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + (head_size + ROPE_DIM) * head_idx + head_size; - q = *reinterpret_cast(&qkv_output[src_q_global_offset + head_dim_idx]); + q = *reinterpret_cast(&q_ptr[src_q_global_offset + head_dim_idx]); k = *reinterpret_cast(&fuse_buf[src_k_global_offset + head_dim_idx]); // Pack two elements into one for gptj rotary embedding. @@ -322,14 +321,12 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* qkv_output, T const* else reinterpret_cast(kDst)[inBlockIdx] = k; } - auto const dst_q_idx - = static_cast(global_token_idx) * head_num * ((head_size + ROPE_DIM) * 2 + head_size) + auto const dst_q_idx = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + head_idx * (head_size + ROPE_DIM) + head_size + head_dim_idx; - auto const dst_k_idx - = static_cast(global_token_idx) * head_num * ((head_size + ROPE_DIM) * 2 + head_size) - + head_num * (head_size + ROPE_DIM) + head_idx * (head_size + ROPE_DIM) + head_size + head_dim_idx; - reinterpret_cast(qkv_output)[dst_q_idx / ELTS_PER_VEC] = q; - reinterpret_cast(qkv_output)[dst_k_idx / ELTS_PER_VEC] = k; + auto const dst_k_idx = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + + head_idx * (head_size + ROPE_DIM) + head_size + head_dim_idx; + reinterpret_cast(q_ptr)[dst_q_idx / ELTS_PER_VEC] = q; + reinterpret_cast(k_ptr)[dst_k_idx / ELTS_PER_VEC] = k; } } } @@ -941,58 +938,130 @@ __global__ void applyMLARopeAppendPagedKVAssignQKernel(KVBlockArray kv_cache, T* } } +template +__global__ void QuantizeCopyInputToFp8Kernel( + T const* q_buf, __nv_fp8_e4m3* quant_q_buf, T const* k_buf, __nv_fp8_e4m3* quant_k_buf, T const* v_buf, + __nv_fp8_e4m3* quant_v_buf, int total_q_len, int total_kv_len, float const* quant_scale_qkv_ptr) +{ + // Constants. + using VecT = typename VecType::Type; + constexpr auto BYTES_PER_ELT = sizeof(T); // 2 + constexpr auto BYTES_PER_LOAD = 16; + constexpr auto ELTS_PER_VEC = BYTES_PER_LOAD / BYTES_PER_ELT; // 8 + constexpr auto QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM; + static_assert((QK_HEAD_DIM * BYTES_PER_ELT) % BYTES_PER_LOAD == 0, "QK head size needs to be multiple of 16 bytes."); + static_assert((V_HEAD_DIM * BYTES_PER_ELT) % BYTES_PER_LOAD == 0, "V head size needs to be multiple of 16 bytes."); + constexpr auto QK_VECS_PER_HEAD = QK_HEAD_DIM * BYTES_PER_ELT / BYTES_PER_LOAD; // 24 + constexpr auto V_VECS_PER_HEAD = V_HEAD_DIM * BYTES_PER_ELT / BYTES_PER_LOAD; //16 + static_assert(BLOCK_SIZE % QK_VECS_PER_HEAD == 0, "Kernel block should be able to handle entire heads."); + static_assert(BLOCK_SIZE % V_VECS_PER_HEAD == 0, "Kernel block should be able to handle entire heads."); + constexpr auto QK_TOKENS_PER_BLOCK = BLOCK_SIZE / QK_VECS_PER_HEAD; + constexpr auto V_TOKENS_PER_BLOCK = BLOCK_SIZE / V_VECS_PER_HEAD; + + size_t const head_idx = blockIdx.z; + size_t const head_num = gridDim.z; + + size_t const qk_head_dim_vec_idx = (threadIdx.x % QK_VECS_PER_HEAD); + size_t const v_head_dim_vec_idx = (threadIdx.x % V_VECS_PER_HEAD); + size_t const qk_head_dim_idx = qk_head_dim_vec_idx * ELTS_PER_VEC; + size_t const v_head_dim_idx = v_head_dim_vec_idx * ELTS_PER_VEC; + + size_t const q_len_loop_end + = size_t((total_q_len + QK_TOKENS_PER_BLOCK - 1) / QK_TOKENS_PER_BLOCK) * QK_TOKENS_PER_BLOCK; + size_t const k_len_loop_end + = size_t((total_kv_len + QK_TOKENS_PER_BLOCK - 1) / QK_TOKENS_PER_BLOCK) * QK_TOKENS_PER_BLOCK; + size_t const v_len_loop_end + = size_t((total_kv_len + V_TOKENS_PER_BLOCK - 1) / V_TOKENS_PER_BLOCK) * V_TOKENS_PER_BLOCK; + float quant_scale_qkv_val = quant_scale_qkv_ptr ? quant_scale_qkv_ptr[0] : 1.f; + + // Quantize Q, both src and dst are contiguous + for (int q_token_idx = (threadIdx.x / QK_VECS_PER_HEAD) + blockIdx.x * QK_TOKENS_PER_BLOCK; + q_token_idx < q_len_loop_end; q_token_idx += QK_TOKENS_PER_BLOCK * gridDim.x) + { + if (q_token_idx < total_q_len) + { + auto const src_q_idx + = static_cast(q_token_idx) * QK_HEAD_DIM * head_num + head_idx * QK_HEAD_DIM + qk_head_dim_idx; + auto const dst_q_idx = src_q_idx; + quantCopy(quant_q_buf + dst_q_idx, &q_buf[src_q_idx], quant_scale_qkv_val); + } + } + + // Quantize K, both src and dst are contiguous + for (int k_token_idx = (threadIdx.x / QK_VECS_PER_HEAD) + blockIdx.x * QK_TOKENS_PER_BLOCK; + k_token_idx < k_len_loop_end; k_token_idx += QK_TOKENS_PER_BLOCK * gridDim.x) + { + if (k_token_idx < total_kv_len) + { + auto const src_k_idx + = static_cast(k_token_idx) * QK_HEAD_DIM * head_num + head_idx * QK_HEAD_DIM + qk_head_dim_idx; + auto const dst_k_idx = src_k_idx; + quantCopy(quant_k_buf + dst_k_idx, &k_buf[src_k_idx], quant_scale_qkv_val); + } + } + + // Quantize V, dst V is contiguous, but src V is not contiguous, so we need to calculate the stride + size_t const src_v_token_stride = (QK_NOPE_HEAD_DIM + V_HEAD_DIM) * head_num; + for (int v_token_idx = (threadIdx.x / V_VECS_PER_HEAD) + blockIdx.x * V_TOKENS_PER_BLOCK; + v_token_idx < v_len_loop_end; v_token_idx += V_TOKENS_PER_BLOCK * gridDim.x) + { + if (v_token_idx < total_kv_len) + { + auto const src_v_idx + = static_cast(v_token_idx) * src_v_token_stride + head_idx * V_HEAD_DIM + v_head_dim_idx; + auto const dst_v_idx + = static_cast(v_token_idx) * V_HEAD_DIM * head_num + head_idx * V_HEAD_DIM + v_head_dim_idx; + quantCopy(quant_v_buf + dst_v_idx, &v_buf[src_v_idx], quant_scale_qkv_val); + } + } +} + template void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, cudaStream_t stream) { dim3 grid(int(tensorrt_llm::common::divUp(params.max_input_seq_len, 32)), params.batch_size, params.head_num + 8); auto head_size = params.meta.qk_nope_head_dim; - applyMLARopeAndAssignQKVKernelOptContext<<>>( - params.attention_input_buf, params.latent_cache, kv_cache_buffer, params.cos_sin_cache, params.head_num, - head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, + applyMLARopeAndAssignQKVKernelOptContext<<>>(params.q_buf, + params.k_buf, params.latent_cache, kv_cache_buffer, params.cos_sin_cache, params.head_num, head_size, + params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, params.bmm1_scale, params.bmm2_scale, params.quant_scale_o, params.quant_scale_kv, params.dequant_scale_q, params.dequant_scale_kv, params.host_bmm1_scale); - if (params.attention_input_buf != nullptr && params.quant_attention_input_buf != nullptr - && params.cache_type == KvCacheDataType::FP8) - { - TLLM_LOG_DEBUG("MLA RoPE Context: Quantizing attention_input_buf to FP8"); - - int const dim_q_per_head = (params.meta.qk_nope_head_dim + params.meta.qk_rope_head_dim); - int const dim_k_per_head = (params.meta.qk_nope_head_dim + params.meta.qk_rope_head_dim); - int const dim_v_per_head = (params.meta.v_head_dim); - - // Total dimension per token across all heads for Q, K, and V components respectively - int const total_q_dim_all_heads = params.head_num * dim_q_per_head; - int const total_k_dim_all_heads - = params.head_num * dim_k_per_head; // Assuming effective num_kv_heads = head_num for layout - int const total_v_dim_all_heads - = params.head_num * dim_v_per_head; // Assuming effective num_kv_heads = head_num for layout - - int const num_total_qkv_elements - = params.acc_q_len * (total_q_dim_all_heads + total_k_dim_all_heads + total_v_dim_all_heads); - size_t headDim = params.meta.kv_lora_rank + params.meta.qk_rope_head_dim; - float const* device_qkv_scale_ptr = params.quant_scale_qkv; +} - if (num_total_qkv_elements > 0) - { - int const threads_per_block = 256; - int const num_blocks = (num_total_qkv_elements + threads_per_block - 1) / threads_per_block; +template +void invokeMLAContextFp8Quantize(MlaParams& params, int total_kv_len, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.cache_type == KvCacheDataType::FP8, "MLA Context: cache_type must be FP8"); + TLLM_CHECK_WITH_INFO(params.q_buf != nullptr, "MLA Context: q_buf must be non-null"); + TLLM_CHECK_WITH_INFO(params.k_buf != nullptr, "MLA Context: k_buf must be non-null"); + TLLM_CHECK_WITH_INFO(params.v_buf != nullptr, "MLA Context: v_buf must be non-null"); + TLLM_CHECK_WITH_INFO(params.quant_q_buf != nullptr, "MLA Context: quant_q_buf must be non-null"); + TLLM_CHECK_WITH_INFO(params.quant_k_buf != nullptr, "MLA Context: quant_k_buf must be non-null"); + TLLM_CHECK_WITH_INFO(params.quant_v_buf != nullptr, "MLA Context: quant_v_buf must be non-null"); + TLLM_CHECK_WITH_INFO(params.quant_scale_qkv != nullptr, "MLA Context: quant_scale_qkv must be non-null"); + + TLLM_LOG_DEBUG("MLA RoPE Context: Quantizing separate qkv to FP8"); + + if (params.acc_q_len > 0) + { + constexpr int threads_per_block = 384; + dim3 grid(int(tensorrt_llm::common::divUp(total_kv_len, 48)), 1, params.head_num); - TLLM_LOG_DEBUG( - "Launching QuantizeCopyInputToFp8Kernel with num_blocks: %d, threads_per_block: %d, elements: %d", - num_blocks, threads_per_block, num_total_qkv_elements); + TLLM_LOG_DEBUG( + "Launching QuantizeCopyInputToFp8Kernel with grid_size: (%d, %d, %d), threads_per_block: %d", + grid.x, grid.y, grid.z, threads_per_block); - tensorrt_llm::kernels::QuantizeCopyInputToFp8Kernel<<>>( - static_cast(params.attention_input_buf), // Source - static_cast<__nv_fp8_e4m3*>(params.quant_attention_input_buf), // Destination - num_total_qkv_elements, device_qkv_scale_ptr); - sync_check_cuda_error(stream); + QuantizeCopyInputToFp8Kernel + <<>>( + params.q_buf, params.quant_q_buf, params.k_buf, params.quant_k_buf, params.v_buf, params.quant_v_buf, + params.acc_q_len, total_kv_len, params.quant_scale_qkv); + sync_check_cuda_error(stream); - cudaStreamSynchronize(stream); - } - else - { - TLLM_LOG_WARNING("MLA RoPE Context: num_total_qkv_elements is 0, skipping quantization."); - } + cudaStreamSynchronize(stream); + } + else + { + TLLM_LOG_WARNING("MLA RoPE Context: acc_q_len is 0, skipping quantization."); } } @@ -1017,8 +1086,8 @@ void invokeMLARopeGeneration(MlaParams& params, KVCacheBuffer kv_cache_buffer attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, kernel_instance, params.attention_input_buf, params.q_pe, params.latent_cache, - params.quant_attention_input_buf, kv_cache_buffer, params.cos_sin_cache, params.head_num, + cudaLaunchKernelEx(&config, kernel_instance, params.q_buf, params.q_pe, params.latent_cache, + params.quant_q_buf, kv_cache_buffer, params.cos_sin_cache, params.head_num, params.meta.kv_lora_rank, params.acc_q_len, seq_len, params.seqQOffset, params.fmha_tile_counter, params.cache_seq_lens, params.cu_kv_seqlens, params.q_pe_ld, params.q_pe_stride, params.cache_type, params.bmm1_scale, params.bmm2_scale, params.quant_scale_o, params.quant_scale_q, params.quant_scale_kv, @@ -1108,17 +1177,6 @@ INSTANTIATE_SET_KVCACHE_MLA(float); INSTANTIATE_SET_KVCACHE_MLA(half); INSTANTIATE_SET_KVCACHE_MLA(__nv_bfloat16); -template -__global__ void QuantizeCopyInputToFp8Kernel( - T_IN const* input_buffer, __nv_fp8_e4m3* output_fp8_buffer, int num_total_elements, float const* device_scale_ptr) -{ - uint element_idx = threadIdx.x + blockDim.x * blockIdx.x; - if (element_idx < num_total_elements) - { - float scale_factor = (device_scale_ptr != nullptr) ? *device_scale_ptr : 1.0f; - output_fp8_buffer[element_idx] = __nv_fp8_e4m3(static_cast(input_buffer[element_idx]) * scale_factor); - } -} } // namespace kernels } // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h index e3472c13d8cc..709548b1e9af 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -52,8 +52,12 @@ template struct MlaParams { T const* latent_cache; // cKV + k_pe - T* attention_input_buf; // [b, s, 3, h, d_h + r] - void* quant_attention_input_buf; + T* q_buf; + T* k_buf = nullptr; // [total_kv_lens, h, d_nope + d_rope], for context MLA, contiguous + T* v_buf = nullptr; // [total_kv_lens, h, d_v], for context MLA, not contiguous + void* quant_q_buf = nullptr; + void* quant_k_buf = nullptr; // [total_kv_lens, h, d_nope + d_rope], for fp8 context MLA, contiguous + void* quant_v_buf = nullptr; // [total_kv_lens, h, d_v], for fp8 context MLA, contiguous T* context_buf; T* q_pe; // [b, h, d_r], strided @@ -83,10 +87,6 @@ struct MlaParams float const* dequant_scale_kv; float host_bmm1_scale; - // for kv cache reuse/chunked context - void* context_paged_kv_ptr = nullptr; - void* context_kv_cache_block_offsets_ptr = nullptr; - int32_t context_paged_kv_max_blocks_per_seq = 0; // for FP8 context qkv quantization float const* quant_scale_qkv = nullptr; }; @@ -94,6 +94,9 @@ struct MlaParams template void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, cudaStream_t stream); +template +void invokeMLAContextFp8Quantize(MlaParams& params, int total_kv_len, cudaStream_t stream); + template void invokeMLARopeGeneration(MlaParams& params, KVCacheBuffer kv_cache_buffer, cudaStream_t stream); @@ -112,10 +115,5 @@ void invokeMLARopeAppendPagedKVAssignQ(KVBlockArray& kv_cache, T* q_ptr, T* late int64_t const* cu_ctx_cached_kv_lens, int64_t const* cu_seq_lens, int const max_input_uncached_seq_len, float2 const* cos_sin_cache, size_t head_num, int nope_size, int rope_size, int lora_size, float const* kv_scale_orig_quant_ptr, cudaStream_t stream); - -template -__global__ void QuantizeCopyInputToFp8Kernel( - T_IN const* input_buffer, __nv_fp8_e4m3* output_fp8_buffer, int num_total_elements, float const* device_scale_ptr); - } // namespace kernels } // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 55c285be7bfa..c23c934d04b4 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -64,10 +64,12 @@ class RunnerBase virtual int64_t getWorkspaceSize(AttentionOp const& op, int const num_tokens, int const max_attention_window_size, int const num_gen_tokens) const = 0; + // typically, we use single qkv input, but for context MLA, we use separate qkv inputs virtual void run(AttentionOp& op, bool const is_context, int32_t const seq_offset, int32_t const num_seqs, int32_t const token_offset, int32_t const num_tokens, int32_t const predicted_tokens_per_seq, - torch::Tensor workspace, torch::Tensor output, torch::optional output_sf, torch::Tensor qkv, - torch::Tensor sequence_length, torch::Tensor host_past_key_value_lengths, torch::Tensor context_lengths, + torch::Tensor workspace, torch::Tensor output, torch::optional output_sf, torch::Tensor qkv_or_q, + torch::optional k, torch::optional v, torch::Tensor sequence_length, + torch::Tensor host_past_key_value_lengths, int32_t const total_kv_len, torch::Tensor context_lengths, torch::Tensor host_context_lengths, torch::optional kv_cache_block_offsets, torch::optional host_kv_cache_block_offsets, torch::optional host_kv_cache_pool_pointers, @@ -77,8 +79,6 @@ class RunnerBase torch::optional rotary_cos_sin, torch::optional latent_cache, torch::optional q_pe, torch::optional block_ids_per_seq, torch::optional mrope_rotary_cos_sin, torch::optional mrope_position_deltas, - torch::optional mla_context_paged_kv, - torch::optional mla_context_kv_cache_block_offsets, torch::optional softmax_stats_tensor, c10::ArrayRef> spec_decoding_tensor_params, torch::optional attention_sinks) const @@ -121,8 +121,9 @@ class Runner : public RunnerBase void run(AttentionOp& op, bool const is_context, int32_t const seq_offset, int32_t const num_seqs, int32_t const token_offset, int32_t const num_tokens, int32_t const predicted_tokens_per_seq, - torch::Tensor workspace, torch::Tensor output, torch::optional output_sf, torch::Tensor qkv, - torch::Tensor sequence_length, torch::Tensor host_past_key_value_lengths, torch::Tensor context_lengths, + torch::Tensor workspace, torch::Tensor output, torch::optional output_sf, torch::Tensor qkv_or_q, + torch::optional k, torch::optional v, torch::Tensor sequence_length, + torch::Tensor host_past_key_value_lengths, int32_t const total_kv_len, torch::Tensor context_lengths, torch::Tensor host_context_lengths, torch::optional kv_cache_block_offsets, torch::optional host_kv_cache_block_offsets, torch::optional host_kv_cache_pool_pointers, @@ -132,14 +133,14 @@ class Runner : public RunnerBase torch::optional rotary_cos_sin, torch::optional latent_cache, torch::optional q_pe, torch::optional block_ids_per_seq, torch::optional mrope_rotary_cos_sin, torch::optional mrope_position_deltas, - torch::optional mla_context_paged_kv, - torch::optional mla_context_kv_cache_block_offsets, torch::optional softmax_stats_tensor, c10::ArrayRef> spec_decoding_tensor_params, torch::optional attention_sinks) const override { - auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); - T* attention_input = static_cast(qkv.slice(0, token_offset).data_ptr()); + auto stream = at::cuda::getCurrentCUDAStream(qkv_or_q.get_device()); + T* attention_input = static_cast(qkv_or_q.slice(0, token_offset).data_ptr()); + T* k_ptr = nullptr; + T* v_ptr = nullptr; AttentionOutT* context_buf = static_cast(output.slice(0, token_offset).data_ptr()); TORCH_CHECK(!op.mFuseFp4Quant || output_sf.has_value()); void* context_buf_sf = op.mFuseFp4Quant ? output_sf->data_ptr() : nullptr; @@ -164,23 +165,36 @@ class Runner : public RunnerBase [[maybe_unused]] MlaParams mla_params; if (op.isMLAEnabled()) { - if (is_context && op.mPagedContextFMHA && op.mPagedKVCache) + if (is_context) { - TORCH_CHECK(mla_context_paged_kv.has_value()); - TORCH_CHECK(mla_context_kv_cache_block_offsets.has_value()); - - mla_params.context_paged_kv_ptr = mla_context_paged_kv->data_ptr(); - mla_params.context_kv_cache_block_offsets_ptr = mla_context_kv_cache_block_offsets->data_ptr(); - mla_params.context_paged_kv_max_blocks_per_seq = mla_context_kv_cache_block_offsets->size(-1); + if (latent_cache.has_value()) + { + mla_params.latent_cache = static_cast(latent_cache->data_ptr()); + } + else + { + // kv cache reuse / chunked context cases, latent_cache is not used + mla_params.latent_cache = nullptr; + } + TORCH_CHECK(k.has_value()); + TORCH_CHECK(v.has_value()); + TORCH_CHECK(k->dim() == 2); + TORCH_CHECK(v->dim() == 2); + TORCH_CHECK(k->strides()[1] == 1); + TORCH_CHECK(v->strides()[1] == 1); + + T* k_buf = static_cast(k->slice(0, token_offset).data_ptr()); + T* v_buf = static_cast(v->slice(0, token_offset).data_ptr()); + mla_params.k_buf = k_buf; + mla_params.v_buf = v_buf; + k_ptr = k_buf; + v_ptr = v_buf; } else { - // assume latent_cache has been written to paged kv cache by the PyTorch backend TORCH_CHECK(latent_cache.has_value()); mla_params.latent_cache = static_cast(latent_cache->data_ptr()); - } - if (!is_context) - { + TORCH_CHECK(q_pe.has_value()); TORCH_CHECK(q_pe->dim() == 3); TORCH_CHECK(q_pe->strides()[2] == 1); @@ -189,7 +203,7 @@ class Runner : public RunnerBase mla_params.q_pe_ld = q_pe->strides()[1]; mla_params.q_pe_stride = q_pe->strides()[0]; } - mla_params.attention_input_buf = attention_input; + mla_params.q_buf = attention_input; mla_params.context_buf = reinterpret_cast(context_buf); mla_params.cos_sin_cache = rotary_cos_sin_ptr; @@ -300,6 +314,7 @@ class Runner : public RunnerBase common_enqueue_params.host_primary_pool_pointer = host_primary_pool_pointer; common_enqueue_params.host_secondary_pool_pointer = host_secondary_pool_pointer; common_enqueue_params.num_tokens = num_tokens; + common_enqueue_params.total_kv_len = total_kv_len; common_enqueue_params.max_blocks_per_sequence = max_blocks_per_sequence; common_enqueue_params.sequence_lengths = sequence_lengths_ptr; common_enqueue_params.context_lengths = context_lengths_ptr; @@ -316,6 +331,8 @@ class Runner : public RunnerBase { enqueue_params.softmaxStatsPtr = static_cast(softmax_stats_tensor.value().data_ptr()); } + enqueue_params.k_ptr = k_ptr; + enqueue_params.v_ptr = v_ptr; if (op.isMLAEnabled()) { @@ -424,7 +441,7 @@ using torch_ext::trtllm::attention::AttentionInputType; void attention(torch::Tensor q, std::optional k, std::optional v, torch::Tensor& output, std::optional output_sf, std::optional out_dtype, std::optional workspace_, torch::Tensor sequence_length, torch::Tensor host_past_key_value_lengths, - torch::Tensor context_lengths, torch::Tensor host_context_lengths, torch::Tensor host_request_types, + torch::Tensor total_kv_lens, torch::Tensor context_lengths, torch::Tensor host_context_lengths, torch::Tensor host_request_types, std::optional kv_cache_block_offsets, std::optional host_kv_cache_block_offsets, std::optional host_kv_cache_pool_pointers, std::optional host_kv_cache_pool_mapping, std::optional cache_indirection, std::optional kv_scale_orig_quant, @@ -452,9 +469,9 @@ void attention(torch::Tensor q, std::optional k, std::optional k, std::optional k, std::optional(); int32_t const num_gen_tokens = is_gen_only ? num_tokens : num_tokens - num_ctx_tokens; + auto const ctx_total_kv_len = total_kv_lens.index({0}).item(); + auto const gen_total_kv_len = total_kv_lens.index({1}).item(); for (int32_t idx = num_contexts; idx < num_seqs; idx++) { @@ -661,7 +680,7 @@ void attention(torch::Tensor q, std::optional k, std::optional 0) && (attn_input_type != AttentionInputType::GenerationOnly)) @@ -671,12 +690,12 @@ void attention(torch::Tensor q, std::optional k, std::optionalrun(*op, /*is_context=*/true, seq_offset, /*num_seqs=*/num_contexts, token_offset, - /*num_tokens=*/num_ctx_tokens, predicted_tokens_per_seq, workspace, output, output_sf, qkv, sequence_length, - host_past_key_value_lengths, context_lengths, host_context_lengths, kv_cache_block_offsets, - host_kv_cache_block_offsets, host_kv_cache_pool_pointers, host_kv_cache_pool_mapping, cache_indirection, - kv_scale_orig_quant, kv_scale_quant_orig, out_scale, rotary_inv_freq, rotary_cos_sin, latent_cache, q_pe, - block_ids_per_seq, mrope_rotary_cos_sin, mrope_position_deltas, mla_context_paged_kv, - mla_context_kv_cache_block_offsets, softmax_stats_tensor, spec_decoding_tensor_params, attention_sinks); + /*num_tokens=*/num_ctx_tokens, predicted_tokens_per_seq, workspace, output, output_sf, qkv_or_q, k, v, + sequence_length, host_past_key_value_lengths, ctx_total_kv_len, context_lengths, host_context_lengths, + kv_cache_block_offsets, host_kv_cache_block_offsets, host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, cache_indirection, kv_scale_orig_quant, kv_scale_quant_orig, out_scale, + rotary_inv_freq, rotary_cos_sin, latent_cache, q_pe, block_ids_per_seq, mrope_rotary_cos_sin, + mrope_position_deltas, softmax_stats_tensor, spec_decoding_tensor_params, attention_sinks); } if ((num_generations > 0) && (attn_input_type != AttentionInputType::ContextOnly)) @@ -687,12 +706,12 @@ void attention(torch::Tensor q, std::optional k, std::optionalrun(*op, /*is_context=*/false, seq_offset, /*num_seqs=*/num_generations, token_offset, - /*num_tokens=*/num_gen_tokens, predicted_tokens_per_seq, workspace, output, output_sf, qkv, sequence_length, - host_past_key_value_lengths, context_lengths, host_context_lengths, kv_cache_block_offsets, - host_kv_cache_block_offsets, host_kv_cache_pool_pointers, host_kv_cache_pool_mapping, cache_indirection, - kv_scale_orig_quant, kv_scale_quant_orig, out_scale, rotary_inv_freq, rotary_cos_sin, latent_cache, q_pe, - block_ids_per_seq, mrope_rotary_cos_sin, mrope_position_deltas, mla_context_paged_kv, - mla_context_kv_cache_block_offsets, softmax_stats_tensor, spec_decoding_tensor_params, attention_sinks); + /*num_tokens=*/num_gen_tokens, predicted_tokens_per_seq, workspace, output, output_sf, qkv_or_q, k, v, + sequence_length, host_past_key_value_lengths, gen_total_kv_len, context_lengths, host_context_lengths, + kv_cache_block_offsets, host_kv_cache_block_offsets, host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, cache_indirection, kv_scale_orig_quant, kv_scale_quant_orig, out_scale, + rotary_inv_freq, rotary_cos_sin, latent_cache, q_pe, block_ids_per_seq, mrope_rotary_cos_sin, + mrope_position_deltas, softmax_stats_tensor, spec_decoding_tensor_params, attention_sinks); } TLLM_LOG_TRACE("Attention op stops at layer %d", layer_idx); diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 8bd4c49bbc94..187c2d14b8e5 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -610,6 +610,10 @@ def support_fused_rope(cls) -> bool: def support_fused_qkv(cls) -> bool: return False + @classmethod + def need_contiguous_qkv(cls) -> bool: + return False + @classmethod def support_mla(cls) -> bool: return False diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index b7a8b5a6fc57..057c2721a082 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -24,6 +24,7 @@ class TrtllmAttentionWrapper: sequence_length: torch.Tensor host_past_key_value_lengths: torch.Tensor + total_kv_lens: torch.Tensor context_lengths: torch.Tensor host_context_lengths: torch.Tensor host_request_types: torch.Tensor @@ -67,6 +68,7 @@ class TrtllmAttentionWrapper: qk_nope_head_dim: Optional[int] v_head_dim: Optional[int] attention_chunk_size: Optional[int] + softmax_stats_tensor: Optional[torch.Tensor] use_spec_decoding: bool is_spec_dec_tree: bool spec_decoding_position_offsets: Optional[torch.Tensor] @@ -154,6 +156,7 @@ def plan( beam_width: int = 1, sequence_length: torch.Tensor = ..., host_past_key_value_lengths: torch.Tensor = ..., + total_kv_lens: torch.Tensor = ..., context_lengths: torch.Tensor = ..., host_context_lengths: torch.Tensor = ..., host_request_types: torch.Tensor = ..., @@ -174,8 +177,6 @@ def plan( latent_cache: Optional[torch.Tensor] = None, q_pe: Optional[torch.Tensor] = None, mrope_config: Optional[dict] = None, - mla_context_paged_kv: Optional[torch.Tensor] = None, - mla_context_kv_cache_block_offsets: Optional[torch.Tensor] = None, softmax_stats_tensor: Optional[torch.Tensor] = None, is_spec_decoding_enabled: bool = False, use_spec_decoding: bool = False, @@ -201,6 +202,7 @@ def plan( beam_width (int): Beam width in beam search. sequence_length (torch.Tensor): The length of each sequence with shape (batch_size) on GPU. host_past_key_value_lengths (torch.Tensor): Same as sequence_length, but on CPU. + total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and generation requests, with shape (2) on CPU. context_lengths (torch.Tensor): The context-phase sequence length of each request with shape (batch_size) on GPU. host_context_lengths (torch.Tensor): Same as context_lengths, but on CPU. host_request_types (torch.Tensor): The tensor that indicates whether a request is in context or generation phase, with shape (batch_size) on CPU. @@ -216,8 +218,6 @@ def plan( out_scale_sf (torch.Tensor): The tensor to store the global scale for NVFP4 scaling factors, with shape (1) on GPU. use_paged_context_fmha (bool): Sets the mPagedContextFMHA attribute in the op runner. mrope_config (dict): The dictionary containing the mRope configuration. - mla_context_paged_kv (torch.Tensor): The paged KV cache for MLA context, for kv cache reuse/chunked context. - mla_context_kv_cache_block_offsets (torch.Tensor): The block offsets for the paged KV cache for MLA context, for kv cache reuse/chunked context. softmax_stats_tensor (torch.Tensor): The tensor to store the softmax statistics (max/sum) attention_sinks (torch.Tensor): The attention sinks (additional value in the denominator of the softmax) with shape of (num_heads_q) on GPU. """ @@ -230,6 +230,7 @@ def plan( self.beam_width = beam_width self.sequence_length = sequence_length self.host_past_key_value_lengths = host_past_key_value_lengths + self.total_kv_lens = total_kv_lens self.context_lengths = context_lengths self.host_context_lengths = host_context_lengths self.host_request_types = host_request_types @@ -253,8 +254,6 @@ def plan( self.mrope_position_deltas = mrope_config.get( 'mrope_position_deltas') if mrope_config is not None else None self.block_ids_per_seq = block_ids_per_seq - self.mla_context_paged_kv = mla_context_paged_kv - self.mla_context_kv_cache_block_offsets = mla_context_kv_cache_block_offsets self.softmax_stats_tensor = softmax_stats_tensor self.attention_sinks = attention_sinks @@ -361,18 +360,12 @@ def run( else: raise ValueError("Unexpected attention mask type") else: - assert is_fused_qkv if self.attention_input_type == AttentionInputType.context_only: - if self.use_paged_context_fmha: - assert self.mla_context_paged_kv is not None - assert self.mla_context_kv_cache_block_offsets is not None - qkv_hidden_size = self.num_heads * (self.qk_nope_head_dim + - self.qk_rope_head_dim) - else: - qkv_hidden_size = self.num_heads * ( - 2 * (self.qk_nope_head_dim + self.qk_rope_head_dim) - ) + self.num_kv_heads * self.v_head_dim + assert not is_fused_qkv + qkv_hidden_size = self.num_heads * (self.qk_nope_head_dim + + self.qk_rope_head_dim) elif self.attention_input_type == AttentionInputType.generation_only: + assert is_fused_qkv qkv_hidden_size = self.num_heads * (self.kv_lora_rank + self.qk_rope_head_dim) else: @@ -430,6 +423,7 @@ def run( self.workspace, self.sequence_length, self.host_past_key_value_lengths, + self.total_kv_lens, self.context_lengths, self.host_context_lengths, self.host_request_types, @@ -479,8 +473,6 @@ def run( self.v_head_dim, self.mrope_rotary_cos_sin, self.mrope_position_deltas, - self.mla_context_paged_kv, - self.mla_context_kv_cache_block_offsets, self.attention_chunk_size, self.softmax_stats_tensor, spec_decoding_bool_params, @@ -623,6 +615,7 @@ def __post_init__(self) -> None: self.kv_lens = torch.empty_like(self.kv_lens_cuda, device='cpu', pin_memory=True) + self.total_kv_lens = torch.empty(2, device='cpu', dtype=torch.int) self.host_request_types = torch.empty_like(self.prompt_lens_cpu) # For debugging, can use it to call the wrapper's plan function @@ -700,6 +693,7 @@ def __post_init__(self) -> None: ) def prepare(self) -> None: + # print("prepare TrtllmAttentionMetadata") extra_attrs = get_model_extra_attrs() # If model extra attrs is set, attention_metadata is setup in executor. if extra_attrs is None: @@ -746,6 +740,10 @@ def prepare(self) -> None: kv_lens + self.kv_cache_params.num_extra_kv_tokens) self.kv_lens_cuda[:self.num_seqs].copy_( kv_lens[:self.num_seqs].pin_memory(), non_blocking=True) + # total kv lens for context requests and generation requests, without extra tokens + self.total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item() + self.total_kv_lens[1] = kv_lens[self.num_contexts:self.num_seqs].sum( + ).item() self.host_request_types[:self.num_contexts].fill_(0) self.host_request_types[self.num_contexts:self.num_seqs].fill_(1) @@ -1101,8 +1099,6 @@ def forward( q_pe: Optional[torch.Tensor] = None, mrope_config: Optional[dict] = None, attention_window_size: Optional[int] = None, - mla_context_paged_kv: Optional[torch.Tensor] = None, - mla_context_kv_cache_block_offsets: Optional[torch.Tensor] = None, softmax_stats_tensor: Optional[torch.Tensor] = None, enable_attn_nvfp4_output: bool = True, output: Optional[torch.Tensor] = None, @@ -1149,6 +1145,7 @@ def forward( beam_width=metadata.beam_width, sequence_length=metadata.kv_lens_cuda_runtime, host_past_key_value_lengths=metadata.kv_lens_runtime, + total_kv_lens=metadata.total_kv_lens, context_lengths=metadata.prompt_lens_cuda_runtime, host_context_lengths=metadata.prompt_lens_cpu_runtime, host_request_types=metadata.host_request_types_runtime, @@ -1169,9 +1166,6 @@ def forward( latent_cache=latent_cache, q_pe=q_pe, mrope_config=mrope_config, - mla_context_paged_kv=mla_context_paged_kv, - mla_context_kv_cache_block_offsets= - mla_context_kv_cache_block_offsets, softmax_stats_tensor=softmax_stats_tensor, is_spec_decoding_enabled=metadata.is_spec_decoding_enabled, use_spec_decoding=metadata.use_spec_decoding, @@ -1217,6 +1211,10 @@ def support_fused_rope(cls) -> bool: def support_fused_qkv(cls) -> bool: return True + @classmethod + def need_contiguous_qkv(cls) -> bool: + return True + @classmethod def support_mla(cls) -> bool: return True @@ -1248,7 +1246,7 @@ def load_paged_kv_cache_for_mla( self, metadata: TrtllmAttentionMetadata, out_dtype: torch.dtype, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor]: assert out_dtype in [torch.float16, torch.bfloat16, torch.float32] assert self.is_mla_enable and self.mla_params is not None assert metadata.kv_cache_manager is not None @@ -1290,15 +1288,19 @@ def load_chunked_kv_cache_for_mla( num_ctx_cached_tokens: int, cu_chunked_seq_len: torch.Tensor, out_dtype: torch.dtype, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor]: assert out_dtype in [torch.float16, torch.bfloat16, torch.float32] assert self.is_mla_enable and self.mla_params is not None assert metadata.kv_cache_manager is not None if metadata.max_ctx_cached_token_len == 0: - return torch.empty((0, metadata.kv_cache_manager.head_dim), - dtype=out_dtype, - device=cu_chunked_seq_len.device) + empty_kv = torch.empty((0, self.mla_params.kv_lora_rank), + dtype=out_dtype, + device=cu_chunked_seq_len.device) + empty_k_pe = torch.empty((0, self.mla_params.qk_rope_head_dim), + dtype=out_dtype, + device=cu_chunked_seq_len.device) + return empty_kv, empty_k_pe sink_token_length = 0 beam_width = 1 diff --git a/tensorrt_llm/_torch/metadata.py b/tensorrt_llm/_torch/metadata.py index 25bb9d9a7eaf..fd076cc4cab0 100644 --- a/tensorrt_llm/_torch/metadata.py +++ b/tensorrt_llm/_torch/metadata.py @@ -28,7 +28,7 @@ class KVCacheParams: # The number of sink tokens for each layer. host_sink_token_length: Optional[torch.Tensor] = None # The number of extra kv for draft tokens - num_extra_kv_tokens: Optional[List[int]] = 0 + num_extra_kv_tokens: Optional[int] = 0 class CacheType(Enum): diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 5210d341d461..8eb9acfada22 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -764,7 +764,6 @@ def forward( enable_allreduce=not (self.disable_attn_allreduce)), **kwargs, ) - if isinstance(self.mlp, Deepseekv3MoE): return self.forward_MoE( hidden_states=hidden_states, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index f9f1eaa1df1d..48e47cff1ca0 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -835,7 +835,7 @@ def yarn_get_mscale(scale=1, mscale=1): self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] self.rope_fusion = self.mha.support_fused_rope() - self.support_fused_qkv = self.mha.support_fused_qkv() + self.need_contiguous_qkv = self.mha.need_contiguous_qkv() self.rotary_emb = None self.apply_rotary_emb = not self.rope_fusion if self.apply_rotary_emb: @@ -1054,10 +1054,11 @@ def forward_impl(self, attn_output_gen = None return output - def _maybe_concat_qkv(self, q, k, v): - if k is not None and v is not None and self.support_fused_qkv: - qkv = torch.concat([q, k, v], dim=-1) - q, k, v = qkv, None, None + def _maybe_contiguous(self, q, k, v): + if self.need_contiguous_qkv: + q, k, v = [ + x.contiguous() if x is not None else None for x in (q, k, v) + ] return q, k, v def forward_context_default( @@ -1068,6 +1069,7 @@ def forward_context_default( attn_metadata: AttentionMetadata, latent_cache: Optional[torch.Tensor] = None, output: Optional[torch.Tensor] = None) -> torch.Tensor: + # print(f"=========forward_context_default layer_idx: {self.layer_idx}, q: {q.shape}, compressed_kv: {compressed_kv.shape}, k_pe: {k_pe.shape}, latent_cache: {latent_cache.shape}") kv = self.kv_b_proj(compressed_kv) k_nope, v = kv.split( [ @@ -1085,8 +1087,8 @@ def forward_context_default( self.qk_rope_head_dim) k = k.view(-1, self.num_heads * self.qk_head_dim) - # May concat q(including q_pe), k + k_pe, v together - q, k, v = self._maybe_concat_qkv(q, k, v) + # q, k and v may should be contiguous + q, k, v = self._maybe_contiguous(q, k, v) # out_scale = getattr(self.o_proj, "inv_input_scale", None) out_scale = None # Currently we use BF16 MHA for context phase @@ -1101,6 +1103,11 @@ def forward_context_default( out_scale=out_scale, output=output, ) + # if self.layer_idx == 0: + # print(f"======= q shape: {q.shape}, value:\n{q}") + # print(f"======= k shape: {k.shape}, value:\n{k}") + # print(f"======= v shape: {v.shape}, value:\n{v}") + # print(f"======= attn_output shape: {attn_output.shape}, value:\n{attn_output}") return attn_output @@ -1111,6 +1118,7 @@ def forward_context_with_cached_kv( attn_metadata: AttentionMetadata, output: Optional[torch.Tensor] = None, ) -> torch.Tensor: + # print(f"=========forward_context_with_cached_kv layer_idx: {self.layer_idx}, q: {q.shape}, latent_cache: {latent_cache.shape}") assert latent_cache is not None trtllm_attention = cast(TrtllmAttention, self.mha) @@ -1139,50 +1147,39 @@ def forward_context_with_cached_kv( ], -1, ) - full_k_nope = full_k_nope.view(-1, self.num_heads, - self.qk_nope_head_dim) - full_v = full_v.view(-1, self.num_heads, self.v_head_dim) - - # build paged_full_kv - tokens_per_block = attn_metadata.kv_cache_manager.tokens_per_block - # paged_full_kv will be initialized to 0 in the kernel to avoid NaN - paged_full_kv = torch.empty([ - attn_metadata.num_contexts, 2, - (attn_metadata.max_ctx_kv_len + tokens_per_block - 1) // - tokens_per_block, self.num_heads, tokens_per_block, - max(self.qk_nope_head_dim + self.qk_rope_head_dim, self.v_head_dim) - ], - dtype=q.dtype, - device=q.device) - mla_context_kv_cache_block_offsets = trtllm_attention.set_paged_kv_cache_for_mla( - paged_full_kv, - full_k_nope, - full_v, - full_k_pe, - attn_metadata, - ) + + full_k = torch.empty( + (full_k_nope.shape[0], self.num_heads, self.qk_head_dim), + dtype=full_k_nope.dtype, + device=full_k_nope.device) + full_k[..., :self.qk_nope_head_dim] = full_k_nope.view( + -1, self.num_heads, self.qk_nope_head_dim) + full_k[..., + self.qk_nope_head_dim:] = full_k_pe.view(-1, 1, + self.qk_rope_head_dim) + full_k = full_k.view(-1, self.num_heads * self.qk_head_dim) + + # q, full_k and full_v may should be contiguous + q, full_k, full_v = self._maybe_contiguous(q, full_k, full_v) # release pytorch activation memory full_compressed_kv = None full_k_pe = None full_kv = None full_k_nope = None - full_v = None # out_scale = getattr(self.o_proj, "inv_input_scale", None) out_scale = None # Currently we use BF16 MHA for context phase attn_output = self.mha.forward( q, - None, - None, + full_k, + full_v, attn_metadata, attention_input_type=AttentionInputType.context_only, - latent_cache=None, + latent_cache= + None, # need to be None to differentiate from normal context phase out_scale=out_scale, - mla_context_paged_kv=paged_full_kv, - mla_context_kv_cache_block_offsets= - mla_context_kv_cache_block_offsets, output=output, ) @@ -1204,7 +1201,7 @@ def forward_context_with_chunked_prefill( # determine the number of loop # currently we assume that the chunk size is the same as the max_num_tokens - chunk_size = attn_metadata.runtime_features.chunk_size + attn_metadata.runtime_features.chunk_size chunked_loop_num = attn_metadata.chunked_loop_num # [toal_token_q, num_heads, 2] -> [toal_token_q, num_heads] float2 @@ -1229,6 +1226,7 @@ def forward_context_with_chunked_prefill( # use fake cached_cu_seq_len for chunked loop origin_kv_lens_cuda_runtime = attn_metadata.kv_lens_cuda_runtime origin_kv_lens_runtime = attn_metadata.kv_lens_runtime + origin_ctx_total_kv_lens = attn_metadata.total_kv_lens[0] for loop_idx in range(chunked_loop_num): # {b, chunked_unit_size, h, kv_lora_rank + qk_rope_head_dim} zero padded @@ -1246,45 +1244,50 @@ def forward_context_with_chunked_prefill( # up proj to uncompressed kv # [tokens, 2, h, kv_dim], without rope_dim chunked_kv = self.kv_b_proj(chunked_compressed_kv) - - # build full_kv - # full_kv {B, 2, chunk_size / tokens_per_block, h, tokens_per_block, kv_dim + rope_dim} - tokens_per_block = attn_metadata.kv_cache_manager.tokens_per_block - full_kv = torch.zeros([ - attn_metadata.num_contexts, 2, - (chunk_size + tokens_per_block - 1) // tokens_per_block, - self.num_heads, tokens_per_block, - max(self.qk_nope_head_dim + self.qk_rope_head_dim, - self.v_head_dim) - ], - dtype=q.dtype, - device=q.device) - mla_kv_cache_block_offsets = trtllm_attention.set_chunked_kv_cache_for_mla( - full_kv, - chunked_kv, - chunked_k_pe, - cu_chunked_seq_len=temp_cu_chunked_seq_len, - cached=True, - metadata=attn_metadata) + chunked_k_nope, chunked_v = chunked_kv.split( + [ + self.num_heads * self.qk_nope_head_dim, + self.num_heads * self.v_head_dim + ], + -1, + ) + chunked_k = torch.empty( + (chunked_k_nope.shape[0], self.num_heads, self.qk_head_dim), + dtype=chunked_k_nope.dtype, + device=chunked_k_nope.device) + chunked_k[..., :self.qk_nope_head_dim] = chunked_k_nope.view( + -1, self.num_heads, self.qk_nope_head_dim) + chunked_k[..., self.qk_nope_head_dim:] = chunked_k_pe.view( + -1, 1, self.qk_rope_head_dim) + chunked_k = chunked_k.view(-1, self.num_heads * self.qk_head_dim) + + q, chunked_k, chunked_v = self._maybe_contiguous( + q, chunked_k, chunked_v) + + # release pytorch activation memory + chunked_compressed_kv = None + chunked_k_pe = None + chunked_kv = None + chunked_k_nope = None # copy chunked_seq_len to replace kv_lens_runtime attn_metadata.kv_lens_runtime = attn_metadata.host_chunked_seq_len[ loop_idx] attn_metadata.kv_lens_cuda_runtime = attn_metadata.chunked_seq_len[ loop_idx] + attn_metadata.total_kv_lens[0] = total_ctx_chunked_tokens + out_scale = None # do not apply mask for attention within loop temp_attn_output = self.mha.forward( q, - None, - None, + chunked_k, + chunked_v, attn_metadata, attention_input_type=AttentionInputType.context_only, latent_cache=None, out_scale=out_scale, attention_mask=PredefinedAttentionMask.FULL, - mla_context_paged_kv=full_kv, - mla_context_kv_cache_block_offsets=mla_kv_cache_block_offsets, softmax_stats_tensor=self.temp_softmax_stats_tensor, output=temp_attn_output, ) @@ -1299,41 +1302,41 @@ def forward_context_with_chunked_prefill( _, k_pe = latent_cache.view([ -1, self.kv_lora_rank + self.qk_rope_head_dim ]).split([self.kv_lora_rank, self.qk_rope_head_dim], -1) - k_pe = k_pe.contiguous() # final round of attention - # out_scale = getattr(self.o_proj, "inv_input_scale", None) - out_scale = None # Currently we use BF16 MHA for context phase + k_nope, v = kv.split( + [ + self.num_heads * self.qk_nope_head_dim, + self.num_heads * self.v_head_dim + ], + -1, + ) + k = torch.empty_like(q).view(-1, self.num_heads, self.qk_head_dim) + k[..., :self.qk_nope_head_dim] = k_nope.view(-1, self.num_heads, + self.qk_nope_head_dim) + k[..., self.qk_nope_head_dim:] = k_pe.view(-1, 1, self.qk_rope_head_dim) + k = k.view(-1, self.num_heads * self.qk_head_dim) + + q, k, v = self._maybe_contiguous(q, k, v) - tokens_per_block = attn_metadata.kv_cache_manager.tokens_per_block - full_kv = torch.zeros([ - attn_metadata.num_contexts, 2, - (attn_metadata.max_ctx_seq_len + tokens_per_block - 1) // - tokens_per_block, self.num_heads, tokens_per_block, - max(self.qk_nope_head_dim + self.qk_rope_head_dim, self.v_head_dim) - ], - dtype=q.dtype, - device=q.device) - mla_kv_cache_block_offsets = trtllm_attention.set_chunked_kv_cache_for_mla( - full_kv, - kv, - k_pe, - cu_chunked_seq_len=None, - cached=False, - metadata=attn_metadata) # copy q_lens to replace kv_lens_runtime attn_metadata.kv_lens_runtime = attn_metadata.prompt_lens_cpu_runtime attn_metadata.kv_lens_cuda_runtime = attn_metadata.prompt_lens_cuda_runtime + attn_metadata.total_kv_lens[ + 0] = attn_metadata.prompt_lens_cpu_runtime[:attn_metadata. + num_contexts].sum().item( + ) + + # out_scale = getattr(self.o_proj, "inv_input_scale", None) + out_scale = None # Currently we use BF16 MHA for context phase temp_attn_output = self.mha.forward( q, - None, - None, + k, + v, attn_metadata, attention_input_type=AttentionInputType.context_only, latent_cache=None, out_scale=out_scale, - mla_context_paged_kv=full_kv, - mla_context_kv_cache_block_offsets=mla_kv_cache_block_offsets, softmax_stats_tensor=self.temp_softmax_stats_tensor, output=temp_attn_output, ) @@ -1345,6 +1348,7 @@ def forward_context_with_chunked_prefill( # copy back kv_lens_runtime and kv_lens_cuda_runtime attn_metadata.kv_lens_runtime = origin_kv_lens_runtime attn_metadata.kv_lens_cuda_runtime = origin_kv_lens_cuda_runtime + attn_metadata.total_kv_lens[0] = origin_ctx_total_kv_lens return attn_output @@ -1484,7 +1488,7 @@ def forward( attn_metadata: AttentionMetadata, all_reduce_params: Optional[AllReduceParams] = None, ) -> torch.Tensor: - + # print(f"=========forward layer_idx: {self.layer_idx}, position_ids: {position_ids.shape}, hidden_states: {hidden_states.shape}") attn_output = self.create_output(hidden_states) if self.register_to_config: torch.ops.trtllm.mla_custom_op_inplace(hidden_states, position_ids, From 36d3b504303f2fdd23fd11b5e7d081ca7a84eb56 Mon Sep 17 00:00:00 2001 From: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> Date: Thu, 31 Jul 2025 20:31:55 -0700 Subject: [PATCH 02/14] support Hopper and remove copy of tensor V Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 6 +- .../fmhaRunner.cpp | 28 ++++++++++ .../fused_multihead_attention_common.h | 6 +- cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp | 25 +-------- .../trtllmGenKernels/fmha/fmhaRunnerParams.h | 2 + .../trtllmGenKernels/fmha/kernelParams.h | 5 ++ .../_torch/attention_backend/interface.py | 4 -- .../_torch/attention_backend/trtllm.py | 4 -- tensorrt_llm/_torch/modules/attention.py | 56 ++++--------------- 9 files changed, 58 insertions(+), 78 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 87daafcbff42..2d077bcf82f9 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -1720,6 +1720,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea // - v_ptr: [B, S, H_kv, D_v], which supports variable sequence length // - cu_q_seqlens: the cumulative query sequence lengths, needed for variable sequence length. // - cu_kv_seqlens: the cumulative kv sequence lengths, needed for variable sequence length. + // - total_kv_len: the total kv sequence lengths, needed for variable sequence length. // Construct the fmha params for running kernels. MHARunnerParams fmhaParams{}; @@ -2620,8 +2621,8 @@ int AttentionOp::initialize() noexcept // the wrong kernel, no matter mIsGenerationMLA is true or false if (mIsMLAEnabled) { - // context MLA always use separate_qkv layout - fmhaParams.attentionInputLayout = AttentionInputLayout::SEPARATE_QKV; + // Context MLA always use separate_q_k_v layout + fmhaParams.attentionInputLayout = AttentionInputLayout::SEPARATE_Q_K_V; // Context attention of MLA is different fmhaParams.numKvHeads = mNumHeads; fmhaParams.headSize = mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim; @@ -2630,6 +2631,7 @@ int AttentionOp::initialize() noexcept // attention op and that could fail to create the FmhaDispatcher for context phase. // Luckily, for deepseek, qk_nope_head_dim is the same as v_head_dim in context phase. fmhaParams.headSizeV = mMLAParams.qk_nope_head_dim; + fmhaParams.headSizeQkNope = mMLAParams.qk_nope_head_dim; } fmhaParams.qScaling = mQScaling; fmhaParams.attnLogitSoftcappingScale = mAttnLogitSoftcappingScale; diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp index cf6e1032007d..1c2ce91fb048 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp @@ -179,6 +179,28 @@ void FusedMHARunnerV2::setupKernelParams(MHARunnerParams runnerParams) // Thus, v_stride_in_bytes always equals to k_stride_in_bytes so far. mKernelParams.v_stride_in_bytes = mKernelParams.k_stride_in_bytes; } + else if (mFixedParams.attentionInputLayout == AttentionInputLayout::SEPARATE_Q_K_V) + { + // Separate QKV input layout, [total_kv_seqlen, H_KV, D] + [total_kv_seqlen, H_KV, DV] + mKernelParams.k_ptr = runnerParams.kPtr; + mKernelParams.v_ptr = runnerParams.vPtr; + // Tensor K is contiguous. + mKernelParams.k_stride_in_bytes + = get_size_in_bytes(mFixedParams.numKvHeads * mFixedParams.headSize, mFixedParams.dataType); + if (mFixedParams.headSizeQkNope > 0) + { + // Tensor V for context MLA is not contiguous, so we need to add the headSizeQkNope to the stride. + mKernelParams.v_stride_in_bytes = get_size_in_bytes( + mFixedParams.numKvHeads * (mFixedParams.headSizeQkNope + mFixedParams.headSizeV), + mFixedParams.dataType); + } + else + { + // Tensor V is contiguous for other cases. + mKernelParams.v_stride_in_bytes + = get_size_in_bytes(mFixedParams.numKvHeads * mFixedParams.headSizeV, mFixedParams.dataType); + } + } } mKernelParams.o_ptr = runnerParams.outputPtr; @@ -623,6 +645,12 @@ void FusedMHARunnerV2::setTmaDescriptors(MHARunnerParams runnerParams) k_ptr = reinterpret_cast(mKernelParams.kv_ptr); v_ptr = k_ptr + h_kv * d_in_bytes; } + else if (layout == AttentionInputLayout::SEPARATE_Q_K_V) + { + // Layout: [total_kv_seqlen, H_KV, D] + [total_kv_seqlen, H_KV, DV] + k_ptr = reinterpret_cast(mKernelParams.k_ptr); + v_ptr = reinterpret_cast(mKernelParams.v_ptr); + } Multiple_tma_descriptor<3> kv_tma_descriptor; // K diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h index d8fe8c49fc68..77bbbba876af 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h @@ -84,7 +84,7 @@ enum class AttentionInputLayout Q_PAGED_KV, // Q has contiguous [B, S, H, D] layout, while K has contiguous [B, S, H_kv, D] layout, and V has // contiguous [B, S, H_kv, D_v] layout. Only used for context MLA now. - SEPARATE_QKV, + SEPARATE_Q_K_V, }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -118,6 +118,8 @@ struct MHARunnerFixedParams int headSize; // The head size of V. int headSizeV = 0; + // The head size of Q/K non-RoPE part, only used for MLA now. + int headSizeQkNope = 0; // The scaling applied to bmm1_scale. float qScaling; // The attention logit softcapping scale. @@ -169,7 +171,7 @@ struct MHARunnerFixedParams case AttentionInputLayout::PACKED_QKV: output += "packed_qkv"; break; case AttentionInputLayout::Q_CONTIGUOUS_KV: output += "q_contiguous_kv"; break; case AttentionInputLayout::Q_PAGED_KV: output += "q_paged_kv"; break; - case AttentionInputLayout::SEPARATE_QKV: output += "separate_qkv"; break; + case AttentionInputLayout::SEPARATE_Q_K_V: output += "separate_q_k_v"; break; default: output += std::to_string(static_cast(attentionInputLayout)) + " (unknown)"; break; } diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp index 5e9adf69cd26..2db0317adc18 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp @@ -36,7 +36,7 @@ QkvLayout AttentionInputLayoutToQkvLayout(AttentionInputLayout layout) { return QkvLayout::PagedKv; } - else if (layout == AttentionInputLayout::SEPARATE_QKV) + else if (layout == AttentionInputLayout::SEPARATE_Q_K_V) { return QkvLayout::SeparateQkv; } @@ -152,7 +152,7 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) maxBlocksPerSeq = pagedKvCache.mMaxBlocksPerSeq; numTokensPerBlock = pagedKvCache.mTokensPerBlock; } - else if (mFixedParams.attentionInputLayout == AttentionInputLayout::SEPARATE_QKV) + else if (mFixedParams.attentionInputLayout == AttentionInputLayout::SEPARATE_Q_K_V) { qkvLayout = kernels::QkvLayout::SeparateQkv; } @@ -189,6 +189,7 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) // Assume same headDim for Qk and V here. tllmRunnerParams.mHeadDimQk = mFixedParams.headSize; tllmRunnerParams.mHeadDimV = mFixedParams.headSizeV; + tllmRunnerParams.mHeadDimQkNope = mFixedParams.headSizeQkNope; tllmRunnerParams.mNumHeadsQ = mFixedParams.numQHeads; tllmRunnerParams.mNumHeadsKv = mFixedParams.numKvHeads; tllmRunnerParams.mNumHeadsQPerKv = tllmRunnerParams.mNumHeadsQ / tllmRunnerParams.mNumHeadsKv; @@ -211,26 +212,6 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) tllmRunnerParams.softmaxStatsPtr = reinterpret_cast(runnerParams.softmaxStatsPtr); tllmRunnerParams.stream = runnerParams.stream; - // if (runnerParams.totalQSeqLen < 8192) { - // std::vector cumSeqLensQ(runnerParams.b + 1); - // std::vector cumSeqLensKv(runnerParams.b + 1); - // cudaMemcpy(cumSeqLensQ.data(), runnerParams.cuQSeqLenPtr, (runnerParams.b + 1) * sizeof(int), - // cudaMemcpyDeviceToHost); cudaMemcpy(cumSeqLensKv.data(), runnerParams.cuKvSeqLenPtr, (runnerParams.b + 1) - // * sizeof(int), cudaMemcpyDeviceToHost); std::stringstream ss; for (int i = 0; i < runnerParams.b + 1; - // i++) { - // ss << cumSeqLensQ[i] << " "; - // } - // TLLM_LOG_INFO("cumSeqLensQ: %s", ss.str().c_str()); - // ss.str(""); - // for (int i = 0; i < runnerParams.b + 1; i++) { - // ss << cumSeqLensKv[i] << " "; - // } - // TLLM_LOG_INFO("cumSeqLensKv: %s", ss.str().c_str()); - // TLLM_LOG_INFO("mMaxSeqLenCacheKv: %d, mMaxSeqLenQ: %d, mMaxSeqLenKv: %d", - // tllmRunnerParams.mMaxSeqLenCacheKv, tllmRunnerParams.mMaxSeqLenQ, tllmRunnerParams.mMaxSeqLenKv); - // TLLM_LOG_INFO("mSumOfSeqLensQ: %d, mSumOfSeqLensKv: %d", tllmRunnerParams.mSumOfSeqLensQ, - // tllmRunnerParams.mSumOfSeqLensKv); - // } mTllmGenFMHARunner->run(tllmRunnerParams); } else diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h index c9c35c228e64..0f042ec74b94 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h @@ -238,6 +238,8 @@ struct TllmGenFmhaRunnerParams int mHeadDimQk; // Head dimension for V. int mHeadDimV; + // Head dimension for Q/K non-RoPE part, only used for MLA now. + int mHeadDimQkNope; // Number of heads for Q and K/V. int mNumHeadsQ, mNumHeadsKv, mNumHeadsQPerKv; // The batch size. diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h index 632ecdb97bdf..c8a9936648ad 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h @@ -357,6 +357,11 @@ struct KernelParams { strideKeysVals = maxHeadDimKv; } + else if (isSeparateQkv(options.mQkvLayout) && !isK && options.mHeadDimQkNope > 0) + { + // Tensor V for context MLA is not contiguous, so we need to add the headSizeQkNope to the stride. + strideKeysVals = options.mNumHeadsKv * (options.mHeadDimQkNope + options.mHeadDimV); + } // The stride between heads. int32_t strideHeads{isK ? options.mHeadDimQk : options.mHeadDimV}; diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 187c2d14b8e5..8bd4c49bbc94 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -610,10 +610,6 @@ def support_fused_rope(cls) -> bool: def support_fused_qkv(cls) -> bool: return False - @classmethod - def need_contiguous_qkv(cls) -> bool: - return False - @classmethod def support_mla(cls) -> bool: return False diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 057c2721a082..d52df69bebc1 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1211,10 +1211,6 @@ def support_fused_rope(cls) -> bool: def support_fused_qkv(cls) -> bool: return True - @classmethod - def need_contiguous_qkv(cls) -> bool: - return True - @classmethod def support_mla(cls) -> bool: return True diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 48e47cff1ca0..ee17c1df091f 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -835,7 +835,6 @@ def yarn_get_mscale(scale=1, mscale=1): self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] self.rope_fusion = self.mha.support_fused_rope() - self.need_contiguous_qkv = self.mha.need_contiguous_qkv() self.rotary_emb = None self.apply_rotary_emb = not self.rope_fusion if self.apply_rotary_emb: @@ -1054,13 +1053,6 @@ def forward_impl(self, attn_output_gen = None return output - def _maybe_contiguous(self, q, k, v): - if self.need_contiguous_qkv: - q, k, v = [ - x.contiguous() if x is not None else None for x in (q, k, v) - ] - return q, k, v - def forward_context_default( self, q: torch.Tensor, @@ -1069,7 +1061,6 @@ def forward_context_default( attn_metadata: AttentionMetadata, latent_cache: Optional[torch.Tensor] = None, output: Optional[torch.Tensor] = None) -> torch.Tensor: - # print(f"=========forward_context_default layer_idx: {self.layer_idx}, q: {q.shape}, compressed_kv: {compressed_kv.shape}, k_pe: {k_pe.shape}, latent_cache: {latent_cache.shape}") kv = self.kv_b_proj(compressed_kv) k_nope, v = kv.split( [ @@ -1087,9 +1078,6 @@ def forward_context_default( self.qk_rope_head_dim) k = k.view(-1, self.num_heads * self.qk_head_dim) - # q, k and v may should be contiguous - q, k, v = self._maybe_contiguous(q, k, v) - # out_scale = getattr(self.o_proj, "inv_input_scale", None) out_scale = None # Currently we use BF16 MHA for context phase @@ -1103,11 +1091,6 @@ def forward_context_default( out_scale=out_scale, output=output, ) - # if self.layer_idx == 0: - # print(f"======= q shape: {q.shape}, value:\n{q}") - # print(f"======= k shape: {k.shape}, value:\n{k}") - # print(f"======= v shape: {v.shape}, value:\n{v}") - # print(f"======= attn_output shape: {attn_output.shape}, value:\n{attn_output}") return attn_output @@ -1118,7 +1101,6 @@ def forward_context_with_cached_kv( attn_metadata: AttentionMetadata, output: Optional[torch.Tensor] = None, ) -> torch.Tensor: - # print(f"=========forward_context_with_cached_kv layer_idx: {self.layer_idx}, q: {q.shape}, latent_cache: {latent_cache.shape}") assert latent_cache is not None trtllm_attention = cast(TrtllmAttention, self.mha) @@ -1148,20 +1130,13 @@ def forward_context_with_cached_kv( -1, ) - full_k = torch.empty( - (full_k_nope.shape[0], self.num_heads, self.qk_head_dim), - dtype=full_k_nope.dtype, - device=full_k_nope.device) - full_k[..., :self.qk_nope_head_dim] = full_k_nope.view( - -1, self.num_heads, self.qk_nope_head_dim) - full_k[..., - self.qk_nope_head_dim:] = full_k_pe.view(-1, 1, - self.qk_rope_head_dim) + full_k_nope = full_k_nope.view(-1, self.num_heads, + self.qk_nope_head_dim) + full_k_pe = full_k_pe.view(-1, 1, self.qk_rope_head_dim) + full_k = torch.cat( + (full_k_nope, full_k_pe.expand(-1, self.num_heads, -1)), dim=-1) full_k = full_k.view(-1, self.num_heads * self.qk_head_dim) - # q, full_k and full_v may should be contiguous - q, full_k, full_v = self._maybe_contiguous(q, full_k, full_v) - # release pytorch activation memory full_compressed_kv = None full_k_pe = None @@ -1251,18 +1226,14 @@ def forward_context_with_chunked_prefill( ], -1, ) - chunked_k = torch.empty( - (chunked_k_nope.shape[0], self.num_heads, self.qk_head_dim), - dtype=chunked_k_nope.dtype, - device=chunked_k_nope.device) - chunked_k[..., :self.qk_nope_head_dim] = chunked_k_nope.view( - -1, self.num_heads, self.qk_nope_head_dim) - chunked_k[..., self.qk_nope_head_dim:] = chunked_k_pe.view( - -1, 1, self.qk_rope_head_dim) - chunked_k = chunked_k.view(-1, self.num_heads * self.qk_head_dim) - q, chunked_k, chunked_v = self._maybe_contiguous( - q, chunked_k, chunked_v) + chunked_k_nope = chunked_k_nope.view(-1, self.num_heads, + self.qk_nope_head_dim) + chunked_k_pe = chunked_k_pe.view(-1, 1, self.qk_rope_head_dim) + chunked_k = torch.cat( + (chunked_k_nope, chunked_k_pe.expand(-1, self.num_heads, -1)), + dim=-1) + chunked_k = chunked_k.view(-1, self.num_heads * self.qk_head_dim) # release pytorch activation memory chunked_compressed_kv = None @@ -1317,8 +1288,6 @@ def forward_context_with_chunked_prefill( k[..., self.qk_nope_head_dim:] = k_pe.view(-1, 1, self.qk_rope_head_dim) k = k.view(-1, self.num_heads * self.qk_head_dim) - q, k, v = self._maybe_contiguous(q, k, v) - # copy q_lens to replace kv_lens_runtime attn_metadata.kv_lens_runtime = attn_metadata.prompt_lens_cpu_runtime attn_metadata.kv_lens_cuda_runtime = attn_metadata.prompt_lens_cuda_runtime @@ -1488,7 +1457,6 @@ def forward( attn_metadata: AttentionMetadata, all_reduce_params: Optional[AllReduceParams] = None, ) -> torch.Tensor: - # print(f"=========forward layer_idx: {self.layer_idx}, position_ids: {position_ids.shape}, hidden_states: {hidden_states.shape}") attn_output = self.create_output(hidden_states) if self.register_to_config: torch.ops.trtllm.mla_custom_op_inplace(hidden_states, position_ids, From 4d51e85d6ef52e41b95bb36bdd935e9e9284b319 Mon Sep 17 00:00:00 2001 From: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> Date: Thu, 31 Jul 2025 21:35:19 -0700 Subject: [PATCH 03/14] remove unused codes and fix some problems Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 1 + cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu | 115 +---------- .../kernels/mlaChunkedPrefill.cuh | 9 - cpp/tensorrt_llm/kernels/mlaKernels.cu | 117 ----------- cpp/tensorrt_llm/kernels/mlaKernels.h | 5 - cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp | 183 ------------------ .../kernels/mlaChunkedPrefillTest.cu | 148 -------------- .../unit_tests/kernels/mlaPreprocessTest.cu | 125 +----------- .../_torch/attention_backend/trtllm.py | 80 -------- tensorrt_llm/_torch/modules/attention.py | 20 +- 10 files changed, 16 insertions(+), 787 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 2d077bcf82f9..96ee78a0b283 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -1652,6 +1652,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea if (mIsMLAEnabled) { + TLLM_CHECK_WITH_INFO(params.mla_param != nullptr, "MLA param is nullptr"); params.mla_param->cache_type = cache_type; params.mla_param->cu_q_seqlens = cu_q_seqlens; params.mla_param->quant_scale_kv = params.kv_scale_orig_quant; diff --git a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu index c6ae3a6d4914..215dd5a9c3c0 100644 --- a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu +++ b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu @@ -71,26 +71,6 @@ struct loadChunkedKVKernelTraits static constexpr int kKVThreadPerHead = (kLoraSize * kBytesPerElem) / kBytesPerLoad; }; -template -struct setChunkedKVKernelTraits -{ - using VecT = uint4; - static constexpr int kQKNopeSize = 128; - static constexpr int kVHeadSize = 128; - static_assert(kQKNopeSize == kVHeadSize); - static constexpr int kRopeSize = 64; - static constexpr int kHeadSize = kQKNopeSize + kRopeSize; - static constexpr int kBytesPerElem = sizeof(T); - static constexpr int kBytesPerLoad = 16; - static constexpr int kElemPerLoad = kBytesPerLoad / kBytesPerElem; - static_assert((kHeadSize * kBytesPerElem) % kBytesPerLoad == 0, - "kHeadSize * kBytesPerElem must be multiple of kBytesPerLoad (16Bytes)"); - static constexpr int kThreadPerHead = (kHeadSize * kBytesPerElem) / kBytesPerLoad; - static constexpr int kKVThreadPerHead = (kQKNopeSize * kBytesPerElem) / kBytesPerLoad; - static constexpr int kCpTokenPerBlock = 16; - static constexpr int kBlockSize = kThreadPerHead * kCpTokenPerBlock; -}; - template inline __device__ void quantCopy( __nv_fp8_e4m3* dst_global_ptr, SrcType const* src_fragment_ptr, float const scale_val = 1.f) @@ -311,76 +291,6 @@ __global__ void loadChunkedKVCacheForMLAKernel(T* output_kv_ptr, T* output_k_pe_ } } -// in the most of cases, chunk_size = max_seq_len -// output_kv {B, 2, ceil(max_seq_len / kv_cache_tokens_per_block), h, kv_cache_tokens_per_block, d}, padding with -// zero -// kv {token_size = B*chunked_unit_size, 2, H=128, uncompressed_h=128}, k_pe {token_size = B*chunked_unit_size, h=1, -// rope_h} -// cu_seq_lens {batch + 1}, fake cu_seq_len, for chunked prefill is {0, chunk_size, chunk_size * 2 ....} -template -__global__ void setChunkedKVCacheForMLAKernel(T* output_kv, T const* kv, T const* k_pe, int const max_seq_len, - int const num_heads, int uncompressed_head_size, int rope_size, int64_t const* cu_seq_lens, - int kv_cache_tokens_per_block) -{ - using KT = setChunkedKVKernelTraits; - int const batch_idx = static_cast(blockIdx.y); - int const head_idx = static_cast(blockIdx.z); - int const head_dim_vec_idx = (threadIdx.x % KT::kThreadPerHead); - int const head_dim_idx = head_dim_vec_idx * KT::kElemPerLoad; - bool const is_valid_kv = head_dim_idx < KT::kQKNopeSize; - - int64_t const global_token_offset = cu_seq_lens[batch_idx]; - int64_t const cache_kv_len = cu_seq_lens[batch_idx + 1] - cu_seq_lens[batch_idx]; - int const kv_cache_block_num = (max_seq_len + kv_cache_tokens_per_block - 1) / kv_cache_tokens_per_block; - int const kv_cache_block_size = num_heads * kv_cache_tokens_per_block * (uncompressed_head_size + rope_size); - int64_t const offset_for_kv_in_mem_pool = kv_cache_block_num * kv_cache_block_size; - int64_t const kv_offset = num_heads * uncompressed_head_size; - size_t const seq_len_loop_end = cache_kv_len; - for (int local_token_idx = (threadIdx.x / KT::kThreadPerHead) + blockIdx.x * KT::kCpTokenPerBlock; - local_token_idx < seq_len_loop_end; local_token_idx += gridDim.x * KT::kCpTokenPerBlock) - { - if (local_token_idx >= cache_kv_len) - { - break; - } - if (is_valid_kv) - { - - int64_t ld_kv_global_offset - = int64_t(global_token_offset + local_token_idx) * 2 * num_heads * uncompressed_head_size - + head_idx * uncompressed_head_size; - int64_t ld_kv_local_offset = head_dim_vec_idx; - auto k_data = (reinterpret_cast(kv + ld_kv_global_offset))[ld_kv_local_offset]; - auto v_data = (reinterpret_cast( - kv + kv_offset + ld_kv_global_offset))[ld_kv_local_offset]; - - int64_t st_k_global_offset = int64_t(batch_idx) * 2 * offset_for_kv_in_mem_pool - + local_token_idx / kv_cache_tokens_per_block * kv_cache_block_size - + head_idx * kv_cache_tokens_per_block * (uncompressed_head_size + rope_size) - + (local_token_idx % kv_cache_tokens_per_block) * (uncompressed_head_size + rope_size); - int64_t st_v_global_offset = st_k_global_offset + offset_for_kv_in_mem_pool; - int64_t st_k_local_offset = head_dim_vec_idx; - int64_t st_v_local_offset = head_dim_vec_idx; - (reinterpret_cast(output_kv + st_k_global_offset))[st_k_local_offset] = k_data; - (reinterpret_cast(output_kv + st_v_global_offset))[st_v_local_offset] = v_data; - } - else - { - // rope h = 1 - int64_t ld_rope_global_offset = int64_t(global_token_offset + local_token_idx) * rope_size; - int64_t ld_rope_local_offset = head_dim_vec_idx - KT::kKVThreadPerHead; - auto rope_data - = (reinterpret_cast(k_pe + ld_rope_global_offset))[ld_rope_local_offset]; - int64_t st_rope_global_offset = int64_t(batch_idx) * 2 * offset_for_kv_in_mem_pool - + local_token_idx / kv_cache_tokens_per_block * kv_cache_block_size - + head_idx * kv_cache_tokens_per_block * (uncompressed_head_size + rope_size) - + (local_token_idx % kv_cache_tokens_per_block) * (uncompressed_head_size + rope_size); - int64_t st_rope_local_offset = head_dim_vec_idx; - (reinterpret_cast(output_kv + st_rope_global_offset))[st_rope_local_offset] = rope_data; - } - } -} - } // namespace namespace tensorrt_llm @@ -427,26 +337,6 @@ void invokeMLALoadChunkedKV(T* output_kv_ptr, T* output_k_pe_ptr, KVBlockArray c kv_cache, cu_ctx_chunked_len, chunked_size, chunked_idx, kv_scale_quant_orig_ptr); } -// output_kv {B, 2, ceil(chunked_size / kv_cache_tokens_per_block), h, kv_cache_tokens_per_block, d}, padding with -// zero -// kv {total_token, 2, H, uncompressed_h=128} 0 for k and 1 for v, k_pe {total_token, h=1, rope_h} -// input kv and k_pe can be cached tokens or uncached tokens -template -void invokeMLASetChunkedKV(T* output_kv, T const* kv, T const* k_pe, int const batch_size, int const max_seq_len, - int const num_heads, int uncompressed_head_size, int rope_size, int64_t const* cu_seq_lens, - int const kv_cache_tokens_per_block, cudaStream_t stream) -{ - using KT = setChunkedKVKernelTraits; - TLLM_CHECK_WITH_INFO( - uncompressed_head_size + rope_size == KT::kHeadSize, "head dim should be equal to %d", KT::kHeadSize); - TLLM_CHECK_WITH_INFO(kv_cache_tokens_per_block % KT::kCpTokenPerBlock == 0, - "kv_cache_tokens_per_block should be multiple of %d", KT::kCpTokenPerBlock); - - dim3 grid(tensorrt_llm::common::divUp(max_seq_len, KT::kCpTokenPerBlock), batch_size, num_heads); - setChunkedKVCacheForMLAKernel<<>>(output_kv, kv, k_pe, max_seq_len, num_heads, - uncompressed_head_size, rope_size, cu_seq_lens, kv_cache_tokens_per_block); -} - #define INSTANTIATE_MLA_CHUNKED_PREFILL_KERNEL(T) \ template void invokeMergeAttnWithSoftmax(T * merged_attn, float* merged_softmax_stats, T const* pre_attn, \ float const* pre_softmax_stats, T const* curr_attn, float const* curr_softmax_stats, int const batch_size, \ @@ -457,10 +347,7 @@ void invokeMLASetChunkedKV(T* output_kv, T const* kv, T const* k_pe, int const b int chunked_idx, float const* kv_scale_quant_orig_ptr, cudaStream_t stream); \ template void invokeMLALoadChunkedKV(T * output_kv_ptr, T * output_k_pe_ptr, \ KVBlockArray const& kv_cache, int const num_contexts, int64_t const* cu_ctx_chunked_len, int lora_size, \ - int rope_size, int chunked_size, int chunked_idx, float const* kv_scale_quant_orig_ptr, cudaStream_t stream); \ - template void invokeMLASetChunkedKV(T * output_kv, T const* kv, T const* k_pe, int const batch_size, \ - int const max_seq_len, int const num_heads, int uncompressed_head_size, int rope_size, \ - int64_t const* cu_seq_lens, int const kv_cache_tokens_per_block, cudaStream_t stream); + int rope_size, int chunked_size, int chunked_idx, float const* kv_scale_quant_orig_ptr, cudaStream_t stream); INSTANTIATE_MLA_CHUNKED_PREFILL_KERNEL(half); INSTANTIATE_MLA_CHUNKED_PREFILL_KERNEL(float); diff --git a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh index 765c75b9a91c..3b339ac1a85b 100644 --- a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh +++ b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh @@ -37,14 +37,5 @@ template void invokeMLALoadChunkedKV(T* output_kv_ptr, T* output_k_pe_ptr, KVBlockArray const& kv_cache, int const num_contexts, int64_t const* cu_ctx_chunked_len, int lora_size, int rope_size, int chunked_size, int chunked_idx, float const* kv_scale_quant_orig_ptr, cudaStream_t stream); - -// output_kv {B, 2, ceil(chunked_size / kv_cache_tokens_per_block), h, kv_cache_tokens_per_block, d}, padding with -// zero -// kv {total_token, 2, H, uncompressed_h=128} 0 for k and 1 for v, k_pe {total_token, h=1, rope_h} -// input kv and k_pe can be cached tokens or uncached tokens -template -void invokeMLASetChunkedKV(T* output_kv, T const* kv, T const* k_pe, int const batch_size, int const max_seq_len, - int const num_heads, int uncompressed_head_size, int rope_size, int64_t const* cu_seq_lens, - int const kv_cache_tokens_per_block, cudaStream_t stream); } // namespace kernels } // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index 2ffee12b788f..a34de5756e5f 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -121,27 +121,6 @@ struct loadPagedKVKernelTraits static constexpr int kKVThreadPerHead = (kLoraSize * kBytesPerElem) / kBytesPerLoad; }; -template -struct setPagedKVKernelTraits -{ - static constexpr int kQKNopeSize = 128; - static constexpr int kVHeadSize = 128; - static_assert(kQKNopeSize == kVHeadSize); - static constexpr int kRopeSize = 64; - static constexpr int kHeadSize = kQKNopeSize + kRopeSize; - using VecT = typename VecType::Type; - static constexpr int kBytesPerElem = sizeof(T); - static constexpr int kBytesPerLoad = 16; - static constexpr int kElemPerLoad = kBytesPerLoad / kBytesPerElem; - static_assert((kHeadSize * kBytesPerElem) % kBytesPerLoad == 0, - "kHeadSize * kBytesPerElem must be multiple of kBytesPerLoad (16Bytes)"); - static constexpr int kNumHeads = 128; - static constexpr int kThreadPerHead = (kHeadSize * kBytesPerElem) / kBytesPerLoad; - static constexpr int kKVThreadPerHead = (kQKNopeSize * kBytesPerElem) / kBytesPerLoad; - static constexpr int kCpTokenPerBlock = 16; - static constexpr int kBlockSize = kThreadPerHead * kCpTokenPerBlock; -}; - template inline __device__ void quantCopy( __nv_fp8_e4m3* dst_global_ptr, SrcType const* src_fragment_ptr, float const scale_val = 1.f) @@ -709,79 +688,6 @@ __global__ void loadPagedKVCacheForMLAKernel(T* compressed_kv_ptr, T* k_pe_ptr, } } -// k {total_token, h, d}, v {total_token, h, d}, k_pe {total_token, h=1, d_rope} -// output {b, 2, ceil(max_seq / kv_cache_tokens_per_block), h, kv_cache_tokens_per_block, d} -template -__global__ void setPagedKVCacheForMLAKernel(T* output, T const* k_ptr, T const* v_ptr, T const* k_pe_ptr, - int64_t const* cu_seq_lens, int const max_input_seq_len, int num_heads, int kv_dim, int rope_dim, - int kv_cache_tokens_per_block, int64_t kv_token_stride) -{ - using KT = typename tensorrt_llm::kernels::setPagedKVKernelTraits; - int const batch_idx = static_cast(blockIdx.y); - int const head_idx = static_cast(blockIdx.z); - int const head_dim_vec_idx = (threadIdx.x % KT::kThreadPerHead); - int const head_dim_idx = head_dim_vec_idx * KT::kElemPerLoad; - bool const is_valid_v = head_dim_idx < KT::kVHeadSize; - - size_t const seq_len_loop_end - = (max_input_seq_len + KT::kCpTokenPerBlock - 1) / KT::kCpTokenPerBlock * KT::kCpTokenPerBlock; - size_t const kv_cache_block_size = num_heads * kv_cache_tokens_per_block * (kv_dim + rope_dim); - size_t const kv_cache_block_num = (max_input_seq_len + kv_cache_tokens_per_block - 1) / kv_cache_tokens_per_block; - int64_t const global_token_offset = cu_seq_lens[batch_idx]; - int64_t const cache_kv_len = cu_seq_lens[batch_idx + 1] - cu_seq_lens[batch_idx]; - - for (int local_token_idx = (threadIdx.x / KT::kThreadPerHead) + blockIdx.x * KT::kCpTokenPerBlock; - local_token_idx < seq_len_loop_end; local_token_idx += KT::kCpTokenPerBlock * gridDim.x) - { - int token_idx_in_kv_cache = local_token_idx; - bool const valid_token = token_idx_in_kv_cache < cache_kv_len; - if (valid_token) - { - // copy k and v - if (is_valid_v) - { - int ld_kv_global_offset = (global_token_offset + local_token_idx) * kv_token_stride + head_idx * kv_dim; - int ld_kv_local_offset = head_dim_vec_idx; - auto k_data - = (reinterpret_cast(k_ptr + ld_kv_global_offset))[ld_kv_local_offset]; - auto v_data - = (reinterpret_cast(v_ptr + ld_kv_global_offset))[ld_kv_local_offset]; - // {b, 0, token / kv_cache_tokens_per_block, h, token % kv_cache_tokens_per_block, ...} - int st_k_global_offset = batch_idx * 2 * kv_cache_block_num * kv_cache_block_size - + local_token_idx / kv_cache_tokens_per_block * kv_cache_block_size - + head_idx * kv_cache_tokens_per_block * (kv_dim + rope_dim) - + (local_token_idx % kv_cache_tokens_per_block) * (kv_dim + rope_dim); - // {b, 1, token / kv_cache_tokens_per_block, h, token % kv_cache_tokens_per_block, ...} - int st_v_global_offset = st_k_global_offset + kv_cache_block_num * kv_cache_block_size; - int st_k_local_offset = head_dim_vec_idx; - int st_v_local_offset = head_dim_vec_idx; - (reinterpret_cast(output + st_k_global_offset))[st_k_local_offset] = k_data; - (reinterpret_cast(output + st_v_global_offset))[st_v_local_offset] = v_data; - } - // copy k_pe, only 1 head - else - { - int ld_rope_global_offset = (global_token_offset + local_token_idx) * rope_dim; - int ld_rope_local_offset = head_dim_vec_idx - KT::kKVThreadPerHead; - auto rope_data = (reinterpret_cast( - k_pe_ptr + ld_rope_global_offset))[ld_rope_local_offset]; - // {b, 0, token / kv_cache_tokens_per_block, h, token % kv_cache_tokens_per_block, ...} - int st_rope_global_offset = batch_idx * 2 * kv_cache_block_num * kv_cache_block_size - + local_token_idx / kv_cache_tokens_per_block * kv_cache_block_size - + head_idx * kv_cache_tokens_per_block * (kv_dim + rope_dim) - + (local_token_idx % kv_cache_tokens_per_block) * (kv_dim + rope_dim); - int st_rope_local_offset = head_dim_vec_idx; - (reinterpret_cast(output + st_rope_global_offset))[st_rope_local_offset] - = rope_data; - } - } - else - { - break; - } - } -} - // q {total_uncached_tokens, h, d_nope + d_rope} // latent_cache {total_uncached_tokens, d_k + d_rope} template @@ -1109,20 +1015,6 @@ void invokeMLALoadPagedKV(T* compressed_kv_ptr, T* k_pe_ptr, KVBlockArray& kv_ca compressed_kv_ptr, k_pe_ptr, kv_cache, cu_ctx_cached_kv_lens, max_input_seq_len, kv_scale_quant_orig_ptr); } -template -void invokeMLASetPagedKV(T* output, T const* k_ptr, T const* v_ptr, T const* k_pe_ptr, int const num_requests, - int64_t const* cu_seq_lens, int const max_input_seq_len, int num_heads, int kv_dim, int rope_dim, - int kv_cache_tokens_per_block, int64_t kv_token_stride, cudaStream_t stream) -{ - using KT = typename tensorrt_llm::kernels::setPagedKVKernelTraits; - TLLM_CHECK_WITH_INFO(kv_dim + rope_dim == KT::kHeadSize, "head dim should be equal to %d", KT::kHeadSize); - TLLM_CHECK_WITH_INFO(kv_cache_tokens_per_block % KT::kCpTokenPerBlock == 0, - "kv_cache_tokens_per_block should be multiple of %d", KT::kCpTokenPerBlock); - dim3 grid(tensorrt_llm::common::divUp(max_input_seq_len, KT::kCpTokenPerBlock), num_requests, num_heads); - setPagedKVCacheForMLAKernel<<>>(output, k_ptr, v_ptr, k_pe_ptr, cu_seq_lens, - max_input_seq_len, num_heads, kv_dim, rope_dim, kv_cache_tokens_per_block, kv_token_stride); -} - template void invokeMLARopeAppendPagedKVAssignQ(KVBlockArray& kv_cache, T* q_ptr, T* latent_cache_ptr, int const num_requests, int64_t const* cu_ctx_cached_kv_lens, int64_t const* cu_seq_lens, int const max_input_uncached_seq_len, @@ -1168,15 +1060,6 @@ INSTANTIATE_RW_KVCACHE_MLA(half, __nv_fp8_e4m3); INSTANTIATE_RW_KVCACHE_MLA(__nv_bfloat16, __nv_bfloat16); INSTANTIATE_RW_KVCACHE_MLA(__nv_bfloat16, __nv_fp8_e4m3); -#define INSTANTIATE_SET_KVCACHE_MLA(T) \ - template void invokeMLASetPagedKV(T* output, T const* k_ptr, T const* v_ptr, T const* k_pe_ptr, \ - int const num_requests, int64_t const* cu_seq_lens, int const max_input_seq_len, int num_heads, int kv_dim, \ - int rope_dim, int kv_cache_tokens_per_block, int64_t kv_token_stride, cudaStream_t stream); - -INSTANTIATE_SET_KVCACHE_MLA(float); -INSTANTIATE_SET_KVCACHE_MLA(half); -INSTANTIATE_SET_KVCACHE_MLA(__nv_bfloat16); - } // namespace kernels } // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h index 709548b1e9af..482dc0392241 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -105,11 +105,6 @@ void invokeMLALoadPagedKV(T* compressed_kv_ptr, T* k_pe_ptr, KVBlockArray& kv_ca int64_t const* cu_ctx_cached_kv_lens, int const max_input_seq_len, int const lora_size, int const rope_size, float const* kv_scale_quant_orig_ptr, cudaStream_t stream); -template -void invokeMLASetPagedKV(T* output, T const* k_ptr, T const* v_ptr, T const* k_pe_ptr, int const num_requests, - int64_t const* cu_seq_lens, int const max_input_seq_len, int num_heads, int kv_dim, int rope_dim, - int kv_cache_tokens_per_block, int64_t kv_token_stride, cudaStream_t stream); - template void invokeMLARopeAppendPagedKVAssignQ(KVBlockArray& kv_cache, T* q_ptr, T* latent_cache_ptr, int const num_requests, int64_t const* cu_ctx_cached_kv_lens, int64_t const* cu_seq_lens, int const max_input_uncached_seq_len, diff --git a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp index 2bac01442438..3e052e3283a7 100644 --- a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp +++ b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp @@ -62,39 +62,6 @@ void loadChunkedKVCacheForMLAHelper(torch::Tensor& output_kv, torch::Tensor& out kv_scale_quant_orig_ptr, stream); } -template -void setPagedKVCacheForMLAHelper(torch::Tensor& output, torch::Tensor const& k, torch::Tensor const& v, - torch::Tensor const& k_pe, int const num_requests, torch::Tensor const& cu_seq_lens, int const max_input_seq_len, - int num_heads, int kv_dim, int rope_dim, int kv_cache_tokens_per_block, int64_t kv_token_stride) -{ - auto stream = at::cuda::getCurrentCUDAStream(output.get_device()); - auto* output_ptr = static_cast(output.data_ptr()); - auto const* k_ptr = static_cast(k.data_ptr()); - auto const* v_ptr = static_cast(v.data_ptr()); - auto const* k_pe_ptr = static_cast(k_pe.data_ptr()); - auto const* cu_seq_lens_ptr = cu_seq_lens.data_ptr(); - - // cudaMemset is faster than torch::zeros - TLLM_CUDA_CHECK(cudaMemsetAsync(output_ptr, 0, output.numel() * torch::elementSize(output.scalar_type()), stream)); - tensorrt_llm::kernels::invokeMLASetPagedKV(output_ptr, k_ptr, v_ptr, k_pe_ptr, num_requests, cu_seq_lens_ptr, - max_input_seq_len, num_heads, kv_dim, rope_dim, kv_cache_tokens_per_block, kv_token_stride, stream); -} - -template -void setChunkedKVCacheForMLAHelper(torch::Tensor& output, torch::Tensor const& kv, torch::Tensor const& k_pe, - int const num_requests, torch::Tensor const& cu_seq_lens, int num_heads, int kv_dim, int rope_dim, - int kv_cache_tokens_per_block, int max_seq_len) -{ - auto stream = at::cuda::getCurrentCUDAStream(output.get_device()); - T* output_ptr = static_cast(output.data_ptr()); - T* kv_ptr = static_cast(kv.data_ptr()); - T* k_pe_ptr = static_cast(k_pe.data_ptr()); - auto* cu_seq_lens_ptr = cu_seq_lens.data_ptr(); - - tensorrt_llm::kernels::invokeMLASetChunkedKV(output_ptr, kv_ptr, k_pe_ptr, num_requests, max_seq_len, num_heads, - kv_dim, rope_dim, cu_seq_lens_ptr, kv_cache_tokens_per_block, stream); -} - template void invokeMLARopeAppendPagedKVAssignQHelper(KVBlockArray& kv_cache, torch::Tensor& q, torch::Tensor& latent_cache, int const num_requests, torch::Tensor const& cu_ctx_cached_kv_lens, torch::Tensor const& cu_seq_lens, @@ -371,111 +338,6 @@ std::vector loadChunkedKVCacheForMLA(torch::ScalarType out_dtype, return outputs; } -torch::Tensor setPagedKVCacheForMLA(torch::Tensor& output, torch::Tensor const& k, torch::Tensor const& v, - torch::Tensor const& k_pe, int64_t const num_requests, torch::Tensor const& cu_seq_lens, - int64_t const max_input_seq_len, int64_t const num_heads, int64_t const kv_dim, int64_t const rope_dim, - int64_t const kv_cache_tokens_per_block) -{ - TORCH_CHECK(output.numel() > 0); - auto output_dtype = output.scalar_type(); - TORCH_CHECK(output_dtype == torch::kFloat16 || output_dtype == torch::kFloat32 || output_dtype == torch::kBFloat16); - CHECK_TH_CUDA(output); - CHECK_CONTIGUOUS(output); - - // k and v can be non-contiguous - CHECK_TH_CUDA(k); - CHECK_TYPE(k, output_dtype); - CHECK_TH_CUDA(v); - CHECK_TYPE(v, output_dtype); - TORCH_CHECK(k.dim() == 3); - TORCH_CHECK(v.dim() == 3); - TORCH_CHECK(k.size(0) == v.size(0)); - TORCH_CHECK(k.size(1) == v.size(1)); - TORCH_CHECK(k.size(2) == v.size(2)); - TORCH_CHECK(k.stride(1) == k.size(2)); - TORCH_CHECK(v.stride(1) == v.size(2)); - TORCH_CHECK(k.stride(2) == 1); - TORCH_CHECK(v.stride(2) == 1); - // k and v should have the same token stride - int64_t k_token_stride = k.stride(0); - int64_t v_token_stride = v.stride(0); - TORCH_CHECK(k_token_stride == v_token_stride); - - // k_pe should be contiguous - CHECK_INPUT(k_pe, output_dtype); - - CHECK_INPUT(cu_seq_lens, torch::kInt64); - TORCH_CHECK(cu_seq_lens.dim() == 1); - TORCH_CHECK(cu_seq_lens.size(0) >= num_requests + 1); - - if (output_dtype == torch::kFloat16) - { - setPagedKVCacheForMLAHelper(output, k, v, k_pe, num_requests, cu_seq_lens, max_input_seq_len, num_heads, - kv_dim, rope_dim, kv_cache_tokens_per_block, k_token_stride); - } - else if (output_dtype == torch::kFloat32) - { - setPagedKVCacheForMLAHelper(output, k, v, k_pe, num_requests, cu_seq_lens, max_input_seq_len, num_heads, - kv_dim, rope_dim, kv_cache_tokens_per_block, k_token_stride); - } - else if (output_dtype == torch::kBFloat16) - { - setPagedKVCacheForMLAHelper<__nv_bfloat16>(output, k, v, k_pe, num_requests, cu_seq_lens, max_input_seq_len, - num_heads, kv_dim, rope_dim, kv_cache_tokens_per_block, k_token_stride); - } - - int64_t max_block_num = (max_input_seq_len + kv_cache_tokens_per_block - 1) / kv_cache_tokens_per_block; - - torch::Tensor faked_kv_cache_block_offsets = torch::arange( - 0, num_requests * 2 * max_block_num, torch::TensorOptions().dtype(torch::kInt32).device(output.device())); - - faked_kv_cache_block_offsets = faked_kv_cache_block_offsets.view({num_requests, 2, max_block_num}); - - return faked_kv_cache_block_offsets; -} - -torch::Tensor setChunkedKVCacheForMLA(torch::Tensor& output, torch::Tensor const& kv, torch::Tensor const& k_pe, - int64_t const num_requests, torch::Tensor const& cu_seq_lens, int64_t const num_heads, int64_t const kv_dim, - int64_t const rope_dim, int64_t const kv_cache_tokens_per_block, int64_t const max_seq_len) -{ - TORCH_CHECK(output.numel() > 0); - TORCH_CHECK(output.scalar_type() == torch::kFloat16 || output.scalar_type() == torch::kFloat32 - || output.scalar_type() == torch::kBFloat16); - CHECK_TH_CUDA(output); - CHECK_CONTIGUOUS(output); - CHECK_INPUT(kv, output.scalar_type()); - CHECK_INPUT(k_pe, output.scalar_type()); - CHECK_INPUT(cu_seq_lens, torch::kInt64); - TORCH_CHECK(cu_seq_lens.dim() == 1); - TORCH_CHECK(cu_seq_lens.size(0) >= num_requests + 1); - - if (output.scalar_type() == torch::kFloat16) - { - setChunkedKVCacheForMLAHelper(output, kv, k_pe, num_requests, cu_seq_lens, num_heads, kv_dim, rope_dim, - kv_cache_tokens_per_block, max_seq_len); - } - else if (output.scalar_type() == torch::kFloat32) - { - setChunkedKVCacheForMLAHelper(output, kv, k_pe, num_requests, cu_seq_lens, num_heads, kv_dim, rope_dim, - kv_cache_tokens_per_block, max_seq_len); - } - else if (output.scalar_type() == torch::kBFloat16) - { - setChunkedKVCacheForMLAHelper<__nv_bfloat16>(output, kv, k_pe, num_requests, cu_seq_lens, num_heads, kv_dim, - rope_dim, kv_cache_tokens_per_block, max_seq_len); - } - - int64_t max_block_num = (max_seq_len + kv_cache_tokens_per_block - 1) / kv_cache_tokens_per_block; - - // TODO: actually this offset is always the same for all requests and all layers. - torch::Tensor faked_kv_cache_block_offsets = torch::arange( - 0, num_requests * 2 * max_block_num, torch::TensorOptions().dtype(torch::kInt32).device(output.device())); - - faked_kv_cache_block_offsets = faked_kv_cache_block_offsets.view({num_requests, 2, max_block_num}); - - return faked_kv_cache_block_offsets; -} - void MLARopeAppendPagedKVAssignQ(torch::Tensor& q, torch::Tensor& latent_cache, int64_t const num_contexts, torch::Tensor const& cu_ctx_cached_kv_lens, torch::Tensor const& cu_seq_lens, int64_t const max_input_uncached_seq_len, torch::Tensor const& cos_sin_cache, int64_t const head_num, @@ -664,51 +526,6 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) m.impl("load_chunked_kv_cache_for_mla", &torch_ext::loadChunkedKVCacheForMLA); } -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "set_paged_kv_cache_for_mla(" - "Tensor output" - ", Tensor k" - ", Tensor v" - ", Tensor k_pe" - ", int num_requests" - ", Tensor cu_seq_lens" - ", int max_input_seq_len" - ", int num_heads" - ", int kv_dim" - ", int rope_dim" - ", int kv_cache_tokens_per_block" - ") -> Tensor"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("set_paged_kv_cache_for_mla", &torch_ext::setPagedKVCacheForMLA); -} - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "set_chunked_kv_cache_for_mla(" - "Tensor output" - ", Tensor kv" - ", Tensor k_pe" - ", int num_requests" - ", Tensor cu_seq_lens" - ", int num_heads" - ", int kv_dim" - ", int rope_dim" - ", int kv_cache_tokens_per_block" - ", int max_seq_len" - ") -> Tensor"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("set_chunked_kv_cache_for_mla", &torch_ext::setChunkedKVCacheForMLA); -} - TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( diff --git a/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu b/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu index 18e183853371..8f4bd25a9a52 100644 --- a/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu +++ b/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu @@ -63,59 +63,6 @@ void loadChunkedKVKernelRef(T* kv_output, T* k_pe_output, tensorrt_llm::kernels: } } -// kv {total_tokens, 2, h, nope_size} -// k_pe {total_tokens, h=1, rope_size} -// output {b, 2, ceil(max_seq / cache_tokens_per_block), h, cache_tokens_per_block, (nope_size + rope_size)} -// max_seq <= chunk_size -template -void setChunkedKVCacheForMLAKernelRef(T* output, T* kv_ptr, T* k_pe_ptr, int num_contexts, int64_t const* cu_seq_len, - int const max_input_seq_len, int num_heads, int nope_size, int rope_size, int cache_tokens_per_block) -{ - int head_size = nope_size + rope_size; - int const kv_cache_size_per_block = num_heads * cache_tokens_per_block * head_size; - int const kv_cache_block_num_per_seq = (max_input_seq_len + cache_tokens_per_block - 1) / cache_tokens_per_block; - for (int b = 0; b < num_contexts; b++) - { - int const global_token_offset = cu_seq_len[b]; - int const current_seq_len = cu_seq_len[b + 1] - cu_seq_len[b]; - for (int s = 0; s < current_seq_len; s++) - { - int const global_token_idx = global_token_offset + s; - int const kv_cache_block_offset_for_k - = (b * 2 * kv_cache_block_num_per_seq + s / cache_tokens_per_block) * kv_cache_size_per_block; - int const kv_cache_block_offset_for_v - = kv_cache_block_offset_for_k + (kv_cache_block_num_per_seq * kv_cache_size_per_block); - for (int h = 0; h < num_heads; h++) - { - int const ld_k_head_offset = (global_token_idx * 2 * num_heads * nope_size) + h * nope_size; - int const ld_v_head_offset = ld_k_head_offset + num_heads * nope_size; - int const ld_k_pe_head_offset = global_token_idx * rope_size; - // copy kv - for (int d = 0; d < nope_size; d++) - { - int const ld_k_idx = ld_k_head_offset + d; - int const ld_v_idx = ld_v_head_offset + d; - int const st_k_idx = kv_cache_block_offset_for_k + h * cache_tokens_per_block * head_size - + (s % cache_tokens_per_block) * head_size + d; - int const st_v_idx = kv_cache_block_offset_for_v + h * cache_tokens_per_block * head_size - + (s % cache_tokens_per_block) * head_size + d; - output[st_k_idx] = kv_ptr[ld_k_idx]; - output[st_v_idx] = kv_ptr[ld_v_idx]; - } - - // copy k_pe - for (int d = 0; d < rope_size; d++) - { - int const ld_k_pe_idx = ld_k_pe_head_offset + d; - int const st_k_pe_idx = kv_cache_block_offset_for_k + h * cache_tokens_per_block * head_size - + (s % cache_tokens_per_block) * head_size + (nope_size + d); - output[st_k_pe_idx] = k_pe_ptr[ld_k_pe_idx]; - } - } - } - } -} - // Q {total_q, H, D} // KV {total_kv, 2, H, D} // softmax_sum {total_q, H, 2} // {max/sum} @@ -322,9 +269,6 @@ protected: d_k_pe_output{nullptr}, h_compressed_kv_output_ref{nullptr}, h_k_pe_output_ref{nullptr}, h_kv_scale_quant_orig{nullptr}, d_kv_scale_quant_orig{nullptr}, - // for kernel 2 - h_kv_tensor{nullptr}, d_kv_tensor{nullptr}, h_k_pe_tensor{nullptr}, d_k_pe_tensor{nullptr}, - // for merge attn {kv_full_tensor = kv + k_pe} m_h_q_tensor{nullptr}, m_h_kv_full_tensor{nullptr}, m_h_chunked_kv_tensor{nullptr}, m_h_output_tensor{nullptr}, m_h_softmax_sum_tensor{nullptr}, m_h_softmax_sum_accum_tensor{nullptr}, m_h_output_tensor_ref{nullptr}, @@ -621,27 +565,6 @@ protected: cudaMemcpyHostToDevice); } - // kv, k_pe for invokeMLASetChunkedKV (kernel 2) - this->h_kv_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mBatchSize * this->mChunkSize, 2, this->mNumHeads, this->mNopeSize}), dtype); - this->h_k_pe_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mBatchSize * this->mChunkSize, 1, this->mRopeSize}), dtype); - this->d_kv_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->h_kv_tensor->getShape(), dtype); - this->d_k_pe_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->h_k_pe_tensor->getShape(), dtype); - { - auto* kv_ptr = bufferCast(*(this->h_kv_tensor)); - auto* k_pe_ptr = bufferCast(*(this->h_k_pe_tensor)); - - fillArrayDataWithMod(kv_ptr, h_kv_tensor->getSize()); - fillArrayDataWithMod(k_pe_ptr, h_k_pe_tensor->getSize()); - - cudaMemcpyAsync(d_kv_tensor->data(), h_kv_tensor->data(), h_kv_tensor->getSizeInBytes(), - cudaMemcpyHostToDevice, mStream->get()); - cudaMemcpyAsync(d_k_pe_tensor->data(), h_k_pe_tensor->data(), h_k_pe_tensor->getSizeInBytes(), - cudaMemcpyHostToDevice, mStream->get()); - cudaStreamSynchronize(mStream->get()); - } - // invokeMergeAttnWithSoftmax, we just ignore rope_size here for simplicity this->m_h_q_tensor = tensorrt_llm::runtime::BufferManager::pinned( @@ -915,39 +838,6 @@ protected: cudaMemcpyDeviceToHost); sync_check_cuda_error(this->mStream->get()); } - - void PerformSetChunkedKVRef() - { - using tensorrt_llm::runtime::bufferCast; - auto* kv_ptr = bufferCast(*(this->h_kv_tensor)); - auto* k_pe_ptr = bufferCast(*(this->h_k_pe_tensor)); - auto* kv_cache_ptr = bufferCast(*(this->h_kv_cache_tensor_ref)); - auto* cu_chunked_seq_lens_ptr = bufferCast(*(this->h_cu_chunk_lens)); - this->PrepareChunkedLen(0); - setChunkedKVCacheForMLAKernelRef(kv_cache_ptr, kv_ptr, k_pe_ptr, this->mBatchSize, cu_chunked_seq_lens_ptr, - this->mChunkSize, this->mNumHeads, this->mNopeSize, this->mRopeSize, this->mTokensPerBlock); - } - - void PerformSetChunkedKV() - { - using tensorrt_llm::runtime::bufferCast; - auto* kv_ptr = bufferCast(*(this->d_kv_tensor)); - auto* k_pe_ptr = bufferCast(*(this->d_k_pe_tensor)); - auto* kv_cache_ptr = bufferCast(*(this->d_kv_cache_tensor)); - auto* cu_chunked_seq_lens_ptr = bufferCast(*(this->d_cu_chunk_lens)); - this->PrepareChunkedLen(0); - // copy cu chunk lens to device - cudaMemcpy(this->d_cu_chunk_lens->data(), this->h_cu_chunk_lens->data(), - this->h_cu_chunk_lens->getSizeInBytes(), cudaMemcpyHostToDevice); - tensorrt_llm::kernels::invokeMLASetChunkedKV(kv_cache_ptr, kv_ptr, k_pe_ptr, this->mBatchSize, this->mChunkSize, - this->mNumHeads, this->mNopeSize, this->mRopeSize, cu_chunked_seq_lens_ptr, this->mTokensPerBlock, - mStream->get()); - cudaStreamSynchronize(this->mStream->get()); - // copy result back to host - cudaMemcpy(this->h_kv_cache_tensor->data(), kv_cache_ptr, this->h_kv_cache_tensor->getSizeInBytes(), - cudaMemcpyDeviceToHost); - sync_check_cuda_error(this->mStream->get()); - } }; using MLATypes @@ -1086,41 +976,3 @@ TYPED_TEST(MlaChunkedPrefillTest, MlaChunkedLoad) } ASSERT_TRUE(allEqual); } - -TYPED_TEST(MlaChunkedPrefillTest, MlaChunkedSet) -{ - using tensorrt_llm::runtime::bufferCast; - using DataType = typename TestFixture::DataType; - using TCache = typename TestFixture::TCache; - if constexpr (std::is_same_v) - { - this->setDefaultParams(); - this->allocateBuffers(); - - sync_check_cuda_error(this->mStream->get()); - bool allEqual{true}; - - this->PerformSetChunkedKVRef(); - sync_check_cuda_error(this->mStream->get()); - this->PerformSetChunkedKV(); - sync_check_cuda_error(this->mStream->get()); - - // check result - auto* kv_cache_ptr = bufferCast(*(this->h_kv_cache_tensor)); - auto* kv_cache_ptr_ref = bufferCast(*(this->h_kv_cache_tensor_ref)); - - for (int i = 0; i < this->h_kv_cache_tensor->getSize(); i++) - { - if (std::abs(static_cast(kv_cache_ptr[i]) - static_cast(kv_cache_ptr_ref[i])) - > getTolerance(kv_cache_ptr[i])) - { - std::cout << "KV cache mismatch at index " << i << ": " - << "expected " << static_cast(kv_cache_ptr_ref[i]) << ", got " - << static_cast(kv_cache_ptr[i]) << std::endl; - allEqual = false; - break; - } - } - ASSERT_TRUE(allEqual); - } -} diff --git a/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu b/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu index df8c79199401..050f3a105e22 100644 --- a/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu +++ b/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu @@ -79,60 +79,6 @@ void loadPagedKvKernelRef(T* compressed_kv_output, T* k_pe_output, } } -// k {total_token, h, uncompressed_h=128}, v {total_token, h, uncompressed_h}, k_pe {total_token, h=1, rope_h} -// output {b, 2, ceil(max_seq / kv_cache_tokens_per_block), h, kv_cache_tokens_per_block, (uncompressed_h + rope_h)} -// copy k, v, k_pe to a continuous memory space (then it will be packed to kv_cache) -template -void setPagedKvCacheForMLAKernelRef(T* output, T* const k_ptr, T* const v_ptr, T* const k_pe_ptr, int num_requests, - int64_t const* cu_seq_lens, int const max_input_seq_len, int num_heads, int uncompressed_head_size, int rope_size, - int kv_cache_tokens_per_block, int64_t kv_token_stride) -{ - int const kv_cache_size_per_block = num_heads * kv_cache_tokens_per_block * (uncompressed_head_size + rope_size); - int const kv_cache_block_num_per_seq - = (max_input_seq_len + kv_cache_tokens_per_block - 1) / kv_cache_tokens_per_block; - for (int b = 0; b < num_requests; b++) - { - int const global_token_offset = cu_seq_lens[b]; - int const current_token_len = cu_seq_lens[b + 1] - cu_seq_lens[b]; - for (int s = 0; s < current_token_len; s++) - { - int const global_token_idx = global_token_offset + s; - int const kv_cache_block_offset_for_k - = ((b * 2 * kv_cache_block_num_per_seq) + (s / kv_cache_tokens_per_block)) * kv_cache_size_per_block; - int const kv_cache_block_offset_for_v - = kv_cache_block_offset_for_k + (kv_cache_block_num_per_seq * kv_cache_size_per_block); - for (int h = 0; h < num_heads; h++) - { - // copy k, v - int const ld_kv_head_offset = (global_token_idx * kv_token_stride) + (h * uncompressed_head_size); - int const ld_k_pe_head_offset = (global_token_idx * rope_size); - for (int d = 0; d < uncompressed_head_size; d++) - { - int const ld_kv_idx = ld_kv_head_offset + d; - int const st_k_idx = kv_cache_block_offset_for_k - + h * kv_cache_tokens_per_block * (uncompressed_head_size + rope_size) - + (s % kv_cache_tokens_per_block) * (uncompressed_head_size + rope_size) + d; - int const st_v_idx = kv_cache_block_offset_for_v - + h * kv_cache_tokens_per_block * (uncompressed_head_size + rope_size) - + (s % kv_cache_tokens_per_block) * (uncompressed_head_size + rope_size) + d; - output[st_k_idx] = k_ptr[ld_kv_idx]; - output[st_v_idx] = v_ptr[ld_kv_idx]; - } - // copy k_pe, head_num = 1 - for (int d = 0; d < rope_size; d++) - { - int const ld_k_pe_idx = ld_k_pe_head_offset + d; - int const st_k_pe_idx = kv_cache_block_offset_for_k - + h * kv_cache_tokens_per_block * (uncompressed_head_size + rope_size) - + (s % kv_cache_tokens_per_block) * (uncompressed_head_size + rope_size) + d - + uncompressed_head_size; - output[st_k_pe_idx] = k_pe_ptr[ld_k_pe_idx]; - } - } - } - } -} - inline bool almostEqual(float a, float b, float atol = 1e-2, float rtol = 1e-3) { if (isnan(a) || isnan(b)) @@ -169,11 +115,8 @@ protected: // for kernel 1 d_compressed_kv_output{nullptr}, h_compressed_kv_output{nullptr}, h_compressed_kv_output_ref{nullptr}, d_k_pe_output{nullptr}, h_k_pe_output{nullptr}, h_k_pe_output_ref{nullptr}, - // for kernel 2 - d_k_tensor{nullptr}, d_v_tensor{nullptr}, d_k_pe_tensor{nullptr}, h_k_tensor{nullptr}, h_v_tensor{nullptr}, - h_k_pe_tensor{nullptr}; - int mNumRequests{}; + int mNumRequests{}; int mMaxSeqLen{}; int mMaxCachedSeqLen{}; int mMaxUncachedSeqLen{}; @@ -463,33 +406,6 @@ protected: cudaMemcpy(this->d_k_pe_output->data(), this->h_k_pe_output->data(), this->h_k_pe_output->getSizeInBytes(), cudaMemcpyHostToDevice); } - // k, v, k_pe for setPagedKvCacheForMLAKernel (kernel 2) - this->h_k_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mTotalTokens, this->mNumHeadsUncompressed, this->mUncompressedHeadSize}), dtype); - this->h_v_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mTotalTokens, this->mNumHeadsUncompressed, this->mUncompressedHeadSize}), dtype); - this->h_k_pe_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mTotalTokens, this->mNumHeadsCompressed, this->mRopeSize}), dtype); - this->d_k_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - ITensor::makeShape({this->mTotalTokens, this->mNumHeadsUncompressed, this->mUncompressedHeadSize}), dtype); - this->d_v_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - ITensor::makeShape({this->mTotalTokens, this->mNumHeadsUncompressed, this->mUncompressedHeadSize}), dtype); - this->d_k_pe_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - ITensor::makeShape({this->mTotalTokens, this->mNumHeadsCompressed, this->mRopeSize}), dtype); - { - auto* k_ptr = bufferCast(*(this->h_k_tensor)); - auto* v_ptr = bufferCast(*(this->h_v_tensor)); - auto* k_pe_ptr = bufferCast(*(this->h_k_pe_tensor)); - fillArrayDataWithMod(k_ptr, this->h_k_tensor->getSize()); - fillArrayDataWithMod(v_ptr, this->h_v_tensor->getSize()); - fillArrayDataWithMod(k_pe_ptr, this->h_k_pe_tensor->getSize()); - cudaMemcpy(this->d_k_tensor->data(), this->h_k_tensor->data(), this->h_k_tensor->getSizeInBytes(), - cudaMemcpyHostToDevice); - cudaMemcpy(this->d_v_tensor->data(), this->h_v_tensor->data(), this->h_v_tensor->getSizeInBytes(), - cudaMemcpyHostToDevice); - cudaMemcpy(this->d_k_pe_tensor->data(), this->h_k_pe_tensor->data(), this->h_k_pe_tensor->getSizeInBytes(), - cudaMemcpyHostToDevice); - } return true; } @@ -539,35 +455,6 @@ protected: cu_ctx_cached_kv_lens_ptr, this->mLoraSize, this->mRopeSize, kv_scale_quant_orig_ptr); } - void PerformSetPagedKV() - { - using tensorrt_llm::runtime::bufferCast; - auto* k_ptr = bufferCast(*(this->d_k_tensor)); - auto* v_ptr = bufferCast(*(this->d_v_tensor)); - auto* k_pe_ptr = bufferCast(*(this->d_k_pe_tensor)); - auto* kv_cache_ptr = bufferCast(*(this->d_kv_cache_tensor)); - auto* cu_seq_lens_ptr = bufferCast(*(this->d_cu_seq_lens)); - tensorrt_llm::kernels::invokeMLASetPagedKV(kv_cache_ptr, k_ptr, v_ptr, k_pe_ptr, this->mNumRequests, - cu_seq_lens_ptr, this->mMaxSeqLen, this->mNumHeadsUncompressed, this->mUncompressedHeadSize, - this->mRopeSize, this->mTokensPerBlock, this->mKvTokenStride, this->mStream->get()); - cudaStreamSynchronize(this->mStream->get()); - cudaMemcpy(this->h_kv_cache_tensor->data(), this->d_kv_cache_tensor->data(), - this->d_kv_cache_tensor->getSizeInBytes(), cudaMemcpyDeviceToHost); - } - - void PerformSetPagedKVRef() - { - using tensorrt_llm::runtime::bufferCast; - auto* k_ptr = bufferCast(*(this->h_k_tensor)); - auto* v_ptr = bufferCast(*(this->h_v_tensor)); - auto* k_pe_ptr = bufferCast(*(this->h_k_pe_tensor)); - auto* kv_cache_ptr = bufferCast(*(this->h_kv_cache_tensor_ref)); - auto* cu_seq_lens_ptr = bufferCast(*(this->h_cu_seq_lens)); - setPagedKvCacheForMLAKernelRef(kv_cache_ptr, k_ptr, v_ptr, k_pe_ptr, this->mNumRequests, cu_seq_lens_ptr, - this->mMaxSeqLen, this->mNumHeadsUncompressed, this->mUncompressedHeadSize, this->mRopeSize, - this->mTokensPerBlock, this->mKvTokenStride); - } - template bool CheckEqual(T const* expected, T const* output, size_t size) { @@ -617,14 +504,4 @@ TYPED_TEST(MlaPreprocessTest, MLAPreprocessDefault) allEqual = this->CheckEqual(k_pe_output_ref_ptr, k_pe_output_ptr, this->h_k_pe_output->getSize()); EXPECT_TRUE(allEqual); } - - { - this->PerformSetPagedKV(); - sync_check_cuda_error(this->mStream->get()); - this->PerformSetPagedKVRef(); - auto* kv_cache_ptr = bufferCast(*(this->h_kv_cache_tensor)); - auto* kv_cache_ref_ptr = bufferCast(*(this->h_kv_cache_tensor_ref)); - allEqual = this->CheckEqual(kv_cache_ref_ptr, kv_cache_ptr, this->h_kv_cache_tensor->getSize()); - EXPECT_TRUE(allEqual); - } } diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index d52df69bebc1..1b6bf066f147 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1324,86 +1324,6 @@ def load_chunked_kv_cache_for_mla( ) return output_kv, output_k_pe - def set_paged_kv_cache_for_mla( - self, - paged_kv: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - k_pe: torch.Tensor, - metadata: TrtllmAttentionMetadata, - ) -> torch.Tensor: - assert self.is_mla_enable and self.mla_params is not None - assert self.mla_params.qk_nope_head_dim == self.mla_params.v_head_dim - assert metadata.kv_cache_manager is not None - assert paged_kv.shape[0] == metadata.num_contexts - assert paged_kv.is_contiguous() - - num_contexts = metadata.num_contexts - max_seq_len = metadata.max_ctx_kv_len - tokens_per_block = metadata.kv_cache_manager.tokens_per_block - - paged_kv_offsets = torch.ops.trtllm.set_paged_kv_cache_for_mla( - paged_kv, - k, - v, - k_pe, - num_contexts, - metadata.ctx_kv_indptr, - max_seq_len, - self.num_heads, - self.mla_params.qk_nope_head_dim, - self.mla_params.qk_rope_head_dim, - tokens_per_block, - ) - - max_block_num = (max_seq_len + tokens_per_block - 1) // tokens_per_block - assert paged_kv_offsets.shape == (num_contexts, 2, max_block_num) - return paged_kv_offsets - - def set_chunked_kv_cache_for_mla( - self, - paged_kv: torch.Tensor, - kv: torch.Tensor, - k_pe: torch.Tensor, - cu_chunked_seq_len: torch.Tensor, - cached: bool, - metadata: TrtllmAttentionMetadata, - ) -> torch.Tensor: - assert self.is_mla_enable and self.mla_params is not None - assert self.mla_params.qk_nope_head_dim == self.mla_params.v_head_dim - assert metadata.kv_cache_manager is not None - assert paged_kv.shape[0] == metadata.num_contexts - assert paged_kv.is_contiguous() - - kv = kv.contiguous() - k_pe = k_pe.contiguous() - - num_contexts = metadata.num_contexts - tokens_per_block = metadata.kv_cache_manager.tokens_per_block - if cached: - # this indptr is the fake. - cu_seq_len = cu_chunked_seq_len - max_seq_len = metadata.runtime_features.chunk_size - else: - cu_seq_len = metadata.ctx_uncached_token_indptr - max_seq_len = metadata.max_ctx_seq_len - paged_kv_offsets = torch.ops.trtllm.set_chunked_kv_cache_for_mla( - paged_kv, - kv, - k_pe, - num_contexts, - cu_seq_len, - self.num_heads, - self.mla_params.qk_nope_head_dim, - self.mla_params.qk_rope_head_dim, - metadata.kv_cache_manager.tokens_per_block, - max_seq_len, - ) - - max_block_num = (max_seq_len + tokens_per_block - 1) // tokens_per_block - assert paged_kv_offsets.shape == (num_contexts, 2, max_block_num) - return paged_kv_offsets - def mla_rope_append_paged_kv_assign_q( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index ee17c1df091f..af53006a9e92 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -1146,14 +1146,15 @@ def forward_context_with_cached_kv( # out_scale = getattr(self.o_proj, "inv_input_scale", None) out_scale = None # Currently we use BF16 MHA for context phase + # latent_cache must be None to differentiate from normal context phase, + # so that we can skip applying RoPE and appending KV cache inside attention op attn_output = self.mha.forward( q, full_k, full_v, attn_metadata, attention_input_type=AttentionInputType.context_only, - latent_cache= - None, # need to be None to differentiate from normal context phase + latent_cache=None, out_scale=out_scale, output=output, ) @@ -1176,7 +1177,6 @@ def forward_context_with_chunked_prefill( # determine the number of loop # currently we assume that the chunk size is the same as the max_num_tokens - attn_metadata.runtime_features.chunk_size chunked_loop_num = attn_metadata.chunked_loop_num # [toal_token_q, num_heads, 2] -> [toal_token_q, num_heads] float2 @@ -1250,6 +1250,8 @@ def forward_context_with_chunked_prefill( out_scale = None # do not apply mask for attention within loop + # latent_cache must be None to differentiate from normal context phase, + # so that we can skip applying RoPE and appending KV cache inside attention op temp_attn_output = self.mha.forward( q, chunked_k, @@ -1282,10 +1284,10 @@ def forward_context_with_chunked_prefill( ], -1, ) - k = torch.empty_like(q).view(-1, self.num_heads, self.qk_head_dim) - k[..., :self.qk_nope_head_dim] = k_nope.view(-1, self.num_heads, - self.qk_nope_head_dim) - k[..., self.qk_nope_head_dim:] = k_pe.view(-1, 1, self.qk_rope_head_dim) + + k_nope = k_nope.view(-1, self.num_heads, self.qk_nope_head_dim) + k_pe = k_pe.view(-1, 1, self.qk_rope_head_dim) + k = torch.cat((k_nope, k_pe.expand(-1, self.num_heads, -1)), dim=-1) k = k.view(-1, self.num_heads * self.qk_head_dim) # copy q_lens to replace kv_lens_runtime @@ -1298,6 +1300,9 @@ def forward_context_with_chunked_prefill( # out_scale = getattr(self.o_proj, "inv_input_scale", None) out_scale = None # Currently we use BF16 MHA for context phase + + # latent_cache must be None to differentiate from normal context phase, + # so that we can skip applying RoPE and appending KV cache inside attention op temp_attn_output = self.mha.forward( q, k, @@ -1457,6 +1462,7 @@ def forward( attn_metadata: AttentionMetadata, all_reduce_params: Optional[AllReduceParams] = None, ) -> torch.Tensor: + attn_output = self.create_output(hidden_states) if self.register_to_config: torch.ops.trtllm.mla_custom_op_inplace(hidden_states, position_ids, From 70b2c04e8227d60b0c965172a371b07f0f657285 Mon Sep 17 00:00:00 2001 From: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> Date: Fri, 1 Aug 2025 02:38:58 -0700 Subject: [PATCH 04/14] tiny fix Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/mlaKernels.cu | 2 +- cpp/tensorrt_llm/thop/attentionOp.cpp | 3 +-- cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu | 4 ++-- tensorrt_llm/_torch/attention_backend/trtllm.py | 1 - 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index a34de5756e5f..2d9da57ee6fc 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index c23c934d04b4..6b1f535c7c1b 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -194,7 +194,6 @@ class Runner : public RunnerBase { TORCH_CHECK(latent_cache.has_value()); mla_params.latent_cache = static_cast(latent_cache->data_ptr()); - TORCH_CHECK(q_pe.has_value()); TORCH_CHECK(q_pe->dim() == 3); TORCH_CHECK(q_pe->strides()[2] == 1); @@ -469,7 +468,7 @@ void attention(torch::Tensor q, std::optional k, std::optional None: ) def prepare(self) -> None: - # print("prepare TrtllmAttentionMetadata") extra_attrs = get_model_extra_attrs() # If model extra attrs is set, attention_metadata is setup in executor. if extra_attrs is None: From 81c0f85efffe087508d5484d9f5dcce54583be1e Mon Sep 17 00:00:00 2001 From: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> Date: Wed, 6 Aug 2025 06:20:31 -0700 Subject: [PATCH 05/14] fix test_attention_mla.py Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> --- tests/unittest/_torch/test_attention_mla.py | 32 ++++++++++++--------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/tests/unittest/_torch/test_attention_mla.py b/tests/unittest/_torch/test_attention_mla.py index 41c37031bf5c..4941216cc3d1 100644 --- a/tests/unittest/_torch/test_attention_mla.py +++ b/tests/unittest/_torch/test_attention_mla.py @@ -442,20 +442,25 @@ def _run_test_for_backend(backend_name, num_heads, num_kv_heads, num_layers, device=device, ).uniform_(-1, 1) for ctx_len in context_sequence_lengths ]) - ctx_k = torch.cat([ - torch.empty( - [ctx_len, num_kv_heads * qk_head_dim], - dtype=dtype, - device=device, - ).uniform_(-1, 1) for ctx_len in context_sequence_lengths - ]) - ctx_v = torch.cat([ + ctx_kv = torch.cat([ torch.empty( - [ctx_len, num_kv_heads * v_head_dim], + [ctx_len, num_kv_heads * (qk_nope_head_dim + v_head_dim)], dtype=dtype, device=device, ).uniform_(-1, 1) for ctx_len in context_sequence_lengths ]) + # ctx_v.stride(0) == num_kv_heads * (qk_nope_head_dim + v_head_dim) + ctx_k_nope, ctx_v = ctx_kv.split( + [num_kv_heads * qk_nope_head_dim, num_kv_heads * v_head_dim], + dim=-1) + ctx_k_nope = ctx_k_nope.view(-1, num_kv_heads, qk_nope_head_dim) + ctx_k = torch.cat([ + ctx_k_nope, + ctx_k_pe.view(-1, 1, qk_rope_head_dim).expand(-1, num_kv_heads, -1) + ], + dim=-1) + ctx_k = ctx_k.view(-1, num_kv_heads * qk_head_dim) + gen_compressed_kv_list = [ torch.cat([ torch.empty( @@ -676,12 +681,13 @@ def yarn_get_mscale(scale=1, mscale=1): v = inputs_per_layer[layer_idx]["ctx_v"] compressed_kv = inputs_per_layer[layer_idx]["ctx_compressed_kv"] k_pe = inputs_per_layer[layer_idx]["ctx_k_pe"] - qkv = torch.cat([q, k, v], dim=-1) latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) + # q/k will be modified in the forward pass, so we need to clone them + # we should not clone v because we need to keep the stride of v result = ctx_layers[layer_idx].forward( - qkv, - None, - None, + q.clone(), + k.clone(), + v, attn_metadata, attention_input_type=AttentionInputType.context_only, latent_cache=latent_cache, From b183d41b40fec479fa66f081d9965a03ce45eb47 Mon Sep 17 00:00:00 2001 From: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> Date: Mon, 11 Aug 2025 20:36:52 -0700 Subject: [PATCH 06/14] support sm120 separate qkv Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com> --- cpp/kernels/fmha_v2/fmha_test.py | 45 +++++--- cpp/kernels/fmha_v2/setup.py | 42 +++++--- .../fmha_v2/src/fmha/gmem_tile_qkv_packed.h | 100 +++++++++++++----- cpp/kernels/fmha_v2/src/fmha/kernel_traits.h | 38 ++++++- .../fmha_v2/src/fused_multihead_attention.cpp | 2 +- .../cubin/fmha_cubin.h | 68 ++---------- 6 files changed, 175 insertions(+), 120 deletions(-) diff --git a/cpp/kernels/fmha_v2/fmha_test.py b/cpp/kernels/fmha_v2/fmha_test.py index 21486b00ea63..96d21dce8a26 100644 --- a/cpp/kernels/fmha_v2/fmha_test.py +++ b/cpp/kernels/fmha_v2/fmha_test.py @@ -155,31 +155,39 @@ def test_trtllm_sage_attention_fmha(d, s): @pytest.mark.parametrize('dtype', ["-bf16", "-e4m3", "-e4m3 -bf16-output"], ids=["bf16", "e4m3", "e4m3-bf16"]) @pytest.mark.parametrize('s', [1024, 4096], ids=["seqlen-1024", "seqlen-4096"]) -@pytest.mark.parametrize( - 'input_layout', ["", "-paged-kv", "-contiguous-q-kv", "-separate-q-k-v"], - ids=["packed-qkv", "paged-kv", "q-contiguous-kv", "separate-q-k-v"]) -def test_trtllm_context_mla_attention_fmha(dtype, s, input_layout): +def test_trtllm_context_mla_attention_fmha(dtype, s): + sm_version = getSMVersion() + if sm_version < 90: + pytest.skip("MLA kernels are only tested on sm90 and above currently.") + # use higher error tolerance for bf16 and s = 4096. epsilon = '' if dtype == "-bf16" and s == 4096: epsilon += ' -epsilon 0.03' - sm_version = getSMVersion() - if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version != 89: - pytest.skip("FP8 MLAs only supported on sm89 currently.") + if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version != 120: + pytest.skip("FP8 MLAs are only supported on sm120 currently.") - # Context phase kernels. + # Context phase kernels, always use separate-q-k-v layout. subprocess.run( - f"bin/fmha.exe -v 0 -runs 1 -min-s 1024 -s {s} -b 8 -h 8 -d 192 -dv 128 {dtype} \ - -force-non-warp-specialization -causal-mask {epsilon}", + f"bin/fmha.exe -v 0 -runs 1 -min-s 1024 -s {s} -b 8 -h 8 -d 192 -dv 128 {dtype} " + f"-causal-mask {epsilon} -separate-q-k-v", shell=True, check=True) - if sm_version == 90: - # Now only hopper-style supports separate-q-k-v + # For chunked prefill, we need to enable -save-softmax (dtype: bf16, layout: separate-q-k-v). + # Currently fp8 kernel doesn't support saving softmax. + if dtype == "-bf16": + # padding mask + subprocess.run( + f"bin/fmha.exe -v 0 -runs 1 -min-s 1024 -s {s} -b 8 -h 8 -d 192 -dv 128 {dtype} " + f"{epsilon} -separate-q-k-v -save-softmax", + shell=True, + check=True) + # causal mask subprocess.run( - f"bin/fmha.exe -v 0 -runs 1 -min-s 1024 -s {s} -b 8 -h 8 -d 192 -dv 128 {dtype} \ - -causal-mask {epsilon} {input_layout}", + f"bin/fmha.exe -v 0 -runs 1 -min-s 1024 -s {s} -b 8 -h 8 -d 192 -dv 128 {dtype} " + f"-causal-mask {epsilon} -separate-q-k-v -save-softmax", shell=True, check=True) @@ -210,14 +218,17 @@ def test_trtllm_context_mla_attention_fmha(dtype, s, input_layout): "num-grouped-heads-64", "num-grouped-heads-128" ]) def test_trtllm_gen_mla_attention_fmha(dtype, s, num_grouped_heads): + sm_version = getSMVersion() + if sm_version < 90: + pytest.skip("MLA kernels are only tested on sm90 and above currently.") + # use higher error tolerance for bf16 and s = 4096. epsilon = '' if dtype == "-bf16" and s == 4096: epsilon += ' -epsilon 0.03' - sm_version = getSMVersion() - if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version != 89: - pytest.skip("FP8 MLAs only supported on sm89 currently.") + if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version != 120: + pytest.skip("FP8 MLAs are only supported on sm120 currently.") # Generation phase kernels. subprocess.run( diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py index 8434d4225dfe..f92e79a087ca 100644 --- a/cpp/kernels/fmha_v2/setup.py +++ b/cpp/kernels/fmha_v2/setup.py @@ -2075,6 +2075,8 @@ def get_kernel_code(kspec, kname, lname): kernel_traits += '_paged_kv_cache' elif kspec.input_layout == InputLayout.CONTIGUOUS_Q_KV: kernel_traits += '_contiguous_kv_cache' + elif kspec.input_layout == InputLayout.SEPARATE_Q_K_V: + kernel_traits += '_q_k_v' flags = 0 if kspec.ldgsts_q: @@ -3183,7 +3185,7 @@ def get_cubin_header(kernel_traits, specs_names): attention_mask_type_value = attention_mask_type.value # Attention input layout: - # packed_qkv (0), contiguous_q_kv (1), q_paged_kv (2). + # packed_qkv (0), contiguous_q_kv (1), q_paged_kv (2), separate_q_k_v (3). attention_input_layout = InputLayout.PACKED_QKV if '_q_kv' in kname: attention_input_layout = InputLayout.CONTIGUOUS_Q_KV @@ -4702,9 +4704,16 @@ def enumerate_hmma_paged_kv_flash_kernels(specs, sm=80, dtype='fp16'): def enumerate_hmma_flash_kernels(specs, sm=80, dtype='fp16', head_size_v=0): - for (input_layout, enable_attn_logit_softcapping) in \ - product([InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV, InputLayout.Q_PAGED_KV], \ - [False, True]): + input_layouts = [ + InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV, + InputLayout.Q_PAGED_KV + ] + # Deepseek MLA (context 192/128 separate-q-k-v) + if head_size_v == 128: + input_layouts.append(InputLayout.SEPARATE_Q_K_V) + for (input_layout, + enable_attn_logit_softcapping) in product(input_layouts, + [False, True]): enumerate_hmma_flash_kernels_base(specs, sm, dtype, input_layout, enable_attn_logit_softcapping, head_size_v) @@ -5080,7 +5089,7 @@ def enumerate_qmma_flash_kernels(specs, ] input_layouts = [ InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV, - InputLayout.Q_PAGED_KV + InputLayout.Q_PAGED_KV, InputLayout.SEPARATE_Q_K_V ] for (head_size_params, (q_loop_step, kv_loop_step), tiled), input_layout in \ product(params_q_kv_step, input_layouts): @@ -5094,6 +5103,9 @@ def enumerate_qmma_flash_kernels(specs, # skip if head_size is not in head_sizes if head_sizes is not None and head_size not in head_sizes: continue + # skip if head_size_v is not 128 for separate-q-k-v + if input_layout == InputLayout.SEPARATE_Q_K_V and head_size_v != 128: + continue specs.append( kernel_spec(sm=sm, sm_mma=89, @@ -6354,28 +6366,30 @@ def enumerate_kernels(): and kspec.version == 2 and kspec.cross_mha == False and kspec.flash_attention == False) - # Deepseek MLA (192/128 packed + 576/512 paged) - or (kspec.sm in [80, 86, 89, 90, 100, 120] + # Deepseek MLA (generation 576/512 paged) + or (kspec.sm in [90, 100, 120] and kspec.dtype in ['bf16', 'e4m3_fp32'] - and (((kspec.head_size, kspec.head_size_v) == (192, 128) and kspec.input_layout in [InputLayout.PACKED_QKV, InputLayout.Q_PAGED_KV]) - or ((kspec.head_size, kspec.head_size_v) == (576, 512) and kspec.input_layout == InputLayout.Q_PAGED_KV)) + and kspec.head_size == 576 + and kspec.head_size_v == 512 + and kspec.input_layout == InputLayout.Q_PAGED_KV and kspec.sage_block_sizes is None and kspec.version == 2 and kspec.cross_mha == False and kspec.flash_attention == True and kspec.warp_specialization == False and kspec.tiled == True) - # Deepseek MLA (hopper-style context 192/128) - or (kspec.sm == 90 - and kspec.dtype == 'bf16' + # Deepseek MLA (context 192/128 separate-q-k-v) + or (kspec.sm in [90, 100, 120] + and kspec.dtype in ['bf16', 'e4m3_fp32'] and kspec.head_size == 192 and kspec.head_size_v == 128 + and kspec.input_layout == InputLayout.SEPARATE_Q_K_V and kspec.sage_block_sizes is None and kspec.version == 2 and kspec.cross_mha == False and kspec.flash_attention == True - and kspec.warp_specialization == True - and kspec.alibi == False + and ((kspec.warp_specialization == True and kspec.alibi == False) # sm90 + or (kspec.warp_specialization == False and kspec.tiled == True)) # non-sm90 and kspec.enable_attn_logit_softcapping == False) # SageAttention (warp_spec, head_size in (80, 128), packed QKV, padding mask) or (kspec.sm == 90 diff --git a/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h b/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h index 8f54c52b0ba7..172131a22fda 100644 --- a/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h +++ b/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h @@ -418,7 +418,7 @@ struct Gmem_tile_qkv //////////////////////////////////////////////////////////////////////////////////////////////////// -// We expect the Q layout to be [B, S, H, D] with variable sequence length support. +// We expect the Q/K/V layout to be [B, S, H, D] with variable sequence length support. template < // The instruction traits. typename Traits, @@ -440,7 +440,7 @@ template < int NUM_MATS = 1, // Is sliding window attention used ? bool SLIDING_WINDOW_ATTENTION = false> -struct Gmem_tile_q +struct Gmem_tile_q_k_v { // The size of each LDG. @@ -523,22 +523,38 @@ struct Gmem_tile_q USE_LDGSTS = USE_LDGSTS_ }; - // Ctor (keep qkv_offset for compatibility) + // Ctor + // qkv_offset: 0 for Q, 1 for K, 2 for V template - inline __device__ Gmem_tile_q(bert::Fused_multihead_attention_params_v2 const& params, int qkv_offset, + inline __device__ Gmem_tile_q_k_v(bert::Fused_multihead_attention_params_v2 const& params, int qkv_offset, Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0) - : Gmem_tile_q(params, binfo, tidx, cta_row_offset, cta_col_offset_in_bytes) { - } - // Ctor. - template - inline __device__ Gmem_tile_q(bert::Fused_multihead_attention_params_v2 const& params, Block_info const& binfo, - int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0) - : params_q_stride_in_bytes_(params.q_stride_in_bytes) - , actual_seqlen_(binfo.actual_q_seqlen) - , q_ptr_(reinterpret_cast(params.q_ptr)) - { + int seq_offset = 0; + if (qkv_offset == 0) + { + // Q tensor + params_q_k_v_stride_in_bytes_ = params.q_stride_in_bytes; + q_k_v_ptr_ = reinterpret_cast(params.q_ptr); + actual_seqlen_ = binfo.actual_q_seqlen; + seq_offset = binfo.sum_s; + } + else if (qkv_offset == 1) + { + // K tensor + params_q_k_v_stride_in_bytes_ = params.k_stride_in_bytes; + q_k_v_ptr_ = reinterpret_cast(params.k_ptr); + actual_seqlen_ = binfo.actual_kv_seqlen; + seq_offset = binfo.sum_s_kv; + } + else if (qkv_offset == 2) + { + // V tensor + params_q_k_v_stride_in_bytes_ = params.v_stride_in_bytes; + q_k_v_ptr_ = reinterpret_cast(params.v_ptr); + actual_seqlen_ = binfo.actual_kv_seqlen; + seq_offset = binfo.sum_s_kv; + } // Compute the position in the sequence (within the CTA for the moment). int row = tidx / THREADS_PER_ROW; @@ -550,17 +566,20 @@ struct Gmem_tile_q // Do not load/store if the thread is in the padded area col_in_bytes_ = cta_col_offset_in_bytes + col * BYTES_PER_LDG; - // The row offset in the batched GEMM. For each seq element, we store QKV in that order. - // We won't consider past_q_length when loading from gmem_q. - int64_t row_offset = (int64_t) (row + cta_row_offset) * params_q_stride_in_bytes_; - // Add the block index. (sum_s * h + hidx). - int64_t idx = binfo.bidx; + // The row offset in the batched GEMM, including the sequence offset. + int64_t row_offset = (int64_t) (row + cta_row_offset + seq_offset) * params_q_k_v_stride_in_bytes_; + // Add the head index. + int64_t idx = binfo.bidh; // Assemble the final pointer. - q_ptr_ += row_offset + idx * VALID_BYTES_PER_ROW + col_in_bytes_; + q_k_v_ptr_ += row_offset + idx * VALID_BYTES_PER_ROW + col_in_bytes_; // Take the CTA offset to modify the sequence length. actual_seqlen_ -= cta_row_offset; + + // Set the initial seq_len and qkv_offset in case of reinterating + actual_seqlen_init_ = actual_seqlen_; + q_k_v_ptr_init_ = q_k_v_ptr_; } // Store data to shared memory. @@ -590,7 +609,7 @@ struct Gmem_tile_q #pragma unroll for (int ii = 0; ii < LDGS; ++ii) { - ptrs[ii] = q_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_q_stride_in_bytes_; + ptrs[ii] = q_k_v_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_q_k_v_stride_in_bytes_; } // Trigger LDGSTS or the LDGs. @@ -598,10 +617,24 @@ struct Gmem_tile_q Ldgsts_helper::load(this, smem_tile, ptrs, preds); } + // Move the pointer to the next row location. + inline __device__ void move(int const steps = 1) + { + q_k_v_ptr_ += (int64_t) ROWS * params_q_k_v_stride_in_bytes_ * steps; + actual_seqlen_ -= (int) ROWS * steps; + } + + // Move the pointer to the next row location by the offset (not step). + inline __device__ void move_by_offset(int const offset) + { + q_k_v_ptr_ = q_k_v_ptr_init_ + (int64_t) offset * params_q_k_v_stride_in_bytes_; + actual_seqlen_ = actual_seqlen_init_ - (int) offset; + } + // Move the pointer to the next column location inline __device__ void move_col() { - q_ptr_ += (int64_t) COLS * (BITS_PER_ELEMENT / 8); + q_k_v_ptr_ += (int64_t) COLS * (BITS_PER_ELEMENT / 8); // Update col_in_bytes_ to ensure load predicates work col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG; } @@ -609,15 +642,29 @@ struct Gmem_tile_q // Rewind the pointer back to previous column location inline __device__ void rewind_col(int const steps) { - q_ptr_ -= COLS * (BITS_PER_ELEMENT / 8) * steps; + q_k_v_ptr_ -= COLS * (BITS_PER_ELEMENT / 8) * steps; // Update col_in_bytes_ to ensure load predicates work col_in_bytes_ -= THREADS_PER_ROW * BYTES_PER_LDG * steps; } - // The stride between rows for the QKV matrice. - int64_t params_q_stride_in_bytes_; + // Move the pointer to the specified step. + inline __device__ void move_to(int const step) + { + q_k_v_ptr_ = q_k_v_ptr_init_ + (int64_t) ROWS * params_q_k_v_stride_in_bytes_ * step; + actual_seqlen_ = actual_seqlen_init_ - (int) ROWS * step; + } + + inline __device__ void reset() + { + q_k_v_ptr_ = q_k_v_ptr_init_; + actual_seqlen_ = actual_seqlen_init_; + } + + // The stride between rows for the Q/K/V matrice. + int64_t params_q_k_v_stride_in_bytes_; // The pointer. - char* q_ptr_; + char* q_k_v_ptr_; + char* q_k_v_ptr_init_; // The register to store predicates. uint32_t preds_[PRED_REGS]; // The fetch registers. @@ -627,6 +674,7 @@ struct Gmem_tile_q int64_t col_in_bytes_; // The sequence length. int actual_seqlen_; + int actual_seqlen_init_; }; //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h b/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h index 872036ca3d29..065ce36869c5 100644 --- a/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h +++ b/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h @@ -988,6 +988,40 @@ using Kernel_traits_v2 = Kernel_traits_