Skip to content

[https://nvbugs/6368562][fix] In KvCacheCreator, compute an upper-bound MLA FMHA context-workspace size… #15695

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

Closed
tensorrt-cicd wants to merge 1 commit into from
Closed

[https://nvbugs/6368562][fix] In KvCacheCreator, compute an upper-bound MLA FMHA context-workspace size… #15695

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
64 changes: 60 additions & 4 deletions tensorrt_llm/_torch/pyexecutor/_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -425,6 +425,35 @@ def _get_kv_size_per_token(self,
num_layers=self._get_num_draft_layers())
return total

def _estimate_mla_context_workspace_bytes(self) -> int:
"""Upper-bound the per-rank MLA context-FMHA workspace.

The estimator's warmup forward at ``max_num_tokens`` allocates this
workspace; on tight configs the allocation OOMs and the OOM is caught,
so ``peak_memory`` under-counts. Reserve it explicitly. See
``getWorkspaceSizeForContext`` in ``cpp/tensorrt_llm/common/attentionOp.cpp``.
Returns 0 for non-MLA models or when required fields are missing.
"""
config = self._model_engine.model.model_config.pretrained_config
if not is_mla(config):
return 0
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
Comment on lines +440 to +451

@coderabbitai coderabbitai Bot Jun 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀 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 -C2

Repository: 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 -C2

Repository: 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 -C2

Repository: 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`.

# 4x slack covers autotuner intermediates (cuBLAS, fp8 GEMM tuning,
# fused_moe scratch) and NCCL symmetric buffers that share this
# headroom during the estimation warmup.
return int(workspace_bytes * 4)

def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
allocated_bytes: int) -> int:
"""
Expand All @@ -434,13 +463,16 @@ def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
"""
kv_size_per_token = self._get_kv_size_per_token()

available_kv_mem = (total_gpu_memory - peak_memory +
allocated_bytes) * fraction
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)
Comment on lines +466 to +469

@coderabbitai coderabbitai Bot Jun 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀 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 -C2

Repository: 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 3

Repository: 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 2

Repository: 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.h

Repository: 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.

logger.info(
f"Peak memory during memory usage profiling (torch + non-torch): {peak_memory / (GB):.2f} GiB, "
f"available KV cache memory when calculating max tokens: {available_kv_mem / (GB):.2f} GiB, "
f"fraction is set {fraction}, kv size per token is {kv_size_per_token}. device total memory {total_gpu_memory / (GB):.2f} GiB, "
f"temporary kv cache memory during profiling {allocated_bytes / (GB):.2f} GiB"
f"temporary kv cache memory during profiling {allocated_bytes / (GB):.2f} GiB, "
f"MLA FMHA workspace reserve {fmha_workspace_reserve / (GB):.2f} GiB"
)
return int(available_kv_mem)

Expand Down Expand Up @@ -648,11 +680,35 @@ def _get_token_num_for_estimation(self) -> int:
return max_num_tokens_for_estimation

free_mem, _ = torch.cuda.mem_get_info()
max_memory = self._kv_cache_config.free_gpu_memory_fraction * free_mem
fmha_workspace_reserve = self._estimate_mla_context_workspace_bytes()
max_memory = max(
self._kv_cache_config.free_gpu_memory_fraction * free_mem -
fmha_workspace_reserve, 0)
kv_size_per_token = self._get_kv_size_per_token()
max_num_tokens_in_memory = (
kv_size_per_token.tokens_for_budget(max_memory) //
self._tokens_per_block * self._tokens_per_block)

# For MLA models the cuda_graph_warmup_block reservation crowds out the
# FMHA workspace and other transient warmup allocations. Cap blocks
# against the reserved budget; configure_kv_cache_capacity computes the
# real final capacity after estimation succeeds.
if fmha_workspace_reserve > 0:
max_blocks_in_memory = (max_num_tokens_in_memory //
self._tokens_per_block)
estimation_min_blocks = ceil_div(
self._max_num_tokens,
self._tokens_per_block) + self._model_engine.batch_size
num_cache_blocks = min(
num_cache_blocks,
max(estimation_min_blocks, max_blocks_in_memory // 2))
max_num_tokens_for_estimation = (
num_cache_blocks * self._tokens_per_block *
self._dummy_reqs[0].sampling_config.beam_width)
logger.info(
f"MLA FMHA context workspace reserve: {fmha_workspace_reserve / (GB):.2f} GiB; "
f"num_cache_blocks (post-cap): {num_cache_blocks}")

return min(max_num_tokens_for_estimation, max_num_tokens_in_memory)

def try_prepare_estimation(self) -> bool:
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,6 @@ accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-cutl
accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-tp4-cutlass-auto] SKIP (https://nvbugs/5596343)
accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[cutlass-auto] SKIP (https://nvbugs/5596343)
accuracy/test_llm_api_pytorch.py::TestKanana_Instruct::test_auto_dtype SKIP (https://nvbugs/6209806)
accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus] SKIP (https://nvbugs/6368562)
accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16[attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6305318)
accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16[attn_backend=TRTLLM-torch_compile=True] SKIP (https://nvbugs/6305318)
accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/5616182)
Expand Down
Loading