[https://nvbugs/6368562][fix] In KvCacheCreator, compute an upper-bound MLA FMHA context-workspace size… - #15695
Closed
tensorrt-cicd wants to merge 1 commit into
Closed
Conversation
Contributor
📝 WalkthroughWalkthroughAdds a helper to estimate per-rank MLA context-FMHA workspace bytes from model config. Both ChangesMLA FMHA Workspace Reservation in KV-Cache Sizing
Integration Test Waiver Removal
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 415-426: The MLA workspace estimate in the helper that computes
reserve bytes is using the global attention head count instead of the per-rank
local count. Update the logic around `num_attention_heads` in `_util.py` so it
derives the local head count as `num_attention_heads // tp_size` when
`enable_attention_dp` is off, and only keeps the global count for attention-DP
runs. Make sure the workspace calculation in the same path uses that local value
before computing `workspace_bytes`.
- Around line 441-444: The KV memory calculation in `_TorchExecutor` is
subtracting the MLA FMHA workspace twice:
`AttentionOp::getWorkspaceSizeForContext()` already accounts for it in
`peak_memory` after warmup, so remove the extra `fmha_workspace_reserve`
subtraction from the final `available_kv_mem` formula in the workspace
estimation path. Keep `fmha_workspace_reserve` only for the warmup OOM fallback
logic, and update the capacity computation in
`_estimate_*`/`_estimate_mla_context_workspace_bytes`-related code so the KV
budget is not artificially reduced.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 26feddb8-0eaa-4a73-9d3c-3a32c0e0cf30
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/_util.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
Comment on lines
+415
to
+426
| num_heads = getattr(config, "num_attention_heads", None) | ||
| qk_rope = getattr(config, "qk_rope_head_dim", None) | ||
| qk_nope = getattr(config, "qk_nope_head_dim", None) | ||
| v_head = getattr(config, "v_head_dim", None) | ||
| kv_lora = getattr(config, "kv_lora_rank", None) | ||
| if None in (num_heads, qk_rope, qk_nope, v_head, kv_lora): | ||
| return 0 | ||
| # Per-token: q_buf_2 (kv_lora+qk_rope) + fp8 q/k (qk_rope+qk_nope each) | ||
| # + fp8 v (v_head) + bf16 staging copy of q_buf_2 (2 bytes). | ||
| per_token_bytes = 3 * (kv_lora + qk_rope) + 2 * (qk_rope + | ||
| qk_nope) + v_head | ||
| workspace_bytes = self._max_num_tokens * num_heads * per_token_bytes |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# How MLA context attention partitions query heads across TP, and whether ADP changes it.
rg -nP '(num_attention_heads|num_heads|head).*(tp_size|tensor_parallel|//)' tensorrt_llm/_torch/attention_backend -C2
rg -nP 'num_attention_heads' tensorrt_llm/_torch/models/modeling_deepseekv3.py -C2Repository: NVIDIA/TensorRT-LLM
Length of output: 26348
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== _util.py around the estimator =="
sed -n '390,455p' tensorrt_llm/_torch/pyexecutor/_util.py
echo
echo "== DeepSeekV3 model config wiring =="
sed -n '730,820p' tensorrt_llm/_torch/attention_backend/trtllm.py
echo
echo "== Search for tensor-parallel-specific MLA head handling =="
rg -n "tp_size|tensor_parallel|num_attention_heads|num_key_value_heads|attention_dp|attn.*dp|all_reduce|local heads|head.*tp" \
tensorrt_llm/_torch/attention_backend tensorrt_llm/_torch/models tensorrt_llm/_torch/pyexecutor -C2Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== _util.py estimator and call sites =="
sed -n '403,455p' tensorrt_llm/_torch/pyexecutor/_util.py
echo
sed -n '640,690p' tensorrt_llm/_torch/pyexecutor/_util.py
echo
sed -n '720,810p' tensorrt_llm/_torch/pyexecutor/_util.py
echo
echo "== DeepSeekV3 attention constructor / config =="
rg -n "num_attention_heads|num_key_value_heads|enable_attention_dp|tensor_parallel|tp_size|qk_rope_head_dim|qk_nope_head_dim|v_head_dim|kv_lora_rank" \
tensorrt_llm/_torch/models/modeling_deepseekv3.py -C2Repository: NVIDIA/TensorRT-LLM
Length of output: 28656
Use the local head count for MLA workspace estimation. config.num_attention_heads is model-wide, but MLA attention is sharded by tp_size when enable_attention_dp is off. This helper should use num_attention_heads // tp_size in that case and keep the global count only for attention-DP runs; otherwise the reserve is inflated by roughly tp_size on every TP rank.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 415 - 426, The MLA
workspace estimate in the helper that computes reserve bytes is using the global
attention head count instead of the per-rank local count. Update the logic
around `num_attention_heads` in `_util.py` so it derives the local head count as
`num_attention_heads // tp_size` when `enable_attention_dp` is off, and only
keeps the global count for attention-DP runs. Make sure the workspace
calculation in the same path uses that local value before computing
`workspace_bytes`.
Comment on lines
+441
to
+444
| fmha_workspace_reserve = self._estimate_mla_context_workspace_bytes() | ||
| available_kv_mem = max( | ||
| (total_gpu_memory - peak_memory + allocated_bytes) * fraction - | ||
| fmha_workspace_reserve, 0) |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Determine whether the MLA context FMHA workspace uses the torch allocator (=> already in peak_memory)
# or a separate allocator freed before measurement.
rg -nP 'getWorkspaceSizeForContext|context.*workspace|fmha.*workspace' cpp/tensorrt_llm/common/attentionOp.cpp -C3
rg -nP 'workspace' tensorrt_llm/_torch/attention_backend/trtllm.py -C2Repository: NVIDIA/TensorRT-LLM
Length of output: 8416
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the KV-capacity calculation and the warmup path around the referenced lines.
sed -n '400,500p' tensorrt_llm/_torch/pyexecutor/_util.py
printf '\n----\n'
sed -n '640,700p' tensorrt_llm/_torch/pyexecutor/_util.py
# Locate the MLA workspace estimator and related callers.
rg -n "_estimate_mla_context_workspace_bytes|configure_kv_cache_capacity|peak_memory|allocated_bytes|fmha_workspace_reserve" tensorrt_llm/_torch/pyexecutor/_util.py -n -C 3Repository: NVIDIA/TensorRT-LLM
Length of output: 16166
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Trace where the attention workspace is allocated and which allocator it uses.
rg -n "AttentionWorkspaceManager|buildContextLayout|getWorkspaceSizeForContext|cudaMalloc|torch::empty|at::empty|c10::cuda|IAllocator|workspaceViews" cpp tensorrt_llm -C 2Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the attention workspace manager's allocation site and whether it caches the tensor.
sed -n '1280,1335p' cpp/tensorrt_llm/thop/attentionOp.cpp
printf '\n----\n'
sed -n '208,320p' cpp/tensorrt_llm/thop/attentionOp.hRepository: NVIDIA/TensorRT-LLM
Length of output: 6147
Avoid subtracting the MLA FMHA workspace twice. AttentionOp::getWorkspaceSizeForContext() allocates this buffer with torch::empty(...), so a successful warmup already includes it in peak_memory. Subtracting fmha_workspace_reserve again here double-counts the same bytes and shrinks the KV budget more than necessary; keep the reserve only for the warmup OOM path or remove it from this final capacity formula.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 441 - 444, The KV
memory calculation in `_TorchExecutor` is subtracting the MLA FMHA workspace
twice: `AttentionOp::getWorkspaceSizeForContext()` already accounts for it in
`peak_memory` after warmup, so remove the extra `fmha_workspace_reserve`
subtraction from the final `available_kv_mem` formula in the workspace
estimation path. Keep `fmha_workspace_reserve` only for the warmup OOM fallback
logic, and update the capacity computation in
`_estimate_*`/`_estimate_mla_context_workspace_bytes`-related code so the KV
budget is not artificially reduced.
tensorrt-cicd
force-pushed
the
repair-bot-bug6368562
branch
2 times, most recently
from
June 29, 2026 07:02
5b9f641 to
070f940
Compare
trtllm-agent
force-pushed
the
repair-bot-bug6368562
branch
from
July 6, 2026 03:22
070f940 to
696e2d0
Compare
…estimation The KV-cache estimator builds a max-size cache (8208 blocks for max_seq_len=262144) during the estimation phase, then runs warmup forwards to measure peak memory. For MLA models near the device-memory limit (e.g. Kimi-K2-Thinking NVFP4 TP=4 on B200 with 155 GiB weights/rank), the (max_num_tokens, 0) context-warmup config needs ~900 MiB of FMHA workspace which cannot fit once the KV pool consumes the remaining headroom. The OOM is caught and skipped, leaving the test to hang on follow-up collectives. Reserve the worst-case MLA context-FMHA workspace explicitly when sizing the estimation-phase KV cache and when computing the final KV-cache budget in ``KvCacheCreator._cal_max_memory``. The reservation only applies to MLA models (is_mla(config)); other configs keep the existing budgeting behaviour. Also remove the now-passing test from waives.txt. Verified: TestKimiK2::test_nvfp4[4gpus] passes (MMLU 88.13 > 84.01, GSM8K 93.75 > 87.64). Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
trtllm-agent
force-pushed
the
repair-bot-bug6368562
branch
from
July 7, 2026 00:37
696e2d0 to
a3c7dae
Compare
Collaborator
|
Duplicate with #16399 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
cuda_graph_warmup_block(8208 blocks → 8.59 GiB) without reserving headroom for the MLA FMHA context workspace, so the (max_num_tokens, 0) warmup OOMed trying to grow it to ~900 MiB on a near-saturated B200.KvCacheCreator, compute an upper-bound MLA FMHA context-workspace size frompretrained_config(num_heads, qk_rope+qk_nope, v_head, kv_lora_rank, max_num_tokens), subtract it frommax_memoryin_get_token_num_for_estimation, capnum_cache_blocksaccordingly, and also deduct it fromavailable_kv_memin_cal_max_memory. Only applies whenis_mla(config); non-MLA configs keep existing behavior. Removed the matching waives.txt line.Test plan
Links
Summary by CodeRabbit
Bug Fixes
Tests