From dbfebe86ecb200e03742864addcbdb7cb2fd3b00 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Mon, 20 Jul 2026 06:59:03 +0000 Subject: [PATCH 01/23] BCG V1 Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../torch_compile_and_piecewise_cuda_graph.md | 28 +- tensorrt_llm/_torch/modules/attention.py | 24 +- .../_torch/modules/mamba/gdn_mixer.py | 14 +- .../breakable_cuda_graph/__init__.py | 22 ++ .../breakable_cuda_graph.py | 298 ++++++++++++++++++ .../breakable_cuda_graph/context.py | 26 ++ .../breakable_cuda_graph/cuda_utils.py | 25 ++ .../pyexecutor/breakable_cuda_graph_runner.py | 212 +++++++++++++ .../_torch/pyexecutor/model_engine.py | 240 +++++++++----- tensorrt_llm/llmapi/__init__.py | 13 +- tensorrt_llm/llmapi/llm_args.py | 83 ++++- .../usage/llm_args_golden_manifest.json | 18 ++ .../defs/accuracy/test_llm_api_pytorch.py | 47 ++- .../executor/test_breakable_cuda_graph.py | 253 +++++++++++++++ .../api_stability/api_stability_core.py | 5 +- .../api_stability/references/llm.yaml | 8 + tests/unittest/llmapi/test_llm_args.py | 83 ++++- 17 files changed, 1288 insertions(+), 111 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.py create mode 100644 tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py create mode 100644 tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py create mode 100644 tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py create mode 100644 tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py create mode 100644 tests/unittest/_torch/executor/test_breakable_cuda_graph.py diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 9511be887b1d..c7e0167d39f4 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -1,4 +1,4 @@ -# Torch Compile & Piecewise CUDA Graph +# Torch Compile & Prefill CUDA Graph In this guide, we show how to enable torch.compile and Piecewise CUDA Graph in TensorRT LLM. TensorRT LLM uses torch.compile for lightweight vertical fusion and Piecewise CUDA Graph. @@ -41,12 +41,29 @@ To enable torch.compile and Piecewise CUDA Graph, add the following configuratio ```yaml ... # Other extra config +prefill_cuda_graph_backend: piecewise +prefill_capture_num_tokens: '${capture_num_tokens}' # e.g. [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072] torch_compile_config: - capture_num_tokens: '${capture_num_tokens}' # List of num tokens to capture. e.g., [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072] enable_userbuffers: false - enable_piecewise_cuda_graph: true ``` +`TorchCompileConfig.enable_piecewise_cuda_graph` and +`TorchCompileConfig.capture_num_tokens` are deprecated aliases for these +prefill-specific options. + +The experimental breakable implementation can capture the model body without +torch.compile: + +```yaml +prefill_cuda_graph_backend: breakable +prefill_capture_num_tokens: [128, 256, 512] +``` + +The first version of the breakable backend supports BF16 Qwen3.5 on one GPU for +context-only, tensor/pipeline parallelism and mixed context/decode batches with KV cache. Speculative +decoding, LoRA, multimodal inputs, and context +logits fall back to eager execution or are rejected during initialization. + ## Tips for Piecewise CUDA Graph ### Piecewise CUDA Graph & Generation Only CUDA Graph @@ -59,9 +76,10 @@ cuda_graph_config: max_batch_size: 1024 # Specify max capture batch size for generation only cuda graph. By default, TensorRT LLM will generate a capture list based on it. torch_compile_config: - capture_num_tokens: '${capture_num_tokens}' # Specify capture_num_tokens for piecewise cuda graph enable_userbuffers: false - enable_piecewise_cuda_graph: true + +prefill_cuda_graph_backend: piecewise +prefill_capture_num_tokens: '${capture_num_tokens}' ``` ### Piecewise CUDA Graph Padding diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 2579a8143451..084272738786 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -23,6 +23,8 @@ cp_allgather, reducescatter) from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType +from ..pyexecutor.breakable_cuda_graph import (eager_on_graph, + is_in_breakable_cuda_graph) from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, is_torch_compiling) from .linear import (Linear, TensorParallelMode, WeightMode, @@ -117,6 +119,9 @@ def attn_custom_op_inplace( ) +breakable_attn_custom_op_inplace = eager_on_graph(True)(attn_custom_op_inplace) + + def _helix_zero_kv_mask( attn_metadata: AttentionMetadata, num_tokens: int, @@ -963,20 +968,23 @@ def forward_impl( if "mrope_position_deltas" in mrope_config: mrope_position_deltas = mrope_config["mrope_position_deltas"] - # Currently only TRTLLM and FLASHINFER are torch compile compatible backends. - # Only enable custom inplace op when torch compiling. - use_custom_inplace_op = (self.register_to_config - and (self.attn_backend == "TRTLLM" - or self.attn_backend == "FLASHINFER") - and is_torch_compiling() - and not self.is_marlin_enabled) + use_breakable_cuda_graph = (not is_torch_compiling() + and is_in_breakable_cuda_graph()) + # Currently only TRTLLM and FLASHINFER support the custom inplace op. + use_custom_inplace_op = ( + self.register_to_config and + (self.attn_backend == "TRTLLM" or self.attn_backend == "FLASHINFER") + and (is_torch_compiling() or use_breakable_cuda_graph) + and not self.is_marlin_enabled) if use_custom_inplace_op: outputs = create_attn_outputs(q, attention_mask, self.layer_idx_str) assert len(outputs) == 1 or len(outputs) == 2 output = outputs[0] output_sf = outputs[1] if len(outputs) == 2 else None - attn_custom_op_inplace( + custom_op = (breakable_attn_custom_op_inplace if + use_breakable_cuda_graph else attn_custom_op_inplace) + custom_op( q, k, v, diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 32aef7e61bcc..9136683df510 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -31,6 +31,7 @@ from ...attention_backend import AttentionMetadata from ...distributed import AllReduceParams from ...model_config import ModelConfig +from ...pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ...speculative import SpecMetadata from ...utils import EventType, get_model_extra_attrs, is_gdn_replay_enabled, is_torch_compiling from ..linear import FP8QDQLinearMethod, Linear, TensorParallelMode @@ -174,6 +175,9 @@ def gdn_custom_op_inplace( ) +breakable_gdn_custom_op_inplace = eager_on_graph(True)(gdn_custom_op_inplace) + + def ensure_divisibility(numerator, denominator): """Ensure that numerator is divisible by the denominator.""" assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator) @@ -1053,11 +1057,17 @@ def forward( ): mixed_qkv, z, a, b = self._compute_tokenwise_inputs(hidden_states) - if self.register_to_config and is_torch_compiling(): + use_breakable_cuda_graph = not is_torch_compiling() and is_in_breakable_cuda_graph() + if self.register_to_config and (is_torch_compiling() or use_breakable_cuda_graph): attn_out = mixed_qkv.new_empty( (1, mixed_qkv.shape[0], self.num_v_heads_per_tp, self.head_v_dim) ) - gdn_custom_op_inplace(mixed_qkv, a, b, self.layer_idx_str, attn_out) + custom_op = ( + breakable_gdn_custom_op_inplace + if use_breakable_cuda_graph + else gdn_custom_op_inplace + ) + custom_op(mixed_qkv, a, b, self.layer_idx_str, attn_out) else: attn_out = self.forward_core( mixed_qkv, diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.py new file mode 100644 index 000000000000..f4ff7bbb24df --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.py @@ -0,0 +1,22 @@ +# Adapted from SGLang's breakable CUDA graph implementation. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, + break_graph, + eager_on_graph, + get_current_replay_token, +) +from .context import enable_breakable_cuda_graph, is_in_breakable_cuda_graph + +__all__ = [ + "BreakableCUDAGraph", + "BreakableCUDAGraphCapture", + "break_graph", + "eager_on_graph", + "enable_breakable_cuda_graph", + "get_current_replay_token", + "is_in_breakable_cuda_graph", +] diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py new file mode 100644 index 000000000000..8136f5e51459 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -0,0 +1,298 @@ +# Adapted from SGLang's breakable CUDA graph implementation. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import functools +import itertools +import logging +import threading +from contextvars import ContextVar +from typing import Any, Callable, Optional + +import torch +from cuda.bindings import runtime as rt + +from ...utils import make_weak_ref +from .cuda_utils import check_cuda_errors + +logger = logging.getLogger(__name__) + +__all__ = [ + "BreakableCUDAGraph", + "BreakableCUDAGraphCapture", + "break_graph", + "eager_on_graph", + "get_current_replay_token", +] + +_current_capture: ContextVar[Optional["BreakableCUDAGraphCapture"]] = ContextVar( + "breakable_cuda_graph_capture", default=None +) +_current_stream: ContextVar[Optional[torch.cuda.Stream]] = ContextVar( + "breakable_cuda_graph_stream", default=None +) +_current_replay_token: ContextVar[Optional[int]] = ContextVar( + "breakable_cuda_graph_replay_token", default=None +) +_forked_streams: ContextVar[Optional[set[torch.cuda.Stream]]] = ContextVar( + "breakable_cuda_graph_forked_streams", default=None +) +_replay_token_counter = itertools.count(1) + +_original_wait_stream: Optional[Callable] = None +_wait_stream_hook_lock = threading.Lock() +_wait_stream_hook_refcount = 0 + + +def get_current_stream(device: Optional[torch.device] = None) -> torch.cuda.Stream: + """Return the active BCG stream or PyTorch's current stream.""" + stream = _current_stream.get() + return torch.cuda.current_stream(device) if stream is None else stream + + +def get_current_replay_token() -> Optional[int]: + """Return a unique token for the active BCG replay.""" + return _current_replay_token.get() + + +def _capture_status(stream_ptr: int) -> rt.cudaStreamCaptureStatus: + status, *_ = check_cuda_errors(rt.cudaStreamGetCaptureInfo(stream_ptr)) + return status + + +def _is_stream_capturing(stream: torch.cuda.Stream) -> bool: + return ( + _capture_status(stream.cuda_stream) + == rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive + ) + + +def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream) -> None: + assert _original_wait_stream is not None + forked = _forked_streams.get() + capturing = _current_stream.get() + if forked is None or capturing is None: + _original_wait_stream(self, other) + return + + capture_ptr = capturing.cuda_stream + self_is_capture = self is capturing or self.cuda_stream == capture_ptr + other_is_capture = other is capturing or other.cuda_stream == capture_ptr + if self_is_capture and not other_is_capture: + if not _is_stream_capturing(other): + return + _original_wait_stream(self, other) + forked.discard(other) + elif other_is_capture and not self_is_capture: + _original_wait_stream(self, other) + forked.add(self) + else: + _original_wait_stream(self, other) + + +def _install_wait_stream_hook() -> None: + global _original_wait_stream, _wait_stream_hook_refcount + with _wait_stream_hook_lock: + if _wait_stream_hook_refcount == 0: + _original_wait_stream = torch.cuda.Stream.wait_stream + torch.cuda.Stream.wait_stream = _hooked_wait_stream + _wait_stream_hook_refcount += 1 + + +def _uninstall_wait_stream_hook() -> None: + global _original_wait_stream, _wait_stream_hook_refcount + with _wait_stream_hook_lock: + _wait_stream_hook_refcount -= 1 + if _wait_stream_hook_refcount == 0: + assert _original_wait_stream is not None + torch.cuda.Stream.wait_stream = _original_wait_stream + _original_wait_stream = None + + +def _weak_ref_if_tensor(value: Any) -> Any: + if torch.is_tensor(value): + return make_weak_ref(value) + if isinstance(value, tuple): + return tuple(_weak_ref_if_tensor(item) for item in value) + if isinstance(value, list): + return [_weak_ref_if_tensor(item) for item in value] + if isinstance(value, dict): + return {key: _weak_ref_if_tensor(item) for key, item in value.items()} + return value + + +def _copy_output(destination: Any, source: Any) -> Any: + if torch.is_tensor(destination) and torch.is_tensor(source): + destination.copy_(source) + return destination + + if ( + isinstance(destination, (tuple, list)) + and isinstance(source, (tuple, list)) + and len(destination) == len(source) + ): + copied = [_copy_output(dst, src) for dst, src in zip(destination, source)] + return tuple(copied) if isinstance(destination, tuple) else copied + + if hasattr(destination, "__dict__") and hasattr(source, "__dict__"): + for key, source_value in source.__dict__.items(): + destination_value = getattr(destination, key, None) + if torch.is_tensor(destination_value) and torch.is_tensor(source_value): + destination_value.copy_(source_value) + else: + setattr(destination, key, source_value) + return destination + + if isinstance(destination, dict) and isinstance(source, dict): + for key, source_value in source.items(): + destination_value = destination.get(key) + if torch.is_tensor(destination_value) and torch.is_tensor(source_value): + destination_value.copy_(source_value) + else: + destination[key] = source_value + return destination + + return source + + +def eager_on_graph(enable: bool) -> Callable[[Callable], Callable]: + """Run a callable eagerly between captured CUDA graph segments.""" + + def decorator(inner: Callable) -> Callable: + if not enable: + return inner + + @functools.wraps(inner) + def wrapper(*args, **kwargs): + capture = _current_capture.get() + if capture is None: + return inner(*args, **kwargs) + + logger.debug( + "Break CUDA graph for function %s", getattr(inner, "__name__", type(inner).__name__) + ) + capture._end_current_segment() + output = inner(*args, **kwargs) + + captured_args = tuple(_weak_ref_if_tensor(arg) for arg in args) + captured_kwargs = {key: _weak_ref_if_tensor(value) for key, value in kwargs.items()} + captured_output = _weak_ref_if_tensor(output) + + def replay_fn() -> Any: + new_output = inner(*captured_args, **captured_kwargs) + return _copy_output(captured_output, new_output) + + capture.cuda_graph._break_functions.append(replay_fn) + capture._begin_new_segment() + return output + + return wrapper + + return decorator + + +class BreakableCUDAGraph: + """A sequence of CUDA graph segments separated by eager functions.""" + + def __init__(self) -> None: + self._segments: list[torch.cuda.CUDAGraph] = [] + self._break_functions: list[Callable[[], Any]] = [] + + @property + def num_segments(self) -> int: + return len(self._segments) + + @property + def num_breaks(self) -> int: + return len(self._break_functions) + + def pool(self): + if not self._segments: + raise RuntimeError("Cannot get the pool of an empty BCG") + return self._segments[0].pool() + + def replay(self) -> None: + stream_token = _current_stream.set(torch.cuda.current_stream()) + replay_token = _current_replay_token.set(next(_replay_token_counter)) + try: + for index, segment in enumerate(self._segments): + segment.replay() + if index < len(self._break_functions): + self._break_functions[index]() + finally: + _current_replay_token.reset(replay_token) + _current_stream.reset(stream_token) + + def reset(self) -> None: + for segment in self._segments: + segment.reset() + self._segments.clear() + self._break_functions.clear() + + +class BreakableCUDAGraphCapture: + """Capture a region as CUDA graph segments separated by eager work.""" + + def __init__( + self, + cuda_graph: BreakableCUDAGraph, + pool=None, + stream: Optional[torch.cuda.Stream] = None, + capture_error_mode: str = "global", + ) -> None: + if not isinstance(cuda_graph, BreakableCUDAGraph): + raise TypeError("cuda_graph must be a BreakableCUDAGraph") + self.cuda_graph = cuda_graph + self._pool = (0, 0) if pool is None else pool + self._stream = stream + self._capture_error_mode = capture_error_mode + self._stream_context = None + self._capture_token = None + self._stream_token = None + self._forked_token = None + + def __enter__(self) -> "BreakableCUDAGraphCapture": + _install_wait_stream_hook() + if self._stream is not None: + self._stream_context = torch.cuda.stream(self._stream) + self._stream_context.__enter__() + self._capture_token = _current_capture.set(self) + self._stream_token = _current_stream.set(self._stream or torch.cuda.current_stream()) + self._forked_token = _forked_streams.set(set()) + self._begin_new_segment() + return self + + def __exit__(self, *args: object) -> bool: + try: + self._end_current_segment() + finally: + _forked_streams.reset(self._forked_token) + _current_stream.reset(self._stream_token) + _current_capture.reset(self._capture_token) + if self._stream_context is not None: + self._stream_context.__exit__(*args) + self._stream_context = None + _uninstall_wait_stream_hook() + return False + + def _begin_new_segment(self) -> None: + segment = torch.cuda.CUDAGraph() + segment.capture_begin(pool=self._pool, capture_error_mode=self._capture_error_mode) + self.cuda_graph._segments.append(segment) + + def _end_current_segment(self) -> None: + main_stream = get_current_stream() + forked = _forked_streams.get() + if forked: + assert _original_wait_stream is not None + for side_stream in list(forked): + if _is_stream_capturing(side_stream): + _original_wait_stream(main_stream, side_stream) + forked.clear() + self.cuda_graph._segments[-1].capture_end() + + +@eager_on_graph(True) +def break_graph() -> None: + """Insert an empty eager break between CUDA graph segments.""" + return None diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py new file mode 100644 index 000000000000..610a4da16a20 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py @@ -0,0 +1,26 @@ +# Adapted from SGLang's breakable CUDA graph implementation. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Iterator + +_breakable_cuda_graph_active: ContextVar[bool] = ContextVar( + "breakable_cuda_graph_active", default=False +) + + +def is_in_breakable_cuda_graph() -> bool: + """Return whether the current context is executing a BCG region.""" + return _breakable_cuda_graph_active.get() + + +@contextmanager +def enable_breakable_cuda_graph() -> Iterator[None]: + """Mark capture or replay work as breakable CUDA graph execution.""" + token = _breakable_cuda_graph_active.set(True) + try: + yield + finally: + _breakable_cuda_graph_active.reset(token) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py new file mode 100644 index 000000000000..6c56ba287218 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py @@ -0,0 +1,25 @@ +# Adapted from SGLang's breakable CUDA graph implementation. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from cuda.bindings import runtime as rt + + +def _cuda_get_error_string(error: rt.cudaError_t) -> str: + result, message = rt.cudaGetErrorString(error) + if result != rt.cudaError_t.cudaSuccess: + return "" + if isinstance(message, bytes): + return message.decode("utf-8", "replace") + return str(message) + + +def check_cuda_errors(result): + """Raise a Python exception for a failed cuda-python runtime call.""" + if result[0] != rt.cudaError_t.cudaSuccess: + raise RuntimeError(f"CUDA error {int(result[0])}({_cuda_get_error_string(result[0])})") + if len(result) == 1: + return None + if len(result) == 2: + return result[1] + return result[1:] diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py new file mode 100644 index 000000000000..e2b27646c52a --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +from enum import Enum +from typing import Any, Callable, Iterator, Optional + +import torch +from torch import nn + +from ..utils import make_weak_ref +from .breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, + enable_breakable_cuda_graph, +) + + +class BreakableCUDAGraphRunnerState(Enum): + IDLE = "idle" + WARMUP = "warmup" + CAPTURE = "capture" + REPLAY = "replay" + + +class BreakableCUDAGraphRunner: + """Capture and replay prefill model bodies as breakable CUDA graphs.""" + + _WARMUP_STEPS = 2 + + def __init__(self, layer_model: nn.Module, logits_processor: nn.Module) -> None: + self.layer_model = layer_model + self.logits_processor = logits_processor + self._graphs: dict[int, BreakableCUDAGraph] = {} + self._outputs: dict[int, torch.Tensor] = {} + self._memory_pool = None + self._capture_stream = torch.cuda.Stream() + self._shared_output: Optional[torch.Tensor] = None + self._state = BreakableCUDAGraphRunnerState.IDLE + self._active_graph: Optional[BreakableCUDAGraph] = None + self._active_num_tokens: Optional[int] = None + + @property + def state(self) -> BreakableCUDAGraphRunnerState: + return self._state + + @property + def is_warming_up(self) -> bool: + return self._state == BreakableCUDAGraphRunnerState.WARMUP + + @property + def is_capturing(self) -> bool: + return self._state == BreakableCUDAGraphRunnerState.CAPTURE + + def has_graph(self, num_tokens: int) -> bool: + return num_tokens in self._graphs + + def warmup(self, engine_forward: Callable[[], Any], steps: int = _WARMUP_STEPS) -> None: + """Run the complete eager engine forward under the warmup state. + model_engine.forward will use state to determine what forward to do.""" + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot warm up BCG while runner is {self._state.value}") + self._state = BreakableCUDAGraphRunnerState.WARMUP + try: + for _ in range(steps): + engine_forward() + finally: + self._state = BreakableCUDAGraphRunnerState.IDLE + + def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None: + """Warm up eagerly, then capture one prefill token bucket.""" + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot capture BCG while runner is {self._state.value}") + if num_tokens in self._graphs: + raise ValueError(f"BCG for num_tokens={num_tokens} is already captured") + + current_stream = torch.cuda.current_stream() + self._capture_stream.wait_stream(current_stream) + graph = None + try: + with torch.cuda.stream(self._capture_stream): + self.warmup(engine_forward) + + self._state = BreakableCUDAGraphRunnerState.CAPTURE + graph = BreakableCUDAGraph() + self._active_graph = graph + self._active_num_tokens = num_tokens + output = engine_forward() + + current_stream.wait_stream(self._capture_stream) + if not torch.is_tensor(output): + raise TypeError( + f"Breakable prefill capture requires a tensor body output, got {type(output)}" + ) + assert graph is not None + self._graphs[num_tokens] = graph + self._outputs[num_tokens] = make_weak_ref(output) + if self._memory_pool is None: + self._memory_pool = graph.pool() + except Exception: + if graph is not None: + graph.reset() + raise + finally: + self._active_graph = None + self._active_num_tokens = None + self._state = BreakableCUDAGraphRunnerState.IDLE + + @contextlib.contextmanager + def capture_context(self) -> Iterator[None]: + """Open the segmented CUDA graph capture for the active bucket.""" + if not self.is_capturing or self._active_graph is None: + raise RuntimeError("BCG capture context requested outside capture") + with ( + enable_breakable_cuda_graph(), + BreakableCUDAGraphCapture( + self._active_graph, pool=self._memory_pool, stream=self._capture_stream + ), + ): + yield + + def capture_output(self, output: torch.Tensor) -> torch.Tensor: + """Route all bucket outputs through the largest capture's buffer. """ + + if not self.is_capturing or self._active_num_tokens is None: + raise RuntimeError("BCG output registered outside capture") + num_tokens = self._active_num_tokens + if self._shared_output is None: + self._shared_output = make_weak_ref(output) + return self._shared_output + if num_tokens > self._shared_output.shape[0]: + raise ValueError( + "BCG buckets must be captured in descending order: " + f"{num_tokens} exceeds shared output size " + f"{self._shared_output.shape[0]}" + ) + self._shared_output[:num_tokens].copy_(output[:num_tokens]) + return self._shared_output[:num_tokens] + + def capture_model_body(self, outer_forward: Callable[[], Any]) -> Any: + """Run the outer model while capturing only its decoder body. + model_engine.forward is too broad and may pollute the CUDA stream + before the actual model forward. We want to reuse the functions + in forward that prepare the data and set the relevant flags.""" + if not self.is_capturing: + raise RuntimeError("BCG body capture requested outside capture") + + original_body_forward = self.layer_model.forward + original_logits_forward = self.logits_processor.forward + captured_output = None + + def capture_forward(*args, **kwargs): + nonlocal captured_output + with self.capture_context(): + captured_output = self.capture_output( + original_body_forward(*args, **kwargs)) + return captured_output + + def passthrough_forward(hidden_states, *args, **kwargs): + del args, kwargs + return hidden_states + + self.layer_model.forward = capture_forward + self.logits_processor.forward = passthrough_forward + try: + outer_forward() + if captured_output is None: + raise RuntimeError("BCG capture did not execute the model body") + return captured_output + finally: + self.logits_processor.forward = original_logits_forward + self.layer_model.forward = original_body_forward + + def replay(self, num_tokens: int) -> torch.Tensor: + if num_tokens not in self._graphs: + raise KeyError(f"No BCG captured for num_tokens={num_tokens}") + self._graphs[num_tokens].replay() + return self._outputs[num_tokens] + + def execute(self, num_tokens: int, outer_forward: Callable[[], Any]) -> Any: + """Patch the body with replay while preserving the outer forward. + this function reuse model_engine._forward_step to set flags. + and just patch the body model forward""" + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot execute BCG while runner is {self._state.value}") + if num_tokens not in self._graphs: + raise KeyError(f"No BCG captured for num_tokens={num_tokens}") + + original_forward = self.layer_model.forward + + def replay_forward(*args, **kwargs): + del args, kwargs + return self.replay(num_tokens) + + self._state = BreakableCUDAGraphRunnerState.REPLAY + self.layer_model.forward = replay_forward + try: + with enable_breakable_cuda_graph(): + return outer_forward() + finally: + self.layer_model.forward = original_forward + self._state = BreakableCUDAGraphRunnerState.IDLE + + def clear(self) -> None: + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot clear BCG while runner is {self._state.value}") + for graph in self._graphs.values(): + graph.reset() + self._graphs.clear() + self._outputs.clear() + self._shared_output = None + self._memory_pool = None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6bfa53c107a5..1a79e0155abc 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -35,6 +35,7 @@ create_input_processor_with_hash) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, + PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) from tensorrt_llm.logger import logger @@ -76,6 +77,7 @@ from ..utils import (get_model_extra_attrs, set_per_request_piecewise_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) +from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner from .config_utils import is_mla from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, @@ -170,14 +172,14 @@ def warmup(self, resource_manager: ResourceManager) -> None: return -def _filter_piecewise_capture_num_tokens( +def _filter_prefill_capture_num_tokens( candidate_num_tokens: list[int], max_num_tokens: int, max_batch_size: int, max_seq_len: int, num_extra_decoding_steps: int = 0, ) -> Tuple[list[int], list[int]]: - """Cap piecewise CUDA graph capture candidates at the engine's reachable + """Cap prefill CUDA graph capture candidates at the engine's reachable `num_tokens` ceiling `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` clamping user-requested sizes above it down to the ceiling. @@ -200,10 +202,10 @@ def _filter_piecewise_capture_num_tokens( """ max_capturable_num_tokens = max( 0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)) - piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens) - if piecewise_capacity_limit > 0: + prefill_capacity_limit = min(max_num_tokens, max_capturable_num_tokens) + if prefill_capacity_limit > 0: kept = sorted({ - min(i, piecewise_capacity_limit) + min(i, prefill_capacity_limit) for i in candidate_num_tokens if 0 < i <= max_num_tokens }) else: @@ -625,15 +627,14 @@ def __init__( and bool(self._cuda_graph_seq_lens)) self.torch_compile_config = self.llm_args.torch_compile_config + self.prefill_cuda_graph_backend = self.llm_args.prefill_cuda_graph_backend torch_compile_enabled = bool(self.torch_compile_config is not None) torch_compile_fullgraph = self.torch_compile_config.enable_fullgraph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_fullgraph'].default torch_compile_inductor_enabled = self.torch_compile_config.enable_inductor if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_inductor'].default - torch_compile_piecewise_cuda_graph = self.torch_compile_config.enable_piecewise_cuda_graph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ - 'enable_piecewise_cuda_graph'].default - torch_compile_piecewise_cuda_graph_num_tokens = self.torch_compile_config.capture_num_tokens if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ - 'capture_num_tokens'].default + torch_compile_piecewise_cuda_graph = (self.prefill_cuda_graph_backend == + PrefillCudaGraphBackend.PIECEWISE) torch_compile_enable_userbuffers = self.torch_compile_config.enable_userbuffers if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_userbuffers'].default torch_compile_max_num_streams = self.torch_compile_config.max_num_streams if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ @@ -642,14 +643,14 @@ def __init__( self._torch_compile_enabled = torch_compile_enabled self._torch_compile_piecewise_cuda_graph = torch_compile_piecewise_cuda_graph - piecewise_cuda_graph_num_tokens = ( - torch_compile_piecewise_cuda_graph_num_tokens - or cuda_graph_batch_sizes or []) + prefill_cuda_graph_num_tokens = self.llm_args.prefill_capture_num_tokens + if prefill_cuda_graph_num_tokens is None: + prefill_cuda_graph_num_tokens = cuda_graph_batch_sizes or [] num_extra_decoding_steps = self._get_num_extra_decoding_steps() - self._piecewise_cuda_graph_num_tokens, unrecordable = ( - _filter_piecewise_capture_num_tokens( - piecewise_cuda_graph_num_tokens, + self._prefill_cuda_graph_num_tokens, unrecordable = ( + _filter_prefill_capture_num_tokens( + prefill_cuda_graph_num_tokens, max_num_tokens=self.max_num_tokens, max_batch_size=self.batch_size, max_seq_len=self.max_seq_len, @@ -657,7 +658,7 @@ def __init__( )) if unrecordable: logger.warning( - f"Skipping piecewise CUDA graph capture for num_tokens=" + f"Skipping prefill CUDA graph capture for num_tokens=" f"{unrecordable}: exceeds reachable ceiling " f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)=" f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. " @@ -682,7 +683,7 @@ def __init__( enable_userbuffers=use_ub, enable_piecewise_cuda_graph=self. _torch_compile_piecewise_cuda_graph, - capture_num_tokens=self._piecewise_cuda_graph_num_tokens, + capture_num_tokens=self._prefill_cuda_graph_num_tokens, max_num_streams=torch_compile_max_num_streams, mapping=self.mapping) apply_llm_torch_compile = getattr(self.model, @@ -923,6 +924,21 @@ def __init__( enable_encoder_decoder_mixed_cuda_graph), ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) + self.breakable_cuda_graph_runner = None + if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: + if self.spec_config is not None: + raise ValueError( + "breakable prefill CUDA graph does not support speculative decoding" + ) + decoder_model = (self.model if isinstance( + self.model, DecoderModelForCausalLM) else getattr( + self.model, "llm", None)) + if not isinstance(decoder_model, DecoderModelForCausalLM): + raise ValueError( + "breakable prefill CUDA graph requires a decoder model body" + ) + self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner( + decoder_model.model, decoder_model.logits_processor) # Initialize CUDA Graph LoRA manager if LoRA is enabled self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None @@ -2043,7 +2059,8 @@ def _get_graphs_to_capture( def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): """Warm up or capture CUDA graphs for the configured graph shapes.""" if not (self.cuda_graph_runner.enabled - or self._torch_compile_piecewise_cuda_graph): + or self.prefill_cuda_graph_backend + != PrefillCudaGraphBackend.DISABLED): return self._capture_generation_cuda_graphs(resource_manager) @@ -2051,7 +2068,7 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. if not self.cuda_graph_runner.is_warmup_only: - self._capture_piecewise_cuda_graphs(resource_manager) + self._capture_prefill_cuda_graphs(resource_manager) @torch.inference_mode() @with_warmup_flag @@ -2447,18 +2464,24 @@ def _capture_mixed_encoder_decoder_cuda_graphs( self.enable_spec_decode = saved_enable_spec_decode self.runtime_draft_len = saved_runtime_draft_len - def _capture_piecewise_cuda_graphs(self, resource_manager: ResourceManager): - """Captures piecewise CUDA graphs for context/prefill steps via torch.compile.""" - if not (self._torch_compile_piecewise_cuda_graph - and self._torch_compile_enabled): + def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): + """Capture configured CUDA graphs for context/prefill steps.""" + if (self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.DISABLED + or (self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.PIECEWISE + and not self._torch_compile_enabled)): return - logger.info("Running piecewise CUDA graph warmup...") - piecewise_cuda_graph_num_tokens = sorted( - self._piecewise_cuda_graph_num_tokens, reverse=True) + logger.info("Running prefill CUDA graph warmup...") + prefill_cuda_graph_num_tokens = sorted( + self._prefill_cuda_graph_num_tokens, reverse=True) - with capture_piecewise_cuda_graph(True), self.no_cuda_graph(): - for num_tokens in piecewise_cuda_graph_num_tokens: + capture_context = (capture_piecewise_cuda_graph(True) + if self._torch_compile_piecewise_cuda_graph else + contextlib.nullcontext()) + with capture_context, self.no_cuda_graph(): + for num_tokens in prefill_cuda_graph_num_tokens: warmup_request = self._create_warmup_request( resource_manager, num_tokens, 0) with self._release_batch_context(warmup_request, @@ -2467,26 +2490,29 @@ def _capture_piecewise_cuda_graphs(self, resource_manager: ResourceManager): continue logger.info( - f"Run piecewise CUDA graph warmup for num tokens={num_tokens}" + f"Run prefill CUDA graph capture for num tokens={num_tokens}" ) - # Run a few times to ensure capture - for _ in range(3): - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) + if self.breakable_cuda_graph_runner is not None: + self.breakable_cuda_graph_runner.capture( + num_tokens, lambda: self.forward( + batch, + new_tensors_device=None, + resource_manager=resource_manager)) + else: + # Run a few times to ensure torch.compile capture. + for _ in range(4): + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) torch.cuda.synchronize() gc.collect() torch.cuda.empty_cache() - # When using piecewise cuda graph, the logits may suffer severe memory fragmentation problem. - # As the number of requests grows, the blocks allocated by torch cannot be reused. - # So after piecewise cuda graph capture, a request with most requests is triggered to make - # sure that large enough blocks are allocated and can be correctly reused. - for num_tokens in piecewise_cuda_graph_num_tokens: + # The logits allocations grow with the number of requests and are not + # part of the captured model body. Warm up the largest request count so + # those allocations can be reused during stable inference. + for num_tokens in prefill_cuda_graph_num_tokens: warmup_request = self._create_warmup_request(resource_manager, num_tokens, 0, @@ -2496,11 +2522,20 @@ def _capture_piecewise_cuda_graphs(self, resource_manager: ResourceManager): if batch is None: continue logger.info( - f"Run piecewise CUDA graph warmup for num tokens={num_tokens} with most requests" + f"Run prefill CUDA graph warmup for num tokens={num_tokens} with most requests" ) - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) + if self.breakable_cuda_graph_runner is not None: + with self.no_cuda_graph(): + self.breakable_cuda_graph_runner.warmup( + lambda: self.forward( + batch, + new_tensors_device=None, + resource_manager=resource_manager), + steps=1) + else: + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) torch.cuda.synchronize() ### Helper methods promoted from the original warmup method ### @@ -3333,6 +3368,9 @@ def _release_cuda_graphs(self): if hasattr(self, 'cuda_graph_runner') and self.cuda_graph_runner is not None: self.cuda_graph_runner.clear() + if (hasattr(self, 'breakable_cuda_graph_runner') + and self.breakable_cuda_graph_runner is not None): + self.breakable_cuda_graph_runner.clear() if hasattr(self, 'encoder_cuda_graph_runner' ) and self.encoder_cuda_graph_runner is not None: self.encoder_cuda_graph_runner.clear() @@ -3516,45 +3554,41 @@ def _set_spec_metadata_all_rank_num_tokens( spec_metadata.subseq_all_rank_num_tokens = all_rank_num_seqs def _get_padding_params( - self, total_num_tokens: int, num_ctx_requests: int, - attn_all_rank_num_tokens: Optional[List[int]] + self, + total_num_tokens: int, + num_ctx_requests: int, + attn_all_rank_num_tokens: Optional[List[int]], ) -> Tuple[int, bool, Optional[List[int]]]: """ Get the padding parameters for tensor padding. Return: padded_num_tokens: the padded number of tokens - can_run_piecewise_cuda_graph: whether the piecewise cuda graph can be run + can_run_prefill_cuda_graph: whether a prefill CUDA graph can run attn_all_rank_num_tokens: the number of tokens for each rank """ - padded_num_tokens = total_num_tokens - all_rank_ctx_requests = self._get_all_rank_ctx_requests( num_ctx_requests) - def get_padded_piecewise_tokens(tokens): - captured_num_tokens = self._torch_compile_backend.capture_num_tokens - return captured_num_tokens[bisect.bisect_left( - captured_num_tokens, tokens)] - - if (self._torch_compile_backend is not None - and self._torch_compile_piecewise_cuda_graph - and self._torch_compile_backend.capture_num_tokens): - max_captured_num_tokens = self._torch_compile_backend.capture_num_tokens[ - -1] - # Torch piecewise cuda graph is enabled. + def get_padded_prefill_tokens(tokens: int) -> int: + return self._prefill_cuda_graph_num_tokens[bisect.bisect_left( + self._prefill_cuda_graph_num_tokens, tokens)] + + if (self.prefill_cuda_graph_backend + != PrefillCudaGraphBackend.DISABLED + and self._prefill_cuda_graph_num_tokens): + max_captured_num_tokens = self._prefill_cuda_graph_num_tokens[-1] if attn_all_rank_num_tokens is not None: - # Any rank has context requests, we enable piecewise cuda graph. has_ctx_requests = num_ctx_requests != 0 or ( all_rank_ctx_requests is not None and any(ctx_requests != 0 for ctx_requests in all_rank_ctx_requests)) - can_run_piecewise_cuda_graph = (has_ctx_requests and - max(attn_all_rank_num_tokens) - <= max_captured_num_tokens) - all_ranks_can_run_piecewise_cuda_graph = list( - self.dist.tp_allgather(can_run_piecewise_cuda_graph)) - if all(all_ranks_can_run_piecewise_cuda_graph): - padded_num_tokens = get_padded_piecewise_tokens( + can_run_prefill_cuda_graph = (has_ctx_requests + and max(attn_all_rank_num_tokens) + <= max_captured_num_tokens) + all_ranks_can_run_prefill_cuda_graph = list( + self.dist.tp_allgather(can_run_prefill_cuda_graph)) + if all(all_ranks_can_run_prefill_cuda_graph): + padded_num_tokens = get_padded_prefill_tokens( max(attn_all_rank_num_tokens)) logger.debug( f"Pad tensor with {total_num_tokens} tokens to {padded_num_tokens} tokens" @@ -3564,19 +3598,18 @@ def get_padded_piecewise_tokens(tokens): ] * len(attn_all_rank_num_tokens) else: logger.debug( - "Not all ranks can run piecewise cuda graph, disable piecewise cuda graph" + "Not all ranks can run prefill CUDA graph, disable prefill CUDA graph" ) return total_num_tokens, False, attn_all_rank_num_tokens elif num_ctx_requests != 0 and total_num_tokens <= max_captured_num_tokens: - padded_num_tokens = get_padded_piecewise_tokens( - total_num_tokens) + padded_num_tokens = get_padded_prefill_tokens(total_num_tokens) logger.debug( f"Pad tensor with {total_num_tokens} tokens to {padded_num_tokens} tokens" ) return padded_num_tokens, True, None else: logger.debug( - f"Piecewise CUDA graph cannot be used with {total_num_tokens} tokens, {num_ctx_requests} context requests" + f"Prefill CUDA graph cannot be used with {total_num_tokens} tokens, {num_ctx_requests} context requests" ) return total_num_tokens, False, None @@ -5778,10 +5811,22 @@ def previous_seq_slots_device(): lora_params = self._get_lora_params_from_requests( scheduled_requests, attn_metadata, peft_cache_manager, maybe_graph) + has_multimodal_input = any( + param.multimodal_data and any( + key != 'mrope_config' for key in param.multimodal_data) + for param in multimodal_params_list) attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) - padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( - total_num_tokens, num_ctx_requests, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + if (self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.BREAKABLE + and (bool(lora_params) or has_multimodal_input)): + padded_num_tokens = total_num_tokens + can_run_prefill_cuda_graph = False + else: + padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( + total_num_tokens, num_ctx_requests, attn_all_rank_num_tokens) + set_per_request_piecewise_cuda_graph_flag( + can_run_prefill_cuda_graph and self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.PIECEWISE) attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != total_num_tokens else None virtual_num_tokens = total_num_tokens @@ -6054,9 +6099,11 @@ def _prepare_tp_inputs_no_cache( attn_metadata.num_contexts = scheduled_requests.num_context_requests attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) - padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( + padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( num_tokens, attn_metadata.num_contexts, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + set_per_request_piecewise_cuda_graph_flag( + can_run_prefill_cuda_graph and self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.PIECEWISE) attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != num_tokens else None if self.enable_attention_dp: @@ -7037,7 +7084,6 @@ def forward(self, moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) - if kv_cache_manager is None: inputs, gather_ids = self._prepare_tp_inputs_no_cache( scheduled_requests, attn_metadata, spec_metadata, @@ -7154,14 +7200,42 @@ def forward(self, self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() + breakable_runner = self.breakable_cuda_graph_runner + has_multimodal_input = any( + param.multimodal_data and any( + key != 'mrope_config' for key in param.multimodal_data) + for param in inputs.get('multimodal_params', ())) + breakable_request_eligible = ( + breakable_runner is not None + and scheduled_requests.num_context_requests > 0 + and spec_metadata is None and not gather_context_logits + and not inputs.get('lora_params') + and not has_multimodal_input) + with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): - if not can_run_graph: - # Fallback to eager execution if graph was not used + + def forward_step(): with MoeLoadBalancerIterContext(moe_load_balancer): - outputs = self._forward_step( + return self._forward_step( inputs, gather_ids=gather_ids, gather_context_logits=gather_context_logits) + if not can_run_graph: + if (breakable_runner is not None + and breakable_runner.is_capturing): + return breakable_runner.capture_model_body( + forward_step) + + num_tokens = inputs['input_ids'].shape[0] + can_run_breakable_graph = ( + breakable_request_eligible + and breakable_runner.has_graph(num_tokens)) + if can_run_breakable_graph and not breakable_runner.is_warming_up: + outputs = breakable_runner.execute( + num_tokens, forward_step) + else: + # PCG or real eager + outputs = forward_step() else: needs_capture = self.cuda_graph_runner.needs_capture(key) if needs_capture: diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 2ddc301eb01a..aff2e9736977 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -20,12 +20,12 @@ MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, NGramDecodingConfig, PARDDecodingConfig, - PrometheusMetricsConfig, ReorderRequestPolicyConfig, - RocketSparseAttentionConfig, SADecodingConfig, - SAEnhancerConfig, SaveHiddenStatesDecodingConfig, - SchedulerConfig, SkipSoftmaxAttentionConfig, - TorchCompileConfig, TorchLlmArgs, - TriAttentionKvCacheCompressionConfig, + PrefillCudaGraphBackend, PrometheusMetricsConfig, + ReorderRequestPolicyConfig, RocketSparseAttentionConfig, + SADecodingConfig, SAEnhancerConfig, + SaveHiddenStatesDecodingConfig, SchedulerConfig, + SkipSoftmaxAttentionConfig, TorchCompileConfig, + TorchLlmArgs, TriAttentionKvCacheCompressionConfig, UserProvidedDecodingConfig) from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder @@ -93,6 +93,7 @@ 'SkipSoftmaxAttentionConfig', 'TriAttentionKvCacheCompressionConfig', 'PrometheusMetricsConfig', + 'PrefillCudaGraphBackend', 'ThinkingBudgetLogitsProcessor', 'add_thinking_budget_logits_processor', 'MultimodalConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1afb2cc25994..fa4d21505a33 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5009,6 +5009,18 @@ class SamplerType(StrEnum): auto = "auto" +class PrefillCudaGraphBackend(StrEnum): + """CUDA graph implementation used for prefill requests.""" + + DISABLED = "disabled" + PIECEWISE = "piecewise" + BREAKABLE = "breakable" + + +_DEFAULT_PREFILL_CAPTURE_NUM_TOKENS = [2**i for i in range(8) + ] + [i for i in range(256, 3073, 256)] + + class TorchCompileConfig(StrictBaseModel): """Configuration for torch.compile.""" enable_fullgraph: bool = Field( @@ -5048,8 +5060,7 @@ def validate_capture_num_tokens(cls, v): @model_validator(mode='after') def set_default_capture_num_tokens(self) -> 'TorchCompileConfig': if self.enable_piecewise_cuda_graph and self.capture_num_tokens is None: - self.capture_num_tokens = [2**i for i in range(8) - ] + [i for i in range(256, 3073, 256)] + self.capture_num_tokens = list(_DEFAULT_PREFILL_CAPTURE_NUM_TOKENS) return self @@ -5297,6 +5308,20 @@ def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': torch_compile_config: Optional[TorchCompileConfig] = Field( default=None, description="Torch compile config.", status="prototype") + prefill_cuda_graph_backend: PrefillCudaGraphBackend = Field( + default=PrefillCudaGraphBackend.DISABLED, + description="CUDA graph implementation used for prefill requests. " + "Defaults to disabled.", + status="prototype", + telemetry=TelemetryField.categorical("disabled", "piecewise", + "breakable")) + + prefill_capture_num_tokens: Optional[List[int]] = Field( + default=None, + description= + "Token-count buckets captured by the selected prefill CUDA graph implementation.", + status="prototype") + enable_autotuner: bool = Field( default=True, description= @@ -5593,6 +5618,60 @@ def validate_encode_only_torch_compile_config(self) -> 'TorchLlmArgs': "graphs or disable enable_piecewise_cuda_graph.") return self + @model_validator(mode="after") + def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': + """Normalize legacy piecewise CUDA graph options into prefill fields.""" + backend_is_explicit = "prefill_cuda_graph_backend" in self.model_fields_set + buckets_are_explicit = "prefill_capture_num_tokens" in self.model_fields_set + compile_config = self.torch_compile_config + + if compile_config is not None and compile_config.enable_piecewise_cuda_graph: + if (backend_is_explicit and self.prefill_cuda_graph_backend + != PrefillCudaGraphBackend.PIECEWISE): + raise ValueError( + "torch_compile_config.enable_piecewise_cuda_graph conflicts " + "with prefill_cuda_graph_backend") + logger.warning( + "TorchCompileConfig.enable_piecewise_cuda_graph is deprecated; " + "use prefill_cuda_graph_backend='piecewise' instead.") + self.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE + + legacy_buckets = (compile_config.capture_num_tokens + if compile_config is not None else None) + if legacy_buckets is not None: + if (buckets_are_explicit + and self.prefill_capture_num_tokens is not None + and sorted(set(legacy_buckets)) != sorted( + set(self.prefill_capture_num_tokens))): + raise ValueError( + "torch_compile_config.capture_num_tokens conflicts with " + "prefill_capture_num_tokens") + if not buckets_are_explicit: + logger.warning( + "TorchCompileConfig.capture_num_tokens is deprecated; use " + "prefill_capture_num_tokens instead.") + self.prefill_capture_num_tokens = list(legacy_buckets) + + if self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED: + if self.prefill_capture_num_tokens is None: + self.prefill_capture_num_tokens = list( + _DEFAULT_PREFILL_CAPTURE_NUM_TOKENS) + if self.encode_only: + raise ValueError( + "encode_only does not support prefill CUDA graphs") + + if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE: + if self.torch_compile_config is None: + self.torch_compile_config = TorchCompileConfig() + elif (self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.BREAKABLE + and self.torch_compile_config is not None): + raise ValueError( + "breakable prefill CUDA graph does not support torch_compile_config" + ) + + return self + @model_validator(mode="after") def validate_speculative_config(self): if self.speculative_config: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 38436f22514a..e16d33f7ea45 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1286,6 +1286,24 @@ "kind": "value", "path": "pp_partition" }, + { + "allowed_values": [], + "annotation": "Optional[List[int]]", + "converter": "", + "kind": "value", + "path": "prefill_capture_num_tokens" + }, + { + "allowed_values": [ + "disabled", + "piecewise", + "breakable" + ], + "annotation": "", + "converter": "allowlist", + "kind": "categorical", + "path": "prefill_cuda_graph_backend" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8b00d737c793..ba0fbbae8380 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -33,9 +33,10 @@ DFlashDecodingConfig, DSparkDecodingConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MambaStateConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, - NGramDecodingConfig, PARDDecodingConfig, RocketSparseAttentionConfig, - SADecodingConfig, SamplingParams, SchedulerConfig, - SkipSoftmaxAttentionConfig, SAEnhancerConfig, TorchCompileConfig) + NGramDecodingConfig, PARDDecodingConfig, PrefillCudaGraphBackend, + RocketSparseAttentionConfig, SADecodingConfig, SamplingParams, + SchedulerConfig, SkipSoftmaxAttentionConfig, SAEnhancerConfig, + TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -6201,6 +6202,46 @@ def test_bf16(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_blackwell + @pytest.mark.threadleak(enabled=False) + def test_bf16_breakable_prefill_cuda_graph(self): + model_path = f"{llm_models_root()}/Qwen3.5-4B" + prompts = [ + [[17] * 128], + [[17] * 129], + # The second request is admitted while the first is decoding, + # exercising BCG replay for a mixed context/decode batch. + [[17] * 64, [23] * 65], + [[31] * 256], + ] + sampling_params = SamplingParams(max_tokens=4) + + def run(backend): + results = [] + with LLM( + model_path, + trust_remote_code=True, + max_seq_len=1024, + max_num_tokens=512, + max_batch_size=4, + disable_overlap_scheduler=True, + kv_cache_config=self.kv_cache_config, + cuda_graph_config=CudaGraphConfig(enable_padding=True, + max_batch_size=4), + prefill_cuda_graph_backend=backend, + prefill_capture_num_tokens=[128, 256, 512], + ) as llm: + for batch in prompts: + results.append([ + output.outputs[0].token_ids for output in llm.generate( + batch, sampling_params=sampling_params) + ]) + return results + + eager_results = run(PrefillCudaGraphBackend.DISABLED) + breakable_results = run(PrefillCudaGraphBackend.BREAKABLE) + assert breakable_results == eager_results + @skip_pre_hopper def test_fp8(self): model_path = f"{llm_models_root()}/Qwen3.5-4B-FP8" diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py new file mode 100644 index 000000000000..957747429b02 --- /dev/null +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -0,0 +1,253 @@ +# Adapted from SGLang's breakable CUDA graph tests. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import weakref + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, + break_graph, + eager_on_graph, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph import ( + _copy_output, + _weak_ref_if_tensor, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import ( + BreakableCUDAGraphRunner, + BreakableCUDAGraphRunnerState, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _capture(body): + graph = BreakableCUDAGraph() + with BreakableCUDAGraphCapture(graph, stream=torch.cuda.Stream()): + body() + return graph + + +def test_no_break_capture_and_repeated_replay(): + x = torch.zeros(4, device="cuda") + output = torch.zeros_like(x) + graph = _capture(lambda: output.copy_(x + 1)) + + assert graph.num_segments == 1 + assert graph.num_breaks == 0 + for value in (5, 11): + x.fill_(value) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, value + 1)) + + +def test_single_and_multiple_breakpoints(): + @eager_on_graph(True) + def add_one(value): + return value + 1 + + @eager_on_graph(True) + def double(value): + return value * 2 + + x = torch.zeros(4, device="cuda") + output = torch.zeros_like(x) + + def body(): + value = add_one(x + 1) + value = double(value + 1) + output.copy_(value) + + graph = _capture(body) + assert graph.num_segments == 3 + assert graph.num_breaks == 2 + + x.fill_(5) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 16)) + + +def test_disabled_and_outside_capture(): + @eager_on_graph(False) + def disabled(value): + return value + 1 + + @eager_on_graph(True) + def outside(value): + return value + 2 + + value = torch.tensor([1.0, 2.0], device="cuda") + torch.testing.assert_close(disabled(value), value + 1) + torch.testing.assert_close(outside(value), value + 2) + + +def test_break_graph_inserts_empty_breakpoint(): + x = torch.zeros(4, device="cuda") + output = torch.zeros_like(x) + + def body(): + value = x + 1 + break_graph() + output.copy_(value + 2) + + graph = _capture(body) + assert graph.num_segments == 2 + assert graph.num_breaks == 1 + x.fill_(10) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 13)) + + +def test_output_writeback_for_tensor_dict_and_object(): + class Output: + def __init__(self, tensor, label): + self.tensor = tensor + self.label = label + + tensor = torch.zeros(4, device="cuda") + assert _copy_output(tensor, torch.full_like(tensor, 3)) is tensor + torch.testing.assert_close(tensor, torch.full_like(tensor, 3)) + + output_dict = {"value": torch.zeros(4, device="cuda")} + assert _copy_output(output_dict, {"value": torch.ones(4, device="cuda")}) is output_dict + torch.testing.assert_close(output_dict["value"], torch.ones(4, device="cuda")) + + output_object = Output(torch.zeros(4, device="cuda"), "old") + assert ( + _copy_output(output_object, Output(torch.full((4,), 2.0, device="cuda"), "new")) + is output_object + ) + torch.testing.assert_close(output_object.tensor, torch.full_like(output_object.tensor, 2)) + assert output_object.label == "new" + + +def test_tensor_capture_uses_non_owning_reference(): + tensor = torch.ones(4, device="cuda") + python_ref = weakref.ref(tensor) + non_owning = _weak_ref_if_tensor(tensor) + assert non_owning.data_ptr() == tensor.data_ptr() + del tensor + gc.collect() + assert python_ref() is None + + +def test_side_stream_is_joined_before_segment_end(): + x = torch.ones(4, device="cuda") + output = torch.zeros_like(x) + side_stream = torch.cuda.Stream() + + def body(): + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + output.copy_((x + 1) * 2) + + graph = _capture(body) + x.fill_(3) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 8)) + + +class _Body(nn.Module): + def __init__(self): + super().__init__() + self.forward_calls = 0 + + def forward(self, value): + self.forward_calls += 1 + return value + 1 + + +class _LogitsProcessor(nn.Module): + def __init__(self): + super().__init__() + self.forward_calls = 0 + + def forward(self, value): + self.forward_calls += 1 + return value * 2 + + +def test_runner_warmup_capture_execute_and_shared_output(): + body = _Body().cuda() + logits_processor = _LogitsProcessor().cuda() + runner = BreakableCUDAGraphRunner(body, logits_processor) + counters = {"outer": 0} + inputs = {} + + def engine_forward(): + counters["outer"] += 1 + if runner.is_capturing: + return runner.capture_model_body( + lambda: {"logits": logits_processor(body(inputs["value"]))} + ) + return {"logits": logits_processor(body(inputs["value"]))} + + inputs["value"] = torch.zeros((8, 4), device="cuda") + runner.capture(8, engine_forward) + first_shared_output = runner._shared_output + inputs["value"] = torch.zeros((4, 4), device="cuda") + runner.capture(4, engine_forward) + + assert runner.state == BreakableCUDAGraphRunnerState.IDLE + assert counters == {"outer": 6} + assert body.forward_calls == 6 + assert logits_processor.forward_calls == 4 + assert runner._shared_output is first_shared_output + + original_forward = body.forward + inputs["value"].fill_(3) + result = runner.execute(4, engine_forward) + torch.cuda.synchronize() + torch.testing.assert_close(result["logits"], torch.full((4, 4), 8.0, device="cuda")) + assert counters == {"outer": 7} + assert body.forward_calls == 6 + assert logits_processor.forward_calls == 5 + assert body.forward == original_forward + + +def test_runner_graph_miss_nested_execute_and_exception_recovery(): + body = _Body().cuda() + runner = BreakableCUDAGraphRunner(body, _LogitsProcessor().cuda()) + with pytest.raises(KeyError, match="No BCG captured"): + runner.execute(4, lambda: None) + + runner._graphs[4] = object() + runner._outputs[4] = torch.zeros(1, device="cuda") + original_forward = body.forward + + def nested(): + return runner.execute(4, lambda: None) + + with pytest.raises(RuntimeError, match="while runner is replay"): + runner.execute(4, nested) + assert runner.state == BreakableCUDAGraphRunnerState.IDLE + assert body.forward == original_forward + + def fail(): + raise ValueError("expected") + + with pytest.raises(ValueError, match="expected"): + runner.execute(4, fail) + assert runner.state == BreakableCUDAGraphRunnerState.IDLE + assert body.forward == original_forward + + +def test_runner_warmup_exception_restores_idle_state(): + runner = BreakableCUDAGraphRunner(_Body().cuda(), _LogitsProcessor().cuda()) + + def fail(): + raise ValueError("expected") + + with pytest.raises(ValueError, match="expected"): + runner.warmup(fail) + assert runner.state == BreakableCUDAGraphRunnerState.IDLE diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index c9b9a42388a7..30f3727f688f 100644 --- a/tests/unittest/api_stability/api_stability_core.py +++ b/tests/unittest/api_stability/api_stability_core.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + # autoflake: skip_file import copy import inspect @@ -29,7 +32,7 @@ from tensorrt_llm.llmapi import (CalibConfig, CompletionOutput, GuidedDecodingParams, QuantConfig, RequestOutput, SamplingParams) -from tensorrt_llm.llmapi.llm_args import SamplerType +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend, SamplerType from tensorrt_llm.llmapi.llm_utils import LlmArgs from tensorrt_llm.logger import Singleton from tensorrt_llm.sampling_params import LogprobMode diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 409eafe66895..c10d9123c4d7 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -231,6 +231,14 @@ methods: annotation: Optional[tensorrt_llm.llmapi.llm_args.TorchCompileConfig] default: null status: prototype + prefill_cuda_graph_backend: + annotation: tensorrt_llm.llmapi.llm_args.PrefillCudaGraphBackend + default: disabled + status: prototype + prefill_capture_num_tokens: + annotation: Optional[List[int]] + default: null + status: prototype enable_autotuner: annotation: bool default: True diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2fa8772fd073..7ef8c3124eba 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -48,7 +48,8 @@ MambaStateConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, - PeftCacheConfig, PybindMirror, + PeftCacheConfig, + PrefillCudaGraphBackend, PybindMirror, RayPlacementConfig, SkipSoftmaxAttentionConfig, SleepConfig, SpeculativeConfig, @@ -1957,6 +1958,86 @@ class TestPiecewiseCudaGraphCaptureDefaults: _EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS = [2**i for i in range(8)] + list( range(256, 3073, 256)) + def test_prefill_capture_num_tokens_uses_plain_int_list(self): + annotation = TorchLlmArgs.model_fields[ + "prefill_capture_num_tokens"].annotation + list_annotation = get_args(annotation)[0] + assert get_origin(list_annotation) is list + assert get_args(list_annotation) == (int, ) + + def test_breakable_uses_default_capture_buckets(self): + args = TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE) + assert args.prefill_capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert args.torch_compile_config is None + + def test_piecewise_new_config_enables_default_torch_compile(self): + args = TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + prefill_capture_num_tokens=[512, 128, 512]) + assert args.torch_compile_config == TorchCompileConfig() + assert args.prefill_capture_num_tokens == [512, 128, 512] + + def test_legacy_piecewise_config_maps_to_new_fields(self): + args = TorchLlmArgs(model=llama_model_path, + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True, + capture_num_tokens=[128, 256])) + assert args.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE + assert args.prefill_capture_num_tokens == [256, 128] + + def test_explicit_legacy_and_new_config_conflicts(self): + with pytest.raises(ValueError, match="conflicts"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True)) + + with pytest.raises(ValueError, match="conflicts"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + prefill_capture_num_tokens=[128], + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True, capture_num_tokens=[256])) + + def test_breakable_rejects_explicit_torch_compile(self): + with pytest.raises(ValueError, match="does not support"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + torch_compile_config=TorchCompileConfig()) + + def test_prefill_filter_sorts_dedupes_and_drops_nonpositive(self): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + _filter_prefill_capture_num_tokens + + kept, unrecordable = _filter_prefill_capture_num_tokens( + [256, 0, -1, 128, 256], + max_num_tokens=512, + max_batch_size=1, + max_seq_len=513, + ) + assert kept == [128, 256, 512] + assert unrecordable == [] + + @pytest.mark.parametrize("backend", [ + PrefillCudaGraphBackend.PIECEWISE, + PrefillCudaGraphBackend.BREAKABLE, + ]) + def test_piecewise_and_breakable_use_identical_padding(self, backend): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + PyTorchModelEngine + + engine = object.__new__(PyTorchModelEngine) + engine.enable_attention_dp = False + engine.prefill_cuda_graph_backend = backend + engine._prefill_cuda_graph_num_tokens = [128, 256, 512] + assert engine._get_padding_params(129, 1, None) == (256, True, None) + def test_torch_compile_config_capture_num_tokens_default_when_piecewise_enabled( self): """Default capture set is the powers-of-2 + 256-stride list. From 93eb4ce21630560a479081124954b36e7143679c Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:02:07 +0000 Subject: [PATCH 02/23] bcg v2, support more model, fix adp bug and dsv4 bug Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../_torch/compilation/piecewise_optimizer.py | 4 +- .../_torch/models/modeling_minimaxm3.py | 12 +++- tensorrt_llm/_torch/modules/attention.py | 10 +--- tensorrt_llm/_torch/modules/mla.py | 8 ++- .../breakable_cuda_graph.py | 24 +++----- .../pyexecutor/breakable_cuda_graph_runner.py | 10 +--- .../_torch/pyexecutor/model_engine.py | 55 +++++-------------- tensorrt_llm/_torch/utils.py | 10 ++-- tensorrt_llm/llmapi/llm_args.py | 2 +- .../executor/test_breakable_cuda_graph.py | 25 ++------- tests/unittest/llmapi/test_llm_args.py | 27 +++++++++ 11 files changed, 85 insertions(+), 102 deletions(-) diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py index 5472c674c448..18ae3d6ae87e 100644 --- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py +++ b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py @@ -12,7 +12,7 @@ from tensorrt_llm.llmapi.utils import enable_llm_debug from ..utils import (get_model_extra_attrs, - get_per_request_piecewise_cuda_graph_flag, + get_per_request_prefill_cuda_graph_flag, get_piecewise_cuda_graph_flag, make_weak_ref, set_piecewise_running) from .multi_stream.auto_multi_stream import multi_stream_schedule @@ -202,7 +202,7 @@ def __call__(self, *args): if (runtime_num_of_token is None or runtime_num_of_token not in self.entries or not get_piecewise_cuda_graph_flag() - or not get_per_request_piecewise_cuda_graph_flag()): + or not get_per_request_prefill_cuda_graph_flag()): return self.default_callable(*args) if self.is_first_runner or self.is_last_runner: diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index e1978a577a3b..9360b03297a2 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -63,6 +63,7 @@ ) from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..utils import ( ActivationType, AuxStreamType, @@ -661,6 +662,11 @@ def minimax_m3_attn_custom_op_inplace( ) +maybe_bcg_minimax_m3_attn_custom_op_inplace = eager_on_graph(True)( + minimax_m3_attn_custom_op_inplace +) + + class MiniMaxM3Attention(Attention): """M3 attention: dense (layers 0-2) or sparse (layers 3-59). @@ -1321,8 +1327,10 @@ def _forward_attention_core( output = q.new_empty( (q.shape[0], self.num_heads * self.head_dim), dtype=self.attn_activation_dtype ) - if self.register_to_config and is_torch_compiling(): - minimax_m3_attn_custom_op_inplace( + if self.register_to_config and ( + is_torch_compiling() or is_in_breakable_cuda_graph() + ): + maybe_bcg_minimax_m3_attn_custom_op_inplace( q, k, v, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 084272738786..1b197a750b49 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -119,7 +119,7 @@ def attn_custom_op_inplace( ) -breakable_attn_custom_op_inplace = eager_on_graph(True)(attn_custom_op_inplace) +maybe_bcg_attn_custom_op_inplace = eager_on_graph(True)(attn_custom_op_inplace) def _helix_zero_kv_mask( @@ -968,13 +968,11 @@ def forward_impl( if "mrope_position_deltas" in mrope_config: mrope_position_deltas = mrope_config["mrope_position_deltas"] - use_breakable_cuda_graph = (not is_torch_compiling() - and is_in_breakable_cuda_graph()) # Currently only TRTLLM and FLASHINFER support the custom inplace op. use_custom_inplace_op = ( self.register_to_config and (self.attn_backend == "TRTLLM" or self.attn_backend == "FLASHINFER") - and (is_torch_compiling() or use_breakable_cuda_graph) + and (is_torch_compiling() or is_in_breakable_cuda_graph()) and not self.is_marlin_enabled) if use_custom_inplace_op: @@ -982,9 +980,7 @@ def forward_impl( assert len(outputs) == 1 or len(outputs) == 2 output = outputs[0] output_sf = outputs[1] if len(outputs) == 2 else None - custom_op = (breakable_attn_custom_op_inplace if - use_breakable_cuda_graph else attn_custom_op_inplace) - custom_op( + maybe_bcg_attn_custom_op_inplace( q, k, v, diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 6bdf977d131e..83f5639041e3 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -44,6 +44,7 @@ from ..attention_backend.utils import create_attention from ..distributed import AllReduceParams from ..model_config import ModelConfig +from ..pyexecutor.breakable_cuda_graph import eager_on_graph from ..utils import ( AuxStreamType, Fp4QuantizedTensor, @@ -167,6 +168,9 @@ def mla_custom_op_inplace( ) +maybe_bcg_mla_custom_op_inplace = eager_on_graph(True)(mla_custom_op_inplace) + + def fp8_block_scaling_bmm_out( mat1: torch.Tensor, mat2_fp8: torch.Tensor, @@ -1733,7 +1737,7 @@ def _forward_custom_op( sparse_output_sf = attn_output[2] if isinstance(hidden_states, Fp4QuantizedTensor): - torch.ops.trtllm.mla_custom_op_inplace( + maybe_bcg_mla_custom_op_inplace( hidden_states.unquantized_hidden_states, position_ids, self.layer_idx_str, @@ -1745,7 +1749,7 @@ def _forward_custom_op( hidden_states.scaling_factor, ) else: - torch.ops.trtllm.mla_custom_op_inplace( + maybe_bcg_mla_custom_op_inplace( hidden_states, position_ids, self.layer_idx_str, diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index 8136f5e51459..fa74c35b0147 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -109,18 +109,6 @@ def _uninstall_wait_stream_hook() -> None: _original_wait_stream = None -def _weak_ref_if_tensor(value: Any) -> Any: - if torch.is_tensor(value): - return make_weak_ref(value) - if isinstance(value, tuple): - return tuple(_weak_ref_if_tensor(item) for item in value) - if isinstance(value, list): - return [_weak_ref_if_tensor(item) for item in value] - if isinstance(value, dict): - return {key: _weak_ref_if_tensor(item) for key, item in value.items()} - return value - - def _copy_output(destination: Any, source: Any) -> Any: if torch.is_tensor(destination) and torch.is_tensor(source): destination.copy_(source) @@ -174,9 +162,15 @@ def wrapper(*args, **kwargs): capture._end_current_segment() output = inner(*args, **kwargs) - captured_args = tuple(_weak_ref_if_tensor(arg) for arg in args) - captured_kwargs = {key: _weak_ref_if_tensor(value) for key, value in kwargs.items()} - captured_output = _weak_ref_if_tensor(output) + # 看下attn的参数 + def make_weak_ref_with_str_none(x): + if isinstance(x, (str, None)): + return x + return make_weak_ref(x) + + captured_args = tuple(make_weak_ref_with_str_none(arg) for arg in args) + captured_kwargs = {key: make_weak_ref_with_str_none(value) for key, value in kwargs.items()} + captured_output = make_weak_ref_with_str_none(output) def replay_fn() -> Any: new_output = inner(*captured_args, **captured_kwargs) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py index e2b27646c52a..5a2e024d9db7 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -28,9 +28,8 @@ class BreakableCUDAGraphRunner: _WARMUP_STEPS = 2 - def __init__(self, layer_model: nn.Module, logits_processor: nn.Module) -> None: + def __init__(self, layer_model: nn.Module) -> None: self.layer_model = layer_model - self.logits_processor = logits_processor self._graphs: dict[int, BreakableCUDAGraph] = {} self._outputs: dict[int, torch.Tensor] = {} self._memory_pool = None @@ -146,7 +145,6 @@ def capture_model_body(self, outer_forward: Callable[[], Any]) -> Any: raise RuntimeError("BCG body capture requested outside capture") original_body_forward = self.layer_model.forward - original_logits_forward = self.logits_processor.forward captured_output = None def capture_forward(*args, **kwargs): @@ -156,19 +154,13 @@ def capture_forward(*args, **kwargs): original_body_forward(*args, **kwargs)) return captured_output - def passthrough_forward(hidden_states, *args, **kwargs): - del args, kwargs - return hidden_states - self.layer_model.forward = capture_forward - self.logits_processor.forward = passthrough_forward try: outer_forward() if captured_output is None: raise RuntimeError("BCG capture did not execute the model body") return captured_output finally: - self.logits_processor.forward = original_logits_forward self.layer_model.forward = original_body_forward def replay(self, num_tokens: int) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 1a79e0155abc..22fbb4fa13b7 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -75,7 +75,8 @@ from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..utils import (get_model_extra_attrs, - set_per_request_piecewise_cuda_graph_flag, + get_per_request_prefill_cuda_graph_flag, + set_per_request_prefill_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner from .config_utils import is_mla @@ -926,10 +927,6 @@ def __init__( self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) self.breakable_cuda_graph_runner = None if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: - if self.spec_config is not None: - raise ValueError( - "breakable prefill CUDA graph does not support speculative decoding" - ) decoder_model = (self.model if isinstance( self.model, DecoderModelForCausalLM) else getattr( self.model, "llm", None)) @@ -937,8 +934,7 @@ def __init__( raise ValueError( "breakable prefill CUDA graph requires a decoder model body" ) - self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner( - decoder_model.model, decoder_model.logits_processor) + self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner(decoder_model.model) # Initialize CUDA Graph LoRA manager if LoRA is enabled self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None @@ -5811,22 +5807,11 @@ def previous_seq_slots_device(): lora_params = self._get_lora_params_from_requests( scheduled_requests, attn_metadata, peft_cache_manager, maybe_graph) - has_multimodal_input = any( - param.multimodal_data and any( - key != 'mrope_config' for key in param.multimodal_data) - for param in multimodal_params_list) attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) - if (self.prefill_cuda_graph_backend - == PrefillCudaGraphBackend.BREAKABLE - and (bool(lora_params) or has_multimodal_input)): - padded_num_tokens = total_num_tokens - can_run_prefill_cuda_graph = False - else: - padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( - total_num_tokens, num_ctx_requests, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag( - can_run_prefill_cuda_graph and self.prefill_cuda_graph_backend - == PrefillCudaGraphBackend.PIECEWISE) + (padded_num_tokens, can_run_prefill_cuda_graph, + attn_all_rank_num_tokens) = self._get_padding_params( + total_num_tokens, num_ctx_requests, attn_all_rank_num_tokens) + set_per_request_prefill_cuda_graph_flag(can_run_prefill_cuda_graph) attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != total_num_tokens else None virtual_num_tokens = total_num_tokens @@ -6101,9 +6086,7 @@ def _prepare_tp_inputs_no_cache( attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( num_tokens, attn_metadata.num_contexts, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag( - can_run_prefill_cuda_graph and self.prefill_cuda_graph_backend - == PrefillCudaGraphBackend.PIECEWISE) + set_per_request_prefill_cuda_graph_flag(can_run_prefill_cuda_graph) attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != num_tokens else None if self.enable_attention_dp: @@ -6585,6 +6568,7 @@ def _prepare_inputs( maybe_graph: bool = False, promoted_context_request_ids: frozenset[int] = frozenset() ) -> Tuple[Dict[str, Any], Optional[torch.Tensor]]: + set_per_request_prefill_cuda_graph_flag(False) if self.mapping is not None and 'cp_type' in self.mapping.cp_config: cp_type = self.mapping.cp_config['cp_type'] if CpType.STAR == cp_type: @@ -7201,16 +7185,6 @@ def forward(self, self._prepare_inputs_event.record() breakable_runner = self.breakable_cuda_graph_runner - has_multimodal_input = any( - param.multimodal_data and any( - key != 'mrope_config' for key in param.multimodal_data) - for param in inputs.get('multimodal_params', ())) - breakable_request_eligible = ( - breakable_runner is not None - and scheduled_requests.num_context_requests > 0 - and spec_metadata is None and not gather_context_logits - and not inputs.get('lora_params') - and not has_multimodal_input) with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): @@ -7221,20 +7195,19 @@ def forward_step(): gather_ids=gather_ids, gather_context_logits=gather_context_logits) if not can_run_graph: - if (breakable_runner is not None - and breakable_runner.is_capturing): - return breakable_runner.capture_model_body( - forward_step) + if (breakable_runner is not None and breakable_runner.is_capturing): + return breakable_runner.capture_model_body(forward_step) num_tokens = inputs['input_ids'].shape[0] can_run_breakable_graph = ( - breakable_request_eligible + breakable_runner is not None + and get_per_request_prefill_cuda_graph_flag() and breakable_runner.has_graph(num_tokens)) if can_run_breakable_graph and not breakable_runner.is_warming_up: outputs = breakable_runner.execute( num_tokens, forward_step) else: - # PCG or real eager + # real eager or BCG warmup or PCG outputs = forward_step() else: needs_capture = self.cuda_graph_runner.needs_capture(key) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index cb62ec99f76b..c65ecadddaf3 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -397,12 +397,14 @@ def piecewise_cuda_graph(enable: bool): set_piecewise_cuda_graph_flag(prev_enable) -def set_per_request_piecewise_cuda_graph_flag(enable: bool): - _global_attrs.per_request_piecewise_cuda_graph_flag = enable +def set_per_request_prefill_cuda_graph_flag(enable: bool): + """Set whether the current batch can use its prefill CUDA graph backend.""" + _global_attrs.per_request_prefill_cuda_graph_flag = enable -def get_per_request_piecewise_cuda_graph_flag() -> bool: - return getattr(_global_attrs, 'per_request_piecewise_cuda_graph_flag', True) +def get_per_request_prefill_cuda_graph_flag() -> bool: + """Return whether the current batch can use its prefill CUDA graph backend.""" + return getattr(_global_attrs, 'per_request_prefill_cuda_graph_flag', True) def create_lm_head_tp_mapping(mapping: Mapping, token_count: int) -> Mapping: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fa4d21505a33..0a2f12bed46a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5053,7 +5053,7 @@ def validate_capture_num_tokens(cls, v): "When torch compile is enabled, userbuffers is enabled by default.") max_num_streams: PositiveInt = Field( - default=1, + default=3, description= "The maximum number of CUDA streams to use for torch.compile.") diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index 957747429b02..966bc252f4b4 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -15,10 +15,7 @@ break_graph, eager_on_graph, ) -from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph import ( - _copy_output, - _weak_ref_if_tensor, -) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph import _copy_output from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import ( BreakableCUDAGraphRunner, BreakableCUDAGraphRunnerState, @@ -130,16 +127,6 @@ def __init__(self, tensor, label): assert output_object.label == "new" -def test_tensor_capture_uses_non_owning_reference(): - tensor = torch.ones(4, device="cuda") - python_ref = weakref.ref(tensor) - non_owning = _weak_ref_if_tensor(tensor) - assert non_owning.data_ptr() == tensor.data_ptr() - del tensor - gc.collect() - assert python_ref() is None - - def test_side_stream_is_joined_before_segment_end(): x = torch.ones(4, device="cuda") output = torch.zeros_like(x) @@ -180,7 +167,7 @@ def forward(self, value): def test_runner_warmup_capture_execute_and_shared_output(): body = _Body().cuda() logits_processor = _LogitsProcessor().cuda() - runner = BreakableCUDAGraphRunner(body, logits_processor) + runner = BreakableCUDAGraphRunner(body) counters = {"outer": 0} inputs = {} @@ -201,7 +188,7 @@ def engine_forward(): assert runner.state == BreakableCUDAGraphRunnerState.IDLE assert counters == {"outer": 6} assert body.forward_calls == 6 - assert logits_processor.forward_calls == 4 + assert logits_processor.forward_calls == 6 assert runner._shared_output is first_shared_output original_forward = body.forward @@ -211,13 +198,13 @@ def engine_forward(): torch.testing.assert_close(result["logits"], torch.full((4, 4), 8.0, device="cuda")) assert counters == {"outer": 7} assert body.forward_calls == 6 - assert logits_processor.forward_calls == 5 + assert logits_processor.forward_calls == 7 assert body.forward == original_forward def test_runner_graph_miss_nested_execute_and_exception_recovery(): body = _Body().cuda() - runner = BreakableCUDAGraphRunner(body, _LogitsProcessor().cuda()) + runner = BreakableCUDAGraphRunner(body) with pytest.raises(KeyError, match="No BCG captured"): runner.execute(4, lambda: None) @@ -243,7 +230,7 @@ def fail(): def test_runner_warmup_exception_restores_idle_state(): - runner = BreakableCUDAGraphRunner(_Body().cuda(), _LogitsProcessor().cuda()) + runner = BreakableCUDAGraphRunner(_Body().cuda()) def fail(): raise ValueError("expected") diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 7ef8c3124eba..cf166601e2f8 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2038,6 +2038,33 @@ def test_piecewise_and_breakable_use_identical_padding(self, backend): engine._prefill_cuda_graph_num_tokens = [128, 256, 512] assert engine._get_padding_params(129, 1, None) == (256, True, None) + def test_attention_dp_prefill_graph_uses_all_rank_decision(self): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + PyTorchModelEngine + + class FakeDist: + def __init__(self, decisions): + self.decisions = decisions + + def tp_allgather(self, value): + del value + return self.decisions + + engine = object.__new__(PyTorchModelEngine) + engine.enable_attention_dp = True + engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.BREAKABLE + engine._prefill_cuda_graph_num_tokens = [128, 256, 512] + engine._get_all_rank_ctx_requests = lambda _: [0, 1, 0, 0] + + all_rank_num_tokens = [1, 129, 1, 1] + engine.dist = FakeDist([True, True, True, True]) + assert engine._get_padding_params( + 1, 0, all_rank_num_tokens) == (256, True, [256] * 4) + + engine.dist = FakeDist([True, False, True, True]) + assert engine._get_padding_params( + 1, 0, all_rank_num_tokens) == (1, False, all_rank_num_tokens) + def test_torch_compile_config_capture_num_tokens_default_when_piecewise_enabled( self): """Default capture set is the powers-of-2 + 256-stride list. From e2bc3a5e0efca13af57135b18cb9530ffaf031f0 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:03:27 -0700 Subject: [PATCH 03/23] [07/23/14:03] share BCG pool across segments Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../breakable_cuda_graph.py | 2 +- .../pyexecutor/breakable_cuda_graph_runner.py | 17 ++++++++-- .../executor/test_breakable_cuda_graph.py | 31 +++++++++++++++++++ 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index fa74c35b0147..28fa114cda27 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -164,7 +164,7 @@ def wrapper(*args, **kwargs): # 看下attn的参数 def make_weak_ref_with_str_none(x): - if isinstance(x, (str, None)): + if isinstance(x, (str, type(None))): return x return make_weak_ref(x) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py index 5a2e024d9db7..807d3b722505 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -14,6 +14,7 @@ BreakableCUDAGraphCapture, enable_breakable_cuda_graph, ) +from .trace_log_utils import log_mem_snapshot class BreakableCUDAGraphRunnerState(Enum): @@ -76,10 +77,20 @@ def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None: current_stream = torch.cuda.current_stream() self._capture_stream.wait_stream(current_stream) graph = None + created_memory_pool = False + log_mem_snapshot(f"bcg/before_capture_{num_tokens}") try: with torch.cuda.stream(self._capture_stream): self.warmup(engine_forward) + # Every segment in the first BCG bucket must receive the same + # explicit pool handle. Passing None lets each CUDAGraph create + # its own private pool, which multiplies the model workspace by + # the number of eager breaks. + if self._memory_pool is None: + self._memory_pool = torch.cuda.graph_pool_handle() + created_memory_pool = True + self._state = BreakableCUDAGraphRunnerState.CAPTURE graph = BreakableCUDAGraph() self._active_graph = graph @@ -94,11 +105,13 @@ def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None: assert graph is not None self._graphs[num_tokens] = graph self._outputs[num_tokens] = make_weak_ref(output) - if self._memory_pool is None: - self._memory_pool = graph.pool() + log_mem_snapshot(f"bcg/after_capture_{num_tokens}") except Exception: if graph is not None: graph.reset() + if created_memory_pool and not self._graphs: + self._memory_pool = None + log_mem_snapshot(f"bcg/capture_failed_{num_tokens}") raise finally: self._active_graph = None diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index 966bc252f4b4..e34bcab1b15a 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -202,6 +202,37 @@ def engine_forward(): assert body.forward == original_forward +def test_runner_first_bucket_segments_share_one_memory_pool(): + class BreakableBody(nn.Module): + @eager_on_graph(True) + def eager_add_one(self, value): + return value + 1 + + @eager_on_graph(True) + def eager_double(self, value): + return value * 2 + + def forward(self, value): + value = self.eager_add_one(value + 1) + return self.eager_double(value + 1) + + body = BreakableBody().cuda() + runner = BreakableCUDAGraphRunner(body) + inputs = {"value": torch.zeros((8, 4), device="cuda")} + + def engine_forward(): + if runner.is_capturing: + return runner.capture_model_body(lambda: body(inputs["value"])) + return body(inputs["value"]) + + runner.capture(8, engine_forward) + + graph = runner._graphs[8] + assert graph.num_segments == 3 + assert runner._memory_pool is not None + assert all(segment.pool() == runner._memory_pool for segment in graph._segments) + + def test_runner_graph_miss_nested_execute_and_exception_recovery(): body = _Body().cuda() runner = BreakableCUDAGraphRunner(body) From 0a983663c2f12c8f1c47d98a742dbd20bf3d8f6e Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Thu, 23 Jul 2026 01:37:01 -0700 Subject: [PATCH 04/23] [07/23/16:37] benchmark DSV4 1P1D AAAgent with BCG Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../sparse/deepseek_v4/module.py | 6 +++++ .../_torch/modules/test_mla_registry.py | 24 ++++++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 7d0ee04aff18..2d1bcbf543d9 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -18,6 +18,7 @@ ) from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph from tensorrt_llm._torch.utils import AuxStreamType from tensorrt_llm._utils import get_sm_version, is_sm_100f @@ -181,6 +182,11 @@ def _should_use_dsv4_epilogue_fusion() -> bool: num_generations = attn_metadata.num_generations if self._disable_dsv4_epilogue_fusion: return False + # BCG eager breaks replay with dynamic, unpadded attention metadata, + # while captured segment tensors use static bucket shapes. The fused + # epilogue buffers cannot safely bridge that shape boundary. + if is_in_breakable_cuda_graph(): + return False if num_contexts == 0 and num_generations == 0: return False if num_contexts > 0 and num_generations > 0: diff --git a/tests/unittest/_torch/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py index 908069a02560..3811d05a3490 100644 --- a/tests/unittest/_torch/modules/test_mla_registry.py +++ b/tests/unittest/_torch/modules/test_mla_registry.py @@ -13,12 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import Mock, patch import torch from torch import nn from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module import ( + prepare_sparse_attn_outputs, +) from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.mla import MLA from tensorrt_llm.functional import PositionEmbeddingType @@ -80,3 +84,21 @@ def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None: assert registry["0"]() is target_mla assert registry["0_0"]() is draft_mla assert registry["0_1"]() is next_mla + + +def test_dsv4_epilogue_fusion_is_disabled_inside_breakable_graph() -> None: + metadata = SimpleNamespace(num_contexts=1, num_generations=0, num_tokens=5) + mla_layer = Mock(spec=MLA) + mla_layer._disable_dsv4_epilogue_fusion = False + mla_layer.create_output.return_value = torch.empty(8, 8) + hidden_states = torch.empty(8, 8) + + with patch( + "tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module." + "is_in_breakable_cuda_graph", + return_value=True, + ): + outputs = prepare_sparse_attn_outputs(mla_layer, hidden_states, metadata) + + assert outputs == [mla_layer.create_output.return_value] + mla_layer.create_output.assert_called_once_with(hidden_states, 1) From c846c64f99917e35f4952ff8e8cd51e1e66feb14 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Fri, 24 Jul 2026 05:09:59 -0700 Subject: [PATCH 05/23] [07/24/20:09] refactor DSv4 BCG epilogue fusion Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../sparse/deepseek_v4/module.py | 195 ++++++++++-------- .../sparse/dsa/custom_ops.py | 4 + .../attention_backend/sparse/dsa/module.py | 2 +- tensorrt_llm/_torch/modules/mla.py | 27 +-- .../defs/accuracy/test_llm_api_pytorch.py | 177 ++++++++++++++++ .../_torch/modules/test_mla_registry.py | 165 ++++++++++++++- 6 files changed, 452 insertions(+), 118 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 2d1bcbf543d9..c0e450f15a1e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -18,7 +18,6 @@ ) from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding -from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph from tensorrt_llm._torch.utils import AuxStreamType from tensorrt_llm._utils import get_sm_version, is_sm_100f @@ -162,18 +161,69 @@ def create_sparse_attn_weights(self) -> None: # Fused epilogue buffer management and output projection. -def _validate_dsv4_epilogue_buffers( +def _create_dsv4_epilogue_buffers( self, + q: torch.Tensor, num_tokens: int, - dsv4_epilogue_output: tuple[torch.Tensor, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: - fp8_o, output_sf = dsv4_epilogue_output + if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: + raise ValueError( + "DSv4 fused epilogue requires num_heads_tp to be divisible by n_local_groups." + ) + heads_per_group = self.num_heads_tp // self.n_local_groups scale_buf_m = (num_tokens + 3) // 4 * 4 - if fp8_o.shape[1] != num_tokens or output_sf.shape[2] != scale_buf_m: - raise RuntimeError("Invalid DSv4 fused epilogue buffers for current token count.") + fp8_o = q.new_empty( + (self.n_local_groups, num_tokens, heads_per_group * self.v_head_dim), + dtype=torch.float8_e4m3fn, + ) + output_sf = q.new_empty( + ( + self.n_local_groups, + heads_per_group * (self.v_head_dim // 128), + scale_buf_m, + ), + dtype=torch.float32, + ) return fp8_o, output_sf +def _run_dsv4_epilogue_bmm( + self, + epilogue_output: tuple[torch.Tensor, torch.Tensor], + output: torch.Tensor, +) -> None: + attn_fp8, attn_scale = epilogue_output + torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( + attn_fp8, + self.o_a_proj, + attn_scale, + self.o_a_proj_scale, + output.transpose(0, 1), + ) + + +def _run_dsv4_epilogue_bmms( + self, + output: torch.Tensor, + num_context_tokens: int, + num_tokens: int, + context_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]], + generation_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]], +) -> None: + if context_epilogue_output is not None: + _run_dsv4_epilogue_bmm( + self, + context_epilogue_output, + output[:num_context_tokens], + ) + if generation_epilogue_output is not None: + _run_dsv4_epilogue_bmm( + self, + generation_epilogue_output, + output[num_context_tokens:num_tokens], + ) + + def prepare_sparse_attn_outputs( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata ) -> list[torch.Tensor]: @@ -182,16 +232,8 @@ def _should_use_dsv4_epilogue_fusion() -> bool: num_generations = attn_metadata.num_generations if self._disable_dsv4_epilogue_fusion: return False - # BCG eager breaks replay with dynamic, unpadded attention metadata, - # while captured segment tensors use static bucket shapes. The fused - # epilogue buffers cannot safely bridge that shape boundary. - if is_in_breakable_cuda_graph(): - return False if num_contexts == 0 and num_generations == 0: return False - if num_contexts > 0 and num_generations > 0: - # The fused buffers do not carry token offsets for a mixed batch. - return False if self.mapping.has_cp_helix() or not is_sm_100f(): return False if not getattr(self.mapping, "enable_attention_dp", False): @@ -212,32 +254,15 @@ def _should_use_dsv4_epilogue_fusion() -> bool: return False return not self.inverse_rotary_emb.is_neox - def _create_dsv4_epilogue_buffers() -> tuple[torch.Tensor, torch.Tensor]: - if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: - raise ValueError( - "DSv4 fused epilogue requires num_heads_tp to be divisible by n_local_groups." - ) - heads_per_group = self.num_heads_tp // self.n_local_groups - num_tokens = attn_metadata.num_tokens - scale_buf_m = (num_tokens + 3) // 4 * 4 - fp8_o = hidden_states.new_empty( - (self.n_local_groups, num_tokens, heads_per_group * self.v_head_dim), - dtype=torch.float8_e4m3fn, - ) - output_sf = hidden_states.new_empty( - ( - self.n_local_groups, - heads_per_group * (self.v_head_dim // 128), - scale_buf_m, - ), - dtype=torch.float32, - ) - return fp8_o, output_sf - if _should_use_dsv4_epilogue_fusion(): - attn_output = [self.create_output(hidden_states[:0], attn_metadata.num_contexts)] - attn_output.extend(_create_dsv4_epilogue_buffers()) - return attn_output + num_tokens = hidden_states.shape[0] + return [ + torch.empty( + [num_tokens, self.n_local_groups, self.o_lora_rank], + device=hidden_states.device, + dtype=self.dtype, + ) + ] return [self.create_output(hidden_states, attn_metadata.num_contexts)] @@ -249,24 +274,10 @@ def project_sparse_attn_output( all_reduce_params: Optional["AllReduceParams"] = None, ) -> torch.Tensor: del attn_metadata, all_reduce_params - if len(attn_output) > 1: - attn_fp8, attn_scale = attn_output[1:] - num_tokens = attn_fp8.shape[1] - o_lora = torch.empty( - [num_tokens, self.n_local_groups, self.o_lora_rank], - device=attn_fp8.device, - dtype=self.dtype, - ) - torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( - attn_fp8, - self.o_a_proj, - attn_scale, - self.o_a_proj_scale, - o_lora.transpose(0, 1), - ) - return self.o_b_proj(o_lora.flatten(1)) - attn_output_tensor = attn_output[0] + if attn_output_tensor.ndim == 3: + return self.o_b_proj(attn_output_tensor.flatten(1)) + assert position_ids is not None num_tokens = attn_output_tensor.shape[0] attn_output_tensor = attn_output_tensor.view(num_tokens, self.num_heads_tp, -1) @@ -353,11 +364,11 @@ def forward_generation_sparse_attn( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, - sparse_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + enable_dsv4_epilogue_fusion: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Run the DeepSeek-V4 generation absorption path.""" if get_sm_version() < 100: @@ -402,10 +413,8 @@ def forward_generation_sparse_attn( attention_output = output output_sf = None inverse_rope_cos_sin = None - if sparse_epilogue_output is not None: - attention_output, output_sf = _validate_dsv4_epilogue_buffers( - self, num_tokens, sparse_epilogue_output - ) + if enable_dsv4_epilogue_fusion: + attention_output, output_sf = _create_dsv4_epilogue_buffers(self, q, num_tokens) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -429,11 +438,13 @@ def forward_generation_sparse_attn( mla_bmm2_scale=mla_bmm2_scale, quant_q_buffer=quant_q_buffer, dsv4_inv_rope_cos_sin_cache=inverse_rope_cos_sin, - enable_dsv4_epilogue_fusion=sparse_epilogue_output is not None, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) - if sparse_epilogue_output is not None: - return attn_out_latent + if enable_dsv4_epilogue_fusion: + assert attention_output is not None and output_sf is not None + return attention_output, output_sf + assert output is not None if self.mapping.has_cp_helix(): raise RuntimeError( "DeepSeek-V4 + CP Helix is not supported because the post-process " @@ -451,11 +462,11 @@ def forward_context_sparse_attn( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, - sparse_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + enable_dsv4_epilogue_fusion: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Run the DeepSeek-V4 context absorption path.""" if get_sm_version() < 100: @@ -484,10 +495,8 @@ def forward_context_sparse_attn( attention_output = output output_sf = None inverse_rope_cos_sin = None - if sparse_epilogue_output is not None: - attention_output, output_sf = _validate_dsv4_epilogue_buffers( - self, num_tokens, sparse_epilogue_output - ) + if enable_dsv4_epilogue_fusion: + attention_output, output_sf = _create_dsv4_epilogue_buffers(self, q, num_tokens) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -507,13 +516,16 @@ def forward_context_sparse_attn( quant_scale_qkv=quant_scale_qkv, sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), dsv4_inv_rope_cos_sin_cache=inverse_rope_cos_sin, - enable_dsv4_epilogue_fusion=sparse_epilogue_output is not None, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) self._fused_quant_q_buffer = None self._fused_q_pe = None - if sparse_epilogue_output is not None: - return attn_out_latent + if enable_dsv4_epilogue_fusion: + assert attention_output is not None and output_sf is not None + return attention_output, output_sf + + assert output is not None if self.mapping.has_cp_helix(): raise RuntimeError( "DeepSeek-V4 + CP Helix is not supported because the post-process " @@ -538,14 +550,11 @@ def forward_sparse_attn( """Run DeepSeek-V4 MLA and write into the algorithm-defined output buffers.""" assert self.mha is None and self.mqa is not None, "DeepSeek-V4 is only supported in MQA mode" output = attn_output[0] - sparse_epilogue_output = (attn_output[1], attn_output[2]) if len(attn_output) > 1 else None + enable_dsv4_epilogue_fusion = output.ndim == 3 num_contexts = attn_metadata.num_contexts num_generations = attn_metadata.num_generations num_ctx_tokens = attn_metadata.num_ctx_tokens num_tokens = attn_metadata.num_tokens - if sparse_epilogue_output is not None and ((num_contexts > 0) == (num_generations > 0)): - raise RuntimeError("DSv4 epilogue fusion requires a context-only or generation-only batch.") - hidden_states = hidden_states[:num_tokens, ...] if position_ids is not None: position_ids = position_ids[..., :num_tokens] @@ -756,6 +765,8 @@ def _indexer_branch(): assert output is not None, "output must be provided" + context_o_lora_bmm_input = None + generation_o_lora_bmm_input = None if num_contexts > 0: q_ctx = q[:num_ctx_tokens, ...] topk_indices_ctx = topk_indices[:num_ctx_tokens, :] if topk_indices is not None else None @@ -767,17 +778,17 @@ def _indexer_branch(): assert ctx_position_ids is not None k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) - forward_context_sparse_attn( + context_o_lora_bmm_input = forward_context_sparse_attn( self, q_ctx, compressed_kv_ctx, k_pe_ctx, attn_metadata, - output[:num_ctx_tokens, :], + None if enable_dsv4_epilogue_fusion else output[:num_ctx_tokens, :], position_ids=ctx_position_ids, latent_cache=latent_cache_ctx, topk_indices=topk_indices_ctx, - sparse_epilogue_output=sparse_epilogue_output, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) if num_generations > 0: @@ -795,17 +806,33 @@ def _indexer_branch(): assert gen_position_ids is not None k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) - forward_generation_sparse_attn( + generation_o_lora_bmm_input = forward_generation_sparse_attn( self, q_gen, compressed_kv_gen, k_pe_gen, attn_metadata, - output[num_ctx_tokens:num_tokens, :], + None if enable_dsv4_epilogue_fusion else output[num_ctx_tokens:num_tokens, :], position_ids=gen_position_ids, latent_cache=latent_cache_gen, topk_indices=topk_indices_gen, - sparse_epilogue_output=sparse_epilogue_output, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, + ) + + if enable_dsv4_epilogue_fusion: + assert context_o_lora_bmm_input is None or isinstance( + context_o_lora_bmm_input, tuple + ) + assert generation_o_lora_bmm_input is None or isinstance( + generation_o_lora_bmm_input, tuple + ) + _run_dsv4_epilogue_bmms( + self, + output, + num_ctx_tokens, + num_tokens, + context_o_lora_bmm_input, + generation_o_lora_bmm_input, ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py index 67c912fba56a..60c56cc22b2d 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py @@ -7,6 +7,7 @@ import torch +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph from tensorrt_llm._torch.utils import Fp4QuantizedTensor from .module import _forward_dsa_attn, forward_dsa_proj @@ -146,3 +147,6 @@ def _mla_dsa_attn_inplace_fake( output: torch.Tensor, ) -> None: """Model the in-place output mutation during fake-tensor propagation.""" + + +maybe_bcg_mla_dsa_attn_inplace = eager_on_graph(True)(mla_dsa_attn_inplace) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py index 5f07977a505d..5202067f14a6 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py @@ -198,7 +198,7 @@ def forward_sparse_attn_custom_op( ) q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] indexer_intermediates = proj_outputs[4:] - torch.ops.trtllm.mla_dsa_attn_inplace( + custom_ops.maybe_bcg_mla_dsa_attn_inplace( q, compressed_kv, k_pe, diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 83f5639041e3..50e1336cedf8 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -129,7 +129,7 @@ def _create_mla_outputs_fake(hidden_states, layer_idx): @torch.library.custom_op( "trtllm::mla_custom_op_inplace", - mutates_args=("output", "sparse_output", "sparse_output_sf"), + mutates_args=("output",), ) def mla_custom_op_inplace( hidden_states: torch.Tensor, @@ -137,8 +137,6 @@ def mla_custom_op_inplace( layer_idx: str, output: torch.Tensor, latent_cache_gen: Optional[torch.Tensor], - sparse_output: Optional[torch.Tensor], - sparse_output_sf: Optional[torch.Tensor], hidden_states_fp4: Optional[torch.Tensor] = None, hidden_states_sf: Optional[torch.Tensor] = None, ) -> None: @@ -152,18 +150,11 @@ def mla_custom_op_inplace( scaling_factor=hidden_states_sf, unquantized_hidden_states=hidden_states, ) - attn_output = [output] - if sparse_output is not None: - attn_output.append(sparse_output) - if sparse_output_sf is not None: - if sparse_output is None: - raise RuntimeError("sparse_output_sf requires sparse_output") - attn_output.append(sparse_output_sf) mla_layer.forward_impl( position_ids, hidden_states, metadata, - attn_output=attn_output, + attn_output=[output], latent_cache_gen=latent_cache_gen, ) @@ -1727,14 +1718,8 @@ def _forward_custom_op( return output = attn_output[0] - sparse_output = None - sparse_output_sf = None - if len(attn_output) > 3: - raise RuntimeError("MLA output hooks may return at most two sparse output buffers.") - if len(attn_output) > 1: - sparse_output = attn_output[1] - if len(attn_output) > 2: - sparse_output_sf = attn_output[2] + if len(attn_output) != 1: + raise RuntimeError("MLA custom ops require exactly one output tensor.") if isinstance(hidden_states, Fp4QuantizedTensor): maybe_bcg_mla_custom_op_inplace( @@ -1743,8 +1728,6 @@ def _forward_custom_op( self.layer_idx_str, output, latent_cache_gen, - sparse_output, - sparse_output_sf, hidden_states.fp4_tensor, hidden_states.scaling_factor, ) @@ -1755,8 +1738,6 @@ def _forward_custom_op( self.layer_idx_str, output, latent_cache_gen, - sparse_output, - sparse_output_sf, ) def _project_output( diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ba0fbbae8380..2ad2f516b8af 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -15,13 +15,16 @@ import asyncio import json import os +import statistics import sys +import time from unittest import mock import pytest import torch from datasets import load_dataset from defs.conftest import get_sm_version, is_sm_100f +from mpi4py import MPI from mpi4py.futures import MPIPoolExecutor from tensorrt_llm import LLM @@ -3934,6 +3937,180 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): eplb_config, mtp_nextn=mtp_nextn) + @pytest.mark.skip_less_mpi_world_size(8) + @pytest.mark.threadleak(enabled=False) + def test_mixed_breakable_cuda_graph_epilogue_fusion_ab(self, mocker): + from transformers import AutoTokenizer + + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + fusion_env = "TRTLLM_DSV4_DISABLE_FMHA_EPILOGUE_FUSION" + + def patched_start_mpi_pool(session): + assert not session.mpi_pool, "MPI session already started" + session.mpi_pool = MPIPoolExecutor( + max_workers=session.n_workers, + path=sys.path, + env={fusion_env: os.environ.get(fusion_env, "0")}, + ) + + mocker.patch.object(MpiPoolSession, "_start_mpi_pool", + patched_start_mpi_pool) + + prompt_lengths = [64, 129, 257, 385, 513, 769] + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + base_prompt_ids = tokenizer.encode( + "TensorRT-LLM accelerates reliable large language model inference " + "with efficient attention, parallelism, and CUDA graphs. ", + add_special_tokens=False, + ) + assert base_prompt_ids + prompts = [ + (base_prompt_ids * ((prompt_length + len(base_prompt_ids) - 1) // + len(base_prompt_ids)))[:prompt_length] + for prompt_length in prompt_lengths + ] + sampling_params = SamplingParams( + max_tokens=8, + min_tokens=8, + seed=42, + temperature=0, + ignore_eos=True, + detokenize=False, + add_special_tokens=False, + ) + llm_kwargs = dict( + tensor_parallel_size=8, + moe_expert_parallel_size=8, + moe_config=MoeConfig(backend="TRTLLM"), + enable_attention_dp=True, + max_batch_size=8, + max_num_tokens=1024, + max_seq_len=2048, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + dtype="fp8", + free_gpu_memory_fraction=0.6, + ), + cuda_graph_config=CudaGraphConfig( + batch_sizes=[1, 2, 4, 6, 8], + enable_padding=True, + ), + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + prefill_capture_num_tokens=[128, 256, 512, 1024], + ) + + def run_variant(disable_fusion: bool) -> dict: + variant_start = time.perf_counter() + with mock.patch.dict( + os.environ, + {fusion_env: "1" if disable_fusion else "0"}, + clear=False, + ): + with LLM( + self.MODEL_PATH, + **llm_kwargs, + env_overrides={ + fusion_env: "1" if disable_fusion else "0" + }, + ) as llm: + init_seconds = time.perf_counter() - variant_start + warmup_start = time.perf_counter() + llm.generate(prompts, + sampling_params=sampling_params, + use_tqdm=False) + warmup_seconds = time.perf_counter() - warmup_start + + rounds = [] + for _ in range(5): + round_start = time.perf_counter() + outputs = llm.generate( + prompts, + sampling_params=sampling_params, + use_tqdm=False, + ) + latency_seconds = time.perf_counter() - round_start + token_ids = [ + output.outputs[0].token_ids for output in outputs + ] + output_tokens = sum(len(ids) for ids in token_ids) + rounds.append({ + "latency_seconds": + latency_seconds, + "output_tokens": + output_tokens, + "output_tokens_per_second": + output_tokens / latency_seconds, + "token_ids": + token_ids, + }) + + latencies = [ + round_result["latency_seconds"] for round_result in rounds + ] + throughputs = [ + round_result["output_tokens_per_second"] + for round_result in rounds + ] + return { + "fusion_disabled": + disable_fusion, + "engine_init_capture_seconds": + init_seconds, + "warmup_seconds": + warmup_seconds, + "rounds": + rounds, + "median_latency_seconds": + statistics.median(latencies), + "p90_latency_seconds": + statistics.quantiles(latencies, n=10, method="inclusive")[8], + "median_output_tokens_per_second": + statistics.median(throughputs), + "p90_output_tokens_per_second": + statistics.quantiles(throughputs, n=10, method="inclusive")[8], + } + + disabled_result = run_variant(disable_fusion=True) + fusion_result = run_variant(disable_fusion=False) + disabled_token_ids = [ + round_result["token_ids"] + for round_result in disabled_result["rounds"] + ] + fusion_token_ids = [ + round_result["token_ids"] + for round_result in fusion_result["rounds"] + ] + disabled_repeatable = all(token_ids == disabled_token_ids[0] + for token_ids in disabled_token_ids[1:]) + fusion_repeatable = all(token_ids == fusion_token_ids[0] + for token_ids in fusion_token_ids[1:]) + token_ids_match = fusion_token_ids == disabled_token_ids + + result = { + "model": self.MODEL_PATH, + "prompt_lengths": prompt_lengths, + "max_tokens": sampling_params.max_tokens, + "disabled": disabled_result, + "fusion": fusion_result, + "disabled_repeatable": disabled_repeatable, + "fusion_repeatable": fusion_repeatable, + "token_ids_match": token_ids_match, + } + result_json = json.dumps(result, sort_keys=True) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"DSV4_BCG_EPILOGUE_AB_RESULT={result_json}") + result_path = os.environ.get("TRTLLM_DSV4_BCG_AB_RESULT_PATH") + if result_path and MPI.COMM_WORLD.Get_rank() == 0: + result_dir = os.path.dirname(result_path) + if result_dir: + os.makedirs(result_dir, exist_ok=True) + with open(result_path, "w") as result_file: + json.dump(result, result_file, indent=2, sort_keys=True) + assert disabled_repeatable + assert fusion_repeatable + assert token_ids_match + _DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT = ( "Solve the problem carefully. End your response with a final line exactly " diff --git a/tests/unittest/_torch/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py index 3811d05a3490..a2ea2906f37b 100644 --- a/tests/unittest/_torch/modules/test_mla_registry.py +++ b/tests/unittest/_torch/modules/test_mla_registry.py @@ -16,12 +16,16 @@ from types import SimpleNamespace from unittest.mock import Mock, patch +import pytest import torch from torch import nn from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module import ( + _create_dsv4_epilogue_buffers, + _run_dsv4_epilogue_bmms, prepare_sparse_attn_outputs, + project_sparse_attn_output, ) from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.mla import MLA @@ -86,19 +90,160 @@ def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None: assert registry["0_1"]() is next_mla -def test_dsv4_epilogue_fusion_is_disabled_inside_breakable_graph() -> None: - metadata = SimpleNamespace(num_contexts=1, num_generations=0, num_tokens=5) - mla_layer = Mock(spec=MLA) - mla_layer._disable_dsv4_epilogue_fusion = False - mla_layer.create_output.return_value = torch.empty(8, 8) - hidden_states = torch.empty(8, 8) +def _make_dsv4_epilogue_layer() -> SimpleNamespace: + return SimpleNamespace( + _disable_dsv4_epilogue_fusion=False, + mapping=SimpleNamespace( + has_cp_helix=lambda: False, + enable_attention_dp=True, + ), + num_heads=128, + num_heads_tp=128, + mqa=SimpleNamespace( + sparse_params=object(), + has_fp8_kv_cache=True, + ), + o_a_proj=SimpleNamespace(dtype=torch.float8_e4m3fn), + kv_lora_rank=448, + qk_rope_head_dim=64, + qk_head_dim=512, + v_head_dim=512, + n_local_groups=8, + o_lora_rank=3, + dtype=torch.bfloat16, + inverse_rotary_emb=SimpleNamespace(is_neox=False), + create_output=Mock(), + ) + + +def test_mla_custom_op_marks_only_final_output_mutable() -> None: + schema = torch.ops.trtllm.mla_custom_op_inplace.default._schema + mutated_args = [ + arg.name + for arg in schema.arguments + if arg.alias_info is not None and arg.alias_info.is_write + ] + assert mutated_args == ["output"] + + +def test_dsv4_epilogue_fusion_supports_mixed_batch() -> None: + mla_layer = _make_dsv4_epilogue_layer() + metadata = SimpleNamespace(num_contexts=1, num_generations=1) + hidden_states = torch.empty(8, 16) with patch( - "tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module." - "is_in_breakable_cuda_graph", + "tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module.is_sm_100f", return_value=True, ): outputs = prepare_sparse_attn_outputs(mla_layer, hidden_states, metadata) - assert outputs == [mla_layer.create_output.return_value] - mla_layer.create_output.assert_called_once_with(hidden_states, 1) + assert len(outputs) == 1 + assert outputs[0].shape == (8, 8, 3) + assert outputs[0].dtype == torch.bfloat16 + mla_layer.create_output.assert_not_called() + + +def test_dsv4_fusion_create_output_uses_bucket_token_count() -> None: + mla_layer = _make_dsv4_epilogue_layer() + metadata = SimpleNamespace(num_contexts=1, num_generations=0) + hidden_states = torch.empty(8, 16) + + with patch( + "tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module.is_sm_100f", + return_value=True, + ): + output = prepare_sparse_attn_outputs(mla_layer, hidden_states, metadata)[0] + + assert output.shape == (8, 8, 3) + assert output.dtype == torch.bfloat16 + + +def test_dsv4_fusion_o_proj_only_flattens_lora_output() -> None: + projected = torch.randn(7, 5) + mla_layer = SimpleNamespace( + n_local_groups=4, + o_lora_rank=3, + o_b_proj=Mock(return_value=projected), + ) + lora_o = torch.randn(7, 4, 3) + + output = project_sparse_attn_output(mla_layer, [lora_o]) + + assert output is projected + mla_layer.o_b_proj.assert_called_once() + torch.testing.assert_close(mla_layer.o_b_proj.call_args.args[0], lora_o.flatten(1)) + + +def test_dsv4_epilogue_buffers_use_real_token_count() -> None: + mla_layer = SimpleNamespace( + n_local_groups=4, + num_heads_tp=128, + v_head_dim=512, + ) + q = torch.empty(8, 16) + + fp8_o, output_sf = _create_dsv4_epilogue_buffers(mla_layer, q, num_tokens=5) + + assert fp8_o.shape == (4, 5, 32 * 512) + assert output_sf.shape == (4, 32 * 4, 8) + + +@pytest.mark.parametrize( + "num_context_tokens,num_generation_tokens,bucket_tokens", + [(5, 0, 8), (0, 3, 4), (5, 3, 12)], +) +def test_dsv4_epilogue_bmm_writes_only_phase_ranges( + num_context_tokens: int, + num_generation_tokens: int, + bucket_tokens: int, +) -> None: + groups = 2 + rank = 3 + output = torch.full((bucket_tokens, groups, rank), -1.0) + mla_layer = SimpleNamespace( + o_a_proj=torch.empty(0), + o_a_proj_scale=torch.empty(0), + ) + + def fake_bmm(_attn_fp8, _weight, attn_scale, _weight_scale, phase_output): + phase_output.fill_(attn_scale.item()) + + with patch.object( + torch.ops.trtllm, + "cute_dsl_fp8_bmm_blackwell", + side_effect=fake_bmm, + ) as bmm: + context_epilogue = None + if num_context_tokens: + context_epilogue = ( + torch.empty(groups, num_context_tokens, 4), + torch.tensor(11.0), + ) + generation_epilogue = None + if num_generation_tokens: + generation_epilogue = ( + torch.empty(groups, num_generation_tokens, 4), + torch.tensor(22.0), + ) + _run_dsv4_epilogue_bmms( + mla_layer, + output, + num_context_tokens, + num_context_tokens + num_generation_tokens, + context_epilogue, + generation_epilogue, + ) + + assert bmm.call_count == bool(num_context_tokens) + bool(num_generation_tokens) + if num_context_tokens: + torch.testing.assert_close( + output[:num_context_tokens], torch.full_like(output[:num_context_tokens], 11.0) + ) + if num_generation_tokens: + generation_end = num_context_tokens + num_generation_tokens + torch.testing.assert_close( + output[num_context_tokens:generation_end], + torch.full_like(output[num_context_tokens:generation_end], 22.0), + ) + real_tokens = num_context_tokens + num_generation_tokens + torch.testing.assert_close(output[real_tokens:], torch.full_like(output[real_tokens:], -1.0)) From d7ad5566720fd5795b348f1126c71009115d35ab Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Mon, 27 Jul 2026 07:40:10 +0000 Subject: [PATCH 06/23] fix bcg and epilogue Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py index 807d3b722505..a8ef22b500c9 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -14,7 +14,6 @@ BreakableCUDAGraphCapture, enable_breakable_cuda_graph, ) -from .trace_log_utils import log_mem_snapshot class BreakableCUDAGraphRunnerState(Enum): @@ -78,7 +77,6 @@ def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None: self._capture_stream.wait_stream(current_stream) graph = None created_memory_pool = False - log_mem_snapshot(f"bcg/before_capture_{num_tokens}") try: with torch.cuda.stream(self._capture_stream): self.warmup(engine_forward) @@ -105,13 +103,11 @@ def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None: assert graph is not None self._graphs[num_tokens] = graph self._outputs[num_tokens] = make_weak_ref(output) - log_mem_snapshot(f"bcg/after_capture_{num_tokens}") except Exception: if graph is not None: graph.reset() if created_memory_pool and not self._graphs: self._memory_pool = None - log_mem_snapshot(f"bcg/capture_failed_{num_tokens}") raise finally: self._active_graph = None From bccfed8a07c96bd12e0b7007d64b21feba1ac01b Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Tue, 28 Jul 2026 05:04:49 +0000 Subject: [PATCH 07/23] rename some var name to make mla more clear Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../sparse/deepseek_v4/module.py | 90 ++++++++++--------- .../sparse/dsa/custom_ops.py | 2 +- .../_torch/models/modeling_minimaxm3.py | 4 +- tensorrt_llm/_torch/modules/attention.py | 2 +- .../_torch/modules/mamba/gdn_mixer.py | 2 +- tensorrt_llm/_torch/modules/mla.py | 19 ++-- .../breakable_cuda_graph.py | 58 +++++------- tensorrt_llm/_torch/utils.py | 18 +++- .../executor/test_breakable_cuda_graph.py | 33 ++++--- .../_torch/modules/test_mla_registry.py | 9 +- 10 files changed, 126 insertions(+), 111 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index c0e450f15a1e..381dc7fc80c5 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -187,40 +187,36 @@ def _create_dsv4_epilogue_buffers( return fp8_o, output_sf -def _run_dsv4_epilogue_bmm( +def _run_dsv4_o_lora_bmms( self, - epilogue_output: tuple[torch.Tensor, torch.Tensor], - output: torch.Tensor, -) -> None: - attn_fp8, attn_scale = epilogue_output - torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( - attn_fp8, - self.o_a_proj, - attn_scale, - self.o_a_proj_scale, - output.transpose(0, 1), - ) - - -def _run_dsv4_epilogue_bmms( - self, - output: torch.Tensor, + o_lora_output: torch.Tensor, num_context_tokens: int, num_tokens: int, - context_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]], - generation_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]], + context_o_lora_bmm_input: Optional[tuple[torch.Tensor, torch.Tensor]], + generation_o_lora_bmm_input: Optional[tuple[torch.Tensor, torch.Tensor]], ) -> None: - if context_epilogue_output is not None: - _run_dsv4_epilogue_bmm( - self, - context_epilogue_output, - output[:num_context_tokens], + def run_o_lora_bmm( + o_lora_bmm_input: tuple[torch.Tensor, torch.Tensor], + phase_o_lora_output: torch.Tensor, + ) -> None: + attn_fp8, attn_scale = o_lora_bmm_input + torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( + attn_fp8, + self.o_a_proj, + attn_scale, + self.o_a_proj_scale, + phase_o_lora_output.transpose(0, 1), ) - if generation_epilogue_output is not None: - _run_dsv4_epilogue_bmm( - self, - generation_epilogue_output, - output[num_context_tokens:num_tokens], + + if context_o_lora_bmm_input is not None: + run_o_lora_bmm( + context_o_lora_bmm_input, + o_lora_output[:num_context_tokens], + ) + if generation_o_lora_bmm_input is not None: + run_o_lora_bmm( + generation_o_lora_bmm_input, + o_lora_output[num_context_tokens:num_tokens], ) @@ -410,11 +406,13 @@ def forward_generation_sparse_attn( quant_q_buffer, ) - attention_output = output - output_sf = None + dsv4_output = output + o_lora_bmm_input_scale = None inverse_rope_cos_sin = None if enable_dsv4_epilogue_fusion: - attention_output, output_sf = _create_dsv4_epilogue_buffers(self, q, num_tokens) + dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers( + self, q, num_tokens + ) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -426,8 +424,8 @@ def forward_generation_sparse_attn( attn_metadata, attention_input_type=AttentionInputType.generation_only, out_scale=self.out_scale, - output=attention_output, - output_sf=output_sf, + output=dsv4_output, + output_sf=o_lora_bmm_input_scale, latent_cache=latent_cache, q_pe=q_pe, sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), @@ -441,8 +439,8 @@ def forward_generation_sparse_attn( enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) if enable_dsv4_epilogue_fusion: - assert attention_output is not None and output_sf is not None - return attention_output, output_sf + assert dsv4_output is not None and o_lora_bmm_input_scale is not None + return dsv4_output, o_lora_bmm_input_scale assert output is not None if self.mapping.has_cp_helix(): @@ -492,11 +490,13 @@ def forward_context_sparse_attn( quant_q_buffer = None quant_scale_qkv = None - attention_output = output - output_sf = None + dsv4_output = output + o_lora_bmm_input_scale = None inverse_rope_cos_sin = None if enable_dsv4_epilogue_fusion: - attention_output, output_sf = _create_dsv4_epilogue_buffers(self, q, num_tokens) + dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers( + self, q, num_tokens + ) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -508,8 +508,8 @@ def forward_context_sparse_attn( attn_metadata, attention_input_type=AttentionInputType.context_only, out_scale=self.out_scale, - output=attention_output, - output_sf=output_sf, + output=dsv4_output, + output_sf=o_lora_bmm_input_scale, latent_cache=latent_cache, q_pe=q_pe, quant_q_buffer=quant_q_buffer, @@ -522,8 +522,8 @@ def forward_context_sparse_attn( self._fused_q_pe = None if enable_dsv4_epilogue_fusion: - assert attention_output is not None and output_sf is not None - return attention_output, output_sf + assert dsv4_output is not None and o_lora_bmm_input_scale is not None + return dsv4_output, o_lora_bmm_input_scale assert output is not None if self.mapping.has_cp_helix(): @@ -826,7 +826,9 @@ def _indexer_branch(): assert generation_o_lora_bmm_input is None or isinstance( generation_o_lora_bmm_input, tuple ) - _run_dsv4_epilogue_bmms( + # The fused kernel output is group-first, which BCG cannot slice on + # dim 0. Write O-LoRA as token-first so replay can slice the bucket. + _run_dsv4_o_lora_bmms( self, output, num_ctx_tokens, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py index 60c56cc22b2d..5e4bc2efc4aa 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py @@ -149,4 +149,4 @@ def _mla_dsa_attn_inplace_fake( """Model the in-place output mutation during fake-tensor propagation.""" -maybe_bcg_mla_dsa_attn_inplace = eager_on_graph(True)(mla_dsa_attn_inplace) +maybe_bcg_mla_dsa_attn_inplace = eager_on_graph(mla_dsa_attn_inplace) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 9360b03297a2..b08a1db57087 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -662,9 +662,7 @@ def minimax_m3_attn_custom_op_inplace( ) -maybe_bcg_minimax_m3_attn_custom_op_inplace = eager_on_graph(True)( - minimax_m3_attn_custom_op_inplace -) +maybe_bcg_minimax_m3_attn_custom_op_inplace = eager_on_graph(minimax_m3_attn_custom_op_inplace) class MiniMaxM3Attention(Attention): diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 1b197a750b49..b6d3d2e7f23f 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -119,7 +119,7 @@ def attn_custom_op_inplace( ) -maybe_bcg_attn_custom_op_inplace = eager_on_graph(True)(attn_custom_op_inplace) +maybe_bcg_attn_custom_op_inplace = eager_on_graph(attn_custom_op_inplace) def _helix_zero_kv_mask( diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 9136683df510..f510e0da8342 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -175,7 +175,7 @@ def gdn_custom_op_inplace( ) -breakable_gdn_custom_op_inplace = eager_on_graph(True)(gdn_custom_op_inplace) +breakable_gdn_custom_op_inplace = eager_on_graph(gdn_custom_op_inplace) def ensure_divisibility(numerator, denominator): diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 50e1336cedf8..646509bb25da 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -112,18 +112,21 @@ def _extract_mla_extra_attrs(layer_idx: str): return metadata, mla_layer -def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> list[torch.Tensor]: +def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) - return mla_layer._create_outputs(hidden_states, metadata) + outputs = mla_layer._create_outputs(hidden_states, metadata) + if len(outputs) != 1: + raise RuntimeError("MLA custom ops require exactly one output tensor.") + return outputs[0] @torch.library.custom_op("trtllm::create_mla_outputs", mutates_args=()) -def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> list[torch.Tensor]: +def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @create_mla_outputs.register_fake -def _create_mla_outputs_fake(hidden_states, layer_idx): +def _create_mla_outputs_fake(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @@ -159,7 +162,7 @@ def mla_custom_op_inplace( ) -maybe_bcg_mla_custom_op_inplace = eager_on_graph(True)(mla_custom_op_inplace) +maybe_bcg_mla_custom_op_inplace = eager_on_graph(mla_custom_op_inplace) def fp8_block_scaling_bmm_out( @@ -1801,9 +1804,9 @@ def forward( "unquantized_hidden_states view" ) output_hidden_states = hidden_states.unquantized_hidden_states - attn_output = torch.ops.trtllm.create_mla_outputs( - output_hidden_states, self.layer_idx_str - ) + attn_output = [ + torch.ops.trtllm.create_mla_outputs(output_hidden_states, self.layer_idx_str) + ] self._forward_custom_op( hidden_states, position_ids, diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index 28fa114cda27..aad92c62da6f 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -143,46 +143,36 @@ def _copy_output(destination: Any, source: Any) -> Any: return source -def eager_on_graph(enable: bool) -> Callable[[Callable], Callable]: +def eager_on_graph(inner: Callable) -> Callable: """Run a callable eagerly between captured CUDA graph segments.""" - def decorator(inner: Callable) -> Callable: - if not enable: - return inner + @functools.wraps(inner) + def wrapper(*args, **kwargs): + capture = _current_capture.get() + if capture is None: + return inner(*args, **kwargs) - @functools.wraps(inner) - def wrapper(*args, **kwargs): - capture = _current_capture.get() - if capture is None: - return inner(*args, **kwargs) + logger.debug( + "Break CUDA graph for function %s", getattr(inner, "__name__", type(inner).__name__) + ) + capture._end_current_segment() + output = inner(*args, **kwargs) - logger.debug( - "Break CUDA graph for function %s", getattr(inner, "__name__", type(inner).__name__) - ) - capture._end_current_segment() - output = inner(*args, **kwargs) + captured_args = tuple(make_weak_ref(arg, preserve_unsupported=True) for arg in args) + captured_kwargs = { + key: make_weak_ref(value, preserve_unsupported=True) for key, value in kwargs.items() + } + captured_output = make_weak_ref(output, preserve_unsupported=True) - # 看下attn的参数 - def make_weak_ref_with_str_none(x): - if isinstance(x, (str, type(None))): - return x - return make_weak_ref(x) + def replay_fn() -> Any: + new_output = inner(*captured_args, **captured_kwargs) + return _copy_output(captured_output, new_output) - captured_args = tuple(make_weak_ref_with_str_none(arg) for arg in args) - captured_kwargs = {key: make_weak_ref_with_str_none(value) for key, value in kwargs.items()} - captured_output = make_weak_ref_with_str_none(output) + capture.cuda_graph._break_functions.append(replay_fn) + capture._begin_new_segment() + return output - def replay_fn() -> Any: - new_output = inner(*captured_args, **captured_kwargs) - return _copy_output(captured_output, new_output) - - capture.cuda_graph._break_functions.append(replay_fn) - capture._begin_new_segment() - return output - - return wrapper - - return decorator + return wrapper class BreakableCUDAGraph: @@ -286,7 +276,7 @@ def _end_current_segment(self) -> None: self.cuda_graph._segments[-1].capture_end() -@eager_on_graph(True) +@eager_on_graph def break_graph() -> None: """Insert an empty eager break between CUDA graph segments.""" return None diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index c65ecadddaf3..cde74352f7ae 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -159,20 +159,30 @@ def wrapper(self, *args, **kwargs): return decorator -def make_weak_ref(x): +def make_weak_ref(x, preserve_unsupported: bool = False): if isinstance(x, torch.Tensor): return convert_to_torch_tensor( TensorWrapper(x.data_ptr(), x.dtype, x.shape, x.stride())) if x.is_cuda else x elif isinstance(x, tuple): - return tuple(make_weak_ref(i) for i in x) + return tuple( + make_weak_ref(i, preserve_unsupported=preserve_unsupported) + for i in x) elif isinstance(x, list): - return [make_weak_ref(i) for i in x] + return [ + make_weak_ref(i, preserve_unsupported=preserve_unsupported) + for i in x + ] elif isinstance(x, dict): - return {k: make_weak_ref(v) for k, v in x.items()} + return { + k: make_weak_ref(v, preserve_unsupported=preserve_unsupported) + for k, v in x.items() + } elif isinstance(x, (int, float, bool)): return x + elif preserve_unsupported: + return x else: raise TypeError(f"Invalid type {type(x)} to make weak ref") diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index e34bcab1b15a..1cb925bd9937 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import gc -import weakref import pytest import torch @@ -20,6 +18,7 @@ BreakableCUDAGraphRunner, BreakableCUDAGraphRunnerState, ) +from tensorrt_llm._torch.utils import make_weak_ref pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -46,11 +45,11 @@ def test_no_break_capture_and_repeated_replay(): def test_single_and_multiple_breakpoints(): - @eager_on_graph(True) + @eager_on_graph def add_one(value): return value + 1 - @eager_on_graph(True) + @eager_on_graph def double(value): return value * 2 @@ -72,20 +71,28 @@ def body(): torch.testing.assert_close(output, torch.full_like(output, 16)) -def test_disabled_and_outside_capture(): - @eager_on_graph(False) - def disabled(value): - return value + 1 - - @eager_on_graph(True) +def test_outside_capture(): + @eager_on_graph def outside(value): return value + 2 value = torch.tensor([1.0, 2.0], device="cuda") - torch.testing.assert_close(disabled(value), value + 1) torch.testing.assert_close(outside(value), value + 2) +def test_make_weak_ref_option_preserves_unsupported_values(): + unsupported = object() + with pytest.raises(TypeError, match="Invalid type"): + make_weak_ref(unsupported) + + captured = make_weak_ref( + {"nested": (unsupported, [None, "value"])}, + preserve_unsupported=True, + ) + assert captured == {"nested": (unsupported, [None, "value"])} + assert captured["nested"][0] is unsupported + + def test_break_graph_inserts_empty_breakpoint(): x = torch.zeros(4, device="cuda") output = torch.zeros_like(x) @@ -204,11 +211,11 @@ def engine_forward(): def test_runner_first_bucket_segments_share_one_memory_pool(): class BreakableBody(nn.Module): - @eager_on_graph(True) + @eager_on_graph def eager_add_one(self, value): return value + 1 - @eager_on_graph(True) + @eager_on_graph def eager_double(self, value): return value * 2 diff --git a/tests/unittest/_torch/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py index a2ea2906f37b..4e40bf419666 100644 --- a/tests/unittest/_torch/modules/test_mla_registry.py +++ b/tests/unittest/_torch/modules/test_mla_registry.py @@ -23,7 +23,7 @@ from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module import ( _create_dsv4_epilogue_buffers, - _run_dsv4_epilogue_bmms, + _run_dsv4_o_lora_bmms, prepare_sparse_attn_outputs, project_sparse_attn_output, ) @@ -126,6 +126,11 @@ def test_mla_custom_op_marks_only_final_output_mutable() -> None: assert mutated_args == ["output"] +def test_create_mla_outputs_custom_op_returns_tensor() -> None: + schema = torch.ops.trtllm.create_mla_outputs.default._schema + assert [str(return_value.type) for return_value in schema.returns] == ["Tensor"] + + def test_dsv4_epilogue_fusion_supports_mixed_batch() -> None: mla_layer = _make_dsv4_epilogue_layer() metadata = SimpleNamespace(num_contexts=1, num_generations=1) @@ -225,7 +230,7 @@ def fake_bmm(_attn_fp8, _weight, attn_scale, _weight_scale, phase_output): torch.empty(groups, num_generation_tokens, 4), torch.tensor(22.0), ) - _run_dsv4_epilogue_bmms( + _run_dsv4_o_lora_bmms( mla_layer, output, num_context_tokens, From d9e373deee642b34d75119eec86129e8375deb6e Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:06:40 +0000 Subject: [PATCH 08/23] [None][test] add focused BCG accuracy coverage Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../accuracy/test_disaggregated_serving.py | 62 +++++ .../defs/accuracy/test_llm_api_pytorch.py | 239 ++++++++---------- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_dgx_b200.yml | 4 + 4 files changed, 168 insertions(+), 138 deletions(-) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 749ed07f50a0..41fc98bd3e4e 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2480,6 +2480,68 @@ class TestDeepSeekV4Flash(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash" + @pytest.mark.skip_less_device(8) + def test_prefill_breakable_cuda_graph(self): + """Disaggregated accuracy with BCG enabled only on the prefill worker.""" + cache_transceiver_config = { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 4096, + } + ctx_server_config = { + "tensor_parallel_size": 4, + "moe_expert_parallel_size": 4, + "enable_attention_dp": True, + "disable_overlap_scheduler": True, + "max_batch_size": 8, + "max_num_tokens": 1024, + "max_seq_len": 4096, + "moe_config": { + "backend": "TRTLLM", + }, + "kv_cache_config": { + "dtype": "fp8", + "free_gpu_memory_fraction": 0.6, + }, + "cache_transceiver_config": cache_transceiver_config, + "prefill_cuda_graph_backend": "breakable", + "prefill_capture_num_tokens": [128, 256, 512, 1024], + } + gen_server_config = { + "tensor_parallel_size": 4, + "moe_expert_parallel_size": 4, + "enable_attention_dp": True, + "disable_overlap_scheduler": True, + "max_batch_size": 8, + "max_num_tokens": 1024, + "max_seq_len": 4096, + "moe_config": { + "backend": "TRTLLM", + }, + "kv_cache_config": { + "dtype": "fp8", + "free_gpu_memory_fraction": 0.6, + }, + "cache_transceiver_config": cache_transceiver_config, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + }, + } + with launch_disaggregated_llm(disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + server_waiting_timeout=3600) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + @pytest.mark.skip_less_device(4) def test_auto_dtype(self): # Disagg smoke test: CTX TP=2 + GEN TP=2 = 4 GPUs. diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 2ad2f516b8af..8ee1a2f833a6 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -15,7 +15,6 @@ import asyncio import json import os -import statistics import sys import time from unittest import mock @@ -24,7 +23,6 @@ import torch from datasets import load_dataset from defs.conftest import get_sm_version, is_sm_100f -from mpi4py import MPI from mpi4py.futures import MPIPoolExecutor from tensorrt_llm import LLM @@ -3600,6 +3598,65 @@ def test_nvfp4_multi_gpus_piecewise_cuda_graph(self, tp_size, pp_size, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip_less_mpi_world_size(8) + @skip_pre_blackwell + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size,mtp_nextn,attention_dp,max_batch_size,moe_backend,fp8kv,chunked_prefill", + [ + (8, 1, 8, 0, True, 24, "CUTLASS", False, False), + (8, 1, 8, 3, False, 16, "TRTLLM", True, True), + ], + ids=["baseline", "mtp3_fp8kv_chunked"]) + def test_nvfp4_multi_gpus_breakable_cuda_graph( + self, tp_size, pp_size, ep_size, mtp_nextn, attention_dp, + max_batch_size, moe_backend, fp8kv, chunked_prefill): + sm_version = get_sm_version() + if moe_backend == "TRTLLM" and sm_version in (120, 121): + pytest.skip(f"{moe_backend} backend does not support SM 120 or 121") + + moe_config = MoeConfig(backend=moe_backend, max_num_tokens=16384) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7) + if fp8kv: + kv_cache_config.dtype = "fp8" + kv_cache_config.enable_block_reuse = True + + pytorch_config = dict( + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, + max_batch_size=max_batch_size, + ), + moe_config=moe_config, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + prefill_capture_num_tokens=[2048, 8192], + ) + + mtp_config = None + if mtp_nextn > 0: + mtp_config = MTPDecodingConfig(max_draft_len=mtp_nextn) + + llm_kwargs = dict( + max_batch_size=max_batch_size, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) + if chunked_prefill: + llm_kwargs.update( + enable_chunked_prefill=True, + max_num_tokens=8192, + ) + + with LLM(f"{llm_models_root()}/DeepSeek-V3.2-Exp-FP4-v2", + **pytorch_config, **llm_kwargs) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.skip_less_mpi_world_size(8) @skip_pre_blackwell @pytest.mark.parametrize( @@ -3939,25 +3996,9 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): @pytest.mark.skip_less_mpi_world_size(8) @pytest.mark.threadleak(enabled=False) - def test_mixed_breakable_cuda_graph_epilogue_fusion_ab(self, mocker): + def test_mixed_breakable_cuda_graph(self): from transformers import AutoTokenizer - from tensorrt_llm.llmapi.mpi_session import MpiPoolSession - - fusion_env = "TRTLLM_DSV4_DISABLE_FMHA_EPILOGUE_FUSION" - - def patched_start_mpi_pool(session): - assert not session.mpi_pool, "MPI session already started" - session.mpi_pool = MPIPoolExecutor( - max_workers=session.n_workers, - path=sys.path, - env={fusion_env: os.environ.get(fusion_env, "0")}, - ) - - mocker.patch.object(MpiPoolSession, "_start_mpi_pool", - patched_start_mpi_pool) - - prompt_lengths = [64, 129, 257, 385, 513, 769] tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) base_prompt_ids = tokenizer.encode( "TensorRT-LLM accelerates reliable large language model inference " @@ -3965,11 +4006,14 @@ def patched_start_mpi_pool(session): add_special_tokens=False, ) assert base_prompt_ids - prompts = [ - (base_prompt_ids * ((prompt_length + len(base_prompt_ids) - 1) // - len(base_prompt_ids)))[:prompt_length] - for prompt_length in prompt_lengths - ] + + def make_prompt(prompt_length): + return (base_prompt_ids * + ((prompt_length + len(base_prompt_ids) - 1) // + len(base_prompt_ids)))[:prompt_length] + + generation_prompt = make_prompt(64) + context_prompt = make_prompt(129) sampling_params = SamplingParams( max_tokens=8, min_tokens=8, @@ -3979,7 +4023,7 @@ def patched_start_mpi_pool(session): detokenize=False, add_special_tokens=False, ) - llm_kwargs = dict( + common_llm_kwargs = dict( tensor_parallel_size=8, moe_expert_parallel_size=8, moe_config=MoeConfig(backend="TRTLLM"), @@ -3996,120 +4040,39 @@ def patched_start_mpi_pool(session): batch_sizes=[1, 2, 4, 6, 8], enable_padding=True, ), - prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, - prefill_capture_num_tokens=[128, 256, 512, 1024], - ) - - def run_variant(disable_fusion: bool) -> dict: - variant_start = time.perf_counter() - with mock.patch.dict( - os.environ, - {fusion_env: "1" if disable_fusion else "0"}, - clear=False, - ): - with LLM( - self.MODEL_PATH, - **llm_kwargs, - env_overrides={ - fusion_env: "1" if disable_fusion else "0" - }, - ) as llm: - init_seconds = time.perf_counter() - variant_start - warmup_start = time.perf_counter() - llm.generate(prompts, - sampling_params=sampling_params, - use_tqdm=False) - warmup_seconds = time.perf_counter() - warmup_start - - rounds = [] - for _ in range(5): - round_start = time.perf_counter() - outputs = llm.generate( - prompts, - sampling_params=sampling_params, - use_tqdm=False, - ) - latency_seconds = time.perf_counter() - round_start - token_ids = [ - output.outputs[0].token_ids for output in outputs - ] - output_tokens = sum(len(ids) for ids in token_ids) - rounds.append({ - "latency_seconds": - latency_seconds, - "output_tokens": - output_tokens, - "output_tokens_per_second": - output_tokens / latency_seconds, - "token_ids": - token_ids, - }) - - latencies = [ - round_result["latency_seconds"] for round_result in rounds - ] - throughputs = [ - round_result["output_tokens_per_second"] - for round_result in rounds - ] - return { - "fusion_disabled": - disable_fusion, - "engine_init_capture_seconds": - init_seconds, - "warmup_seconds": - warmup_seconds, - "rounds": - rounds, - "median_latency_seconds": - statistics.median(latencies), - "p90_latency_seconds": - statistics.quantiles(latencies, n=10, method="inclusive")[8], - "median_output_tokens_per_second": - statistics.median(throughputs), - "p90_output_tokens_per_second": - statistics.quantiles(throughputs, n=10, method="inclusive")[8], - } + ) - disabled_result = run_variant(disable_fusion=True) - fusion_result = run_variant(disable_fusion=False) - disabled_token_ids = [ - round_result["token_ids"] - for round_result in disabled_result["rounds"] - ] - fusion_token_ids = [ - round_result["token_ids"] - for round_result in fusion_result["rounds"] - ] - disabled_repeatable = all(token_ids == disabled_token_ids[0] - for token_ids in disabled_token_ids[1:]) - fusion_repeatable = all(token_ids == fusion_token_ids[0] - for token_ids in fusion_token_ids[1:]) - token_ids_match = fusion_token_ids == disabled_token_ids - - result = { - "model": self.MODEL_PATH, - "prompt_lengths": prompt_lengths, - "max_tokens": sampling_params.max_tokens, - "disabled": disabled_result, - "fusion": fusion_result, - "disabled_repeatable": disabled_repeatable, - "fusion_repeatable": fusion_repeatable, - "token_ids_match": token_ids_match, - } - result_json = json.dumps(result, sort_keys=True) - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"DSV4_BCG_EPILOGUE_AB_RESULT={result_json}") - result_path = os.environ.get("TRTLLM_DSV4_BCG_AB_RESULT_PATH") - if result_path and MPI.COMM_WORLD.Get_rank() == 0: - result_dir = os.path.dirname(result_path) - if result_dir: - os.makedirs(result_dir, exist_ok=True) - with open(result_path, "w") as result_file: - json.dump(result, result_file, indent=2, sort_keys=True) - assert disabled_repeatable - assert fusion_repeatable - assert token_ids_match + def run(backend): + with LLM( + self.MODEL_PATH, + **common_llm_kwargs, + prefill_cuda_graph_backend=backend, + prefill_capture_num_tokens=[128, 256, 512, 1024], + ) as llm: + generation_request = llm.generate_async( + generation_prompt, + sampling_params=sampling_params, + streaming=True, + ) + next(generation_request) + assert not generation_request.finished + + # Admit a context request while the first request is decoding. + context_request = llm.generate_async( + context_prompt, + sampling_params=sampling_params, + streaming=False, + ) + generation_output = generation_request.result() + context_output = context_request.result() + return [ + generation_output.outputs[0].token_ids, + context_output.outputs[0].token_ids, + ] + + eager_token_ids = run(PrefillCudaGraphBackend.DISABLED) + breakable_token_ids = run(PrefillCudaGraphBackend.BREAKABLE) + assert breakable_token_ids == eager_token_ids _DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT = ( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 406b01f25807..e3932cba01f3 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -57,6 +57,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a16_mxfp4[latency-TRTLLM] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention[target_sparsity_0.9-fp8kv=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=True] + - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16_breakable_prefill_cuda_graph - accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[TRTLLM] - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 74d7936cab15..b329d698f1e9 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -176,12 +176,16 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[baseline] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[baseline] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[mtp3_fp8kv_chunked] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4_mtp_index_share[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_mixed_breakable_cuda_graph TIMEOUT (120) ISOLATION + - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_prefill_breakable_cuda_graph TIMEOUT (120) ISOLATION - examples/test_deepseek_v4_pro.py::test_short_token_boundary_smoke TIMEOUT (120) - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) From 24830d86f9c634a978860e6e12e183ef2c4d5ac5 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 29 Jul 2026 02:03:51 +0000 Subject: [PATCH 09/23] [None][chore] apply pre-commit fixes to BCG changes Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../pyexecutor/breakable_cuda_graph_runner.py | 13 +++++------ .../_torch/pyexecutor/model_engine.py | 23 ++++++++++--------- .../defs/accuracy/test_llm_api_pytorch.py | 9 ++++---- tests/unittest/llmapi/test_llm_args.py | 6 +++-- 4 files changed, 27 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py index a8ef22b500c9..d1b8562ff568 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -56,7 +56,7 @@ def has_graph(self, num_tokens: int) -> bool: def warmup(self, engine_forward: Callable[[], Any], steps: int = _WARMUP_STEPS) -> None: """Run the complete eager engine forward under the warmup state. - model_engine.forward will use state to determine what forward to do.""" + model_engine.forward will use state to determine what forward to do.""" if self._state != BreakableCUDAGraphRunnerState.IDLE: raise RuntimeError(f"Cannot warm up BCG while runner is {self._state.value}") self._state = BreakableCUDAGraphRunnerState.WARMUP @@ -128,7 +128,7 @@ def capture_context(self) -> Iterator[None]: yield def capture_output(self, output: torch.Tensor) -> torch.Tensor: - """Route all bucket outputs through the largest capture's buffer. """ + """Route all bucket outputs through the largest capture's buffer.""" if not self.is_capturing or self._active_num_tokens is None: raise RuntimeError("BCG output registered outside capture") @@ -147,9 +147,9 @@ def capture_output(self, output: torch.Tensor) -> torch.Tensor: def capture_model_body(self, outer_forward: Callable[[], Any]) -> Any: """Run the outer model while capturing only its decoder body. - model_engine.forward is too broad and may pollute the CUDA stream - before the actual model forward. We want to reuse the functions - in forward that prepare the data and set the relevant flags.""" + model_engine.forward is too broad and may pollute the CUDA stream + before the actual model forward. We want to reuse the functions + in forward that prepare the data and set the relevant flags.""" if not self.is_capturing: raise RuntimeError("BCG body capture requested outside capture") @@ -159,8 +159,7 @@ def capture_model_body(self, outer_forward: Callable[[], Any]) -> Any: def capture_forward(*args, **kwargs): nonlocal captured_output with self.capture_context(): - captured_output = self.capture_output( - original_body_forward(*args, **kwargs)) + captured_output = self.capture_output(original_body_forward(*args, **kwargs)) return captured_output self.layer_model.forward = capture_forward diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 22fbb4fa13b7..756b0aec351a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -934,7 +934,8 @@ def __init__( raise ValueError( "breakable prefill CUDA graph requires a decoder model body" ) - self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner(decoder_model.model) + self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner( + decoder_model.model) # Initialize CUDA Graph LoRA manager if LoRA is enabled self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None @@ -2462,8 +2463,7 @@ def _capture_mixed_encoder_decoder_cuda_graphs( def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): """Capture configured CUDA graphs for context/prefill steps.""" - if (self.prefill_cuda_graph_backend - == PrefillCudaGraphBackend.DISABLED + if (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.DISABLED or (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE and not self._torch_compile_enabled)): @@ -2523,10 +2523,10 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): if self.breakable_cuda_graph_runner is not None: with self.no_cuda_graph(): self.breakable_cuda_graph_runner.warmup( - lambda: self.forward( - batch, - new_tensors_device=None, - resource_manager=resource_manager), + lambda: self.forward(batch, + new_tensors_device=None, + resource_manager= + resource_manager), steps=1) else: self.forward(batch, @@ -3569,8 +3569,7 @@ def get_padded_prefill_tokens(tokens: int) -> int: return self._prefill_cuda_graph_num_tokens[bisect.bisect_left( self._prefill_cuda_graph_num_tokens, tokens)] - if (self.prefill_cuda_graph_backend - != PrefillCudaGraphBackend.DISABLED + if (self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED and self._prefill_cuda_graph_num_tokens): max_captured_num_tokens = self._prefill_cuda_graph_num_tokens[-1] if attn_all_rank_num_tokens is not None: @@ -4676,7 +4675,7 @@ def _apply_steady_gen_fast_prepare( attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = \ self._get_padding_params(num_requests, 0, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + set_per_request_prefill_cuda_graph_flag(can_run_piecewise_cuda_graph) attn_metadata.padded_num_tokens = ( padded_num_tokens if padded_num_tokens != num_requests else None) virtual_num_tokens = num_requests @@ -7194,8 +7193,10 @@ def forward_step(): inputs, gather_ids=gather_ids, gather_context_logits=gather_context_logits) + if not can_run_graph: - if (breakable_runner is not None and breakable_runner.is_capturing): + if (breakable_runner is not None + and breakable_runner.is_capturing): return breakable_runner.capture_model_body(forward_step) num_tokens = inputs['input_ids'].shape[0] diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8ee1a2f833a6..da7d40f26bb8 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -16,7 +16,6 @@ import json import os import sys -import time from unittest import mock import pytest @@ -3607,9 +3606,11 @@ def test_nvfp4_multi_gpus_piecewise_cuda_graph(self, tp_size, pp_size, (8, 1, 8, 3, False, 16, "TRTLLM", True, True), ], ids=["baseline", "mtp3_fp8kv_chunked"]) - def test_nvfp4_multi_gpus_breakable_cuda_graph( - self, tp_size, pp_size, ep_size, mtp_nextn, attention_dp, - max_batch_size, moe_backend, fp8kv, chunked_prefill): + def test_nvfp4_multi_gpus_breakable_cuda_graph(self, tp_size, pp_size, + ep_size, mtp_nextn, + attention_dp, max_batch_size, + moe_backend, fp8kv, + chunked_prefill): sm_version = get_sm_version() if moe_backend == "TRTLLM" and sm_version in (120, 121): pytest.skip(f"{moe_backend} backend does not support SM 120 or 121") diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index cf166601e2f8..8392112d72bf 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2043,6 +2043,7 @@ def test_attention_dp_prefill_graph_uses_all_rank_decision(self): PyTorchModelEngine class FakeDist: + def __init__(self, decisions): self.decisions = decisions @@ -2058,8 +2059,9 @@ def tp_allgather(self, value): all_rank_num_tokens = [1, 129, 1, 1] engine.dist = FakeDist([True, True, True, True]) - assert engine._get_padding_params( - 1, 0, all_rank_num_tokens) == (256, True, [256] * 4) + assert engine._get_padding_params(1, 0, + all_rank_num_tokens) == (256, True, + [256] * 4) engine.dist = FakeDist([True, False, True, True]) assert engine._get_padding_params( From acc3ef232e07ab4b2ef2cb9253f1830b505e8f41 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:13:24 +0000 Subject: [PATCH 10/23] [None][fix] preserve PCG behavior with breakable CUDA graphs Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../breakable_cuda_graph/breakable_cuda_graph.py | 3 +++ tensorrt_llm/_torch/pyexecutor/model_engine.py | 14 +++++++++----- .../_torch/executor/test_breakable_cuda_graph.py | 11 +++++++++++ tests/unittest/llmapi/test_llm_args.py | 2 +- 4 files changed, 24 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index aad92c62da6f..790cbce368f6 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -148,6 +148,9 @@ def eager_on_graph(inner: Callable) -> Callable: @functools.wraps(inner) def wrapper(*args, **kwargs): + if torch.compiler.is_compiling(): + return inner(*args, **kwargs) + capture = _current_capture.get() if capture is None: return inner(*args, **kwargs) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 756b0aec351a..607bddc6f7bd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -173,14 +173,14 @@ def warmup(self, resource_manager: ResourceManager) -> None: return -def _filter_prefill_capture_num_tokens( +def _filter_piecewise_capture_num_tokens( candidate_num_tokens: list[int], max_num_tokens: int, max_batch_size: int, max_seq_len: int, num_extra_decoding_steps: int = 0, ) -> Tuple[list[int], list[int]]: - """Cap prefill CUDA graph capture candidates at the engine's reachable + """Cap piecewise CUDA graph capture candidates at the engine's reachable `num_tokens` ceiling `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` clamping user-requested sizes above it down to the ceiling. @@ -203,10 +203,10 @@ def _filter_prefill_capture_num_tokens( """ max_capturable_num_tokens = max( 0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)) - prefill_capacity_limit = min(max_num_tokens, max_capturable_num_tokens) - if prefill_capacity_limit > 0: + piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens) + if piecewise_capacity_limit > 0: kept = sorted({ - min(i, prefill_capacity_limit) + min(i, piecewise_capacity_limit) for i in candidate_num_tokens if 0 < i <= max_num_tokens }) else: @@ -219,6 +219,10 @@ def _filter_prefill_capture_num_tokens( return kept, unrecordable +# BCG uses the same capture-bucket filtering semantics as PCG. +_filter_prefill_capture_num_tokens = _filter_piecewise_capture_num_tokens + + def _filter_cuda_graph_batch_sizes(cuda_graph_batch_sizes: list[int], max_batch_size: int, max_num_tokens: int, max_total_draft_tokens: int, diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index 1cb925bd9937..c2fbbab105b2 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -80,6 +80,17 @@ def outside(value): torch.testing.assert_close(outside(value), value + 2) +def test_eager_on_graph_during_torch_compile(): + @eager_on_graph + def add_one(value): + return value + 1 + + compiled_add_one = torch.compile(add_one, backend="eager", fullgraph=True) + value = torch.ones(4, device="cuda") + + torch.testing.assert_close(compiled_add_one(value), value + 1) + + def test_make_weak_ref_option_preserves_unsupported_values(): unsupported = object() with pytest.raises(TypeError, match="Invalid type"): diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 8392112d72bf..627a69664213 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2021,7 +2021,7 @@ def test_prefill_filter_sorts_dedupes_and_drops_nonpositive(self): max_batch_size=1, max_seq_len=513, ) - assert kept == [128, 256, 512] + assert kept == [128, 256] assert unrecordable == [] @pytest.mark.parametrize("backend", [ From 9fa740342d2eca7dd33c794dcbc39f10821915ff Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Thu, 30 Jul 2026 06:28:49 +0000 Subject: [PATCH 11/23] [None][test] remove DeepSeek V4 BCG accuracy tests Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../accuracy/test_disaggregated_serving.py | 62 -------------- .../defs/accuracy/test_llm_api_pytorch.py | 80 ------------------- .../test_lists/test-db/l0_dgx_b200.yml | 2 - 3 files changed, 144 deletions(-) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 41fc98bd3e4e..749ed07f50a0 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2480,68 +2480,6 @@ class TestDeepSeekV4Flash(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash" - @pytest.mark.skip_less_device(8) - def test_prefill_breakable_cuda_graph(self): - """Disaggregated accuracy with BCG enabled only on the prefill worker.""" - cache_transceiver_config = { - "backend": "NIXL", - "transceiver_runtime": "PYTHON", - "max_tokens_in_buffer": 4096, - } - ctx_server_config = { - "tensor_parallel_size": 4, - "moe_expert_parallel_size": 4, - "enable_attention_dp": True, - "disable_overlap_scheduler": True, - "max_batch_size": 8, - "max_num_tokens": 1024, - "max_seq_len": 4096, - "moe_config": { - "backend": "TRTLLM", - }, - "kv_cache_config": { - "dtype": "fp8", - "free_gpu_memory_fraction": 0.6, - }, - "cache_transceiver_config": cache_transceiver_config, - "prefill_cuda_graph_backend": "breakable", - "prefill_capture_num_tokens": [128, 256, 512, 1024], - } - gen_server_config = { - "tensor_parallel_size": 4, - "moe_expert_parallel_size": 4, - "enable_attention_dp": True, - "disable_overlap_scheduler": True, - "max_batch_size": 8, - "max_num_tokens": 1024, - "max_seq_len": 4096, - "moe_config": { - "backend": "TRTLLM", - }, - "kv_cache_config": { - "dtype": "fp8", - "free_gpu_memory_fraction": 0.6, - }, - "cache_transceiver_config": cache_transceiver_config, - } - disaggregated_server_config = { - "hostname": "localhost", - "backend": "pytorch", - "context_servers": { - "num_instances": 1 - }, - "generation_servers": { - "num_instances": 1 - }, - } - with launch_disaggregated_llm(disaggregated_server_config, - ctx_server_config, - gen_server_config, - self.MODEL_PATH, - server_waiting_timeout=3600) as llm: - task = MMLU(self.MODEL_NAME) - task.evaluate(llm, is_integration_test=True) - @pytest.mark.skip_less_device(4) def test_auto_dtype(self): # Disagg smoke test: CTX TP=2 + GEN TP=2 = 4 GPUs. diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index da7d40f26bb8..e81f27b165f5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3995,86 +3995,6 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): eplb_config, mtp_nextn=mtp_nextn) - @pytest.mark.skip_less_mpi_world_size(8) - @pytest.mark.threadleak(enabled=False) - def test_mixed_breakable_cuda_graph(self): - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) - base_prompt_ids = tokenizer.encode( - "TensorRT-LLM accelerates reliable large language model inference " - "with efficient attention, parallelism, and CUDA graphs. ", - add_special_tokens=False, - ) - assert base_prompt_ids - - def make_prompt(prompt_length): - return (base_prompt_ids * - ((prompt_length + len(base_prompt_ids) - 1) // - len(base_prompt_ids)))[:prompt_length] - - generation_prompt = make_prompt(64) - context_prompt = make_prompt(129) - sampling_params = SamplingParams( - max_tokens=8, - min_tokens=8, - seed=42, - temperature=0, - ignore_eos=True, - detokenize=False, - add_special_tokens=False, - ) - common_llm_kwargs = dict( - tensor_parallel_size=8, - moe_expert_parallel_size=8, - moe_config=MoeConfig(backend="TRTLLM"), - enable_attention_dp=True, - max_batch_size=8, - max_num_tokens=1024, - max_seq_len=2048, - kv_cache_config=KvCacheConfig( - enable_block_reuse=False, - dtype="fp8", - free_gpu_memory_fraction=0.6, - ), - cuda_graph_config=CudaGraphConfig( - batch_sizes=[1, 2, 4, 6, 8], - enable_padding=True, - ), - ) - - def run(backend): - with LLM( - self.MODEL_PATH, - **common_llm_kwargs, - prefill_cuda_graph_backend=backend, - prefill_capture_num_tokens=[128, 256, 512, 1024], - ) as llm: - generation_request = llm.generate_async( - generation_prompt, - sampling_params=sampling_params, - streaming=True, - ) - next(generation_request) - assert not generation_request.finished - - # Admit a context request while the first request is decoding. - context_request = llm.generate_async( - context_prompt, - sampling_params=sampling_params, - streaming=False, - ) - generation_output = generation_request.result() - context_output = context_request.result() - return [ - generation_output.outputs[0].token_ids, - context_output.outputs[0].token_ids, - ] - - eager_token_ids = run(PrefillCudaGraphBackend.DISABLED) - breakable_token_ids = run(PrefillCudaGraphBackend.BREAKABLE) - assert breakable_token_ids == eager_token_ids - _DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT = ( "Solve the problem carefully. End your response with a final line exactly " diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index b329d698f1e9..0780ae814890 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -184,8 +184,6 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4_mtp_index_share[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240) - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_mixed_breakable_cuda_graph TIMEOUT (120) ISOLATION - - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_prefill_breakable_cuda_graph TIMEOUT (120) ISOLATION - examples/test_deepseek_v4_pro.py::test_short_token_boundary_smoke TIMEOUT (120) - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) From 7650e7bbd016be981198de1253706be1aa96bba4 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Mon, 3 Aug 2026 06:42:16 +0000 Subject: [PATCH 12/23] [None][fix] repair BCG tests after rebase Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../compilation/test_remove_copy_pass.py | 47 +------------------ .../executor/test_pytorch_model_engine.py | 1 + 2 files changed, 2 insertions(+), 46 deletions(-) diff --git a/tests/unittest/_torch/compilation/test_remove_copy_pass.py b/tests/unittest/_torch/compilation/test_remove_copy_pass.py index bf2794e65230..974906f485a2 100644 --- a/tests/unittest/_torch/compilation/test_remove_copy_pass.py +++ b/tests/unittest/_torch/compilation/test_remove_copy_pass.py @@ -138,7 +138,7 @@ def test_remove_copy_for_mutates_tensor_list( graph.lint() -def test_remove_copy_for_mutates_args_restores_optional_none() -> None: +def test_remove_copy_for_mla_restores_final_output_mutation() -> None: graph = Graph() hidden_states = graph.placeholder("hidden_states") output = graph.placeholder("output") @@ -153,8 +153,6 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: "latent_cache_gen": None, "_all_bases": (output,), "_output_base_index": 0, - "_sparse_output_base_index": None, - "_sparse_output_sf_base_index": None, }, ) mutated_output = graph.call_function(getitem, args=(functionalized, 1)) @@ -166,48 +164,5 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: inplace_nodes = [node for node in graph.nodes if node.target == inplace_func] assert len(inplace_nodes) == 1 assert inplace_nodes[0].kwargs["output"] is output - assert inplace_nodes[0].kwargs["sparse_output"] is None - assert inplace_nodes[0].kwargs["sparse_output_sf"] is None assert clone.args[0] is output graph.lint() - - -def test_remove_copy_for_mutates_args_rejects_getitem_for_optional_none( - monkeypatch: pytest.MonkeyPatch, -) -> None: - graph = Graph() - hidden_states = graph.placeholder("hidden_states") - output = graph.placeholder("output") - inplace_func = torch.ops.trtllm.mla_custom_op_inplace.default - functionalized = graph.call_function( - auto_functionalized_v2, - args=(inplace_func,), - kwargs={ - "hidden_states": hidden_states, - "position_ids": None, - "layer_idx": "0", - "latent_cache_gen": None, - "_all_bases": (output,), - "_output_base_index": 0, - "_sparse_output_base_index": None, - "_sparse_output_sf_base_index": None, - }, - ) - optional_output = graph.call_function(getitem, args=(functionalized, 2)) - clone = graph.call_function(torch.ops.aten.clone.default, args=(optional_output,)) - graph.output(clone) - - monkeypatch.setattr( - remove_copy_pass, - "inplace_info", - lambda: {inplace_func: {1: "output", 2: "sparse_output"}}, - ) - - with pytest.raises( - AssertionError, - match=( - "getitem user for optional output 'sparse_output' has no " - "base tensor -- graph is malformed" - ), - ): - remove_copy_pass.remove_copy_for_mutates_args(graph) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 03d76b437123..14563d7215f4 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -242,6 +242,7 @@ def _make_forward_only_engine( outputs = {"logits": object()} engine._forward_step = Mock(return_value=outputs) engine._execute_logit_post_processors = Mock() + engine.breakable_cuda_graph_runner = None runner = Mock() runner.enabled = runner_enabled From 003a9131bf6837dc1e2efb17297c87e6966af358 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Tue, 4 Aug 2026 06:24:08 +0000 Subject: [PATCH 13/23] [None][fix] avoid implicit legacy PCG bucket conflicts Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- tensorrt_llm/llmapi/llm_args.py | 28 +++++++---------- tests/unittest/llmapi/test_llm_args.py | 43 ++++++++++++-------------- 2 files changed, 32 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0a2f12bed46a..514117b7f076 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5032,13 +5032,12 @@ class TorchCompileConfig(StrictBaseModel): enable_piecewise_cuda_graph: bool = Field( default=False, - description="Enable piecewise CUDA graph in torch.compile.") + description="Deprecated. Use prefill_cuda_graph_backend='piecewise' " + "instead.") capture_num_tokens: Optional[List[PositiveInt]] = Field( default=None, - description= - "List of num of tokens to capture the piecewise CUDA graph for. If not provided, the number of tokens will be the same as cuda_graph_config.batch_sizes." - ) + description="Deprecated. Use prefill_capture_num_tokens instead.") @field_validator('capture_num_tokens') @classmethod @@ -5057,12 +5056,6 @@ def validate_capture_num_tokens(cls, v): description= "The maximum number of CUDA streams to use for torch.compile.") - @model_validator(mode='after') - def set_default_capture_num_tokens(self) -> 'TorchCompileConfig': - if self.enable_piecewise_cuda_graph and self.capture_num_tokens is None: - self.capture_num_tokens = list(_DEFAULT_PREFILL_CAPTURE_NUM_TOKENS) - return self - class TorchLlmArgs(BaseLlmArgs): # PyTorch backend specific configurations @@ -5624,6 +5617,9 @@ def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': backend_is_explicit = "prefill_cuda_graph_backend" in self.model_fields_set buckets_are_explicit = "prefill_capture_num_tokens" in self.model_fields_set compile_config = self.torch_compile_config + legacy_buckets_are_explicit = (compile_config is not None + and "capture_num_tokens" + in compile_config.model_fields_set) if compile_config is not None and compile_config.enable_piecewise_cuda_graph: if (backend_is_explicit and self.prefill_cuda_graph_backend @@ -5638,18 +5634,18 @@ def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': legacy_buckets = (compile_config.capture_num_tokens if compile_config is not None else None) - if legacy_buckets is not None: - if (buckets_are_explicit + if legacy_buckets_are_explicit: + logger.warning( + "TorchCompileConfig.capture_num_tokens is deprecated; use " + "prefill_capture_num_tokens instead.") + if (legacy_buckets is not None and buckets_are_explicit and self.prefill_capture_num_tokens is not None and sorted(set(legacy_buckets)) != sorted( set(self.prefill_capture_num_tokens))): raise ValueError( "torch_compile_config.capture_num_tokens conflicts with " "prefill_capture_num_tokens") - if not buckets_are_explicit: - logger.warning( - "TorchCompileConfig.capture_num_tokens is deprecated; use " - "prefill_capture_num_tokens instead.") + if not buckets_are_explicit and legacy_buckets is not None: self.prefill_capture_num_tokens = list(legacy_buckets) if self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 627a69664213..573187bb921c 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1936,13 +1936,13 @@ class TestPiecewiseCudaGraphCaptureDefaults: Three invariants are exercised: - 1. `TorchCompileConfig.capture_num_tokens` defaults to a fixed - powers-of-2 + 256-stride list when `enable_piecewise_cuda_graph` - is True (and stays `None` otherwise). The fixed list keeps the - capture set small to bound startup time and CUDA graph memory; - the model-engine filter (invariants 2 and 3) clamps out-of-range - entries to the reachable ceiling and never invents sizes beyond - this list. + 1. `TorchLlmArgs.prefill_capture_num_tokens` defaults to a fixed + powers-of-2 + 256-stride list when a prefill CUDA graph backend is + enabled. The deprecated `TorchCompileConfig.capture_num_tokens` stays + `None` unless explicitly set. The fixed list keeps the capture set small + to bound startup time and CUDA graph memory; the model-engine filter + (invariants 2 and 3) clamps out-of-range entries to the reachable ceiling + and never invents sizes beyond this list. 2. `_filter_piecewise_capture_num_tokens` caps the candidate list at `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` -- the largest forward-pass `num_tokens` the warmup builder can @@ -1988,6 +1988,14 @@ def test_legacy_piecewise_config_maps_to_new_fields(self): assert args.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE assert args.prefill_capture_num_tokens == [256, 128] + def test_explicit_new_buckets_with_legacy_piecewise_enable(self): + args = TorchLlmArgs(model=llama_model_path, + prefill_capture_num_tokens=[128, 256], + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True)) + assert args.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE + assert args.prefill_capture_num_tokens == [128, 256] + def test_explicit_legacy_and_new_config_conflicts(self): with pytest.raises(ValueError, match="conflicts"): TorchLlmArgs( @@ -2067,18 +2075,10 @@ def tp_allgather(self, value): assert engine._get_padding_params( 1, 0, all_rank_num_tokens) == (1, False, all_rank_num_tokens) - def test_torch_compile_config_capture_num_tokens_default_when_piecewise_enabled( + def test_torch_compile_config_does_not_populate_legacy_capture_buckets( self): - """Default capture set is the powers-of-2 + 256-stride list. - - Keeps the capture set bounded (~20 entries) so server startup - time and CUDA graph memory stay predictable. The model engine - further filters and appends the reachable ceiling, so - out-of-range entries (e.g. > max_seq_len-1) are never recorded - and gap ISLs still get a graph. - """ config = TorchCompileConfig(enable_piecewise_cuda_graph=True) - assert config.capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert config.capture_num_tokens is None def test_torch_compile_config_capture_num_tokens_stays_none_when_piecewise_disabled( self): @@ -2099,12 +2099,8 @@ def test_torch_compile_config_capture_num_tokens_user_override_preserved( # `validate_capture_num_tokens` dedupes and reverse-sorts. assert config.capture_num_tokens == sorted(set(user_list), reverse=True) - def test_torch_llm_args_capture_num_tokens_default_when_piecewise_enabled( + def test_torch_llm_args_prefill_buckets_default_when_piecewise_enabled( self): - """Same default applies when reached through `TorchLlmArgs` construction. - - This is the path real users hit via `trtllm-serve` YAML. - """ args = TorchLlmArgs( model=llama_model_path, max_batch_size=1, @@ -2116,7 +2112,8 @@ def test_torch_llm_args_capture_num_tokens_default_when_piecewise_enabled( torch_compile_config=TorchCompileConfig( enable_piecewise_cuda_graph=True), ) - assert args.torch_compile_config.capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert args.prefill_capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert args.torch_compile_config.capture_num_tokens is None def test_piecewise_filter_never_invents_far_ceiling(self): """A ceiling far above the largest candidate is NOT added. From 85f10f78fde5246f1a68cbf10bf87b71f30a0837 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Tue, 4 Aug 2026 06:24:46 +0000 Subject: [PATCH 14/23] [None][fix] restrict eager graph captured values Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../breakable_cuda_graph.py | 8 +++----- tensorrt_llm/_torch/utils.py | 20 +++++-------------- .../executor/test_breakable_cuda_graph.py | 19 +++++++++--------- 3 files changed, 18 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index 790cbce368f6..2cfc4593219e 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -161,11 +161,9 @@ def wrapper(*args, **kwargs): capture._end_current_segment() output = inner(*args, **kwargs) - captured_args = tuple(make_weak_ref(arg, preserve_unsupported=True) for arg in args) - captured_kwargs = { - key: make_weak_ref(value, preserve_unsupported=True) for key, value in kwargs.items() - } - captured_output = make_weak_ref(output, preserve_unsupported=True) + captured_args = tuple(make_weak_ref(arg) for arg in args) + captured_kwargs = {key: make_weak_ref(value) for key, value in kwargs.items()} + captured_output = make_weak_ref(output) def replay_fn() -> Any: new_output = inner(*captured_args, **captured_kwargs) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index cde74352f7ae..2697b60f4077 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -159,29 +159,19 @@ def wrapper(self, *args, **kwargs): return decorator -def make_weak_ref(x, preserve_unsupported: bool = False): +def make_weak_ref(x): if isinstance(x, torch.Tensor): return convert_to_torch_tensor( TensorWrapper(x.data_ptr(), x.dtype, x.shape, x.stride())) if x.is_cuda else x elif isinstance(x, tuple): - return tuple( - make_weak_ref(i, preserve_unsupported=preserve_unsupported) - for i in x) + return tuple(make_weak_ref(i) for i in x) elif isinstance(x, list): - return [ - make_weak_ref(i, preserve_unsupported=preserve_unsupported) - for i in x - ] + return [make_weak_ref(i) for i in x] elif isinstance(x, dict): - return { - k: make_weak_ref(v, preserve_unsupported=preserve_unsupported) - for k, v in x.items() - } - elif isinstance(x, (int, float, bool)): - return x - elif preserve_unsupported: + return {make_weak_ref(k): make_weak_ref(v) for k, v in x.items()} + elif x is None or isinstance(x, (int, float, str, bool)): return x else: raise TypeError(f"Invalid type {type(x)} to make weak ref") diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index c2fbbab105b2..1c342aa61875 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -91,17 +91,16 @@ def add_one(value): torch.testing.assert_close(compiled_add_one(value), value + 1) -def test_make_weak_ref_option_preserves_unsupported_values(): +def test_make_weak_ref_supports_value_types_and_rejects_objects(): unsupported = object() with pytest.raises(TypeError, match="Invalid type"): make_weak_ref(unsupported) - captured = make_weak_ref( - {"nested": (unsupported, [None, "value"])}, - preserve_unsupported=True, - ) - assert captured == {"nested": (unsupported, [None, "value"])} - assert captured["nested"][0] is unsupported + value = {"nested": (None, "value", [1, 2.0, True])} + assert make_weak_ref(value) == value + + with pytest.raises(TypeError, match="Invalid type"): + make_weak_ref({unsupported: "value"}) def test_break_graph_inserts_empty_breakpoint(): @@ -222,12 +221,14 @@ def engine_forward(): def test_runner_first_bucket_segments_share_one_memory_pool(): class BreakableBody(nn.Module): + @staticmethod @eager_on_graph - def eager_add_one(self, value): + def eager_add_one(value): return value + 1 + @staticmethod @eager_on_graph - def eager_double(self, value): + def eager_double(value): return value * 2 def forward(self, value): From 23318fb903359f6c79a35cf9c92b36461f306e01 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Tue, 4 Aug 2026 07:01:08 +0000 Subject: [PATCH 15/23] reorganize tests Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../defs/accuracy/test_llm_api_pytorch.py | 83 ++++++++++++++++++- .../test_lists/test-db/l0_dgx_b200.yml | 4 +- 2 files changed, 83 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index e81f27b165f5..19312e402d0c 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3603,9 +3603,8 @@ def test_nvfp4_multi_gpus_piecewise_cuda_graph(self, tp_size, pp_size, "tp_size,pp_size,ep_size,mtp_nextn,attention_dp,max_batch_size,moe_backend,fp8kv,chunked_prefill", [ (8, 1, 8, 0, True, 24, "CUTLASS", False, False), - (8, 1, 8, 3, False, 16, "TRTLLM", True, True), ], - ids=["baseline", "mtp3_fp8kv_chunked"]) + ids=["baseline"]) def test_nvfp4_multi_gpus_breakable_cuda_graph(self, tp_size, pp_size, ep_size, mtp_nextn, attention_dp, max_batch_size, @@ -3995,6 +3994,86 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): eplb_config, mtp_nextn=mtp_nextn) + @pytest.mark.skip_less_mpi_world_size(8) + @pytest.mark.threadleak(enabled=False) + def test_mixed_breakable_cuda_graph(self): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + base_prompt_ids = tokenizer.encode( + "TensorRT-LLM accelerates reliable large language model inference " + "with efficient attention, parallelism, and CUDA graphs. ", + add_special_tokens=False, + ) + assert base_prompt_ids + + def make_prompt(prompt_length): + return (base_prompt_ids * + ((prompt_length + len(base_prompt_ids) - 1) // + len(base_prompt_ids)))[:prompt_length] + + generation_prompt = make_prompt(64) + context_prompt = make_prompt(129) + sampling_params = SamplingParams( + max_tokens=8, + min_tokens=8, + seed=42, + temperature=0, + ignore_eos=True, + detokenize=False, + add_special_tokens=False, + ) + common_llm_kwargs = dict( + tensor_parallel_size=8, + moe_expert_parallel_size=8, + moe_config=MoeConfig(backend="TRTLLM"), + enable_attention_dp=True, + max_batch_size=8, + max_num_tokens=1024, + max_seq_len=2048, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + dtype="fp8", + free_gpu_memory_fraction=0.6, + ), + cuda_graph_config=CudaGraphConfig( + batch_sizes=[1, 2, 4, 6, 8], + enable_padding=True, + ), + ) + + def run(backend): + with LLM( + self.MODEL_PATH, + **common_llm_kwargs, + prefill_cuda_graph_backend=backend, + prefill_capture_num_tokens=[128, 256, 512, 1024], + ) as llm: + generation_request = llm.generate_async( + generation_prompt, + sampling_params=sampling_params, + streaming=True, + ) + next(generation_request) + assert not generation_request.finished + + # Admit a context request while the first request is decoding. + context_request = llm.generate_async( + context_prompt, + sampling_params=sampling_params, + streaming=False, + ) + generation_output = generation_request.result() + context_output = context_request.result() + return [ + generation_output.outputs[0].token_ids, + context_output.outputs[0].token_ids, + ] + + eager_token_ids = run(PrefillCudaGraphBackend.DISABLED) + breakable_token_ids = run(PrefillCudaGraphBackend.BREAKABLE) + assert breakable_token_ids == eager_token_ids + _DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT = ( "Solve the problem carefully. End your response with a final line exactly " diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 0780ae814890..089c9cb70d04 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -176,14 +176,13 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[baseline] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] TIMEOUT (60) - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[baseline] TIMEOUT (60) - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[mtp3_fp8kv_chunked] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4_mtp_index_share[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_mixed_breakable_cuda_graph TIMEOUT (120) ISOLATION - examples/test_deepseek_v4_pro.py::test_short_token_boundary_smoke TIMEOUT (120) - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) @@ -256,6 +255,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp3] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_pp4_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[baseline_fp8kv] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[baseline] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8_attn_dp] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[ep8] TIMEOUT (60) From 2eaa26669aa0d08f8794d9fd10a0b70831a2b97e Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 5 Aug 2026 03:12:47 +0000 Subject: [PATCH 16/23] [None][test] provide valid prepared graph inputs Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../_torch/executor/test_pytorch_model_engine.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 14563d7215f4..a28b3ad79c7e 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -238,7 +238,11 @@ def _make_forward_only_engine( ) engine.spec_metadata = spec_metadata engine._set_up_spec_metadata = Mock(return_value=spec_metadata) - engine._prepare_inputs = Mock(return_value=({"prepared": True}, None)) + prepared_inputs = { + "prepared": True, + "input_ids": torch.zeros(2, dtype=torch.int32), + } + engine._prepare_inputs = Mock(return_value=(prepared_inputs, None)) outputs = {"logits": object()} engine._forward_step = Mock(return_value=outputs) engine._execute_logit_post_processors = Mock() @@ -762,7 +766,8 @@ def test_forward_commits_candidate_only_on_graph_hit(self) -> None: prepare_args = engine._prepare_inputs.call_args.args self.assertIs(prepare_args[0], graph_batch) self.assertEqual(prepare_args[-1], frozenset({1})) - runner.replay.assert_called_once_with(key, {"prepared": True}) + prepared_inputs = engine._prepare_inputs.return_value[0] + runner.replay.assert_called_once_with(key, prepared_inputs) engine._forward_step.assert_not_called() engine._execute_logit_post_processors.assert_called_once_with( batch, outputs) @@ -836,7 +841,8 @@ def test_zero_runtime_draft_speculation_commits_graph_candidate( self.assertEqual( semantic_attn_metadata.update_spec_dec_param.call_args. kwargs["num_contexts"], 1) - runner.replay.assert_called_once_with(key, {"prepared": True}) + prepared_inputs = engine._prepare_inputs.return_value[0] + runner.replay.assert_called_once_with(key, prepared_inputs) def test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager( self) -> None: @@ -917,7 +923,8 @@ def test_forward_allows_guided_context_logits_on_graph_hit(self) -> None: prepare_args = engine._prepare_inputs.call_args.args self.assertIs(prepare_args[0], graph_batch) self.assertEqual(prepare_args[-1], frozenset({context.py_request_id})) - runner.replay.assert_called_once_with(key, {"prepared": True}) + prepared_inputs = engine._prepare_inputs.return_value[0] + runner.replay.assert_called_once_with(key, prepared_inputs) def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: engine, runner, resource_manager, _, _ = _make_forward_only_engine(None) From f4cc607cef53d9e1deb75019dafa07dab8c3830d Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:43:57 +0000 Subject: [PATCH 17/23] [None][fix] keep prefill warmup ranks aligned Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 607bddc6f7bd..020c01799cd2 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2486,6 +2486,8 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): resource_manager, num_tokens, 0) with self._release_batch_context(warmup_request, resource_manager) as batch: + self._assert_all_tp_ranks_have_warmup_batch( + batch, num_tokens) if batch is None: continue @@ -2519,6 +2521,7 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): least_requests=False) with self._release_batch_context(warmup_request, resource_manager) as batch: + self._assert_all_tp_ranks_have_warmup_batch(batch, num_tokens) if batch is None: continue logger.info( From eaafc4e91db04db14ae0920b2c784333dcd44428 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:48:07 +0000 Subject: [PATCH 18/23] [None][fix] simplify breakable CUDA graph helpers Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../_torch/modules/mamba/gdn_mixer.py | 9 ++----- .../breakable_cuda_graph.py | 5 ++-- .../breakable_cuda_graph/cuda_utils.py | 25 ------------------- 3 files changed, 5 insertions(+), 34 deletions(-) delete mode 100644 tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index f510e0da8342..331bfe442976 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -175,7 +175,7 @@ def gdn_custom_op_inplace( ) -breakable_gdn_custom_op_inplace = eager_on_graph(gdn_custom_op_inplace) +maybe_bcg_gdn_custom_op_inplace = eager_on_graph(gdn_custom_op_inplace) def ensure_divisibility(numerator, denominator): @@ -1062,12 +1062,7 @@ def forward( attn_out = mixed_qkv.new_empty( (1, mixed_qkv.shape[0], self.num_v_heads_per_tp, self.head_v_dim) ) - custom_op = ( - breakable_gdn_custom_op_inplace - if use_breakable_cuda_graph - else gdn_custom_op_inplace - ) - custom_op(mixed_qkv, a, b, self.layer_idx_str, attn_out) + maybe_bcg_gdn_custom_op_inplace(mixed_qkv, a, b, self.layer_idx_str, attn_out) else: attn_out = self.forward_core( mixed_qkv, diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index 2cfc4593219e..3b429c3856f3 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -12,8 +12,9 @@ import torch from cuda.bindings import runtime as rt +from tensorrt_llm._utils import CUASSERT + from ...utils import make_weak_ref -from .cuda_utils import check_cuda_errors logger = logging.getLogger(__name__) @@ -56,7 +57,7 @@ def get_current_replay_token() -> Optional[int]: def _capture_status(stream_ptr: int) -> rt.cudaStreamCaptureStatus: - status, *_ = check_cuda_errors(rt.cudaStreamGetCaptureInfo(stream_ptr)) + status, *_ = CUASSERT(rt.cudaStreamGetCaptureInfo(stream_ptr)) return status diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py deleted file mode 100644 index 6c56ba287218..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py +++ /dev/null @@ -1,25 +0,0 @@ -# Adapted from SGLang's breakable CUDA graph implementation. -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from cuda.bindings import runtime as rt - - -def _cuda_get_error_string(error: rt.cudaError_t) -> str: - result, message = rt.cudaGetErrorString(error) - if result != rt.cudaError_t.cudaSuccess: - return "" - if isinstance(message, bytes): - return message.decode("utf-8", "replace") - return str(message) - - -def check_cuda_errors(result): - """Raise a Python exception for a failed cuda-python runtime call.""" - if result[0] != rt.cudaError_t.cudaSuccess: - raise RuntimeError(f"CUDA error {int(result[0])}({_cuda_get_error_string(result[0])})") - if len(result) == 1: - return None - if len(result) == 2: - return result[1] - return result[1:] From d1a650df524fb14994156ca1501353a9eff4776d Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Sat, 8 Aug 2026 08:16:24 -0700 Subject: [PATCH 19/23] chore: apply post-rebase formatting Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../sparse/deepseek_v4/module.py | 16 ++++------------ tensorrt_llm/_torch/models/modeling_minimaxm3.py | 4 +--- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 381dc7fc80c5..9abe0f7a77a4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -410,9 +410,7 @@ def forward_generation_sparse_attn( o_lora_bmm_input_scale = None inverse_rope_cos_sin = None if enable_dsv4_epilogue_fusion: - dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers( - self, q, num_tokens - ) + dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers(self, q, num_tokens) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -494,9 +492,7 @@ def forward_context_sparse_attn( o_lora_bmm_input_scale = None inverse_rope_cos_sin = None if enable_dsv4_epilogue_fusion: - dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers( - self, q, num_tokens - ) + dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers(self, q, num_tokens) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -820,12 +816,8 @@ def _indexer_branch(): ) if enable_dsv4_epilogue_fusion: - assert context_o_lora_bmm_input is None or isinstance( - context_o_lora_bmm_input, tuple - ) - assert generation_o_lora_bmm_input is None or isinstance( - generation_o_lora_bmm_input, tuple - ) + assert context_o_lora_bmm_input is None or isinstance(context_o_lora_bmm_input, tuple) + assert generation_o_lora_bmm_input is None or isinstance(generation_o_lora_bmm_input, tuple) # The fused kernel output is group-first, which BCG cannot slice on # dim 0. Write O-LoRA as token-first so replay can slice the bucket. _run_dsv4_o_lora_bmms( diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index b08a1db57087..f9c2ed8dbe42 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1325,9 +1325,7 @@ def _forward_attention_core( output = q.new_empty( (q.shape[0], self.num_heads * self.head_dim), dtype=self.attn_activation_dtype ) - if self.register_to_config and ( - is_torch_compiling() or is_in_breakable_cuda_graph() - ): + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): maybe_bcg_minimax_m3_attn_custom_op_inplace( q, k, From 187964c370842f2b9f1c074e1548f6fad0d277c9 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Mon, 10 Aug 2026 07:43:37 +0000 Subject: [PATCH 20/23] [None][fix] address breakable CUDA graph review feedback Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../torch_compile_and_piecewise_cuda_graph.md | 19 +++++++-- .../breakable_cuda_graph.py | 6 +++ .../_torch/pyexecutor/model_engine.py | 9 ++++ tensorrt_llm/llmapi/llm_args.py | 14 ++++--- .../executor/test_breakable_cuda_graph.py | 18 ++++++++ .../executor/test_pytorch_model_engine.py | 41 +++++++++++++++++++ tests/unittest/llmapi/test_llm_args.py | 24 +++++++++++ 7 files changed, 121 insertions(+), 10 deletions(-) diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index c7e0167d39f4..6d3ee0eceaf7 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -59,10 +59,21 @@ prefill_cuda_graph_backend: breakable prefill_capture_num_tokens: [128, 256, 512] ``` -The first version of the breakable backend supports BF16 Qwen3.5 on one GPU for -context-only, tensor/pipeline parallelism and mixed context/decode batches with KV cache. Speculative -decoding, LoRA, multimodal inputs, and context -logits fall back to eager execution or are rejected during initialization. +The breakable backend is experimental. The integration coverage in this change +includes BF16 Qwen3.5 on one GPU and NVFP4 DeepSeek models on multiple GPUs, +with both context-only and mixed context/decode batches using the KV cache. + +The following restrictions are enforced: + +- `torch_compile_config`, LoRA, and multimodal models are rejected during + engine initialization. +- Speculative decoding is supported. +- Context-logit requests run eagerly instead of replaying a breakable CUDA + graph. + +Other model families, quantization modes, and parallel configurations are not +yet covered by this experimental backend's integration tests and should be +validated before use. ## Tips for Piecewise CUDA Graph diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index 3b429c3856f3..83e2b27e25fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -81,6 +81,12 @@ def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream) -> No other_is_capture = other is capturing or other.cuda_stream == capture_ptr if self_is_capture and not other_is_capture: if not _is_stream_capturing(other): + logger.warning( + "Dropping a wait from the breakable CUDA graph capture stream " + "to a non-capturing stream; this dependency will not be " + "preserved during replay.", + stacklevel=2, + ) return _original_wait_stream(self, other) forked.discard(other) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 020c01799cd2..f0de51782173 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -491,6 +491,7 @@ def __init__( setattr(self, "moe_load_balancer", moe_load_balancer) else: self.model = model + self._validate_breakable_cuda_graph_compatibility() pretrained_config = self.model.model_config.pretrained_config model_type = getattr(pretrained_config, "model_type", None) self._enable_dsv4_adp_dummy_fixes = should_enable_dsv4_adp_dummy_fixes( @@ -3125,6 +3126,13 @@ def is_multimodal(self) -> bool: return True return isinstance(self.input_processor, BaseMultimodalInputProcessor) + def _validate_breakable_cuda_graph_compatibility(self) -> None: + if (self.llm_args.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.BREAKABLE and self.is_multimodal): + raise ValueError( + "breakable prefill CUDA graph does not support multimodal models" + ) + def _set_up_multimodal_encoder_attn_metadata(self) -> None: """Construct AttentionMetadata for any multimodal encoders inside the loaded model, using the engine's encoder runtime sizes @@ -7210,6 +7218,7 @@ def forward_step(): can_run_breakable_graph = ( breakable_runner is not None and get_per_request_prefill_cuda_graph_flag() + and not gather_context_logits and breakable_runner.has_graph(num_tokens)) if can_run_breakable_graph and not breakable_runner.is_warming_up: outputs = breakable_runner.execute( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 514117b7f076..131d1302c30b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5659,12 +5659,14 @@ def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE: if self.torch_compile_config is None: self.torch_compile_config = TorchCompileConfig() - elif (self.prefill_cuda_graph_backend - == PrefillCudaGraphBackend.BREAKABLE - and self.torch_compile_config is not None): - raise ValueError( - "breakable prefill CUDA graph does not support torch_compile_config" - ) + elif self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: + if self.torch_compile_config is not None: + raise ValueError( + "breakable prefill CUDA graph does not support torch_compile_config" + ) + if self.enable_lora or self.lora_config is not None: + raise ValueError( + "breakable prefill CUDA graph does not support LoRA") return self diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index 1c342aa61875..70d6fdbddfc1 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 +import logging + import pytest import torch from torch import nn @@ -161,6 +163,22 @@ def body(): torch.testing.assert_close(output, torch.full_like(output, 8)) +def test_dropped_wait_stream_logs_callsite(caplog): + side_stream = torch.cuda.Stream() + + def body(): + torch.cuda.current_stream().wait_stream(side_stream) + + with caplog.at_level( + logging.WARNING, + logger=("tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph"), + ): + _capture(body) + + record = next(record for record in caplog.records if "Dropping a wait" in record.message) + assert record.funcName == "body" + + class _Body(nn.Module): def __init__(self): super().__init__() diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index a28b3ad79c7e..8f33ced306e9 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -25,6 +25,7 @@ PyTorchModelEngine, _build_request_multimodal_input, _make_single_token_context_graph_batch) from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchLlmArgs) @@ -954,6 +955,35 @@ def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: self.assertIs(context.py_multimodal_data, multimodal_data) self.assertIn("multimodal_embedding", multimodal_data) + def test_breakable_graph_falls_back_for_context_logits(self) -> None: + engine, _, resource_manager, _, outputs = \ + _make_forward_only_engine(None) + breakable_runner = Mock() + breakable_runner.is_capturing = False + breakable_runner.is_warming_up = False + breakable_runner.has_graph.return_value = True + breakable_runner.execute.return_value = outputs + engine.breakable_cuda_graph_runner = breakable_runner + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [_make_request_stub(1)] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock() + ), patch( + "tensorrt_llm._torch.pyexecutor.model_engine.get_per_request_prefill_cuda_graph_flag", + return_value=True): + actual_outputs = engine.forward( + batch, + resource_manager, + gather_context_logits=True, + ) + + self.assertIs(actual_outputs, outputs) + breakable_runner.execute.assert_not_called() + engine._forward_step.assert_called_once() + def test_generation_only_forward_does_not_call_new_selector(self) -> None: key = KeyType(batch_size=1, draft_len=0, is_first_draft=False) engine, runner, resource_manager, _, _ = _make_forward_only_engine(key) @@ -1091,6 +1121,17 @@ def test_encoder_cuda_graph_stages_and_restores_fixed_sequence_slots( (fixed_slot_output[511:512], fixed_slot_output[:400])) torch.testing.assert_close(restored_output, expected_output) + def test_breakable_rejects_multimodal_models(self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = DummyLegacyMultimodalIndexModel() + engine.input_processor = None + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE) + + self.assertTrue(engine.is_multimodal) + with self.assertRaisesRegex(ValueError, "multimodal models"): + engine._validate_breakable_cuda_graph_compatibility() + def test_prepare_multimodal_indices_uses_mixin_token_ids(self) -> None: engine = object.__new__(PyTorchModelEngine) engine.model = DummyMultimodalIndexModel() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 573187bb921c..0a2aeb12c5f2 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2019,6 +2019,30 @@ def test_breakable_rejects_explicit_torch_compile(self): prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, torch_compile_config=TorchCompileConfig()) + def test_breakable_allows_speculative_decoding(self): + speculative_config = MTPDecodingConfig(max_draft_len=1) + args = TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + speculative_config=speculative_config) + assert args.speculative_config == speculative_config + + @pytest.mark.parametrize("lora_kwargs", [ + { + "enable_lora": True + }, + { + "lora_config": LoraConfig(lora_target_modules=["attn_q"]) + }, + ]) + def test_breakable_rejects_lora(self, lora_kwargs): + with pytest.raises(ValueError, match="LoRA"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + **lora_kwargs, + ) + def test_prefill_filter_sorts_dedupes_and_drops_nonpositive(self): from tensorrt_llm._torch.pyexecutor.model_engine import \ _filter_prefill_capture_num_tokens From a6569453b36dc85950028f8ad6f2189f0a53ec6b Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Mon, 10 Aug 2026 09:50:03 +0000 Subject: [PATCH 21/23] [None][fix] update prefill graph flag after rebase Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f0de51782173..1a9375068182 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4056,7 +4056,7 @@ def _prepare_encoder_decoder_inputs_fast( attn_all_rank_num_tokens) = self._get_padding_params( total_num_tokens, scheduled_requests.num_context_requests, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + set_per_request_prefill_cuda_graph_flag(can_run_piecewise_cuda_graph) attn_metadata.padded_num_tokens = (padded_num_tokens if padded_num_tokens != total_num_tokens else None) From efa156ab288e1fa8705b73af495e983b1db46696 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:16:33 +0000 Subject: [PATCH 22/23] [None][fix] repair breakable graph CI regressions Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../sparse/deepseek_v4/module.py | 2 ++ .../_torch/pyexecutor/model_engine.py | 14 ++++++-- .../defs/accuracy/test_llm_api_pytorch.py | 1 + .../deepseek_v4/test_deepseek_v4_o_proj.py | 6 ++-- .../executor/test_pytorch_model_engine.py | 33 +++++++++++++++++-- 5 files changed, 49 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 9abe0f7a77a4..58625b057419 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -271,6 +271,8 @@ def project_sparse_attn_output( ) -> torch.Tensor: del attn_metadata, all_reduce_params attn_output_tensor = attn_output[0] + # BCG/mixed-batch epilogue fusion runs o_a_proj at the end of attention, + # so this 3D tensor is O-LoRA output and only o_b_proj remains. if attn_output_tensor.ndim == 3: return self.o_b_proj(attn_output_tensor.flatten(1)) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 1a9375068182..a7c790f481af 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3127,8 +3127,18 @@ def is_multimodal(self) -> bool: return isinstance(self.input_processor, BaseMultimodalInputProcessor) def _validate_breakable_cuda_graph_compatibility(self) -> None: - if (self.llm_args.prefill_cuda_graph_backend - == PrefillCudaGraphBackend.BREAKABLE and self.is_multimodal): + if self.llm_args.prefill_cuda_graph_backend != PrefillCudaGraphBackend.BREAKABLE: + return + + if isinstance(self.model, DecoderModelForCausalLM): + return + decoder_model = getattr(self.model, "llm", None) + if (self.llm_args.disable_mm_encoder + and isinstance(decoder_model, DecoderModelForCausalLM) + and getattr(self.model, "mm_encoder", None) is None): + return + if (isinstance(self.model, MultimodalModelMixin) or isinstance( + self.input_processor, BaseMultimodalInputProcessor)): raise ValueError( "breakable prefill CUDA graph does not support multimodal models" ) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 19312e402d0c..4cf15a585ae9 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6365,6 +6365,7 @@ def run(backend): max_num_tokens=512, max_batch_size=4, disable_overlap_scheduler=True, + disable_mm_encoder=True, kv_cache_config=self.kv_cache_config, cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=4), diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py index 196c06531831..887e5f66555c 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -276,9 +276,9 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): attn_out_latent = torch.randn(num_tokens, num_heads, qk_head_dim, dtype=dtype, device=device) position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) - # Call the deepseek_v4 output projection (mla_rope_inplace modifies attn_out_latent - # in-place, so clone before passing to preserve original for reference) - output = project_sparse_attn_output(mla, [attn_out_latent.clone()], position_ids) + # The non-fused MLA path stores attention output as a flattened 2D buffer. + # mla_rope_inplace modifies it in place, so preserve the 3D reference input. + output = project_sparse_attn_output(mla, [attn_out_latent.clone().flatten(1)], position_ids) # Calculate reference output if dtype_str == "bf16": diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 8f33ced306e9..1832adb40477 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -15,6 +15,7 @@ MultimodalEncoderMixin from tensorrt_llm._torch.models.modeling_multimodal_mixin import \ MultimodalModelMixin +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( @@ -42,7 +43,8 @@ from tensorrt_llm._torch.speculative.spec_sampler_base import \ SampleStateTensorsSpec from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.inputs.registry import BaseMultimodalDummyInputsBuilder +from tensorrt_llm.inputs.registry import (BaseMultimodalDummyInputsBuilder, + BaseMultimodalInputProcessor) from tensorrt_llm.llmapi import (CudaGraphConfig, SADecodingConfig, SamplingParams) from tensorrt_llm.mapping import CpType, Mapping @@ -1126,12 +1128,39 @@ def test_breakable_rejects_multimodal_models(self) -> None: engine.model = DummyLegacyMultimodalIndexModel() engine.input_processor = None engine.llm_args = SimpleNamespace( - prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE) + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + disable_mm_encoder=False) self.assertTrue(engine.is_multimodal) with self.assertRaisesRegex(ValueError, "multimodal models"): engine._validate_breakable_cuda_graph_compatibility() + def test_breakable_allows_text_decoder_with_multimodal_processor( + self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = Mock(spec=DecoderModelForCausalLM) + engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + disable_mm_encoder=False) + + self.assertTrue(engine.is_multimodal) + engine._validate_breakable_cuda_graph_compatibility() + + def test_breakable_allows_multimodal_wrapper_in_text_only_mode( + self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = DummyLegacyMultimodalIndexModel() + engine.model.llm = Mock(spec=DecoderModelForCausalLM) + engine.model.mm_encoder = None + engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + disable_mm_encoder=True) + + self.assertTrue(engine.is_multimodal) + engine._validate_breakable_cuda_graph_compatibility() + def test_prepare_multimodal_indices_uses_mixin_token_ids(self) -> None: engine = object.__new__(PyTorchModelEngine) engine.model = DummyMultimodalIndexModel() From 4bf70af0def2fb5f9b63f7b71bd36930a736cdf6 Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:58:01 +0000 Subject: [PATCH 23/23] [None][fix] retain eager bridge output storage Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../breakable_cuda_graph.py | 2 +- .../executor/test_breakable_cuda_graph.py | 35 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py index 83e2b27e25fb..c70f02480927 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -170,7 +170,7 @@ def wrapper(*args, **kwargs): captured_args = tuple(make_weak_ref(arg) for arg in args) captured_kwargs = {key: make_weak_ref(value) for key, value in kwargs.items()} - captured_output = make_weak_ref(output) + captured_output = output def replay_fn() -> Any: new_output = inner(*captured_args, **captured_kwargs) diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index 70d6fdbddfc1..7c38f9e328ad 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -4,6 +4,7 @@ import logging +import weakref import pytest import torch @@ -73,6 +74,40 @@ def body(): torch.testing.assert_close(output, torch.full_like(output, 16)) +def test_eager_output_storage_survives_allocator_churn(): + eager_output_ptrs = [] + eager_output_refs = [] + + @eager_on_graph + def add_one(value): + output = value + 1 + eager_output_ptrs.append(output.data_ptr()) + eager_output_refs.append(weakref.ref(output)) + return output + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + x = torch.zeros(1024, device="cuda") + output = torch.zeros_like(x) + graph = BreakableCUDAGraph() + + with BreakableCUDAGraphCapture(graph, stream=stream): + output.copy_(add_one(x) * 2) + + captured_output_ptr = eager_output_ptrs[0] + assert eager_output_refs[0]() is not None + churn = [torch.full_like(x, value) for value in range(4)] + assert all(tensor.data_ptr() != captured_output_ptr for tensor in churn) + + x.fill_(5) + graph.replay() + + stream.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 12)) + for value, tensor in enumerate(churn): + torch.testing.assert_close(tensor, torch.full_like(tensor, value)) + + def test_outside_capture(): @eager_on_graph def outside(value):