Skip to content

[https://nvbugs/6029882][fix] Fix attentionOp fp8 mla kvreuse workspace calculation #14852

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pengbowang-nv merged 2 commits into from
Jun 15, 2026
Merged

[https://nvbugs/6029882][fix] Fix attentionOp fp8 mla kvreuse workspace calculation #14852

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
be31457
fix attentionOp fp8 mla kvreuse workspace calc
pengbowang-nv Jun 10, 2026
615aee7
unwaive all possibly related tests
pengbowang-nv Jun 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions cpp/tensorrt_llm/common/attentionOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -748,7 +748,7 @@ size_t AttentionOp::getFmhaMultiCtasKvScratchSize() const noexcept
}

size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t max_num_seq, int32_t input_seq_length,
int32_t cross_kv_length, int32_t max_num_tokens) const noexcept
int32_t cross_kv_length, int32_t max_num_tokens, int32_t total_kv_len) const noexcept
{
if (max_num_tokens == 0)
{
Expand Down Expand Up @@ -828,8 +828,12 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
}
else
{
fp8_k_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast<size_t>(total_k_dim_all_heads);
fp8_v_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast<size_t>(total_v_dim_all_heads);
// Use total_kv_len when available (KV cache reuse causes total_kv_len >> max_num_tokens).
// enqueueContext sizes these buffers by total_kv_len, so workspace must match.
size_t const kv_buf_tokens = std::max(
static_cast<size_t>(total_kv_len), static_cast<size_t>(mChunkPrefillBufferBatchSize) * max_num_tokens);
fp8_k_buf_size = kv_buf_tokens * static_cast<size_t>(total_k_dim_all_heads);
fp8_v_buf_size = kv_buf_tokens * static_cast<size_t>(total_v_dim_all_heads);
}
}
else if (useSageAttnSeparateQkv)
Expand Down
2 changes: 1 addition & 1 deletion cpp/tensorrt_llm/common/attentionOp.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ class AttentionOp
[[nodiscard]] int getHeadSize(bool checkInit = true) const;
[[nodiscard]] int getMaxNumSeqLenTile(int batch_beam_size = 1) const;
[[nodiscard]] size_t getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t nbReq, int32_t max_input_length,
int32_t cross_kv_length = 0, int32_t max_num_tokens = 0) const noexcept;
int32_t cross_kv_length = 0, int32_t max_num_tokens = 0, int32_t total_kv_len = 0) const noexcept;
// total_num_seq is the sum of beam_width for multiple requests
[[nodiscard]] size_t getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32_t total_num_seq,
int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept;
Expand Down
12 changes: 6 additions & 6 deletions cpp/tensorrt_llm/thop/attentionOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -342,7 +342,7 @@ class RunnerBase
virtual ~RunnerBase() = default;
virtual void prepare(AttentionOp& op) const = 0;
virtual int64_t getWorkspaceSize(AttentionOp const& op, int const num_tokens, int const max_attention_window_size,
int const num_gen_tokens, int const max_blocks_per_sequence) const
int const num_gen_tokens, int const max_blocks_per_sequence, int const ctx_total_kv_len = 0) 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,
Expand Down Expand Up @@ -403,10 +403,10 @@ class Runner : public RunnerBase
}

int64_t getWorkspaceSize(AttentionOp const& op, int const num_tokens, int const max_attention_window_size,
int const num_gen_tokens, int const max_blocks_per_sequence) const override
int const num_gen_tokens, int const max_blocks_per_sequence, int const ctx_total_kv_len = 0) const override
{
size_t const context_workspace_size
= op.getWorkspaceSizeForContext(op.mType, max_num_requests, op.mMaxContextLength, 0, num_tokens);
size_t const context_workspace_size = op.getWorkspaceSizeForContext(
op.mType, max_num_requests, op.mMaxContextLength, 0, num_tokens, ctx_total_kv_len);
size_t const generation_workspace_size = op.getWorkspaceSizeForGeneration(
op.mType, max_num_requests, max_attention_window_size, num_gen_tokens, max_blocks_per_sequence);

Expand Down Expand Up @@ -1218,8 +1218,8 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
= beam_width == 1 ? attention_window_size : cache_indirection.value().size(2);
int32_t const max_blocks_per_sequence
= use_kv_cache && kv_cache_block_offsets.has_value() ? kv_cache_block_offsets.value().size(-1) : 0;
int64_t const workspace_size
= runner->getWorkspaceSize(*op, num_tokens, max_attention_window_size, num_gen_tokens, max_blocks_per_sequence);
int64_t const workspace_size = runner->getWorkspaceSize(
*op, num_tokens, max_attention_window_size, num_gen_tokens, max_blocks_per_sequence, ctx_total_kv_len);
TLLM_LOG_TRACE("Expected workspace size is %ld bytes", workspace_size);

torch::Tensor workspace;
Expand Down
5 changes: 0 additions & 5 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,6 @@ accuracy/test_llm_api.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[xg
accuracy/test_llm_api.py::TestMistralNemo12B::test_fp8 SKIP (https://nvbugs/5413197)
accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6281818)
accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6281818)
accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] SKIP (https://nvbugs/6084775)
accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_mtp] SKIP (https://nvbugs/6029882)
accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_pp4_mtp] SKIP (https://nvbugs/6018046)
accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_tp4] SKIP (https://nvbugs/6215793)
accuracy/test_llm_api_pytorch.py::TestDeepSeekR1LongBenchV2::test_fp8_8gpus SKIP (https://nvbugs/6193778)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload] SKIP (https://nvbugs/6185196)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] SKIP (https://nvbugs/6185196)
Expand Down Expand Up @@ -63,7 +59,6 @@ accuracy/test_llm_api_pytorch.py::TestKanana_Instruct::test_auto_dtype SKIP (htt
accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[dep8] SKIP (https://nvbugs/6260890)
accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8] SKIP (https://nvbugs/6248837)
accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8_attn_dp] SKIP (https://nvbugs/6144270)
accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus] SKIP (https://nvbugs/6261793)
accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[pp4-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6075556)
accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/5616182)
accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6211191)
Expand Down
Loading