From 603c8a0aa26b49e3d2fbb7576f0b5674bb648c4a 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/25] 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 284eee0b0574..4214325f5b72 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, add_lora_result +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 0b624cb6362d..33f28ca148bb 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: @@ -632,15 +634,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[ @@ -649,14 +650,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, @@ -664,7 +665,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))}. " @@ -689,7 +690,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, @@ -933,6 +934,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 @@ -2053,7 +2069,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) @@ -2061,7 +2078,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 @@ -2457,18 +2474,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, @@ -2477,26 +2500,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, @@ -2506,11 +2532,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 ### @@ -3343,6 +3378,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() @@ -3526,45 +3564,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" @@ -3574,19 +3608,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 @@ -5829,10 +5862,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 @@ -6105,9 +6150,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: @@ -7091,7 +7138,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, @@ -7215,14 +7261,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 bf1afd2b97f7..64d5b0f58c7f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5016,6 +5016,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( @@ -5055,8 +5067,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 @@ -5304,6 +5315,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= @@ -5600,6 +5625,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 2146edc5a267..6b079fd1092d 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 a65e2c70ef3e..e779f2cb35bd 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 @@ -6172,6 +6173,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 2d3d96f0aabb..8ab1e4b21b51 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -49,7 +49,8 @@ MambaStateConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, - PeftCacheConfig, PybindMirror, + PeftCacheConfig, + PrefillCudaGraphBackend, PybindMirror, RayPlacementConfig, SkipSoftmaxAttentionConfig, SleepConfig, SpeculativeConfig, @@ -1985,6 +1986,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 fe9372855ca4494cb2affc4b0709d2af45e756b6 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/25] 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 73164f885660..53ee6d35edf7 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 4214325f5b72..fac575173f44 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 33f28ca148bb..5e326ccef461 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 @@ -936,10 +937,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)) @@ -947,8 +944,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 @@ -5862,22 +5858,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 @@ -6152,9 +6137,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: @@ -6639,6 +6622,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: @@ -7262,16 +7246,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()): @@ -7282,20 +7256,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 64d5b0f58c7f..7581dfcd95d5 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5060,7 +5060,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 8ab1e4b21b51..3f773fbb4c05 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2066,6 +2066,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 383885617ed9bf2149cbc545f415054b3913b40d 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/25] [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 73d4da9a6dcd946c9b055130d2aab0469051e698 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/25] [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 f7b3e9fbb9d27ffd71f2175b19c33bfa479255e9 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/25] [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 e779f2cb35bd..2fa1b62d379e 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 @@ -3920,6 +3923,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 fbb0147bf527c3061291f40634f434595af859f5 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/25] 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 980525e645f3b5cf06766cb4f56145b4712ab0f3 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/25] 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 fac575173f44..b6fea37271ba 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 2f46386ad10d1fd8716c3bddbad49971a26b7c1a 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/25] [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 6511f0c43425..92c555a79396 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2337,6 +2337,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 2fa1b62d379e..6d00ed8f55a3 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 @@ -3586,6 +3584,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( @@ -3925,25 +3982,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 " @@ -3951,11 +3992,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, @@ -3965,7 +4009,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"), @@ -3982,120 +4026,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 bc6a2427170f..da4b7a2fb1b4 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 842127604988..e75a7a555f65 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -177,12 +177,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 73fb1ed30ba3e6fe7136e797c0d0505d4840dd13 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/25] [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 5e326ccef461..445cc20c576f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -944,7 +944,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 @@ -2472,8 +2473,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)): @@ -2533,10 +2533,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, @@ -3579,8 +3579,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: @@ -4686,7 +4685,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 @@ -7255,8 +7254,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 6d00ed8f55a3..884af7c31c02 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 @@ -3593,9 +3592,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 3f773fbb4c05..cb45e0239889 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2071,6 +2071,7 @@ def test_attention_dp_prefill_graph_uses_all_rank_decision(self): PyTorchModelEngine class FakeDist: + def __init__(self, decisions): self.decisions = decisions @@ -2086,8 +2087,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 8897e8762bba204fa0179dd4032abd846d10ebd4 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/25] [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 445cc20c576f..195e386c49e6 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 cb45e0239889..7fbc2191966c 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2049,7 +2049,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 06f13138507f55b8c956240173c4b2dc07d3e5b8 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/25] [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 92c555a79396..6511f0c43425 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2337,68 +2337,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 884af7c31c02..462f8c6db5ef 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3981,86 +3981,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 e75a7a555f65..c0be51973468 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -185,8 +185,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 6dacfc5efb1ba51d3d69c88de0a1ae700fdcafec 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/25] [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 84be2e7db3f1..8eeadcf7fb78 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 e4c21f1c66c6da6b7b8f26bcacf9738d3d288641 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/25] [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 7581dfcd95d5..fede158f0200 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5039,13 +5039,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 @@ -5064,12 +5063,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 @@ -5631,6 +5624,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 @@ -5645,18 +5641,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 7fbc2191966c..3a011343ff49 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1964,13 +1964,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 @@ -2016,6 +2016,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( @@ -2095,18 +2103,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): @@ -2127,12 +2127,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, @@ -2144,7 +2140,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 557406f222da486e8be24c46ff428c95460b6730 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/25] [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 12a879e19bb5bd3ad8c4b5aa4e38abc9bd26f186 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/25] 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 462f8c6db5ef..0f33f73f43aa 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3589,9 +3589,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, @@ -3981,6 +3980,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 c0be51973468..a46a74690d77 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -177,14 +177,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 5aff7d3ec88ea30227d06e4e6edb5e12c9e6bb40 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/25] [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 8eeadcf7fb78..5f9945ac1988 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() @@ -818,7 +822,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) @@ -892,7 +897,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: @@ -973,7 +979,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 89fe680b0a758c4b4c3c6da5cfa50b251467fe88 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/25] [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 195e386c49e6..1dcca90253c8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2496,6 +2496,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 @@ -2529,6 +2531,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 b401fd5ec2d582d302b3c451d4607c73fc415199 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/25] [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 9a3ff89e9ed1d167762fc1d1de9a8ee6202e5395 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/25] 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 ab3d2b75c3b973622e4388fcbab1f2ad12d33a32 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/25] [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 1dcca90253c8..8dd961def36c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -494,6 +494,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_scheduler_aware_adp_dummy = ( @@ -3135,6 +3136,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 @@ -7271,6 +7279,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 fede158f0200..f22963ba77b0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5666,12 +5666,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 5f9945ac1988..6ba61e7d6068 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) @@ -1010,6 +1011,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) @@ -1147,6 +1177,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 3a011343ff49..8fc215e4e99b 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2047,6 +2047,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 f328f84d51f0d9c24cbb87b30909bfa8617b1cf3 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/25] [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 8dd961def36c..e2f21226ca05 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4066,7 +4066,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 5ad889cf676b758fcfa3d6c441bc1ffb0d63f863 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/25] [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 e2f21226ca05..18579a909e11 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3137,8 +3137,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 0f33f73f43aa..3feb0e04bc0a 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6336,6 +6336,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 6ba61e7d6068..1de09f2595b8 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 @@ -1182,12 +1184,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 e4d0e14a808ef6bac9cc4e3f0828cf5b6b740125 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/25] [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): From 64104515bfa2df6e3a061b3d527fd8e17e14dc3d Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 12 Aug 2026 14:35:26 +0000 Subject: [PATCH 24/25] [None][fix] clean up failed BCG capture Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py | 6 +++++- .../_torch/pyexecutor/breakable_cuda_graph_runner.py | 1 + 2 files changed, 6 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 c70f02480927..e29147f92162 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 @@ -251,7 +251,11 @@ def __enter__(self) -> "BreakableCUDAGraphCapture": 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() + try: + self._begin_new_segment() + except Exception: + _uninstall_wait_stream_hook() + raise return self def __exit__(self, *args: object) -> bool: diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py index d1b8562ff568..9eccb529196d 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -108,6 +108,7 @@ def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None: graph.reset() if created_memory_pool and not self._graphs: self._memory_pool = None + self._shared_output = None raise finally: self._active_graph = None From e0164283da2dde3dce76ea8662902d6cdec6c46b Mon Sep 17 00:00:00 2001 From: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> Date: Wed, 12 Aug 2026 14:36:58 +0000 Subject: [PATCH 25/25] [None][docs] clarify DeepSeek-V4 output contract Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com> --- .../sparse/deepseek_v4/module.py | 2 ++ .../_torch/modules/ATTENTION_DEVELOPER_GUIDE.md | 16 ++++++++-------- 2 files changed, 10 insertions(+), 8 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 58625b057419..740932094880 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -548,6 +548,8 @@ 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] + # A 3D token-major output is the internal fusion marker, avoiding + # algorithm-specific parameters in the shared MLA custom-op schema. enable_dsv4_epilogue_fusion = output.ndim == 3 num_contexts = attn_metadata.num_contexts num_generations = attn_metadata.num_generations diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 10cc9a9ae872..44fa4ece561d 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -134,14 +134,14 @@ both. The separate adapter types make the `Attention` and `MLA` signatures statically checkable without runtime signature inspection. Ordinary sparse variants use `attention_output_hidden_size` and the shared -output allocation. DeepSeek-V4's fused epilogue is the exception: it requires -two typed output buffers before the attention op runs, so it implements the -optional output-preparation hook. `_create_outputs()` always returns a tensor -list whose first entry is the standard attention output. The same list flows -through forward, the registered custom op, and output projection; algorithms -own the meaning of any additional entries. Context- and generation-phase -helpers stay within each algorithm module and are not part of the generic hook -facade. +output allocation. DeepSeek-V4's fused epilogue instead uses the optional +output-preparation hook to create one token-major O-LoRA output tensor. Its +context- and generation-phase helpers allocate the private FP8 attention and +scale buffers, then write the O-LoRA result into the corresponding token range. +The shared MLA custom-op contract exposes exactly one mutable output tensor; +`_create_outputs()` keeps that tensor in a single-entry list through forward +and output projection. Phase-specific scratch buffers remain inside the +DeepSeek-V4 algorithm module and do not widen the generic hook facade. Sparse prediction inputs stay out of shared MLA APIs. Algorithm modules wrap their module-to-backend inputs in a `SparseBackendForwardArgs` subclass and