From a3c7daed567570b7a4f5f3e8fd0b37882aed35b4 Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Sun, 28 Jun 2026 01:47:15 -0700 Subject: [PATCH] [nvbugs/6368562][fix] Reserve MLA FMHA context workspace in KV-cache 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> --- tensorrt_llm/_torch/pyexecutor/_util.py | 64 +++++++++++++++++++++++-- tests/integration/test_lists/waives.txt | 1 - 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b2d45dec50aa..5738a85306e3 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -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 + # 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: """ @@ -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) 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) @@ -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: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 54ee277060a7..a937867ca4bc 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -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)