From 2c6095774ce6eca8d60bef252fae26b854830467 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 7 Jul 2026 06:18:15 -0700 Subject: [PATCH] [nvbugs/6422315][fix] Guard mamba add_dummy_requests with inference_mode PythonMambaCacheManager.add_dummy_requests writes into self.mamba_cache.prev_num_accepted_tokens / cache_buf_idx on the MTP replay path. Those tensors are inference tensors, so any in-place update must run inside torch.inference_mode(). The sibling helpers _prepare_mamba_cache_blocks and _refresh_dummy_request_mask are already decorated; add_dummy_requests was missed by the replay-checkpoint change, and under attention-DP it is invoked from _pad_attention_dp_dummy_request outside any inference-mode guard, producing 'Inplace update to inference tensor outside InferenceMode is not allowed' on all ranks. Add @torch.inference_mode() to PythonMambaCacheManager.add_dummy_requests and unwaive TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm]. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py | 1 + tests/integration/test_lists/waives.txt | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index d4d54d7df905..54be9fd150d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -725,6 +725,7 @@ def _is_padding_sentinel(self, request_id: int) -> bool: return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= CUDA_GRAPH_DUMMY_REQUEST_ID) + @torch.inference_mode() def add_dummy_requests(self, request_ids: List[int], **kwargs): # Sentinels alias to the shared _padding_slot; non-sentinel # dummies (warmup, attention-DP idle padding) get their own diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index aa78edb371ad..fd88980c548d 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -15,7 +15,6 @@ accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[fp8-4-att accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-attn_dp_off-trtllm] SKIP (temporary ToT main waive; B200 AutoDeploy NVFP4 GSM8K accuracy below threshold) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-4-attn_dp_off-trtllm] SKIP (https://nvbugs/6367792) accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-flashinfer] SKIP (https://nvbugs/6418090) -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm] SKIP (https://nvbugs/6422315) accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6422094) accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws8_80gb-trtllm] SKIP (https://nvbugs/6418090) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8] SKIP (https://nvbugs/6412108)