From 68e9eab3b7084795fd6bf3bae74279c52329f432 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:41:37 -0700 Subject: [PATCH 1/4] [nvbugs/6631019][fix] Scope the compile flag to multimodal encoder graph capture Multimodal encoder CUDA graphs are captured from load_weights, before the engine applies its own torch.compile decision. is_torch_compiling_flag is a plain module global -- unlike its threading.local / ContextVar neighbours -- so in a reused worker the capture observed the previous engine's raised flag and Attention.forward_impl took its registered custom-op path, which resolves attention metadata from extra_attrs that only an engine forward binds. Give the runner's capture region the same treatment it already gives grad mode: establish the value rather than inherit it. Adds a torch_compiling() contextmanager beside the existing unbalanced setter, since no caller could scope the flag before. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/models/multimodal_encoder_graph.py | 10 ++- tensorrt_llm/_torch/utils.py | 16 +++++ tests/integration/test_lists/waives.txt | 1 - .../modeling/test_multimodal_encoder_graph.py | 69 +++++++++++++------ 4 files changed, 72 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py b/tensorrt_llm/_torch/models/multimodal_encoder_graph.py index a5c4164f68ab..9e32aae15ab0 100644 --- a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py +++ b/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() diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 20f4b754fa9c..b7402f7dc4b1 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/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 diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 7ac341da26d8..8b29d1bfbaa5 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -179,7 +179,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:B300/unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6610548) 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_B200/perf/test_perf_sanity.py::test_e2e[aggr_upload-gemma4_26b_a4b_nvfp4_blackwell-gemma4_26b_a4b_nvfp4_tp1_1k1k] SKIP (https://nvbugs/6571410) full:DGX_B200/perf/test_perf_sanity.py::test_e2e[aggr_upload-host_perf_llama8b_spec_decode-llama8b_spec_bs1_128_128] SKIP (https://nvbugs/6571408) diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py index 43e71972d233..371ab38868a4 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py +++ b/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 == [False, False, False] + + +@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 == [False, False, False] + assert not is_torch_compiling() class _ReallocatingProvider: From 1246721711523020a827f6dceeef680154f9a391 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Wed, 19 Aug 2026 21:29:18 +0000 Subject: [PATCH 2/4] [None][fix] protect RADIO eager attention fallback Signed-off-by: Allison Lim --- tensorrt_llm/_torch/models/modeling_radio.py | 23 +++++++++------- .../_torch/modeling/test_modeling_radio.py | 27 ++++++++++++++++++- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_radio.py b/tensorrt_llm/_torch/models/modeling_radio.py index 573ff7a7b2c5..4bccf16a7841 100644 --- a/tensorrt_llm/_torch/models/modeling_radio.py +++ b/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: diff --git a/tests/unittest/_torch/modeling/test_modeling_radio.py b/tests/unittest/_torch/modeling/test_modeling_radio.py index fb4bc1a4b932..cd7c43eb24a9 100644 --- a/tests/unittest/_torch/modeling/test_modeling_radio.py +++ b/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. From 7ed1fe981a3028cffc820b44e6200b22dd430782 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Wed, 19 Aug 2026 14:39:03 -0700 Subject: [PATCH 3/4] [None][test] harden compile state regression coverage Signed-off-by: Allison Lim --- .../modeling/test_multimodal_encoder_graph.py | 4 +-- tests/unittest/_torch/test_utils.py | 27 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/_torch/test_utils.py diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py index 371ab38868a4..0636481c12e3 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py @@ -571,7 +571,7 @@ def test_capture_uses_inference_mode_when_grad_enabled(cuda_device, make_cuda_ru assert torch.is_grad_enabled() runner.capture_all(cuda_device) - assert observed == [False, False, False] + assert observed and not any(observed) @pytestmark_cuda @@ -596,7 +596,7 @@ def test_capture_lowers_torch_compiling_and_restores_it(cuda_device, make_cuda_r runner.capture_all(cuda_device) assert is_torch_compiling(), "capture must restore the caller's value" - assert observed == [False, False, False] + assert observed and not any(observed) assert not is_torch_compiling() diff --git a/tests/unittest/_torch/test_utils.py b/tests/unittest/_torch/test_utils.py new file mode 100644 index 000000000000..1cc75f8b80db --- /dev/null +++ b/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() From 76cd91252c9beb8ad441bbf3600d2061efa9b1aa Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Wed, 19 Aug 2026 15:08:52 -0700 Subject: [PATCH 4/4] [None][test] include torch utils test in CPU CI Signed-off-by: Allison Lim --- tests/integration/test_lists/test-db/l0_cpu.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 40f884624ba9..21c831211446 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -38,6 +38,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