[https://nvbugs/6631019][fix] Add a torch_compiling(enable) contextmanager to _torch/utils.py and lower…

tensorrt_llm/_torch/models/modeling_radio.py

@@ -28,6 +28,7 @@
     MultimodalEncoderGraphRunner)
 from tensorrt_llm._torch.modules import attention as trtllm_attention
 from tensorrt_llm._torch.modules import mlp as trtllm_mlp
+from tensorrt_llm._torch.utils import torch_compiling
 from tensorrt_llm._utils import prefer_pinned
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
@@ -916,15 +917,19 @@ def _encoder_graph_fn(
 
     def _run_blocks(self, x: torch.Tensor,
                     attn_metadata: AttentionMetadata) -> torch.Tensor:
-        if self._blocks_graph_runner is not None:
-            seq_lengths = attn_metadata.seq_lens.tolist()
-            output = self._blocks_graph_runner.maybe_run(
-                seq_lengths=seq_lengths,
-                inputs={"x": x},
-            )
-            if output is not None:
-                return output["x"]
-        return self._run_blocks_eager(x, attn_metadata)
+        # RADIO runs outside the compiled LM region and receives its vision
+        # metadata explicitly. Keep both graph replay and eager fallback off
+        # the custom-op path, which resolves the LM metadata from extra_attrs.
+        with torch_compiling(False):
+            if self._blocks_graph_runner is not None:
+                seq_lengths = attn_metadata.seq_lens.tolist()
+                output = self._blocks_graph_runner.maybe_run(
+                    seq_lengths=seq_lengths,
+                    inputs={"x": x},
+                )
+                if output is not None:
+                    return output["x"]
+            return self._run_blocks_eager(x, attn_metadata)
 
     def _run_blocks_eager(self, x: torch.Tensor,
                           attn_metadata: AttentionMetadata) -> torch.Tensor:

tensorrt_llm/_torch/models/multimodal_encoder_graph.py

@@ -39,7 +39,7 @@
 import torch
 
 from ...logger import logger
-from ..utils import make_weak_ref
+from ..utils import make_weak_ref, torch_compiling
 
 if TYPE_CHECKING:
     from ...llmapi.llm_args import MultimodalEncoderCudaGraphConfig
@@ -383,7 +383,13 @@ def _capture_key(
             capture_kwargs["pool"] = self._memory_pool
 
         graph = torch.cuda.CUDAGraph()
-        with torch.inference_mode():
+        # `torch_compiling(False)` for the same reason as `inference_mode`: ambient state the encoder
+        # region must not inherit from whoever called us. Capture can be driven from `load_weights`,
+        # outside any engine forward, so the process-global compile flag may still carry a previous
+        # engine's value. `encoder_fn` is handed its metadata explicitly, so its layers must take
+        # the eager path rather than the custom op, which resolves metadata from `extra_attrs` that
+        # only an engine forward binds.
+        with torch.inference_mode(), torch_compiling(False):
             for _ in range(self._config.warmup_steps):
                 self._encoder_fn(static_inputs, metadata)
             torch.cuda.synchronize()

tensorrt_llm/_torch/utils.py

@@ -102,6 +102,22 @@ def is_torch_compiling() -> bool:
     return is_torch_compiling_flag
 
 
+@contextlib.contextmanager
+def torch_compiling(enable: bool):
+    """Scope `is_torch_compiling()` to a region, restoring the prior value.
+
+    The flag is a plain module global, not thread- or context-local, so it
+    outlives the engine that set it. Code running a model region outside an
+    engine forward must establish the value rather than inherit it.
+    """
+    prev_enable = is_torch_compiling()
+    set_torch_compiling(enable)
+    try:
+        yield
+    finally:
+        set_torch_compiling(prev_enable)
+
+
 def set_piecewise_running(enable: bool):
     global is_piecewise_running_flag
     is_piecewise_running_flag = enable

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -39,6 +39,7 @@ l0_cpu:
   - unittest/_torch/speculative/hw_agnostic
   - unittest/_torch/test_mmap_utils.py
   - unittest/_torch/test_model_config.py
+  - unittest/_torch/test_utils.py
   - unittest/_torch/thop/parallel_hw_agnostic
   - unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py
   - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py

tests/integration/test_lists/waives.txt

@@ -188,7 +188,6 @@ full:B300/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_servi
 full:B300/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6475623)
 full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633268)
 full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy SKIP (https://nvbugs/6571418)
-full:DGX_B200/accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] SKIP (https://nvbugs/6631019)
 full:DGX_B200/disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] SKIP (https://nvbugs/6594241)
 full:DGX_H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=False] SKIP (https://nvbugs/6633927)
 full:DGX_H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633927)

tests/unittest/_torch/modeling/test_modeling_radio.py

@@ -8,7 +8,8 @@
 from tensorrt_llm._torch import model_config as model_config_lib
 from tensorrt_llm._torch.models import modeling_radio
 from tensorrt_llm._torch.models.modeling_multimodal_encoder import MultimodalEncoderMixin
-from tensorrt_llm._torch.models.modeling_radio import RADIOVisionModel
+from tensorrt_llm._torch.models.modeling_radio import RADIOVisionModel, VisionTransformer
+from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling
 from tensorrt_llm.llmapi.llm_args import MultimodalEncoderCudaGraphConfig
 from tensorrt_llm.models.modeling_utils import QuantConfig
 from tensorrt_llm.quantization.mode import QuantAlgo
@@ -134,6 +135,30 @@ def _init_finite_weights(model: torch.nn.Module) -> None:
                 param.zero_()
 
 
+def test_radio_run_blocks_lowers_torch_compiling_for_eager_fallback():
+    """A graph miss must keep RADIO attention off the compiled LM path."""
+    x = torch.zeros(1, 4)
+    attn_metadata = mock.Mock()
+    attn_metadata.seq_lens.tolist.return_value = [1]
+
+    observed = []
+    vision_tower = mock.Mock()
+    vision_tower._blocks_graph_runner.maybe_run.side_effect = lambda **_: observed.append(
+        is_torch_compiling()
+    )
+    vision_tower._run_blocks_eager.side_effect = (
+        lambda x, _: observed.append(is_torch_compiling()) or x
+    )
+
+    with torch_compiling(True):
+        output = VisionTransformer._run_blocks(vision_tower, x, attn_metadata)
+        assert is_torch_compiling(), "RADIO must restore the engine's compile flag"
+
+    assert output is x
+    assert observed == [False, False]
+    assert not is_torch_compiling()
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_radio_blocks_cuda_graph_matches_eager(tiny_vit_config):
     """Block-loop CUDA graph wiring must produce the same output as eager.

tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py

@@ -10,7 +10,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Sequence
+from typing import Callable, Dict, List, Optional, Sequence
 from unittest import mock
 
 import pytest
@@ -22,6 +22,7 @@
     MultimodalEncoderGraphRunner,
     _CapturedGraph,
 )
+from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling
 from tensorrt_llm.llmapi.llm_args import MultimodalEncoderCudaGraphConfig
 
 HIDDEN = 8
@@ -387,13 +388,14 @@ def _factory(
         *,
         buckets: List[EncoderGraphKey],
         enable_padding: bool = True,
+        encoder_fn: Optional[Callable] = None,
     ) -> MultimodalEncoderGraphRunner:
         config = MultimodalEncoderCudaGraphConfig(
             buckets=[(bucket.total_tokens, bucket.num_contexts) for bucket in buckets],
             enable_padding=enable_padding,
         )
         return MultimodalEncoderGraphRunner(
-            encoder_fn=_make_toy_encoder_fn(SCALE),
+            encoder_fn=encoder_fn or _make_toy_encoder_fn(SCALE),
             metadata_provider=_ToyMetadataProvider(cuda_device),
             input_specs={
                 "x": EncoderGraphTensorSpec(shape=(HIDDEN,), dtype=torch.float32, token_dim=0),
@@ -538,39 +540,64 @@ def test_replay_matches_eager_across_sizes(cuda_device, make_cuda_runner, real_t
     torch.testing.assert_close(out["y"], expected, rtol=0, atol=0)
 
 
-@pytestmark_cuda
-def test_capture_uses_inference_mode_when_grad_enabled(cuda_device):
-    bucket = EncoderGraphKey(num_contexts=1, total_tokens=128)
-    grad_enabled_states = []
+def _probing_encoder_fn(probe: Callable[[], bool], sink: List[bool]):
+    """Toy encoder that records `probe()` on every call the runner makes.
+
+    Used to assert what ambient state capture establishes for the encoder region,
+    across the warmup steps and the capture itself.
+    """
 
     def encoder_fn(
         inputs: Dict[str, torch.Tensor], metadata: _ToyMetadata
     ) -> Dict[str, torch.Tensor]:
-        grad_enabled_states.append(torch.is_grad_enabled())
+        sink.append(probe())
         x = inputs["x"]
         bias = metadata.seq_lens_cuda.sum().to(x.dtype)
         return {"y": x + bias}
 
-    config = MultimodalEncoderCudaGraphConfig(
-        buckets=[(bucket.total_tokens, bucket.num_contexts)],
-        enable_padding=True,
-        warmup_steps=2,
-    )
-    runner = MultimodalEncoderGraphRunner(
-        encoder_fn=encoder_fn,
-        metadata_provider=_ToyMetadataProvider(cuda_device),
-        input_specs={
-            "x": EncoderGraphTensorSpec(shape=(HIDDEN,), dtype=torch.float32, token_dim=0),
-        },
-        output_specs={"y": 0},
-        config=config,
+    return encoder_fn
+
+
+@pytestmark_cuda
+def test_capture_uses_inference_mode_when_grad_enabled(cuda_device, make_cuda_runner):
+    bucket = EncoderGraphKey(num_contexts=1, total_tokens=128)
+    observed: List[bool] = []
+    runner = make_cuda_runner(
+        buckets=[bucket],
+        encoder_fn=_probing_encoder_fn(torch.is_grad_enabled, observed),
     )
 
     with torch.enable_grad():
         assert torch.is_grad_enabled()
         runner.capture_all(cuda_device)
 
-    assert grad_enabled_states == [False, False, False]
+    assert observed and not any(observed)
+
+
+@pytestmark_cuda
+def test_capture_lowers_torch_compiling_and_restores_it(cuda_device, make_cuda_runner):
+    """Capture must not inherit a raised compile flag from a previous engine.
+
+    `is_torch_compiling_flag` is a plain module global, so in a reused worker it
+    still holds whatever the last engine set. Layers gate their custom-op path on
+    it (`Attention.forward_impl`), and that op resolves attention metadata from
+    `extra_attrs`, which nothing binds during `load_weights`-driven capture. The
+    caller's value must be restored on exit so an engine mid-compile is unaffected.
+    """
+    bucket = EncoderGraphKey(num_contexts=1, total_tokens=128)
+    observed: List[bool] = []
+    runner = make_cuda_runner(
+        buckets=[bucket],
+        encoder_fn=_probing_encoder_fn(is_torch_compiling, observed),
+    )
+
+    with torch_compiling(True):
+        assert is_torch_compiling()
+        runner.capture_all(cuda_device)
+        assert is_torch_compiling(), "capture must restore the caller's value"
+
+    assert observed and not any(observed)
+    assert not is_torch_compiling()
 
 
 class _ReallocatingProvider:

tests/unittest/_torch/test_utils.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling
+
+pytestmark = pytest.mark.cpu_only
+
+
+def test_torch_compiling_restores_flag_after_exception() -> None:
+    with torch_compiling(False):
+        with pytest.raises(RuntimeError), torch_compiling(True):
+            raise RuntimeError
+        assert not is_torch_compiling()