diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 8b179030177f..5c83fe79082a 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -59,14 +59,18 @@ prefill_cuda_graph_backend: breakable prefill_capture_num_tokens: [128, 256, 512] ``` -The breakable backend is experimental. The integration coverage in this change -includes BF16 Qwen3.5 on one GPU and NVFP4 DeepSeek models on multiple GPUs, -with both context-only and mixed context/decode batches using the KV cache. +The breakable backend is experimental. Integration coverage includes BF16 +Qwen3.5 on one GPU, NVFP4 DeepSeek models on multiple GPUs, and Qwen3.8 +Flash-Next NVFP4 with TP4/EP4, attention DP, MTP3, and PLE host offload. The +Qwen3.8 test covers 2,048- and 8,192-token graph buckets, including native QSA +sparse-attention dispatch for a long prompt; the other tests cover both +context-only and mixed context/decode batches using the KV cache. The following restrictions are enforced: -- `torch_compile_config`, LoRA, and multimodal models are rejected during - engine initialization. +- `torch_compile_config` and LoRA are rejected during engine initialization. + Multimodal wrappers are also rejected unless the multimodal encoder is + explicitly disabled and the wrapper exposes a decoder-only text model. - Speculative decoding is supported. - Context-logit requests run eagerly instead of replaying a breakable CUDA graph. diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/custom_ops.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/custom_ops.py new file mode 100644 index 000000000000..835ff56235ef --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/custom_ops.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""QSA custom operations for breakable CUDA graph capture.""" + +from typing import Optional + +import torch + +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph + +from ...interface import PredefinedAttentionMask +from .metadata import QSAAttentionMetadata +from .module import QSASparseHooks + + +def _extract_qsa_extra_attrs(layer_idx: str): + from tensorrt_llm._torch.attention.attention import Attention, extract_extra_attrs + + metadata, attention = extract_extra_attrs(layer_idx, "attn") + if not isinstance(metadata, QSAAttentionMetadata): + raise TypeError("QSA breakable CUDA graph received incompatible metadata") + if not isinstance(attention, Attention): + raise TypeError("QSA breakable CUDA graph requires an Attention layer") + if not isinstance(attention.sparse_attn_hooks, QSASparseHooks): + raise TypeError("QSA breakable CUDA graph requires QSA sparse hooks") + return metadata, attention + + +@torch.library.custom_op("trtllm::qsa_attn_inplace", mutates_args=("output",)) +def qsa_attn_inplace( + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + q_index: torch.Tensor, + token_k: torch.Tensor, + position_coordinates: torch.Tensor, + mrope_rotary_cos_sin: Optional[torch.Tensor], + mrope_position_deltas: Optional[torch.Tensor], + output_gate: Optional[torch.Tensor], + layer_idx: str, + output: torch.Tensor, +) -> None: + """Run live QSA dispatch and write the fixed physical output in place.""" + metadata, attention = _extract_qsa_extra_attrs(layer_idx) + output.zero_() + sparse_output = attention.sparse_attn_hooks.forward( + attention, + q, + k, + v, + metadata, + PredefinedAttentionMask.CAUSAL, + None, + None, + None, + None, + None, + 0, + False, + output_gate, + qsa_index_projection=(q_index, token_k, position_coordinates), + ) + if sparse_output is not None: + output[: sparse_output.shape[0]].copy_(sparse_output) + return + + dense_output, output_sf = attention._attn_impl( + q, + k, + v, + metadata, + PredefinedAttentionMask.CAUSAL, + mrope_rotary_cos_sin, + mrope_position_deltas, + None, + None, + output=output, + ) + if output_sf is not None: + raise RuntimeError("QSA output gating requires an unquantized attention output") + if output_gate is not None: + dense_output = attention.apply_output_gate( + dense_output, output_gate[: dense_output.shape[0]] + ) + output[: dense_output.shape[0]].copy_(dense_output) + + +@qsa_attn_inplace.register_fake +def _qsa_attn_inplace_fake( + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + q_index: torch.Tensor, + token_k: torch.Tensor, + position_coordinates: torch.Tensor, + mrope_rotary_cos_sin: Optional[torch.Tensor], + mrope_position_deltas: Optional[torch.Tensor], + output_gate: Optional[torch.Tensor], + layer_idx: str, + output: torch.Tensor, +) -> None: + """Model the in-place output mutation during fake-tensor propagation.""" + + +maybe_bcg_qsa_attn_inplace = eager_on_graph(qsa_attn_inplace) + + +__all__ = ["maybe_bcg_qsa_attn_inplace", "qsa_attn_inplace"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index 2473f5d1cc04..89ca3d472026 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -8,6 +8,7 @@ import torch from tensorrt_llm._torch.modules.top_k import TopK +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph from tensorrt_llm.logger import logger from ...interface import AttentionMask, PredefinedAttentionMask @@ -403,6 +404,18 @@ def select_qsa_paged_tokens( class QSASparseHooks(AttentionSparseHooks): """Keep the QSA side cache current and replace dense attention when useful.""" + def __init__(self) -> None: + self._num_sparse_prefill_dispatches = 0 + + @property + def num_sparse_prefill_dispatches(self) -> int: + """Return exact sparse prefill calls observed by this layer. + + The count includes graph capture and warmup calls. Runtime diagnostics + should compare snapshots when attributing calls to live requests. + """ + return self._num_sparse_prefill_dispatches + def forward( self, attention: "Attention", @@ -440,6 +453,11 @@ def forward( raise NotImplementedError("QSA sparse attention does not support LoRA") # The model's QKNormRoPEAttention preprocesses q/k before this hook; # mrope_config is therefore backend state, not work for the QSA kernel. + mrope_rotary_cos_sin = None + mrope_position_deltas = None + if mrope_config is not None: + mrope_rotary_cos_sin = mrope_config.get("mrope_rotary_cos_sin") + mrope_position_deltas = mrope_config.get("mrope_position_deltas") del mrope_config if not isinstance(attn_metadata, QSAAttentionMetadata): raise TypeError("QSA sparse attention received incompatible metadata") @@ -450,9 +468,40 @@ def forward( return None hidden_states = kwargs.get("qsa_index_hidden_states") position_ids = kwargs.get("qsa_position_ids") - if hidden_states is None or position_ids is None: + index_projection = kwargs.get("qsa_index_projection") + if index_projection is None and (hidden_states is None or position_ids is None): raise ValueError("QSA sparse attention requires hidden states and position IDs") + needs_fixed_output = is_in_breakable_cuda_graph() or q.shape[0] > num_tokens + if needs_fixed_output and index_projection is None: + # Keep the token-wise index projection captured. Cache mutation and + # live batch dispatch execute in the eager bridge below. The same + # bridge preserves physical padding during graph warmup, which runs + # before capture and must still return a bucket-sized output. + q_index, token_k, position_coordinates = attention.indexer.project( + hidden_states, position_ids + ) + output = q.new_empty((q.shape[0], attention.num_heads * attention.head_dim)) + from . import custom_ops + + custom_ops.maybe_bcg_qsa_attn_inplace( + q, + k, + v, + q_index, + token_k, + position_coordinates, + mrope_rotary_cos_sin, + mrope_position_deltas, + output_gate, + attention.layer_idx_str, + output, + ) + return output + + if output_gate is not None: + output_gate = output_gate[:num_tokens] + # The standard TRT-LLM backend owns quantized formats with auxiliary # scale pages, such as NVFP4. Use it until QSA kernels consume the data # and scale pages together; the cache manager itself remains generic. @@ -478,13 +527,23 @@ def forward( ) return None - hidden_states = hidden_states[:num_tokens] - q_index = attention.indexer.project_and_update_cache( - hidden_states, - position_ids, - attention.layer_idx, - attn_metadata, - ) + if index_projection is None: + hidden_states = hidden_states[:num_tokens] + q_index = attention.indexer.project_and_update_cache( + hidden_states, + position_ids, + attention.layer_idx, + attn_metadata, + ) + else: + q_index, token_k, position_coordinates = index_projection + q_index = q_index[:num_tokens] + attention.indexer.update_cache_and_compress( + attention.layer_idx, + token_k[:num_tokens], + position_coordinates[:num_tokens], + attn_metadata, + ) params = attention.sparse_params max_kv_len = int(attn_metadata.kv_lens_runtime[: attn_metadata.num_seqs].max()) @@ -613,6 +672,7 @@ def forward( output = attention.apply_output_gate(output, output_gate) return output + self._num_sparse_prefill_dispatches += 1 index_cache = attn_metadata.kv_cache_manager.get_index_k_buffer(attention.layer_idx) if index_cache is None: raise RuntimeError(f"QSA index cache is unavailable for layer {attention.layer_idx}") diff --git a/tensorrt_llm/_torch/models/modeling_qwen4_exp.py b/tensorrt_llm/_torch/models/modeling_qwen4_exp.py index c5350231a38f..376d797f8825 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen4_exp.py +++ b/tensorrt_llm/_torch/models/modeling_qwen4_exp.py @@ -14,6 +14,7 @@ # limitations under the License. """TensorRT-LLM text implementation for Qwen3.8-Flash-Next checkpoints.""" +import weakref from collections.abc import Sequence from typing import TYPE_CHECKING, Dict, List, Literal, Optional @@ -45,9 +46,10 @@ from ..modules.qwen4_exp.hyper_connection import HCResidual, Qwen4ExpHyperConnection from ..modules.qwen4_exp.ple import PLEMetadata, Qwen4ExpPLE from ..modules.rms_norm import RMSNorm +from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..pyexecutor.config_utils import get_qwen3_hybrid_layer_types, get_qwen4_exp_ple_layer_mask from ..speculative import SpecMetadata -from ..utils import AuxStreamType, EventType, create_lm_head_tp_mapping +from ..utils import AuxStreamType, EventType, create_lm_head_tp_mapping, get_model_extra_attrs from .checkpoints.base_weight_mapper import BaseWeightMapper from .modeling_qwen3 import Qwen3Attention from .modeling_qwen3_5 import _normalize_qwen35_exclude_modules, _normalize_qwen35_quant_config_dict @@ -70,6 +72,67 @@ _PendingCombine = tuple[_PendingBlockOutput, HCResidual] +def _extract_qwen4_exp_model(model_key: str): + extra_attrs = get_model_extra_attrs() + if extra_attrs is None: + raise RuntimeError("Model extra attrs are unavailable for Qwen4-Exp PLE") + model_refs = extra_attrs.get("qwen4_exp_models") + if model_refs is None or model_key not in model_refs: + raise RuntimeError(f"Qwen4-Exp model {model_key!r} is not registered") + model = model_refs[model_key]() + if model is None: + raise RuntimeError(f"Qwen4-Exp model {model_key!r} is no longer available") + metadata_ref = extra_attrs.get("attention_metadata") + if metadata_ref is None: + raise RuntimeError("Attention metadata is unavailable for Qwen4-Exp PLE") + return model, metadata_ref(), extra_attrs.get("spec_metadata") + + +@torch.library.custom_op("trtllm::qwen4_exp_ple_inplace", mutates_args=("output",)) +def qwen4_exp_ple_inplace( + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + model_key: str, + output: torch.Tensor, +) -> None: + """Prepare live PLE state and write its fixed physical output in place.""" + model, attn_metadata, spec_metadata = _extract_qwen4_exp_model(model_key) + ple_layer_idx = model.ple_layer_index + if ple_layer_idx is None: + raise RuntimeError("Qwen4-Exp PLE bridge requires a local PLE layer") + ple_module = model.layers[ple_layer_idx].ple + if ple_module is None: + raise RuntimeError("Qwen4-Exp PLE layer has no PLE module") + ple_state = model._prepare_ple_state( + attn_metadata, + input_ids, + getattr(attn_metadata, "mamba_metadata", None), + spec_metadata, + ) + if ple_state is None: + raise RuntimeError("Qwen4-Exp PLE state preparation returned no state") + ple_meta, conv_state, ngram_context = ple_state + + # Host-offload prefetch uses a side stream. Run the lookup synchronously in + # this eager bridge so replay does not depend on a dropped cross-stream wait. + output.zero_() + ple_output = ple_module(hidden_states, ple_meta, conv_state, ngram_context) + output[: ple_output.shape[0]].copy_(ple_output) + + +@qwen4_exp_ple_inplace.register_fake +def _qwen4_exp_ple_inplace_fake( + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + model_key: str, + output: torch.Tensor, +) -> None: + """Model the in-place output mutation during fake-tensor propagation.""" + + +maybe_bcg_qwen4_exp_ple_inplace = eager_on_graph(qwen4_exp_ple_inplace) + + def _qwen4_exp_tp_output_reduction_enabled(mapping) -> bool: """Return whether decoder outputs require a tensor-parallel reduction.""" return mapping.tp_size > 1 and not mapping.enable_attention_dp @@ -289,6 +352,8 @@ def forward( attn_metadata: AttentionMetadata, mamba_metadata: Optional[Mamba2Metadata] = None, ple_state: Optional[tuple] = None, + ple_input_ids: Optional[torch.Tensor] = None, + ple_model_key: Optional[str] = None, spec_metadata=None, lora_params: Optional[dict] = None, pending_combine: Optional[_PendingCombine] = None, @@ -336,11 +401,23 @@ def forward( hidden_states = torch.cat([hidden_states] * self.hc_count, dim=-1) # PLE n-gram short-conv side path, injected into the bundle before mix. - if self.ple is not None and ple_state is not None: - ple_meta, conv_state, ngram_context = ple_state - hidden_states = hidden_states + self.ple( - hidden_states, ple_meta, conv_state, ngram_context - ) + if self.ple is not None: + if ple_state is not None: + ple_meta, conv_state, ngram_context = ple_state + hidden_states = hidden_states + self.ple( + hidden_states, ple_meta, conv_state, ngram_context + ) + elif ple_input_ids is not None and is_in_breakable_cuda_graph(): + if ple_model_key is None: + raise RuntimeError("Qwen4-Exp PLE bridge requires a registered model key") + ple_output = torch.empty_like(hidden_states) + maybe_bcg_qwen4_exp_ple_inplace( + hidden_states, + ple_input_ids, + ple_model_key, + ple_output, + ) + hidden_states = hidden_states + ple_output # Attention block: mix -> mixer -> attn_tp_all_reduce -> combine. mixed, residual = self.attn_hyper_connection.mix(hidden_states) @@ -410,6 +487,13 @@ class Qwen4ExpModel(DecoderModel): def __init__(self, model_config: ModelConfig[PretrainedConfig]): super().__init__(model_config) + model_refs = model_config.extra_attrs.setdefault("qwen4_exp_models", {}) + self._bcg_model_key = "qwen4_exp" + suffix = 0 + while self._bcg_model_key in model_refs: + self._bcg_model_key = f"qwen4_exp_{suffix}" + suffix += 1 + model_refs[self._bcg_model_key] = weakref.ref(self) config = model_config.pretrained_config dtype = config.torch_dtype @@ -660,7 +744,7 @@ def _prepare_ple_state( state_indices, is_decode=is_decode, eos_token_id=self.eos_token_id, - physical_tokens=attn_metadata.num_tokens, + physical_tokens=(attn_metadata.padded_num_tokens or attn_metadata.num_tokens), num_contexts=num_contexts, use_spec_decoding=use_spec_decoding, uniform_row_width=uniform_row_width, @@ -816,9 +900,10 @@ def forward( raise ValueError( "Qwen4-Exp PLE requires original token IDs when inputs_embeds are provided" ) + use_breakable_cuda_graph = is_in_breakable_cuda_graph() ple_state = ( self._prepare_ple_state(attn_metadata, ple_input_ids, mamba_metadata, spec_metadata) - if ple_input_ids is not None + if ple_input_ids is not None and not use_breakable_cuda_graph else None ) prefetched_ple_module = None @@ -841,6 +926,12 @@ def forward( attn_metadata=attn_metadata, mamba_metadata=mamba_metadata, ple_state=ple_state if self.ple_layer_mask[layer_idx] else None, + ple_input_ids=( + ple_input_ids + if use_breakable_cuda_graph and self.ple_layer_mask[layer_idx] + else None + ), + ple_model_key=getattr(self, "_bcg_model_key", None), spec_metadata=spec_metadata, lora_params=lora_params, pending_combine=pending_combine, diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/ple.py b/tensorrt_llm/_torch/modules/qwen4_exp/ple.py index 0b7fb4b5f0c6..da575cdd53fd 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/ple.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/ple.py @@ -1048,6 +1048,12 @@ def _prepare_ngram_lookup( """Build recurrent contexts and the hashed per-head global row IDs.""" history = ngram_context.index_select(0, metadata.state_indices).to(torch.long) combined = torch.cat([history, metadata.padded_tokens.to(torch.long)], dim=1) + if metadata.processed_tokens == 0: + # An attention-DP rank can have no live requests while another rank + # selects a prefill graph bucket. It must still enter the row-sharded + # embedding collectives, but there are no local contexts to hash. + ngram_ids = combined.new_empty((0, self.ple_embedding.ngram_heads)) + return combined, ngram_ids windows = combined.unfold(1, self.ngram_size, 1) contexts = windows[metadata.req_indices, metadata.token_offsets] return combined, self.ple_embedding.hash_contexts( @@ -1379,7 +1385,7 @@ def forward( ): raise RuntimeError("PLE activations and recurrent-state pools must share a device") hidden_states = hidden_states[: m.processed_tokens] - if m.processed_tokens == 0: + if m.processed_tokens == 0 and not self.ple_embedding.use_attention_dp_sharding: # No recurrent state advances on an empty logical step. Preserve # the graph's physical row count with zero PLE contributions. return hidden_states.new_zeros((m.physical_tokens, hc_dim)) @@ -1394,6 +1400,12 @@ def forward( all_rank_num_tokens=m.all_rank_num_tokens, ).flatten(start_dim=-2) + if m.processed_tokens == 0: + # The lookup above pairs this empty rank with its attention-DP peers' + # AllGather/ReduceScatter. Discard the empty local result without + # advancing recurrent state. + return hidden_states.new_zeros((m.physical_tokens, hc_dim)) + key = self.key_proj(embeddings) value = self.value_proj(embeddings) token_count = hidden_states.shape[0] diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py index 9eccb529196d..1272d8f3ab9b 100644 --- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -38,6 +38,7 @@ def __init__(self, layer_model: nn.Module) -> None: self._state = BreakableCUDAGraphRunnerState.IDLE self._active_graph: Optional[BreakableCUDAGraph] = None self._active_num_tokens: Optional[int] = None + self._replay_counts: dict[int, int] = {} @property def state(self) -> BreakableCUDAGraphRunnerState: @@ -54,6 +55,20 @@ def is_capturing(self) -> bool: def has_graph(self, num_tokens: int) -> bool: return num_tokens in self._graphs + @property + def captured_token_buckets(self) -> list[int]: + """Return the token buckets with a successfully captured graph.""" + return sorted(self._graphs) + + @property + def replay_counts(self) -> dict[int, int]: + """Return replay calls submitted without a synchronous error. + + CUDA execution remains asynchronous; callers that require completion + must synchronize separately. Clearing the runner resets these counts. + """ + return self._replay_counts.copy() + def warmup(self, engine_forward: Callable[[], Any], steps: int = _WARMUP_STEPS) -> None: """Run the complete eager engine forward under the warmup state. model_engine.forward will use state to determine what forward to do.""" @@ -176,6 +191,7 @@ def replay(self, num_tokens: int) -> torch.Tensor: if num_tokens not in self._graphs: raise KeyError(f"No BCG captured for num_tokens={num_tokens}") self._graphs[num_tokens].replay() + self._replay_counts[num_tokens] = self._replay_counts.get(num_tokens, 0) + 1 return self._outputs[num_tokens] def execute(self, num_tokens: int, outer_forward: Callable[[], Any]) -> Any: @@ -209,5 +225,6 @@ def clear(self) -> None: graph.reset() self._graphs.clear() self._outputs.clear() + self._replay_counts.clear() self._shared_output = None self._memory_pool = None diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 4d8fb4b6343f..1d7114749719 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -1084,6 +1084,39 @@ def get_startup_metrics(self) -> dict: return startup_metrics + def get_prefill_cuda_graph_stats(self) -> dict: + """Return rank-local breakable prefill CUDA graph diagnostics. + + The snapshot covers the primary model only and is not atomic with + inference. QSA dispatches are summed across layers and include capture + and warmup, so callers should use snapshot deltas for live requests. + """ + stats = { + "captured_token_buckets": [], + "replay_counts": {}, + "qsa_sparse_prefill_dispatches": 0, + } + if not self._is_pytorch_backend or self.engine is None: + return stats + + model_engine = getattr(self.engine, "model_engine", None) + runner = getattr(model_engine, "breakable_cuda_graph_runner", None) + if runner is None: + return stats + + sparse_prefill_dispatches = 0 + model = getattr(model_engine, "model", None) + if model is not None: + for module in model.modules(): + hooks = getattr(module, "sparse_attn_hooks", None) + sparse_prefill_dispatches += getattr( + hooks, "num_sparse_prefill_dispatches", 0) + + stats["captured_token_buckets"] = runner.captured_token_buckets + stats["replay_counts"] = runner.replay_counts + stats["qsa_sparse_prefill_dispatches"] = sparse_prefill_dispatches + return stats + def start_profile(self, output_dir: Optional[str] = None, num_steps: Optional[int] = None, diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index b2e82b6e8a40..5617fe0fce23 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -61,9 +61,14 @@ # Methods that are explicitly implemented for multi-rank MPI/IPC executor # deployments and may be called via collective_rpc even when world_size > 1. # sleep() and wakeup() coordinate across all ranks via a dedicated control -# communicator started in PyExecutor.start_worker(); rank-0 is the entry -# point, so routing through the rank-0 RPC shim is correct. -_MULTI_RANK_ALLOWED_METHODS: frozenset[str] = frozenset({"sleep", "wakeup"}) +# communicator started in PyExecutor.start_worker(). Prefill graph diagnostics +# are explicitly rank-local. Rank 0 is the entry point for all three, so routing +# through the rank-0 RPC shim is correct. +_MULTI_RANK_COORDINATED_METHODS: frozenset[str] = frozenset({"sleep", "wakeup"}) +_MULTI_RANK_RANK_LOCAL_METHODS: frozenset[str] = frozenset( + {"get_prefill_cuda_graph_stats"}) +_MULTI_RANK_ALLOWED_METHODS = (_MULTI_RANK_COORDINATED_METHODS + | _MULTI_RANK_RANK_LOCAL_METHODS) def _check_collective_rpc_guard( @@ -90,8 +95,9 @@ def _check_collective_rpc_guard( raise NotImplementedError( f"MPI collective_rpc does not support model_world_size > 1 for " f"method '{method}'; use the Ray executor for general multi-rank " - "deployments, or use sleep()/wakeup() which handle multi-rank " - "coordination internally.") + "deployments. The MPI shim supports only sleep()/wakeup(), which " + "coordinate internally, and get_prefill_cuda_graph_stats(), " + "which reads rank 0 only.") if unique_reply_rank is not None or target_ranks is not None: raise NotImplementedError( "unique_reply_rank and target_ranks are not supported; " @@ -1186,11 +1192,13 @@ def collective_rpc( Rank-0 RPC shim used for uniform dispatch from :meth:`~tensorrt_llm.llmapi.llm.LLM._collective_rpc`. Most methods - require ``model_world_size == 1``; ``sleep`` and ``wakeup`` are + require ``model_world_size == 1``. ``sleep`` and ``wakeup`` are explicitly allowed for multi-rank deployments because :meth:`~tensorrt_llm.executor.base_worker.BaseWorker.sleep` and :meth:`~tensorrt_llm.executor.base_worker.BaseWorker.wakeup` handle - cross-rank coordination internally via the control communicator. + cross-rank coordination internally via the control communicator; + ``get_prefill_cuda_graph_stats`` is allowed because it is explicitly + rank-local. Args: method: Name of the worker method to invoke. @@ -1208,9 +1216,8 @@ def collective_rpc( Raises: RuntimeError: If the RPC client has not been initialised yet. NotImplementedError: If ``model_world_size > 1`` and the method - is not in the allowed-methods set (currently ``sleep`` and - ``wakeup``), or if ``unique_reply_rank`` or ``target_ranks`` - are provided. + is not in the allowed-methods set, or if + ``unique_reply_rank`` or ``target_ranks`` are provided. """ if self.rpc_client is None: raise RuntimeError( diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 2b0bb26a6d02..176b54d4304d 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -226,9 +226,9 @@ def collective_rpc( ) -> List: """Execute a method call on the rank-0 RpcWorker via the RPC client. - Rank-0 RPC shim. ``sleep`` and ``wakeup`` are allowed for - ``model_world_size > 1``; all other methods require - ``model_world_size == 1``. See + Rank-0 RPC shim. ``sleep``, ``wakeup``, and rank-local prefill graph + diagnostics are allowed for ``model_world_size > 1``; all other + methods require ``model_world_size == 1``. See :meth:`~tensorrt_llm.executor.proxy.GenerationExecutorProxy.collective_rpc` for details. @@ -249,9 +249,8 @@ def collective_rpc( Raises: NotImplementedError: If ``model_world_size > 1`` and the method - is not in the allowed-methods set (currently ``sleep`` and - ``wakeup``), or if ``unique_reply_rank`` or ``target_ranks`` - are provided. + is not in the allowed-methods set, or if + ``unique_reply_rank`` or ``target_ranks`` are provided. """ _check_collective_rpc_guard(self.model_world_size, unique_reply_rank, diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 46afc55d3e90..9d9441053763 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -54,6 +54,9 @@ TestQwen3_5_35B_A3B::test_bf16_mtp: TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload: ref_al: 1.890148 min_al: 1.7956406 +TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload_breakable_prefill_cuda_graph: + ref_al: 3.34331714337111 + min_al: 3.1761512862025545 # The usual 95% tolerance for these low values is below the no-acceptance # floor of 1.0, so require both modes to accept at least some draft tokens. TestLlama4SpeculativeDecoding::test_llama4_eagle3[linear]: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index f6db6870ab5a..345721aea9da 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -35,7 +35,7 @@ AttentionDpConfig, CudaGraphConfig, DFlashDecodingConfig, DSparkDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MambaStateConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, - PrefillCudaGraphBackend, SamplingParams, SchedulerConfig, + PrefillCudaGraphBackend, SamplingParams, SchedulerConfig, SchedulingParams, SkipSoftmaxAttentionConfig, SAEnhancerConfig, TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -5817,7 +5817,11 @@ def _build_llm(self, moe_expert_parallel_size: int = 1, enable_attention_dp: bool = False, cover_guided_decoding: bool = False, - enable_block_reuse: bool = False) -> LLM: + enable_block_reuse: bool = False, + prefill_cuda_graph_backend: + PrefillCudaGraphBackend = PrefillCudaGraphBackend.DISABLED, + prefill_capture_num_tokens: Optional[list[int]] = None, + disable_mm_encoder: bool = False) -> LLM: """Construct the engine shared by both evaluation tasks.""" kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5, enable_block_reuse=enable_block_reuse, @@ -5835,6 +5839,13 @@ def _build_llm(self, guided_decoding_args = ({ "guided_decoding_backend": "xgrammar" } if cover_guided_decoding else {}) + prefill_cuda_graph_args = {} + if prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED: + assert prefill_capture_num_tokens + prefill_cuda_graph_args = dict( + prefill_cuda_graph_backend=prefill_cuda_graph_backend, + prefill_capture_num_tokens=prefill_capture_num_tokens, + ) return LLM(model_path, trust_remote_code=True, tensor_parallel_size=tensor_parallel_size, @@ -5847,8 +5858,10 @@ def _build_llm(self, cuda_graph_config=cuda_graph_config, moe_config=MoeConfig(backend=moe_backend), speculative_config=mtp_config, + disable_mm_encoder=disable_mm_encoder, **stats_args, - **guided_decoding_args) + **guided_decoding_args, + **prefill_cuda_graph_args) def _run_evals(self, model_path: str, @@ -5861,27 +5874,94 @@ def _run_evals(self, moe_expert_parallel_size: int = 1, enable_attention_dp: bool = False, cover_guided_decoding: bool = False, - enable_block_reuse: bool = False) -> None: + enable_block_reuse: bool = False, + prefill_cuda_graph_backend: + PrefillCudaGraphBackend = PrefillCudaGraphBackend.DISABLED, + prefill_capture_num_tokens: Optional[list[int]] = None, + disable_mm_encoder: bool = False) -> None: if not os.path.exists(model_path): pytest.skip(f"Model directory {model_path} does not exist") monkeypatch.setenv("TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD", "1") - check_acceptance_length = ( + mtp_adp_case = (tensor_parallel_size == 4 and moe_backend == "TRTLLM" + and max_draft_len == 3 and moe_expert_parallel_size == 4 + and enable_attention_dp and not enable_block_reuse) + check_acceptance_length = mtp_adp_case and ( expected_quant_algo == QuantAlgo.FP8_BLOCK_SCALES - and tensor_parallel_size == 4 and moe_backend == "TRTLLM" - and max_draft_len == 3 and moe_expert_parallel_size == 4 - and enable_attention_dp and not enable_block_reuse) - with self._build_llm(model_path, - tensor_parallel_size, - moe_backend, - max_draft_len, - check_acceptance_length, - moe_expert_parallel_size=moe_expert_parallel_size, - enable_attention_dp=enable_attention_dp, - cover_guided_decoding=cover_guided_decoding, - enable_block_reuse=enable_block_reuse) as llm: + or prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE) + + parity_sampling_params = SamplingParams(max_tokens=8, + seed=42, + temperature=0, + ignore_eos=True, + detokenize=False, + add_special_tokens=False) + + scheduling_params = (SchedulingParams(attention_dp_rank=0, + attention_dp_relax=False) + if enable_attention_dp else None) + + def parity_output(llm: LLM, token_id: int, + prompt_length: int) -> list[int]: + output = llm.generate([[token_id] * prompt_length], + sampling_params=parity_sampling_params, + scheduling_params=scheduling_params)[0] + return output.outputs[0].token_ids + + def parity_outputs(llm: LLM) -> list[list[int]]: + return [ + parity_output(llm, token_id, prompt_length) + for token_id, prompt_length in ((17, 1536), (23, 6144)) + ] + + def prefill_graph_stats(llm: LLM) -> dict: + replies = llm._collective_rpc("get_prefill_cuda_graph_stats") + assert len(replies) == 1 + return replies[0] + + eager_parity = None + if prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: + with self._build_llm( + model_path, + tensor_parallel_size, + moe_backend, + max_draft_len, + False, + moe_expert_parallel_size=moe_expert_parallel_size, + enable_attention_dp=enable_attention_dp, + disable_mm_encoder=disable_mm_encoder) as eager_llm: + eager_parity = parity_outputs(eager_llm) + with self._build_llm( + model_path, + tensor_parallel_size, + moe_backend, + max_draft_len, + check_acceptance_length, + moe_expert_parallel_size=moe_expert_parallel_size, + enable_attention_dp=enable_attention_dp, + cover_guided_decoding=cover_guided_decoding, + enable_block_reuse=enable_block_reuse, + prefill_cuda_graph_backend=prefill_cuda_graph_backend, + prefill_capture_num_tokens=prefill_capture_num_tokens, + disable_mm_encoder=disable_mm_encoder) as llm: assert llm.args.quant_config.quant_algo == expected_quant_algo + if eager_parity is not None: + before = prefill_graph_stats(llm) + short_output = parity_output(llm, 17, 1536) + after_short = prefill_graph_stats(llm) + long_output = parity_output(llm, 23, 6144) + after_long = prefill_graph_stats(llm) + assert [short_output, long_output] == eager_parity + assert after_long["captured_token_buckets"] == [2048, 8192] + assert (after_short["replay_counts"][2048] + > before["replay_counts"].get(2048, 0)) + assert (after_long["replay_counts"][8192] + > after_short["replay_counts"].get(8192, 0)) + assert (after_short["qsa_sparse_prefill_dispatches"] == + before["qsa_sparse_prefill_dispatches"]) + assert (after_long["qsa_sparse_prefill_dispatches"] + > after_short["qsa_sparse_prefill_dispatches"]) if cover_guided_decoding: assert_guided_decoding_regex(llm) mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", @@ -5890,9 +5970,12 @@ def _run_evals(self, task.evaluate(llm, extra_evaluator_kwargs=self.GSM8K_EVALUATOR_KWARGS) if check_acceptance_length: + test_name = ("test_fp8_adp4_mtp3_trtllm_ple_offload" if + expected_quant_algo == QuantAlgo.FP8_BLOCK_SCALES + else "test_nvfp4_adp4_mtp3_trtllm_ple_offload_" + "breakable_prefill_cuda_graph") assert_acceptance_length_for_llm( - "TestQwen3_8_Flash_Next::" - "test_fp8_adp4_mtp3_trtllm_ple_offload", + f"TestQwen3_8_Flash_Next::{test_name}", llm, ) task = MMLU(self.MODEL_NAME) @@ -5961,19 +6044,22 @@ def test_fp8_adp4_mtp3_prefix_cache(self, monkeypatch: pytest.MonkeyPatch, @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(70000) @pytest.mark.skip_less_host_memory(131072) - def test_nvfp4_adp4_mtp3_trtllm_ple_offload(self, - monkeypatch: pytest.MonkeyPatch, - mocker) -> None: - """NVFP4 on four GPUs with attention DP, MTP3 and PLE offload.""" - self._run_evals(f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4", - tensor_parallel_size=4, - moe_backend="TRTLLM", - max_draft_len=3, - expected_quant_algo=QuantAlgo.MIXED_PRECISION, - monkeypatch=monkeypatch, - mocker=mocker, - moe_expert_parallel_size=4, - enable_attention_dp=True) + def test_nvfp4_adp4_mtp3_trtllm_ple_offload_breakable_prefill_cuda_graph( + self, monkeypatch: pytest.MonkeyPatch, mocker) -> None: + """NVFP4, MTP3 and PLE offload with breakable prefill CUDA graphs.""" + self._run_evals( + f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4", + tensor_parallel_size=4, + moe_backend="TRTLLM", + max_draft_len=3, + expected_quant_algo=QuantAlgo.MIXED_PRECISION, + monkeypatch=monkeypatch, + mocker=mocker, + moe_expert_parallel_size=4, + enable_attention_dp=True, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + prefill_capture_num_tokens=[2048, 8192], + disable_mm_encoder=True) @skip_pre_blackwell @pytest.mark.skip_less_device_memory(100000) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 24818d8a8ef4..152c5d4187b0 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -610,7 +610,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_bf16_tep4_cutlass accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_prefix_cache accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_1gpu_cutedsl_ple_offload -accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload +accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload_breakable_prefill_cuda_graph accuracy/test_llm_api_pytorch.py::TestSeedOss_36B::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestStep3p7SpeculativeDecoding::test_guided_decoding TIMEOUT (60) accuracy/test_llm_api_pytorch.py::TestStep3p7SpeculativeDecoding::test_step3p7_mtp TIMEOUT (60) diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index 120a5e82064e..c25667418fb4 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -21,7 +21,7 @@ l0_gb300: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v2_kv_cache-True-True-cutlass-auto] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] # Cover nvbugs 5461712 and 5505402 - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload - - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload + - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload_breakable_prefill_cuda_graph - accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache] - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py index 7c38f9e328ad..fef5afea1362 100644 --- a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -260,6 +260,8 @@ def engine_forward(): assert body.forward_calls == 6 assert logits_processor.forward_calls == 6 assert runner._shared_output is first_shared_output + assert runner.captured_token_buckets == [4, 8] + assert runner.replay_counts == {} original_forward = body.forward inputs["value"].fill_(3) @@ -270,6 +272,11 @@ def engine_forward(): assert body.forward_calls == 6 assert logits_processor.forward_calls == 7 assert body.forward == original_forward + assert runner.replay_counts == {4: 1} + + runner.clear() + assert runner.captured_token_buckets == [] + assert runner.replay_counts == {} def test_runner_first_bucket_segments_share_one_memory_pool(): @@ -311,10 +318,19 @@ def test_runner_graph_miss_nested_execute_and_exception_recovery(): with pytest.raises(KeyError, match="No BCG captured"): runner.execute(4, lambda: None) - runner._graphs[4] = object() + class FailingGraph: + @staticmethod + def replay(): + raise RuntimeError("replay failed") + + runner._graphs[4] = FailingGraph() runner._outputs[4] = torch.zeros(1, device="cuda") original_forward = body.forward + with pytest.raises(RuntimeError, match="replay failed"): + runner.replay(4) + assert runner.replay_counts == {} + def nested(): return runner.execute(4, lambda: None) diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index b98e515c1711..9354a7d96595 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -165,6 +165,281 @@ def test_qsa_empty_batch_keeps_the_regular_backend_path() -> None: assert output is None +@pytest.mark.parametrize( + "num_contexts,max_kv_len,expected_dispatches,expect_sparse_output", + ((1, 4, 0, False), (0, 8, 0, True), (1, 8, 1, True)), + ids=("below-threshold", "decode", "exact-sparse-prefill"), +) +def test_qsa_sparse_prefill_dispatch_counter( + monkeypatch: pytest.MonkeyPatch, + num_contexts: int, + max_kv_len: int, + expected_dispatches: int, + expect_sparse_output: bool, +) -> None: + from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels, module + + head_dim = 4 + num_heads = 4 + metadata = object.__new__(QSAAttentionMetadata) + metadata._num_tokens = 1 + metadata._num_contexts = num_contexts + metadata._seq_lens = torch.ones(1, dtype=torch.int32) + metadata.kv_lens_runtime = torch.tensor([max_kv_len], dtype=torch.int32) + metadata.qsa_req_idx_per_token = torch.zeros(1, dtype=torch.int32) + metadata.qsa_logical_positions = torch.tensor([max_kv_len - 1], dtype=torch.int64) + metadata.qsa_sequence_lengths = torch.tensor([max_kv_len], dtype=torch.int32) + metadata.qsa_visible_blocks = torch.ones(1, dtype=torch.int32) + metadata.qsa_topk_indices = torch.zeros((1, 7), dtype=torch.int32) + metadata.qsa_topk_row_starts = torch.zeros(1, dtype=torch.int32) + metadata.qsa_block_table = torch.arange(max_kv_len, dtype=torch.int32).unsqueeze(0) + kv_pool = torch.zeros((max_kv_len, 2, 1, 1, head_dim)) + index_cache = torch.zeros((max_kv_len, 1, 1, head_dim)) + metadata.kv_cache_manager = SimpleNamespace( + dtype=DataType.BF16, + tokens_per_block=1, + get_buffers=lambda *args, **kwargs: kv_pool, + get_index_k_buffer=lambda *args, **kwargs: index_cache, + ) + + indexer = SimpleNamespace( + project_and_update_cache=lambda *args, **kwargs: torch.zeros((1, num_heads, head_dim)), + top_k=None, + ) + attention = SimpleNamespace( + head_dim=head_dim, + num_heads=num_heads, + num_key_value_heads=1, + q_scaling=1.0, + layer_idx=0, + indexer=indexer, + sparse_params=QSASparseParams( + index_n_heads=num_heads, + index_kv_heads=1, + index_head_dim=head_dim, + token_topk=4, + compress_ratio=4, + seq_len_threshold=4, + ), + split_qkv=lambda *args: ( + torch.zeros((1, num_heads, head_dim)), + torch.zeros((1, 1, head_dim)), + torch.zeros((1, 1, head_dim)), + ), + ) + monkeypatch.setattr(kernels, "qsa_supports_head_dims", lambda *args: True) + monkeypatch.setattr( + module, + "select_qsa_paged_tokens", + lambda *args, **kwargs: torch.zeros((1, 7), dtype=torch.int32), + ) + monkeypatch.setattr(module, "qsa_sparse_gqa", lambda **kwargs: kwargs["q"]) + hooks = QSASparseHooks() + + output = hooks.forward( + attention=attention, + q=torch.zeros((1, head_dim)), + k=None, + v=None, + attn_metadata=metadata, + attention_mask=PredefinedAttentionMask.CAUSAL, + attention_window_size=None, + attention_mask_data=None, + mrope_config=None, + attention_sinks=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + has_lora=False, + output_gate=None, + qsa_index_hidden_states=torch.zeros((1, head_dim)), + qsa_position_ids=torch.zeros(1, dtype=torch.int32), + ) + + assert (output is not None) == expect_sparse_output + assert hooks.num_sparse_prefill_dispatches == expected_dispatches + + +@pytest.mark.parametrize("in_graph", [False, True], ids=["padded-warmup", "capture"]) +def test_qsa_fixed_output_bridge_captures_index_projection( + monkeypatch: pytest.MonkeyPatch, + in_graph: bool, +) -> None: + from tensorrt_llm._torch.attention.backends.sparse.qsa import custom_ops, module + + metadata = object.__new__(QSAAttentionMetadata) + metadata._num_tokens = 2 + hidden_states = torch.randn(4, 16) + position_ids = torch.arange(4, dtype=torch.int32) + q_index = torch.randn(4, 4, 8) + token_k = torch.randn(4, 1, 8) + position_coordinates = torch.arange(12, dtype=torch.int32).reshape(4, 3) + projection_calls = [] + + def project(hidden: torch.Tensor, positions: torch.Tensor): + projection_calls.append((hidden, positions)) + return q_index, token_k, position_coordinates + + attention = SimpleNamespace( + head_dim=8, + indexer=SimpleNamespace(project=project), + layer_idx_str="3", + num_heads=4, + ) + bridge_calls = [] + + def bridge(*args: object) -> None: + bridge_calls.append(args) + output = args[-1] + assert isinstance(output, torch.Tensor) + output.fill_(7) + + monkeypatch.setattr(module, "is_in_breakable_cuda_graph", lambda: in_graph) + monkeypatch.setattr(custom_ops, "maybe_bcg_qsa_attn_inplace", bridge) + + output = QSASparseHooks().forward( + attention=attention, + q=torch.empty((4, 64)), + k=None, + v=None, + attn_metadata=metadata, + attention_mask=PredefinedAttentionMask.CAUSAL, + attention_window_size=None, + attention_mask_data=None, + mrope_config=None, + attention_sinks=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + has_lora=False, + output_gate=None, + qsa_index_hidden_states=hidden_states, + qsa_position_ids=position_ids, + ) + + assert projection_calls == [(hidden_states, position_ids)] + assert bridge_calls[0][3] is q_index + assert bridge_calls[0][4] is token_k + assert bridge_calls[0][5] is position_coordinates + assert output.shape == (4, 32) + torch.testing.assert_close(output, torch.full_like(output, 7)) + + +def test_qsa_breakable_cuda_graph_bridge_uses_live_dense_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.attention.backends.sparse.qsa import custom_ops + + metadata = object() + q_index = torch.randn(4, 4, 8) + token_k = torch.randn(4, 1, 8) + position_coordinates = torch.arange(12, dtype=torch.int32).reshape(4, 3) + calls = {} + + class _Hooks: + def forward(self, *args: object, **kwargs: object): + calls["metadata"] = args[4] + calls["sparse_gate"] = args[13] + calls["projection"] = kwargs["qsa_index_projection"] + return None + + def dense_attention(*args: object, **kwargs: object): + calls["dense_metadata"] = args[3] + assert kwargs["output"].shape == (4, 32) + return torch.full((2, 32), 5.0), None + + def apply_output_gate(output: torch.Tensor, gate: torch.Tensor): + calls["dense_gate"] = gate + return output + gate + + attention = SimpleNamespace( + sparse_attn_hooks=_Hooks(), + _attn_impl=dense_attention, + apply_output_gate=apply_output_gate, + ) + monkeypatch.setattr( + custom_ops, + "_extract_qsa_extra_attrs", + lambda layer_idx: (metadata, attention), + ) + output = torch.full((4, 32), float("nan")) + output_gate = torch.full((4, 32), 2.0) + + custom_ops.qsa_attn_inplace( + torch.empty((4, 64)), + None, + None, + q_index, + token_k, + position_coordinates, + None, + None, + output_gate, + "3", + output, + ) + + assert calls["metadata"] is metadata + assert calls["dense_metadata"] is metadata + assert calls["sparse_gate"] is output_gate + torch.testing.assert_close(calls["dense_gate"], output_gate[:2]) + assert calls["projection"] == (q_index, token_k, position_coordinates) + torch.testing.assert_close(output[:2], torch.full((2, 32), 7.0)) + torch.testing.assert_close(output[2:], torch.zeros((2, 32))) + + +def test_qsa_breakable_cuda_graph_bridge_copies_live_sparse_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.attention.backends.sparse.qsa import custom_ops + + metadata = object() + q_index = torch.randn(4, 4, 8) + token_k = torch.randn(4, 1, 8) + position_coordinates = torch.arange(12, dtype=torch.int32).reshape(4, 3) + output_gate = torch.full((4, 32), 2.0) + calls = {} + + class _Hooks: + def forward(self, *args: object, **kwargs: object): + calls["metadata"] = args[4] + calls["gate"] = args[13] + calls["projection"] = kwargs["qsa_index_projection"] + return torch.full((2, 32), 6.0) + + def unexpected_dense_attention(*args: object, **kwargs: object): + raise AssertionError("dense attention must not run after sparse dispatch") + + attention = SimpleNamespace( + sparse_attn_hooks=_Hooks(), + _attn_impl=unexpected_dense_attention, + ) + monkeypatch.setattr( + custom_ops, + "_extract_qsa_extra_attrs", + lambda layer_idx: (metadata, attention), + ) + output = torch.full((4, 32), float("nan")) + + custom_ops.qsa_attn_inplace( + torch.empty((4, 64)), + None, + None, + q_index, + token_k, + position_coordinates, + None, + None, + output_gate, + "3", + output, + ) + + assert calls["metadata"] is metadata + assert calls["gate"] is output_gate + assert calls["projection"] == (q_index, token_k, position_coordinates) + torch.testing.assert_close(output[:2], torch.full((2, 32), 6.0)) + torch.testing.assert_close(output[2:], torch.zeros((2, 32))) + + def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers(monkeypatch) -> None: monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) manager = object.__new__(QSAMambaHybridCacheManagerV2) diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 5fa2f5323f8f..c16155a20db7 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -301,6 +301,104 @@ def capture_ple_ids( assert output is inputs_embeds +def test_breakable_cuda_graph_defers_ple_state_to_layer_bridge(monkeypatch) -> None: + from tensorrt_llm._torch.models import modeling_qwen4_exp + from tensorrt_llm._torch.models.modeling_qwen4_exp import Qwen4ExpModel + + captured = {} + + class _PLELayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.ple = nn.Identity() + + def forward(self, **kwargs: object): + captured.update(kwargs) + return kwargs["hidden_states"], None + + model = object.__new__(Qwen4ExpModel) + nn.Module.__init__(model) + model.model_config = SimpleNamespace(mapping=SimpleNamespace(has_pp=lambda: False)) + model.layers = nn.ModuleList((_PLELayer(),)) + model.num_hidden_layers = 1 + model.ple_layer_mask = [True] + model.ple_layer_index = 0 + model.has_ple = True + model.defer_combine_mask = [False] + model._bcg_model_key = "target" + + monkeypatch.setattr(modeling_qwen4_exp, "is_in_breakable_cuda_graph", lambda: True) + monkeypatch.setattr( + Qwen4ExpModel, + "_prepare_ple_state", + lambda *args, **kwargs: pytest.fail("PLE state must be prepared by the replay bridge"), + ) + input_ids = torch.tensor([10, 11, 12]) + inputs_embeds = torch.randn(3, 8) + + output = model( + attn_metadata=SimpleNamespace(mamba_metadata=None, max_num_requests=1), + inputs_embeds=inputs_embeds, + orig_input_ids=input_ids, + ) + + assert output is inputs_embeds + assert captured["ple_state"] is None + assert captured["ple_input_ids"] is input_ids + assert captured["ple_model_key"] == "target" + + +def test_breakable_cuda_graph_ple_bridge_uses_live_state_and_fixed_output( + monkeypatch, +) -> None: + from tensorrt_llm._torch.models import modeling_qwen4_exp + + metadata = SimpleNamespace(mamba_metadata=object()) + spec_metadata = object() + ple_state = (object(), object(), object()) + calls = {} + + class _PLE: + def __call__(self, hidden_states, *state): + calls["ple_state"] = state + return hidden_states[:2].mul(3) + + def prepare_state(attn_metadata, input_ids, mamba_metadata, spec): + calls["prepare"] = (attn_metadata, input_ids, mamba_metadata, spec) + return ple_state + + model = SimpleNamespace( + ple_layer_index=0, + layers=[SimpleNamespace(ple=_PLE())], + _prepare_ple_state=prepare_state, + ) + monkeypatch.setattr( + modeling_qwen4_exp, + "_extract_qwen4_exp_model", + lambda model_key: (model, metadata, spec_metadata), + ) + hidden_states = torch.randn(4, 8) + input_ids = torch.tensor([10, 11, 12, 13]) + output = torch.full_like(hidden_states, float("nan")) + + modeling_qwen4_exp.qwen4_exp_ple_inplace( + hidden_states, + input_ids, + "target", + output, + ) + + assert calls["prepare"] == ( + metadata, + input_ids, + metadata.mamba_metadata, + spec_metadata, + ) + assert calls["ple_state"] == ple_state + torch.testing.assert_close(output[:2], hidden_states[:2] * 3) + torch.testing.assert_close(output[2:], torch.zeros_like(output[2:])) + + def test_multimodal_forward_requires_original_token_ids_on_a_ple_rank() -> None: from tensorrt_llm._torch.models.modeling_qwen4_exp import Qwen4ExpModel @@ -1420,6 +1518,7 @@ def fake_build( num_contexts=2, num_seqs=2, num_tokens=5, + padded_num_tokens=8, seq_lens=seq_lens, seq_lens_cuda=seq_lens_cuda, all_rank_num_tokens=None, @@ -1445,6 +1544,7 @@ def fake_build( assert captured["state_indices"].data_ptr() == state_indices_long.data_ptr() assert captured["state_indices"].dtype == torch.long assert captured["kwargs"]["host_seq_lens"] == [2, 3] + assert captured["kwargs"]["physical_tokens"] == 8 assert captured["input_ids"].shape == (5,) torch.testing.assert_close(conv_state.flatten(), torch.tensor([1.0, 0.0], device=device)) torch.testing.assert_close( diff --git a/tests/unittest/_torch/modules/test_qwen4_exp_ple.py b/tests/unittest/_torch/modules/test_qwen4_exp_ple.py index 8949a28e1d6b..b685698f1ec5 100644 --- a/tests/unittest/_torch/modules/test_qwen4_exp_ple.py +++ b/tests/unittest/_torch/modules/test_qwen4_exp_ple.py @@ -765,3 +765,113 @@ def fake_reducescatter(partial, actual_mapping, dim, sizes): ) torch.testing.assert_close(output, full_weight[local_ids]) + + +def test_empty_attention_dp_rank_participates_in_embedding_collectives(monkeypatch) -> None: + from tensorrt_llm._torch.modules.qwen4_exp import ple as qwen4_exp_ple + + config = SimpleNamespace( + hidden_size=4, + hc_count=2, + ngram_size=3, + heads_per_ngram=1, + ple_embed_dim=4, + ple_conv_kernel_size=2, + vocab_size=32, + eos_token_id=2, + seed=1234, + ngram_vocab_size_base=11, + make_ngram_vocab_size_divisible_by=8, + rms_norm_eps=1e-6, + ) + mapping = SimpleNamespace(tp_size=2, tp_rank=0, cp_size=1, enable_attention_dp=True) + module = Qwen4ExpPLE(config, dtype=torch.float32, mapping=mapping) + physical_tokens = 8 + metadata = PLEMetadata.build( + torch.empty(0, dtype=torch.long), + torch.empty(0, dtype=torch.long), + torch.empty(0, dtype=torch.long), + is_decode=False, + eos_token_id=config.eos_token_id, + physical_tokens=physical_tokens, + all_rank_num_tokens=[physical_tokens, physical_tokens], + is_cuda_graph=True, + ) + embedding = module.ple_embedding + with torch.no_grad(): + embedding.ngram_embedding.weight.copy_( + torch.arange( + embedding.ngram_embedding.weight.numel(), + dtype=embedding.ngram_embedding.weight.dtype, + ).reshape_as(embedding.ngram_embedding.weight) + + 1 + ) + remote_ids = torch.full( + (physical_tokens, embedding.ngram_heads), + embedding.vocab_end_index, + dtype=torch.long, + ) + collective_calls = [] + + def fake_allgather(input_ids, actual_mapping, dim, sizes): + assert actual_mapping is mapping + assert dim == 0 + assert sizes is None + assert input_ids.shape == (physical_tokens, embedding.ngram_heads) + assert torch.count_nonzero(input_ids) == 0 + collective_calls.append(("allgather", input_ids.shape)) + return torch.cat((input_ids, remote_ids)) + + def fake_reducescatter(partial, actual_mapping, dim, sizes): + assert actual_mapping is mapping + assert dim == 0 + assert sizes is None + assert partial.shape == ( + 2 * physical_tokens, + embedding.ngram_heads, + embedding.head_dim_per_ngram, + ) + expected_local = embedding.ngram_embedding.weight[0].expand( + physical_tokens, + embedding.ngram_heads, + embedding.head_dim_per_ngram, + ) + torch.testing.assert_close(partial[:physical_tokens], expected_local) + assert torch.count_nonzero(partial[physical_tokens:]) == 0 + collective_calls.append(("reducescatter", partial.shape)) + return partial[:physical_tokens] + + monkeypatch.setattr(qwen4_exp_ple, "allgather", fake_allgather) + monkeypatch.setattr(qwen4_exp_ple, "reducescatter", fake_reducescatter) + hidden_states = torch.empty((physical_tokens, config.hidden_size * config.hc_count)) + conv_state = torch.zeros((1, *module.conv_state_shape)) + ngram_context = torch.full( + (1, module.ngram_context_len), + config.eos_token_id, + dtype=torch.long, + ) + conv_state_before = conv_state.clone() + ngram_context_before = ngram_context.clone() + + output = module(hidden_states, metadata, conv_state, ngram_context) + + assert collective_calls == [ + ( + "allgather", + torch.Size([physical_tokens, embedding.ngram_heads]), + ), + ( + "reducescatter", + torch.Size( + [ + 2 * physical_tokens, + embedding.ngram_heads, + embedding.head_dim_per_ngram, + ] + ), + ), + ] + torch.testing.assert_close(conv_state, conv_state_before) + torch.testing.assert_close(ngram_context, ngram_context_before) + assert output.shape == hidden_states.shape + assert torch.count_nonzero(output) == 0 diff --git a/tests/unittest/executor/test_sleep_collective_rpc_guards.py b/tests/unittest/executor/test_sleep_collective_rpc_guards.py index 360f0e2856aa..e0fe2fe884c7 100644 --- a/tests/unittest/executor/test_sleep_collective_rpc_guards.py +++ b/tests/unittest/executor/test_sleep_collective_rpc_guards.py @@ -12,7 +12,7 @@ # 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. -"""Pure-Python guard tests for MPI sleep/wakeup and collective_rpc. +"""Pure-Python guard tests for worker controls and collective_rpc. No GPU or model weights required; all CUDA/MPI/ZMQ machinery is bypassed via object.__new__ + manual attribute injection. @@ -1602,10 +1602,43 @@ def _noop_control_action(): # --------------------------------------------------------------------------- -# GenerationExecutorProxy / GenerationExecutorRpcProxy collective_rpc() +# BaseWorker diagnostics and proxy collective_rpc() # --------------------------------------------------------------------------- +def test_worker_reports_stable_empty_prefill_graph_stats(): + """Disabled and non-PyTorch workers preserve the diagnostics schema.""" + expected = { + "captured_token_buckets": [], + "replay_counts": {}, + "qsa_sparse_prefill_dispatches": 0, + } + worker = _make_worker() + assert worker.get_prefill_cuda_graph_stats() == expected + + worker._is_pytorch_backend = False + assert worker.get_prefill_cuda_graph_stats() == expected + + +def test_worker_reports_rank_local_prefill_graph_stats(): + """The worker aggregates BCG and sparse-QSA counters without GPU work.""" + worker = _make_worker() + runner = SimpleNamespace(captured_token_buckets=[2048, 8192], replay_counts={2048: 2, 8192: 1}) + model = MagicMock() + model.modules.return_value = [ + SimpleNamespace(sparse_attn_hooks=SimpleNamespace(num_sparse_prefill_dispatches=3)), + SimpleNamespace(sparse_attn_hooks=SimpleNamespace(num_sparse_prefill_dispatches=5)), + SimpleNamespace(), + ] + worker.engine.model_engine = SimpleNamespace(breakable_cuda_graph_runner=runner, model=model) + + assert worker.get_prefill_cuda_graph_stats() == { + "captured_token_buckets": [2048, 8192], + "replay_counts": {2048: 2, 8192: 1}, + "qsa_sparse_prefill_dispatches": 8, + } + + @pytest.mark.parametrize("cls", ["ipc", "rpc"]) class TestProxyCollectiveRpcGuards: """Guard-path tests for both IPC and RPC proxy collective_rpc() shims.""" @@ -1627,6 +1660,17 @@ def test_multirank_allowed_for_sleep_wakeup(self, cls): result = p.collective_rpc(method_name, args=(["kv_cache"],)) assert result == ["ok"] + def test_multirank_allows_rank_local_prefill_graph_stats(self, cls): + """Rank-local prefill graph diagnostics may be queried through rank 0.""" + expected = {"captured_token_buckets": [2048, 8192]} + mock_call = MagicMock() + mock_call.remote.return_value = expected + mock_client = MagicMock() + mock_client.get_prefill_cuda_graph_stats.return_value = mock_call + proxy = _make_proxy(cls, model_world_size=4, rpc_client=mock_client) + + assert proxy.collective_rpc("get_prefill_cuda_graph_stats") == [expected] + def test_multirank_raises_for_non_allowlisted_method(self, cls): """Non-allowlisted methods still raise NotImplementedError for world_size > 1.""" p = _make_proxy(cls, model_world_size=2, rpc_client=MagicMock())