diff --git a/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml b/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml index 9d8e16234d4d..79090f279767 100644 --- a/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml +++ b/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml @@ -4,5 +4,19 @@ max_seq_len: 4096 enable_chunked_prefill: true cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 32, 64] transforms: + match_swiglu_pattern: + enabled: true + match_nvfp4_swiglu_pattern: + enabled: true fuse_nvfp4_moe: allow_different_input_scales: true + fuse_nvfp4_swiglu: + enabled: true + fuse_swiglu: + enabled: true + multi_stream_moe: + stage: compile + enabled: true + multi_stream_mla_attn: + stage: compile + enabled: true diff --git a/examples/auto_deploy/super_v3.yaml b/examples/auto_deploy/super_v3.yaml index 13b536a630d9..19ee522c2fc9 100644 --- a/examples/auto_deploy/super_v3.yaml +++ b/examples/auto_deploy/super_v3.yaml @@ -3,7 +3,7 @@ compile_backend: torch-cudagraph max_batch_size: 384 max_seq_len: 65536 # tunable enable_chunked_prefill: true -attn_backend: flashinfer +attn_backend: trtllm model_factory: AutoModelForCausalLM skip_loading_weights: false cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 24, 32, 64, 128, 256, 320, 384] @@ -37,7 +37,7 @@ transforms: "fc2_latent_proj": "gather" multi_stream_moe: stage: compile - enabled: false + enabled: true gather_logits_before_lm_head: # TODO: fix https://github.com/NVIDIA/TensorRT-LLM/issues/9878 to enable by default enabled: true diff --git a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py index 87434c48e9ca..1d326479aaa6 100644 --- a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py +++ b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py @@ -1,10 +1,22 @@ -"""Compile backend with cudagraph.""" +"""Compile backend with cudagraph. -from typing import Any, Dict, List, Optional, Tuple +1. Monolithic CUDA graph: captures entire model as one graph for decode-only. +2. Piecewise CUDA graph: splits model at dynamic ops, captures static segments + individually. Used for prefill/mixed batches when piecewise_enabled=True. + +When piecewise_enabled=True, a DualModeCapturedGraph is returned that dispatches: + - Decode-only batches → monolithic CapturedGraph (fastest, single graph replay) + - Prefill/mixed batches → PiecewiseCapturedGraph (per-segment replay + eager dynamic ops) +""" + +import copy # noqa: I001 +import operator +from typing import Any, Callable, Dict, List, Optional, Tuple import torch import torch.nn as nn from torch.cuda import CUDAGraph +from torch.fx import GraphModule from torch.fx._pytree import tree_flatten_spec from torch.utils._pytree import PyTree, TreeSpec, tree_flatten @@ -13,6 +25,52 @@ from ...utils.cuda_graph import CudaGraphWarmUpPhase from ...utils.logger import ad_logger from ..compiler import CompileBackendRegistry, CompilerBackend, GetArgsKwargsForBatchSize +from ..piecewise_runner import ADPiecewiseRunner +from ..piecewise_utils import SplitInfo, split_graph_at_dynamic_ops + + +# Trivial FX ops that are metadata-only or typically no-ops — used to identify +# static segments with no meaningful GPU compute (e.g., between adjacent dynamic ops). +# NOTE: reshape, contiguous, and to *can* launch kernels in edge cases (non-contiguous +# tensors, dtype/device casts), but in practice these appear only as lightweight +# plumbing in empty partitions. +_TRIVIAL_CALL_FUNCTIONS = {operator.getitem, getattr} +_TRIVIAL_CALL_METHODS = { + "view", + "reshape", + "contiguous", + "permute", + "transpose", + "unsqueeze", + "squeeze", + "expand", + "size", + "dim", + "to", +} + + +def _submod_has_cuda_ops(submod: nn.Module) -> bool: + """Check if a submodule has ops beyond those in _TRIVIAL_CALL_FUNCTIONS/METHODS.""" + if not isinstance(submod, GraphModule): + return True # Conservative: non-FX modules assumed to have GPU ops + + for node in submod.graph.nodes: + if node.op == "call_module": + # nn.Module calls (Linear, LayerNorm, etc.) launch CUDA kernels + return True + if node.op == "call_function": + if node.target in _TRIVIAL_CALL_FUNCTIONS: + continue + # Any non-trivial call_function is potentially a CUDA op + return True + if node.op == "call_method": + if node.target in _TRIVIAL_CALL_METHODS: + continue + # Non-trivial method call — could launch a kernel + return True + + return False def _args_kwargs_flatten_spec(in_spec: TreeSpec, *args, **kwargs) -> List[Any]: @@ -167,9 +225,260 @@ def forward(self, *args, **kwargs) -> Any: return self._out_spec.unflatten(out_flat) +class PiecewiseCapturedGraph(nn.Module): + """Manages piecewise CUDA graph capture/replay for prefill/mixed batches. + + The model is split at dynamic op boundaries (attention, SSM, conv, delta). + Static segments are wrapped in ADPiecewiseRunner for CUDA graph capture. + Dynamic segments run eagerly. The split_gm orchestrates the flow. + """ + + def __init__( + self, + model: nn.Module, + piecewise_num_tokens: Optional[List[int]] = None, + ): + super().__init__() + self.original_model = model + self.piecewise_num_tokens = piecewise_num_tokens or [] + self.split_info: Optional[SplitInfo] = None + self.split_gm: Optional[GraphModule] = None + self._is_prepared = False + + def prepare(self) -> None: + """Prepare the piecewise graph: swap to inplace ops, split, wrap static segments.""" + if self._is_prepared: + return + + model = self.original_model + if not isinstance(model, GraphModule): + ad_logger.warning( + "PiecewiseCapturedGraph: model is not a GraphModule, " + "piecewise CUDA graph requires an FX GraphModule. " + "Falling back to eager execution." + ) + self._is_prepared = True + return + + # Create a new GraphModule that shares all parameters/buffers/submodules + # with the original (zero-copy) but has its OWN copy of the FX graph + # (so split_graph_at_dynamic_ops mutations don't affect the original). + gm = GraphModule(model, copy.deepcopy(model.graph)) + + # Split graph at dynamic op boundaries + self.split_info = split_graph_at_dynamic_ops(gm) + self.split_gm = self.split_info.split_gm + + # Skip trivial submodules that have no CUDA ops (only contain getitem/reshape plumbing). + # Capturing these as CUDA graphs produces empty graphs and triggers PyTorch warnings. + # Create a shared pool upfront so all runners share memory allocations. + graph_pool = torch.cuda.graph_pool_handle() + num_wrapped = 0 + num_skipped = 0 + for idx in self.split_info.static_submod_indices: + submod_name = f"submod_{idx}" + if hasattr(self.split_gm, submod_name): + original_submod = getattr(self.split_gm, submod_name) + + if not _submod_has_cuda_ops(original_submod): + ad_logger.info( + f"PiecewiseCapturedGraph: skipping {submod_name} " + f"(no CUDA ops, will run eagerly)" + ) + num_skipped += 1 + continue + + runner = ADPiecewiseRunner( + submodule=original_submod, + piecewise_num_tokens=self.piecewise_num_tokens, + graph_pool=graph_pool, + ) + setattr(self.split_gm, submod_name, runner) + num_wrapped += 1 + + self._is_prepared = True + ad_logger.info( + f"PiecewiseCapturedGraph: prepared with " + f"{self.split_info.num_submodules} submodules " + f"({num_wrapped} wrapped for CUDA graph, {num_skipped} trivial skipped, " + f"{len(self.split_info.dynamic_submod_indices)} dynamic eager), " + f"piecewise_num_tokens={self.piecewise_num_tokens}" + ) + + def warmup_and_capture( + self, + get_args_kwargs: Callable[[int], Any], + warmup_iters: int = 3, + ) -> None: + """Warmup and capture CUDA graphs for all configured num_tokens values. + + Follows the same pattern as monolithic CapturedGraph._capture_one_graph: + the orchestrator controls the warmup → capture transition explicitly. + + Args: + get_args_kwargs: Callable that takes num_tokens and returns (args, kwargs). + warmup_iters: Number of eager warmup iterations before capture (default: 3, + matching monolithic CapturedGraph._capture_one_graph). + """ + if not self._is_prepared: + self.prepare() + + if self.split_gm is None: + return + + # Sort num_tokens in descending order (largest first for memory allocation) + num_tokens_list = sorted(self.piecewise_num_tokens, reverse=True) + for nt in num_tokens_list: + ad_logger.info(f"PiecewiseCapturedGraph: warming up for num_tokens={nt}") + args, kwargs = get_args_kwargs(nt) + + # Set the num_tokens context so ALL ADPiecewiseRunners use the correct value. + # This is critical: in piecewise-split models, some submodules receive + # intermediate tensors (SSM metadata, chunk indices) whose dim0 != num_tokens, + # so inferring from arg shapes is unreliable. + ADPiecewiseRunner.set_current_num_tokens(nt) + + with CudaGraphWarmUpPhase(): + ADPiecewiseRunner.set_current_phase("warmup") + for _ in range(warmup_iters): + self.split_gm(*args, **kwargs) + + # Capture phase: capture CUDA graphs for all static segments + ADPiecewiseRunner.set_current_phase("capture") + self.split_gm(*args, **kwargs) + + ad_logger.info(f"PiecewiseCapturedGraph: captured graphs for num_tokens={nt}") + + # Clear contexts after warmup/capture phase + ADPiecewiseRunner.set_current_num_tokens(None) + ADPiecewiseRunner.set_current_phase("replay") + + def forward(self, *args, num_tokens: Optional[int] = None, **kwargs) -> Any: + """Forward pass through the piecewise graph. + + Each submodule handles its own capture/replay: + - Static submodules (ADPiecewiseRunner): replay CUDA graph if available + - Dynamic submodules: run eagerly + + Args: + num_tokens: The total number of tokens in this batch. Must be provided + by the caller (DualModeCapturedGraph) — we cannot reliably infer it + from arg shapes because kwargs like input_ids may be [1, num_tokens] + (shape[0]=1, not num_tokens) and the first kwarg might not be input_ids. + """ + if self.split_gm is not None: + # Set num_tokens context for all ADPiecewiseRunners. + ADPiecewiseRunner.set_current_num_tokens(num_tokens) + result = self.split_gm(*args, **kwargs) + return result + else: + # Fallback: model is not a GraphModule, run eagerly + return self.original_model(*args, **kwargs) + + +class DualModeCapturedGraph(nn.Module): + """Dispatches between monolithic CG (decode) and piecewise CG (prefill/mixed). + + At runtime: + - If batch is decode-only (num_prefill == 0) -> use monolithic CapturedGraph + - If batch has prefill/mixed tokens and total num_tokens <= largest pre-captured + bucket -> use PiecewiseCapturedGraph with the smallest bucket >= num_tokens + - Otherwise -> fall back to eager + + Padding is handled upstream by SequenceInfo._padded_num_tokens: input_ids and + position_ids are shaped to the bucket size via _shape_for_forward, while all + metadata (batch_info_host, cu_seqlens, etc.) remains unchanged so dynamic ops + process only real tokens. Output logits are truncated back to the real token + count after the forward pass. + """ + + def __init__( + self, + monolithic: CapturedGraph, + piecewise: PiecewiseCapturedGraph, + batch_info_kwarg_name: str = "batch_info_host", + batched_input_names: Optional[List[str]] = None, + ): + super().__init__() + self.monolithic = monolithic + self.piecewise = piecewise + self.batch_info_kwarg_name = batch_info_kwarg_name + # Names of kwargs used to infer total num_tokens + self.batched_input_names = batched_input_names or ["input_ids", "position_ids"] + + # Sorted list of pre-captured bucket sizes for nearest-bucket lookup + self._captured_num_tokens_sorted: List[int] = sorted(piecewise.piecewise_num_tokens) + + def _is_decode_only(self, **kwargs) -> bool: + """Check if the current batch is decode-only using batch_info_host. + + batch_info_host = [num_prefill, num_prefill_tokens, num_decode] + Decode-only means num_prefill == 0. + """ + batch_info = kwargs.get(self.batch_info_kwarg_name) + if batch_info is not None and isinstance(batch_info, torch.Tensor): + # batch_info_host[0] = num_prefill + num_prefill = batch_info[0].item() + return num_prefill == 0 + + # Fallback heuristic: check if first batched input has sequence dim == 1 + # (decode = 1 token per sequence) + for name in self.batched_input_names: + v = kwargs.get(name) + if v is not None and isinstance(v, torch.Tensor) and v.ndim >= 2: + return v.shape[1] == 1 + + # Default to monolithic (decode) path + return True + + def _get_num_tokens(self, **kwargs) -> int: + """Extract total num_tokens from the batched inputs. + + For prefill/mixed with flattened layout: input_ids shape = [1, total_num_tokens] + We use numel() which works for both [1, N] and [N] layouts. + """ + for name in self.batched_input_names: + v = kwargs.get(name) + if v is not None and isinstance(v, torch.Tensor): + return v.numel() + return 0 + + def _find_nearest_bucket(self, num_tokens: int) -> Optional[int]: + """Find smallest captured bucket >= num_tokens, or None.""" + for bucket in self._captured_num_tokens_sorted: + if bucket >= num_tokens: + return bucket + return None + + def forward(self, *args, **kwargs) -> Any: + # NOTE: AD calls model(**named_args) so everything is in kwargs, args is empty + if self._is_decode_only(**kwargs): + return self.monolithic(*args, **kwargs) + + # ── PREFILL/MIXED PATH ── + num_tokens = self._get_num_tokens(**kwargs) + bucket = self._find_nearest_bucket(num_tokens) + if bucket is not None: + # Piecewise CG path -- padding is handled upstream by SequenceInfo + return self.piecewise(*args, num_tokens=bucket, **kwargs) + + # No bucket large enough -- eager fallback + ad_logger.debug( + f"DualModeCapturedGraph: num_tokens={num_tokens} exceeds largest bucket " + f"{self._captured_num_tokens_sorted[-1] if self._captured_num_tokens_sorted else 'N/A'}" + f", falling back to eager" + ) + return self.piecewise.original_model(*args, **kwargs) + + @CompileBackendRegistry.register("torch-cudagraph") class TorchCudagraphCompiler(CompilerBackend): - """Compiler that uses only CUDA graphs.""" + """Compiler that uses CUDA graphs. + + Supports two modes: + - piecewise_enabled=False (default): monolithic CG only (decode-only batches) + - piecewise_enabled=True: dual-mode (monolithic for decode + piecewise for prefill/mixed) + """ def __init__( self, @@ -177,21 +486,47 @@ def __init__( cuda_graph_batch_sizes: Optional[List[int]] = None, num_batched_inputs: int = 1, get_args_kwargs_for_compile: GetArgsKwargsForBatchSize = None, + piecewise_enabled: bool = False, + piecewise_num_tokens: Optional[List[int]] = None, + get_mixed_args_kwargs_for_compile: Optional[Callable[[int], Any]] = None, **kwargs_for_init, ): super().__init__(*args_for_init, **kwargs_for_init) self.num_batched_inputs = num_batched_inputs self.cuda_graph_batch_sizes = cuda_graph_batch_sizes or [] self.get_args_kwargs_for_compile = get_args_kwargs_for_compile + self.piecewise_enabled = piecewise_enabled + self.piecewise_num_tokens = piecewise_num_tokens or [] + self.get_mixed_args_kwargs_for_compile = get_mixed_args_kwargs_for_compile @torch.inference_mode() - def compile(self) -> CapturedGraph: - captured_model = CapturedGraph(self.model, num_batched_inputs=self.num_batched_inputs) - - # try capturing cudagraph + def compile(self) -> nn.Module: assert self.get_args_kwargs_for_compile is not None, ( "get_args_kwargs_for_compile must be provided" ) - captured_model.capture_graph(self.get_args_kwargs_for_compile, self.cuda_graph_batch_sizes) - return captured_model + # wrap get_args_kwargs_for_compile with CudaGraphWarmUpPhase. Note that host-side prepare + # functions may be called as part of get_args_kwargs. We want to let these functions know it's + # a warm-up phase. + def get_args_kwargs_warmup(batch_size: int): + with CudaGraphWarmUpPhase(): + return self.get_args_kwargs_for_compile(batch_size) + + monolithic = CapturedGraph(self.model, num_batched_inputs=self.num_batched_inputs) + monolithic.capture_graph(get_args_kwargs_warmup, self.cuda_graph_batch_sizes) + + piecewise = None + if self.piecewise_enabled: + ad_logger.info("TorchCudagraphCompiler: dual-mode enabled (monolithic + piecewise)") + piecewise = PiecewiseCapturedGraph( + model=self.model, + piecewise_num_tokens=self.piecewise_num_tokens, + ) + piecewise.prepare() + + if self.get_mixed_args_kwargs_for_compile is not None and self.piecewise_num_tokens: + piecewise.warmup_and_capture(self.get_mixed_args_kwargs_for_compile) + + if piecewise is not None: + return DualModeCapturedGraph(monolithic, piecewise) + return monolithic diff --git a/tensorrt_llm/_torch/auto_deploy/compile/piecewise_runner.py b/tensorrt_llm/_torch/auto_deploy/compile/piecewise_runner.py new file mode 100644 index 000000000000..6540796a9d59 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/compile/piecewise_runner.py @@ -0,0 +1,358 @@ +"""ADPiecewiseRunner: manages warmup → capture → replay for a single static CUDA graph segment. + +Each static submodule in a piecewise-split model is wrapped in an ADPiecewiseRunner. +The runner's behavior is controlled by two class-level contexts set by the orchestrator +(PiecewiseCapturedGraph) before each split_gm forward pass: + + - `_current_phase`: determines execution mode ("warmup", "capture", or "replay") + - `_current_num_tokens`: identifies which bucket entry to use + +Phase semantics: + 1. WARMUP: Run the submodule eagerly. (Data-ptr tracking runs but is NOT relied on + for correctness — see note on dynamic-index identification below.) + 2. CAPTURE: Capture the submodule as a CUDA graph. All non-weight tensor args are + treated as dynamic. For those that came from a previous static runner (found in + the _static_output_registry), we reuse the same buffer (zero-copy). Others + (model inputs, dynamic-segment outputs) are referenced directly and refreshed + via _prepare_replay_inputs during replay. + 3. REPLAY: Copy only dynamic inputs into the static buffers, then replay the + captured graph. + +Dynamic-index identification: + We do NOT rely on data_ptr() change detection during warmup, because PyTorch's + caching allocator can reuse the same address for activation tensors across warmup + iterations, making them falsely appear "static." Instead, we mark ALL non-weight + tensor args as dynamic. Weights/buffers are identified by matching against + data_ptrs collected from `submodule.parameters()` and `submodule.buffers()`. + +Each runner maintains entries keyed by `num_tokens`. +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Set, Tuple + +import torch +import torch.nn as nn +from torch.utils._pytree import tree_flatten, tree_unflatten + +from ..utils.logger import ad_logger + + +@dataclass +class SegmentEntry: + """State for a single (num_tokens) configuration of a segment.""" + + cuda_graph: Optional[torch.cuda.CUDAGraph] = None + # Static input list — each element is a direct reference to a tensor at a fixed address. + # During replay, _prepare_replay_inputs refreshes activation buffers as needed. + # + # Three categories: + # - Weight tensors: referenced directly (already at fixed addresses, never change). + # - Activation tensors from a previous static runner's output: reused from the + # static output registry. During replay the previous runner's CUDA graph writes to + # the same address, so _prepare_replay_inputs skips the copy (zero-copy). + # - Activation tensors from model inputs or dynamic segment outputs: referenced + # directly from the capture iteration. During replay, the dynamic segment produces + # output at a new address, so _prepare_replay_inputs copies into this buffer. + static_inputs: Optional[List[Any]] = None + # Indices of dynamic (activation) tensor args that need copy during replay + dynamic_indices: Optional[Set[int]] = None + # Static output — the output tensor(s) produced during capture. + # During replay, the CUDA graph writes to the same addresses, so returning + # this object gives the caller the updated data. + static_output: Any = None + # Tracks data_ptr() of tensor args during warmup to identify static vs dynamic + _warmup_data_ptrs: Optional[List[Optional[int]]] = None + + +class ADPiecewiseRunner(nn.Module): + """Wraps a static submodule and manages its CUDA graph capture/replay. + + Behavior is controlled by two class-level contexts set by the orchestrator: + - `_current_phase`: "warmup" (eager + ptr tracking), "capture" (CUDA graph + capture), or "replay" (graph replay / eager fallback at runtime) + - `_current_num_tokens`: identifies which bucket entry to use + + If `num_tokens` doesn't match any pre-configured bucket, falls back to eager. + Bucket resolution (nearest bucket >= real token count) is handled upstream by + DualModeCapturedGraph, so the runner always sees an exact bucket value. + """ + + # Class-level contexts: the orchestrator sets these before each split_gm forward pass + # so ALL runners in the graph use the same correct num_tokens and phase. + _current_num_tokens: Optional[int] = None + _current_phase: str = "replay" # "warmup", "capture", or "replay" + + # Class-level registry of output tensors produced during CUDA graph capture. + # Key: (num_tokens, data_ptr) -> output tensor at a fixed address. + # During capture, a runner checks if any of its activation inputs match a + # registered output (by data_ptr). If so, it references that buffer directly — + # enabling zero-copy during replay (the producer's graph writes, the consumer's + # graph reads, same address). + # Note: runners capture in sequential order, so all registry entries are from + # earlier runners — no need to track runner_id. + _static_output_registry: Dict[Tuple[int, int], torch.Tensor] = {} + + @classmethod + def set_current_num_tokens(cls, num_tokens: Optional[int]) -> None: + """Set the current num_tokens context for all runners. + + Called by PiecewiseCapturedGraph before each forward pass through the split graph. + """ + cls._current_num_tokens = num_tokens + + @classmethod + def set_current_phase(cls, phase: str) -> None: + """Set the current execution phase for all runners. + + Called by PiecewiseCapturedGraph to control warmup → capture → replay transitions. + Valid phases: "warmup", "capture", "replay". + """ + assert phase in ("warmup", "capture", "replay"), f"Invalid phase: {phase}" + cls._current_phase = phase + + @classmethod + def clear_static_output_registry(cls) -> None: + """Clear the static output registry. + + Called when switching between different graph configurations or resetting state. + """ + cls._static_output_registry.clear() + + def __init__( + self, + submodule: nn.Module, + piecewise_num_tokens: Optional[List[int]] = None, + graph_pool: Optional[Tuple[int, ...]] = None, + ): + super().__init__() + self.submodule = submodule + self._graph_pool = graph_pool + + # Collect data_ptrs of all parameters and buffers in this submodule. + # These are weight tensors with stable addresses that NEVER need copying. + # Everything else that appears in flat_args is a cross-partition activation + # (from a previous static runner or a dynamic segment) and must be treated + # as dynamic for correctness during CUDA graph replay. + self._weight_ptrs: Set[int] = set() + for p in submodule.parameters(): + self._weight_ptrs.add(p.data_ptr()) + for b in submodule.buffers(): + self._weight_ptrs.add(b.data_ptr()) + + # Pre-populate entries for each bucket size + self.entries: Dict[int, SegmentEntry] = {} + if piecewise_num_tokens: + for nt in piecewise_num_tokens: + self.entries[nt] = SegmentEntry() + + def _find_entry(self, num_tokens: int) -> Optional[SegmentEntry]: + """Find the SegmentEntry for the given num_tokens. + + Expects an exact match — bucket resolution (nearest bucket >= real token count) + is handled upstream by DualModeCapturedGraph._find_nearest_bucket before + num_tokens reaches the runner. + + Returns None if num_tokens doesn't match any pre-configured bucket (eager fallback). + """ + return self.entries.get(num_tokens) + + def _track_warmup_ptrs(self, entry: SegmentEntry, flat_args: List[Any]) -> None: + """Track data_ptr() during warmup to identify static (weight) vs dynamic (activation) args. + + On the first warmup call, record all data_ptrs. On subsequent calls, mark args whose + data_ptr changed as "dynamic" (by setting their tracked ptr to None). + """ + if entry._warmup_data_ptrs is None: + # First warmup: record all data_ptrs + entry._warmup_data_ptrs = [ + a.data_ptr() if isinstance(a, torch.Tensor) else None for a in flat_args + ] + else: + # Subsequent warmup: check for changes + for i, a in enumerate(flat_args): + if isinstance(a, torch.Tensor): + if ( + entry._warmup_data_ptrs[i] is not None + and a.data_ptr() != entry._warmup_data_ptrs[i] + ): + # data_ptr changed → this is a dynamic (activation) tensor + entry._warmup_data_ptrs[i] = None + + def _identify_dynamic_indices(self, entry: SegmentEntry, flat_args: List[Any]) -> Set[int]: + """Mark all non-weight tensor args as dynamic. + + Weight/buffer tensors (matched via _weight_ptrs) are static. + Everything else is dynamic — the capture code will further check + _static_output_registry for zero-copy reuse where possible. + """ + dynamic_indices: Set[int] = set() + for i, a in enumerate(flat_args): + if not isinstance(a, torch.Tensor): + continue + if a.data_ptr() in self._weight_ptrs: + continue # Weight/buffer — stable address, no copy needed + dynamic_indices.add(i) + return dynamic_indices + + def _prepare_replay_inputs(self, entry: SegmentEntry, flat_inputs: List[Any]) -> None: + """Refresh dynamic activation buffers before CUDA graph replay. + + For each dynamic tensor input, this copies runtime data into the captured + static buffer unless both tensors already share the same data_ptr() (no-copy + fast path, common for static segment chaining). + + When runtime input is smaller than the bucketed static buffer (padding case), + copy the valid prefix and clear the padded tail. Clearing avoids stale values + from prior warmup/capture executions leaking into downstream ops. + """ + for idx in entry.dynamic_indices: + new_inp = flat_inputs[idx] + static_inp = entry.static_inputs[idx] + + if not isinstance(new_inp, torch.Tensor) or not isinstance(static_inp, torch.Tensor): + continue + + # Fast path: no copy needed when producer already wrote into the + # captured static buffer (segment N output -> segment N+1 input). + if new_inp.data_ptr() == static_inp.data_ptr(): + continue + + if static_inp.shape == new_inp.shape: + static_inp.copy_(new_inp, non_blocking=True) + elif ( + new_inp.shape[0] < static_inp.shape[0] and new_inp.shape[1:] == static_inp.shape[1:] + ): + # Padded case: runtime input is smaller along dim 0. + n = new_inp.shape[0] + static_inp[:n].copy_(new_inp, non_blocking=True) + static_inp[n:].zero_() + elif ( + new_inp.ndim >= 2 + and new_inp.shape[1] < static_inp.shape[1] + and new_inp.shape[0] == static_inp.shape[0] + ): + # Padded case: runtime input is smaller along dim 1 + # (e.g., [1, real, D] vs [1, bucket, D]). + n = new_inp.shape[1] + static_inp[:, :n].copy_(new_inp, non_blocking=True) + static_inp[:, n:].zero_() + else: + # Fallback: shapes are incompatible — this is a real error + static_inp.copy_(new_inp, non_blocking=True) + + def forward(self, *args, **kwargs) -> Any: + # Use the class-level contexts set by the orchestrator + num_tokens = ADPiecewiseRunner._current_num_tokens + phase = ADPiecewiseRunner._current_phase + entry = self._find_entry(num_tokens) if num_tokens is not None else None + + if entry is None: + # Unknown num_tokens or exceeds all buckets — fallback to eager + return self.submodule(*args, **kwargs) + + # Flatten inputs once (used by all phases) + flat_args, args_spec = tree_flatten((args, kwargs)) + + # --- WARMUP PHASE --- + if phase == "warmup": + # Track data_ptr() to distinguish weights from activations + self._track_warmup_ptrs(entry, flat_args) + return self.submodule(*args, **kwargs) + + # --- CAPTURE PHASE --- + if phase == "capture": + ad_logger.debug(f"ADPiecewiseRunner: capturing CUDA graph for num_tokens={num_tokens}") + + # Identify which args are dynamic (activations) vs static (weights) + entry.dynamic_indices = self._identify_dynamic_indices(entry, flat_args) + + # Build static_inputs list for this entry. Every element is a direct + # reference (no cloning) — we just need each tensor at a persistent address. + # + # For activation tensors, we check the static output registry to find + # outputs from previous static runners. During replay, those runners' + # CUDA graphs write to the same address, so _prepare_replay_inputs can skip the + # copy (zero-copy). All other activation tensors (model inputs, dynamic + # segment outputs) are referenced directly from this capture iteration; + # _prepare_replay_inputs will copy new data into them during replay. + entry.static_inputs = [] + num_reused = 0 + num_referenced = 0 + for i, a in enumerate(flat_args): + if isinstance(a, torch.Tensor) and i in entry.dynamic_indices: + # Check if this activation is a previous static runner's output + # (if so, record the registry reference for zero-copy during replay) + prev_output = ADPiecewiseRunner._static_output_registry.get( + (num_tokens, a.data_ptr()) + ) + if prev_output is not None: + entry.static_inputs.append(prev_output) + num_reused += 1 + else: + # Model input or dynamic segment output — reference directly. + # During replay, _prepare_replay_inputs will copy new data into this buffer. + entry.static_inputs.append(a) + else: + # Weight tensor — reference directly (fixed address, never changes) + entry.static_inputs.append(a) + if isinstance(a, torch.Tensor): + num_referenced += 1 + + # Unflatten back to get the static args/kwargs + static_args_kwargs = tree_unflatten(entry.static_inputs, args_spec) + static_args, static_kwargs = static_args_kwargs + + # Capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._graph_pool): + output = self.submodule(*static_args, **static_kwargs) + + torch.cuda.synchronize() + + # Fallback: if no pool was provided at construction time, store the + # auto-created pool so subsequent captures within this runner reuse it. + if self._graph_pool is None: + self._graph_pool = graph.pool() + + entry.cuda_graph = graph + entry.static_output = output + + # Register outputs in the static output registry so next runners can reuse them + flat_output, _ = tree_flatten(output) + for out_tensor in flat_output: + if isinstance(out_tensor, torch.Tensor): + ADPiecewiseRunner._static_output_registry[ + (num_tokens, out_tensor.data_ptr()) + ] = out_tensor + + num_dynamic = len(entry.dynamic_indices) - num_reused + ad_logger.debug( + f"ADPiecewiseRunner: captured graph for num_tokens={num_tokens} — " + f"{num_dynamic} dynamic activation buffers, " + f"{num_reused} reused from previous static segments, " + f"{num_referenced} weight tensors (zero-copy)" + ) + + return output + + # --- REPLAY PHASE --- + # Copy only dynamic inputs into static buffers. + # _prepare_replay_inputs skips copy if input is already at static buffer address + # (common case: segment N's output is segment N+1's input, so addresses match) + self._prepare_replay_inputs(entry, flat_args) + + # Replay the captured graph + entry.cuda_graph.replay() + + return entry.static_output + + @property + def graph_pool(self): + """Return the CUDA graph memory pool (for sharing across runners).""" + return self._graph_pool + + @graph_pool.setter + def graph_pool(self, pool): + self._graph_pool = pool diff --git a/tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py b/tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py new file mode 100644 index 000000000000..e4daf670ce4a --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py @@ -0,0 +1,219 @@ +"""Utilities for piecewise CUDA graph: graph splitting at dynamic op boundaries. + +This module provides the logic to: +1. Identify dynamic (uncapturable) custom ops in the FX graph (attention, SSM, conv, delta). +2. Split the FX GraphModule at those boundaries using torch.fx.passes.split_module. +3. Return the split GraphModule and metadata about which submodules are dynamic vs static. +""" + +from dataclasses import dataclass, field +from typing import Dict, List, Set + +from torch.fx import GraphModule, Node +from torch.fx.passes.split_module import split_module + +from ..utils.logger import ad_logger + +# --------------------------------------------------------------------------- +# Dynamic ops registry: these ops cannot be captured in CUDA graphs for +# mixed/prefill batches because they have data-dependent control flow or +# dynamic kernel configurations. +# --------------------------------------------------------------------------- + +# Cached attention ops (grid depends on per-sequence lengths) +_CACHED_ATTENTION_OPS = [ + "auto_deploy::flashinfer_attention_mha_with_cache", + "auto_deploy::triton_attention_flattened_mha_with_cache", + "auto_deploy::torch_cached_attention_with_cache", + "auto_deploy::trtllm_attention_mha_with_cache", +] + +# Cached SSM ops (Python-level branching on batch_info_host) +_CACHED_SSM_OPS = [ + "auto_deploy::triton_cached_ssm", + "auto_deploy::torch_cached_ssm", + "auto_deploy::flashinfer_cached_ssm", +] + +# Cached causal conv ops (branching on prefill vs decode) +_CACHED_CONV_OPS = [ + "auto_deploy::triton_cached_causal_conv1d", + "auto_deploy::cuda_cached_causal_conv1d", +] + +# Cached delta rule ops (branching on prefill vs decode) +_CACHED_DELTA_OPS = [ + "auto_deploy::fla_cached_delta_rule", +] + +# Metadata preparation ops (branch on batch_info_host, do CPU math on CUDA tensors) +_METADATA_PREP_OPS = [ + "auto_deploy::flashinfer_attention_prepare_metadata", + "auto_deploy::mamba_ssm_prepare_metadata", +] + +# Logits gather ops (CPU branching on host tensor + shape-dependent logic) +_LOGITS_GATHER_OPS = [ + "auto_deploy::gather_logits_before_lm_head", +] + + +def _get_all_dynamic_op_names() -> Set[str]: + """Return the full set of dynamic op qualified names.""" + return set( + _CACHED_ATTENTION_OPS + + _CACHED_SSM_OPS + + _CACHED_CONV_OPS + + _CACHED_DELTA_OPS + + _METADATA_PREP_OPS + + _LOGITS_GATHER_OPS + ) + + +def is_dynamic_cached_op(node: Node) -> bool: + """Check if a node is a dynamic (uncapturable) cached op. + + These are ops that cannot be captured inside a CUDA graph for mixed/prefill + batches due to data-dependent control flow or dynamic kernel grids. + """ + if node.op != "call_function": + return False + + target = node.target + # Handle OpOverload: get the qualified name + if hasattr(target, "name"): + # torch._ops.OpOverload has .name() method + op_name = target.name() + elif hasattr(target, "__qualname__"): + op_name = target.__qualname__ + else: + op_name = str(target) + + # Strip the ".default" suffix if present for matching + dynamic_ops = _get_all_dynamic_op_names() + # Check with namespace::name format AND base name (for wrapper functions + for dyn_op in dynamic_ops: + if dyn_op in op_name: + return True + # Also check by base op name without namespace prefix + base_name = dyn_op.split("::")[-1] if "::" in dyn_op else dyn_op + if base_name in op_name: + return True + + return False + + +@dataclass +class SplitInfo: + """Metadata about a split GraphModule.""" + + # The split GraphModule with submod_0, submod_1, ... submodules + split_gm: GraphModule + # Total number of submodules + num_submodules: int + # Indices of dynamic (uncapturable) submodules — these run eagerly + dynamic_submod_indices: List[int] = field(default_factory=list) + # Indices of static (capturable) submodules — these get CUDA graph captured + static_submod_indices: List[int] = field(default_factory=list) + + +def split_graph_at_dynamic_ops(gm: GraphModule) -> SplitInfo: + """Split an FX GraphModule at dynamic op boundaries. + + Each dynamic op (attention, SSM, conv, delta) becomes its own submodule. + Static regions between dynamic ops are grouped into separate submodules. + + The split produces submodules named `submod_0`, `submod_1`, etc. + Dynamic submodules contain exactly one dynamic op. + Static submodules contain everything else (norms, linears, MLPs, etc.). + + Args: + gm: The FX GraphModule to split. + + Returns: + SplitInfo with the split GraphModule and metadata. + """ + # Assign partition IDs: each dynamic op gets its own partition, + # static ops between dynamic ops share a partition. + partition_counter = [0] # mutable counter + node_to_partition: Dict[Node, int] = {} + dynamic_partitions: Set[int] = set() + + # First pass: identify dynamic nodes and assign them unique partitions + for node in gm.graph.nodes: + if node.op in ("placeholder", "output"): + continue + + if is_dynamic_cached_op(node): + # Dynamic op gets its own partition + partition_counter[0] += 1 + node_to_partition[node] = partition_counter[0] + dynamic_partitions.add(partition_counter[0]) + # Next static region gets a new partition + partition_counter[0] += 1 + else: + # Static op joins the current static partition + node_to_partition[node] = partition_counter[0] + + if not dynamic_partitions: + ad_logger.info("No dynamic ops found in graph — no splitting needed.") + return SplitInfo( + split_gm=gm, + num_submodules=1, + dynamic_submod_indices=[], + static_submod_indices=[0], + ) + + # Use torch.fx split_module to perform the actual split + def partition_fn(node: Node) -> int: + return node_to_partition.get(node, 0) + + split_gm = split_module( + gm, + gm, # root_module + partition_fn, + keep_original_order=True, + ) + + # Analyze the split result to identify dynamic vs static submodules + submod_names = [] + for name, _ in split_gm.named_children(): + if name.startswith("submod_"): + submod_names.append(name) + + # Sort by index + submod_names.sort(key=lambda n: int(n.split("_")[1])) + + # Build a mapping from partition ID to submod index + # The split_module assigns submod_N names in order of first-seen partition IDs + partition_ids_in_order = [] + seen = set() + for node in gm.graph.nodes: + if node.op in ("placeholder", "output"): + continue + pid = node_to_partition.get(node, 0) + if pid not in seen: + seen.add(pid) + partition_ids_in_order.append(pid) + + dynamic_indices = [] + static_indices = [] + for idx, pid in enumerate(partition_ids_in_order): + if idx >= len(submod_names): + break + if pid in dynamic_partitions: + dynamic_indices.append(idx) + else: + static_indices.append(idx) + + ad_logger.info( + f"Piecewise split: {len(submod_names)} submodules " + f"({len(static_indices)} static, {len(dynamic_indices)} dynamic)" + ) + + return SplitInfo( + split_gm=split_gm, + num_submodules=len(submod_names), + dynamic_submod_indices=dynamic_indices, + static_submod_indices=static_indices, + ) diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 84aca711e363..ca2bd8342bde 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -50,6 +50,7 @@ transforms: expected_layout: bsnd match_rmsnorm_pattern: stage: pattern_matcher + run_shape_prop: true match_l2norm_pattern: stage: pattern_matcher ############################################################################################ @@ -75,6 +76,18 @@ transforms: stage: pattern_matcher quantize_nvfp4_from_graph: stage: pattern_matcher + # SwiGLU pattern matching must run AFTER quantization transforms. For pre-quantized + # checkpoints (e.g., NVFP4), quantization converts torch_linear_simple ops to quantized + # ops first, and then match_nvfp4_swiglu_pattern captures the NVFP4 SwiGLU pattern. + # For non-quantized models, quantization transforms are no-ops, so match_swiglu_pattern + # proceeds normally. + match_swiglu_pattern: + stage: pattern_matcher + enabled: false + match_nvfp4_swiglu_pattern: + stage: pattern_matcher + requires_shape_prop: true + enabled: false quantize_fp8_moe: stage: pattern_matcher quantize_nvfp4_moe: @@ -126,6 +139,8 @@ transforms: fuse_nvfp4_linear: stage: post_load_fusion backend: trtllm + fuse_nvfp4_swiglu: + stage: post_load_fusion fuse_moe: stage: post_load_fusion expect_mem_change: true @@ -149,6 +164,9 @@ transforms: fuse_l2norm: stage: post_load_fusion backend: fla + fuse_swiglu: + stage: post_load_fusion + enabled: false fuse_add_rms_norm: stage: post_load_fusion enabled: true @@ -200,9 +218,14 @@ transforms: multi_stream_moe: stage: compile enabled: false + multi_stream_mla_attn: + stage: compile + enabled: false compile_model: stage: compile expect_mem_change: true run_per_gm: false cuda_graph_batch_sizes: null backend: torch-compile + piecewise_enabled: true + piecewise_num_tokens: null diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py index ead9ba122dab..438271839738 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.py @@ -19,6 +19,7 @@ - torch_attention: PyTorch reference implementations - torch_backend_attention: PyTorch-based attention backend - flashinfer_attention: FlashInfer-based optimized attention +- trtllm_attention: TRT-LLM thop.attention-based optimized attention - triton_attention: Triton-based attention implementations - triton_attention_with_kv_cache: Triton attention with KV cache support - triton_attention_with_paged_kv_cache: Triton attention with paged KV cache @@ -29,6 +30,7 @@ "torch_attention", "torch_backend_attention", "flashinfer_attention", + "trtllm_attention", "triton_attention", "triton_attention_with_kv_cache", "triton_attention_with_paged_kv_cache", diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py index 4183f5148ca8..7a6b23e4c447 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py @@ -366,10 +366,10 @@ def flashinfer_mha_with_cache( v = v.to(torch.float8_e4m3fn) flashinfer.page.append_paged_kv_cache( - append_key=k, - append_value=v, - batch_indices=flashinfer_batch_indices, - positions=flashinfer_positions, + append_key=k[:num_total_tokens], + append_value=v[:num_total_tokens], + batch_indices=flashinfer_batch_indices[:num_total_tokens], + positions=flashinfer_positions[:num_total_tokens], paged_kv_cache=kv_cache, kv_indices=cache_loc, kv_indptr=cu_num_pages[: num_seq + 1], @@ -453,7 +453,21 @@ def flashinfer_mha_with_cache( else: y = y_decode - return y.view(q_shape_og) # [b,s,n*h_d] or [b,s, n, h_d] + # Reshape to match input shape [b, s, ...] + # y has shape [num_total_tokens, n_heads, head_dim] (pure) or [b*s, n_heads, head_dim] (mixed). + # q_shape_og is [b, s, ...] which may be padded (bucketed) for piecewise CG. + # Pad with zeros if y is smaller than the padded shape, so downstream ops + # (o_proj, residual add, LayerNorm) don't see garbage in padding positions. + bs = b * s + if y.shape[0] < bs: + y_padded = torch.zeros((bs, y.shape[1], y.shape[2]), dtype=y.dtype, device=y.device) + y_padded[: y.shape[0]] = y + y = y_padded + elif num_total_tokens < bs: + # Mixed batch: y is already [b*s, ...] but positions [num_total_tokens:] are uninitialized + y[num_total_tokens:].zero_() + + return y.view(q_shape_og) @flashinfer_mha_with_cache.register_fake diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py index 36c8d54d4e4d..f3ac6a0aff65 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py @@ -341,7 +341,7 @@ def torch_backend_mha_with_cache( scale = 1.0 / math.sqrt(qk_head_dim) if scale is None else scale - # Create output tensor + # Preallocate output tensor y = q.new_empty(*bs_view, num_heads, v_head_dim).contiguous() # Compute attention @@ -380,6 +380,12 @@ def torch_backend_mha_with_cache( sinks, ) + # Zero padding positions so downstream ops don't see garbage (piecewise CG) + num_total_tokens = num_prefill_tokens + num_decode + bs = b * s + if num_total_tokens < bs: + y[num_total_tokens:].zero_() + return y.view(*output_shape) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py index 70eb07e50d44..7d65b143b579 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py @@ -287,6 +287,10 @@ def flattened_mha_with_cache( sliding_window, ) + # Zero padding positions so downstream ops don't see garbage (piecewise CG) + if num_total_tokens < bs: + y[num_total_tokens:].zero_() + return y.view(*output_shape) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py new file mode 100644 index 000000000000..55ae6f9fd5b2 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py @@ -0,0 +1,601 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TRT-LLM attention backend for Auto-Deploy. + +This module wraps TRT-LLM's optimized ``thop.attention`` kernel for use in Auto-Deploy, +following the same design pattern as the FlashInfer backend: + +- Minimal module-level state (``_TrtllmPlanner``, analogous to ``_FlashInferPlanner``) +- SequenceInfo fields used directly as thop.attention metadata +- Pool pointers derived lazily from ``kv_cache.data_ptr()`` +- Workspace managed as module-level state (not a ResourceHandler / graph input) +- All possible "constants" inferred from tensor shapes at runtime +""" + +from typing import List, Optional + +import torch +from torch._ops import OpOverloadPacket +from torch._subclasses import FakeTensor +from torch.fx import Node + +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import AttentionMaskType +from tensorrt_llm.quantization import QuantMode + +from .....llmapi.llm_args import KvCacheConfig +from ...utils.cuda_graph import cuda_graph_state +from ...utils.logger import ad_logger +from ...utils.node_utils import extract_op_args +from ..attention_interface import ( + AttentionDescriptor, + AttentionLayout, + AttentionRegistry, + Constant, + KVPagedResourceHandler, + MHACallable, + PrepareMetadataHostCallable, + ResourceHandlerDict, +) + +# ============================================================================= +# Module-level planner (analogous to _GlobalFlashInferPlanner) +# ============================================================================= + + +class _TrtllmPlanner: + """Minimal planner for TRT-LLM attention backend. + + Analogous to ``_FlashInferPlanner`` in the FlashInfer backend. Only stores + data that cannot be derived from SequenceInfo or tensor shapes. + + Two main entry points: + - ``reset()``: one-time allocation of ALL persistent buffers. + - ``plan()``: per-forward host metadata (host_request_types, block_offsets, host_total_kv_lens). + + Pool pointer management: + Each attention layer needs its own ``host_pool_pointers`` tensor so that CUDA graph + replay reads the correct (layer-specific) pool base from a stable tensor address. + Tensors are created lazily on first access per layer and never re-allocated. + """ + + def __init__(self): + self.workspace: Optional[torch.Tensor] = None + # Per-layer pool pointer tensors, keyed by kv_cache.data_ptr(). + # Each is [1, 2] int64 pinned, created lazily on first access per layer. + # This ensures each layer's attention kernel in a CUDA graph is captured + # with its own stable tensor address, avoiding the issue where a shared + # tensor would hold only the last-set layer's pointer during graph replay. + self._per_layer_pool_ptrs: dict = {} + # pool_mapping: fixed [1, 2] all zeros since we always pass layer_idx=0 + # and pool_pointers already encodes the layer offset via kv_cache.data_ptr() + self.host_pool_mapping: Optional[torch.Tensor] = None # [1, 2] int32 pinned + # thop-specific host metadata NOT available from SequenceInfo + self.host_request_types: Optional[torch.Tensor] = None # [max_batch] int32 pinned + self.host_total_kv_lens: Optional[torch.Tensor] = None # [2] int64 pinned + # thop variant of input_pos_host and seq_len_host + # keeping a separate copy here since we sometimes have to overwrite the original values + self.host_past_kv_lengths: Optional[torch.Tensor] = None # [max_batch] int32 pinned + self.host_context_lengths: Optional[torch.Tensor] = None # [max_batch] int32 pinned + # Persistent block_offsets buffer for CUDA graph compatibility. + # Pre-allocated to max size so the tensor address is stable across replays. + self.block_offsets: Optional[torch.Tensor] = None + # FP8 scale tensors (lazily initialized from constants on first FP8 use) + self.kv_scale_orig_quant: Optional[torch.Tensor] = None + self.kv_scale_quant_orig: Optional[torch.Tensor] = None + + def reset(self, device: torch.device, max_batch: int, max_blocks_per_seq: int) -> None: + """One-time allocation of ALL persistent buffers. + + Guards against double-init. Called lazily from ``prepare_trtllm_metadata_host`` + on the first forward pass after cache initialization. + """ + if self.workspace is not None: + return # already initialized + + # Workspace: pre-allocate a modest initial buffer (like flashinfer's 320MB). + # thop.attention auto-resizes via resize_() if more space is needed during warm-up. + self.workspace = torch.empty(256 * 1024 * 1024, dtype=torch.uint8, device=device) + self.host_pool_mapping = torch.zeros(1, 2, dtype=torch.int32, device="cpu", pin_memory=True) + self.host_total_kv_lens = torch.zeros(2, dtype=torch.int64, device="cpu", pin_memory=True) + self.host_request_types = torch.zeros( + max_batch, dtype=torch.int32, device="cpu", pin_memory=True + ) + self.block_offsets = torch.zeros( + 1, max_batch, 2, max_blocks_per_seq, dtype=torch.int32, device=device + ) + self.host_past_kv_lengths = torch.zeros( + max_batch, dtype=torch.int32, device="cpu", pin_memory=True + ) + self.host_context_lengths = torch.zeros( + max_batch, dtype=torch.int32, device="cpu", pin_memory=True + ) + + def plan( + self, + num_prefill: int, + num_decode: int, + max_context_length: int, + block_offset_multiplier: int, + seq_len_with_cache_host: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc: torch.Tensor, + page_seq_indices: torch.Tensor, + page_in_seq: torch.Tensor, + input_pos_host: torch.Tensor, + seq_len_host: torch.Tensor, + ) -> None: + """Per-forward host metadata: fills host_request_types, block_offsets, host_total_kv_lens. + + Called from ``prepare_trtllm_metadata_host`` before every forward (including replays). + """ + num_seq = num_prefill + num_decode + + # host_request_types: 0 = prefill (context), 1 = decode (generation) + self.host_request_types[:num_prefill].fill_(0) + self.host_request_types[num_prefill:num_seq].fill_(1) + + # Compute block_offsets for thop.attention using pre-computed page indices. + block_offsets = self.block_offsets + total_pages = int(cu_num_pages_host[num_seq]) + base_offsets = cache_loc[:total_pages] * block_offset_multiplier + seq_idx = page_seq_indices[:total_pages] + pg_idx = page_in_seq[:total_pages] + block_offsets[0, seq_idx, 0, pg_idx] = base_offsets # K + block_offsets[0, seq_idx, 1, pg_idx] = base_offsets + 1 # V + + # host_total_kv_lens: [context_total_kv, gen_total_kv] + is_capturing = torch.cuda.is_current_stream_capturing() or cuda_graph_state.in_warm_up() + if is_capturing: + # CUDA graph capture: set host tensors to MAX values so the kernel captures + # the worst-case execution pattern. + self.host_total_kv_lens[0] = max_context_length * num_prefill + self.host_total_kv_lens[1] = max_context_length * num_decode + self.host_past_kv_lengths[:num_seq].fill_(max_context_length) + self.host_context_lengths[:num_seq].fill_(max_context_length) + else: + self.host_total_kv_lens[0] = seq_len_with_cache_host[:num_prefill].sum() + self.host_total_kv_lens[1] = seq_len_with_cache_host[num_prefill:num_seq].sum() + self.host_past_kv_lengths[:num_seq] = input_pos_host[:num_seq] + self.host_context_lengths[:num_seq] = seq_len_host[:num_seq] + + def get_pool_pointers_for_layer(self, kv_cache: torch.Tensor) -> torch.Tensor: + """Return a per-layer ``host_pool_pointers`` tensor for this kv_cache view. + + Each attention layer receives a different ``kv_cache`` tensor (a strided view + into the pool). We create one pinned [1, 2] int64 tensor per unique + ``data_ptr`` and cache it forever. This guarantees that each layer's + ``thop.attention`` call in a CUDA graph is captured with a *stable, distinct* + tensor address, so graph replay reads the correct pool base for every layer. + """ + ptr = kv_cache.data_ptr() + t = self._per_layer_pool_ptrs.get(ptr) + if t is not None: + return t + + t = torch.zeros(1, 2, dtype=torch.int64, device="cpu", pin_memory=True) + t[0, 0] = ptr + self._per_layer_pool_ptrs[ptr] = t + return t + + +_GlobalTrtllmPlanner = _TrtllmPlanner() + + +# ============================================================================= +# Host-side prepare function (analogous to prepare_flashinfer_metadata_host) +# ============================================================================= + + +def prepare_trtllm_metadata_host( + batch_info_host: torch.Tensor, + max_seq_info_host: torch.Tensor, + seq_len_with_cache_host: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc: torch.Tensor, + page_seq_indices: torch.Tensor, + page_in_seq: torch.Tensor, + input_pos_host: torch.Tensor, + seq_len_host: torch.Tensor, +) -> None: + """Fill thop-specific host metadata and compute block_offsets. + + This runs OUTSIDE the CUDA graph before every forward (including replays). + Block offsets MUST be computed here (not in the device-side prepare_metadata op) + because they are batch-dependent and need to be updated before each replay. + + All max-size constants are read from ``max_seq_info_host`` which is set once via + ``SequenceInfo.update_cache_information()`` after cache initialization: + ``[max_context_length, max_blocks_per_seq, block_offset_multiplier, max_batch_size]`` + + ``page_seq_indices`` and ``page_in_seq`` are pre-computed in SequenceInfo from + ``pages_per_seq`` and avoid the expensive GPU searchsorted that was previously needed. + """ + num_prefill, _, num_decode = batch_info_host.tolist() + + # Read all max-size constants from max_seq_info_host (set at cache init time) + max_context_length, max_blocks_per_seq, block_offset_multiplier, max_batch_size = ( + max_seq_info_host.tolist() + ) + + # One-time allocation of all persistent buffers (lazy, guards against double-init) + _GlobalTrtllmPlanner.reset(cache_loc.device, max_batch_size, max_blocks_per_seq) + + # Per-forward: fill host_request_types, block_offsets, host_total_kv_lens + _GlobalTrtllmPlanner.plan( + num_prefill=num_prefill, + num_decode=num_decode, + max_context_length=max_context_length, + block_offset_multiplier=block_offset_multiplier, + seq_len_with_cache_host=seq_len_with_cache_host, + cu_num_pages_host=cu_num_pages_host, + cache_loc=cache_loc, + page_seq_indices=page_seq_indices, + page_in_seq=page_in_seq, + input_pos_host=input_pos_host, + seq_len_host=seq_len_host, + ) + + +# ============================================================================= +# Cached attention op (analogous to flashinfer_mha_with_cache) +# ============================================================================= + + +@torch.library.custom_op("auto_deploy::trtllm_attention_mha_with_cache", mutates_args=("kv_cache",)) +def trtllm_mha_with_cache( + # Q, K, V inputs + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + # STANDARD METADATA (SequenceInfo fields used directly by thop) + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + seq_len_host: torch.Tensor, + input_pos_host: torch.Tensor, + seq_len_with_cache: torch.Tensor, + max_seq_info_host: torch.Tensor, + # CACHE + kv_cache: torch.Tensor, + # CONSTANTS (only truly un-inferable values) + scale: Optional[float], + sliding_window: Optional[int] = None, + kv_scale_orig_quant: float = 1.0, + kv_scale_quant_orig: float = 1.0, +) -> torch.Tensor: + """TRT-LLM attention with paged KV cache for Auto-Deploy. + + Infers num_heads, num_kv_heads, head_dim, and tokens_per_block from tensor shapes. + All max-size constants (max_num_requests, max_context_length) are read from + ``max_seq_info_host`` which is set once via ``SequenceInfo.update_cache_information()``. + + Note: ``prepare_trtllm_metadata_host`` is guaranteed to be called before this op, + so all persistent planner buffers are already initialized. + + Note: layer_idx is always passed as 0 to thop.attention because + the kv_cache tensor is already a strided view for the correct layer, + pool_pointers encodes kv_cache.data_ptr() (layer-specific), and + pool_mapping is all zeros. See module docstring for details. + """ + # Infer dimensions from tensor shapes (bsnd layout) + num_heads = q.shape[2] + num_kv_heads = k.shape[2] + head_dim = q.shape[3] + tokens_per_block = kv_cache.shape[3] # HND: [blocks, 2, heads, tpb, head_dim] + + # Get batch dimensions and model-level constants from host tensors (no device sync) + num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_seq = num_prefill + num_decode + num_tokens = num_prefill_tokens + num_decode + max_context_length = int(max_seq_info_host[0]) + max_num_requests = int(max_seq_info_host[3]) + # Use sliding_window for attention_window_size if provided, else full context length + attention_window_size = ( + sliding_window + if isinstance(sliding_window, int) and sliding_window > 0 + else max_context_length + ) + + # Get per-layer pool pointer tensor (stable address for CUDA graph replay) + host_kv_cache_pool_pointers = _GlobalTrtllmPlanner.get_pool_pointers_for_layer(kv_cache) + + # FP8 KV cache: lazily create scale tensors from float constants on first use + if kv_cache.dtype == torch.float8_e4m3fn: + if _GlobalTrtllmPlanner.kv_scale_orig_quant is None: + _GlobalTrtllmPlanner.kv_scale_orig_quant = torch.tensor( + [kv_scale_orig_quant], dtype=torch.float32, device=q.device + ) + _GlobalTrtllmPlanner.kv_scale_quant_orig = torch.tensor( + [kv_scale_quant_orig], dtype=torch.float32, device=q.device + ) + quant_mode = int(QuantMode.FP8_KV_CACHE) + else: + quant_mode = 0 + + # Reshape Q, K, V to [num_tokens, num_heads * head_dim] and fuse. + # Input is [bs, 1] (generate-only) or [1, total_seq_len] (prefill/mixed). + # With piecewise CUDA graphs the tensor may be padded to a bucket size + # (b*s > num_tokens), so flatten first and slice to the real token count. + q_shape_og = q.shape + q_flat = q.reshape(-1, num_heads * head_dim)[:num_tokens] + k_flat = k.reshape(-1, num_kv_heads * head_dim)[:num_tokens] + v_flat = v.reshape(-1, num_kv_heads * head_dim)[:num_tokens] + qkv_fused = torch.cat([q_flat, k_flat, v_flat], dim=-1).contiguous() + + # Prepare output + output = torch.empty(num_tokens, num_heads * head_dim, dtype=q.dtype, device=q.device) + + # Map SequenceInfo fields to thop.attention args + sequence_length = seq_len_with_cache[:num_seq] # device + context_lengths = seq_len[:num_seq] # device + host_past_kv_lengths = _GlobalTrtllmPlanner.host_past_kv_lengths[:num_seq] # host (pinned) + host_context_lengths = _GlobalTrtllmPlanner.host_context_lengths[:num_seq] # host (pinned) + + # thop-specific metadata from _GlobalTrtllmPlanner + host_request_types = _GlobalTrtllmPlanner.host_request_types[:num_seq] + host_total_kv_lens = _GlobalTrtllmPlanner.host_total_kv_lens + + # Block offsets from host_prepare + kv_cache_block_offsets = _GlobalTrtllmPlanner.block_offsets + + # Pool mapping (shared, always zeros since layer offset is in pool_pointers) + host_kv_cache_pool_mapping = _GlobalTrtllmPlanner.host_pool_mapping + + # Pack parameters for thop.attention + rotary_embedding_scales = [1.0, 1.0, 1.0] + rotary_embedding_max_position_info = [max_context_length, max_context_length] + spec_decoding_bool_params = [False, False, False] + spec_decoding_tensor_params = [None, None, None] + + sm_version = get_sm_version() + if sm_version >= 89: # Ada/Hopper + spec_decoding_tensor_params.extend([None, None, None]) + + mla_tensor_params = [None, None] + + thop.attention( + qkv_fused, # q (actually fused QKV) + None, # k (None when using fused QKV) + None, # v (None when using fused QKV) + output, # output + None, # output_sf (NVFP4) + _GlobalTrtllmPlanner.workspace, # workspace (module-level, like flashinfer) + sequence_length, # sequence_length + host_past_kv_lengths, # host_past_key_value_lengths + host_total_kv_lens, # host_total_kv_lens + context_lengths, # context_lengths + host_context_lengths, # host_context_lengths + host_request_types, # host_request_types + kv_cache_block_offsets, # kv_cache_block_offsets + host_kv_cache_pool_pointers, # host_kv_cache_pool_pointers + host_kv_cache_pool_mapping, # host_kv_cache_pool_mapping + None, # cache_indirection (beam search) + _GlobalTrtllmPlanner.kv_scale_orig_quant, # kv_scale_orig_quant + _GlobalTrtllmPlanner.kv_scale_quant_orig, # kv_scale_quant_orig + None, # out_scale + None, # rotary_inv_freq + None, # rotary_cos_sin + None, # latent_cache (MLA) + None, # q_pe (MLA) + None, # block_ids_per_seq + None, # attention_sinks + True, # is_fused_qkv + True, # update_kv_cache + 1, # predicted_tokens_per_seq + 0, # layer_idx (always 0; pool_pointers already encodes the layer offset) + num_heads, # num_heads + num_kv_heads, # num_kv_heads + head_dim, # head_size + tokens_per_block, # tokens_per_block + max_num_requests, # max_num_requests + max_context_length, # max_context_length + attention_window_size, # attention_window_size + 0, # sink_token_length + 1, # beam_width + int(AttentionMaskType.causal), # mask_type + quant_mode, # quant_mode + 1.0, # q_scaling + 0, # position_embedding_type + 0, # rotary_embedding_dim + 10000.0, # rotary_embedding_base + 0, # rotary_embedding_scale_type + rotary_embedding_scales, # rotary_embedding_scales + rotary_embedding_max_position_info, # rotary_embedding_max_position_info + True, # use_paged_context_fmha + 0, # attention_input_type + False, # is_mla_enable + max_num_requests, # chunked_prefill_buffer_batch_size + None, # q_lora_rank (MLA) + None, # kv_lora_rank (MLA) + None, # qk_nope_head_dim (MLA) + None, # qk_rope_head_dim (MLA) + None, # v_head_dim (MLA) + None, # mrope_rotary_cos_sin + None, # mrope_position_deltas + mla_tensor_params, # mla_tensor_params + None, # attention_chunk_size + None, # softmax_stats_tensor + spec_decoding_bool_params, # spec_decoding_bool_params + spec_decoding_tensor_params, # spec_decoding_tensor_params + None, # sparse_kv_indices + None, # sparse_kv_offsets + None, # sparse_attn_indices + None, # sparse_attn_offsets + 1, # sparse_attn_indices_block_size + 0, # sparse_mla_topk + None, # skip_softmax_threshold_scale_factor_prefill + None, # skip_softmax_threshold_scale_factor_decode + None, # skip_softmax_stat + None, # cu_q_seqlens + None, # cu_kv_seqlens + None, # fmha_scheduler_counter + None, # mla_bmm1_scale + None, # mla_bmm2_scale + None, # quant_q_buffer + ) + + # If input was padded (piecewise CG), embed the real output into a padded + # tensor so downstream static segments see the expected bucket-sized shape. + total_padded_tokens = q_shape_og[0] * q_shape_og[1] + if total_padded_tokens > num_tokens: + padded_output = torch.zeros( + total_padded_tokens, num_heads * head_dim, dtype=q.dtype, device=q.device + ) + padded_output[:num_tokens] = output + return padded_output.view(*q_shape_og) + return output.view(*q_shape_og) + + +@trtllm_mha_with_cache.register_fake +def trtllm_mha_with_cache_fake( + # Q, K, V inputs + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + # STANDARD METADATA (SequenceInfo fields used directly by thop) + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + seq_len_host: torch.Tensor, + input_pos_host: torch.Tensor, + seq_len_with_cache: torch.Tensor, + max_seq_info_host: torch.Tensor, + # CACHE + kv_cache: torch.Tensor, + # CONSTANTS (only truly un-inferable values) + scale: Optional[float], + sliding_window: Optional[int] = None, + kv_scale_orig_quant: float = 1.0, + kv_scale_quant_orig: float = 1.0, +) -> torch.Tensor: + """Fake implementation for torch.compile tracing.""" + return torch.empty_like(q.contiguous()) + + +# ============================================================================= +# AttentionDescriptor (analogous to FlashInferAttention) +# ============================================================================= + + +@AttentionRegistry.register("trtllm") +class TrtllmAttention(AttentionDescriptor): + """TRT-LLM attention backend for Auto-Deploy. + + Follows the same stateless descriptor pattern as ``FlashInferAttention``. + """ + + @classmethod + def get_attention_layout(cls) -> AttentionLayout: + """Get the attention layout expected by the backend.""" + return "bsnd" + + @classmethod + def get_num_qkv_args(cls) -> int: + """Get the number of qkv arguments expected by the source op.""" + return 3 + + @classmethod + def get_source_attention_op(cls) -> OpOverloadPacket: + """Get the source attention op that we target for replacement.""" + return torch.ops.auto_deploy.torch_attention + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + """Get the cached attention op.""" + return torch.ops.auto_deploy.trtllm_attention_mha_with_cache.default + + @classmethod + def get_standard_metadata_args(cls) -> List[str]: + """Get the list of standard metadata arguments from SequenceInfo.""" + return [ + "batch_info_host", + "seq_len", + "seq_len_host", + "input_pos_host", + "seq_len_with_cache", + "max_seq_info_host", + ] + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: KvCacheConfig + ) -> ResourceHandlerDict: + """Return only KV cache handler (no workspace handler, managed like flashinfer).""" + k_fake: FakeTensor = source_attn_node.args[1].meta["val"] + num_kv_heads = k_fake.shape[2] + head_dim = k_fake.shape[3] + + return { + "kv_cache": KVPagedResourceHandler( + num_kv_heads, + head_dim, + dtype=cls.resolve_cache_dtype(cache_config.dtype, k_fake.dtype), + kv_factor=2, + kv_layout="HND", + ) + } + + @classmethod + def get_host_prepare_metadata_function(cls) -> Optional[PrepareMetadataHostCallable]: + """Return host-side prepare function for thop-specific metadata.""" + return prepare_trtllm_metadata_host + + @classmethod + def get_constants(cls, source_attn_node: Node) -> List[Constant]: + """Extract constants from the source attention node. + + Returns scale, sliding_window, kv_scale_orig_quant, and kv_scale_quant_orig. + Everything else (num_heads, head_dim, max_context_length, etc.) is inferred + from tensor shapes or SequenceInfo metadata at runtime. + """ + # Sanity check: layout == "bsnd" + layout = source_attn_node.kwargs.get("layout", None) + if ( + layout is None + and len(source_attn_node.args) > 0 + and isinstance(source_attn_node.args[-1], str) + ): + layout = source_attn_node.args[-1] + if layout != "bsnd": + raise RuntimeError( + f"Expected torch_attention layout='bsnd' but got {layout!r} " + f"for node: {source_attn_node.format_node()}" + ) + + # Check other arguments + _attn_mask, _dropout_p, _is_causal = extract_op_args( + source_attn_node, "attn_mask", "dropout_p", "is_causal" + ) + + # Get scale + if len(source_attn_node.args) > 6: + scale = source_attn_node.args[6] + else: + scale = source_attn_node.kwargs.get("scale", None) + + if not (isinstance(scale, float) or scale is None): + ad_logger.warning(f"Provided {scale=}, is not a float. Using default scale instead.") + scale = None + + # Get sliding_window from source attention node + sliding_window = extract_op_args(source_attn_node, "sliding_window")[0] + + return [ + scale, + sliding_window, + 1.0, # kv_scale_orig_quant (hard-coded, same as FlashInfer) + 1.0, # kv_scale_quant_orig (hard-coded, same as FlashInfer) + ] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 4af2891c0aa2..502c6b526466 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -28,6 +28,7 @@ from abc import ABC, abstractmethod from typing import Dict, List, Literal, Optional, Protocol, Sequence, Set, Tuple, Type, Union +import numpy as np import torch from torch._ops import OpOverloadPacket from torch.fx import Node @@ -40,45 +41,88 @@ Constant = Union[int, float, str, None] +# Torch dtype → numpy dtype for fast list-to-tensor conversion. +# numpy's list→array conversion is ~2-3x faster than torch.tensor(list) for large lists. +_TORCH_TO_NUMPY_DTYPE: Dict[torch.dtype, np.dtype] = { + torch.int: np.int32, + torch.int32: np.int32, + torch.int64: np.int64, + torch.long: np.int64, + torch.float: np.float32, + torch.float32: np.float32, + torch.float64: np.float64, + torch.double: np.float64, + torch.float16: np.float16, + torch.bool: np.bool_, +} + + +def _list_to_tensor(data: list, dtype: torch.dtype) -> torch.Tensor: + """Convert a Python list to a tensor, using numpy for speed.""" + np_dtype = _TORCH_TO_NUMPY_DTYPE.get(dtype) + if np_dtype is not None: + return torch.from_numpy(np.array(data, dtype=np_dtype)) + return torch.tensor(data, dtype=dtype) + class PrepareMetadataHostCallable(Protocol): def __call__(self, **sequence_info_args: torch.Tensor) -> None: ... class InputBuffer: - """Manages contiguous memory buffers for efficient host-to-device transfers. + """Manages memory buffers for efficient host-to-device transfers. - This class consolidates multiple tensors into a single contiguous buffer on both - host (pinned memory) and device. This enables efficient bulk transfers with a - single async H2D copy instead of multiple small copies. + Supports two categories of tensors: - The buffer layout places the truncatable tensor (typically cache_loc) last, - allowing partial copies when the full buffer isn't needed. + - **Contiguous tensors** (default): packed into a single contiguous buffer on both + host (pinned) and device. Copied in one bulk async H2D transfer. + - **Truncatable tensors** (``truncatable=True`` in spec): each gets its own separate + host+device buffer pair, copied independently with truncation to actual length. + Use this for large, variable-length tensors (e.g., ``cache_loc``) to avoid + copying unused capacity. Usage: - 1. Create InputBuffer with tensor specifications (name, max_numel, dtype) + 1. Create InputBuffer with tensor specifications 2. Use store() to write data to the pinned host buffer - 3. Call copy_to_device() to perform a single async H2D transfer + 3. Call copy_to_device() to perform async H2D transfers 4. Access device tensors via get_view() """ - def __init__(self, tensor_specs: List[Tuple[str, int, torch.dtype]]): + def __init__(self, tensor_specs: List[Tuple]): """Initialize the InputBuffer. Args: - tensor_specs: Ordered list of (name, max_numel, dtype) tuples. - The last tensor is treated as truncatable during copy. + tensor_specs: Ordered list of tensor specs. Each element is either: + - ``(name, max_numel, dtype)`` for contiguous tensors (default) + - ``(name, max_numel, dtype, True)`` for truncatable tensors """ - self._tensor_specs = {name: (numel, dtype) for name, numel, dtype in tensor_specs} - self._tensor_order = [name for name, _, _ in tensor_specs] + # Parse specs into canonical form: (name, numel, dtype, truncatable) + parsed = [] + for spec in tensor_specs: + if len(spec) == 4: + name, numel, dtype, truncatable = spec + else: + name, numel, dtype = spec + truncatable = False + parsed.append((name, numel, dtype, truncatable)) + + self._tensor_specs: Dict[str, Tuple[int, torch.dtype]] = { + name: (numel, dtype) for name, numel, dtype, _ in parsed + } + self._tensor_order = [name for name, _, _, _ in parsed] + self._truncatable_names: Set[str] = {name for name, _, _, t in parsed if t} + self._contiguous_names = [name for name, _, _, t in parsed if not t] - # Calculate offsets for each tensor (aligned to dtype's element size) + # Track current lengths for each tensor (for truncation optimization) + self._current_lengths: Dict[str, int] = {name: 0 for name in self._tensor_order} + + # === CONTIGUOUS BUFFER (small, fixed-size tensors) === self._offsets: Dict[str, int] = {} self._byte_sizes: Dict[str, int] = {} current_offset = 0 - for name, numel, dtype in tensor_specs: - # Align to the tensor's element size for proper memory access + for name in self._contiguous_names: + numel, dtype = self._tensor_specs[name] alignment = dtype.itemsize aligned_offset = (current_offset + alignment - 1) // alignment * alignment byte_size = numel * dtype.itemsize @@ -86,46 +130,53 @@ def __init__(self, tensor_specs: List[Tuple[str, int, torch.dtype]]): self._byte_sizes[name] = byte_size current_offset = aligned_offset + byte_size - # Total buffer size self._total_bytes = current_offset - # Allocate contiguous buffers (device buffer starts on default device, use to() to move) - self._device_buffer = torch.empty(self._total_bytes, dtype=torch.uint8) - self._host_buffer = torch.empty( - self._total_bytes, dtype=torch.uint8, device="cpu", pin_memory=True - ) - - # Create persistent views into device and host buffers - # Persistent views help us identify the arguments as static during graph capture. - self._device_views = self._create_views(self._device_buffer) - self._host_views = self._create_views(self._host_buffer) - - # Track current lengths for each tensor (for truncation optimization) - self._current_lengths: Dict[str, int] = {name: 0 for name in self._tensor_order} + if self._total_bytes > 0: + self._device_buffer = torch.empty(self._total_bytes, dtype=torch.uint8) + self._host_buffer = torch.empty( + self._total_bytes, dtype=torch.uint8, device="cpu", pin_memory=True + ) + else: + self._device_buffer = torch.empty(0, dtype=torch.uint8) + self._host_buffer = torch.empty(0, dtype=torch.uint8, device="cpu") + + # Create persistent views into contiguous buffers + self._device_views: Dict[str, torch.Tensor] = {} + self._host_views: Dict[str, torch.Tensor] = {} + self._create_contiguous_views() + + # === TRUNCATABLE BUFFERS (large, variable-length tensors) === + self._trunc_device_bufs: Dict[str, torch.Tensor] = {} + self._trunc_host_bufs: Dict[str, torch.Tensor] = {} + for name in self._truncatable_names: + numel, dtype = self._tensor_specs[name] + byte_size = numel * dtype.itemsize + self._trunc_device_bufs[name] = torch.empty(byte_size, dtype=torch.uint8) + self._trunc_host_bufs[name] = torch.empty( + byte_size, dtype=torch.uint8, device="cpu", pin_memory=True + ) + # Create typed views + self._device_views[name] = self._trunc_device_bufs[name].view(dtype) + self._host_views[name] = self._trunc_host_bufs[name].view(dtype) - def _create_views(self, buffer: torch.Tensor) -> Dict[str, torch.Tensor]: - """Create views into the given buffer for each tensor.""" - views = {} - for name in self._tensor_order: + def _create_contiguous_views(self) -> None: + """Create typed views into the contiguous host and device buffers.""" + for name in self._contiguous_names: offset = self._offsets[name] byte_size = self._byte_sizes[name] _, dtype = self._tensor_specs[name] - views[name] = buffer[offset : offset + byte_size].view(dtype) - return views + self._device_views[name] = self._device_buffer[offset : offset + byte_size].view(dtype) + self._host_views[name] = self._host_buffer[offset : offset + byte_size].view(dtype) @property def tensor_names(self) -> List[str]: - """Return the list of tensor names in buffer order.""" + """Return the list of tensor names in spec order.""" return self._tensor_order.copy() - @property - def _truncatable_name(self) -> str: - """Return the name of the truncatable tensor.""" - return self._tensor_order[-1] - @property def total_bytes(self) -> int: - """Total size of the buffer in bytes.""" + """Total size of the contiguous buffer in bytes.""" return self._total_bytes @property @@ -134,74 +185,38 @@ def device(self) -> torch.device: return self._device_buffer.device def get_view(self, name: str) -> torch.Tensor: - """Get the device tensor view for the specified name. - - Args: - name: Name of the tensor. - - Returns: - A view into the device buffer for the specified tensor. - """ + """Get the device tensor view for the specified name.""" return self._device_views[name] def get_view_at_current_length(self, name: str) -> torch.Tensor: - """Get the device tensor view for the specified name at the current length. - - Args: - name: Name of the tensor. - - Returns: - A view into the device buffer for the specified tensor at the current length. - """ + """Get the device tensor view truncated to the current stored length.""" return self._device_views[name][: self._current_lengths[name]] def get_host_view(self, name: str) -> torch.Tensor: - """Get the host tensor view for the specified name. - - Args: - name: Name of the tensor. - - Returns: - A view into the pinned host buffer for the specified tensor. - """ + """Get the host tensor view for the specified name.""" return self._host_views[name] def get_capacity(self, name: str) -> int: - """Get the maximum number of elements for the specified tensor. - - Args: - name: Name of the tensor. - - Returns: - Maximum number of elements that can be stored. - """ + """Get the maximum number of elements for the specified tensor.""" numel, _ = self._tensor_specs[name] return numel def get_current_length(self, name: str) -> int: - """Get the current stored length for the specified tensor. - - Args: - name: Name of the tensor. - - Returns: - Number of elements currently stored in the tensor. - """ + """Get the current stored length for the specified tensor.""" return self._current_lengths[name] def store( self, name: str, - data: List[Number], + data: torch.Tensor, fill_value: Optional[Number] = None, ) -> int: - """Store data into the host buffer. + """Store a tensor into the pinned host buffer. Args: name: Name of the tensor to store to. - data: List of values to store. - fill_value: Optional value to fill the entire tensor with before storing. - If None, only the provided data is written. + data: 1-D torch.Tensor to store. + fill_value: Optional value to fill the entire buffer with before storing. Returns: Number of elements stored. @@ -209,95 +224,99 @@ def store( numel, dtype = self._tensor_specs[name] host_view = self.get_host_view(name) - # Fill with default value if specified if fill_value is not None: host_view.fill_(fill_value) - # Convert list to tensor and copy to host buffer - length = len(data) + length = data.numel() assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}" - - temp_tensor = torch.tensor(data, dtype=dtype) - host_view[:length].copy_(temp_tensor) + # Use numpy for the memcpy into pinned memory — avoids torch dispatcher overhead + dst = host_view[:length].numpy() + src = (data if data.dtype == dtype else data.to(dtype)).numpy() + np.copyto(dst, src) self._current_lengths[name] = length return length def copy_to_device(self) -> None: - """Copy from host buffer to device buffer. + """Copy from host buffers to device buffers. - Uses the current length of the truncatable tensor (last in spec) to minimize - transfer size. All tensors before the truncatable one are fully copied. + Contiguous tensors are copied in a single bulk transfer. + Truncatable tensors are each copied independently, truncated to actual length. """ - # Calculate bytes to copy based on truncatable tensor's current length - truncatable_len = self._current_lengths[self._truncatable_name] - truncatable_offset = self._offsets[self._truncatable_name] - truncatable_dtype = self._tensor_specs[self._truncatable_name][1] - copy_bytes = truncatable_offset + truncatable_len * truncatable_dtype.itemsize - - # Single async copy with nvtx_range("ad_input_buffer_h2d_copy"): - self._device_buffer[:copy_bytes].copy_( - self._host_buffer[:copy_bytes], non_blocking=True - ) + # Copy contiguous buffer in one shot + if self._total_bytes > 0: + self._device_buffer[: self._total_bytes].copy_( + self._host_buffer[: self._total_bytes], non_blocking=True + ) + + # Copy each truncatable tensor independently, truncated to current length + for name in self._truncatable_names: + length = self._current_lengths[name] + if length > 0: + _, dtype = self._tensor_specs[name] + copy_bytes = length * dtype.itemsize + self._trunc_device_bufs[name][:copy_bytes].copy_( + self._trunc_host_bufs[name][:copy_bytes], non_blocking=True + ) def resize(self, name: str, new_capacity: int) -> None: - """Resize a tensor's capacity. + """Resize a truncatable tensor's capacity. - This operation is only supported for the last tensor in the buffer to avoid - complex offset recalculations. + Only truncatable tensors can be resized (they have independent buffers). Args: name: Name of the tensor to resize. new_capacity: New maximum number of elements for the tensor. """ - assert name == self._truncatable_name, ( - f"Can only resize the last tensor in the buffer ('{self._truncatable_name}'). " - f"Attempted to resize '{name}'." + assert name in self._truncatable_names, ( + f"Can only resize truncatable tensors. '{name}' is not truncatable. " + f"Truncatable tensors: {self._truncatable_names}" ) old_numel, dtype = self._tensor_specs[name] if new_capacity <= old_numel: - return # No need to resize if new capacity is smaller or equal + return - # Update tensor specs self._tensor_specs[name] = (new_capacity, dtype) - - # Calculate new byte size for this tensor new_byte_size = new_capacity * dtype.itemsize - self._byte_sizes[name] = new_byte_size - - # Update total bytes (offset stays the same since it's the last tensor) - self._total_bytes = self._offsets[name] + new_byte_size # Resize device buffer in-place - self._device_buffer.resize_(self._total_bytes) + self._trunc_device_bufs[name].resize_(new_byte_size) - # Host buffer must be re-allocated to ensure we have pinned memory - old_host_buffer = self._host_buffer - self._host_buffer = torch.empty( - self._total_bytes, dtype=torch.uint8, device="cpu", pin_memory=True + # Host buffer must be re-allocated for pinned memory + old_host = self._trunc_host_bufs[name] + self._trunc_host_bufs[name] = torch.empty( + new_byte_size, dtype=torch.uint8, device="cpu", pin_memory=True ) - self._host_buffer[: old_host_buffer.numel()].copy_(old_host_buffer) - del old_host_buffer + self._trunc_host_bufs[name][: old_host.numel()].copy_(old_host) + del old_host - # Recreate views after the update - self._device_views = self._create_views(self._device_buffer) - self._host_views = self._create_views(self._host_buffer) + # Recreate typed views + self._device_views[name] = self._trunc_device_bufs[name].view(dtype) + self._host_views[name] = self._trunc_host_bufs[name].view(dtype) def to(self, *args, **kwargs) -> None: - """Move the device buffer to a new device/dtype. - - Note: This recreates the device views after moving. - """ + """Move all device buffers to a new device/dtype.""" old_device = self._device_buffer.device + + # Move contiguous buffer self._device_buffer = self._device_buffer.to(*args, **kwargs) + # Move truncatable buffers + for name in self._truncatable_names: + self._trunc_device_bufs[name] = self._trunc_device_bufs[name].to(*args, **kwargs) + # Recreate views if device changed if old_device != self._device_buffer.device: - self._device_views = self._create_views(self._device_buffer) + self._create_contiguous_views() + for name in self._truncatable_names: + _, dtype = self._tensor_specs[name] + self._device_views[name] = self._trunc_device_bufs[name].view(dtype) +# TODO (lucaslie): as this list is growing we may want to "upstream" the active arguments to +# nest_sequences and _prepare_inputs to skip on unnecessary computations. class SequenceInfo: """An interface to hold information about how the sequence is laid out and stored in cache. @@ -348,6 +367,16 @@ class SequenceInfo: - batch_info: [num_prefill, num_prefill_tokens, num_decode] Batch metadata containing the number of prefill sequences, total prefill tokens, and number of decode sequences. + - max_seq_info: [max_context_length, max_blocks_per_seq, block_offset_multiplier, max_batch_size] + Model-level constants for the attention kernel: maximum context length (equal to max_seq_len), + maximum number of KV cache blocks per sequence (ceil(max_seq_len / tokens_per_block)), + block offset multiplier derived from kv_cache strides, and maximum batch size. These are + set once via update_cache_information() after cache initialization and remain constant. + - page_seq_indices: [si_0, si_1, ..., si_{np-1}] where si_j is the sequence index that page j + belongs to. For example, if seq 0 has 2 pages and seq 1 has 3, then + page_seq_indices = [0, 0, 1, 1, 1]. + - page_in_seq: [pi_0, pi_1, ..., pi_{np-1}] where pi_j is the page index within its sequence. + For example, if seq 0 has 2 pages and seq 1 has 3, then page_in_seq = [0, 1, 0, 1, 2]. - cache_loc: [c_0, c_1, ..., c_{np-1}] where np is total number of pages allocated to describe all sequences in the batch. Each value is a page index in the cache. - logits_gather_indices: [g_0, g_1, ..., g_{s_total-1}] @@ -405,6 +434,7 @@ def __init__( self.max_seq_len = max_seq_len self.max_batch_size = max_batch_size self.tokens_per_block = tokens_per_block or max_seq_len + self.max_blocks_per_seq = math.ceil(max_seq_len / self.tokens_per_block) # NOTE (lucaslie): +1 is a WAR to address issue when using flashinfer attention with # (max_batch_size, max_seq_len) input in trtllm runtime. # see https://github.com/NVIDIA/TensorRT-LLM/issues/4504 @@ -427,7 +457,8 @@ def __init__( # log parameters ad_logger.info( - f"[SequenceInfo:] {self.max_seq_len=}, {self.max_batch_size=}, {self.max_num_tokens=}" + f"[SequenceInfo:] {self.max_seq_len=}, {self.max_batch_size=}, {self.max_num_tokens=}, " + f"{self.max_blocks_per_seq=}, {self.tokens_per_block=}" ) # indicator if extra args are activated that are needed for cached attention backends @@ -452,15 +483,20 @@ def __init__( ("slot_idx", self.max_batch_size, torch.long), ("use_initial_states", self.max_batch_size, torch.bool), ("batch_info", 3, torch.int), + # [max_context_length, max_blocks_per_seq, block_offset_multiplier, max_batch_size] + ("max_seq_info", 4, torch.int), ("logits_gather_indices", self.max_num_tokens, torch.long), ("logits_gather_info", 2, torch.int), # OTHER FIELDS WHERE WE NEED EFFICIENT HOST<>DEVICE TRANSFER ("_gather_idx", self.max_num_tokens, torch.int), ("_mask_scatter_indices", self.max_num_tokens, torch.int), - # cache_loc is LAST for truncation optimization (it can be the largest tensor) + # TRUNCATABLE TENSORS: large, variable-length, each independently truncated during + # H2D copy to avoid copying unused capacity. # NOTE: sufficient for max_num_tokens forward pass. will be resized when KVCacheManager # is created. - ("cache_loc", self.max_num_tokens, torch.int), + ("cache_loc", self.max_num_tokens, torch.int, True), + ("page_seq_indices", self.max_num_tokens, torch.int, True), + ("page_in_seq", self.max_num_tokens, torch.int, True), ] # Create the InputBuffer that manages contiguous host and device memory @@ -470,13 +506,25 @@ def __init__( f"{name}_host" for name in self._input_buffer.tensor_names } - # Initialize args_list from tensor specs - self._args_list: Dict[str, List[int]] = { - name: [0] * numel for name, numel, _ in tensor_specs + # Initialize args_list from tensor specs (all entries are tensors) + self._args_list: Dict[str, torch.Tensor] = { + spec[0]: torch.zeros(spec[1], dtype=spec[2]) for spec in tensor_specs } self._active_args = ("input_ids", "position_ids") self._shapeable_args = ("input_ids", "position_ids", "input_ids_host", "position_ids_host") + + # Args that require copy to InputBuffer but are NOT included in named_args / graph inputs. + # Pre-populated with args that are always needed regardless of graph inclusion. + self._requires_copy: Set[str] = { + "max_seq_info", # written once at cache init (update_cache_information) + "page_seq_indices", # used by host-prepare (not a graph input) + "page_in_seq", # used by host-prepare (not a graph input) + "logits_gather_indices", # always needed for logits gathering + "logits_gather_info", # always needed for logits gathering + "_gather_idx", # overlap scheduler metadata + "_mask_scatter_indices", # overlap scheduler metadata + } ############################################################################################ # EXTRA TENSOR FIELDS ###################################################################### @@ -486,9 +534,45 @@ def __init__( # HOST PREPARE FOR ATTENTION FORWARD ####################################################### self._host_prepare_functions: List[Tuple[PrepareMetadataHostCallable, List[str]]] = [] + # Padded num_tokens for piecewise CUDA graph bucket alignment. + # When set, _shape_for_forward uses this instead of total_num_tokens so that + # input tensors match the captured CUDA graph bucket size. + self._padded_num_tokens: Optional[int] = None + + # Sorted list of piecewise CUDA graph bucket sizes. + # Populated during compilation; used at runtime to find nearest bucket. + self._piecewise_bucket_sizes: List[int] = [] + # call reset once to set a consistent initial state self.reset() + @property + def padded_num_tokens(self) -> Optional[int]: + """Return the padded num_tokens for piecewise CG bucket alignment, or None.""" + return self._padded_num_tokens + + @padded_num_tokens.setter + def padded_num_tokens(self, value: Optional[int]): + """Set the padded num_tokens for piecewise CG bucket alignment.""" + self._padded_num_tokens = value + + @property + def piecewise_bucket_sizes(self) -> List[int]: + """Return the sorted list of piecewise CG bucket sizes.""" + return self._piecewise_bucket_sizes + + @piecewise_bucket_sizes.setter + def piecewise_bucket_sizes(self, value: List[int]): + """Set the sorted piecewise CG bucket sizes.""" + self._piecewise_bucket_sizes = sorted(value) + + def find_nearest_piecewise_bucket(self, num_tokens: int) -> Optional[int]: + """Find smallest piecewise bucket >= num_tokens, or None.""" + for bucket in self._piecewise_bucket_sizes: + if bucket >= num_tokens: + return bucket + return None + @property def device(self) -> torch.device: return self._input_buffer.device @@ -510,10 +594,13 @@ def _shape_for_forward(self, tnsr: torch.Tensor) -> torch.Tensor: sl = self.seq_len[0] # use [1,total_len] shape to indicate non-generate-only batch for cached attention else: - bs, sl = 1, self.total_num_tokens + # Use padded size if set (for piecewise CG bucket alignment), + # otherwise use real total_num_tokens. + effective_num_tokens = self._padded_num_tokens or self.total_num_tokens + bs, sl = 1, effective_num_tokens - # truncate to total tokens now, reshape, and return - return tnsr[: self.total_num_tokens].view(bs, sl, *tnsr.shape[1:]) + # truncate to effective tokens now, reshape, and return + return tnsr[: bs * sl].view(bs, sl, *tnsr.shape[1:]) def _get_arg(self, name: str) -> torch.Tensor: """Get the argument from the input buffer either on device or host.""" @@ -556,19 +643,19 @@ def args(self) -> Tuple[torch.Tensor, ...]: @property def seq_len(self) -> List[int]: - return self._args_list["seq_len"].copy() + return self._args_list["seq_len"].tolist() @property def input_pos(self) -> List[int]: - return self._args_list["input_pos"].copy() + return self._args_list["input_pos"].tolist() @property def cache_loc(self) -> List[int]: - return self._args_list["cache_loc"].copy() + return self._args_list["cache_loc"].tolist() @property def pages_per_seq(self) -> List[int]: - return self._args_list["pages_per_seq"].copy() + return self._args_list["pages_per_seq"].tolist() @property def num_sequences(self) -> int: @@ -602,11 +689,24 @@ def estimate_cache_tokens_per_forward(self) -> int: num_blocks_estimate = num_blocks_estimate_per_seq * self.max_batch_size return num_blocks_estimate * self.tokens_per_block - def estimate_cache_loc_capacity(self, num_blocks: int) -> None: - """Estimate needed capacity of cache_loc based on available blocks and resize.""" - # set num_blocks + def update_cache_information(self, num_blocks: int, block_offset_multiplier: int = 0) -> None: + """Update cache information after cache manager creation. + + Sets num_blocks and block_offset_multiplier, writes max_seq_info to the host buffer + (constant after this call), and resizes cache_loc if needed. + """ + # set num_blocks and block_offset_multiplier self._num_blocks = num_blocks + # write max_seq_info once (constant after this call) + max_seq_info = [ + self.max_seq_len, + self.max_blocks_per_seq, + block_offset_multiplier, + self.max_batch_size, + ] + self._store_arg("max_seq_info", max_seq_info) + # get current capacity cache_loc_capacity = self._input_buffer.get_capacity("cache_loc") @@ -624,10 +724,10 @@ def estimate_cache_loc_capacity(self, num_blocks: int) -> None: estimated_capacity = estimated_capacity + 1 if estimated_capacity > cache_loc_capacity: - self._input_buffer.resize("cache_loc", estimated_capacity) - # Also resize the args_list to match - old_size = len(self._args_list["cache_loc"]) - self._args_list["cache_loc"].extend([0] * (estimated_capacity - old_size)) + # Resize all truncatable page tensors together (they share the same max size) + for tensor_name in ("cache_loc", "page_seq_indices", "page_in_seq"): + self._input_buffer.resize(tensor_name, estimated_capacity) + self._args_list[tensor_name] = torch.zeros(estimated_capacity, dtype=torch.int) @staticmethod def _get_page_assignments( @@ -688,6 +788,25 @@ def activate_arg(self, arg_name: str) -> bool: return True return False + def require_copy(self, arg_name: str) -> bool: + """Mark an argument as requiring copy to InputBuffer. + + Unlike activate_arg, this does NOT add the arg to _named_args / graph inputs. + It only ensures _store_arg writes to InputBuffer so _get_arg returns updated values. + + Use cases: + - Host-prepare functions that need args outside the CUDA graph + - Args that are always needed in InputBuffer regardless of graph inclusion + + Returns: + True if the argument was newly marked, False if already marked. + """ + assert arg_name in self.available_args, f"{arg_name=} not found in {self.available_args}" + if arg_name not in self._requires_copy: + self._requires_copy.add(arg_name) + return True + return False + def to(self, *args, **kwargs) -> None: # Move the InputBuffer (which recreates views automatically) self._input_buffer.to(*args, **kwargs) @@ -733,9 +852,13 @@ def set_example_sequence( ) def set_max_num_tokens_sample(self) -> None: - """Set an example sequence with max_num_tokens.""" - # TODO (lucaslie): understand what this implies for extra arguments - seq_len = self.max_num_tokens // self.max_batch_size + """Set an example sequence with max_num_tokens. + + The per-sequence length is capped to the maximum that fits in the paged KV cache + (max_blocks_per_seq * tokens_per_block) to avoid exceeding block_offsets capacity. + """ + max_cache_tokens_per_seq = self.max_blocks_per_seq * self.tokens_per_block + seq_len = min(self.max_num_tokens // self.max_batch_size, max_cache_tokens_per_seq) input_ids = torch.ones(self.max_batch_size, seq_len, dtype=torch.int).tolist() self.set_example_sequence(input_ids) @@ -766,31 +889,38 @@ def _flatten(nested_seqs: Sequence[Sequence[int]]) -> List[int]: def _store_arg( self, name: str, - tnsr_like: List[Number], + data: "Union[List[Number], torch.Tensor]", reset_val: Optional[Number] = None, - force_copy: bool = False, ) -> None: """Store the argument into the pinned host buffer for later batch transfer to device. The data is stored in the host-side pinned memory buffer managed by InputBuffer. The actual H2D transfer happens in a single batch at the end of nest_sequences(). + Lists are converted to tensors at the boundary so the rest of the pipeline is + tensor-only. + Args: name: Name of the argument to store. - tnsr_like: List of values to store. + data: List of values or a 1-D torch.Tensor to store. reset_val: Value to reset/fill the tensor with before writing data. - force_copy: Whether to force immediate copy to device (for use outside nest_sequences). """ with nvtx_range(f"ad_store_on_host_seq_info_arg_{name}"): - # Always store list object for Python access - self._args_list[name] = tnsr_like.copy() + # Convert to tensor at the boundary (numpy is ~2-3x faster than torch.tensor for large lists) + if not isinstance(data, torch.Tensor): + _, dtype = self._input_buffer._tensor_specs[name] + data = _list_to_tensor(data, dtype) - # Only store to buffer when the argument is active or force_copy is True - if not (name in self._active_args or f"{name}_host" in self._active_args or force_copy): + self._args_list[name] = data + + # Only store to buffer when the argument is active or requires copy + is_active = name in self._active_args or f"{name}_host" in self._active_args + is_required = name in self._requires_copy or f"{name}_host" in self._requires_copy + if not (is_active or is_required): return # Store to the InputBuffer's pinned host memory - self._input_buffer.store(name, tnsr_like, fill_value=reset_val) + self._input_buffer.store(name, data, fill_value=reset_val) def _store_extra_arg( self, name: str, tnsr_like: Optional[Union[torch.Tensor, Sequence[torch.Tensor]]] @@ -835,6 +965,7 @@ def nest_sequences( logits_gather_info: Optional[Sequence[int]] = None, _gather_idx: Optional[Sequence[int]] = None, _mask_scatter_indices: Optional[Sequence[int]] = None, + _ungathered_input_ids: Optional[torch.Tensor] = None, **extra_args: Dict[str, Union[torch.Tensor, Sequence[torch.Tensor]]], ) -> None: """Create and store sequence information for the next forward pass. @@ -867,6 +998,8 @@ def nest_sequences( logits_gather_info: Info list containing [num_tokens_to_gather, gather_required]. _gather_idx: Gather indices for the overlap scheduler to reorder input tokens. _mask_scatter_indices: Mask scatter indices for the overlap scheduler. + _ungathered_input_ids: Optional tensor of ungathered input ids from the overlap + scheduler. If provided, triggers rescatter_input_ids after H2D copy. extra_args: Extra arguments to be stored in the interface. This i/f will ensure that all sequence info args are updated accordingly. Reset values are @@ -920,6 +1053,23 @@ def nest_sequences( self._store_arg("cache_loc", cache_loc) self._store_arg("pages_per_seq", pages_per_seq) + # Auto-compute page_seq_indices and page_in_seq from pages_per_seq using + # vectorized torch ops. This avoids the slow Python-list-to-tensor conversion + # that dominates host time for large page counts (e.g., 50k ISL models). + # page_seq_indices[j] = which sequence page j belongs to + # page_in_seq[j] = which page within that sequence (0-indexed) + pages_per_seq_t = torch.tensor(pages_per_seq, dtype=torch.int) + seq_indices = torch.arange(len(pages_per_seq), dtype=torch.int) + page_seq_indices_t = torch.repeat_interleave(seq_indices, pages_per_seq_t) + cu_pages = torch.zeros(len(pages_per_seq) + 1, dtype=torch.int) + cu_pages[1:] = pages_per_seq_t.cumsum(0) + total_pages = cu_pages[-1].item() + page_in_seq_t = torch.arange(total_pages, dtype=torch.int) - torch.repeat_interleave( + cu_pages[:-1], pages_per_seq_t + ) + self._store_arg("page_seq_indices", page_seq_indices_t) + self._store_arg("page_in_seq", page_in_seq_t) + # update cumulative number of pages if cu_num_pages is None: pages_per_seq = self.pages_per_seq @@ -951,21 +1101,21 @@ def nest_sequences( if logits_gather_indices is None: # default is to gather all logits logits_gather_indices = list(range(self.total_num_tokens)) - self._store_arg("logits_gather_indices", logits_gather_indices, force_copy=True) + self._store_arg("logits_gather_indices", logits_gather_indices) # check for updated logits_gather_info if logits_gather_info is None: logits_gather_info = [len(logits_gather_indices), 1] - self._store_arg("logits_gather_info", logits_gather_info, force_copy=True) + self._store_arg("logits_gather_info", logits_gather_info) ### UPDATE OVERLAP SCHEDULER METADATA ###################################################### # check for updated _gather_idx if _gather_idx is not None: - self._store_arg("_gather_idx", _gather_idx, force_copy=True) + self._store_arg("_gather_idx", _gather_idx) # check for updated _mask_scatter_indices if _mask_scatter_indices is not None: - self._store_arg("_mask_scatter_indices", _mask_scatter_indices, force_copy=True) + self._store_arg("_mask_scatter_indices", _mask_scatter_indices) ### UPDATE EXTRA INPUTS #################################################################### self._extra_args = {} @@ -977,6 +1127,14 @@ def nest_sequences( # The copy is truncated at the end of cache_loc to minimize transfer size self._input_buffer.copy_to_device() + ### RESCATTER + HOST PREPARE ############################################################### + # Rescatter input_ids if ungathered tokens are provided (overlap scheduler) + if _ungathered_input_ids is not None: + self.rescatter_input_ids(_ungathered_input_ids) + + # Run host-prepare functions for attention forward (e.g. trtllm block_offsets computation) + self.run_host_prepare_for_attention_forward() + @nvtx_range("ad_rescatter_input_ids") def rescatter_input_ids(self, ungathered_input_ids: torch.Tensor): """Re-scatter the provided ungathered input ids into the input_ids tensor. @@ -1025,6 +1183,9 @@ def register_host_prepare_for_attention_forward( self, host_function: PrepareMetadataHostCallable, args: List[str] ): self._host_prepare_functions.append((host_function, args)) + # Ensure all host-prepare args are stored to InputBuffer via _requires_copy + for arg in args: + self.require_copy(arg) def run_host_prepare_for_attention_forward(self) -> None: for host_function, args in self._host_prepare_functions: diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py index b11ccc7ee61b..17a5d87480d7 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py @@ -18,9 +18,11 @@ This module provides linear layer implementations: - linear: Linear layer operations - torch_router: MoE router operations +- swiglu: SwiGLU MLP custom operations """ __all__ = [ "linear", "torch_router", + "swiglu", ] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py new file mode 100644 index 000000000000..e13acb6f79fa --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SwiGLU MLP custom operations for graph transformation. + +This module provides custom operators for SwiGLU MLP fusion: +- torch_swiglu_mlp: Intermediate representation after pattern matching +- fused_swiglu_mlp: Fused implementation with concatenated gate+up weights +""" + +from typing import Optional + +import torch +import torch.nn.functional as F + +try: + from flashinfer.activation import silu_and_mul as _flashinfer_silu_and_mul +except ImportError: + _flashinfer_silu_and_mul = None + + +def _silu_and_mul(x: torch.Tensor) -> torch.Tensor: + """SwiGLU activation: split x in half, apply silu to first half, multiply with second half. + + Uses FlashInfer's fused kernel when available, falls back to manual implementation. + """ + if _flashinfer_silu_and_mul is not None: + return _flashinfer_silu_and_mul(x) + gate, up = x.chunk(2, dim=-1) + return F.silu(gate) * up + + +@torch.library.custom_op("auto_deploy::torch_swiglu_mlp", mutates_args=()) +def torch_swiglu_mlp( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_bias: Optional[torch.Tensor], + up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Standardized SwiGLU MLP operation. + + Computes: silu(x @ gate.T + gate_bias) * (x @ up.T + up_bias) @ down.T + down_bias + + This is the intermediate representation used after pattern matching, + before weight fusion is applied. + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_weight: Gate projection weight of shape [intermediate_size, hidden_size]. + up_weight: Up projection weight of shape [intermediate_size, hidden_size]. + down_weight: Down projection weight of shape [hidden_size, intermediate_size]. + gate_bias: Optional gate projection bias of shape [intermediate_size]. + up_bias: Optional up projection bias of shape [intermediate_size]. + down_bias: Optional down projection bias of shape [hidden_size]. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + gate_out = F.linear(input, gate_weight, gate_bias) + up_out = F.linear(input, up_weight, up_bias) + hidden = F.silu(gate_out) * up_out + return F.linear(hidden, down_weight, down_bias) + + +@torch_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_bias: Optional[torch.Tensor], + up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape is [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) + + +@torch.library.custom_op("auto_deploy::fused_swiglu_mlp", mutates_args=()) +def fused_swiglu_mlp( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Fused SwiGLU MLP with concatenated gate+up weights. + + Performs a single matmul for gate and up projections, then splits the result. + Computes: silu(gate_out) * up_out @ down.T + down_bias + where gate_out, up_out = split(x @ gate_up.T + gate_up_bias) + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_up_weight: Concatenated gate+up weight of shape [2*intermediate_size, hidden_size]. + down_weight: Down projection weight of shape [hidden_size, intermediate_size]. + gate_up_bias: Optional concatenated gate+up bias of shape [2*intermediate_size]. + down_bias: Optional down projection bias of shape [hidden_size]. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + # Single matmul for both gate and up projections + gate_up_out = F.linear(input, gate_up_weight, gate_up_bias) + + # Apply SwiGLU activation: split, silu(gate) * up (uses FlashInfer when available) + hidden = _silu_and_mul(gate_up_out) + + # Down projection + return F.linear(hidden, down_weight, down_bias) + + +@fused_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape is [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) + + +# ── NVFP4 quantized SwiGLU ops ────────────────────────────────────────────── + + +@torch.library.custom_op("auto_deploy::torch_nvfp4_swiglu_mlp", mutates_args=()) +def torch_nvfp4_swiglu_mlp( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_input_scale: torch.Tensor, + gate_weight_scale: torch.Tensor, + gate_alpha: torch.Tensor, + up_input_scale: torch.Tensor, + up_weight_scale: torch.Tensor, + up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """NVFP4 quantized SwiGLU MLP operation (intermediate representation). + + Computes: silu(nvfp4_linear(x, gate)) * nvfp4_linear(x, up) -> nvfp4_linear(down) + + This is the intermediate representation used after pattern matching for NVFP4 + quantized checkpoints, before gate+up weight fusion is applied. + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_weight: FP4 packed gate weight [intermediate_size, hidden_size/2] uint8. + up_weight: FP4 packed up weight [intermediate_size, hidden_size/2] uint8. + down_weight: FP4 packed down weight [hidden_size, intermediate_size/2] uint8. + gate_input_scale: Input scale for gate projection. + gate_weight_scale: Per-block weight scale for gate projection. + gate_alpha: Alpha (combined scale) for gate projection. + up_input_scale: Input scale for up projection. + up_weight_scale: Per-block weight scale for up projection. + up_alpha: Alpha (combined scale) for up projection. + down_input_scale: Input scale for down projection. + down_weight_scale: Per-block weight scale for down projection. + down_alpha: Alpha (combined scale) for down projection. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + gate_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + input, + gate_weight, + None, + input_scale=[gate_input_scale], + weight_scale=[gate_weight_scale, gate_alpha], + input_zp=[], + weight_zp=[], + ) + up_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + input, + up_weight, + None, + input_scale=[up_input_scale], + weight_scale=[up_weight_scale, up_alpha], + input_zp=[], + weight_zp=[], + ) + hidden = F.silu(gate_out) * up_out + return torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + hidden, + down_weight, + None, + input_scale=[down_input_scale], + weight_scale=[down_weight_scale, down_alpha], + input_zp=[], + weight_zp=[], + ) + + +@torch_nvfp4_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_input_scale: torch.Tensor, + gate_weight_scale: torch.Tensor, + gate_alpha: torch.Tensor, + up_input_scale: torch.Tensor, + up_weight_scale: torch.Tensor, + up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape: [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) + + +@torch.library.custom_op("auto_deploy::fused_nvfp4_swiglu_mlp", mutates_args=()) +def fused_nvfp4_swiglu_mlp( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_input_scale: torch.Tensor, + gate_up_weight_scale: torch.Tensor, + gate_up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """Fused NVFP4 SwiGLU MLP with concatenated gate+up weights. + + Performs a single NVFP4 matmul for gate and up projections, then splits, + applies SwiGLU activation, and does the down NVFP4 matmul. + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_up_weight: Concatenated FP4 packed gate+up weight + [2*intermediate_size, hidden_size/2] uint8. + down_weight: FP4 packed down weight [hidden_size, intermediate_size/2] uint8. + gate_up_input_scale: Shared input scale for gate+up projection. + gate_up_weight_scale: Concatenated per-block weight scale for gate+up. + gate_up_alpha: Shared alpha for gate+up projection. + down_input_scale: Input scale for down projection. + down_weight_scale: Per-block weight scale for down projection. + down_alpha: Alpha for down projection. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + # Single NVFP4 linear for both gate and up projections + gate_up_out = torch.ops.auto_deploy.torch_quant_nvfp4_linear( + input, + gate_up_weight, + bias=None, + input_scale=gate_up_input_scale, + weight_scale=gate_up_weight_scale, + alpha=gate_up_alpha, + ) + + # Apply SwiGLU activation: split, silu(gate) * up (uses FlashInfer when available) + hidden = _silu_and_mul(gate_up_out) + + # Down projection + return torch.ops.auto_deploy.torch_quant_nvfp4_linear( + hidden, + down_weight, + bias=None, + input_scale=down_input_scale, + weight_scale=down_weight_scale, + alpha=down_alpha, + ) + + +@fused_nvfp4_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_input_scale: torch.Tensor, + gate_up_weight_scale: torch.Tensor, + gate_up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape: [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.py index ebaefbf963c9..10942196a44e 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.py @@ -119,6 +119,10 @@ def _cuda_cached_causal_conv1d( pad_slot_id=PAD_SLOT_ID, ) + # Zero padding positions beyond valid tokens (for piecewise CUDA graph) + if num_total_tokens < bs: + inp_flat[num_total_tokens:].zero_() + @_cuda_cached_causal_conv1d.register_fake def _cuda_cached_causal_conv1d_fake( diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 15d46a329d3f..1fbb22836683 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -140,11 +140,13 @@ def _flashinfer_cached_ssm( ) preallocated_ssm_out[num_prefill_tokens:num_total_tokens].copy_(y_decode) if num_total_tokens > 0: - return ( - preallocated_ssm_out[:num_total_tokens] - .view(b, s, num_heads, head_dim) - .to(hidden_states.dtype) - ) + # Cast to input dtype if needed (prefill may compute in higher precision) + if preallocated_ssm_out.dtype != hidden_states.dtype: + preallocated_ssm_out = preallocated_ssm_out.to(hidden_states.dtype) + # Zero padding positions so downstream ops don't see garbage (piecewise CG) + if num_total_tokens < bs: + preallocated_ssm_out[num_total_tokens:].zero_() + return preallocated_ssm_out.view(b, s, num_heads, head_dim) else: return torch.empty_like(hidden_states) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py index 5da162f0f1df..fe9d832e387b 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py @@ -142,6 +142,8 @@ def _run_ssm_prefill( C_prefill = C_flat[:num_prefill_tokens].unsqueeze(0) # [1, S_p, G, N] dt_prefill = dt_flat[:num_prefill_tokens].unsqueeze(0) # [1, S_p, H] + seq_idx_prefill = seq_idx_prefill[:, :num_prefill_tokens] + initial_states = None if torch.any(use_initial_states[:num_prefill]): initial_states = torch.where( diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py index 993d061248e2..28e42f102367 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py @@ -130,6 +130,10 @@ def _triton_cached_causal_conv1d( ) inp_flat[num_prefill_tokens:num_total_tokens] = y_decode + # Zero padding positions beyond valid tokens (for piecewise CUDA graph) + if num_total_tokens < bs: + inp_flat[num_total_tokens:].zero_() + @_triton_cached_causal_conv1d.register_fake def _triton_cached_causal_conv1d_fake( diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py index 35937d50cfdf..f8453a248f06 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py @@ -137,11 +137,13 @@ def _triton_cached_ssm( ) if num_total_tokens > 0: - return ( - preallocated_ssm_out[:num_total_tokens] - .view(b, s, num_heads, head_dim) - .to(hidden_states.dtype) - ) + # Cast to input dtype if needed (prefill may compute in higher precision) + if preallocated_ssm_out.dtype != hidden_states.dtype: + preallocated_ssm_out = preallocated_ssm_out.to(hidden_states.dtype) + # Zero padding positions so downstream ops don't see garbage (piecewise CG) + if num_total_tokens < bs: + preallocated_ssm_out[num_total_tokens:].zero_() + return preallocated_ssm_out.view(b, s, num_heads, head_dim) else: return torch.empty_like(hidden_states) diff --git a/tensorrt_llm/_torch/auto_deploy/export/export.py b/tensorrt_llm/_torch/auto_deploy/export/export.py index b76a72bc393b..79313a28bb93 100644 --- a/tensorrt_llm/_torch/auto_deploy/export/export.py +++ b/tensorrt_llm/_torch/auto_deploy/export/export.py @@ -171,7 +171,7 @@ def _find_moe_module_lists( def _reduce_moe_experts( model: nn.Module, - min_num_experts: int, + num_moe_experts_for_export: int, args: Optional[Tuple[Any, ...]] = None, kwargs: Optional[Dict[str, Any]] = None, ) -> List[Dict[str, Any]]: @@ -179,19 +179,21 @@ def _reduce_moe_experts( Uses a probe forward pass to identify which ``nn.ModuleList`` instances feed into ``torch_moe``-family custom ops (see :func:`_find_moe_module_lists`), - then truncates each to *min_num_experts* entries. The returned list of dicts - carries the metadata needed by :func:`_restore_moe_experts` and - :func:`_expand_moe_experts_in_graph`. + then truncates each to *num_moe_experts_for_export* entries. The returned + list of dicts carries the metadata needed by :func:`_restore_moe_experts` + and :func:`_expand_moe_experts_in_graph`. """ - if min_num_experts < 1: - raise ValueError(f"min_num_experts must be >= 1, got {min_num_experts}") + if num_moe_experts_for_export < 1: + raise ValueError( + f"num_moe_experts_for_export must be >= 1, got {num_moe_experts_for_export}" + ) moe_lists = _find_moe_module_lists(model, args, kwargs) reductions: List[Dict[str, Any]] = [] for path, (parent, attr_name, mod_list) in moe_lists.items(): orig_count = len(mod_list) - if orig_count <= min_num_experts: + if orig_count <= num_moe_experts_for_export: continue reductions.append( @@ -203,10 +205,10 @@ def _reduce_moe_experts( "expert_prefix": path, } ) - setattr(parent, attr_name, nn.ModuleList(list(mod_list[:min_num_experts]))) + setattr(parent, attr_name, nn.ModuleList(list(mod_list[:num_moe_experts_for_export]))) ad_logger.info( f"Reduced MOE experts in '{path}' from {orig_count} to " - f"{min_num_experts} for faster export" + f"{num_moe_experts_for_export} for faster export" ) return reductions @@ -252,7 +254,12 @@ def _expand_moe_experts_in_graph( if not reductions: return - # MOE ops whose arguments include per-expert weight lists (from index 3 onward) + # MOE ops whose arguments include per-expert weight lists. + # All these ops share the same first 3 positional args (x, selected_experts, + # routing_weights) which are plain Tensors, followed by one or more + # List[Tensor] args that hold per-expert weights/scales. We use the op + # schema to discover which arguments are Tensor[] rather than hard-coding + # the starting index. moe_ops = { torch.ops.auto_deploy.torch_moe, torch.ops.auto_deploy.torch_quant_fp8_moe, @@ -266,11 +273,18 @@ def _expand_moe_experts_in_graph( if not is_op(node, moe_ops): continue - # Collect indices of list-of-node arguments (expert weight/scale lists) + # Collect indices of List[Tensor] arguments from the op schema – these + # are the per-expert weight / scale lists. + op = node.target + schema = op._schema if hasattr(op, "_schema") else next(iter(op._schemas.values())) + _tensor_list_types = ("Tensor[]", "List[Tensor]") list_arg_indices = [ i - for i in range(3, len(node.args)) - if isinstance(node.args[i], (list, tuple)) and len(node.args[i]) > 0 + for i, arg_meta in enumerate(schema.arguments) + if any(t in str(arg_meta.type) for t in _tensor_list_types) + and i < len(node.args) + and isinstance(node.args[i], (list, tuple)) + and len(node.args[i]) > 0 ] if not list_arg_indices: continue @@ -673,6 +687,7 @@ def _capture_fn(model, args, kwargs): return egm # Optionally reduce MOE experts for faster export tracing + # TODO (https://github.com/NVIDIA/TensorRT-LLM/issues/7547): Reuse the export patch system moe_reductions: List[Dict[str, Any]] = [] if num_moe_experts_for_export is not None: moe_reductions = _reduce_moe_experts(model, num_moe_experts_for_export, args, kwargs) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py index 17aec049406a..05145efa63ae 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py @@ -444,7 +444,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.shared_experts is not None: final_hidden_states = final_hidden_states + self.shared_experts(identity) - return final_hidden_states.to(hidden_states.dtype) + return final_hidden_states class Glm4MoeLiteAttention(nn.Module): diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index a7e1d8525fa3..ef0eb2dbca80 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -746,7 +746,7 @@ def _build_input_ids(request) -> Tuple[List[int], List[int], bool]: if use_overlap: mask_scatter_indices.extend(list(range(cu_seqlen[-2], cu_seqlen[-1]))) - # get cache indices + # get cache indices and truncate the number of blocks according to total tokens cache_indices = kv_cache_manager.get_cache_indices(request) cache_loc.extend(cache_indices) pages_per_seq.append(len(cache_indices)) @@ -764,12 +764,24 @@ def _build_input_ids(request) -> Tuple[List[int], List[int], bool]: gather_required = len(context_requests) > 0 and not gather_context_logits logits_gather_info = [len(logits_gather_indices), int(gather_required)] - # update the sequence info object now + # Compute batch_info explicitly based on actual request ordering rather than + # relying on the seq_len > 1 heuristic in nest_sequences. With chunked prefill, + # a context request may have context_chunk_size=1, giving seq_len=1. The heuristic + # would misclassify it as decode, causing host_request_types to be inconsistent + # with the actual request ordering (context + extend first, then generation). + # This mismatch leads to incorrect token splitting in thop.attention. + num_prefill_seqs = num_ctx_requests + len(extend_requests) + num_prefill_tokens = sum(seq_len[:num_prefill_seqs]) + num_decode_seqs = len(generation_requests) + batch_info = [num_prefill_seqs, num_prefill_tokens, num_decode_seqs] + + # update the sequence info object now (also triggers rescatter + host_prepare internally) self.cache_seq_interface.info.nest_sequences( input_ids, position_ids=position_ids, seq_len=seq_len, input_pos=input_pos, + batch_info=batch_info, cu_seqlen=cu_seqlen, cache_loc=cache_loc, pages_per_seq=pages_per_seq, @@ -782,12 +794,27 @@ def _build_input_ids(request) -> Tuple[List[int], List[int], bool]: logits_gather_info=logits_gather_info, _gather_idx=None if new_tokens is None else flat_gather_indices, _mask_scatter_indices=None if new_tokens is None else mask_scatter_indices, + _ungathered_input_ids=new_tokens.flatten() if new_tokens is not None else None, **extra_args, ) # scatter the new tokens into the input_ids tensor if provided if new_tokens is not None: self.cache_seq_interface.info.rescatter_input_ids(new_tokens.flatten()) + # Set padded_num_tokens for piecewise CG bucket alignment (prefill/mixed only). + # This makes _shape_for_forward return tensors sized to the bucket, while + # batch_info_host and other metadata stay unchanged for correct dynamic op behavior. + seq_info = self.cache_seq_interface.info + if not seq_info.is_generate and seq_info.piecewise_bucket_sizes: + total_tokens = seq_info.total_num_tokens + bucket = seq_info.find_nearest_piecewise_bucket(total_tokens) + if bucket is not None and bucket > total_tokens: + seq_info.padded_num_tokens = bucket + # Clear padded tails so bucketed piecewise replay never sees stale + # values (e.g. overlap-scheduler dummy tokens like -1) in embedding. + seq_info.named_args["input_ids"][:, total_tokens:bucket].fill_(0) + seq_info.named_args["position_ids"][:, total_tokens:bucket].fill_(0) + self.cache_seq_interface.info.run_host_prepare_for_attention_forward() if spec_resource_manager is not None and isinstance( @@ -806,6 +833,16 @@ def _build_input_ids(request) -> Tuple[List[int], List[int], bool]: def _compute_logits(self) -> List[torch.Tensor]: # run the model logits: torch.Tensor = self.model(**self.cache_seq_interface.named_args)[0] + + # Reset piecewise padding after forward and truncate logits to real token count. + seq_info = self.cache_seq_interface.info + padded = seq_info.padded_num_tokens + if padded is not None: + real_tokens = seq_info.total_num_tokens + # Truncate logits from [1, padded, vocab] to [1, real, vocab] + logits = logits[:, :real_tokens, :] + seq_info.padded_num_tokens = None + logits = self.cache_seq_interface.info.maybe_gather_and_squeeze_logits(logits) # TRTLLMSampler expects float32 logits. PyTorchModelEngine always casts to float32 regardless. diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 4cb3f7c02af2..9db9273ff6b4 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -453,17 +453,29 @@ def _create_and_assign_state_views( return manager, num_managed_mamba_layers - def _assign_kv_cache_views(self, kv_managed: Dict[str, KVPagedResourceHandler]) -> None: + def _assign_kv_cache_views(self, kv_managed: Dict[str, KVPagedResourceHandler]) -> int: """Retrieve and assign buffer views for managed KV paged resources. Args: kv_managed: Dict of KV resources managed by the cache manager. + + Returns: + block_offset_multiplier derived from the first KV cache view's strides. """ + block_offset_multiplier = 0 for idx, (name, h) in enumerate(kv_managed.items()): view = self._kv_cache_manager.get_buffers(idx, kv_layout=h.kv_layout) assert view[0].is_contiguous(), f"Non-contiguous kv cache resource for {name}" self._caches[name] = view + # Compute block_offset_multiplier from the first layer's kv_cache strides. + # This is stride(0)/stride(1) which equals kv_factor for per-layer views + # or num_layers*kv_factor for interleaved pools. + if idx == 0: + block_offset_multiplier = view.stride(0) // view.stride(1) + + return block_offset_multiplier + def _allocate_unmanaged_resources(self) -> None: """Allocate resources not managed by cache managers. @@ -520,27 +532,32 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: # No typed state resources - use pure KVCacheManager self._kv_cache_manager = KVCacheManager(**kv_cache_kwargs) - # 4. Store tuned config and ensure capacity + # 4. Store tuned config self._kv_cache_config_tuned = kv_cache_config - self.info.estimate_cache_loc_capacity(self._kv_cache_manager.blocks_in_primary_pool) - # 5. Assign KV views - self._assign_kv_cache_views(kv_managed) + # 5. Assign KV views (compute block_offset_multiplier from first view's strides) + block_offset_multiplier = self._assign_kv_cache_views(kv_managed) + + # 6. Update cache information (resize cache_loc, set max_seq_info with all max sizes) + self.info.update_cache_information( + num_blocks=self._kv_cache_manager.blocks_in_primary_pool, + block_offset_multiplier=block_offset_multiplier, + ) - # 6. Allocate remaining unmanaged resources + # 7. Allocate remaining unmanaged resources self._allocate_unmanaged_resources() - # 7. Patch shutdown + # 8. Patch shutdown self._kv_cache_manager.shutdown = with_pre_callback( self._kv_cache_manager.shutdown, self._clear_caches, ) - # 8. Compute final token count and cache statistics + # 9. Compute final token count and cache statistics max_resource_count = self._kv_cache_manager.get_max_resource_count() max_tokens_final = max_resource_count * self._kv_cache_manager.tokens_per_block - # 9. Collect statistics of different types of resources + # 10. Collect statistics of different types of resources num_state_total = sum( 1 for h in self._resource_lookup.values() if isinstance(h, StateResourceHandler) ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py b/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py index 376abc8902b6..bcf7a24325e1 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py @@ -6,6 +6,7 @@ from ...compile import ArgsKwargs, CompileBackendRegistry from ...models.factory import ModelFactory from ...shim.interface import CachedSequenceInterface +from ...utils.logger import ad_logger from ..interface import ( BaseTransform, SharedConfig, @@ -15,6 +16,30 @@ ) +def _generate_default_piecewise_num_tokens(max_num_tokens: int) -> List[int]: + """Generate default piecewise bucket sizes when none are specified. + + Uses powers-of-2 from 64 up to max_num_tokens. This provides ~log2(max/64) + bucket sizes with at most 2x padding overhead per bucket. + + For example, max_num_tokens=8192 → [64, 128, 256, 512, 1024, 2048, 4096, 8192] + """ + if max_num_tokens <= 0: + return [] + + buckets = [] + nt = 64 + while nt <= max_num_tokens: + buckets.append(nt) + nt *= 2 + + # Always include max_num_tokens as the largest bucket + if not buckets or buckets[-1] != max_num_tokens: + buckets.append(max_num_tokens) + + return sorted(buckets) + + class CompileModelConfig(TransformConfig): """Configuration for the compile model transform.""" @@ -27,6 +52,18 @@ class CompileModelConfig(TransformConfig): backend: Literal["torch-simple", "torch-compile", "torch-cudagraph", "torch-opt"] = Field( description="The backend to use for compiling the model." ) + piecewise_enabled: bool = Field( + default=False, + description="Enable piecewise CUDA graph for prefill/mixed batches (dual-mode).", + ) + piecewise_num_tokens: Optional[List[int]] = Field( + default=None, + description=( + "Total token counts to pre-capture piecewise CUDA graphs for. " + "If null and piecewise_enabled=true, auto-generates power-of-2 buckets " + "up to max_num_tokens (e.g. [64, 128, 256, ..., max_num_tokens])." + ), + ) @TransformRegistry.register("compile_model") @@ -52,10 +89,69 @@ def _get_args_kwargs(bs: int) -> ArgsKwargs: cm.info.set_generate_only_batch(bs) return (), cm.named_args + def _get_mixed_args_kwargs(num_tokens: int) -> ArgsKwargs: + """Generate synthetic mixed-batch args for piecewise CG capture. + + Always creates 1 prefill sequence + 1 decode sequence to exercise + both code paths in dynamic ops (attention, SSM). The static CUDA + graph segments are agnostic to the prefill/decode split -- they only + see total_num_tokens. + """ + assert num_tokens >= 3, ( + f"Piecewise bucket {num_tokens} too small for mixed batch. " + f"Minimum is 3 (1 prefill seq with len>=2 + 1 decode seq)." + ) + cm.info.set_example_sequence( + input_ids=[ + [1] * (num_tokens - 1), # prefill: len > 1 + [1], # decode: len == 1 + ], + ) + return (), cm.named_args + + extra_kwargs = {} + config_overrides = {} + + if self.config.piecewise_enabled: + extra_kwargs["get_mixed_args_kwargs_for_compile"] = _get_mixed_args_kwargs + + # Auto-generate piecewise_num_tokens if not explicitly specified + if self.config.piecewise_num_tokens is None: + max_num_tokens = cm.info.max_num_tokens + auto_buckets = _generate_default_piecewise_num_tokens(max_num_tokens) + config_overrides["piecewise_num_tokens"] = auto_buckets + ad_logger.info( + f"Auto-generated piecewise_num_tokens from max_num_tokens={max_num_tokens}: " + f"{auto_buckets}" + ) + else: + # Filter out buckets < 3 (mixed batch needs at least 3 tokens) + valid_buckets = [nt for nt in self.config.piecewise_num_tokens if nt >= 3] + dropped = [nt for nt in self.config.piecewise_num_tokens if nt < 3] + if dropped: + ad_logger.warning( + f"Dropping piecewise_num_tokens {dropped} (too small for mixed batch, " + f"minimum is 3). Remaining: {valid_buckets}" + ) + config_overrides["piecewise_num_tokens"] = valid_buckets + + # Only populate when piecewise is enabled so that the runtime padding + # guard (padded_num_tokens) stays inactive when piecewise_enabled=false. + if self.config.piecewise_enabled: + final_buckets = config_overrides.get( + "piecewise_num_tokens", self.config.piecewise_num_tokens or [] + ) + cm.info.piecewise_bucket_sizes = final_buckets + + # Merge config with any overrides + config_dict = self.config.model_dump() + config_dict.update(config_overrides) + compiler_backend = CompileBackendRegistry.get(self.config.backend)( mod, get_args_kwargs_for_compile=_get_args_kwargs, - **self.config.model_dump(), + **extra_kwargs, + **config_dict, ) mod_compiled = compiler_backend.compile() diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py new file mode 100644 index 000000000000..268cdad6f58c --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py @@ -0,0 +1,618 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Graph transforms for SwiGLU MLP fusion. + +This module provides two-stage transformation for SwiGLU MLP: +1. MatchSwiGLUPattern: Detects SwiGLU patterns and replaces with torch_swiglu_mlp +2. FuseSwiGLU: Fuses gate+up weights into a single concatenated matmul + +The SwiGLU pattern is: silu(x @ gate.T) * (x @ up.T) @ down.T +""" + +from typing import Tuple, Type + +import torch +from pydantic import Field +from torch.fx import GraphModule, Node + +# Import the custom ops to ensure they are registered and for use in replacements +from ...custom_ops.linear.swiglu import torch_swiglu_mlp +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils._graph import ( + del_attr_by_name, + delete_all_unused_submodules, + eliminate_dead_code, + get_attr_by_name, +) +from ...utils.node_utils import is_op +from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + + +def _try_free_attr_node(gm: GraphModule, graph, attr_node: Node) -> None: + """Erase a get_attr node and eagerly delete its module attribute if it has no users. + + This is used to free unfused weight tensors as soon as they are no longer + referenced in the graph, avoiding a temporary memory spike that would occur + if cleanup were deferred until after all nodes are processed. + """ + if attr_node is not None and attr_node.op == "get_attr" and len(attr_node.users) == 0: + target = attr_node.target + graph.erase_node(attr_node) + del_attr_by_name(gm, target) + + +def _swiglu_pattern_no_bias(x, gate_weight, up_weight, down_weight): + """Pattern for SwiGLU MLP without biases. + + Matches: silu(linear(x, gate_weight, None)) * linear(x, up_weight, None) -> linear(down_weight, None) + """ + gate_out = torch.ops.auto_deploy.torch_linear_simple.default(x, gate_weight, None) + up_out = torch.ops.auto_deploy.torch_linear_simple.default(x, up_weight, None) + silu_out = torch.ops.aten.silu.default(gate_out) + mul_out = torch.ops.aten.mul.Tensor(silu_out, up_out) + down_out = torch.ops.auto_deploy.torch_linear_simple.default(mul_out, down_weight, None) + return down_out + + +def _swiglu_replacement_no_bias(x, gate_weight, up_weight, down_weight): + """Replacement for SwiGLU pattern without biases.""" + # Call the Python wrapper directly, not via torch.ops.auto_deploy + # This ensures proper FakeTensor mode handling during tracing + return torch_swiglu_mlp(x, gate_weight, up_weight, down_weight, None, None, None) + + +def _swiglu_pattern_with_bias( + x, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias +): + """Pattern for SwiGLU MLP with biases. + + Matches: silu(linear(x, gate_weight, gate_bias)) * linear(x, up_weight, up_bias) -> linear(down_weight, down_bias) + """ + gate_out = torch.ops.auto_deploy.torch_linear_simple.default(x, gate_weight, gate_bias) + up_out = torch.ops.auto_deploy.torch_linear_simple.default(x, up_weight, up_bias) + silu_out = torch.ops.aten.silu.default(gate_out) + mul_out = torch.ops.aten.mul.Tensor(silu_out, up_out) + down_out = torch.ops.auto_deploy.torch_linear_simple.default(mul_out, down_weight, down_bias) + return down_out + + +def _swiglu_replacement_with_bias( + x, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias +): + """Replacement for SwiGLU pattern with biases.""" + # Call the Python wrapper directly, not via torch.ops.auto_deploy + # This ensures proper FakeTensor mode handling during tracing + return torch_swiglu_mlp(x, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias) + + +@TransformRegistry.register("match_swiglu_pattern") +class MatchSwiGLUPattern(BaseTransform): + """Matches SwiGLU MLP patterns and replaces with torch_swiglu_mlp op. + + This transform runs in the pattern_matcher stage and detects the following pattern: + silu(x @ gate.T) * (x @ up.T) @ down.T + + And replaces it with a single torch_swiglu_mlp op that can be fused later. + + Uses ADPatternMatcherPass for declarative pattern matching. + """ + + config: TransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return TransformConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + patterns = ADPatternMatcherPass() + + # Dummy shapes for tracing - shapes don't matter for matching + hidden, intermediate = 128, 256 + + # Pattern 1: SwiGLU without biases (most common case) + dummy_args_no_bias = [ + torch.randn(2, hidden, device="meta", dtype=torch.float16), # x + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # gate_weight + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # up_weight + torch.randn(hidden, intermediate, device="meta", dtype=torch.float16), # down_weight + ] + register_ad_pattern( + search_fn=_swiglu_pattern_no_bias, + replace_fn=_swiglu_replacement_no_bias, + patterns=patterns, + dummy_args=dummy_args_no_bias, + ) + + # Pattern 2: SwiGLU with biases + dummy_args_with_bias = [ + torch.randn(2, hidden, device="meta", dtype=torch.float16), # x + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # gate_weight + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # up_weight + torch.randn(hidden, intermediate, device="meta", dtype=torch.float16), # down_weight + torch.randn(intermediate, device="meta", dtype=torch.float16), # gate_bias + torch.randn(intermediate, device="meta", dtype=torch.float16), # up_bias + torch.randn(hidden, device="meta", dtype=torch.float16), # down_bias + ] + register_ad_pattern( + search_fn=_swiglu_pattern_with_bias, + replace_fn=_swiglu_replacement_with_bias, + patterns=patterns, + dummy_args=dummy_args_with_bias, + ) + + num_matches = patterns.apply(gm.graph) + + if num_matches > 0: + gm.recompile() + + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + + return gm, info + + +class FuseSwiGLUConfig(TransformConfig): + """Configuration for the SwiGLU fusion transform.""" + + enabled: bool = Field( + default=True, + description="Whether to enable SwiGLU fusion.", + ) + + +@TransformRegistry.register("fuse_swiglu") +class FuseSwiGLU(BaseTransform): + """Fuses torch_swiglu_mlp ops by concatenating gate and up weights. + + This transform runs in the post_load_fusion stage and replaces torch_swiglu_mlp ops + with fused_swiglu_mlp ops that use a single concatenated gate+up weight matrix. + + This reduces memory bandwidth by performing a single matmul instead of two + separate matmuls for gate and up projections. + """ + + config: FuseSwiGLUConfig + + @classmethod + def get_config_class(cls) -> Type[FuseSwiGLUConfig]: + return FuseSwiGLUConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + if not self.config.enabled: + return gm, TransformInfo(skipped=True, num_matches=0) + + graph = gm.graph + cnt = 0 + fused_weight_idx = 0 + + for node in list(graph.nodes): + if not is_op(node, torch.ops.auto_deploy.torch_swiglu_mlp.default): + continue + + # Extract args: (input, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias) + input_node = node.args[0] + gate_weight_node = node.args[1] + up_weight_node = node.args[2] + down_weight_node = node.args[3] + gate_bias_node = node.args[4] if len(node.args) > 4 else None + up_bias_node = node.args[5] if len(node.args) > 5 else None + down_bias_node = node.args[6] if len(node.args) > 6 else None + + # Get the actual weight tensors + gate_weight = get_attr_by_name(gm, gate_weight_node.target) + up_weight = get_attr_by_name(gm, up_weight_node.target) + + # Concatenate gate and up weights: [intermediate, hidden] -> [2*intermediate, hidden] + gate_up_weight = torch.cat([gate_weight, up_weight], dim=0) + + # Create new attribute for the fused weight + fused_weight_name = f"fused_swiglu_gate_up_{fused_weight_idx}" + gm.register_buffer(fused_weight_name, gate_up_weight) + + # Handle biases + gate_up_bias_node = None + fused_bias_name = None + + if gate_bias_node is not None and gate_bias_node.op == "get_attr": + gate_bias = get_attr_by_name(gm, gate_bias_node.target) + up_bias = get_attr_by_name(gm, up_bias_node.target) if up_bias_node else None + + if up_bias is not None: + gate_up_bias = torch.cat([gate_bias, up_bias], dim=0) + fused_bias_name = f"fused_swiglu_gate_up_bias_{fused_weight_idx}" + gm.register_buffer(fused_bias_name, gate_up_bias) + + # Create get_attr node for the fused weight + with graph.inserting_before(node): + fused_weight_node = graph.get_attr(fused_weight_name) + + if fused_bias_name is not None: + gate_up_bias_node = graph.get_attr(fused_bias_name) + + # Create the fused_swiglu_mlp node + with graph.inserting_after(node): + fused_node: Node = graph.call_function( + torch.ops.auto_deploy.fused_swiglu_mlp.default, + args=( + input_node, + fused_weight_node, + down_weight_node, + gate_up_bias_node, + down_bias_node, + ), + ) + + # Replace uses and erase old node + node.replace_all_uses_with(fused_node) + graph.erase_node(node) + + # Eagerly free unfused weight/bias tensors that are no longer referenced + # to avoid a temporary memory spike from holding both fused and unfused + # copies simultaneously across all layers. + _try_free_attr_node(gm, graph, gate_weight_node) + _try_free_attr_node(gm, graph, up_weight_node) + _try_free_attr_node(gm, graph, gate_bias_node) + _try_free_attr_node(gm, graph, up_bias_node) + + fused_weight_idx += 1 + cnt += 1 + + if cnt > 0: + gm.recompile() + + # Clean up any remaining dead code and unused submodules + eliminate_dead_code(gm) + delete_all_unused_submodules(gm) + + info = TransformInfo( + skipped=False, num_matches=cnt, is_clean=cnt == 0, has_valid_shapes=cnt == 0 + ) + + return gm, info + + +# ── NVFP4 quantized SwiGLU pattern matching and fusion ────────────────────── + +from ...custom_ops.linear.swiglu import torch_nvfp4_swiglu_mlp # noqa: E402 + + +def _nvfp4_swiglu_pattern_no_bias( + x, + gate_weight, + gate_input_scale, + gate_weight_scale, + gate_alpha, + up_weight, + up_input_scale, + up_weight_scale, + up_alpha, + down_weight, + down_input_scale, + down_weight_scale, + down_alpha, +): + """Pattern for NVFP4 quantized SwiGLU MLP without biases. + + Matches: silu(nvfp4_linear(x, gate)) * nvfp4_linear(x, up) -> nvfp4_linear(down) + """ + gate_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default( + x, + gate_weight, + None, + input_scale=[gate_input_scale], + weight_scale=[gate_weight_scale, gate_alpha], + input_zp=[], + weight_zp=[], + ) + up_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default( + x, + up_weight, + None, + input_scale=[up_input_scale], + weight_scale=[up_weight_scale, up_alpha], + input_zp=[], + weight_zp=[], + ) + silu_out = torch.ops.aten.silu.default(gate_out) + mul_out = torch.ops.aten.mul.Tensor(silu_out, up_out) + down_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default( + mul_out, + down_weight, + None, + input_scale=[down_input_scale], + weight_scale=[down_weight_scale, down_alpha], + input_zp=[], + weight_zp=[], + ) + return down_out + + +def _nvfp4_swiglu_replacement_no_bias( + x, + gate_weight, + gate_input_scale, + gate_weight_scale, + gate_alpha, + up_weight, + up_input_scale, + up_weight_scale, + up_alpha, + down_weight, + down_input_scale, + down_weight_scale, + down_alpha, +): + """Replacement for NVFP4 quantized SwiGLU pattern without biases.""" + return torch_nvfp4_swiglu_mlp( + x, + gate_weight, + up_weight, + down_weight, + gate_input_scale, + gate_weight_scale, + gate_alpha, + up_input_scale, + up_weight_scale, + up_alpha, + down_input_scale, + down_weight_scale, + down_alpha, + ) + + +@TransformRegistry.register("match_nvfp4_swiglu_pattern") +class MatchNVFP4SwiGLUPattern(BaseTransform): + """Matches NVFP4 quantized SwiGLU MLP patterns and replaces with torch_nvfp4_swiglu_mlp. + + This transform runs in the pattern_matcher stage AFTER quantize_nvfp4_linear_from_config + has converted torch_linear_simple ops to torch_fake_quant_nvfp4_linear ops. + + It detects the following NVFP4 pattern: + silu(nvfp4_linear(x, gate)) * nvfp4_linear(x, up) -> nvfp4_linear(down) + + And replaces it with a single torch_nvfp4_swiglu_mlp op that can be fused later. + """ + + config: TransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return TransformConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + patterns = ADPatternMatcherPass() + + # FP4 shape params for dummy args (shapes don't matter for matching) + N = 32 # intermediate_size + K_packed = 32 # hidden_size / 2 (FP4 packing) + K_eff = 2 * K_packed # actual hidden_size + N_down = K_eff # hidden_size (output of down proj) + K_down_packed = N // 2 # intermediate_size / 2 (down proj input) + + # Weight scale sizes (per-block scale: N * K / 16) + gate_cutlass_len = N * (K_eff // 16) + down_cutlass_len = N_down * (N // 16) + + x = torch.randn(2, K_eff, device="meta", dtype=torch.float16) + + # Gate args + gate_w = torch.randint(0, 255, (N, K_packed), device="meta", dtype=torch.uint8) + gate_is = torch.tensor(0.01, device="meta", dtype=torch.float32) + gate_ws = torch.randint(0, 255, (gate_cutlass_len,), device="meta", dtype=torch.uint8) + gate_a = torch.tensor(1.2345, device="meta", dtype=torch.float32) + + # Up args (same shapes as gate) + up_w = torch.randint(0, 255, (N, K_packed), device="meta", dtype=torch.uint8) + up_is = torch.tensor(0.02, device="meta", dtype=torch.float32) + up_ws = torch.randint(0, 255, (gate_cutlass_len,), device="meta", dtype=torch.uint8) + up_a = torch.tensor(2.3456, device="meta", dtype=torch.float32) + + # Down args + down_w = torch.randint(0, 255, (N_down, K_down_packed), device="meta", dtype=torch.uint8) + down_is = torch.tensor(0.03, device="meta", dtype=torch.float32) + down_ws = torch.randint(0, 255, (down_cutlass_len,), device="meta", dtype=torch.uint8) + down_a = torch.tensor(3.4567, device="meta", dtype=torch.float32) + + dummy_args = [ + x, + gate_w, + gate_is, + gate_ws, + gate_a, + up_w, + up_is, + up_ws, + up_a, + down_w, + down_is, + down_ws, + down_a, + ] + + register_ad_pattern( + search_fn=_nvfp4_swiglu_pattern_no_bias, + replace_fn=_nvfp4_swiglu_replacement_no_bias, + patterns=patterns, + dummy_args=dummy_args, + ) + + num_matches = patterns.apply(gm.graph) + + if num_matches > 0: + gm.recompile() + + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + + return gm, info + + +@TransformRegistry.register("fuse_nvfp4_swiglu") +class FuseNVFP4SwiGLU(BaseTransform): + """Fuses torch_nvfp4_swiglu_mlp ops by concatenating gate and up FP4 weights. + + This transform runs in the post_load_fusion stage and replaces torch_nvfp4_swiglu_mlp + ops with fused_nvfp4_swiglu_mlp ops that use a single concatenated gate+up weight matrix. + + FP4 weight fusion: + - gate+up packed weights are concatenated along dim=0 + - gate+up per-block weight scales are concatenated along dim=0 + - gate+up input_scale and alpha must match (shared input) + """ + + config: TransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return TransformConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + graph = gm.graph + cnt = 0 + fused_weight_idx = 0 + + for node in list(graph.nodes): + if not is_op(node, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default): + continue + + # Extract args: + # (input, gate_weight, up_weight, down_weight, + # gate_input_scale, gate_weight_scale, gate_alpha, + # up_input_scale, up_weight_scale, up_alpha, + # down_input_scale, down_weight_scale, down_alpha) + input_node = node.args[0] + gate_weight_node = node.args[1] + up_weight_node = node.args[2] + down_weight_node = node.args[3] + gate_input_scale_node = node.args[4] + gate_weight_scale_node = node.args[5] + gate_alpha_node = node.args[6] + up_input_scale_node = node.args[7] + up_weight_scale_node = node.args[8] + up_alpha_node = node.args[9] + down_input_scale_node = node.args[10] + down_weight_scale_node = node.args[11] + down_alpha_node = node.args[12] + + # Get the actual weight tensors + gate_weight = get_attr_by_name(gm, gate_weight_node.target) + up_weight = get_attr_by_name(gm, up_weight_node.target) + + # Concatenate gate and up FP4 packed weights along dim=0 + gate_up_weight = torch.cat([gate_weight, up_weight], dim=0) + + # Get and concatenate weight scales + gate_weight_scale = get_attr_by_name(gm, gate_weight_scale_node.target) + up_weight_scale = get_attr_by_name(gm, up_weight_scale_node.target) + gate_up_weight_scale = torch.cat([gate_weight_scale, up_weight_scale], dim=0) + + # Register fused buffers + prefix = f"fused_nvfp4_swiglu_{fused_weight_idx}" + gm.register_buffer(f"{prefix}_gate_up_weight", gate_up_weight) + gm.register_buffer(f"{prefix}_gate_up_weight_scale", gate_up_weight_scale) + + # Create get_attr nodes for fused weights/scales + with graph.inserting_before(node): + fused_gate_up_weight_node = graph.get_attr(f"{prefix}_gate_up_weight") + fused_gate_up_weight_scale_node = graph.get_attr(f"{prefix}_gate_up_weight_scale") + + # Create the fused_nvfp4_swiglu_mlp node + # Use gate's input_scale and alpha (same as up's since they share input) + with graph.inserting_after(node): + fused_node: Node = graph.call_function( + torch.ops.auto_deploy.fused_nvfp4_swiglu_mlp.default, + args=( + input_node, + fused_gate_up_weight_node, + down_weight_node, + gate_input_scale_node, # shared input_scale for gate+up + fused_gate_up_weight_scale_node, + gate_alpha_node, # shared alpha for gate+up + down_input_scale_node, + down_weight_scale_node, + down_alpha_node, + ), + ) + + # Replace uses and erase old node + node.replace_all_uses_with(fused_node) + graph.erase_node(node) + + # Eagerly free unfused weight/scale tensors that are no longer referenced + # to avoid a temporary memory spike from holding both fused and unfused + # copies simultaneously across all layers. + _try_free_attr_node(gm, graph, gate_weight_node) + _try_free_attr_node(gm, graph, up_weight_node) + _try_free_attr_node(gm, graph, gate_weight_scale_node) + _try_free_attr_node(gm, graph, up_weight_scale_node) + _try_free_attr_node(gm, graph, up_input_scale_node) + _try_free_attr_node(gm, graph, up_alpha_node) + + fused_weight_idx += 1 + cnt += 1 + + if cnt > 0: + gm.recompile() + + # Clean up any remaining dead code and unused submodules + eliminate_dead_code(gm) + delete_all_unused_submodules(gm) + + info = TransformInfo( + skipped=False, num_matches=cnt, is_clean=cnt == 0, has_valid_shapes=cnt == 0 + ) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py index 103578ca01f0..dee221e742be 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py @@ -9,31 +9,43 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Transformation for fusing Add + Cast + RMSNorm.""" +"""Transformation for fusing Add + (optional Cast) + RMSNorm via direct FX graph manipulation.""" -from typing import Tuple +import operator +from typing import List, Optional, Tuple import torch -from torch.fx import GraphModule +from torch.fx import GraphModule, Node from ...custom_ops.normalization.flashinfer_fused_add_rms_norm import flashinfer_fused_add_rms_norm from ...models.factory import ModelFactory from ...shim.interface import CachedSequenceInterface -from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern +from ...utils._graph import eliminate_dead_code +from ...utils.node_utils import is_op from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry @TransformRegistry.register("fuse_add_rms_norm") class FuseAddRMSNorm(BaseTransform): - """Fuse (add + cast + RMSNorm) into one fused op. + """Fuse (add + optional cast + RMSNorm) into one fused op. - Matches: - x = add(input, residual) - y = x.to(dtype) - z = flashinfer_rms_norm(y, weight, eps) + Uses direct FX graph manipulation instead of the inductor pattern matcher + to correctly handle patterns where intermediate nodes (add, rms_norm) have + multiple users in the graph. - Replaces with: - z, x = flashinfer_fused_add_rms_norm(input, residual, weight, eps) + Pattern 1 (without cast): + %add = aten.add(%x, %residual) + %norm = flashinfer_rms_norm(%add, %weight, eps) + + Pattern 2 (with cast): + %add = aten.add(%x, %residual) + %cast = aten.to.dtype(%add, bfloat16) + %norm = flashinfer_rms_norm(%cast, %weight, eps) + + Both are replaced with: + %fused = flashinfer_fused_add_rms_norm(%x, %residual, %weight, eps) + %norm_out = getitem(%fused, 0) # norm result (replaces %norm) + %add_out = getitem(%fused, 1) # add result (replaces %add) """ def _apply( @@ -43,42 +55,102 @@ def _apply( factory: ModelFactory, shared_config: SharedConfig, ) -> Tuple[GraphModule, TransformInfo]: - patterns = ADPatternMatcherPass() - - # Dummy shapes for tracing - bsz, hidden = 2, 128 - dummy_args = [ - torch.randn(bsz, hidden, device="meta", dtype=torch.bfloat16), # x (bf16) - torch.randn(bsz, hidden, device="meta", dtype=torch.bfloat16), # residual (bf16) - torch.randn(hidden, device="meta", dtype=torch.bfloat16), # weight - 1e-5, # eps - ] - - op_ignore_types = {torch.ops.aten.to.dtype: (torch.dtype,)} - scalar_workaround = {"eps": 1e-5} - - def _fused_add_norm_pattern(x, residual, weight, eps): - added = torch.ops.aten.add.Tensor(x, residual) - cast = torch.ops.aten.to.dtype(added, torch.bfloat16) - # Note: we assume flashinfer_rms_norm is the target - norm = torch.ops.auto_deploy.flashinfer_rms_norm.default(cast, weight, eps) - return norm, added - - def _fused_add_norm_replacement(x, residual, weight, eps): - # Use the python wrapper directly, not via torch.ops.auto_deploy - return flashinfer_fused_add_rms_norm(x, residual, weight, eps) - - # Register pattern - register_ad_pattern( - search_fn=_fused_add_norm_pattern, - replace_fn=_fused_add_norm_replacement, - patterns=patterns, - dummy_args=dummy_args, - op_ignore_types=op_ignore_types, - scalar_workaround=scalar_workaround, - ) - - num_matches = patterns.apply(gm.graph) + graph = gm.graph + num_matches = 0 + + # --- Step 1: collect (add_node, optional cast_node, norm_node) triples --- + matches: List[Tuple[Node, Optional[Node], Node]] = [] + + for node in graph.nodes: + # Match flashinfer_rms_norm (handles both overload packet and .default) + if not is_op(node, torch.ops.auto_deploy.flashinfer_rms_norm): + continue + + input_to_norm = node.args[0] + cast_node: Optional[Node] = None + + # Check for an optional aten.to.dtype cast between add and norm + if isinstance(input_to_norm, Node) and is_op(input_to_norm, torch.ops.aten.to.dtype): + cast_node = input_to_norm + input_to_norm = cast_node.args[0] + + # The (possibly unwrapped) input must be an aten.add.Tensor + if not isinstance(input_to_norm, Node) or not is_op( + input_to_norm, torch.ops.aten.add.Tensor + ): + continue + + add_node = input_to_norm + matches.append((add_node, cast_node, node)) + + # --- Step 2: apply fusions --- + erased: set = set() # track erased node ids to skip stale matches + + for add_node, cast_node, norm_node in matches: + # Safety: skip if a node in this match was already consumed + if id(add_node) in erased or id(norm_node) in erased: + continue + + # Original operands + add_lhs = add_node.args[0] # e.g. previous residual + add_rhs = add_node.args[1] # e.g. attention/MoE output + weight = norm_node.args[1] + eps = norm_node.args[2] + + # Insert the fused call right before the norm node. Using + # inserting_before(norm_node) ensures correct topological order: + # fused_node → norm_out → add_out all appear before norm_node. + with graph.inserting_before(norm_node): + # flashinfer_fused_add_rms_norm(x, residual, weight, eps): + # residual += x → residual becomes add result + # x = rms_norm(residual) → x becomes norm result + # returns (x, residual) = (norm_result, add_result) + fused_node = graph.call_function( + flashinfer_fused_add_rms_norm, + args=(add_rhs, add_lhs, weight, eps), + ) + norm_out = graph.call_function(operator.getitem, args=(fused_node, 0)) + add_out = graph.call_function(operator.getitem, args=(fused_node, 1)) + + # Rewire all consumers of the original norm → norm_out + norm_node.replace_all_uses_with(norm_out) + + # Erase norm first so cast_node (if present) loses its only user + graph.erase_node(norm_node) + erased.add(id(norm_node)) + + # Erase cast_node *before* replacing add's uses, otherwise + # replace_all_uses_with would rewrite cast's input to add_out + # which sits after cast in the graph → topological violation. + if cast_node is not None: + if len(cast_node.users) == 0: + graph.erase_node(cast_node) + erased.add(id(cast_node)) + else: + # Rare: cast has users besides norm. Redirect them to a + # new cast placed after add_out so the ordering is valid. + with graph.inserting_before(list(cast_node.users)[0]): + new_cast = graph.call_function( + cast_node.target, + args=(add_out, *cast_node.args[1:]), + kwargs=cast_node.kwargs, + ) + cast_node.replace_all_uses_with(new_cast) + graph.erase_node(cast_node) + erased.add(id(cast_node)) + + # Rewire all consumers of the original add → add_out + # (includes the residual connection to the *next* layer's add) + add_node.replace_all_uses_with(add_out) + + graph.erase_node(add_node) + erased.add(id(add_node)) + + num_matches += 1 + + # Clean up any remaining dead code + if num_matches > 0: + eliminate_dead_code(gm) info = TransformInfo( skipped=False, diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py index 998f15062fe1..5f6454ad66d3 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py @@ -122,16 +122,12 @@ def _process_metadata_host(self, cm: CachedSequenceInterface): if prep_meta_host_op is None: return - # analyze the args of the host-side prepare metadata function using inspect + # Register the host-side prepare metadata function with SequenceInfo. + # Arg availability is validated by require_copy() inside register_host_prepare. sig = inspect.signature(prep_meta_host_op) - args = sig.parameters.keys() - - # check if all args are available in the cached sequence interface - unavailable_args = args - cm.info.available_args - assert not unavailable_args, f"Missing args in SequenceInfo: {unavailable_args=}" - - # add the host-side prepare metadata function to the graph - cm.info.register_host_prepare_for_attention_forward(prep_meta_host_op, list(args)) + cm.info.register_host_prepare_for_attention_forward( + prep_meta_host_op, list(sig.parameters.keys()) + ) def _process_cache_node(self, gm: GraphModule, cache_name: str) -> Node: """Process the cache nodes by inserting a cached attention replacement op.""" diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py new file mode 100644 index 000000000000..50f572cd0920 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py @@ -0,0 +1,215 @@ +"""Transform for multi-stream execution of Q and KV projection chains in MLA attention. + +In DeepSeek-style MLA (Multi-head Latent Attention), the input layernorm output +forks into two independent projection chains that merge at the RoPE + attention op: + + - **Q chain** (heavier): q_a_proj -> rms_norm -> q_b_proj -> view -> split + - **KV chain** (lighter): kv_a_proj_with_mqa -> split -> rms_norm + view + +The Q chain is ~9x heavier than the KV chain. This transform moves the KV +projection linear onto the auxiliary CUDA stream so it executes concurrently +with the Q chain on the main stream. +""" + +from collections import deque +from typing import Callable, List, Tuple + +import torch +from torch.fx import GraphModule, Node + +from ...models.factory import ModelFactory + +# Reuse CachedSequenceInterface for the _apply signature. +from ...shim.interface import CachedSequenceInterface +from ...utils._graph import create_derived_custom_op +from ...utils.multi_stream_utils import ( + _make_aux_stream_impl, + cuda_stream_manager, + record_event_passthrough, +) +from ...utils.node_utils import is_op +from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry + +# --------------------------------------------------------------------------- +# Supported linear op targets. Extend this list to cover quantised variants. +# --------------------------------------------------------------------------- +_LINEAR_OPS: List[Callable] = [ + torch.ops.auto_deploy.torch_linear_simple, + torch.ops.aten.linear, +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _is_linear(node: Node) -> bool: + """Return ``True`` if *node* is a call to one of the supported linear ops.""" + return is_op(node, _LINEAR_OPS) + + +def _has_downstream_linear(start: Node, max_depth: int = 3) -> bool: + """BFS from *start* through its users and return ``True`` if a linear op is reachable. + + The search only follows *user* edges (downstream in the data-flow graph) + and stops after *max_depth* hops. ``start`` itself is **not** checked. + """ + visited: set[Node] = {start} + queue: deque[Tuple[Node, int]] = deque() + + for user in start.users: + queue.append((user, 1)) + + while queue: + node, depth = queue.popleft() + if node in visited: + continue + visited.add(node) + + if _is_linear(node): + return True + + if depth < max_depth: + for user in node.users: + queue.append((user, depth + 1)) + + return False + + +def _find_kv_proj_linears(gm: GraphModule) -> List[Tuple[Node, Node]]: + """Find (fork_point, kv_linear) pairs suitable for aux-stream execution. + + A *fork point* is a node that directly feeds two or more supported linear + ops. Among these linears the one that does **not** lead to another linear + within a small BFS depth is the KV projection candidate (the lighter + branch). + + Returns a list of ``(fork_point, kv_linear_node)`` tuples. + """ + results: List[Tuple[Node, Node]] = [] + + for node in gm.graph.nodes: + # Collect direct linear users of this node. + linear_users = [u for u in node.users if _is_linear(u)] + if len(linear_users) < 2: + continue + + # Separate into "has downstream linear" (Q-like) and "does not" (KV-like). + kv_candidates = [ln for ln in linear_users if not _has_downstream_linear(ln)] + q_candidates = [ln for ln in linear_users if _has_downstream_linear(ln)] + + if not kv_candidates or not q_candidates: + continue + + # Pick the KV candidate(s). In MLA there is exactly one per fork point. + for kv_linear in kv_candidates: + results.append((node, kv_linear)) + + return results + + +def _create_aux_op(base_op: Callable) -> Callable: + """Create an ``_aux`` variant of a linear op that runs on the auxiliary CUDA stream. + + Uses a custom ``make_fake`` that delegates to the base op's registered fake + so that output shapes are computed correctly (linear output shape != input shape). + """ + return create_derived_custom_op( + base_op, + "_aux", + _make_aux_stream_impl, + make_fake=lambda base: lambda *a, **kw: base(*a, **kw), + ) + + +def _execute_kv_proj_in_aux_stream(gm: GraphModule) -> Tuple[GraphModule, int]: + """Replace KV projection linears with aux-stream variants. + + For each matched ``(fork_point, kv_linear)`` the rewriter: + + 1. Inserts ``record_event_passthrough(fork_point)`` so the main-stream + event is recorded *before* the Q-chain kernels are submitted. + 2. Replaces the KV linear's target with its ``_aux`` variant and wires the + ``record_event_passthrough`` output as the hidden-state input + (creating a true data dependency). + + The remaining KV-chain ops (split, rms_norm, view) stay on the main + stream — they are lightweight and run after the aux wait that is built + into the derived op. + + Aux-stream variants are created lazily — only for base ops that actually + appear in the matched KV positions. + """ + pairs = _find_kv_proj_linears(gm) + if not pairs: + return gm, 0 + + graph = gm.graph + node_order = {n: i for i, n in enumerate(graph.nodes)} + + # Create aux ops lazily for whatever linear op types are found. + ops_in_graph = {kv_linear.target for _, kv_linear in pairs} + op_dict = {op: _create_aux_op(op) for op in ops_in_graph} + + num_replaced = 0 + + for fork_point, kv_linear in pairs: + # Find the Q-chain linear(s) so we can insert the event record + # *before* the earliest Q-chain op in graph order. + q_linears = [u for u in fork_point.users if _is_linear(u) and u is not kv_linear] + earliest_q = min(q_linears, key=lambda n: node_order.get(n, 0)) + + # Insert record_event_passthrough right before the first Q-chain + # linear so the event is recorded before Q kernels hit the GPU. + with graph.inserting_before(earliest_q): + rec_node = graph.call_function( + record_event_passthrough, + args=(fork_point,), + ) + + # Replace KV linear with its aux-stream variant. The hidden-state + # input (args[0]) is rewired to ``rec_node`` to create a data + # dependency that ensures the event is recorded first. + new_args = tuple(rec_node if arg is fork_point else arg for arg in kv_linear.args) + + with graph.inserting_after(kv_linear): + new_node = graph.call_function( + op_dict[kv_linear.target], args=new_args, kwargs=kv_linear.kwargs + ) + + kv_linear.replace_all_uses_with(new_node) + graph.erase_node(kv_linear) + num_replaced += 1 + + return gm, num_replaced + + +# --------------------------------------------------------------------------- +# Transform class +# --------------------------------------------------------------------------- + + +@TransformRegistry.register("multi_stream_mla_attn") +class MultiStreamMLAAttn(BaseTransform): + """Multi-stream Q/KV projection parallelism for MLA attention blocks.""" + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + # Ensure aux stream and events are set up for the current device. + cuda_stream_manager.add_device(torch.cuda.current_device()) + + gm, num_matches = _execute_kv_proj_in_aux_stream(gm) + + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py index 5cd02553d958..bc1e0fe9f342 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py @@ -1,295 +1,195 @@ """Transform for multi-stream execution of MoE layers that have shared experts and routed experts.""" -from threading import RLock -from typing import Any, Callable, Dict, List, Tuple +from typing import Callable, List, Optional, Set, Tuple import torch -from torch.fx import GraphModule - -from tensorrt_llm._torch.utils import ActivationType +from torch.fx import GraphModule, Node from ...models.factory import ModelFactory from ...shim.interface import CachedSequenceInterface from ...utils.logger import ad_logger +from ...utils.multi_stream_utils import ( + begin_aux_stream_passthrough, + cuda_stream_manager, + end_aux_stream_passthrough, + wait_aux_stream_passthrough, +) from ...utils.node_utils import is_op from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry -# Previously, CudaStreamManager and the custom ops that use the cuda streams and events were -# placed in custom_ops folder. However doing so resulted in CudaStreamManager -# being created only in the parent process, but we need each rank to have its own CudaStreamManager that -# manages the cuda streams and events for that rank. Placing the logic to instantiate -# CudaStreamManager and the custom ops that use the cuda streams and events at the transform level ensures that -# each rank has its own CudaStreamManager since each rank applies the transform independently. -class _Singleton(type): - _instances: Dict[type, Any] = {} - _lock = RLock() - - def __call__(cls, *args: Any, **kwargs: Any) -> Any: - if cls not in cls._instances: - with cls._lock: - if cls not in cls._instances: # double-checked locking - cls._instances[cls] = super().__call__(*args, **kwargs) - return cls._instances[cls] - - -# A singleton that holds the pointers to the cuda streams and events. -# Each device has its own cuda streams and events. -class CudaStreamManager(metaclass=_Singleton): - AUX_STREAM_NAME = "aux" - MAIN_STREAM_NAME = "main" - devices: List[torch.device] = [] - events: Dict[torch.device, Dict[str, Any]] = {} - streams: Dict[torch.device, Dict[str, Any]] = {} - - def __init__(self) -> None: - # In case __init__ ever gets called twice, guard against re-init - if hasattr(self, "streams"): - return - - self._lock = RLock() - self.add_device(torch.cuda.current_device()) - - def add_device(self, device: int) -> None: - if device not in self.devices: - self.devices.append(device) - with torch.cuda.device(device): - self.events[device] = { - self.AUX_STREAM_NAME: torch.cuda.Event(), - self.MAIN_STREAM_NAME: torch.cuda.Event(), - } - self.streams[device] = { - self.AUX_STREAM_NAME: torch.cuda.Stream(), - self.MAIN_STREAM_NAME: torch.cuda.default_stream(), - } - else: - ad_logger.warning(f"CudaStreamManager: Device {device} already added") - - def get_stream(self, device: int, stream_name: str) -> torch.cuda.Stream: - return self.streams[device][stream_name] - - def get_event(self, device: int, event_name: str) -> torch.cuda.Event: - return self.events[device][event_name] - - -# Every device will have a singleton instance of CudaStreamManager. -cuda_stream_manager = CudaStreamManager() - - -@torch.library.custom_op("auto_deploy::record_event", mutates_args=()) -def record_event(device: int, stream_name: str) -> None: - event = cuda_stream_manager.get_event(device, stream_name) - event.record() - - -@torch.library.custom_op("auto_deploy::wait_event", mutates_args=()) -def wait_event(device: int, stream_name: str) -> None: - event = cuda_stream_manager.get_event(device, stream_name) - event.wait() - - -# skip during compilation -@torch._dynamo.disable -def record_event_wrapper( - fn: Callable, - *args: Tuple[Any, ...], - **kwargs: Dict[str, Any], -) -> torch.Tensor: - device = kwargs.pop("device", torch.cuda.current_device()) - output = fn(*args, **kwargs) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - return output - - -@torch._dynamo.disable -def aux_stream_wrapper( - fn: Callable, - *args: Tuple[Any, ...], - **kwargs: Dict[str, Any], -) -> torch.Tensor: - stream_name = cuda_stream_manager.AUX_STREAM_NAME - device = kwargs.pop("device", torch.cuda.current_device()) - with torch.cuda.stream(cuda_stream_manager.get_stream(device, stream_name)): - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - output = fn(*args, **kwargs) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) - return output - - -# trtllm bf16 -@torch.library.custom_op("auto_deploy::trtllm_moe_fused_aux", mutates_args=()) -def trtllm_moe_fused_aux( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w3_w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - device = torch.cuda.current_device() - with torch.cuda.stream( - cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) - ): - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - output = torch.ops.auto_deploy.trtllm_moe_fused( - x, - selected_experts, - routing_weights, - w3_w1_stacked_weight, - w2_stacked_weight, - is_gated_mlp, - act_fn, - ) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) - return output - - -@trtllm_moe_fused_aux.register_fake -def trtllm_moe_fused_aux_fake( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w3_w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - return torch.empty_like(x) - - -# triton bf16 -@torch.library.custom_op("auto_deploy::triton_moe_fused_aux", mutates_args=()) -def triton_moe_fused_aux( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, -) -> torch.Tensor: - device = torch.cuda.current_device() - with torch.cuda.stream( - cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) - ): - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - output = torch.ops.auto_deploy.triton_moe_fused( - x, - selected_experts, - routing_weights, - w1_stacked_weight, - w2_stacked_weight, - ) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) - return output - - -@triton_moe_fused_aux.register_fake -def triton_moe_fused_aux_fake( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, -) -> torch.Tensor: - return torch.empty_like(x) - - -@torch.library.custom_op("auto_deploy::trtllm_quant_fp8_moe_fused_aux", mutates_args=()) -def trtllm_quant_fp8_moe_fused_aux( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - fc1_expert_weights: torch.Tensor, - fc2_expert_weights: torch.Tensor, - fc1_act_scale: torch.Tensor, - fc1_dequant_scale: torch.Tensor, - fc2_act_scale_reciprocal: torch.Tensor, - fc2_dequant_scale: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - device = torch.cuda.current_device() - with torch.cuda.stream( - cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) - ): - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - output = torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused( - x, - selected_experts, - routing_weights, - fc1_expert_weights, - fc2_expert_weights, - fc1_act_scale, - fc1_dequant_scale, - fc2_act_scale_reciprocal, - fc2_dequant_scale, - is_gated_mlp, - act_fn, - ) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) - return output - - -@trtllm_quant_fp8_moe_fused_aux.register_fake -def trtllm_quant_fp8_moe_fused_aux_fake( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - fc1_expert_weights: torch.Tensor, - fc2_expert_weights: torch.Tensor, - fc1_act_scale: torch.Tensor, - fc1_dequant_scale: torch.Tensor, - fc2_act_scale_reciprocal: torch.Tensor, - fc2_dequant_scale: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - return torch.empty_like(x) - - -def _execute_op_in_aux_stream( - gm: GraphModule, op_dict: Dict[Callable, Callable] +def _find_merge_add(moe_node: Node) -> Optional[Node]: + """Walk forward from a MoE op through users to find the ``aten.add.Tensor`` merge node. + + The merge ``add`` is the node where shared-expert output and routed-expert + output are combined. The search is a breadth-first traversal of the user + graph starting from the MoE node. + """ + visited: Set[Node] = set() + queue = list(moe_node.users.keys()) + while queue: + n = queue.pop(0) + if n in visited: + continue + visited.add(n) + if is_op(n, torch.ops.aten.add.Tensor): + return n + queue.extend(n.users.keys()) + return None + + +def _get_ancestors(node: Node) -> Set[Node]: + """Return the set of all nodes reachable by walking backwards from *node*.""" + ancestors: Set[Node] = set() + queue = list(node.all_input_nodes) + while queue: + n = queue.pop() + if n in ancestors: + continue + ancestors.add(n) + queue.extend(n.all_input_nodes) + return ancestors + + +def _execute_shared_expert_in_aux_stream( + gm: GraphModule, moe_ops: List[Callable] ) -> Tuple[GraphModule, int]: + """Move shared-expert computation to the auxiliary CUDA stream. + + For each MoE fused op in the graph: + 1. Walk forward to find the ``aten.add.Tensor`` that merges the + shared-expert output and the routed-expert output. + 2. Identify which ``add`` input is the routed branch (descended from + the MoE node) and which is the shared-expert branch. + 3. Trace the shared-expert branch backwards to collect all its + computation nodes and identify the fork point (the latest common + ancestor shared with the MoE / routing path). + 4. Insert ``begin_aux_stream_passthrough`` before the first shared-expert + op to switch to the auxiliary CUDA stream. + 5. Insert ``end_aux_stream_passthrough`` after the last shared-expert op + to switch back to the main stream. + 6. Insert ``wait_aux_stream_passthrough`` on the routed-branch input + just before the ``add`` so the main stream waits for the auxiliary + stream to finish before merging outputs. + """ graph = gm.graph num_replaced = 0 - # Collect targets first to avoid mutating while iterating - target_nodes = [n for n in graph.nodes if is_op(n, op_dict.keys())] - - for n in target_nodes: - target_input_node = None - for input_node in n.all_input_nodes: - if input_node.target == torch.ops.aten.view.default: - target_input_node = input_node - break - # Look through dtype cast nodes (aten.to) to find the view node - if input_node.target == torch.ops.aten.to: - for nested_input in input_node.all_input_nodes: - if nested_input.target == torch.ops.aten.view.default: - target_input_node = nested_input - break - if target_input_node is not None: - break - - assert target_input_node is not None, f"Target input node not found for node {n}" - with graph.inserting_before(target_input_node): - kwargs = target_input_node.kwargs.copy() - kwargs["device"] = torch.cuda.current_device() - new_node = graph.call_function( - record_event_wrapper, - args=(target_input_node.target, *target_input_node.args), - kwargs=kwargs, + # Collect targets first to avoid mutating while iterating. + target_nodes = [n for n in graph.nodes if is_op(n, moe_ops)] + if not target_nodes: + return gm, 0 + + node_order = {node: i for i, node in enumerate(graph.nodes)} + + for moe_node in target_nodes: + # ---- Step 1: Find the merge ``add`` node. ---- + add_node = _find_merge_add(moe_node) + if add_node is None: + ad_logger.warning( + f"No merge add found downstream of MoE node {moe_node.name}; " + "skipping multi-stream transform for this node." + ) + continue + + # ---- Step 2: Determine which ``add`` input is routed vs. shared. ---- + arg0, arg1 = add_node.args[0], add_node.args[1] + arg0_ancestors = _get_ancestors(arg0) + + if moe_node in arg0_ancestors or arg0 is moe_node: + routed_output, shared_output = arg0, arg1 + else: + routed_output, shared_output = arg1, arg0 + + # ---- Step 3: Collect shared-expert nodes & find fork point. ---- + moe_ancestors = _get_ancestors(moe_node) + moe_ancestors.add(moe_node) + + shared_nodes: List[Node] = [] + fork_point: Optional[Node] = None + visited: Set[Node] = set() + queue = [shared_output] + + while queue: + n = queue.pop(0) + if n in visited: + continue + visited.add(n) + + # Skip static weight / parameter nodes. + if n.op == "get_attr": + continue + + if n in moe_ancestors: + # This node is on the MoE / routing path — candidate fork point. + if fork_point is None or node_order.get(n, 0) > node_order.get(fork_point, 0): + fork_point = n + continue + + shared_nodes.append(n) + for inp in n.all_input_nodes: + queue.append(inp) + + if not shared_nodes or fork_point is None: + ad_logger.warning( + f"Could not identify shared-expert subgraph for MoE node " + f"{moe_node.name}; skipping multi-stream transform for this node." + ) + continue + + # Order shared nodes by their position in the graph. + shared_nodes.sort(key=lambda n: node_order.get(n, 0)) + first_shared = shared_nodes[0] + + # Sanity check: the first shared op must directly consume the fork + # point so we can wire begin_aux_stream_passthrough into it. + if fork_point not in first_shared.all_input_nodes: + ad_logger.warning( + f"First shared-expert op ({first_shared.name}) does not directly " + f"consume fork point ({fork_point.name}); skipping." + ) + continue + + # ---- Step 4: Insert begin_aux before the first shared-expert op. ---- + # NOTE: do NOT bake ``torch.cuda.current_device()`` into the graph — + # that would hard-code device 0 and break on other ranks in a + # multi-GPU setup. Omitting ``device`` lets the passthrough + # functions resolve the device at **runtime** (default ``-1``). + with graph.inserting_before(first_shared): + begin_aux_node = graph.call_function( + begin_aux_stream_passthrough, + args=(fork_point,), + ) + + # Create a data dependency: first_shared reads begin_aux output + # instead of fork_point. + first_shared.args = tuple( + begin_aux_node if arg is fork_point else arg for arg in first_shared.args + ) + + # ---- Step 5: Insert end_aux after the last shared-expert op. ---- + with graph.inserting_after(shared_output): + end_aux_node = graph.call_function( + end_aux_stream_passthrough, + args=(shared_output,), ) - target_input_node.replace_all_uses_with(new_node) - graph.erase_node(target_input_node) - with graph.inserting_after(n): - new_node = graph.call_function(op_dict[n.target], args=n.args, kwargs=n.kwargs) - n.replace_all_uses_with(new_node) - graph.erase_node(n) + + # Replace shared-expert input to ``add`` with end_aux output. + add_node.args = tuple( + end_aux_node if arg is shared_output else arg for arg in add_node.args + ) + + # ---- Step 6: Insert wait_aux before the ``add``. ---- + with graph.inserting_before(add_node): + wait_aux_node = graph.call_function( + wait_aux_stream_passthrough, + args=(routed_output,), + ) + + add_node.args = tuple( + wait_aux_node if arg is routed_output else arg for arg in add_node.args + ) + num_replaced += 1 return gm, num_replaced @@ -306,14 +206,16 @@ def _apply( factory: ModelFactory, shared_config: SharedConfig, ) -> Tuple[GraphModule, TransformInfo]: - op_dict = { - torch.ops.auto_deploy.trtllm_moe_fused: torch.ops.auto_deploy.trtllm_moe_fused_aux, - torch.ops.auto_deploy.triton_moe_fused: torch.ops.auto_deploy.triton_moe_fused_aux, - torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused: torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused_aux, - } + base_ops = [ + torch.ops.auto_deploy.trtllm_moe_fused, + torch.ops.auto_deploy.triton_moe_fused, + torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused, + torch.ops.auto_deploy.trtllm_quant_nvfp4_moe_fused, + ] + # Ensure that aux stream and events for the current device are added to the CudaStreamManager. cuda_stream_manager.add_device(torch.cuda.current_device()) - gm, num_matches = _execute_op_in_aux_stream(gm, op_dict) + gm, num_matches = _execute_shared_expert_in_aux_stream(gm, base_ops) info = TransformInfo( skipped=False, @@ -321,5 +223,4 @@ def _apply( is_clean=num_matches == 0, has_valid_shapes=num_matches == 0, ) - return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py index 749b2cd130af..d54c86d0c37d 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py @@ -20,6 +20,100 @@ from .logger import ad_logger from .node_utils import get_weight_tensor, is_op +# --------------------------------------------------------------------------- +# Dynamic custom-op derivation helpers +# --------------------------------------------------------------------------- +# These are used to create new custom ops that share the schema of an existing +# op but wrap it with additional logic (e.g. stream management). A single +# module-level dict of ``Library`` objects (keyed by namespace) is used so that +# the registrations persist and are visible via ``torch.ops..*``. +# --------------------------------------------------------------------------- + +_derived_op_libs: Dict[str, torch.library.Library] = {} +_derived_op_registry: Dict[str, Callable] = {} + + +def _get_lib(namespace: str) -> torch.library.Library: + """Return (and lazily create) a ``FRAGMENT`` Library for *namespace*.""" + if namespace not in _derived_op_libs: + _derived_op_libs[namespace] = torch.library.Library(namespace, "FRAGMENT") + return _derived_op_libs[namespace] + + +def create_derived_custom_op( + base_op: Callable, + suffix: str, + make_impl: Callable[[Callable], Callable], + make_fake: Optional[Callable[[Callable], Callable]] = None, +) -> Callable: + """Dynamically create a new custom op derived from an existing one. + + The new op has the **same** schema (arguments, default values, and return + type) as *base_op* but with a different name (````) and a + custom implementation produced by *make_impl*. + + Args: + base_op: The base custom op — either an ``OpOverloadPacket`` + (e.g. ``torch.ops.auto_deploy.trtllm_moe_fused``) or an + ``OpOverload`` (e.g. ``…trtllm_moe_fused.default``). + suffix: Suffix appended to the base op name to form the new op name + (e.g. ``"_aux"``). + make_impl: A factory ``(base_overload) -> impl_fn`` that receives the + resolved base ``OpOverload`` and returns the *implementation* + function for the new op. ``impl_fn`` will be called with the + same positional/keyword arguments as *base_op*. + make_fake: Optional factory ``(base_overload) -> fake_fn`` that returns + the *Meta / fake-tensor* implementation. When ``None`` the + default fake implementation ``torch.empty_like(args[0])`` is used. + + Returns: + The newly registered op as an ``OpOverloadPacket`` + (e.g. ``torch.ops.auto_deploy.``). Repeated calls with + the same *base_op* and *suffix* return the cached op. + """ + base_overload = base_op.default if hasattr(base_op, "default") else base_op + schema = base_overload._schema + + # e.g. "auto_deploy::trtllm_moe_fused" + qualified_name = schema.name + namespace, base_name = qualified_name.split("::") + new_name = f"{base_name}{suffix}" + new_qualified = f"{namespace}::{new_name}" + + # Return the cached op if it was already created. + if new_qualified in _derived_op_registry: + return _derived_op_registry[new_qualified] + + # Build the schema string for the derived op. ``str(schema)`` produces a + # fully-qualified string such as + # auto_deploy::trtllm_moe_fused(Tensor x, …) -> Tensor + # We replace the qualified name with the bare new name (the Library already + # knows its namespace). + new_schema_str = str(schema).replace(qualified_name, new_name, 1) + + lib = _get_lib(namespace) + lib.define(new_schema_str) + + # Register the real implementation for all devices. + # We use "CompositeExplicitAutograd" so that we can provide a separate + # Meta / fake kernel for shape inference. + lib.impl(new_name, make_impl(base_overload), "CompositeExplicitAutograd") + + # Register the Meta / fake implementation. + if make_fake is not None: + lib.impl(new_name, make_fake(base_overload), "Meta") + else: + + def _default_fake(*args, **kwargs): + return torch.empty_like(args[0]) + + lib.impl(new_name, _default_fake, "Meta") + + new_op = getattr(getattr(torch.ops, namespace), new_name) + _derived_op_registry[new_qualified] = new_op + return new_op + + _NoValType = type("_NoValType", (), {}) diff --git a/tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py new file mode 100644 index 000000000000..2c7b17da9030 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py @@ -0,0 +1,242 @@ +"""Shared CUDA multi-stream utilities for multi-stream transforms. + +This module provides the core infrastructure for executing parts of an FX graph +on auxiliary CUDA streams. It is consumed by the multi-stream MoE and MLA +attention transforms in ``..transform.library``. + +Key components: + - ``CudaStreamManager``: per-device singleton managing auxiliary streams/events. + - Custom ops ``record_event`` / ``wait_event``: graph-safe event primitives. + - Passthrough helpers that switch streams while preserving the data-flow edges + required by FX graph execution and CUDA graph capture. + - ``_make_aux_stream_impl``: factory for building an implementation that runs + a base op on the auxiliary CUDA stream. +""" + +from threading import RLock +from typing import Any, Callable, Dict, List + +import torch + +from .logger import ad_logger + +# --------------------------------------------------------------------------- +# Singleton metaclass +# --------------------------------------------------------------------------- + + +class _Singleton(type): + _instances: Dict[type, Any] = {} + _lock = RLock() + + def __call__(cls, *args: Any, **kwargs: Any) -> Any: + if cls not in cls._instances: + with cls._lock: + if cls not in cls._instances: # double-checked locking + cls._instances[cls] = super().__call__(*args, **kwargs) + return cls._instances[cls] + + +# --------------------------------------------------------------------------- +# CudaStreamManager +# --------------------------------------------------------------------------- + +# Previously, CudaStreamManager and the custom ops that use the cuda streams and events were +# placed in custom_ops folder. However doing so resulted in CudaStreamManager +# being created only in the parent process, but we need each rank to have its own CudaStreamManager that +# manages the cuda streams and events for that rank. Placing the logic to instantiate +# CudaStreamManager and the custom ops that use the cuda streams and events at the transform level ensures that +# each rank has its own CudaStreamManager since each rank applies the transform independently. + + +class CudaStreamManager(metaclass=_Singleton): + AUX_STREAM_NAME = "aux" + MAIN_STREAM_NAME = "main" + devices: List[torch.device] = [] + events: Dict[torch.device, Dict[str, Any]] = {} + streams: Dict[torch.device, Dict[str, Any]] = {} + # Per-device save slot for the caller's stream. ``begin_aux_stream_passthrough`` + # saves the real current stream here so that ``end_aux_stream_passthrough`` can + # restore it — this is critical during CUDA graph capture where the capture stream + # differs from ``torch.cuda.default_stream()``. + _caller_streams: Dict[int, Any] = {} + + def __init__(self) -> None: + # In case __init__ ever gets called twice, guard against re-init + if hasattr(self, "streams"): + return + + self._lock = RLock() + self.add_device(torch.cuda.current_device()) + + def add_device(self, device: int) -> None: + if device not in self.devices: + self.devices.append(device) + with torch.cuda.device(device): + self.events[device] = { + self.AUX_STREAM_NAME: torch.cuda.Event(), + self.MAIN_STREAM_NAME: torch.cuda.Event(), + } + self.streams[device] = { + self.AUX_STREAM_NAME: torch.cuda.Stream(), + self.MAIN_STREAM_NAME: torch.cuda.default_stream(), + } + else: + ad_logger.warning(f"CudaStreamManager: Device {device} already added") + + def get_stream(self, device: int, stream_name: str) -> torch.cuda.Stream: + return self.streams[device][stream_name] + + def get_event(self, device: int, event_name: str) -> torch.cuda.Event: + return self.events[device][event_name] + + +# Every device will have a singleton instance of CudaStreamManager. +cuda_stream_manager = CudaStreamManager() + + +# --------------------------------------------------------------------------- +# Custom ops — graph-safe CUDA event primitives +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("auto_deploy::record_event", mutates_args=()) +def record_event(device: int, stream_name: str) -> None: + event = cuda_stream_manager.get_event(device, stream_name) + event.record() + + +@torch.library.custom_op("auto_deploy::wait_event", mutates_args=()) +def wait_event(device: int, stream_name: str) -> None: + event = cuda_stream_manager.get_event(device, stream_name) + event.wait() + + +# --------------------------------------------------------------------------- +# Passthrough helpers +# --------------------------------------------------------------------------- + + +@torch._dynamo.disable +def record_event_passthrough( + x: torch.Tensor, + *, + device: int = -1, +) -> torch.Tensor: + """Record a CUDA event on the main stream and return the input unchanged. + + Inserted after the gating/routing computation to mark a synchronization + point. The aux stream waits for this event before starting the MoE + computation, enabling overlap between the shared expert (main stream) + and routed experts (aux stream). + """ + if device < 0: + device = torch.cuda.current_device() + torch.ops.auto_deploy.record_event(device, cuda_stream_manager.MAIN_STREAM_NAME) + return x + + +@torch._dynamo.disable +def begin_aux_stream_passthrough( + x: torch.Tensor, + *, + device: int = -1, +) -> torch.Tensor: + """Record a CUDA event on the main stream, switch to aux, and wait for it. + + After this function returns the thread-local current stream is the + auxiliary stream. All subsequent GPU ops dispatched by the FX graph + interpreter will be recorded on aux until ``end_aux_stream_passthrough`` + switches back to main. + """ + if device < 0: + device = torch.cuda.current_device() + # Save the *actual* current stream so ``end_aux`` can restore it. + # During CUDA graph capture the current stream is the capture stream, + # which is NOT ``torch.cuda.default_stream()``. + caller_stream = torch.cuda.current_stream(device) + cuda_stream_manager._caller_streams[device] = caller_stream + # Record where the caller's stream has reached so aux knows when data is ready. + main_event = cuda_stream_manager.get_event(device, cuda_stream_manager.MAIN_STREAM_NAME) + main_event.record(caller_stream) + # Switch the thread-local current stream to aux. + aux_stream = cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) + torch.cuda.set_stream(aux_stream) + # Make aux wait for the main-stream event before executing any work. + aux_stream.wait_event(main_event) + return x + + +@torch._dynamo.disable +def end_aux_stream_passthrough( + x: torch.Tensor, + *, + device: int = -1, +) -> torch.Tensor: + """Record a CUDA event on the aux stream and switch back to the caller's stream. + + This does **not** make the caller's stream wait for aux. The caller must + insert ``wait_aux_stream_passthrough`` at the point where both branches + need to be synchronised (typically right before the ``add`` that merges + shared-expert and routed-expert outputs). + """ + if device < 0: + device = torch.cuda.current_device() + # Record the aux-stream progress so the caller's stream can wait for it later. + aux_event = cuda_stream_manager.get_event(device, cuda_stream_manager.AUX_STREAM_NAME) + aux_event.record() + # Restore the caller's stream saved by ``begin_aux_stream_passthrough``. + # This is critical during CUDA graph capture where the capture stream + # differs from ``torch.cuda.default_stream()``. + caller_stream = cuda_stream_manager._caller_streams.pop(device, None) + if caller_stream is not None: + torch.cuda.set_stream(caller_stream) + else: + torch.cuda.set_stream( + cuda_stream_manager.get_stream(device, cuda_stream_manager.MAIN_STREAM_NAME) + ) + return x + + +@torch._dynamo.disable +def wait_aux_stream_passthrough( + x: torch.Tensor, + *, + device: int = -1, +) -> torch.Tensor: + """Make the current stream wait for the auxiliary stream's last recorded event. + + This is a GPU-side wait (non-blocking on the CPU). Insert this right + before the ``add`` that merges shared-expert output (computed on aux) + with routed-expert output (computed on main). + + Uses ``torch.cuda.current_stream()`` rather than the stored default stream + so that the correct stream is waited on during CUDA graph capture. + """ + if device < 0: + device = torch.cuda.current_device() + aux_event = cuda_stream_manager.get_event(device, cuda_stream_manager.AUX_STREAM_NAME) + torch.cuda.current_stream(device).wait_event(aux_event) + return x + + +# --------------------------------------------------------------------------- +# Aux-stream implementation factory +# --------------------------------------------------------------------------- + + +def _make_aux_stream_impl(base_overload: Callable) -> Callable: + """Build an implementation that runs *base_overload* on the auxiliary CUDA stream.""" + + def _impl(*args, **kwargs): + device = torch.cuda.current_device() + with torch.cuda.stream( + cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) + ): + torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) + output = base_overload(*args, **kwargs) + torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) + torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) + return output + + return _impl diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 905880b06cf6..ef6a0ab72dab 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -65,6 +65,11 @@ class TestLlama3_1_8B(LlmapiAccuracyTestHarness): "max_seq_len": 8192, "compile_backend": "torch-cudagraph", }, + "trtllm": { + "max_batch_size": 512, + "max_seq_len": 8192, + "compile_backend": "torch-cudagraph", + }, "torch": { "max_batch_size": 128, "max_seq_len": 2048, @@ -117,7 +122,7 @@ def get_default_sampling_params(self): @pytest.mark.skip_less_device_memory(32000) @pytest.mark.parametrize("world_size", [1, 2, 4]) @pytest.mark.parametrize("enable_chunked_prefill", [False, True]) - @pytest.mark.parametrize("attn_backend", ["flashinfer", "torch"]) + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm", "torch"]) def test_auto_dtype(self, world_size, enable_chunked_prefill, attn_backend): kwargs = self.get_default_kwargs(enable_chunked_prefill, attn_backend) sampling_params = self.get_default_sampling_params() @@ -154,10 +159,13 @@ class TestNemotronH(LlmapiAccuracyTestHarness): MODEL_NAME = "nvidia/Nemotron-H-8B-Base-8K" MODEL_PATH = f"{llm_models_root()}/Nemotron-H-8B-Base-8K" - def get_default_kwargs(self, enable_chunked_prefill=False): + def get_default_kwargs(self, + enable_chunked_prefill=False, + attn_backend="flashinfer"): config = { "skip_tokenizer_init": False, "trust_remote_code": True, + "attn_backend": attn_backend, # SSMs do not support cache reuse. "kv_cache_config": { "enable_block_reuse": False, @@ -194,8 +202,10 @@ def get_default_sampling_params(self): @pytest.mark.skip_less_device_memory(32000) @pytest.mark.parametrize("enable_chunked_prefill", [False, True]) @pytest.mark.parametrize("ssm_backend", ["triton_ssm", "flashinfer_ssm"]) - def test_auto_dtype(self, enable_chunked_prefill, ssm_backend): - kwargs = self.get_default_kwargs(enable_chunked_prefill) + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_auto_dtype(self, enable_chunked_prefill, ssm_backend, + attn_backend): + kwargs = self.get_default_kwargs(enable_chunked_prefill, attn_backend) kwargs.setdefault("transforms", {}) insert_ssm_cfg = {"backend": ssm_backend} if ssm_backend == "flashinfer_ssm": @@ -217,10 +227,11 @@ class TestNemotronMOE(LlmapiAccuracyTestHarness): MODEL_PATH_FP8 = f"{llm_models_root()}/Nemotron-Nano-3-30B-A3.5B-FP8-KVFP8-dev" MODEL_PATH_NVFP4 = f"{llm_models_root()}/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4" - def get_default_kwargs(self, world_size=1): + def get_default_kwargs(self, world_size=1, attn_backend="flashinfer"): return { "skip_tokenizer_init": False, "trust_remote_code": True, + "attn_backend": attn_backend, # SSMs do not support cache reuse. "kv_cache_config": { "enable_block_reuse": False, @@ -261,8 +272,10 @@ def get_default_sampling_params(self): @pytest.mark.skip_less_device_memory(32000) @pytest.mark.parametrize("world_size", [1, 4]) - def test_bf16(self, world_size): - kwargs = self.get_default_kwargs(world_size=world_size) + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_bf16(self, world_size, attn_backend): + kwargs = self.get_default_kwargs(world_size=world_size, + attn_backend=attn_backend) # TODO: multi-stream MOE seems to increase the memory usage kwargs["max_batch_size"] = 32 kwargs["kv_cache_config"] = {"free_gpu_memory_fraction": 0.4} @@ -279,8 +292,10 @@ def test_bf16(self, world_size): @pytest.mark.skip_less_device_memory(32000) @pytest.mark.parametrize("world_size", [1, 4]) - def test_fp8(self, world_size): - kwargs = self.get_default_kwargs(world_size=world_size) + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_fp8(self, world_size, attn_backend): + kwargs = self.get_default_kwargs(world_size=world_size, + attn_backend=attn_backend) with AutoDeployLLM(model=self.MODEL_PATH_FP8, tokenizer=self.MODEL_PATH_FP8, world_size=world_size, @@ -296,8 +311,9 @@ def test_fp8(self, world_size): @skip_pre_blackwell @pytest.mark.parametrize("world_size", [1, 2, 4]) - def test_nvfp4(self, world_size): - kwargs = self.get_default_kwargs() + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_nvfp4(self, world_size, attn_backend): + kwargs = self.get_default_kwargs(attn_backend=attn_backend) with AutoDeployLLM(model=self.MODEL_PATH_NVFP4, tokenizer=self.MODEL_PATH_NVFP4, world_size=world_size, @@ -328,10 +344,11 @@ class TestNemotronSuperV3(LlmapiAccuracyTestHarness): MAX_SEQ_LEN = max(MMLU.MAX_INPUT_LEN + MMLU.MAX_OUTPUT_LEN, GSM8K.MAX_INPUT_LEN + GSM8K.MAX_OUTPUT_LEN) - def get_default_kwargs(self): + def get_default_kwargs(self, attn_backend="flashinfer"): return { "skip_tokenizer_init": False, "trust_remote_code": True, + "attn_backend": attn_backend, "skip_loading_weights": False, "compile_backend": "torch-cudagraph", "max_batch_size": 128, @@ -343,6 +360,10 @@ def get_default_kwargs(self): "sharding_source": ['factory', 'heuristic'], "sharding_dims": ['ep', 'bmm'], }, + "multi_stream_moe": { + "stage": "compile", + "enabled": True, + }, } } @@ -357,8 +378,9 @@ def get_default_sampling_params(self): # 180GB works, might be able to go lower @pytest.mark.skip_less_device_memory(180000) @pytest.mark.skip_less_device(4) - def test_bf16(self): - kwargs = self.get_default_kwargs() + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_bf16(self, attn_backend): + kwargs = self.get_default_kwargs(attn_backend=attn_backend) sampling_params = self.get_default_sampling_params() print_memory_usage("Before evaluation") with AutoDeployLLM(model=self.MODEL_PATH_BF16, @@ -373,10 +395,11 @@ def test_bf16(self): @pytest.mark.skip_less_device_memory(180000) @pytest.mark.parametrize("world_size", [1, 4, 8]) - def test_fp8(self, world_size): + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_fp8(self, world_size, attn_backend): if get_device_count() < world_size: pytest.skip("Not enough devices for world size, skipping test") - kwargs = self.get_default_kwargs() + kwargs = self.get_default_kwargs(attn_backend=attn_backend) sampling_params = self.get_default_sampling_params() with AutoDeployLLM(model=self.MODEL_PATH_FP8, tokenizer=self.MODEL_PATH_FP8, @@ -394,10 +417,11 @@ def test_fp8(self, world_size): @pytest.mark.skip("Skipping FP4 test until it is supported") @pytest.mark.skip_less_device_memory(180000) @pytest.mark.parametrize("world_size", [4, 8]) - def test_fp4(self, world_size): + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_fp4(self, world_size, attn_backend): if get_device_count() < world_size: pytest.skip("Not enough devices for world size, skipping test") - kwargs = self.get_default_kwargs() + kwargs = self.get_default_kwargs(attn_backend=attn_backend) sampling_params = self.get_default_sampling_params() with AutoDeployLLM(model=self.MODEL_PATH_FP4, tokenizer=self.MODEL_PATH_FP4, @@ -433,10 +457,13 @@ class TestGLM4Flash(LlmapiAccuracyTestHarness): GSM8K.MAX_INPUT_LEN + GSM8K.MAX_OUTPUT_LEN) MAX_NUM_TOKENS = MAX_SEQ_LEN - def get_default_kwargs(self, enable_chunked_prefill=False): + def get_default_kwargs(self, + enable_chunked_prefill=False, + attn_backend="flashinfer"): config = { "skip_tokenizer_init": False, "trust_remote_code": True, + "attn_backend": attn_backend, "compile_backend": "torch-cudagraph", "max_batch_size": 128, "max_seq_len": self.MAX_SEQ_LEN, @@ -455,7 +482,15 @@ def get_default_kwargs(self, enable_chunked_prefill=False): "fuse_nvfp4_moe": { "allow_different_input_scales": True, }, - }, + "multi_stream_moe": { + "stage": "compile", + "enabled": True, + }, + "multi_stream_mla_attn": { + "stage": "compile", + "enabled": True, + }, + } } if enable_chunked_prefill: config["enable_chunked_prefill"] = True @@ -473,8 +508,9 @@ def get_default_sampling_params(self): @pytest.mark.skip_less_device_memory(32000) @pytest.mark.parametrize("enable_chunked_prefill", [True, False]) - def test_auto_dtype(self, enable_chunked_prefill): - kwargs = self.get_default_kwargs(enable_chunked_prefill) + @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) + def test_auto_dtype(self, enable_chunked_prefill, attn_backend): + kwargs = self.get_default_kwargs(enable_chunked_prefill, attn_backend) sampling_params = self.get_default_sampling_params() with AutoDeployLLM(model=self.MODEL_PATH, tokenizer=self.MODEL_PATH, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 30a48da2cb9c..3a0e609a4b25 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -188,7 +188,7 @@ l0_b200: stage: pre_merge backend: autodeploy tests: - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[flashinfer-False-1] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-1] - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[torch-True-1] - - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8[1] + - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8[trtllm-1] - unittest/_torch/auto_deploy/unit/singlegpu 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 fc4bbd824a7b..fae728b022f5 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -231,11 +231,11 @@ l0_dgx_b200: orchestrator: mpi tests: - unittest/_torch/auto_deploy/unit/multigpu - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[flashinfer-False-4] - - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8[4] - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_bf16 - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[4] - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[8] - - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_nvfp4[1] - - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_nvfp4[2] - - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_nvfp4[4] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-4] + - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8[trtllm-4] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_bf16[trtllm] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[trtllm-4] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[trtllm-8] + - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_nvfp4[trtllm-1] + - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_nvfp4[trtllm-2] + - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_nvfp4[trtllm-4] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index d58b93bfe1dd..c0a355719807 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -337,8 +337,8 @@ l0_dgx_h100: orchestrator: mpi tests: - unittest/_torch/auto_deploy/unit/multigpu - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[flashinfer-False-4] - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_bf16 - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[4] - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[8] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-4] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_bf16[trtllm] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[trtllm-4] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_fp8[trtllm-8] - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_attention_dp[4] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 80b26351622c..219db0914aeb 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -438,13 +438,13 @@ l0_h100: orchestrator: mpi tests: - unittest/_torch/auto_deploy/unit/singlegpu - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[flashinfer-False-1] - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[flashinfer-True-1] - - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[triton_ssm-False] - - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[flashinfer_ssm-False] - - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[triton_ssm-True] - - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8[1] - - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_bf16[1] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-1] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-True-1] + - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-triton_ssm-False] + - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-flashinfer_ssm-False] + - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-triton_ssm-True] + - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8[trtllm-1] + - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_bf16[trtllm-1] - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec_output[draft_target] - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec_output[eagle3] - examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_acceptance_rate diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6ca72b125300..ad7aae0ffbeb 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -303,7 +303,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_4B::test_eagle3 SKIP (https://nvbugs accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/5800646) accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664) examples/test_llama.py::test_llama_3_x_with_bf16_lora_torch[llama-3.2-1b-instruct] SKIP (https://nvbugs/5838178) -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_bf16 SKIP (https://nvbugs/5838184) +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_bf16[trtllm] SKIP (https://nvbugs/5838184) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] SKIP (https://nvbugs/5838199) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/5838211) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/5838211) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_captured_graph.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_captured_graph.py index c300dcd8e41b..5f20c084cf02 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_captured_graph.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_captured_graph.py @@ -1,17 +1,28 @@ +import operator +from unittest.mock import MagicMock + import pytest import torch +import torch.nn as nn from _model_test_utils import ( TransformerLikeModel, VisionTransformerLikeModel, generate_dynamic_shapes, ) +from torch.fx import Graph, GraphModule from tensorrt_llm._torch.auto_deploy.compile.backends.torch_cudagraph import ( CapturedGraph, + DualModeCapturedGraph, + PiecewiseCapturedGraph, _args_kwargs_flatten_spec, + _submod_has_cuda_ops, ) from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm from tensorrt_llm._torch.auto_deploy.shim.ad_executor import _round_up_to_closest +from tensorrt_llm._torch.auto_deploy.transform.library.compile_model import ( + _generate_default_piecewise_num_tokens, +) class ModelWithMultipleInputs(torch.nn.Module): @@ -159,3 +170,255 @@ def get_args_kwargs(bs): assert torch.allclose(original_output, replay_output, atol=atol), ( "CUDAGraph replay output mismatch" ) + + +# ============================================================================ +# Helpers for piecewise / _submod_has_cuda_ops tests +# ============================================================================ + + +def _build_trivial_graphmodule(): + """Build a GraphModule with only trivial ops (getitem, view).""" + graph = Graph() + x = graph.placeholder("x") + # getitem is trivial + item = graph.call_function(operator.getitem, args=(x, 0)) + # view is a trivial call_method + viewed = graph.call_method("view", args=(item, -1)) + graph.output(viewed) + root = nn.Module() + return GraphModule(root, graph) + + +def _build_graphmodule_with_linear(): + """Build a GraphModule that calls a Linear submodule (has CUDA ops).""" + + class SmallModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + + def forward(self, x): + return self.linear(x) + + model = SmallModel() + from torch.fx import symbolic_trace + + gm = symbolic_trace(model) + return gm + + +# ============================================================================ +# Tests for _submod_has_cuda_ops +# ============================================================================ + + +class TestSubmodHasCudaOps: + """Tests for _submod_has_cuda_ops.""" + + def test_trivial_graphmodule_returns_false(self): + gm = _build_trivial_graphmodule() + assert _submod_has_cuda_ops(gm) is False + + def test_graphmodule_with_linear_returns_true(self): + gm = _build_graphmodule_with_linear() + assert _submod_has_cuda_ops(gm) is True + + def test_non_graphmodule_returns_true(self): + """Non-FX modules are conservatively treated as having CUDA ops.""" + module = nn.Linear(4, 4) + assert _submod_has_cuda_ops(module) is True + + def test_graphmodule_with_nontrivial_call_function(self): + """A graph with torch.add (non-trivial call_function) should return True.""" + graph = Graph() + x = graph.placeholder("x") + y = graph.call_function(torch.add, args=(x, x)) + graph.output(y) + gm = GraphModule(nn.Module(), graph) + assert _submod_has_cuda_ops(gm) is True + + def test_graphmodule_with_nontrivial_call_method(self): + """A graph with 'matmul' method call should return True.""" + graph = Graph() + x = graph.placeholder("x") + y = graph.call_method("matmul", args=(x, x)) + graph.output(y) + gm = GraphModule(nn.Module(), graph) + assert _submod_has_cuda_ops(gm) is True + + def test_graphmodule_with_only_trivial_methods(self): + """A graph with only trivial call_methods should return False.""" + graph = Graph() + x = graph.placeholder("x") + y = graph.call_method("view", args=(x, -1)) + z = graph.call_method("contiguous", args=(y,)) + graph.output(z) + gm = GraphModule(nn.Module(), graph) + assert _submod_has_cuda_ops(gm) is False + + +# ============================================================================ +# Tests for DualModeCapturedGraph routing logic +# ============================================================================ + + +class TestDualModeCapturedGraphRouting: + """Tests for DualModeCapturedGraph routing logic (no actual graph capture).""" + + def _make_dual_mode(self, piecewise_num_tokens=None): + """Create a DualModeCapturedGraph with mock monolithic and piecewise.""" + if piecewise_num_tokens is None: + piecewise_num_tokens = [64, 128, 256] + + monolithic = MagicMock(spec=nn.Module) + monolithic.return_value = torch.tensor([1.0]) + + piecewise = MagicMock(spec=PiecewiseCapturedGraph) + piecewise.piecewise_num_tokens = piecewise_num_tokens + piecewise.original_model = MagicMock(return_value=torch.tensor([2.0])) + piecewise.return_value = torch.tensor([3.0]) + + dual = DualModeCapturedGraph(monolithic, piecewise) + return dual + + def test_is_decode_only_with_batch_info_host_zero(self): + dual = self._make_dual_mode() + # num_prefill=0 → decode-only + batch_info = torch.tensor([0, 0, 4]) # [num_prefill, num_prefill_tokens, num_decode] + assert dual._is_decode_only(batch_info_host=batch_info) is True + + def test_is_decode_only_with_batch_info_host_nonzero(self): + dual = self._make_dual_mode() + # num_prefill=2 → not decode-only + batch_info = torch.tensor([2, 100, 3]) + assert dual._is_decode_only(batch_info_host=batch_info) is False + + def test_is_decode_only_fallback_heuristic_decode(self): + dual = self._make_dual_mode() + # No batch_info_host; input_ids shape [4, 1] → decode (seq_dim == 1) + input_ids = torch.randint(0, 100, (4, 1)) + assert dual._is_decode_only(input_ids=input_ids) is True + + def test_is_decode_only_fallback_heuristic_prefill(self): + dual = self._make_dual_mode() + # No batch_info_host; input_ids shape [1, 128] → prefill (seq_dim > 1) + input_ids = torch.randint(0, 100, (1, 128)) + assert dual._is_decode_only(input_ids=input_ids) is False + + def test_is_decode_only_default_no_info(self): + dual = self._make_dual_mode() + # No batch_info_host and no batched inputs → defaults to True + assert dual._is_decode_only() is True + + def test_get_num_tokens_flat_layout(self): + dual = self._make_dual_mode() + input_ids = torch.randint(0, 100, (1, 200)) + assert dual._get_num_tokens(input_ids=input_ids) == 200 + + def test_get_num_tokens_1d_layout(self): + dual = self._make_dual_mode() + input_ids = torch.randint(0, 100, (150,)) + assert dual._get_num_tokens(input_ids=input_ids) == 150 + + def test_get_num_tokens_no_input(self): + dual = self._make_dual_mode() + assert dual._get_num_tokens() == 0 + + @pytest.mark.parametrize( + "num_tokens, expected_bucket", + [ + (10, 64), + (64, 64), + (65, 128), + (128, 128), + (200, 256), + (256, 256), + (257, None), # exceeds largest bucket + ], + ) + def test_find_nearest_bucket(self, num_tokens, expected_bucket): + dual = self._make_dual_mode(piecewise_num_tokens=[64, 128, 256]) + assert dual._find_nearest_bucket(num_tokens) == expected_bucket + + def test_find_nearest_bucket_empty(self): + dual = self._make_dual_mode(piecewise_num_tokens=[]) + assert dual._find_nearest_bucket(100) is None + + +# ============================================================================ +# Tests for PiecewiseCapturedGraph.prepare +# ============================================================================ + + +class TestPiecewiseCapturedGraphPrepare: + """Tests for PiecewiseCapturedGraph.prepare.""" + + def test_non_graphmodule_sets_split_gm_none(self): + """When model is not a GraphModule, split_gm should remain None.""" + model = nn.Linear(4, 4) + pcg = PiecewiseCapturedGraph(model, piecewise_num_tokens=[8, 16]) + pcg.prepare() + + assert pcg._is_prepared is True + assert pcg.split_gm is None + + def test_prepare_is_idempotent(self): + """Calling prepare() twice should not re-split.""" + model = nn.Linear(4, 4) + pcg = PiecewiseCapturedGraph(model, piecewise_num_tokens=[8]) + pcg.prepare() + pcg.prepare() # Should be a no-op + assert pcg._is_prepared is True + + +# ============================================================================ +# Tests for _generate_default_piecewise_num_tokens (compile_model.py) +# ============================================================================ + + +class TestGenerateDefaultPiecewiseNumTokens: + """Tests for _generate_default_piecewise_num_tokens.""" + + def test_power_of_two_max(self): + result = _generate_default_piecewise_num_tokens(8192) + assert result == [64, 128, 256, 512, 1024, 2048, 4096, 8192] + + def test_non_power_of_two_appended(self): + result = _generate_default_piecewise_num_tokens(100) + assert result == [64, 100] + + def test_zero_returns_empty(self): + result = _generate_default_piecewise_num_tokens(0) + assert result == [] + + def test_negative_returns_empty(self): + result = _generate_default_piecewise_num_tokens(-10) + assert result == [] + + def test_exactly_64(self): + result = _generate_default_piecewise_num_tokens(64) + assert result == [64] + + def test_less_than_64(self): + result = _generate_default_piecewise_num_tokens(32) + assert result == [32] + + def test_256(self): + result = _generate_default_piecewise_num_tokens(256) + assert result == [64, 128, 256] + + def test_large_non_power_of_two(self): + result = _generate_default_piecewise_num_tokens(5000) + # Powers of 2 from 64: 64, 128, 256, 512, 1024, 2048, 4096 + # Then append 5000 + assert result == [64, 128, 256, 512, 1024, 2048, 4096, 5000] + + def test_result_is_sorted(self): + result = _generate_default_piecewise_num_tokens(10000) + assert result == sorted(result) + + def test_no_duplicates_when_max_is_power_of_two(self): + result = _generate_default_piecewise_num_tokens(4096) + # 4096 is already a power of 2, should not be duplicated + assert result.count(4096) == 1 diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_piecewise_runner.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_piecewise_runner.py new file mode 100644 index 000000000000..9f11a474ad78 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_piecewise_runner.py @@ -0,0 +1,405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for piecewise_runner: ADPiecewiseRunner and SegmentEntry.""" + +import pytest +import torch +import torch.nn as nn + +from tensorrt_llm._torch.auto_deploy.compile.piecewise_runner import ADPiecewiseRunner, SegmentEntry + +# ============================================================================ +# Context management tests +# ============================================================================ + + +class TestADPiecewiseRunnerContextManagement: + """Tests for class-level context management on ADPiecewiseRunner.""" + + def setup_method(self): + """Reset class-level state before each test.""" + ADPiecewiseRunner._current_num_tokens = None + ADPiecewiseRunner._current_phase = "replay" + ADPiecewiseRunner._static_output_registry.clear() + + def test_set_current_num_tokens(self): + ADPiecewiseRunner.set_current_num_tokens(128) + assert ADPiecewiseRunner._current_num_tokens == 128 + + ADPiecewiseRunner.set_current_num_tokens(None) + assert ADPiecewiseRunner._current_num_tokens is None + + def test_set_current_phase_valid(self): + for phase in ("warmup", "capture", "replay"): + ADPiecewiseRunner.set_current_phase(phase) + assert ADPiecewiseRunner._current_phase == phase + + def test_set_current_phase_invalid_raises(self): + with pytest.raises(AssertionError, match="Invalid phase"): + ADPiecewiseRunner.set_current_phase("invalid_phase") + + def test_clear_static_output_registry(self): + # Populate with some dummy data + t = torch.tensor([1.0]) + ADPiecewiseRunner._static_output_registry[(8, 12345)] = t + assert len(ADPiecewiseRunner._static_output_registry) == 1 + + ADPiecewiseRunner.clear_static_output_registry() + assert len(ADPiecewiseRunner._static_output_registry) == 0 + + +# ============================================================================ +# Initialization tests +# ============================================================================ + + +class TestADPiecewiseRunnerInit: + """Tests for ADPiecewiseRunner initialization.""" + + def test_entries_pre_populated(self): + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8, 16, 32]) + assert set(runner.entries.keys()) == {8, 16, 32} + for entry in runner.entries.values(): + assert isinstance(entry, SegmentEntry) + + def test_weight_ptrs_collected(self): + submod = nn.Linear(4, 4, bias=True) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + # Should have weight and bias data_ptrs + assert submod.weight.data_ptr() in runner._weight_ptrs + assert submod.bias.data_ptr() in runner._weight_ptrs + + def test_no_piecewise_num_tokens(self): + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=None) + assert len(runner.entries) == 0 + + +# ============================================================================ +# _find_entry tests +# ============================================================================ + + +class TestADPiecewiseRunnerFindEntry: + """Tests for _find_entry.""" + + def test_exact_match(self): + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8, 16]) + assert runner._find_entry(8) is not None + assert runner._find_entry(16) is not None + + def test_no_match_returns_none(self): + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8, 16]) + assert runner._find_entry(32) is None + assert runner._find_entry(4) is None + + +# ============================================================================ +# _identify_dynamic_indices tests +# ============================================================================ + + +class TestIdentifyDynamicIndices: + """Tests for _identify_dynamic_indices.""" + + def test_weight_tensors_excluded(self): + submod = nn.Linear(4, 4, bias=False) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + entry = runner.entries[8] + + # flat_args: [weight_tensor, activation_tensor] + weight = submod.weight + activation = torch.randn(2, 4) + flat_args = [weight, activation] + + dynamic = runner._identify_dynamic_indices(entry, flat_args) + assert 0 not in dynamic # weight is not dynamic + assert 1 in dynamic # activation is dynamic + + def test_non_tensor_args_ignored(self): + submod = nn.Linear(4, 4, bias=False) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + entry = runner.entries[8] + + flat_args = [42, "hello", torch.randn(2, 4)] + dynamic = runner._identify_dynamic_indices(entry, flat_args) + # Only the tensor at index 2 should be dynamic + assert dynamic == {2} + + def test_all_activations_marked_dynamic(self): + submod = nn.Linear(4, 4, bias=False) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + entry = runner.entries[8] + + act1 = torch.randn(2, 4) + act2 = torch.randn(3, 4) + flat_args = [act1, act2] + + dynamic = runner._identify_dynamic_indices(entry, flat_args) + assert dynamic == {0, 1} + + +# ============================================================================ +# _track_warmup_ptrs tests +# ============================================================================ + + +class TestTrackWarmupPtrs: + """Tests for _track_warmup_ptrs.""" + + def test_first_call_records_ptrs(self): + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + entry = runner.entries[8] + + t1 = torch.randn(2, 4) + t2 = torch.randn(3, 4) + flat_args = [t1, t2] + + runner._track_warmup_ptrs(entry, flat_args) + assert entry._warmup_data_ptrs is not None + assert len(entry._warmup_data_ptrs) == 2 + assert entry._warmup_data_ptrs[0] == t1.data_ptr() + assert entry._warmup_data_ptrs[1] == t2.data_ptr() + + def test_second_call_detects_ptr_change(self): + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + entry = runner.entries[8] + + t1 = torch.randn(2, 4) + t2 = torch.randn(3, 4) + + # First warmup + runner._track_warmup_ptrs(entry, [t1, t2]) + + # Second warmup with a new tensor at index 0 (simulating changed activation) + t1_new = torch.randn(2, 4) # new tensor, different data_ptr + runner._track_warmup_ptrs(entry, [t1_new, t2]) + + # Index 0 should be marked None (dynamic), index 1 unchanged + assert entry._warmup_data_ptrs[0] is None + assert entry._warmup_data_ptrs[1] == t2.data_ptr() + + def test_non_tensor_args_tracked_as_none(self): + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + entry = runner.entries[8] + + flat_args = [42, torch.randn(2, 4)] + runner._track_warmup_ptrs(entry, flat_args) + assert entry._warmup_data_ptrs[0] is None # int -> None + assert entry._warmup_data_ptrs[1] is not None # tensor -> data_ptr + + +# ============================================================================ +# _prepare_replay_inputs tests +# ============================================================================ + + +class TestPrepareReplayInputs: + """Tests for _prepare_replay_inputs.""" + + def _make_entry_with_dynamic(self, static_inputs, dynamic_indices): + entry = SegmentEntry() + entry.static_inputs = static_inputs + entry.dynamic_indices = dynamic_indices + return entry + + def test_same_shape_copy(self): + """When shapes match, static buffer should be updated.""" + static_buf = torch.zeros(4, 8) + new_inp = torch.ones(4, 8) + entry = self._make_entry_with_dynamic([static_buf], {0}) + + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + runner._prepare_replay_inputs(entry, [new_inp]) + + assert torch.equal(static_buf, new_inp) + + def test_same_data_ptr_skips_copy(self): + """When data_ptr matches (zero-copy path), no copy should happen.""" + shared_tensor = torch.zeros(4, 8) + entry = self._make_entry_with_dynamic([shared_tensor], {0}) + + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + + # Pass the same tensor -- same data_ptr + runner._prepare_replay_inputs(entry, [shared_tensor]) + # Should still be zeros (no copy from a different source) + assert torch.equal(shared_tensor, torch.zeros(4, 8)) + + def test_padded_dim0(self): + """When new_inp is smaller along dim 0, only prefix should be copied.""" + static_buf = torch.zeros(8, 4) + new_inp = torch.ones(5, 4) + entry = self._make_entry_with_dynamic([static_buf], {0}) + + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + runner._prepare_replay_inputs(entry, [new_inp]) + + # First 5 rows should be ones, remaining 3 should be zeros + assert torch.equal(static_buf[:5], torch.ones(5, 4)) + assert torch.equal(static_buf[5:], torch.zeros(3, 4)) + + def test_padded_dim1(self): + """When new_inp is smaller along dim 1, only prefix columns should be copied.""" + static_buf = torch.zeros(1, 16, 4) + new_inp = torch.ones(1, 10, 4) + entry = self._make_entry_with_dynamic([static_buf], {0}) + + submod = nn.Linear(4, 4) + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]) + runner._prepare_replay_inputs(entry, [new_inp]) + + # First 10 along dim 1 should be ones, rest zeros + assert torch.equal(static_buf[:, :10, :], torch.ones(1, 10, 4)) + assert torch.equal(static_buf[:, 10:, :], torch.zeros(1, 6, 4)) + + +# ============================================================================ +# Full cycle tests (CUDA required) +# ============================================================================ + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestADPiecewiseRunnerFullCycle: + """End-to-end warmup -> capture -> replay test on CUDA.""" + + def setup_method(self): + """Reset class-level state before each test.""" + ADPiecewiseRunner._current_num_tokens = None + ADPiecewiseRunner._current_phase = "replay" + ADPiecewiseRunner._static_output_registry.clear() + + def test_warmup_capture_replay_linear(self): + """Full cycle: warmup -> capture -> replay with a simple Linear.""" + device = "cuda" + submod = nn.Linear(16, 16, bias=True).to(device) + submod.eval() + + num_tokens = 8 + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[num_tokens], graph_pool=None).to( + device + ) + + # Fixed input for warmup/capture + x = torch.randn(num_tokens, 16, device=device) + + with torch.inference_mode(): + # --- WARMUP --- + ADPiecewiseRunner.set_current_num_tokens(num_tokens) + ADPiecewiseRunner.set_current_phase("warmup") + for _ in range(3): + _ = runner(x) + + # --- CAPTURE --- + ADPiecewiseRunner.set_current_phase("capture") + _ = runner(x) + + # --- REPLAY --- + ADPiecewiseRunner.set_current_phase("replay") + + # New input for replay + x_new = torch.randn(num_tokens, 16, device=device) + replay_out = runner(x_new) + + # Compare with eager output + eager_out = submod(x_new) + + torch.cuda.synchronize() + assert torch.allclose(replay_out, eager_out, atol=1e-5), ( + "Replay output should match eager output" + ) + + def test_eager_fallback_for_unknown_num_tokens(self): + """Runner should fall back to eager for num_tokens not in entries.""" + device = "cuda" + submod = nn.Linear(16, 16).to(device) + submod.eval() + + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]).to(device) + + x = torch.randn(4, 16, device=device) + + with torch.inference_mode(): + # num_tokens=4 is not configured -- should fall back to eager + ADPiecewiseRunner.set_current_num_tokens(4) + ADPiecewiseRunner.set_current_phase("replay") + out = runner(x) + + eager_out = submod(x) + assert torch.allclose(out, eager_out, atol=1e-6) + + def test_eager_fallback_for_none_num_tokens(self): + """Runner should fall back to eager when num_tokens is None.""" + device = "cuda" + submod = nn.Linear(16, 16).to(device) + submod.eval() + + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=[8]).to(device) + + x = torch.randn(4, 16, device=device) + + with torch.inference_mode(): + ADPiecewiseRunner.set_current_num_tokens(None) + ADPiecewiseRunner.set_current_phase("replay") + out = runner(x) + + eager_out = submod(x) + assert torch.allclose(out, eager_out, atol=1e-6) + + def test_multiple_bucket_sizes(self): + """Capture and replay with multiple bucket sizes.""" + device = "cuda" + submod = nn.Linear(16, 16).to(device) + submod.eval() + + buckets = [4, 8, 16] + runner = ADPiecewiseRunner(submod, piecewise_num_tokens=buckets).to(device) + + with torch.inference_mode(): + for nt in buckets: + x = torch.randn(nt, 16, device=device) + + ADPiecewiseRunner.set_current_num_tokens(nt) + + # Warmup + ADPiecewiseRunner.set_current_phase("warmup") + for _ in range(3): + runner(x) + + # Capture + ADPiecewiseRunner.set_current_phase("capture") + runner(x) + + # Replay each bucket + ADPiecewiseRunner.set_current_phase("replay") + for nt in buckets: + x_new = torch.randn(nt, 16, device=device) + ADPiecewiseRunner.set_current_num_tokens(nt) + replay_out = runner(x_new) + eager_out = submod(x_new) + + torch.cuda.synchronize() + assert torch.allclose(replay_out, eager_out, atol=1e-5), ( + f"Replay mismatch for bucket {nt}" + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_piecewise_utils.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_piecewise_utils.py new file mode 100644 index 000000000000..f8e88f5895af --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/compile/test_piecewise_utils.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for piecewise_utils: is_dynamic_cached_op and split_graph_at_dynamic_ops.""" + +from types import SimpleNamespace + +import torch +import torch.nn as nn +from torch.fx import Graph, GraphModule + +from tensorrt_llm._torch.auto_deploy.compile.piecewise_utils import ( + _CACHED_ATTENTION_OPS, + _CACHED_CONV_OPS, + _CACHED_DELTA_OPS, + _CACHED_SSM_OPS, + _LOGITS_GATHER_OPS, + _METADATA_PREP_OPS, + _get_all_dynamic_op_names, + is_dynamic_cached_op, + split_graph_at_dynamic_ops, +) + +# ============================================================================ +# Helpers +# ============================================================================ + + +def _make_mock_node(op: str, target=None): + """Create a lightweight mock FX Node for testing is_dynamic_cached_op.""" + node = SimpleNamespace(op=op, target=target) + return node + + +class _FakeOpOverload: + """Mimics torch._ops.OpOverload with a .name() method. + + Must be callable with __name__/__module__/__qualname__ because torch.fx + validates call_function targets and generates Python code referencing them. + """ + + def __init__(self, qualified_name: str): + self._name = qualified_name + # Attributes required by torch.fx for codegen + short = qualified_name.split("::")[-1] + self.__name__ = short + self.__qualname__ = qualified_name + self.__module__ = "test_piecewise_utils" + + def name(self): + return self._name + + def __call__(self, *args, **kwargs): + # Identity pass-through for graph execution + return args[0] if args else None + + +def _build_graphmodule_with_ops(dynamic_op_names=None): + """Build a simple FX GraphModule with relu ops interspersed with fake dynamic ops. + + The graph looks like: x -> relu -> [dyn_op_0] -> relu -> [dyn_op_1] -> ... -> output + Dynamic ops are simulated by inserting call_function nodes whose target is a + _FakeOpOverload with a name matching one of the dynamic op registries. + """ + if dynamic_op_names is None: + dynamic_op_names = [] + + # Build graph manually + graph = Graph() + x = graph.placeholder("x") + + # First static op: relu + relu_node = graph.call_function(torch.relu, args=(x,)) + prev = relu_node + + for idx, dyn_name in enumerate(dynamic_op_names): + # Insert a fake dynamic op. We use graph.create_node directly because + # graph.call_function tries _target_to_str which asserts isinstance(target, str). + fake_target = _FakeOpOverload(dyn_name) + dyn_node = graph.create_node( + "call_function", fake_target, args=(prev,), name=f"dyn_op_{idx}" + ) + # Follow with another static op + relu_after = graph.call_function(torch.relu, args=(dyn_node,)) + prev = relu_after + + graph.output(prev) + + # We need a root module -- a simple nn.Module suffices + root = nn.Module() + gm = GraphModule(root, graph) + return gm + + +# ============================================================================ +# Tests for is_dynamic_cached_op +# ============================================================================ + + +class TestIsDynamicCachedOp: + """Tests for is_dynamic_cached_op.""" + + def test_known_attention_op_returns_true(self): + target = _FakeOpOverload("auto_deploy::flashinfer_attention_mha_with_cache") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + + def test_known_ssm_op_returns_true(self): + target = _FakeOpOverload("auto_deploy::triton_cached_ssm") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + + def test_known_conv_op_returns_true(self): + target = _FakeOpOverload("auto_deploy::triton_cached_causal_conv1d") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + + def test_known_delta_op_returns_true(self): + target = _FakeOpOverload("auto_deploy::fla_cached_delta_rule") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + + def test_known_metadata_prep_op_returns_true(self): + target = _FakeOpOverload("auto_deploy::flashinfer_attention_prepare_metadata") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + + def test_known_logits_gather_op_returns_true(self): + target = _FakeOpOverload("auto_deploy::gather_logits_before_lm_head") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + + def test_static_op_returns_false(self): + # torch.relu is not a dynamic op + node = _make_mock_node("call_function", target=torch.relu) + assert is_dynamic_cached_op(node) is False + + def test_non_call_function_returns_false(self): + target = _FakeOpOverload("auto_deploy::flashinfer_attention_mha_with_cache") + # Even with a dynamic target, non-call_function ops return False + for op_type in ("placeholder", "call_method", "call_module", "output", "get_attr"): + node = _make_mock_node(op_type, target=target) + assert is_dynamic_cached_op(node) is False, f"Should be False for op={op_type}" + + def test_op_with_default_suffix_still_matches(self): + """Dynamic op name with .default suffix should still match (substring check).""" + target = _FakeOpOverload("auto_deploy::triton_cached_ssm.default") + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True + + def test_all_registry_entries_recognized(self): + """Every op in every registry list should be recognized as dynamic.""" + all_ops = ( + _CACHED_ATTENTION_OPS + + _CACHED_SSM_OPS + + _CACHED_CONV_OPS + + _CACHED_DELTA_OPS + + _METADATA_PREP_OPS + + _LOGITS_GATHER_OPS + ) + for op_name in all_ops: + target = _FakeOpOverload(op_name) + node = _make_mock_node("call_function", target=target) + assert is_dynamic_cached_op(node) is True, f"{op_name} should be recognized as dynamic" + + def test_get_all_dynamic_op_names_returns_full_set(self): + all_names = _get_all_dynamic_op_names() + assert isinstance(all_names, set) + # Should include all registries + for op in _CACHED_ATTENTION_OPS: + assert op in all_names + for op in _CACHED_SSM_OPS: + assert op in all_names + for op in _LOGITS_GATHER_OPS: + assert op in all_names + + +# ============================================================================ +# Tests for split_graph_at_dynamic_ops +# ============================================================================ + + +class TestSplitGraphAtDynamicOps: + """Tests for split_graph_at_dynamic_ops.""" + + def test_no_dynamic_ops_returns_original(self): + """Graph with no dynamic ops should not be split.""" + gm = _build_graphmodule_with_ops(dynamic_op_names=[]) + info = split_graph_at_dynamic_ops(gm) + + assert info.num_submodules == 1 + assert info.dynamic_submod_indices == [] + assert info.static_submod_indices == [0] + # split_gm is the original gm + assert info.split_gm is gm + + def test_single_dynamic_op_produces_3_submodules(self): + """One dynamic op should produce 3 partitions: static -> dynamic -> static.""" + gm = _build_graphmodule_with_ops( + dynamic_op_names=["auto_deploy::flashinfer_attention_mha_with_cache"] + ) + info = split_graph_at_dynamic_ops(gm) + + # Expected: submod_0 (static: relu), submod_1 (dynamic: attn), submod_2 (static: relu) + assert info.num_submodules == 3 + assert len(info.dynamic_submod_indices) == 1 + assert len(info.static_submod_indices) == 2 + + def test_two_dynamic_ops_produces_5_submodules(self): + """Two dynamic ops → 5 partitions: S D S D S.""" + gm = _build_graphmodule_with_ops( + dynamic_op_names=[ + "auto_deploy::flashinfer_attention_mha_with_cache", + "auto_deploy::triton_cached_ssm", + ] + ) + info = split_graph_at_dynamic_ops(gm) + + assert info.num_submodules == 5 + assert len(info.dynamic_submod_indices) == 2 + assert len(info.static_submod_indices) == 3 + + def test_dynamic_and_static_indices_are_disjoint(self): + """Dynamic and static indices should not overlap and should cover all submodules.""" + gm = _build_graphmodule_with_ops( + dynamic_op_names=[ + "auto_deploy::flashinfer_attention_mha_with_cache", + "auto_deploy::triton_cached_ssm", + ] + ) + info = split_graph_at_dynamic_ops(gm) + + all_indices = set(info.dynamic_submod_indices) | set(info.static_submod_indices) + assert len(all_indices) == info.num_submodules + # No overlap + assert len(set(info.dynamic_submod_indices) & set(info.static_submod_indices)) == 0 + + def test_split_submodules_are_named_correctly(self): + """Split submodules should be named submod_0, submod_1, etc.""" + gm = _build_graphmodule_with_ops( + dynamic_op_names=["auto_deploy::triton_cached_causal_conv1d"] + ) + info = split_graph_at_dynamic_ops(gm) + + for i in range(info.num_submodules): + assert hasattr(info.split_gm, f"submod_{i}"), f"Missing submod_{i}" diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py new file mode 100644 index 000000000000..71daa139a279 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py @@ -0,0 +1,601 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the TRT-LLM attention backend custom op. + +Mirrors the structure of test_flashinfer_attention_op.py but exercises the +``trtllm_attention_mha_with_cache`` custom op via ``thop.attention``. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch.auto_deploy.custom_ops.attention.trtllm_attention import ( + _GlobalTrtllmPlanner, + prepare_trtllm_metadata_host, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _reset_trtllm_planner(): + """Force a full reset of the global TRT-LLM planner so buffers are re-allocated.""" + _GlobalTrtllmPlanner.__init__() + + +def _prepare_and_run( + q, + k, + v, + kv_cache, + seq_lens, + input_positions, + cache_locs, + pages_per_seq, + max_seq_len, + max_batch_size, + num_prefill, + num_prefill_tokens, + num_decode, + device, + scale=None, +): + """Build metadata, call host-prepare, and invoke the TRT-LLM attention op.""" + tokens_per_block = kv_cache.shape[3] + max_blocks_per_seq = math.ceil(max_seq_len / tokens_per_block) + block_offset_multiplier = kv_cache.stride(0) // kv_cache.stride(1) + + seq_len_with_cache = [ip + sl for ip, sl in zip(input_positions, seq_lens)] + + # --- tensors for the op --------------------------------------------------- + batch_info_host = torch.tensor( + [num_prefill, num_prefill_tokens, num_decode], dtype=torch.int32, device=device + ) + seq_len_d = torch.tensor(seq_lens, dtype=torch.int32, device=device) + seq_len_h = torch.tensor(seq_lens, dtype=torch.int32).pin_memory() + input_pos_h = torch.tensor(input_positions, dtype=torch.int32).pin_memory() + slwc_d = torch.tensor(seq_len_with_cache, dtype=torch.int32, device=device) + slwc_h = torch.tensor(seq_len_with_cache, dtype=torch.int32).pin_memory() + max_seq_info_h = torch.tensor( + [max_seq_len, max_blocks_per_seq, block_offset_multiplier, max_batch_size], + dtype=torch.int32, + ).pin_memory() + + # --- paging metadata for host prepare ------------------------------------- + cache_loc_d = torch.tensor(cache_locs, dtype=torch.int32, device=device) + + cu_num_pages = [0] + for pps in pages_per_seq: + cu_num_pages.append(cu_num_pages[-1] + pps) + cu_num_pages_h = torch.tensor(cu_num_pages, dtype=torch.int32).pin_memory() + + page_seq_indices_list = [] + page_in_seq_list = [] + for i, np_ in enumerate(pages_per_seq): + page_seq_indices_list.extend([i] * np_) + page_in_seq_list.extend(range(np_)) + page_seq_indices_d = torch.tensor(page_seq_indices_list, dtype=torch.int32, device=device) + page_in_seq_d = torch.tensor(page_in_seq_list, dtype=torch.int32, device=device) + + # --- host prepare --------------------------------------------------------- + prepare_trtllm_metadata_host( + batch_info_host=batch_info_host.cpu(), + max_seq_info_host=max_seq_info_h, + seq_len_with_cache_host=slwc_h, + cu_num_pages_host=cu_num_pages_h, + cache_loc=cache_loc_d, + page_seq_indices=page_seq_indices_d, + page_in_seq=page_in_seq_d, + input_pos_host=input_pos_h, + seq_len_host=seq_len_h, + ) + + # --- call op -------------------------------------------------------------- + return torch.ops.auto_deploy.trtllm_attention_mha_with_cache( + q, + k, + v, + batch_info_host, + seq_len_d, + seq_len_h, + input_pos_h, + slwc_d, + max_seq_info_h, + kv_cache, + scale, + ) + + +# --------------------------------------------------------------------------- +# Test 1 – Context (prefill) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("seq_length", [8, 32, 2048]) +@pytest.mark.parametrize("n_heads", [8]) +@pytest.mark.parametrize("batch_size", [1, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_trtllm_attention_op_context(seq_length, n_heads, batch_size, dtype, device): + D_HEAD = 64 + MAX_SEQ_LEN = 2048 + MAX_BATCH_SIZE = 32 + + _reset_trtllm_planner() + + # Q, K, V in bsnd layout + q = torch.randn(batch_size, seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + k = torch.randn(batch_size, seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + v = torch.randn(batch_size, seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + + # KV cache – HND: [num_blocks, 2, num_heads, tokens_per_block, head_dim] + # Unpaged: 1 block per sequence, tokens_per_block = MAX_SEQ_LEN + kv_cache = torch.zeros( + MAX_BATCH_SIZE, 2, n_heads, MAX_SEQ_LEN, D_HEAD, dtype=dtype, device=device + ) + + output = _prepare_and_run( + q, + k, + v, + kv_cache, + seq_lens=[seq_length] * batch_size, + input_positions=[0] * batch_size, + cache_locs=list(range(batch_size)), + pages_per_seq=[1] * batch_size, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=batch_size, + num_prefill_tokens=batch_size * seq_length, + num_decode=0, + device=device, + ) + + # Reference: SDPA causal + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True + ).transpose(1, 2) + + assert torch.allclose( + output.cpu().to(torch.float32), ref.cpu().to(torch.float32), atol=1e-2, rtol=1e-2 + ) + + +# --------------------------------------------------------------------------- +# Test 2 – Decode (generate) with pre-filled cache +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("prefill_seq_length", [1, 4, 2047]) +@pytest.mark.parametrize("n_heads", [8]) +@pytest.mark.parametrize("batch_size", [1, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_trtllm_attention_op_decode(prefill_seq_length, batch_size, n_heads, dtype, device): + D_HEAD = 64 + MAX_SEQ_LEN = 2048 + MAX_BATCH_SIZE = 32 + + _reset_trtllm_planner() + + # --- Step 1: prefill to populate the cache -------------------------------- + q_pf = torch.randn(batch_size, prefill_seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + k_pf = torch.randn(batch_size, prefill_seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + v_pf = torch.randn(batch_size, prefill_seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + + kv_cache = torch.zeros( + MAX_BATCH_SIZE, 2, n_heads, MAX_SEQ_LEN, D_HEAD, dtype=dtype, device=device + ) + + _prepare_and_run( + q_pf, + k_pf, + v_pf, + kv_cache, + seq_lens=[prefill_seq_length] * batch_size, + input_positions=[0] * batch_size, + cache_locs=list(range(batch_size)), + pages_per_seq=[1] * batch_size, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=batch_size, + num_prefill_tokens=batch_size * prefill_seq_length, + num_decode=0, + device=device, + ) + + # --- Step 2: decode one token --------------------------------------------- + q_dec = torch.randn(batch_size, 1, n_heads, D_HEAD, dtype=dtype, device=device) + k_dec = torch.randn(batch_size, 1, n_heads, D_HEAD, dtype=dtype, device=device) + v_dec = torch.randn(batch_size, 1, n_heads, D_HEAD, dtype=dtype, device=device) + + _reset_trtllm_planner() + + output = _prepare_and_run( + q_dec, + k_dec, + v_dec, + kv_cache, + seq_lens=[1] * batch_size, + input_positions=[prefill_seq_length] * batch_size, + cache_locs=list(range(batch_size)), + pages_per_seq=[1] * batch_size, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=0, + num_prefill_tokens=0, + num_decode=batch_size, + device=device, + ) + + # Reference: Q_dec attends to full K, V (prefill + new) + k_full = torch.cat([k_pf, k_dec], dim=1) + v_full = torch.cat([v_pf, v_dec], dim=1) + ref = torch.nn.functional.scaled_dot_product_attention( + q_dec.transpose(1, 2), + k_full.transpose(1, 2), + v_full.transpose(1, 2), + is_causal=False, + ).transpose(1, 2) + + assert torch.allclose( + output.cpu().to(torch.float32), ref.cpu().to(torch.float32), atol=1e-2, rtol=1e-2 + ) + + +# --------------------------------------------------------------------------- +# Test 3 – Context then Generate (full cycle, verify both outputs) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("prefill_seq_length", [4, 10, 2047]) +@pytest.mark.parametrize("n_heads", [8]) +@pytest.mark.parametrize("batch_size", [1, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_trtllm_attention_context_and_generate( + prefill_seq_length, n_heads, batch_size, dtype, device +): + D_HEAD = 64 + MAX_SEQ_LEN = 2048 + MAX_BATCH_SIZE = 32 + + _reset_trtllm_planner() + + # --- Prefill -------------------------------------------------------------- + q_pf = torch.randn(batch_size, prefill_seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + k_pf = torch.randn(batch_size, prefill_seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + v_pf = torch.randn(batch_size, prefill_seq_length, n_heads, D_HEAD, dtype=dtype, device=device) + + kv_cache = torch.zeros( + MAX_BATCH_SIZE, 2, n_heads, MAX_SEQ_LEN, D_HEAD, dtype=dtype, device=device + ) + + output_pf = _prepare_and_run( + q_pf, + k_pf, + v_pf, + kv_cache, + seq_lens=[prefill_seq_length] * batch_size, + input_positions=[0] * batch_size, + cache_locs=list(range(batch_size)), + pages_per_seq=[1] * batch_size, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=batch_size, + num_prefill_tokens=batch_size * prefill_seq_length, + num_decode=0, + device=device, + ) + + # Verify prefill output + ref_pf = torch.nn.functional.scaled_dot_product_attention( + q_pf.transpose(1, 2), k_pf.transpose(1, 2), v_pf.transpose(1, 2), is_causal=True + ).transpose(1, 2) + + assert torch.allclose( + output_pf.cpu().to(torch.float32), ref_pf.cpu().to(torch.float32), atol=1e-2, rtol=1e-2 + ) + + # --- Generate one token --------------------------------------------------- + q_gen = torch.randn(batch_size, 1, n_heads, D_HEAD, dtype=dtype, device=device) + k_gen = torch.randn(batch_size, 1, n_heads, D_HEAD, dtype=dtype, device=device) + v_gen = torch.randn(batch_size, 1, n_heads, D_HEAD, dtype=dtype, device=device) + + _reset_trtllm_planner() + + output_gen = _prepare_and_run( + q_gen, + k_gen, + v_gen, + kv_cache, + seq_lens=[1] * batch_size, + input_positions=[prefill_seq_length] * batch_size, + cache_locs=list(range(batch_size)), + pages_per_seq=[1] * batch_size, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=0, + num_prefill_tokens=0, + num_decode=batch_size, + device=device, + ) + + # Verify decode output – Q_gen attends to full (prefill + gen) K, V + k_full = torch.cat([k_pf, k_gen], dim=1) + v_full = torch.cat([v_pf, v_gen], dim=1) + ref_gen = torch.nn.functional.scaled_dot_product_attention( + q_gen.transpose(1, 2), + k_full.transpose(1, 2), + v_full.transpose(1, 2), + is_causal=False, + ).transpose(1, 2) + + assert torch.allclose( + output_gen.cpu().to(torch.float32), ref_gen.cpu().to(torch.float32), atol=1e-2, rtol=1e-2 + ) + + +# --------------------------------------------------------------------------- +# Test 4 – Context with non-zero input_pos (chunked prefill) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "seq", + [ + (2, 1), + (2, 64), + (2, 2046), + (16, 1), + (16, 2022), + (1984, 64), + (1024, 1024), + ], +) +@pytest.mark.parametrize("n_heads", [8]) +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_trtllm_attention_op_context_input_pos(seq, batch_size, n_heads, dtype, device): + D_HEAD = 64 + MAX_SEQ_LEN = 2048 + MAX_BATCH_SIZE = 32 + SEQ_LEN = seq[0] + PREFILL_SEQ_LEN = seq[1] + + _reset_trtllm_planner() + + # --- Step 1: prefill the first chunk to populate cache -------------------- + q_1 = torch.randn(batch_size, PREFILL_SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + k_1 = torch.randn(batch_size, PREFILL_SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + v_1 = torch.randn(batch_size, PREFILL_SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + + kv_cache = torch.zeros( + MAX_BATCH_SIZE, 2, n_heads, MAX_SEQ_LEN, D_HEAD, dtype=dtype, device=device + ) + + _prepare_and_run( + q_1, + k_1, + v_1, + kv_cache, + seq_lens=[PREFILL_SEQ_LEN] * batch_size, + input_positions=[0] * batch_size, + cache_locs=list(range(batch_size)), + pages_per_seq=[1] * batch_size, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=batch_size, + num_prefill_tokens=batch_size * PREFILL_SEQ_LEN, + num_decode=0, + device=device, + ) + + # --- Step 2: second context chunk at non-zero input_pos ------------------- + q_2 = torch.randn(batch_size, SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + k_2 = torch.randn(batch_size, SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + v_2 = torch.randn(batch_size, SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + + _reset_trtllm_planner() + + output = _prepare_and_run( + q_2, + k_2, + v_2, + kv_cache, + seq_lens=[SEQ_LEN] * batch_size, + input_positions=[PREFILL_SEQ_LEN] * batch_size, + cache_locs=list(range(batch_size)), + pages_per_seq=[1] * batch_size, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=batch_size, + num_prefill_tokens=batch_size * SEQ_LEN, + num_decode=0, + device=device, + ) + + # NOTE: "chunked prefill" (context tokens with non-zero `input_pos`) can behave + # differently depending on which TRT-LLM kernel path is selected. + # + # - On some paths (notably SM100 fallback), the context-stage output corresponds to + # causal attention over the *current chunk only* (cached prefix is not attended), + # though the KV cache is still updated. + # - On others (e.g. SM80/A30), the output incorporates the cached prefix and matches + # full causal attention over (prefix + current chunk) with the appropriate offset. + # + # We accept either behavior here, but still require the output to match one of the + # two well-defined PyTorch SDPA references. + ref_chunk_only = torch.nn.functional.scaled_dot_product_attention( + q_2.transpose(1, 2), + k_2.transpose(1, 2), + v_2.transpose(1, 2), + is_causal=True, + ).transpose(1, 2) + + out_f32 = output.cpu().to(torch.float32) + if torch.allclose(out_f32, ref_chunk_only.cpu().to(torch.float32), atol=1e-2, rtol=1e-2): + return + + # Full reference: Q_2 attends to full K (chunk1 + chunk2) with causal mask offset. + k_full = torch.cat([k_1, k_2], dim=1) + v_full = torch.cat([v_1, v_2], dim=1) + mask = torch.cat( + [ + torch.ones(SEQ_LEN, PREFILL_SEQ_LEN, device=device, dtype=torch.bool), + torch.tril(torch.ones(SEQ_LEN, SEQ_LEN, device=device, dtype=torch.bool)), + ], + dim=1, + ) + ref_with_prefix = torch.nn.functional.scaled_dot_product_attention( + q_2.transpose(1, 2), + k_full.transpose(1, 2), + v_full.transpose(1, 2), + attn_mask=mask, + ).transpose(1, 2) + + assert torch.allclose(out_f32, ref_with_prefix.cpu().to(torch.float32), atol=1e-2, rtol=1e-2) + + +# --------------------------------------------------------------------------- +# Test 5 – Paged KV cache (prefill + decode with small page_size) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("seq_lengths", [[8, 14], [11, 19, 22, 49]]) +@pytest.mark.parametrize("n_heads", [8]) +@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_trtllm_attention_with_paged_kvcache(seq_lengths, n_heads, dtype, device): + PAGE_SIZE = 8 + D_HEAD = 64 + MAX_SEQ_LEN = 128 + MAX_BATCH_SIZE = 32 + BATCH_SIZE = len(seq_lengths) + TOTAL_SEQ_LEN = sum(seq_lengths) + + MAX_NUM_PAGES = MAX_BATCH_SIZE * MAX_SEQ_LEN // PAGE_SIZE + + _reset_trtllm_planner() + + # Q, K, V – flattened across batch: [1, total_tokens, n_heads, d_head] + q = torch.randn(1, TOTAL_SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + k = torch.randn(1, TOTAL_SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + v = torch.randn(1, TOTAL_SEQ_LEN, n_heads, D_HEAD, dtype=dtype, device=device) + + # Paged KV cache – HND: [num_pages, 2, n_heads, page_size, d_head] + kv_cache = torch.zeros(MAX_NUM_PAGES, 2, n_heads, PAGE_SIZE, D_HEAD, dtype=dtype, device=device) + + # Assign pages randomly + free_pages = torch.randperm(MAX_NUM_PAGES).int().tolist() + pages_per_seq_list = [math.ceil(s / PAGE_SIZE) for s in seq_lengths] + page_assignments = [[free_pages.pop() for _ in range(np)] for np in pages_per_seq_list] + cache_locs = [p for ps in page_assignments for p in ps] + + # Prefill + output_pf = _prepare_and_run( + q, + k, + v, + kv_cache, + seq_lens=seq_lengths, + input_positions=[0] * BATCH_SIZE, + cache_locs=cache_locs, + pages_per_seq=pages_per_seq_list, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=BATCH_SIZE, + num_prefill_tokens=TOTAL_SEQ_LEN, + num_decode=0, + device=device, + ) + + # Reference – per-sequence causal SDPA + cu = [0] + list(torch.cumsum(torch.tensor(seq_lengths), 0).tolist()) + ref_parts = [] + for i, s in enumerate(seq_lengths): + qq = q[0, cu[i] : cu[i + 1], :, :].unsqueeze(0) + kk = k[0, cu[i] : cu[i + 1], :, :].unsqueeze(0) + vv = v[0, cu[i] : cu[i + 1], :, :].unsqueeze(0) + oo = torch.nn.functional.scaled_dot_product_attention( + qq.transpose(1, 2), kk.transpose(1, 2), vv.transpose(1, 2), is_causal=True + ).transpose(1, 2) + ref_parts.append(oo.squeeze(0)) + ref = torch.cat(ref_parts, dim=0) + + assert torch.allclose( + output_pf.squeeze(0).cpu().to(torch.float32), + ref.cpu().to(torch.float32), + atol=1e-2, + rtol=1e-2, + ) + + # --- Now generate one token per sequence ---------------------------------- + _reset_trtllm_planner() + + q_gen = torch.randn(BATCH_SIZE, 1, n_heads, D_HEAD, dtype=dtype, device=device) + k_gen = torch.randn(BATCH_SIZE, 1, n_heads, D_HEAD, dtype=dtype, device=device) + v_gen = torch.randn(BATCH_SIZE, 1, n_heads, D_HEAD, dtype=dtype, device=device) + + # Update page assignments – may need a new page for sequences whose last page is full + for i, (pages, s) in enumerate(zip(page_assignments, seq_lengths)): + last_page_occupancy = s % PAGE_SIZE + if last_page_occupancy == 0: # last page was full → need a new page + pages.append(free_pages.pop()) + pages_per_seq_list[i] = len(pages) + cache_locs_gen = [p for ps in page_assignments for p in ps] + + output_gen = _prepare_and_run( + q_gen, + k_gen, + v_gen, + kv_cache, + seq_lens=[1] * BATCH_SIZE, + input_positions=seq_lengths, + cache_locs=cache_locs_gen, + pages_per_seq=pages_per_seq_list, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + num_prefill=0, + num_prefill_tokens=0, + num_decode=BATCH_SIZE, + device=device, + ) + + # Reference for decode + ref_gen_parts = [] + for i, s in enumerate(seq_lengths): + qq = q_gen[i : i + 1, :, :, :] # [1, 1, n_heads, d_head] + kk = k[0, cu[i] : cu[i + 1], :, :].unsqueeze(0) # [1, s, n_heads, d_head] + kk = torch.cat([kk, k_gen[i : i + 1, :, :, :]], dim=1) # [1, s+1, ...] + vv = v[0, cu[i] : cu[i + 1], :, :].unsqueeze(0) + vv = torch.cat([vv, v_gen[i : i + 1, :, :, :]], dim=1) + oo = torch.nn.functional.scaled_dot_product_attention( + qq.transpose(1, 2), kk.transpose(1, 2), vv.transpose(1, 2), is_causal=False + ).transpose(1, 2) + ref_gen_parts.append(oo.squeeze(0)) + ref_gen = torch.cat(ref_gen_parts, dim=0) + + assert torch.allclose( + output_gen.reshape(-1, n_heads, D_HEAD).cpu().to(torch.float32), + ref_gen.cpu().to(torch.float32), + atol=1e-2, + rtol=1e-2, + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py index f7a8a5972e36..d9879562e8e3 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py @@ -4,6 +4,9 @@ import torch from _custom_op_utils import torch_rope_reference +# Import after we've imported torch (to ensure custom ops are registered) +from tensorrt_llm._torch.auto_deploy.custom_ops.rope import triton_rope # noqa: F401 + def _precompute_freqs_cis( seq_len: int, head_dim: int, rope_theta: Optional[float] = None diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream.py deleted file mode 100644 index 182c6a0a31bd..000000000000 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream.py +++ /dev/null @@ -1,134 +0,0 @@ -from typing import Tuple - -import torch -import torch.nn as nn -from torch.fx import GraphModule, Node - -from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_moe import ( - aux_stream_wrapper, - cuda_stream_manager, - record_event_wrapper, -) -from tensorrt_llm._torch.auto_deploy.utils._graph import canonicalize_graph -from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op - - -@torch.library.custom_op("auto_deploy::multi_stream_linear", mutates_args=()) -def multi_stream_linear( - input: torch.Tensor, weight0: torch.Tensor, weight1: torch.Tensor -) -> torch.Tensor: - output = torch.ops.aten.linear(input, weight0) - output = torch.ops.aten.linear(output, weight1) - return output - - -@multi_stream_linear.register_fake -def multi_stream_linear_fake(input, weight0, weight1): - """Fake implementation of multi_stream_linear.""" - output = torch.ops.aten.linear(input, weight0) - return torch.ops.aten.linear(output, weight1) - - -def replace_multi_stream_linear_with_aux_stream_wrapper(gm: GraphModule) -> Tuple[GraphModule, int]: - """Traverse ``gm`` and replace all ``auto_deploy::multi_stream_linear`` ops with ``aux_stream_wrapper``. - - The replacement preserves the original args/kwargs of the node. - After rewriting, the graph is cleaned and recompiled. - - Args: - gm: The FX graph module to transform. - aux_stream_wrapper: A callable to replace the custom op with. - - Returns: - A tuple of (gm, num_replaced) - """ - graph = gm.graph - num_replaced = 0 - - # Collect targets first to avoid mutating while iterating - target_nodes: list[Node] = [] - target_nodes = [n for n in graph.nodes if is_op(n, torch.ops.auto_deploy.multi_stream_linear)] - - for n in target_nodes: - target_input_node = None - for input_node in n.all_input_nodes: - if len(input_node.users) > 1: - target_input_node = input_node - break - if target_input_node is None: - raise ValueError(f"Target input node not found for node {n}") - with graph.inserting_before(target_input_node): - kwargs = target_input_node.kwargs.copy() - kwargs["device"] = torch.cuda.current_device() - new_node = graph.call_function( - record_event_wrapper, - args=(target_input_node.target, *target_input_node.args), - kwargs=kwargs, - ) - target_input_node.replace_all_uses_with(new_node) - graph.erase_node(target_input_node) - with graph.inserting_after(n): - new_node = graph.call_function( - aux_stream_wrapper, args=(n.target, *n.args), kwargs=n.kwargs - ) - n.replace_all_uses_with(new_node) - graph.erase_node(n) - num_replaced += 1 - - if num_replaced: - canonicalize_graph(gm) - - return gm, num_replaced - - -class ParallelTwoLinear(nn.Module): - def __init__(self, in_dim: int, out_dim: int): - super().__init__() - self.fc10 = nn.Linear(in_dim, in_dim) - self.fc11 = nn.Linear(in_dim, out_dim) - self.fc2 = nn.Linear(in_dim, out_dim) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = torch.nn.functional.relu(x) - y0 = self.fc2(x) - y1 = torch.ops.auto_deploy.multi_stream_linear(x, self.fc10.weight, self.fc11.weight) - return y0 + y1 - - -def test_multi_stream_linear(): - in_dim, out_dim = 128, 256 - cuda_stream_manager.add_device(torch.cuda.current_device()) - model = ( - nn.Sequential(ParallelTwoLinear(in_dim, out_dim), ParallelTwoLinear(out_dim, out_dim)) - .eval() - .to("cuda") - ) - - # Example input used for export - example_input = torch.randn(4, in_dim).to("cuda") - - # Export the graph - egm = torch.export.export(model, (example_input,)) - gm = egm.module() - - test_x = torch.randn(4, in_dim).to("cuda") - ref_output = model(test_x) - - # pattern matching and replace - gm, num_replaced = replace_multi_stream_linear_with_aux_stream_wrapper(gm) - - assert num_replaced == 2 - y = gm(test_x) - assert torch.allclose(y, ref_output) - - static_x = torch.randn(4, in_dim).to("cuda") - static_output = torch.randn(4, out_dim).to("cuda") - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - static_output.copy_(gm(static_x)) - - static_x.copy_(test_x) - graph.replay() - - assert torch.allclose(static_output, ref_output) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py new file mode 100644 index 000000000000..f64fb1ebaf26 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py @@ -0,0 +1,230 @@ +"""Tests for multi-stream Q/KV projection parallelism in MLA attention. + +The test builds a minimal mock model that mirrors the MLA fork pattern: +a shared input feeds two parallel linear chains (one heavier "Q-like", +one lighter "KV-like") whose outputs are combined with an add. + +The transform should: + 1. Detect the fork point (shared input with 2+ linear users). + 2. Identify the lighter KV-like linear (no downstream linear within + a few hops) vs. the heavier Q-like chain (has a downstream linear). + 3. Move the KV linear onto the auxiliary CUDA stream. + 4. Preserve numerical correctness. + 5. Be compatible with CUDA graph capture & replay. +""" + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_attn import ( + _execute_kv_proj_in_aux_stream, + _find_kv_proj_linears, +) +from tensorrt_llm._torch.auto_deploy.utils.multi_stream_utils import cuda_stream_manager + +# --------------------------------------------------------------------------- +# Helpers -- mock MLA-like module +# --------------------------------------------------------------------------- + + +class MockMLABlock(nn.Module): + """Simplified MLA-like attention block with Q and KV projection chains. + + Q chain (heavier): q_a_proj -> relu (stand-in for rms_norm) -> q_b_proj + KV chain (lighter): kv_a_proj + Merge: add(q_b_proj_output, kv_a_proj_output) + + The layernorm at the output simulates the inter-layer distance in a real + transformer (output projection, residual add, layernorm) so that the + next layer's fork point is beyond the BFS max_depth from this layer's + KV linear. + """ + + def __init__(self, hidden_dim: int, q_inner_dim: int, kv_out_dim: int): + super().__init__() + # Q chain: two linears with a non-linearity in between + self.q_a_proj = nn.Linear(hidden_dim, q_inner_dim, bias=False) + self.q_b_proj = nn.Linear(q_inner_dim, kv_out_dim, bias=False) + # KV chain: single linear + self.kv_a_proj = nn.Linear(hidden_dim, kv_out_dim, bias=False) + # Inter-layer distance (layernorm + relu simulate residual + norm) + self.layernorm = nn.LayerNorm(kv_out_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Q chain: q_a_proj -> relu -> q_b_proj + q = self.q_a_proj(x) + q = torch.nn.functional.relu(q) + q = self.q_b_proj(q) + # KV chain: kv_a_proj + kv = self.kv_a_proj(x) + out = q + kv + # Inter-layer distance to push next layer's linears beyond BFS depth + return self.layernorm(torch.nn.functional.relu(out)) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def _build_gm(model, example_input): + """Export *model* to an FX GraphModule.""" + egm = torch.export.export(model, (example_input,)) + return egm.module() + + +def test_pattern_matching_single_block(): + """The pattern matcher should find exactly one pair for a single MLA block.""" + model = MockMLABlock(128, 64, 128).eval().to("cuda") + example_input = torch.randn(4, 128, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 1, f"Expected 1 fork-point pair, got {len(pairs)}" + + +def test_pattern_matching_multi_block(): + """Multiple layers with sufficient inter-layer distance should all be matched.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + model = ( + nn.Sequential( + MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim), + MockMLABlock(kv_out_dim, q_inner_dim, kv_out_dim), + ) + .eval() + .to("cuda") + ) + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 2, f"Expected 2 fork-point pairs, got {len(pairs)}" + + +def test_numerical_correctness(): + """After the transform the GraphModule must produce the same output as the original model.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim).eval().to("cuda") + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + test_x = torch.randn(4, hidden_dim, device="cuda") + ref_output = model(test_x) + + gm, num_replaced = _execute_kv_proj_in_aux_stream(gm) + + assert num_replaced == 1, f"Expected 1 replacement, got {num_replaced}" + + y = gm(test_x) + assert torch.allclose(y, ref_output, atol=1e-5), ( + f"Output mismatch: max diff = {(y - ref_output).abs().max().item()}" + ) + + +def test_numerical_correctness_multi_block(): + """Multi-block correctness test.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = ( + nn.Sequential( + MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim), + MockMLABlock(kv_out_dim, q_inner_dim, kv_out_dim), + ) + .eval() + .to("cuda") + ) + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + test_x = torch.randn(4, hidden_dim, device="cuda") + ref_output = model(test_x) + + gm, num_replaced = _execute_kv_proj_in_aux_stream(gm) + + assert num_replaced == 2, f"Expected 2 replacements, got {num_replaced}" + + y = gm(test_x) + assert torch.allclose(y, ref_output, atol=1e-5), ( + f"Output mismatch: max diff = {(y - ref_output).abs().max().item()}" + ) + + +def test_cuda_graph_compatibility(): + """The transformed GraphModule must work under CUDA graph capture and replay.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim).eval().to("cuda") + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + test_x = torch.randn(4, hidden_dim, device="cuda") + ref_output = model(test_x) + + gm, num_replaced = _execute_kv_proj_in_aux_stream(gm) + assert num_replaced == 1 + + # Allocate static buffers for CUDA graph capture. + static_x = torch.randn(4, hidden_dim, device="cuda") + static_output = torch.randn(4, kv_out_dim, device="cuda") + + # Warm up (required before capture). + for _ in range(3): + static_output.copy_(gm(static_x)) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output.copy_(gm(static_x)) + + static_x.copy_(test_x) + graph.replay() + + assert torch.allclose(static_output, ref_output, atol=1e-5), ( + f"CUDA graph output mismatch: max diff = {(static_output - ref_output).abs().max().item()}" + ) + + +def test_no_match_on_single_linear(): + """A node with only one linear user should not be matched.""" + + class SingleLinear(nn.Module): + def __init__(self, dim): + super().__init__() + self.fc = nn.Linear(dim, dim, bias=False) + + def forward(self, x): + return self.fc(x) + + model = SingleLinear(64).eval().to("cuda") + example_input = torch.randn(4, 64, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 0, f"Expected 0 matches, got {len(pairs)}" + + +def test_no_match_when_both_have_downstream_linear(): + """When *both* branches have downstream linears the pattern should not match.""" + + class BothHeavy(nn.Module): + def __init__(self, dim, inner): + super().__init__() + self.fc_a1 = nn.Linear(dim, inner, bias=False) + self.fc_a2 = nn.Linear(inner, dim, bias=False) + self.fc_b1 = nn.Linear(dim, inner, bias=False) + self.fc_b2 = nn.Linear(inner, dim, bias=False) + + def forward(self, x): + a = self.fc_a2(torch.relu(self.fc_a1(x))) + b = self.fc_b2(torch.relu(self.fc_b1(x))) + return a + b + + model = BothHeavy(64, 32).eval().to("cuda") + example_input = torch.randn(4, 64, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 0, f"Expected 0 matches, got {len(pairs)}" diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py new file mode 100644 index 000000000000..87e30ef26147 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py @@ -0,0 +1,502 @@ +"""Tests for multi-stream MoE shared-expert transform across model architectures. + +Verifies that ``_execute_shared_expert_in_aux_stream`` correctly identifies the +shared-expert branch and moves it to the auxiliary CUDA stream for the MoE +patterns used in DeepSeek V3, GLM4 MoE Lite, Mixtral, and Nemotron-H (with +and without latent projections). + +Architecture patterns tested: + + **DeepSeek V3 / GLM4 MoE Lite** — Gated-MLP shared expert + (gate_proj + up_proj → SiLU gate → down_proj). + Routed MoE dispatched first; shared expert on ``identity``. + Merge: ``moe_out + shared_out``. + + **Mixtral** — Pure routed MoE, *no* shared expert. + The transform must produce zero matches (no-op). + + **Nemotron-H (no latent)** — Simple-MLP shared expert + (up_proj → ReLU² → down_proj). + Shared expert dispatched first; routed MoE second. + Merge: ``shared_out + routed_out``. + + **Nemotron-H (with latent projection)** — Same shared expert as above, + but the routed path wraps the MoE op with + ``fc1_latent_proj → MoE → fc2_latent_proj``. + Tests that the BFS from MoE to the merge ``add`` traverses the extra + projection nodes correctly. + +Each architecture is tested for: + 1. Pattern matching — correct number of replacements. + 2. Graph structure — ``begin_aux``, ``end_aux``, ``wait_aux`` nodes present. + 3. Numerical correctness — output matches eager reference within tolerance. + 4. CUDA graph compatibility — capture + replay produces correct output. + 5. Multi-layer stacking — multiple MoE layers handled independently. +""" + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_moe import ( + _execute_shared_expert_in_aux_stream, +) +from tensorrt_llm._torch.auto_deploy.utils.multi_stream_utils import ( + begin_aux_stream_passthrough, + cuda_stream_manager, + end_aux_stream_passthrough, + wait_aux_stream_passthrough, +) + +# --------------------------------------------------------------------------- +# Mock fused-MoE custom op (distinct name to avoid conflicts with other tests) +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("auto_deploy::mock_fused_moe_moe_test", mutates_args=()) +def mock_fused_moe( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + expert_weight: torch.Tensor, +) -> torch.Tensor: + """Mock fused MoE: a simple linear transform standing in for the real kernel.""" + return torch.ops.aten.linear(x, expert_weight) + + +@mock_fused_moe.register_fake +def _mock_fused_moe_fake(x, selected_experts, routing_weights, expert_weight): + return torch.ops.aten.linear(x, expert_weight) + + +_MOE_OPS = [torch.ops.auto_deploy.mock_fused_moe_moe_test] + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _build_gm(model, example_input): + """Export *model* to an FX ``GraphModule``.""" + return torch.export.export(model, (example_input,)).module() + + +def _stream_targets(gm): + """Return the set of ``call_function`` targets present in *gm*.""" + return {n.target for n in gm.graph.nodes if n.op == "call_function"} + + +def _assert_stream_nodes_present(gm): + """Assert that the three stream-management passthrough nodes are in the graph.""" + targets = _stream_targets(gm) + assert begin_aux_stream_passthrough in targets, "begin_aux_stream_passthrough not in graph" + assert end_aux_stream_passthrough in targets, "end_aux_stream_passthrough not in graph" + assert wait_aux_stream_passthrough in targets, "wait_aux_stream_passthrough not in graph" + + +def _assert_numerical_correctness(gm, model, test_x, *, atol=1e-5): + """Assert that *gm* and *model* produce the same output on *test_x*.""" + ref = model(test_x) + out = gm(test_x) + assert torch.allclose(out, ref, atol=atol), ( + f"Output mismatch: max diff = {(out - ref).abs().max().item()}" + ) + + +def _assert_cuda_graph_correctness(gm, model, test_x, *, atol=1e-5): + """Assert correctness under CUDA graph capture + replay.""" + ref = model(test_x) + out_shape = ref.shape + + static_x = torch.randn_like(test_x) + static_out = torch.empty(out_shape, device="cuda", dtype=ref.dtype) + + # Warm-up (required before capture). + for _ in range(3): + static_out.copy_(gm(static_x)) + + cuda_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph): + static_out.copy_(gm(static_x)) + + static_x.copy_(test_x) + cuda_graph.replay() + + assert torch.allclose(static_out, ref, atol=atol), ( + f"CUDA graph output mismatch: max diff = {(static_out - ref).abs().max().item()}" + ) + + +# --------------------------------------------------------------------------- +# Mock modules — shared expert variants +# --------------------------------------------------------------------------- + + +class _GatedMLP(nn.Module): + """DeepSeek / GLM4 shared expert: ``down_proj(silu(gate_proj(x)) * up_proj(x))``.""" + + def __init__(self, hidden_dim: int, intermediate_dim: int): + super().__init__() + self.gate_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False) + self.up_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False) + self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _SimpleMLP(nn.Module): + """Nemotron-H shared expert: ``down_proj(relu(up_proj(x)) ** 2)``.""" + + def __init__(self, hidden_dim: int, intermediate_dim: int): + super().__init__() + self.up_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False) + self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(torch.relu(self.up_proj(x)) ** 2) + + +# --------------------------------------------------------------------------- +# Mock MoE layer modules — one per architecture pattern +# --------------------------------------------------------------------------- + + +class MockDeepSeekGLM4MoELayer(nn.Module): + """DeepSeek V3 / GLM4 MoE Lite pattern. + + Graph topology:: + + hidden_states ─┬─ gate ─ topk ─────────────────────────┐ + ├─ shared_experts (gated MLP) ─ shared_out │ + └─ mock_fused_moe ──────────── moe_out ─┘ + moe_out + shared_out → layernorm → out + + Routed MoE is dispatched *before* the shared expert in graph order. + The ``add`` has the routed output on the left. + """ + + def __init__(self, hidden_dim: int, intermediate_dim: int, num_experts: int = 8): + super().__init__() + self.gate = nn.Linear(hidden_dim, num_experts, bias=False) + self.shared_experts = _GatedMLP(hidden_dim, intermediate_dim) + self.expert_weight = nn.Parameter(torch.randn(hidden_dim, hidden_dim)) + self.layernorm = nn.LayerNorm(hidden_dim) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + logits = self.gate(hidden_states) + routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1) + + # Routed path first (matches DeepSeek / GLM4 dispatch order). + moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test( + hidden_states, selected_experts, routing_weights, self.expert_weight + ) + # Shared expert on original input. + shared_out = self.shared_experts(identity) + + return self.layernorm(moe_out + shared_out) + + +class MockMixtralMoELayer(nn.Module): + """Mixtral pattern — pure routed MoE, **no** shared expert. + + The transform must return 0 replacements for this topology. + """ + + def __init__(self, hidden_dim: int, num_experts: int = 8): + super().__init__() + self.gate = nn.Linear(hidden_dim, num_experts, bias=False) + self.expert_weight = nn.Parameter(torch.randn(hidden_dim, hidden_dim)) + self.layernorm = nn.LayerNorm(hidden_dim) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.gate(hidden_states) + routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1) + moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test( + hidden_states, selected_experts, routing_weights, self.expert_weight + ) + return self.layernorm(moe_out) + + +class MockNemotronHMoELayer(nn.Module): + """Nemotron-H pattern *without* latent projections. + + Graph topology:: + + hidden_states ─┬─ gate ─ topk ─────────────────────────┐ + ├─ shared_experts (simple MLP) ─ shared_out │ + └─ mock_fused_moe ──────────── moe_out ─┘ + shared_out + moe_out → layernorm → out + + Shared expert is dispatched *before* the MoE in graph order. + The ``add`` has the shared output on the left. + """ + + def __init__(self, hidden_dim: int, intermediate_dim: int, num_experts: int = 8): + super().__init__() + self.gate = nn.Linear(hidden_dim, num_experts, bias=False) + self.shared_experts = _SimpleMLP(hidden_dim, intermediate_dim) + self.expert_weight = nn.Parameter(torch.randn(hidden_dim, hidden_dim)) + self.layernorm = nn.LayerNorm(hidden_dim) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residuals = hidden_states + logits = self.gate(hidden_states) + routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1) + + # Shared expert dispatched first (matches Nemotron-H dispatch order). + shared_out = self.shared_experts(residuals) + # Routed path. + moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test( + hidden_states, selected_experts, routing_weights, self.expert_weight + ) + + return self.layernorm(shared_out + moe_out) + + +class MockNemotronHLatentMoELayer(nn.Module): + """Nemotron-H pattern *with* latent projections. + + Graph topology:: + + hidden_states ─┬─ gate ─ topk ──────────────────────────────────────┐ + ├─ shared_experts (simple MLP) ────────── shared_out │ + └─ fc1_latent ─ mock_fused_moe ─ fc2_latent ─ routed_out ┘ + shared_out + routed_out → ln → out + + The latent projections add nodes between the MoE op and the merge ``add``, + testing that the forward BFS from MoE correctly traverses extra projection + nodes. + """ + + def __init__( + self, + hidden_dim: int, + intermediate_dim: int, + latent_dim: int, + num_experts: int = 8, + ): + super().__init__() + self.gate = nn.Linear(hidden_dim, num_experts, bias=False) + self.shared_experts = _SimpleMLP(hidden_dim, intermediate_dim) + self.fc1_latent_proj = nn.Linear(hidden_dim, latent_dim, bias=False) + self.fc2_latent_proj = nn.Linear(latent_dim, hidden_dim, bias=False) + # Expert weight operates in latent space. + self.expert_weight = nn.Parameter(torch.randn(latent_dim, latent_dim)) + self.layernorm = nn.LayerNorm(hidden_dim) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residuals = hidden_states + logits = self.gate(hidden_states) + routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1) + + # Shared expert dispatched first. + shared_out = self.shared_experts(residuals) + + # Latent projection → MoE → back-projection. + x_latent = self.fc1_latent_proj(hidden_states) + moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test( + x_latent, selected_experts, routing_weights, self.expert_weight + ) + routed_out = self.fc2_latent_proj(moe_out) + + return self.layernorm(shared_out + routed_out) + + +# =================================================================== +# Tests — DeepSeek V3 / GLM4 MoE Lite (gated-MLP shared expert) +# =================================================================== + + +def test_deepseek_glm4_pattern_and_correctness(): + """Single-layer: pattern match + graph structure + numerical correctness.""" + hidden_dim, intermediate_dim = 128, 256 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockDeepSeekGLM4MoELayer(hidden_dim, intermediate_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 1, f"Expected 1 replacement, got {num}" + _assert_stream_nodes_present(gm) + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +def test_deepseek_glm4_cuda_graph(): + """CUDA graph capture + replay for DeepSeek / GLM4 pattern.""" + hidden_dim, intermediate_dim = 128, 256 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockDeepSeekGLM4MoELayer(hidden_dim, intermediate_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + assert num == 1 + + _assert_cuda_graph_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +def test_deepseek_glm4_multi_layer(): + """Two stacked DeepSeek/GLM4 MoE layers — both should be transformed.""" + hidden_dim, intermediate_dim = 128, 256 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = ( + nn.Sequential( + MockDeepSeekGLM4MoELayer(hidden_dim, intermediate_dim), + MockDeepSeekGLM4MoELayer(hidden_dim, intermediate_dim), + ) + .eval() + .to("cuda") + ) + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 2, f"Expected 2 replacements, got {num}" + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +# =================================================================== +# Tests — Mixtral (no shared expert → no-op) +# =================================================================== + + +def test_mixtral_no_shared_expert_no_match(): + """Mixtral has no shared expert; the transform must produce zero matches.""" + hidden_dim = 128 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockMixtralMoELayer(hidden_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 0, f"Expected 0 replacements for Mixtral (no shared expert), got {num}" + + # Graph should NOT contain any stream-management nodes. + targets = _stream_targets(gm) + assert begin_aux_stream_passthrough not in targets + assert end_aux_stream_passthrough not in targets + assert wait_aux_stream_passthrough not in targets + + # Numerical correctness should still hold (graph unchanged). + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +# =================================================================== +# Tests — Nemotron-H without latent projections +# =================================================================== + + +def test_nemotron_h_pattern_and_correctness(): + """Single-layer Nemotron-H (no latent): pattern + graph + correctness.""" + hidden_dim, intermediate_dim = 128, 256 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockNemotronHMoELayer(hidden_dim, intermediate_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 1, f"Expected 1 replacement, got {num}" + _assert_stream_nodes_present(gm) + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +def test_nemotron_h_cuda_graph(): + """CUDA graph capture + replay for Nemotron-H (no latent) pattern.""" + hidden_dim, intermediate_dim = 128, 256 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockNemotronHMoELayer(hidden_dim, intermediate_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + assert num == 1 + + _assert_cuda_graph_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +def test_nemotron_h_multi_layer(): + """Two stacked Nemotron-H (no latent) layers — both should be transformed.""" + hidden_dim, intermediate_dim = 128, 256 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = ( + nn.Sequential( + MockNemotronHMoELayer(hidden_dim, intermediate_dim), + MockNemotronHMoELayer(hidden_dim, intermediate_dim), + ) + .eval() + .to("cuda") + ) + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 2, f"Expected 2 replacements, got {num}" + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +# =================================================================== +# Tests — Nemotron-H with latent projections +# =================================================================== + + +def test_nemotron_h_latent_pattern_and_correctness(): + """Single-layer Nemotron-H (with latent): pattern + graph + correctness.""" + hidden_dim, intermediate_dim, latent_dim = 128, 256, 64 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockNemotronHLatentMoELayer(hidden_dim, intermediate_dim, latent_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 1, f"Expected 1 replacement, got {num}" + _assert_stream_nodes_present(gm) + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +def test_nemotron_h_latent_cuda_graph(): + """CUDA graph capture + replay for Nemotron-H (with latent) pattern.""" + hidden_dim, intermediate_dim, latent_dim = 128, 256, 64 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockNemotronHLatentMoELayer(hidden_dim, intermediate_dim, latent_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + assert num == 1 + + _assert_cuda_graph_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +def test_nemotron_h_latent_multi_layer(): + """Two stacked Nemotron-H (with latent) layers — both should be transformed.""" + hidden_dim, intermediate_dim, latent_dim = 128, 256, 64 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = ( + nn.Sequential( + MockNemotronHLatentMoELayer(hidden_dim, intermediate_dim, latent_dim), + MockNemotronHLatentMoELayer(hidden_dim, intermediate_dim, latent_dim), + ) + .eval() + .to("cuda") + ) + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 2, f"Expected 2 replacements, got {num}" + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_resource_handlers.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_resource_handlers.py index 8252ed7f3bc0..ec2d4e23d506 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_resource_handlers.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_resource_handlers.py @@ -49,8 +49,8 @@ def test_paged_handler_allocate_with_blocks(kv_layout): tokens_per_block = 32 seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=tokens_per_block) seq_info.to("cuda") - # Set up num_blocks via estimate_cache_loc_capacity - seq_info.estimate_cache_loc_capacity(num_blocks=10) + # Set up num_blocks via update_cache_information + seq_info.update_cache_information(num_blocks=10) tensor = handler.allocate(seq_info) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_cached_sequence_interface.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_cached_sequence_interface.py index 49464ce6bb4e..59348c28607d 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_cached_sequence_interface.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_cached_sequence_interface.py @@ -655,8 +655,8 @@ def test_sequence_info_estimate_cache_tokens_per_forward_with_overflow(): assert result == 128 -def test_sequence_info_estimate_cache_loc_capacity_no_resize(): - """Test estimate_cache_loc_capacity() when capacity is sufficient.""" +def test_sequence_info_update_cache_information_no_resize(): + """Test update_cache_information() when capacity is sufficient.""" seq_info = SequenceInfo( max_seq_len=128, max_batch_size=4, @@ -667,15 +667,15 @@ def test_sequence_info_estimate_cache_loc_capacity_no_resize(): initial_capacity = seq_info._input_buffer.get_capacity("cache_loc") # Request a small capacity that should already be available - seq_info.estimate_cache_loc_capacity(num_blocks=4) + seq_info.update_cache_information(num_blocks=4) # Capacity should not have changed if it was already sufficient if initial_capacity >= 4 * 4 + 1: # num_blocks * max_batch_size + 1 assert seq_info._input_buffer.get_capacity("cache_loc") == initial_capacity -def test_sequence_info_estimate_cache_loc_capacity_resizes(): - """Test estimate_cache_loc_capacity() resizes buffer when needed.""" +def test_sequence_info_update_cache_information_resizes(): + """Test update_cache_information() resizes buffer when needed.""" seq_info = SequenceInfo( max_seq_len=128, max_batch_size=4, @@ -687,7 +687,7 @@ def test_sequence_info_estimate_cache_loc_capacity_resizes(): # Request a large capacity large_num_blocks = 1000 - seq_info.estimate_cache_loc_capacity(num_blocks=large_num_blocks) + seq_info.update_cache_information(num_blocks=large_num_blocks) expected_capacity = large_num_blocks * 4 + 1 # num_blocks * max_batch_size + 1 if expected_capacity > initial_capacity: @@ -886,3 +886,116 @@ def test_generic_state_handler_allocated_locally(paged_kv_cache_config): assert interface._caches["generic_state"] is not None # Without typed handlers, should use plain KVCacheManager assert isinstance(interface.kv_cache_manager, KVCacheManager) + + +# ============================================================================= +# _requires_copy Tests +# ============================================================================= + + +def test_requires_copy_pre_populated(): + """Verify _requires_copy is pre-populated with expected args.""" + seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=32) + + expected = { + "max_seq_info", + "page_seq_indices", + "page_in_seq", + "logits_gather_indices", + "logits_gather_info", + "_gather_idx", + "_mask_scatter_indices", + } + assert expected.issubset(seq_info._requires_copy) + + +def test_requires_copy_args_not_in_named_args(): + """Verify that _requires_copy args do NOT appear in named_args.""" + seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=32) + + named_args = seq_info.named_args + for rc_arg in seq_info._requires_copy: + assert rc_arg not in named_args, f"{rc_arg} should not be in named_args" + assert f"{rc_arg}_host" not in named_args, f"{rc_arg}_host should not be in named_args" + + +def test_requires_copy_args_stored_to_input_buffer(): + """Verify that _requires_copy args are written to InputBuffer by _store_arg.""" + seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=32) + + # logits_gather_indices is in _requires_copy, so nest_sequences should store it + seq_info.nest_sequences( + input_ids=[[1, 2, 3]], + input_pos=0, + cache_loc=[0], + pages_per_seq=[1], + ) + + # Should be accessible via _get_arg (reads from InputBuffer) + logits_gather_indices = seq_info._get_arg("logits_gather_indices") + assert logits_gather_indices is not None + + logits_gather_info = seq_info._get_arg("logits_gather_info") + assert logits_gather_info is not None + + +def test_require_copy_marks_new_arg(): + """Verify require_copy() adds an arg to _requires_copy.""" + seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=32) + + # cu_seqlen is available but not in _requires_copy by default + assert "cu_seqlen" not in seq_info._requires_copy + + result = seq_info.require_copy("cu_seqlen") + assert result is True + assert "cu_seqlen" in seq_info._requires_copy + + # Second call should return False (already marked) + result = seq_info.require_copy("cu_seqlen") + assert result is False + + +def test_require_copy_rejects_unavailable_arg(): + """Verify require_copy() rejects args not in available_args.""" + seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=32) + + with pytest.raises(AssertionError): + seq_info.require_copy("nonexistent_arg") + + +def test_register_host_prepare_populates_requires_copy(): + """Verify register_host_prepare_for_attention_forward auto-populates _requires_copy.""" + seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=32) + + # Define a dummy host prepare function + def dummy_host_prepare(batch_info_host: torch.Tensor, cu_num_pages_host: torch.Tensor): + pass + + # batch_info_host and cu_num_pages_host should be marked in _requires_copy + seq_info.register_host_prepare_for_attention_forward( + dummy_host_prepare, ["batch_info_host", "cu_num_pages_host"] + ) + + assert "batch_info_host" in seq_info._requires_copy + assert "cu_num_pages_host" in seq_info._requires_copy + + +def test_requires_copy_host_suffix_enables_base_storage(): + """Verify that marking a _host arg enables storage of the base arg.""" + seq_info = SequenceInfo(max_seq_len=128, max_batch_size=4, tokens_per_block=32) + + # Mark batch_info_host as requires_copy + seq_info.require_copy("batch_info_host") + + # nest_sequences stores batch_info (base name) -- the _host suffix check in + # _store_arg should allow it through because batch_info_host is in _requires_copy + seq_info.nest_sequences( + input_ids=[[1, 2, 3]], + input_pos=0, + cache_loc=[0], + pages_per_seq=[1], + ) + + # batch_info_host should be accessible via _get_arg + batch_info_host = seq_info._get_arg("batch_info_host") + assert batch_info_host is not None diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py index 50bbc715e1f7..82a01eedaa58 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py @@ -249,6 +249,13 @@ def __init__(self, tokens_per_block: int, num_slots: int = 8): self.mamba_cache_free_blocks = num_slots def get_cache_indices(self, request): + # For generation requests (which have max_beam_num_tokens), return the + # correct number of blocks since ADEngine no longer truncates for them. + # For context requests, return many dummy page IDs; ADEngine truncates. + num_tokens = getattr(request, "max_beam_num_tokens", None) + if num_tokens is not None: + num_blocks = self.get_num_kv_blocks(num_tokens) + return list(range(num_blocks)) return list(range(1024)) def get_num_kv_blocks(self, num_tokens: int) -> int: diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py new file mode 100644 index 000000000000..6ea74d972123 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py @@ -0,0 +1,188 @@ +import pytest +import torch +from torch.export import Dim + +from tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import * # noqa +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer +from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op + + +class SwiGLUMLP(torch.nn.Module): + """SwiGLU MLP module: silu(x @ gate.T) * (x @ up.T) @ down.T""" + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class SwiGLUMLPWithBias(torch.nn.Module): + """SwiGLU MLP module with biases.""" + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=True) + self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=True) + self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, x): + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class SwiGLUTestModel(torch.nn.Module): + """Test model with SwiGLU MLP sandwiched between linear layers.""" + + def __init__( + self, + hidden_size: int = 256, + intermediate_size: int = 512, + with_bias: bool = False, + ): + super().__init__() + self.linear1 = torch.nn.Linear(hidden_size, hidden_size, device="cuda", dtype=torch.float16) + if with_bias: + self.mlp = SwiGLUMLPWithBias(hidden_size, intermediate_size) + else: + self.mlp = SwiGLUMLP(hidden_size, intermediate_size) + self.mlp = self.mlp.to(device="cuda", dtype=torch.float16) + self.linear2 = torch.nn.Linear(hidden_size, hidden_size, device="cuda", dtype=torch.float16) + + def forward(self, x): + x = self.linear1(x) + x = self.mlp(x) + x = self.linear2(x) + return x + + +class SwiGLUTestModelMultipleMLP(torch.nn.Module): + """Test model with multiple SwiGLU MLPs to test multiple pattern matches.""" + + def __init__( + self, + hidden_size: int = 256, + intermediate_size: int = 512, + num_layers: int = 2, + ): + super().__init__() + self.layers = torch.nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + torch.nn.ModuleDict( + { + "linear": torch.nn.Linear( + hidden_size, hidden_size, device="cuda", dtype=torch.float16 + ), + "mlp": SwiGLUMLP(hidden_size, intermediate_size).to( + device="cuda", dtype=torch.float16 + ), + } + ) + ) + + def forward(self, x): + for layer in self.layers: + x = layer["linear"](x) + x = layer["mlp"](x) + return x + + +def _run_fusion_test(model, expected_op, expected_num_matches=1): + """Run the SwiGLU fusion test. + + Args: + model: The test model to transform. + expected_op: The expected fused op to find in the transformed graph. + expected_num_matches: Expected number of fused ops. + """ + x = torch.randn(2, 256, device="cuda", dtype=torch.float16) + dynamic_shapes = {0: Dim.DYNAMIC} + gm = torch_export_to_gm(model, args=(x,), dynamic_shapes=(dynamic_shapes,), clone=True) + + # Apply transforms + gm_transformed = InferenceOptimizer( + None, + { + "match_swiglu_pattern": { + "stage": "pattern_matcher", + }, + "fuse_swiglu": { + "stage": "post_load_fusion", + "enabled": True, + }, + }, + )(None, gm) + + # Move to CUDA if needed + gm_transformed = gm_transformed.to("cuda") + + # Check that the expected op is present + count = sum(1 for n in gm_transformed.graph.nodes if is_op(n, expected_op)) + assert count == expected_num_matches, ( + f"Expected {expected_num_matches} {expected_op} ops, got {count}" + ) + + # Verify numerical correctness + y_transformed = gm_transformed(x) + y_model = model(x) + torch.testing.assert_close(y_transformed, y_model, atol=1e-2, rtol=1e-2) + + # Test with a different batch size + new_input = torch.randn(4, 256, device="cuda", dtype=torch.float16) + y_transformed_2 = gm_transformed(new_input) + y_model_2 = model(new_input) + torch.testing.assert_close(y_transformed_2, y_model_2, atol=1e-2, rtol=1e-2) + + +def test_swiglu_fusion_basic(): + """Test basic SwiGLU fusion without biases.""" + model = SwiGLUTestModel(with_bias=False) + _run_fusion_test(model, torch.ops.auto_deploy.fused_swiglu_mlp.default) + + +def test_swiglu_fusion_with_bias(): + """Test SwiGLU fusion with biases.""" + model = SwiGLUTestModel(with_bias=True) + _run_fusion_test(model, torch.ops.auto_deploy.fused_swiglu_mlp.default) + + +@pytest.mark.parametrize("num_layers", [2, 3]) +def test_swiglu_fusion_multiple_layers(num_layers): + """Test that multiple SwiGLU patterns are fused correctly.""" + model = SwiGLUTestModelMultipleMLP(num_layers=num_layers) + _run_fusion_test( + model, torch.ops.auto_deploy.fused_swiglu_mlp.default, expected_num_matches=num_layers + ) + + +def test_swiglu_pattern_match_only(): + """Test pattern matching stage only (without fusion).""" + model = SwiGLUTestModel() + x = torch.randn(2, 256, device="cuda", dtype=torch.float16) + dynamic_shapes = {0: Dim.DYNAMIC} + gm = torch_export_to_gm(model, args=(x,), dynamic_shapes=(dynamic_shapes,), clone=True) + + # Only run pattern matching, not fusion + gm_matched = InferenceOptimizer( + None, + { + "match_swiglu_pattern": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + + # Check that the intermediate op is present + has_swiglu_op = any( + is_op(n, torch.ops.auto_deploy.torch_swiglu_mlp.default) for n in gm_matched.graph.nodes + ) + assert has_swiglu_op, "Pattern matcher should produce torch_swiglu_mlp op" + + # Verify numerical correctness + y_matched = gm_matched(x) + y_model = model(x) + torch.testing.assert_close(y_matched, y_model, atol=1e-3, rtol=1e-3) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py index dc353e49936d..4dfd3f683ce5 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py @@ -1,14 +1,26 @@ +import operator + import torch from torch.export import Dim -from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import * # noqa +from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import ( # noqa + flashinfer_fused_add_rms_norm, +) from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.rms_norm import * # noqa from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.transform.interface import TransformConfig +from tensorrt_llm._torch.auto_deploy.transform.library.fused_add_rms_norm import FuseAddRMSNorm from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op +# --------------------------------------------------------------------------- +# Models +# --------------------------------------------------------------------------- + + +class AddCastNormModel(torch.nn.Module): + """Pattern 1: add + cast(to.dtype) + rms_norm.""" -class TestModel(torch.nn.Module): def __init__(self, hidden_size=128, eps=1e-5): super().__init__() self.weight = torch.nn.Parameter( @@ -23,55 +35,263 @@ def forward(self, x, residual): return norm, added -def _run_test(model): - # The replacement uses flashinfer_fused_add_rms_norm python wrapper which calls the inplace op - # auto_deploy::flashinfer_fused_add_rms_norm_inplace - op = torch.ops.auto_deploy.flashinfer_fused_add_rms_norm_inplace +class AddNormModel(torch.nn.Module): + """Pattern 2: add + rms_norm (no intermediate cast).""" + + def __init__(self, hidden_size=128, eps=1e-5): + super().__init__() + self.weight = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.eps = eps + + def forward(self, x, residual): + added = x + residual + norm = torch.ops.auto_deploy.flashinfer_rms_norm(added, self.weight, self.eps) + return norm, added + + +class MultiUserModel(torch.nn.Module): + """Both add and rms_norm outputs have multiple users (DeepSeek V3 MoE pattern). + + add_result has 2 users: rms_norm + next residual add + norm_result has 2 users: linear1 + linear2 + """ + + def __init__(self, hidden_size=128, eps=1e-5): + super().__init__() + self.weight = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.linear1 = torch.nn.Linear( + hidden_size, hidden_size, bias=False, device="cuda", dtype=torch.bfloat16 + ) + self.linear2 = torch.nn.Linear( + hidden_size, hidden_size, bias=False, device="cuda", dtype=torch.bfloat16 + ) + self.eps = eps + + def forward(self, residual, attn_output, moe_output): + # add with 2 users (norm + next_add) + add_result = residual + attn_output + # rms_norm with 2 users (linear1, linear2) + norm_result = torch.ops.auto_deploy.flashinfer_rms_norm(add_result, self.weight, self.eps) + out1 = self.linear1(norm_result) + out2 = self.linear2(norm_result) + combined = out1 + out2 + # add_result also feeds into next residual add + next_residual = add_result + moe_output + return combined, next_residual + + +class ChainedModel(torch.nn.Module): + """Two consecutive add+norm pairs sharing residual (like transformer layers). + + Layer 1: add1 = embed + attn_out, norm1 = rms_norm(add1) -- add1 has 2 users + Layer 2: add2 = add1 + mlp_out, norm2 = rms_norm(add2) -- add2 has 2 users + """ + + def __init__(self, hidden_size=128, eps=1e-5): + super().__init__() + self.weight1 = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.weight2 = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.linear = torch.nn.Linear( + hidden_size, hidden_size, bias=False, device="cuda", dtype=torch.bfloat16 + ) + self.eps = eps + + def forward(self, embed, attn_out, mlp_out): + add1 = embed + attn_out + norm1 = torch.ops.auto_deploy.flashinfer_rms_norm(add1, self.weight1, self.eps) + branch1 = self.linear(norm1) + + add2 = add1 + mlp_out + norm2 = torch.ops.auto_deploy.flashinfer_rms_norm(add2, self.weight2, self.eps) + branch2 = self.linear(norm2) + + return branch1 + branch2, add2 - def checker(gm): - return any(is_op(n, op) for n in gm.graph.nodes) +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _count_fused_ops(gm): + """Count flashinfer_fused_add_rms_norm wrapper calls in the graph.""" + return sum( + 1 + for n in gm.graph.nodes + if n.op == "call_function" and n.target is flashinfer_fused_add_rms_norm + ) + + +def _count_rms_norm_ops(gm): + """Count flashinfer_rms_norm calls in the graph.""" + return sum(1 for n in gm.graph.nodes if is_op(n, torch.ops.auto_deploy.flashinfer_rms_norm)) + + +def _count_add_ops(gm): + """Count aten.add.Tensor calls in the graph.""" + return sum(1 for n in gm.graph.nodes if is_op(n, torch.ops.aten.add.Tensor)) + + +def _export_model(model, *inputs, dynamic_dim0=True): + """Export a model to a GraphModule, optionally with a dynamic batch dimension.""" + if dynamic_dim0: + dyn = Dim.DYNAMIC + ds = tuple({0: dyn} for _ in inputs) + else: + ds = None + return torch_export_to_gm(model, args=inputs, dynamic_shapes=ds, clone=True) + + +def _apply_transform(gm): + """Apply fuse_add_rms_norm via InferenceOptimizer (integration-style).""" + return InferenceOptimizer( + None, + {"fuse_add_rms_norm": {"stage": "post_load_fusion"}}, + )(None, gm) + + +def _apply_transform_direct(gm): + """Apply the transform directly (unit-test style).""" + config = TransformConfig(stage="post_load_fusion") + transform = FuseAddRMSNorm(config=config) + gm, info = transform._apply(gm, None, None, None) + return gm, info + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_fuse_add_cast_rms_norm(): + """Original test: add + cast(bf16) + rms_norm → fused op.""" + model = AddCastNormModel() bsz, seq_len, hidden = 2, 8, 128 - # Inputs should be bfloat16 x = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + + gm = _export_model(model, x, residual) + gm_t = _apply_transform(gm) + + # Structure check + assert _count_fused_ops(gm_t) >= 1, "fused op not found in graph" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + + # Numerical check + y_fused = gm_t(x.clone(), residual.clone()) + y_ref = model(x.clone(), residual.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) + + +def test_fuse_add_rms_norm_no_cast(): + """Pattern 2: add + rms_norm (no cast) → fused op.""" + model = AddNormModel() + bsz, seq_len, hidden = 2, 8, 128 + x = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + + gm = _export_model(model, x, residual) + gm_t = _apply_transform(gm) + + # Structure check + assert _count_fused_ops(gm_t) >= 1, "fused op not found in graph" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + + # Numerical check + y_fused = gm_t(x.clone(), residual.clone()) + y_ref = model(x.clone(), residual.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) + + +def test_fuse_add_rms_norm_multi_user(): + """Multi-user: both add (2 users) and rms_norm (2 users) → fused op. + + This is the key pattern from the DeepSeek V3 / GLM4-MoE graph that failed + with the old inductor-based pattern matcher due to num_users constraints. + """ + model = MultiUserModel() + bsz, seq_len, hidden = 2, 8, 128 residual = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + attn_out = torch.randn_like(residual) + moe_out = torch.randn_like(residual) - # Dynamic shapes - dyn_batch_size = Dim.DYNAMIC - ds_x = {0: dyn_batch_size} - ds_res = {0: dyn_batch_size} + gm = _export_model(model, residual, attn_out, moe_out) - gm = torch_export_to_gm(model, args=(x, residual), dynamic_shapes=(ds_x, ds_res), clone=True) + # Before: 1 add+norm fusible pair, add has 2 users, norm has 2 users + assert _count_rms_norm_ops(gm) == 1 - gm_transformed = InferenceOptimizer( - None, - { - "fuse_add_rms_norm": { - "stage": "post_load_fusion", - }, - }, - )(None, gm) + gm_t, info = _apply_transform_direct(gm) - # Check if transform happened - if not checker(gm_transformed): - raise AssertionError( - "flashinfer_fused_add_rms_norm_inplace op not found in transformed graph" - ) + # Structure check + assert info.num_matches == 1, f"Expected 1 match, got {info.num_matches}" + assert _count_fused_ops(gm_t) == 1, "fused op not found in graph" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + + # Verify getitem nodes for both outputs + getitems = [ + n + for n in gm_t.graph.nodes + if n.op == "call_function" + and n.target is operator.getitem + and isinstance(n.args[0], torch.fx.Node) + and n.args[0].target is flashinfer_fused_add_rms_norm + ] + assert len(getitems) == 2, f"Expected 2 getitem nodes, got {len(getitems)}" + + # Numerical check + y_fused = gm_t(residual.clone(), attn_out.clone(), moe_out.clone()) + y_ref = model(residual.clone(), attn_out.clone(), moe_out.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) + + +def test_fuse_add_rms_norm_chained(): + """Chained: two consecutive add+norm pairs across transformer layers.""" + model = ChainedModel() + bsz, seq_len, hidden = 2, 8, 128 + embed = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + attn_out = torch.randn_like(embed) + mlp_out = torch.randn_like(embed) + + gm = _export_model(model, embed, attn_out, mlp_out) - # Validation - # Clone inputs because the fused op is inplace - x_in = x.clone() - res_in = residual.clone() + # Before: 2 add+norm fusible pairs + assert _count_rms_norm_ops(gm) == 2 - # The fused op is inplace, so inputs x_in and res_in will be modified. - # gm_transformed returns (x_in, res_in) which are the modified tensors. - y_transformed = gm_transformed(x_in, res_in) + gm_t, info = _apply_transform_direct(gm) - y_model = model(x.clone(), residual.clone()) - torch.testing.assert_close(y_transformed[0], y_model[0], atol=1e-2, rtol=1e-2) - torch.testing.assert_close(y_transformed[1], y_model[1], atol=1e-2, rtol=1e-2) + # Structure check + assert info.num_matches == 2, f"Expected 2 matches, got {info.num_matches}" + assert _count_fused_ops(gm_t) == 2, "Expected 2 fused ops" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + # Verify second fused op receives add_out from first fused op (residual chain) + fused_nodes = [ + n + for n in gm_t.graph.nodes + if n.op == "call_function" and n.target is flashinfer_fused_add_rms_norm + ] + assert len(fused_nodes) == 2 + # The second fused op's residual arg should be a getitem from the first fused op + second_residual_arg = fused_nodes[1].args[1] + assert ( + second_residual_arg.op == "call_function" + and second_residual_arg.target is operator.getitem + and second_residual_arg.args[0] is fused_nodes[0] + ), "Second fused op's residual should come from first fused op's add_out" -def test_fuse_add_rms_norm(): - model = TestModel() - _run_test(model) + # Numerical check + y_fused = gm_t(embed.clone(), attn_out.clone(), mlp_out.clone()) + y_ref = model(embed.clone(), attn_out.clone(), mlp_out.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py new file mode 100644 index 000000000000..0687b1fd53d6 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py @@ -0,0 +1,378 @@ +"""Tests for NVFP4 quantized SwiGLU pattern matching and fusion transforms. + +Tests the parallel NVFP4 SwiGLU path: +1. match_nvfp4_swiglu_pattern: Matches torch_fake_quant_nvfp4_linear SwiGLU -> torch_nvfp4_swiglu_mlp +2. fuse_nvfp4_swiglu: Fuses gate+up FP4 weights -> fused_nvfp4_swiglu_mlp +""" + +import pytest +import torch +import torch.nn as nn +from _torch_test_utils import fp4_compatible, trtllm_ops_available +from torch.export import Dim + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer +from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op +from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import fp4_global_scale + +_skip_reason = "Requires NVFP4 (Blackwell+) and TRT-LLM ops" +_skip_condition = not (fp4_compatible() and trtllm_ops_available()) + + +class NVFP4SwiGLUMLP(nn.Module): + """SwiGLU MLP using NVFP4 quantized linear ops. + + Mimics the graph structure produced by quantize_nvfp4_linear_from_config + applied to a standard SwiGLU MLP: silu(gate(x)) * up(x) -> down(hidden). + """ + + def __init__(self, hidden_size: int = 128, intermediate_size: int = 128): + super().__init__() + device = torch.device("cuda") + scaling_vector_size = 16 + + # Create random weights and quantize them to FP4 + gate_weight = ( + torch.randn(intermediate_size, hidden_size, dtype=torch.half, device=device) * 0.05 + ) + up_weight = ( + torch.randn(intermediate_size, hidden_size, dtype=torch.half, device=device) * 0.05 + ) + down_weight = ( + torch.randn(hidden_size, intermediate_size, dtype=torch.half, device=device) * 0.05 + ) + + # Quantize gate projection + s_w_gate = fp4_global_scale(gate_weight) + gate_fp4, gate_cutlass = torch.ops.trtllm.fp4_quantize( + gate_weight, s_w_gate, scaling_vector_size, False + ) + # Use a shared input scale for gate and up (same input x) + s_in = fp4_global_scale(torch.randn(1, hidden_size, dtype=torch.half, device=device)) + gate_alpha = (1.0 / (s_in * s_w_gate)).to(torch.float32) + + self.register_buffer("gate_weight", gate_fp4) + self.register_buffer("gate_input_scale", s_in.to(torch.float32)) + self.register_buffer("gate_weight_scale", gate_cutlass) + self.register_buffer("gate_alpha", gate_alpha) + + # Quantize up projection (same input scale as gate) + s_w_up = fp4_global_scale(up_weight) + up_fp4, up_cutlass = torch.ops.trtllm.fp4_quantize( + up_weight, s_w_up, scaling_vector_size, False + ) + up_alpha = (1.0 / (s_in * s_w_up)).to(torch.float32) + + self.register_buffer("up_weight", up_fp4) + self.register_buffer("up_input_scale", s_in.to(torch.float32)) + self.register_buffer("up_weight_scale", up_cutlass) + self.register_buffer("up_alpha", up_alpha) + + # Quantize down projection (different input: the hidden state) + s_in_down = fp4_global_scale( + torch.randn(1, intermediate_size, dtype=torch.half, device=device) + ) + s_w_down = fp4_global_scale(down_weight) + down_fp4, down_cutlass = torch.ops.trtllm.fp4_quantize( + down_weight, s_w_down, scaling_vector_size, False + ) + down_alpha = (1.0 / (s_in_down * s_w_down)).to(torch.float32) + + self.register_buffer("down_weight", down_fp4) + self.register_buffer("down_input_scale", s_in_down.to(torch.float32)) + self.register_buffer("down_weight_scale", down_cutlass) + self.register_buffer("down_alpha", down_alpha) + + def forward(self, x): + gate_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + x, + self.gate_weight, + None, + [self.gate_input_scale], + [self.gate_weight_scale, self.gate_alpha], + [], + [], + ) + up_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + x, + self.up_weight, + None, + [self.up_input_scale], + [self.up_weight_scale, self.up_alpha], + [], + [], + ) + hidden = torch.nn.functional.silu(gate_out) * up_out + return torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + hidden, + self.down_weight, + None, + [self.down_input_scale], + [self.down_weight_scale, self.down_alpha], + [], + [], + ) + + +class NVFP4SwiGLUTestModel(nn.Module): + """Test model wrapping NVFP4 SwiGLU MLP between linear layers.""" + + def __init__(self, hidden_size: int = 128, intermediate_size: int = 128): + super().__init__() + device = torch.device("cuda") + self.linear_in = nn.Linear(hidden_size, hidden_size, device=device, dtype=torch.float16) + self.mlp = NVFP4SwiGLUMLP(hidden_size, intermediate_size) + self.linear_out = nn.Linear(hidden_size, hidden_size, device=device, dtype=torch.float16) + + def forward(self, x): + x = self.linear_in(x) + x = self.mlp(x) + x = self.linear_out(x) + return x + + +class NVFP4SwiGLUMultiLayerModel(nn.Module): + """Test model with multiple NVFP4 SwiGLU MLP layers.""" + + def __init__( + self, + hidden_size: int = 128, + intermediate_size: int = 128, + num_layers: int = 2, + ): + super().__init__() + device = torch.device("cuda") + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + nn.ModuleDict( + { + "linear": nn.Linear( + hidden_size, + hidden_size, + device=device, + dtype=torch.float16, + ), + "mlp": NVFP4SwiGLUMLP(hidden_size, intermediate_size), + } + ) + ) + + def forward(self, x): + for layer in self.layers: + x = layer["linear"](x) + x = layer["mlp"](x) + return x + + +# -- Test helpers -------------------------------------------------------------- + + +def _count_ops(gm, op): + """Count how many nodes in the graph match the given op.""" + return sum(1 for n in gm.graph.nodes if is_op(n, op)) + + +def _has_no_fake_quant_nvfp4(gm): + """Verify no torch_fake_quant_nvfp4_linear ops remain.""" + return _count_ops(gm, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear) == 0 + + +# -- Tests --------------------------------------------------------------------- + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +def test_nvfp4_swiglu_pattern_match_only(): + """Test that match_nvfp4_swiglu_pattern produces torch_nvfp4_swiglu_mlp op.""" + torch.manual_seed(0) + model = NVFP4SwiGLUMLP().to("cuda") + x = torch.randn(2, 128, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + + # Verify the graph has torch_fake_quant_nvfp4_linear ops before transform + assert _count_ops(gm, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear) == 3, ( + "Expected 3 torch_fake_quant_nvfp4_linear ops (gate, up, down) before transform" + ) + + # Apply only pattern matching + gm_matched = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + + # Check the intermediate op is present + nvfp4_swiglu_count = _count_ops( + gm_matched, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default + ) + assert nvfp4_swiglu_count == 1, ( + f"Expected 1 torch_nvfp4_swiglu_mlp op, got {nvfp4_swiglu_count}" + ) + + # All 3 fake_quant_nvfp4 ops should be consumed + assert _has_no_fake_quant_nvfp4(gm_matched), ( + "torch_fake_quant_nvfp4_linear ops should be consumed by pattern matcher" + ) + + # Verify numerical correctness + gm_matched = gm_matched.to("cuda") + y_matched = gm_matched(x) + y_model = model(x) + torch.testing.assert_close(y_matched, y_model, atol=1e-3, rtol=1e-3) + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +def test_nvfp4_swiglu_full_fusion(): + """Test full pipeline: pattern match -> fuse -> fused_nvfp4_swiglu_mlp.""" + torch.manual_seed(0) + model = NVFP4SwiGLUTestModel().to("cuda") + x = torch.randn(2, 128, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True, dynamic_shapes=({0: Dim.DYNAMIC},)) + + # Apply pattern matching + fusion + gm_fused = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + "fuse_nvfp4_swiglu": { + "stage": "post_load_fusion", + }, + }, + )(None, gm) + + gm_fused = gm_fused.to("cuda") + + # Check the fused op is present + fused_count = _count_ops(gm_fused, torch.ops.auto_deploy.fused_nvfp4_swiglu_mlp.default) + assert fused_count == 1, f"Expected 1 fused_nvfp4_swiglu_mlp op, got {fused_count}" + + # No intermediate or unfused ops should remain + assert _count_ops(gm_fused, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default) == 0, ( + "Intermediate torch_nvfp4_swiglu_mlp should be replaced by fused version" + ) + assert _has_no_fake_quant_nvfp4(gm_fused), ( + "No torch_fake_quant_nvfp4_linear ops should remain after fusion" + ) + + # Verify numerical correctness (fused uses TRT-LLM kernel, allow wider tolerance) + y_fused = gm_fused(x) + y_model = model(x) + torch.testing.assert_close(y_fused, y_model, atol=0.15, rtol=0.05) + + # Test with a different batch size to verify dynamic shapes work + x2 = torch.randn(4, 128, device="cuda", dtype=torch.float16) + y_fused_2 = gm_fused(x2) + y_model_2 = model(x2) + torch.testing.assert_close(y_fused_2, y_model_2, atol=0.15, rtol=0.05) + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +@pytest.mark.parametrize("num_layers", [2, 3]) +def test_nvfp4_swiglu_fusion_multiple_layers(num_layers): + """Test that multiple NVFP4 SwiGLU patterns are fused correctly.""" + torch.manual_seed(0) + model = NVFP4SwiGLUMultiLayerModel(num_layers=num_layers).to("cuda") + x = torch.randn(2, 128, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + + # Apply pattern matching + fusion + gm_fused = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + "fuse_nvfp4_swiglu": { + "stage": "post_load_fusion", + }, + }, + )(None, gm) + + gm_fused = gm_fused.to("cuda") + + # Check that all layers are fused + fused_count = _count_ops(gm_fused, torch.ops.auto_deploy.fused_nvfp4_swiglu_mlp.default) + assert fused_count == num_layers, ( + f"Expected {num_layers} fused_nvfp4_swiglu_mlp ops, got {fused_count}" + ) + + # Verify numerical correctness + y_fused = gm_fused(x) + y_model = model(x) + torch.testing.assert_close(y_fused, y_model, atol=0.2, rtol=0.1) + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +def test_nvfp4_swiglu_does_not_match_non_swiglu(): + """Test that the NVFP4 SwiGLU matcher does not match non-SwiGLU NVFP4 linears.""" + torch.manual_seed(0) + device = torch.device("cuda") + hidden_size = 128 + + # Model with two sequential NVFP4 linears + relu (NOT a SwiGLU pattern) + class NonSwiGLUModel(nn.Module): + def __init__(self): + super().__init__() + w1 = torch.randn(hidden_size, hidden_size, dtype=torch.half, device=device) * 0.05 + w2 = torch.randn(hidden_size, hidden_size, dtype=torch.half, device=device) * 0.05 + + s_in = fp4_global_scale(torch.randn(1, hidden_size, dtype=torch.half, device=device)) + s_w1 = fp4_global_scale(w1) + s_w2 = fp4_global_scale(w2) + + w1_fp4, w1_cutlass = torch.ops.trtllm.fp4_quantize(w1, s_w1, 16, False) + w2_fp4, w2_cutlass = torch.ops.trtllm.fp4_quantize(w2, s_w2, 16, False) + + self.register_buffer("w1", w1_fp4) + self.register_buffer("w1_is", s_in.to(torch.float32)) + self.register_buffer("w1_ws", w1_cutlass) + self.register_buffer("w1_a", (1.0 / (s_in * s_w1)).to(torch.float32)) + + self.register_buffer("w2", w2_fp4) + self.register_buffer("w2_is", s_in.to(torch.float32)) + self.register_buffer("w2_ws", w2_cutlass) + self.register_buffer("w2_a", (1.0 / (s_in * s_w2)).to(torch.float32)) + + def forward(self, x): + # Sequential linears without SwiGLU pattern + y = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + x, self.w1, None, [self.w1_is], [self.w1_ws, self.w1_a], [], [] + ) + y = torch.nn.functional.relu(y) + return torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + y, self.w2, None, [self.w2_is], [self.w2_ws, self.w2_a], [], [] + ) + + model = NonSwiGLUModel().to("cuda") + x = torch.randn(2, hidden_size, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + + gm_result = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + + # No SwiGLU ops should be found + assert _count_ops(gm_result, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default) == 0, ( + "Non-SwiGLU NVFP4 pattern should not match" + ) + + # Original NVFP4 linear ops should still be present + assert _count_ops(gm_result, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear) == 2, ( + "Original NVFP4 linear ops should be unchanged" + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py index 72a1f3ee10af..55e6a2d9563a 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py @@ -356,106 +356,98 @@ def _count_moe_experts(gm): # Real-model MOE export: GLM4 MoE Lite # --------------------------------------------------------------------------- -try: - from tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import ( - Glm4MoeLiteConfig, - Glm4MoeLiteForCausalLM, +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import ( # noqa: E402 + Glm4MoeLiteConfig, + Glm4MoeLiteForCausalLM, +) + + +def _make_tiny_glm4_config(n_routed_experts: int = 8) -> Glm4MoeLiteConfig: + """Create a minimal ``Glm4MoeLiteConfig`` suitable for unit tests.""" + return Glm4MoeLiteConfig( + vocab_size=256, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + q_lora_rank=32, + kv_lora_rank=32, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=16, + n_routed_experts=n_routed_experts, + n_shared_experts=1, + num_experts_per_tok=2, + moe_intermediate_size=64, + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + norm_topk_prob=True, + first_k_dense_replace=1, # layer 0 = dense MLP, layer 1 = MoE + max_position_embeddings=128, + rope_scaling=None, + pad_token_id=0, ) - _HAS_GLM4 = True -except ImportError: - _HAS_GLM4 = False - - -if _HAS_GLM4: - - def _make_tiny_glm4_config(n_routed_experts: int = 8) -> Glm4MoeLiteConfig: - """Create a minimal ``Glm4MoeLiteConfig`` suitable for unit tests.""" - return Glm4MoeLiteConfig( - vocab_size=256, - hidden_size=64, - intermediate_size=128, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=4, - q_lora_rank=32, - kv_lora_rank=32, - qk_nope_head_dim=12, - qk_rope_head_dim=4, - v_head_dim=16, - n_routed_experts=n_routed_experts, - n_shared_experts=1, - num_experts_per_tok=2, - moe_intermediate_size=64, - n_group=1, - topk_group=1, - routed_scaling_factor=1.0, - norm_topk_prob=True, - first_k_dense_replace=1, # layer 0 = dense MLP, layer 1 = MoE - max_position_embeddings=128, - rope_scaling=None, - pad_token_id=0, - ) - def _count_moe_experts_in_graph(gm: GraphModule) -> int: - """Return the number of experts in the first ``torch_moe`` call in *gm*.""" - for node in gm.graph.nodes: - if node.op == "call_function" and "torch_moe" in str(node.target): - return len(node.args[3]) # w1_weight list length - return 0 +def _count_moe_experts_in_graph(gm: GraphModule) -> int: + """Return the number of experts in the first ``torch_moe`` call in *gm*.""" + for node in gm.graph.nodes: + if node.op == "call_function" and "torch_moe" in str(node.target): + return len(node.args[3]) # w1_weight list length + return 0 + - @pytest.mark.skipif(not _HAS_GLM4, reason="GLM4 MoE Lite model not available on this branch") - @pytest.mark.skipif( - not torch.cuda.is_available(), reason="GLM4 MoE Lite requires CUDA (uses noaux_tc_op)" +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="GLM4 MoE Lite requires CUDA (uses noaux_tc_op)" +) +@pytest.mark.parametrize("n_routed_experts", [8, 16]) +@pytest.mark.parametrize("num_moe_experts_for_export", [2]) +def test_glm4_moe_lite_export_with_reduced_experts(n_routed_experts, num_moe_experts_for_export): + """Export a tiny ``Glm4MoeLiteForCausalLM`` with reduced experts and verify + that the expanded graph has the correct structure and accepts the original + state dict. + """ + # GLM4 MoE Lite uses noaux_tc_op which is CUDA-only, so we must use CUDA device + device = "cuda" + config = _make_tiny_glm4_config(n_routed_experts=n_routed_experts) + model = Glm4MoeLiteForCausalLM(config).to(device) + model.eval() + + input_ids = torch.randint(0, config.vocab_size, (1, 8), device=device) + position_ids = torch.arange(8, device=device).unsqueeze(0) + sample_kwargs = {"input_ids": input_ids, "position_ids": position_ids} + + # --- full export (baseline) --- + gm_full = torch_export_to_gm(model, kwargs=sample_kwargs) + + # --- export with reduced experts --- + gm_reduced = torch_export_to_gm( + model, + kwargs=sample_kwargs, + num_moe_experts_for_export=num_moe_experts_for_export, ) - @pytest.mark.parametrize("n_routed_experts", [8, 16]) - @pytest.mark.parametrize("num_moe_experts_for_export", [2]) - def test_glm4_moe_lite_export_with_reduced_experts( - n_routed_experts, num_moe_experts_for_export - ): - """Export a tiny ``Glm4MoeLiteForCausalLM`` with reduced experts and verify - that the expanded graph has the correct structure and accepts the original - state dict. - """ - # GLM4 MoE Lite uses noaux_tc_op which is CUDA-only, so we must use CUDA device - device = "cuda" - config = _make_tiny_glm4_config(n_routed_experts=n_routed_experts) - model = Glm4MoeLiteForCausalLM(config).to(device) - model.eval() - - input_ids = torch.randint(0, config.vocab_size, (1, 8), device=device) - position_ids = torch.arange(8, device=device).unsqueeze(0) - sample_kwargs = {"input_ids": input_ids, "position_ids": position_ids} - - # --- full export (baseline) --- - gm_full = torch_export_to_gm(model, kwargs=sample_kwargs) - - # --- export with reduced experts --- - gm_reduced = torch_export_to_gm( - model, - kwargs=sample_kwargs, - num_moe_experts_for_export=num_moe_experts_for_export, - ) - # Structural: both graphs must expose all experts - assert _count_moe_experts_in_graph(gm_full) == n_routed_experts - assert _count_moe_experts_in_graph(gm_reduced) == n_routed_experts - - # State-dict keys must match between full and reduced exports - full_keys = set(gm_full.state_dict().keys()) - reduced_keys = set(gm_reduced.state_dict().keys()) - assert full_keys == reduced_keys, ( - f"State-dict key mismatch.\n" - f" Only in full: {full_keys - reduced_keys}\n" - f" Only in reduced: {reduced_keys - full_keys}" - ) + # Structural: both graphs must expose all experts + assert _count_moe_experts_in_graph(gm_full) == n_routed_experts + assert _count_moe_experts_in_graph(gm_reduced) == n_routed_experts + + # State-dict keys must match between full and reduced exports + full_keys = set(gm_full.state_dict().keys()) + reduced_keys = set(gm_reduced.state_dict().keys()) + assert full_keys == reduced_keys, ( + f"State-dict key mismatch.\n" + f" Only in full: {full_keys - reduced_keys}\n" + f" Only in reduced: {reduced_keys - full_keys}" + ) - # Load the original model weights into the reduced export graph - gm_reduced.load_state_dict(model.state_dict(), strict=False) + # Load the original model weights into the reduced export graph + gm_reduced.load_state_dict(model.state_dict(), strict=False) - # Source model must be fully restored - for name, mod in model.named_modules(): - if hasattr(mod, "experts") and isinstance(mod.experts, nn.ModuleList): - assert len(mod.experts) == n_routed_experts, ( - f"Expert list in '{name}' was not restored to {n_routed_experts}" - ) + # Source model must be fully restored + for name, mod in model.named_modules(): + if hasattr(mod, "experts") and isinstance(mod.experts, nn.ModuleList): + assert len(mod.experts) == n_routed_experts, ( + f"Expert list in '{name}' was not restored to {n_routed_experts}" + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py new file mode 100644 index 000000000000..14eed45a63f9 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py @@ -0,0 +1,163 @@ +"""Tests for ``create_derived_custom_op`` in ``_graph.py``.""" + +import torch +from torch._subclasses import FakeTensorMode + +from tensorrt_llm._torch.auto_deploy.utils._graph import create_derived_custom_op + +# --------------------------------------------------------------------------- +# Helpers – tiny custom ops used as base ops for the tests +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("ad_test_derived::double", mutates_args=()) +def _double(x: torch.Tensor) -> torch.Tensor: + return x * 2 + + +@_double.register_fake +def _double_fake(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + +@torch.library.custom_op("ad_test_derived::weighted_add", mutates_args=()) +def _weighted_add(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor: + return x + alpha * y + + +@_weighted_add.register_fake +def _weighted_add_fake(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor: + return torch.empty_like(x) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestCreateDerivedCustomOp: + """Tests for the ``create_derived_custom_op`` utility.""" + + def test_basic_derived_op(self): + """A derived op should be callable and produce correct results.""" + + def make_impl(base_overload): + # Wrapper that calls the base op then negates the result. + def impl(*args, **kwargs): + return -base_overload(*args, **kwargs) + + return impl + + base_op = torch.ops.ad_test_derived.double + derived = create_derived_custom_op(base_op, "_neg", make_impl) + + x = torch.tensor([1.0, 2.0, 3.0]) + result = derived(x) + expected = -(x * 2) + torch.testing.assert_close(result, expected) + + def test_derived_op_is_registered(self): + """The derived op must be accessible via ``torch.ops``.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + create_derived_custom_op(torch.ops.ad_test_derived.double, "_registered", make_impl) + assert hasattr(torch.ops.ad_test_derived, "double_registered") + + def test_caching(self): + """Repeated calls with the same base_op + suffix must return the same object.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + op1 = create_derived_custom_op(torch.ops.ad_test_derived.double, "_cached", make_impl) + op2 = create_derived_custom_op(torch.ops.ad_test_derived.double, "_cached", make_impl) + assert op1 is op2 + + def test_different_suffix_produces_different_op(self): + """Different suffixes must create distinct ops.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + op_a = create_derived_custom_op(torch.ops.ad_test_derived.double, "_sfx_a", make_impl) + op_b = create_derived_custom_op(torch.ops.ad_test_derived.double, "_sfx_b", make_impl) + assert op_a is not op_b + + def test_default_fake_implementation(self): + """When *make_fake* is None the default (empty_like) must be used.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + derived = create_derived_custom_op( + torch.ops.ad_test_derived.double, "_dflt_fake", make_impl + ) + # Calling the Meta implementation via FakeTensorMode + with FakeTensorMode(): + x = torch.empty(4) + out = derived(x) + assert out.shape == x.shape + + def test_custom_fake_implementation(self): + """A user-supplied *make_fake* must override the default.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + # Fake that always returns shape (1,) regardless of input shape. + def make_fake(base_overload): + def fake(*args, **kwargs): + return args[0].new_empty(1) + + return fake + + derived = create_derived_custom_op( + torch.ops.ad_test_derived.double, + "_custom_fake", + make_impl, + make_fake=make_fake, + ) + + with FakeTensorMode(): + x = torch.empty(10) + out = derived(x) + assert out.shape == (1,) + + def test_preserves_schema_with_defaults(self): + """Derived op must preserve the base op's argument defaults.""" + + def make_impl(base_overload): + def impl(*args, **kwargs): + return base_overload(*args, **kwargs) * 10 + + return impl + + base_op = torch.ops.ad_test_derived.weighted_add + derived = create_derived_custom_op(base_op, "_x10", make_impl) + + x = torch.ones(3) + y = torch.ones(3) * 2.0 + + # With default alpha=1.0 → (x + 1.0*y) * 10 = 30 + result_default = derived(x, y) + torch.testing.assert_close(result_default, torch.full((3,), 30.0)) + + # With explicit alpha=0.5 → (x + 0.5*y) * 10 = 20 + result_alpha = derived(x, y, alpha=0.5) + torch.testing.assert_close(result_alpha, torch.full((3,), 20.0)) + + def test_accepts_op_overload(self): + """The function should accept an OpOverload (e.g. ``.default``) as well.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + 1 + + derived = create_derived_custom_op( + torch.ops.ad_test_derived.double.default, "_from_overload", make_impl + ) + + x = torch.tensor([5.0]) + # double → 10, +1 → 11 + torch.testing.assert_close(derived(x), torch.tensor([11.0]))