diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 9511be887b1d..c7e0167d39f4 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -1,4 +1,4 @@ -# Torch Compile & Piecewise CUDA Graph +# Torch Compile & Prefill CUDA Graph In this guide, we show how to enable torch.compile and Piecewise CUDA Graph in TensorRT LLM. TensorRT LLM uses torch.compile for lightweight vertical fusion and Piecewise CUDA Graph. @@ -41,12 +41,29 @@ To enable torch.compile and Piecewise CUDA Graph, add the following configuratio ```yaml ... # Other extra config +prefill_cuda_graph_backend: piecewise +prefill_capture_num_tokens: '${capture_num_tokens}' # e.g. [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072] torch_compile_config: - capture_num_tokens: '${capture_num_tokens}' # List of num tokens to capture. e.g., [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072] enable_userbuffers: false - enable_piecewise_cuda_graph: true ``` +`TorchCompileConfig.enable_piecewise_cuda_graph` and +`TorchCompileConfig.capture_num_tokens` are deprecated aliases for these +prefill-specific options. + +The experimental breakable implementation can capture the model body without +torch.compile: + +```yaml +prefill_cuda_graph_backend: breakable +prefill_capture_num_tokens: [128, 256, 512] +``` + +The first version of the breakable backend supports BF16 Qwen3.5 on one GPU for +context-only, tensor/pipeline parallelism and mixed context/decode batches with KV cache. Speculative +decoding, LoRA, multimodal inputs, and context +logits fall back to eager execution or are rejected during initialization. + ## Tips for Piecewise CUDA Graph ### Piecewise CUDA Graph & Generation Only CUDA Graph @@ -59,9 +76,10 @@ cuda_graph_config: max_batch_size: 1024 # Specify max capture batch size for generation only cuda graph. By default, TensorRT LLM will generate a capture list based on it. torch_compile_config: - capture_num_tokens: '${capture_num_tokens}' # Specify capture_num_tokens for piecewise cuda graph enable_userbuffers: false - enable_piecewise_cuda_graph: true + +prefill_cuda_graph_backend: piecewise +prefill_capture_num_tokens: '${capture_num_tokens}' ``` ### Piecewise CUDA Graph Padding diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py index 5472c674c448..18ae3d6ae87e 100644 --- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py +++ b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py @@ -12,7 +12,7 @@ from tensorrt_llm.llmapi.utils import enable_llm_debug from ..utils import (get_model_extra_attrs, - get_per_request_piecewise_cuda_graph_flag, + get_per_request_prefill_cuda_graph_flag, get_piecewise_cuda_graph_flag, make_weak_ref, set_piecewise_running) from .multi_stream.auto_multi_stream import multi_stream_schedule @@ -202,7 +202,7 @@ def __call__(self, *args): if (runtime_num_of_token is None or runtime_num_of_token not in self.entries or not get_piecewise_cuda_graph_flag() - or not get_per_request_piecewise_cuda_graph_flag()): + or not get_per_request_prefill_cuda_graph_flag()): return self.default_callable(*args) if self.is_first_runner or self.is_last_runner: diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 8d3e0ddcac1e..1cb82f736272 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -62,6 +62,7 @@ ) from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..utils import ( ActivationType, AuxStreamType, @@ -660,6 +661,9 @@ def minimax_m3_attn_custom_op_inplace( ) +maybe_bcg_minimax_m3_attn_custom_op_inplace = eager_on_graph(minimax_m3_attn_custom_op_inplace) + + class MiniMaxM3Attention(Attention): """M3 attention: dense (layers 0-2) or sparse (layers 3-59). @@ -1238,8 +1242,8 @@ def _forward_attention_core( attn_metadata: AttentionMetadata, ) -> torch.Tensor: output = q.new_empty((q.shape[0], self.num_heads * self.head_dim)) - if self.register_to_config and is_torch_compiling(): - minimax_m3_attn_custom_op_inplace( + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): + maybe_bcg_minimax_m3_attn_custom_op_inplace( q, k, v, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 53ab4b86bd44..ac957e60ebb4 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -22,6 +22,8 @@ cp_allgather, reducescatter) from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType +from ..pyexecutor.breakable_cuda_graph import (eager_on_graph, + is_in_breakable_cuda_graph) from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, is_nvfp4_marlin_enabled, is_torch_compiling) from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig @@ -115,6 +117,9 @@ def attn_custom_op_inplace( ) +maybe_bcg_attn_custom_op_inplace = eager_on_graph(attn_custom_op_inplace) + + def _helix_zero_kv_mask( attn_metadata: AttentionMetadata, num_tokens: int, @@ -932,20 +937,19 @@ def forward_impl( if "mrope_position_deltas" in mrope_config: mrope_position_deltas = mrope_config["mrope_position_deltas"] - # Currently only TRTLLM and FLASHINFER are torch compile compatible backends. - # Only enable custom inplace op when torch compiling. - use_custom_inplace_op = (self.register_to_config - and (self.attn_backend == "TRTLLM" - or self.attn_backend == "FLASHINFER") - and is_torch_compiling() - and not self.is_marlin_enabled) + # Currently only TRTLLM and FLASHINFER support the custom inplace op. + use_custom_inplace_op = ( + self.register_to_config and + (self.attn_backend == "TRTLLM" or self.attn_backend == "FLASHINFER") + and (is_torch_compiling() or is_in_breakable_cuda_graph()) + and not self.is_marlin_enabled) if use_custom_inplace_op: outputs = create_attn_outputs(q, attention_mask, self.layer_idx_str) assert len(outputs) == 1 or len(outputs) == 2 output = outputs[0] output_sf = outputs[1] if len(outputs) == 2 else None - attn_custom_op_inplace( + maybe_bcg_attn_custom_op_inplace( q, k, v, diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index f786417a4e46..efcc5ff6c8d5 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -31,6 +31,7 @@ from ...attention_backend import AttentionMetadata from ...distributed import AllReduceParams from ...model_config import ModelConfig +from ...pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ...speculative import SpecMetadata from ...utils import EventType, get_model_extra_attrs, is_gdn_replay_enabled, is_torch_compiling from ..linear import FP8QDQLinearMethod, Linear, TensorParallelMode @@ -174,6 +175,9 @@ def gdn_custom_op_inplace( ) +maybe_bcg_gdn_custom_op_inplace = eager_on_graph(gdn_custom_op_inplace) + + def ensure_divisibility(numerator, denominator): """Ensure that numerator is divisible by the denominator.""" assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator) @@ -1053,11 +1057,12 @@ def forward( ): mixed_qkv, z, a, b = self._compute_tokenwise_inputs(hidden_states) - if self.register_to_config and is_torch_compiling(): + use_breakable_cuda_graph = not is_torch_compiling() and is_in_breakable_cuda_graph() + if self.register_to_config and (is_torch_compiling() or use_breakable_cuda_graph): attn_out = mixed_qkv.new_empty( (1, mixed_qkv.shape[0], self.num_v_heads_per_tp, self.head_v_dim) ) - gdn_custom_op_inplace(mixed_qkv, a, b, self.layer_idx_str, attn_out) + maybe_bcg_gdn_custom_op_inplace(mixed_qkv, a, b, self.layer_idx_str, attn_out) else: attn_out = self.forward_core( mixed_qkv, diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index d454f833726b..9cf400020dee 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -48,6 +48,7 @@ from ..attention_backend.utils import create_attention from ..distributed import AllReduceParams from ..model_config import ModelConfig +from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..utils import ( AuxStreamType, Fp4QuantizedTensor, @@ -134,32 +135,31 @@ def _extract_mla_extra_attrs(layer_idx: str): return metadata, mla_layer -def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> List[torch.Tensor]: +def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) enable_dsv4_epilogue_fusion = mla_layer._should_use_dsv4_epilogue_fusion( metadata.num_contexts, metadata.num_generations ) - output_input = hidden_states[:0] if enable_dsv4_epilogue_fusion else hidden_states - attn_output = mla_layer.create_output(output_input, metadata.num_contexts) - outputs = [attn_output] - if enable_dsv4_epilogue_fusion: - outputs.extend(mla_layer._create_dsv4_epilogue_buffers(hidden_states, metadata.num_tokens)) - return outputs + return mla_layer.create_output( + hidden_states, + metadata.num_contexts, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, + ) @torch.library.custom_op("trtllm::create_mla_outputs", mutates_args=()) -def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> List[torch.Tensor]: +def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @create_mla_outputs.register_fake -def _create_mla_outputs_fake(hidden_states, layer_idx): +def _create_mla_outputs_fake(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @torch.library.custom_op( "trtllm::mla_custom_op_inplace", - mutates_args=("output", "dsv4_output", "dsv4_output_sf"), + mutates_args=("output",), ) def mla_custom_op_inplace( hidden_states: torch.Tensor, @@ -167,9 +167,6 @@ def mla_custom_op_inplace( layer_idx: str, output: torch.Tensor, latent_cache_gen: Optional[torch.Tensor], - dsv4_output: Optional[torch.Tensor], - dsv4_output_sf: Optional[torch.Tensor], - enable_dsv4_epilogue_fusion: bool, hidden_states_fp4: Optional[torch.Tensor] = None, hidden_states_sf: Optional[torch.Tensor] = None, ) -> None: @@ -192,33 +189,25 @@ def mla_custom_op_inplace( # DeepSeek-V4 uses MQA mode and has no residual-less RMSNorm+quant # fusion entry point, so it cannot be reached with a pre-quantized # Fp4QuantizedTensor input; the call site passes plain hidden_states. - if enable_dsv4_epilogue_fusion: - if dsv4_output is None or dsv4_output_sf is None: - raise RuntimeError( - "DSv4 fused epilogue requires caller-provided output and output_sf buffers." - ) - dsv4_epilogue_output = (dsv4_output, dsv4_output_sf) - else: - if dsv4_output is not None or dsv4_output_sf is not None: - raise RuntimeError( - "DSv4 fused epilogue buffers require epilogue fusion to be enabled." - ) - dsv4_epilogue_output = None + enable_dsv4_epilogue_fusion = mla_layer._should_use_dsv4_epilogue_fusion( + metadata.num_contexts, metadata.num_generations + ) mla_layer.forward_impl_with_deepseek_v4( position_ids, hidden_states, metadata, output=output, - dsv4_epilogue_output=dsv4_epilogue_output, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) else: - if enable_dsv4_epilogue_fusion: - raise RuntimeError("DSv4 fused epilogue cannot be enabled for non-DeepSeek-V4 MLA.") mla_layer.forward_impl( position_ids, hidden_states, metadata, output=output, latent_cache_gen=latent_cache_gen ) +maybe_bcg_mla_custom_op_inplace = eager_on_graph(mla_custom_op_inplace) + + @torch.library.custom_op("trtllm::mla_dsa_proj", mutates_args=()) def mla_dsa_proj( hidden_states: torch.Tensor, @@ -337,6 +326,9 @@ def mla_dsa_attn_inplace( ) +maybe_bcg_mla_dsa_attn_inplace = eager_on_graph(mla_dsa_attn_inplace) + + def fp8_block_scaling_bmm_out( mat1: torch.Tensor, mat2_fp8: torch.Tensor, @@ -1135,7 +1127,12 @@ def _attn_forward_gen( ) return attn_output - def create_output(self, hidden_states: torch.Tensor, num_contexts: int): + def create_output( + self, + hidden_states: torch.Tensor, + num_contexts: int, + enable_dsv4_epilogue_fusion: bool = False, + ) -> torch.Tensor: # Upstream POST_MoE/MLP fusion (or attention-DP no-fusion fold) may pass # an Fp4QuantizedTensor here; unpack to the BF16 view for sizing. The # producing fold must have requested return_norm_out so the BF16 view @@ -1148,6 +1145,17 @@ def create_output(self, hidden_states: torch.Tensor, num_contexts: int): ) hidden_states = hidden_states.unquantized_hidden_states num_tokens = hidden_states.shape[0] + if enable_dsv4_epilogue_fusion: + # BCG replays smaller token counts by slicing captured outputs along + # dim 0. The original DSv4 epilogue outputs were not sliceable this + # way, so expose the O-LoRA activation as the attention output. Its + # token-first layout gives all DSv4 epilogue-fusion batch types the + # same single-Tensor, dim-0-sliceable output contract. + return torch.empty( + [num_tokens, self.n_local_groups, self.o_lora_rank], + device=hidden_states.device, + dtype=self.dtype, + ) if self.is_deepseek_v4: hidden_size = self.num_heads_tp_cp * self.v_head_dim else: @@ -1172,10 +1180,6 @@ def _should_use_dsv4_epilogue_fusion(self, num_contexts: int, num_generations: i return False if num_contexts == 0 and num_generations == 0: return False - if num_contexts > 0 and num_generations > 0: - # Context and generation use separate FMHA calls, but the fused - # buffers do not carry token offsets for a mixed batch. - return False if self.mapping.has_cp_helix(): return False if not is_sm_100f(): @@ -1221,38 +1225,15 @@ def _create_dsv4_epilogue_buffers( ) return fp8_o, output_sf - def _validate_dsv4_epilogue_buffers( - self, - num_tokens: int, - dsv4_epilogue_output: tuple[torch.Tensor, torch.Tensor], - ) -> tuple[torch.Tensor, torch.Tensor]: - fp8_o, output_sf = dsv4_epilogue_output - scale_buf_m = (num_tokens + 3) // 4 * 4 - if fp8_o.shape[1] != num_tokens or output_sf.shape[2] != scale_buf_m: - raise RuntimeError("Invalid DSv4 fused epilogue buffers for current token count.") - return fp8_o, output_sf - def _deepseek_v4_o_proj( self, - attn_out_latent: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + attn_out_latent: torch.Tensor, position_ids: Optional[torch.Tensor] = None, + *, + enable_dsv4_epilogue_fusion: bool, ) -> torch.Tensor: - if isinstance(attn_out_latent, tuple): - attn_fp8, attn_scale = attn_out_latent - num_tokens = attn_fp8.shape[1] - o_lora = torch.empty( - [num_tokens, self.n_local_groups, self.o_lora_rank], - device=attn_fp8.device, - dtype=self.dtype, - ) - torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( - attn_fp8, - self.o_a_proj, - attn_scale, - self.o_a_proj_scale, - o_lora.transpose(0, 1), - ) - return self.o_b_proj(o_lora.flatten(1)) + if enable_dsv4_epilogue_fusion: + return self.o_b_proj(attn_out_latent.flatten(1)) assert position_ids is not None num_tokens = attn_out_latent.shape[0] @@ -1333,6 +1314,38 @@ def _deepseek_v4_o_proj( output = self.o_b_proj(o_lora) return output + def _run_dsv4_o_lora_bmms( + self, + o_lora_output: torch.Tensor, + num_context_tokens: int, + num_tokens: int, + context_o_lora_bmm_input: Optional[tuple[torch.Tensor, torch.Tensor]], + generation_o_lora_bmm_input: Optional[tuple[torch.Tensor, torch.Tensor]], + ) -> None: + def run_o_lora_bmm( + o_lora_bmm_input: tuple[torch.Tensor, torch.Tensor], + phase_o_lora_output: torch.Tensor, + ) -> None: + attn_fp8, attn_scale = o_lora_bmm_input + torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( + attn_fp8, + self.o_a_proj, + attn_scale, + self.o_a_proj_scale, + phase_o_lora_output.transpose(0, 1), + ) + + if context_o_lora_bmm_input is not None: + run_o_lora_bmm( + context_o_lora_bmm_input, + o_lora_output[:num_context_tokens], + ) + if generation_o_lora_bmm_input is not None: + run_o_lora_bmm( + generation_o_lora_bmm_input, + o_lora_output[num_context_tokens:num_tokens], + ) + def _resolve_qa_fused_scale(self): """Lazily decide whether the residual-less q_a_layernorm -> q_b_proj NVFP4 fusion can run, caching q_b_proj's static input_scale. @@ -1733,7 +1746,7 @@ def forward_impl_with_deepseek_v4( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + enable_dsv4_epilogue_fusion: bool, ) -> None: """ Forward pass for the MLA module with DeepSeek-V4 (always in MQA mode). @@ -1742,10 +1755,10 @@ def forward_impl_with_deepseek_v4( position_ids (Optional[torch.IntTensor]): The position IDs. hidden_states (torch.Tensor): The hidden states. attn_metadata (AttentionMetadata): The attention metadata. - output (torch.Tensor): Pre-allocated output tensor, written in-place - when epilogue fusion is disabled. - dsv4_epilogue_output: Caller-provided ``(fp8_o, output_sf)`` - buffers, written in-place when epilogue fusion is enabled. + output (torch.Tensor): Pre-allocated attention output, or the final + three-dimensional LoRA output when epilogue fusion is enabled. + enable_dsv4_epilogue_fusion (bool): Whether to use the fused + DeepSeek-V4 epilogue. """ assert self.mha is None and self.mqa is not None, ( "DeepSeek-V4 is only supported in MQA mode" @@ -1755,12 +1768,6 @@ def forward_impl_with_deepseek_v4( num_generations = attn_metadata.num_generations num_ctx_tokens = attn_metadata.num_ctx_tokens num_tokens = attn_metadata.num_tokens - enable_dsv4_epilogue_fusion = dsv4_epilogue_output is not None - if enable_dsv4_epilogue_fusion and ((num_contexts > 0) == (num_generations > 0)): - raise RuntimeError( - "DSv4 epilogue fusion requires a context-only or generation-only batch." - ) - hidden_states = hidden_states[:num_tokens, ...] if position_ids is not None: position_ids = position_ids[..., :num_tokens] @@ -1938,6 +1945,8 @@ def _indexer_branch(): assert output is not None, "output must be provided" + context_o_lora_bmm_input = None + generation_o_lora_bmm_input = None if num_contexts > 0: q_ctx = q[:num_ctx_tokens, ...] topk_indices_ctx = ( @@ -1953,17 +1962,16 @@ def _indexer_branch(): assert ctx_position_ids is not None k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) - self.forward_context_sparse_mla( + context_o_lora_bmm_input = self.forward_context_sparse_mla( q_ctx, compressed_kv_ctx, k_pe_ctx, attn_metadata, - output[:num_ctx_tokens, :], + None if enable_dsv4_epilogue_fusion else output[:num_ctx_tokens, :], position_ids=ctx_position_ids, latent_cache=latent_cache_ctx, topk_indices=topk_indices_ctx, enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, ) if num_generations > 0: @@ -1981,17 +1989,26 @@ def _indexer_branch(): assert gen_position_ids is not None k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) - self.forward_generation_sparse_mla( + generation_o_lora_bmm_input = self.forward_generation_sparse_mla( q_gen, compressed_kv_gen, k_pe_gen, attn_metadata, - output[num_ctx_tokens:num_tokens, :], + None if enable_dsv4_epilogue_fusion else output[num_ctx_tokens:num_tokens, :], position_ids=gen_position_ids, latent_cache=latent_cache_gen, topk_indices=topk_indices_gen, enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, + ) + if enable_dsv4_epilogue_fusion: + # The fused output is [groups, tokens, hidden], which BCG cannot slice on dim 0. + # Write O-LoRA as [tokens, groups, rank] to make it BCG-sliceable. + self._run_dsv4_o_lora_bmms( + output, + num_ctx_tokens, + num_tokens, + context_o_lora_bmm_input, + generation_o_lora_bmm_input, ) def forward_context_default( @@ -2074,12 +2091,11 @@ def forward_context_sparse_mla( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Run context-phase attention for DSA models. @@ -2110,6 +2126,7 @@ def forward_context_sparse_mla( if not enable_dsv4_epilogue_fusion and self._should_use_short_mha( attn_metadata, position_ids ): + assert output is not None return self.forward_context( q, compressed_kv, k_pe, position_ids, attn_metadata, output, latent_cache ) @@ -2125,10 +2142,10 @@ def forward_context_sparse_mla( latent_cache=latent_cache, topk_indices=topk_indices, enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, ) else: assert not self.is_deepseek_v4, "DeepSeek-V4 is not supported on pre-blackwell GPUs." + assert output is not None return self.forward_sparse_mla_kvcache_bf16( q, latent_cache, attn_metadata, output, topk_indices, is_generation=False ) @@ -2139,12 +2156,11 @@ def forward_generation_sparse_mla( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if get_sm_version() >= 100: return self.forward_absorption_generation( @@ -2157,10 +2173,10 @@ def forward_generation_sparse_mla( latent_cache=latent_cache, topk_indices=topk_indices, enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, ) else: assert not self.is_deepseek_v4, "DeepSeek-V4 is not supported on pre-blackwell GPUs." + assert output is not None return self.forward_sparse_mla_kvcache_bf16( q, latent_cache, attn_metadata, output, topk_indices, is_generation=True ) @@ -2510,12 +2526,11 @@ def forward_absorption_generation( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: num_tokens = q.shape[0] q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split( @@ -2649,14 +2664,11 @@ def _mla_gen_rope(): # Use generation_only for generation phase and context_only for context phase in DSA attention attention_input_type = AttentionInputType.generation_only dsv4_output = output if self.is_deepseek_v4 else None - dsv4_output_sf = None + o_lora_bmm_input_scale = None dsv4_cos_sin_cache = None if enable_dsv4_epilogue_fusion: assert self.is_deepseek_v4 - assert dsv4_epilogue_output is not None - dsv4_output, dsv4_output_sf = self._validate_dsv4_epilogue_buffers( - num_tokens, dsv4_epilogue_output - ) + dsv4_output, o_lora_bmm_input_scale = self._create_dsv4_epilogue_buffers(q, num_tokens) dsv4_cos_sin_cache = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -2669,7 +2681,7 @@ def _mla_gen_rope(): attention_input_type=attention_input_type, out_scale=self.out_scale, output=dsv4_output, - output_sf=dsv4_output_sf, + output_sf=o_lora_bmm_input_scale, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by `invokeMLARopeGeneration` topk_indices=topk_indices, # used by DSA attention @@ -2685,9 +2697,10 @@ def _mla_gen_rope(): fused_q = None if enable_dsv4_epilogue_fusion: - return attn_out_latent + return dsv4_output, o_lora_bmm_input_scale if self.is_deepseek_v4: + assert output is not None if self.mapping.has_cp_helix(): raise RuntimeError( "DeepSeek-V4 + CP Helix is not supported: " @@ -2699,6 +2712,7 @@ def _mla_gen_rope(): ) return output + assert output is not None # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert ( attn_out_latent.shape[0] == q.shape[0] @@ -2739,12 +2753,11 @@ def forward_absorption_context( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: num_tokens = q.shape[0] @@ -2826,14 +2839,11 @@ def forward_absorption_context( quant_scale_qkv = None dsv4_output = output if self.is_deepseek_v4 else None - dsv4_output_sf = None + o_lora_bmm_input_scale = None dsv4_cos_sin_cache = None if enable_dsv4_epilogue_fusion: assert self.is_deepseek_v4 - assert dsv4_epilogue_output is not None - dsv4_output, dsv4_output_sf = self._validate_dsv4_epilogue_buffers( - num_tokens, dsv4_epilogue_output - ) + dsv4_output, o_lora_bmm_input_scale = self._create_dsv4_epilogue_buffers(q, num_tokens) dsv4_cos_sin_cache = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -2846,7 +2856,7 @@ def forward_absorption_context( attention_input_type=attention_input_type, out_scale=self.out_scale, output=dsv4_output, - output_sf=dsv4_output_sf, + output_sf=o_lora_bmm_input_scale, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by applyMLARopeAndAssignQKVKernelOptContext quant_q_buffer=quant_q_buffer, # fused-FP8 path only @@ -2860,9 +2870,10 @@ def forward_absorption_context( self._fused_q_pe = None if enable_dsv4_epilogue_fusion: - return attn_out_latent + return dsv4_output, o_lora_bmm_input_scale if self.is_deepseek_v4: + assert output is not None if self.mapping.has_cp_helix(): raise RuntimeError( "DeepSeek-V4 + CP Helix is not supported: " @@ -2874,6 +2885,7 @@ def forward_absorption_context( ) return output + assert output is not None # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert ( attn_out_latent.shape[0] == q.shape[0] @@ -3068,31 +3080,22 @@ def forward( hidden_states, attn_metadata, self.mapping, self.layer_idx ) - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None - if self.register_to_config: + enable_dsv4_epilogue_fusion = self.is_deepseek_v4 and self._should_use_dsv4_epilogue_fusion( + attn_metadata.num_contexts, attn_metadata.num_generations + ) + use_custom_op = self.register_to_config and ( + is_torch_compiling() or is_in_breakable_cuda_graph() + ) + if use_custom_op: if self.is_deepseek_v4: - outputs = torch.ops.trtllm.create_mla_outputs(hidden_states, self.layer_idx_str) - attn_output = outputs[0] - dsv4_output = None - dsv4_output_sf = None - if len(outputs) == 3: - dsv4_output, dsv4_output_sf = outputs[1], outputs[2] - dsv4_epilogue_output = (dsv4_output, dsv4_output_sf) - elif len(outputs) != 1: - raise RuntimeError( - "create_mla_outputs must return either legacy output or " - "legacy output plus DSv4 fused epilogue buffers." - ) + attn_output = torch.ops.trtllm.create_mla_outputs(hidden_states, self.layer_idx_str) - torch.ops.trtllm.mla_custom_op_inplace( + maybe_bcg_mla_custom_op_inplace( hidden_states, position_ids, self.layer_idx_str, attn_output, latent_cache_gen, - dsv4_output, - dsv4_output_sf, - dsv4_epilogue_output is not None, ) else: attn_output = self.create_output(hidden_states, attn_metadata.num_contexts) @@ -3115,7 +3118,7 @@ def forward( ) q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] indexer_intermediates = proj_outputs[4:] - torch.ops.trtllm.mla_dsa_attn_inplace( + maybe_bcg_mla_dsa_attn_inplace( q, compressed_kv, k_pe, @@ -3132,42 +3135,29 @@ def forward( # take dataclasses, so pass the BF16 + FP4 + SF views as # explicit tensors. if isinstance(hidden_states, Fp4QuantizedTensor): - torch.ops.trtllm.mla_custom_op_inplace( + maybe_bcg_mla_custom_op_inplace( hidden_states.unquantized_hidden_states, position_ids, self.layer_idx_str, attn_output, latent_cache_gen, - None, - None, - False, hidden_states.fp4_tensor, hidden_states.scaling_factor, ) else: - torch.ops.trtllm.mla_custom_op_inplace( + maybe_bcg_mla_custom_op_inplace( hidden_states, position_ids, self.layer_idx_str, attn_output, latent_cache_gen, - None, - None, - False, ) else: - enable_dsv4_epilogue_fusion = ( - self.is_deepseek_v4 - and self._should_use_dsv4_epilogue_fusion( - attn_metadata.num_contexts, attn_metadata.num_generations - ) + attn_output = self.create_output( + hidden_states, + attn_metadata.num_contexts, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) - if enable_dsv4_epilogue_fusion: - dsv4_epilogue_output = self._create_dsv4_epilogue_buffers( - hidden_states, attn_metadata.num_tokens - ) - output_input = hidden_states[:0] if enable_dsv4_epilogue_fusion else hidden_states - attn_output = self.create_output(output_input, attn_metadata.num_contexts) if self.is_dsa: self.forward_impl_with_dsa( position_ids, hidden_states, attn_metadata, output=attn_output @@ -3178,7 +3168,7 @@ def forward( hidden_states, attn_metadata, output=attn_output, - dsv4_epilogue_output=dsv4_epilogue_output, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) else: self.forward_impl( @@ -3190,10 +3180,11 @@ def forward( ) if self.is_deepseek_v4: - if dsv4_epilogue_output is not None: - attn_output = self._deepseek_v4_o_proj(dsv4_epilogue_output) - else: - attn_output = self._deepseek_v4_o_proj(attn_output, position_ids) + attn_output = self._deepseek_v4_o_proj( + attn_output, + position_ids, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, + ) else: attn_output = _helix_cp_output_projection( self.o_proj, diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.py new file mode 100644 index 000000000000..f4ff7bbb24df --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.py @@ -0,0 +1,22 @@ +# Adapted from SGLang's breakable CUDA graph implementation. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, + break_graph, + eager_on_graph, + get_current_replay_token, +) +from .context import enable_breakable_cuda_graph, is_in_breakable_cuda_graph + +__all__ = [ + "BreakableCUDAGraph", + "BreakableCUDAGraphCapture", + "break_graph", + "eager_on_graph", + "enable_breakable_cuda_graph", + "get_current_replay_token", + "is_in_breakable_cuda_graph", +] diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py new file mode 100644 index 000000000000..3b429c3856f3 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -0,0 +1,284 @@ +# Adapted from SGLang's breakable CUDA graph implementation. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import functools +import itertools +import logging +import threading +from contextvars import ContextVar +from typing import Any, Callable, Optional + +import torch +from cuda.bindings import runtime as rt + +from tensorrt_llm._utils import CUASSERT + +from ...utils import make_weak_ref + +logger = logging.getLogger(__name__) + +__all__ = [ + "BreakableCUDAGraph", + "BreakableCUDAGraphCapture", + "break_graph", + "eager_on_graph", + "get_current_replay_token", +] + +_current_capture: ContextVar[Optional["BreakableCUDAGraphCapture"]] = ContextVar( + "breakable_cuda_graph_capture", default=None +) +_current_stream: ContextVar[Optional[torch.cuda.Stream]] = ContextVar( + "breakable_cuda_graph_stream", default=None +) +_current_replay_token: ContextVar[Optional[int]] = ContextVar( + "breakable_cuda_graph_replay_token", default=None +) +_forked_streams: ContextVar[Optional[set[torch.cuda.Stream]]] = ContextVar( + "breakable_cuda_graph_forked_streams", default=None +) +_replay_token_counter = itertools.count(1) + +_original_wait_stream: Optional[Callable] = None +_wait_stream_hook_lock = threading.Lock() +_wait_stream_hook_refcount = 0 + + +def get_current_stream(device: Optional[torch.device] = None) -> torch.cuda.Stream: + """Return the active BCG stream or PyTorch's current stream.""" + stream = _current_stream.get() + return torch.cuda.current_stream(device) if stream is None else stream + + +def get_current_replay_token() -> Optional[int]: + """Return a unique token for the active BCG replay.""" + return _current_replay_token.get() + + +def _capture_status(stream_ptr: int) -> rt.cudaStreamCaptureStatus: + status, *_ = CUASSERT(rt.cudaStreamGetCaptureInfo(stream_ptr)) + return status + + +def _is_stream_capturing(stream: torch.cuda.Stream) -> bool: + return ( + _capture_status(stream.cuda_stream) + == rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive + ) + + +def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream) -> None: + assert _original_wait_stream is not None + forked = _forked_streams.get() + capturing = _current_stream.get() + if forked is None or capturing is None: + _original_wait_stream(self, other) + return + + capture_ptr = capturing.cuda_stream + self_is_capture = self is capturing or self.cuda_stream == capture_ptr + other_is_capture = other is capturing or other.cuda_stream == capture_ptr + if self_is_capture and not other_is_capture: + if not _is_stream_capturing(other): + return + _original_wait_stream(self, other) + forked.discard(other) + elif other_is_capture and not self_is_capture: + _original_wait_stream(self, other) + forked.add(self) + else: + _original_wait_stream(self, other) + + +def _install_wait_stream_hook() -> None: + global _original_wait_stream, _wait_stream_hook_refcount + with _wait_stream_hook_lock: + if _wait_stream_hook_refcount == 0: + _original_wait_stream = torch.cuda.Stream.wait_stream + torch.cuda.Stream.wait_stream = _hooked_wait_stream + _wait_stream_hook_refcount += 1 + + +def _uninstall_wait_stream_hook() -> None: + global _original_wait_stream, _wait_stream_hook_refcount + with _wait_stream_hook_lock: + _wait_stream_hook_refcount -= 1 + if _wait_stream_hook_refcount == 0: + assert _original_wait_stream is not None + torch.cuda.Stream.wait_stream = _original_wait_stream + _original_wait_stream = None + + +def _copy_output(destination: Any, source: Any) -> Any: + if torch.is_tensor(destination) and torch.is_tensor(source): + destination.copy_(source) + return destination + + if ( + isinstance(destination, (tuple, list)) + and isinstance(source, (tuple, list)) + and len(destination) == len(source) + ): + copied = [_copy_output(dst, src) for dst, src in zip(destination, source)] + return tuple(copied) if isinstance(destination, tuple) else copied + + if hasattr(destination, "__dict__") and hasattr(source, "__dict__"): + for key, source_value in source.__dict__.items(): + destination_value = getattr(destination, key, None) + if torch.is_tensor(destination_value) and torch.is_tensor(source_value): + destination_value.copy_(source_value) + else: + setattr(destination, key, source_value) + return destination + + if isinstance(destination, dict) and isinstance(source, dict): + for key, source_value in source.items(): + destination_value = destination.get(key) + if torch.is_tensor(destination_value) and torch.is_tensor(source_value): + destination_value.copy_(source_value) + else: + destination[key] = source_value + return destination + + return source + + +def eager_on_graph(inner: Callable) -> Callable: + """Run a callable eagerly between captured CUDA graph segments.""" + + @functools.wraps(inner) + def wrapper(*args, **kwargs): + if torch.compiler.is_compiling(): + return inner(*args, **kwargs) + + capture = _current_capture.get() + if capture is None: + return inner(*args, **kwargs) + + logger.debug( + "Break CUDA graph for function %s", getattr(inner, "__name__", type(inner).__name__) + ) + capture._end_current_segment() + output = inner(*args, **kwargs) + + captured_args = tuple(make_weak_ref(arg) for arg in args) + captured_kwargs = {key: make_weak_ref(value) for key, value in kwargs.items()} + captured_output = make_weak_ref(output) + + def replay_fn() -> Any: + new_output = inner(*captured_args, **captured_kwargs) + return _copy_output(captured_output, new_output) + + capture.cuda_graph._break_functions.append(replay_fn) + capture._begin_new_segment() + return output + + return wrapper + + +class BreakableCUDAGraph: + """A sequence of CUDA graph segments separated by eager functions.""" + + def __init__(self) -> None: + self._segments: list[torch.cuda.CUDAGraph] = [] + self._break_functions: list[Callable[[], Any]] = [] + + @property + def num_segments(self) -> int: + return len(self._segments) + + @property + def num_breaks(self) -> int: + return len(self._break_functions) + + def pool(self): + if not self._segments: + raise RuntimeError("Cannot get the pool of an empty BCG") + return self._segments[0].pool() + + def replay(self) -> None: + stream_token = _current_stream.set(torch.cuda.current_stream()) + replay_token = _current_replay_token.set(next(_replay_token_counter)) + try: + for index, segment in enumerate(self._segments): + segment.replay() + if index < len(self._break_functions): + self._break_functions[index]() + finally: + _current_replay_token.reset(replay_token) + _current_stream.reset(stream_token) + + def reset(self) -> None: + for segment in self._segments: + segment.reset() + self._segments.clear() + self._break_functions.clear() + + +class BreakableCUDAGraphCapture: + """Capture a region as CUDA graph segments separated by eager work.""" + + def __init__( + self, + cuda_graph: BreakableCUDAGraph, + pool=None, + stream: Optional[torch.cuda.Stream] = None, + capture_error_mode: str = "global", + ) -> None: + if not isinstance(cuda_graph, BreakableCUDAGraph): + raise TypeError("cuda_graph must be a BreakableCUDAGraph") + self.cuda_graph = cuda_graph + self._pool = (0, 0) if pool is None else pool + self._stream = stream + self._capture_error_mode = capture_error_mode + self._stream_context = None + self._capture_token = None + self._stream_token = None + self._forked_token = None + + def __enter__(self) -> "BreakableCUDAGraphCapture": + _install_wait_stream_hook() + if self._stream is not None: + self._stream_context = torch.cuda.stream(self._stream) + self._stream_context.__enter__() + self._capture_token = _current_capture.set(self) + self._stream_token = _current_stream.set(self._stream or torch.cuda.current_stream()) + self._forked_token = _forked_streams.set(set()) + self._begin_new_segment() + return self + + def __exit__(self, *args: object) -> bool: + try: + self._end_current_segment() + finally: + _forked_streams.reset(self._forked_token) + _current_stream.reset(self._stream_token) + _current_capture.reset(self._capture_token) + if self._stream_context is not None: + self._stream_context.__exit__(*args) + self._stream_context = None + _uninstall_wait_stream_hook() + return False + + def _begin_new_segment(self) -> None: + segment = torch.cuda.CUDAGraph() + segment.capture_begin(pool=self._pool, capture_error_mode=self._capture_error_mode) + self.cuda_graph._segments.append(segment) + + def _end_current_segment(self) -> None: + main_stream = get_current_stream() + forked = _forked_streams.get() + if forked: + assert _original_wait_stream is not None + for side_stream in list(forked): + if _is_stream_capturing(side_stream): + _original_wait_stream(main_stream, side_stream) + forked.clear() + self.cuda_graph._segments[-1].capture_end() + + +@eager_on_graph +def break_graph() -> None: + """Insert an empty eager break between CUDA graph segments.""" + return None diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py new file mode 100644 index 000000000000..610a4da16a20 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py @@ -0,0 +1,26 @@ +# Adapted from SGLang's breakable CUDA graph implementation. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Iterator + +_breakable_cuda_graph_active: ContextVar[bool] = ContextVar( + "breakable_cuda_graph_active", default=False +) + + +def is_in_breakable_cuda_graph() -> bool: + """Return whether the current context is executing a BCG region.""" + return _breakable_cuda_graph_active.get() + + +@contextmanager +def enable_breakable_cuda_graph() -> Iterator[None]: + """Mark capture or replay work as breakable CUDA graph execution.""" + token = _breakable_cuda_graph_active.set(True) + try: + yield + finally: + _breakable_cuda_graph_active.reset(token) diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py new file mode 100644 index 000000000000..d1b8562ff568 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +from enum import Enum +from typing import Any, Callable, Iterator, Optional + +import torch +from torch import nn + +from ..utils import make_weak_ref +from .breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, + enable_breakable_cuda_graph, +) + + +class BreakableCUDAGraphRunnerState(Enum): + IDLE = "idle" + WARMUP = "warmup" + CAPTURE = "capture" + REPLAY = "replay" + + +class BreakableCUDAGraphRunner: + """Capture and replay prefill model bodies as breakable CUDA graphs.""" + + _WARMUP_STEPS = 2 + + def __init__(self, layer_model: nn.Module) -> None: + self.layer_model = layer_model + self._graphs: dict[int, BreakableCUDAGraph] = {} + self._outputs: dict[int, torch.Tensor] = {} + self._memory_pool = None + self._capture_stream = torch.cuda.Stream() + self._shared_output: Optional[torch.Tensor] = None + self._state = BreakableCUDAGraphRunnerState.IDLE + self._active_graph: Optional[BreakableCUDAGraph] = None + self._active_num_tokens: Optional[int] = None + + @property + def state(self) -> BreakableCUDAGraphRunnerState: + return self._state + + @property + def is_warming_up(self) -> bool: + return self._state == BreakableCUDAGraphRunnerState.WARMUP + + @property + def is_capturing(self) -> bool: + return self._state == BreakableCUDAGraphRunnerState.CAPTURE + + def has_graph(self, num_tokens: int) -> bool: + return num_tokens in self._graphs + + def warmup(self, engine_forward: Callable[[], Any], steps: int = _WARMUP_STEPS) -> None: + """Run the complete eager engine forward under the warmup state. + model_engine.forward will use state to determine what forward to do.""" + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot warm up BCG while runner is {self._state.value}") + self._state = BreakableCUDAGraphRunnerState.WARMUP + try: + for _ in range(steps): + engine_forward() + finally: + self._state = BreakableCUDAGraphRunnerState.IDLE + + def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None: + """Warm up eagerly, then capture one prefill token bucket.""" + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot capture BCG while runner is {self._state.value}") + if num_tokens in self._graphs: + raise ValueError(f"BCG for num_tokens={num_tokens} is already captured") + + current_stream = torch.cuda.current_stream() + self._capture_stream.wait_stream(current_stream) + graph = None + created_memory_pool = False + try: + with torch.cuda.stream(self._capture_stream): + self.warmup(engine_forward) + + # Every segment in the first BCG bucket must receive the same + # explicit pool handle. Passing None lets each CUDAGraph create + # its own private pool, which multiplies the model workspace by + # the number of eager breaks. + if self._memory_pool is None: + self._memory_pool = torch.cuda.graph_pool_handle() + created_memory_pool = True + + self._state = BreakableCUDAGraphRunnerState.CAPTURE + graph = BreakableCUDAGraph() + self._active_graph = graph + self._active_num_tokens = num_tokens + output = engine_forward() + + current_stream.wait_stream(self._capture_stream) + if not torch.is_tensor(output): + raise TypeError( + f"Breakable prefill capture requires a tensor body output, got {type(output)}" + ) + assert graph is not None + self._graphs[num_tokens] = graph + self._outputs[num_tokens] = make_weak_ref(output) + except Exception: + if graph is not None: + graph.reset() + if created_memory_pool and not self._graphs: + self._memory_pool = None + raise + finally: + self._active_graph = None + self._active_num_tokens = None + self._state = BreakableCUDAGraphRunnerState.IDLE + + @contextlib.contextmanager + def capture_context(self) -> Iterator[None]: + """Open the segmented CUDA graph capture for the active bucket.""" + if not self.is_capturing or self._active_graph is None: + raise RuntimeError("BCG capture context requested outside capture") + with ( + enable_breakable_cuda_graph(), + BreakableCUDAGraphCapture( + self._active_graph, pool=self._memory_pool, stream=self._capture_stream + ), + ): + yield + + def capture_output(self, output: torch.Tensor) -> torch.Tensor: + """Route all bucket outputs through the largest capture's buffer.""" + + if not self.is_capturing or self._active_num_tokens is None: + raise RuntimeError("BCG output registered outside capture") + num_tokens = self._active_num_tokens + if self._shared_output is None: + self._shared_output = make_weak_ref(output) + return self._shared_output + if num_tokens > self._shared_output.shape[0]: + raise ValueError( + "BCG buckets must be captured in descending order: " + f"{num_tokens} exceeds shared output size " + f"{self._shared_output.shape[0]}" + ) + self._shared_output[:num_tokens].copy_(output[:num_tokens]) + return self._shared_output[:num_tokens] + + def capture_model_body(self, outer_forward: Callable[[], Any]) -> Any: + """Run the outer model while capturing only its decoder body. + model_engine.forward is too broad and may pollute the CUDA stream + before the actual model forward. We want to reuse the functions + in forward that prepare the data and set the relevant flags.""" + if not self.is_capturing: + raise RuntimeError("BCG body capture requested outside capture") + + original_body_forward = self.layer_model.forward + captured_output = None + + def capture_forward(*args, **kwargs): + nonlocal captured_output + with self.capture_context(): + captured_output = self.capture_output(original_body_forward(*args, **kwargs)) + return captured_output + + self.layer_model.forward = capture_forward + try: + outer_forward() + if captured_output is None: + raise RuntimeError("BCG capture did not execute the model body") + return captured_output + finally: + self.layer_model.forward = original_body_forward + + def replay(self, num_tokens: int) -> torch.Tensor: + if num_tokens not in self._graphs: + raise KeyError(f"No BCG captured for num_tokens={num_tokens}") + self._graphs[num_tokens].replay() + return self._outputs[num_tokens] + + def execute(self, num_tokens: int, outer_forward: Callable[[], Any]) -> Any: + """Patch the body with replay while preserving the outer forward. + this function reuse model_engine._forward_step to set flags. + and just patch the body model forward""" + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot execute BCG while runner is {self._state.value}") + if num_tokens not in self._graphs: + raise KeyError(f"No BCG captured for num_tokens={num_tokens}") + + original_forward = self.layer_model.forward + + def replay_forward(*args, **kwargs): + del args, kwargs + return self.replay(num_tokens) + + self._state = BreakableCUDAGraphRunnerState.REPLAY + self.layer_model.forward = replay_forward + try: + with enable_breakable_cuda_graph(): + return outer_forward() + finally: + self.layer_model.forward = original_forward + self._state = BreakableCUDAGraphRunnerState.IDLE + + def clear(self) -> None: + if self._state != BreakableCUDAGraphRunnerState.IDLE: + raise RuntimeError(f"Cannot clear BCG while runner is {self._state.value}") + for graph in self._graphs.values(): + graph.reset() + self._graphs.clear() + self._outputs.clear() + self._shared_output = None + self._memory_pool = None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e71c09e2ca6a..9c2f5b98c281 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -32,6 +32,7 @@ create_input_processor_with_hash) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, + PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) from tensorrt_llm.logger import logger @@ -71,8 +72,10 @@ from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..utils import (get_model_extra_attrs, - set_per_request_piecewise_cuda_graph_flag, + get_per_request_prefill_cuda_graph_flag, + set_per_request_prefill_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) +from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner from .config_utils import is_mla from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, @@ -212,6 +215,10 @@ def _filter_piecewise_capture_num_tokens( return kept, unrecordable +# BCG uses the same capture-bucket filtering semantics as PCG. +_filter_prefill_capture_num_tokens = _filter_piecewise_capture_num_tokens + + def _filter_cuda_graph_batch_sizes(cuda_graph_batch_sizes: list[int], max_batch_size: int, max_num_tokens: int, max_total_draft_tokens: int, @@ -540,15 +547,14 @@ def __init__( f"512], max_seq_len=128, enable_padding=True).") self.torch_compile_config = self.llm_args.torch_compile_config + self.prefill_cuda_graph_backend = self.llm_args.prefill_cuda_graph_backend torch_compile_enabled = bool(self.torch_compile_config is not None) torch_compile_fullgraph = self.torch_compile_config.enable_fullgraph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_fullgraph'].default torch_compile_inductor_enabled = self.torch_compile_config.enable_inductor if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_inductor'].default - torch_compile_piecewise_cuda_graph = self.torch_compile_config.enable_piecewise_cuda_graph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ - 'enable_piecewise_cuda_graph'].default - torch_compile_piecewise_cuda_graph_num_tokens = self.torch_compile_config.capture_num_tokens if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ - 'capture_num_tokens'].default + torch_compile_piecewise_cuda_graph = (self.prefill_cuda_graph_backend == + PrefillCudaGraphBackend.PIECEWISE) torch_compile_enable_userbuffers = self.torch_compile_config.enable_userbuffers if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_userbuffers'].default torch_compile_max_num_streams = self.torch_compile_config.max_num_streams if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ @@ -557,14 +563,14 @@ def __init__( self._torch_compile_enabled = torch_compile_enabled self._torch_compile_piecewise_cuda_graph = torch_compile_piecewise_cuda_graph - piecewise_cuda_graph_num_tokens = ( - torch_compile_piecewise_cuda_graph_num_tokens - or cuda_graph_batch_sizes or []) + prefill_cuda_graph_num_tokens = self.llm_args.prefill_capture_num_tokens + if prefill_cuda_graph_num_tokens is None: + prefill_cuda_graph_num_tokens = cuda_graph_batch_sizes or [] num_extra_decoding_steps = self._get_num_extra_decoding_steps() - self._piecewise_cuda_graph_num_tokens, unrecordable = ( - _filter_piecewise_capture_num_tokens( - piecewise_cuda_graph_num_tokens, + self._prefill_cuda_graph_num_tokens, unrecordable = ( + _filter_prefill_capture_num_tokens( + prefill_cuda_graph_num_tokens, max_num_tokens=self.max_num_tokens, max_batch_size=self.batch_size, max_seq_len=self.max_seq_len, @@ -572,7 +578,7 @@ def __init__( )) if unrecordable: logger.warning( - f"Skipping piecewise CUDA graph capture for num_tokens=" + f"Skipping prefill CUDA graph capture for num_tokens=" f"{unrecordable}: exceeds reachable ceiling " f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)=" f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. " @@ -597,7 +603,7 @@ def __init__( enable_userbuffers=use_ub, enable_piecewise_cuda_graph=self. _torch_compile_piecewise_cuda_graph, - capture_num_tokens=self._piecewise_cuda_graph_num_tokens, + capture_num_tokens=self._prefill_cuda_graph_num_tokens, max_num_streams=torch_compile_max_num_streams, mapping=self.mapping) apply_llm_torch_compile = getattr(self.model, @@ -821,6 +827,17 @@ def __init__( sparse_attention_config=self.sparse_attention_config, ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) + self.breakable_cuda_graph_runner = None + if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: + decoder_model = (self.model if isinstance( + self.model, DecoderModelForCausalLM) else getattr( + self.model, "llm", None)) + if not isinstance(decoder_model, DecoderModelForCausalLM): + raise ValueError( + "breakable prefill CUDA graph requires a decoder model body" + ) + self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner( + decoder_model.model) # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( @@ -1857,14 +1874,15 @@ def _get_graphs_to_capture( def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): """Warm up or capture CUDA graphs for the configured graph shapes.""" if not (self.cuda_graph_runner.enabled - or self._torch_compile_piecewise_cuda_graph): + or self.prefill_cuda_graph_backend + != PrefillCudaGraphBackend.DISABLED): return self._capture_generation_cuda_graphs(resource_manager) # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. if not self.cuda_graph_runner.is_warmup_only: - self._capture_piecewise_cuda_graphs(resource_manager) + self._capture_prefill_cuda_graphs(resource_manager) def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager): @@ -2068,60 +2086,80 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: if self.spec_metadata is not None: self.spec_metadata.is_all_greedy_sample = True - def _capture_piecewise_cuda_graphs(self, resource_manager: ResourceManager): - """Captures piecewise CUDA graphs for context/prefill steps via torch.compile.""" - if not (self._torch_compile_piecewise_cuda_graph - and self._torch_compile_enabled): + def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): + """Capture configured CUDA graphs for context/prefill steps.""" + if (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.DISABLED + or (self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.PIECEWISE + and not self._torch_compile_enabled)): return - logger.info("Running piecewise CUDA graph warmup...") - piecewise_cuda_graph_num_tokens = sorted( - self._piecewise_cuda_graph_num_tokens, reverse=True) + logger.info("Running prefill CUDA graph warmup...") + prefill_cuda_graph_num_tokens = sorted( + self._prefill_cuda_graph_num_tokens, reverse=True) - with capture_piecewise_cuda_graph(True), self.no_cuda_graph(): - for num_tokens in piecewise_cuda_graph_num_tokens: + capture_context = (capture_piecewise_cuda_graph(True) + if self._torch_compile_piecewise_cuda_graph else + contextlib.nullcontext()) + with capture_context, self.no_cuda_graph(): + for num_tokens in prefill_cuda_graph_num_tokens: warmup_request = self._create_warmup_request( resource_manager, num_tokens, 0) with self._release_batch_context(warmup_request, resource_manager) as batch: + self._assert_all_tp_ranks_have_warmup_batch( + batch, num_tokens) if batch is None: continue logger.info( - f"Run piecewise CUDA graph warmup for num tokens={num_tokens}" + f"Run prefill CUDA graph capture for num tokens={num_tokens}" ) - # Run a few times to ensure capture - for _ in range(3): - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) + if self.breakable_cuda_graph_runner is not None: + self.breakable_cuda_graph_runner.capture( + num_tokens, lambda: self.forward( + batch, + new_tensors_device=None, + resource_manager=resource_manager)) + else: + # Run a few times to ensure torch.compile capture. + for _ in range(4): + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) torch.cuda.synchronize() gc.collect() torch.cuda.empty_cache() - # When using piecewise cuda graph, the logits may suffer severe memory fragmentation problem. - # As the number of requests grows, the blocks allocated by torch cannot be reused. - # So after piecewise cuda graph capture, a request with most requests is triggered to make - # sure that large enough blocks are allocated and can be correctly reused. - for num_tokens in piecewise_cuda_graph_num_tokens: + # The logits allocations grow with the number of requests and are not + # part of the captured model body. Warm up the largest request count so + # those allocations can be reused during stable inference. + for num_tokens in prefill_cuda_graph_num_tokens: warmup_request = self._create_warmup_request(resource_manager, num_tokens, 0, least_requests=False) with self._release_batch_context(warmup_request, resource_manager) as batch: + self._assert_all_tp_ranks_have_warmup_batch(batch, num_tokens) if batch is None: continue logger.info( - f"Run piecewise CUDA graph warmup for num tokens={num_tokens} with most requests" + f"Run prefill CUDA graph warmup for num tokens={num_tokens} with most requests" ) - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) + if self.breakable_cuda_graph_runner is not None: + with self.no_cuda_graph(): + self.breakable_cuda_graph_runner.warmup( + lambda: self.forward(batch, + new_tensors_device=None, + resource_manager= + resource_manager), + steps=1) + else: + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) torch.cuda.synchronize() ### Helper methods promoted from the original warmup method ### @@ -2862,6 +2900,9 @@ def _release_cuda_graphs(self): if hasattr(self, 'cuda_graph_runner') and self.cuda_graph_runner is not None: self.cuda_graph_runner.clear() + if (hasattr(self, 'breakable_cuda_graph_runner') + and self.breakable_cuda_graph_runner is not None): + self.breakable_cuda_graph_runner.clear() if hasattr(self, 'encoder_cuda_graph_runner' ) and self.encoder_cuda_graph_runner is not None: self.encoder_cuda_graph_runner.clear() @@ -3045,45 +3086,40 @@ def _set_spec_metadata_all_rank_num_tokens( spec_metadata.subseq_all_rank_num_tokens = all_rank_num_seqs def _get_padding_params( - self, total_num_tokens: int, num_ctx_requests: int, - attn_all_rank_num_tokens: Optional[List[int]] + self, + total_num_tokens: int, + num_ctx_requests: int, + attn_all_rank_num_tokens: Optional[List[int]], ) -> Tuple[int, bool, Optional[List[int]]]: """ Get the padding parameters for tensor padding. Return: padded_num_tokens: the padded number of tokens - can_run_piecewise_cuda_graph: whether the piecewise cuda graph can be run + can_run_prefill_cuda_graph: whether a prefill CUDA graph can run attn_all_rank_num_tokens: the number of tokens for each rank """ - padded_num_tokens = total_num_tokens - all_rank_ctx_requests = self._get_all_rank_ctx_requests( num_ctx_requests) - def get_padded_piecewise_tokens(tokens): - captured_num_tokens = self._torch_compile_backend.capture_num_tokens - return captured_num_tokens[bisect.bisect_left( - captured_num_tokens, tokens)] - - if (self._torch_compile_backend is not None - and self._torch_compile_piecewise_cuda_graph - and self._torch_compile_backend.capture_num_tokens): - max_captured_num_tokens = self._torch_compile_backend.capture_num_tokens[ - -1] - # Torch piecewise cuda graph is enabled. + def get_padded_prefill_tokens(tokens: int) -> int: + return self._prefill_cuda_graph_num_tokens[bisect.bisect_left( + self._prefill_cuda_graph_num_tokens, tokens)] + + if (self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED + and self._prefill_cuda_graph_num_tokens): + max_captured_num_tokens = self._prefill_cuda_graph_num_tokens[-1] if attn_all_rank_num_tokens is not None: - # Any rank has context requests, we enable piecewise cuda graph. has_ctx_requests = num_ctx_requests != 0 or ( all_rank_ctx_requests is not None and any(ctx_requests != 0 for ctx_requests in all_rank_ctx_requests)) - can_run_piecewise_cuda_graph = (has_ctx_requests and - max(attn_all_rank_num_tokens) - <= max_captured_num_tokens) - all_ranks_can_run_piecewise_cuda_graph = list( - self.dist.tp_allgather(can_run_piecewise_cuda_graph)) - if all(all_ranks_can_run_piecewise_cuda_graph): - padded_num_tokens = get_padded_piecewise_tokens( + can_run_prefill_cuda_graph = (has_ctx_requests + and max(attn_all_rank_num_tokens) + <= max_captured_num_tokens) + all_ranks_can_run_prefill_cuda_graph = list( + self.dist.tp_allgather(can_run_prefill_cuda_graph)) + if all(all_ranks_can_run_prefill_cuda_graph): + padded_num_tokens = get_padded_prefill_tokens( max(attn_all_rank_num_tokens)) logger.debug( f"Pad tensor with {total_num_tokens} tokens to {padded_num_tokens} tokens" @@ -3093,19 +3129,18 @@ def get_padded_piecewise_tokens(tokens): ] * len(attn_all_rank_num_tokens) else: logger.debug( - "Not all ranks can run piecewise cuda graph, disable piecewise cuda graph" + "Not all ranks can run prefill CUDA graph, disable prefill CUDA graph" ) return total_num_tokens, False, attn_all_rank_num_tokens elif num_ctx_requests != 0 and total_num_tokens <= max_captured_num_tokens: - padded_num_tokens = get_padded_piecewise_tokens( - total_num_tokens) + padded_num_tokens = get_padded_prefill_tokens(total_num_tokens) logger.debug( f"Pad tensor with {total_num_tokens} tokens to {padded_num_tokens} tokens" ) return padded_num_tokens, True, None else: logger.debug( - f"Piecewise CUDA graph cannot be used with {total_num_tokens} tokens, {num_ctx_requests} context requests" + f"Prefill CUDA graph cannot be used with {total_num_tokens} tokens, {num_ctx_requests} context requests" ) return total_num_tokens, False, None @@ -3906,7 +3941,7 @@ def _apply_steady_gen_fast_prepare( attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = \ self._get_padding_params(num_requests, 0, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + set_per_request_prefill_cuda_graph_flag(can_run_piecewise_cuda_graph) attn_metadata.padded_num_tokens = ( padded_num_tokens if padded_num_tokens != num_requests else None) virtual_num_tokens = num_requests @@ -5027,9 +5062,10 @@ def previous_seq_slots_device(): scheduled_requests, attn_metadata, peft_cache_manager, maybe_graph) attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) - padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( - total_num_tokens, num_ctx_requests, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + (padded_num_tokens, can_run_prefill_cuda_graph, + attn_all_rank_num_tokens) = self._get_padding_params( + total_num_tokens, num_ctx_requests, attn_all_rank_num_tokens) + set_per_request_prefill_cuda_graph_flag(can_run_prefill_cuda_graph) attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != total_num_tokens else None virtual_num_tokens = total_num_tokens @@ -5302,9 +5338,9 @@ def _prepare_tp_inputs_no_cache( attn_metadata.num_contexts = scheduled_requests.num_context_requests attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) - padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( + padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( num_tokens, attn_metadata.num_contexts, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + set_per_request_prefill_cuda_graph_flag(can_run_prefill_cuda_graph) attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != num_tokens else None if self.enable_attention_dp: @@ -5786,6 +5822,7 @@ def _prepare_inputs( maybe_graph: bool = False, promoted_context_request_ids: frozenset[int] = frozenset() ) -> Tuple[Dict[str, Any], Optional[torch.Tensor]]: + set_per_request_prefill_cuda_graph_flag(False) if self.mapping is not None and 'cp_type' in self.mapping.cp_config: cp_type = self.mapping.cp_config['cp_type'] if CpType.STAR == cp_type: @@ -6324,7 +6361,6 @@ def forward(self, moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) - if kv_cache_manager is None: inputs, gather_ids = self._prepare_tp_inputs_no_cache( scheduled_requests, attn_metadata, spec_metadata, @@ -6441,14 +6477,33 @@ def forward(self, self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() + breakable_runner = self.breakable_cuda_graph_runner + with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): - if not can_run_graph: - # Fallback to eager execution if graph was not used + + def forward_step(): with MoeLoadBalancerIterContext(moe_load_balancer): - outputs = self._forward_step( + return self._forward_step( inputs, gather_ids=gather_ids, gather_context_logits=gather_context_logits) + + if not can_run_graph: + if (breakable_runner is not None + and breakable_runner.is_capturing): + return breakable_runner.capture_model_body(forward_step) + + num_tokens = inputs['input_ids'].shape[0] + can_run_breakable_graph = ( + breakable_runner is not None + and get_per_request_prefill_cuda_graph_flag() + and breakable_runner.has_graph(num_tokens)) + if can_run_breakable_graph and not breakable_runner.is_warming_up: + outputs = breakable_runner.execute( + num_tokens, forward_step) + else: + # real eager or BCG warmup or PCG + outputs = forward_step() else: needs_capture = self.cuda_graph_runner.needs_capture(key) if needs_capture: diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 641b948a47b0..64cac4c80d4e 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import contextlib import functools import os @@ -163,8 +177,8 @@ def make_weak_ref(x): elif isinstance(x, list): return [make_weak_ref(i) for i in x] elif isinstance(x, dict): - return {k: make_weak_ref(v) for k, v in x.items()} - elif isinstance(x, (int, float, bool)): + return {make_weak_ref(k): make_weak_ref(v) for k, v in x.items()} + elif x is None or isinstance(x, (int, float, str, bool)): return x else: raise TypeError(f"Invalid type {type(x)} to make weak ref") @@ -390,12 +404,14 @@ def piecewise_cuda_graph(enable: bool): set_piecewise_cuda_graph_flag(prev_enable) -def set_per_request_piecewise_cuda_graph_flag(enable: bool): - _global_attrs.per_request_piecewise_cuda_graph_flag = enable +def set_per_request_prefill_cuda_graph_flag(enable: bool): + """Set whether the current batch can use its prefill CUDA graph backend.""" + _global_attrs.per_request_prefill_cuda_graph_flag = enable -def get_per_request_piecewise_cuda_graph_flag() -> bool: - return getattr(_global_attrs, 'per_request_piecewise_cuda_graph_flag', True) +def get_per_request_prefill_cuda_graph_flag() -> bool: + """Return whether the current batch can use its prefill CUDA graph backend.""" + return getattr(_global_attrs, 'per_request_prefill_cuda_graph_flag', True) def create_lm_head_tp_mapping(mapping: Mapping, token_count: int) -> Mapping: diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 7a3583907041..12f52206c354 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -20,12 +20,12 @@ MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, NGramDecodingConfig, PARDDecodingConfig, - PrometheusMetricsConfig, ReorderRequestPolicyConfig, - RocketSparseAttentionConfig, SADecodingConfig, - SAEnhancerConfig, SaveHiddenStatesDecodingConfig, - SchedulerConfig, SkipSoftmaxAttentionConfig, - TorchCompileConfig, TorchLlmArgs, - TriAttentionKvCacheCompressionConfig, + PrefillCudaGraphBackend, PrometheusMetricsConfig, + ReorderRequestPolicyConfig, RocketSparseAttentionConfig, + SADecodingConfig, SAEnhancerConfig, + SaveHiddenStatesDecodingConfig, SchedulerConfig, + SkipSoftmaxAttentionConfig, TorchCompileConfig, + TorchLlmArgs, TriAttentionKvCacheCompressionConfig, UserProvidedDecodingConfig) from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder @@ -92,6 +92,7 @@ 'SkipSoftmaxAttentionConfig', 'TriAttentionKvCacheCompressionConfig', 'PrometheusMetricsConfig', + 'PrefillCudaGraphBackend', 'ThinkingBudgetLogitsProcessor', 'add_thinking_budget_logits_processor', 'MultimodalConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f1f382dc8c0d..538c108914e1 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4985,6 +4985,18 @@ class SamplerType(StrEnum): auto = "auto" +class PrefillCudaGraphBackend(StrEnum): + """CUDA graph implementation used for prefill requests.""" + + DISABLED = "disabled" + PIECEWISE = "piecewise" + BREAKABLE = "breakable" + + +_DEFAULT_PREFILL_CAPTURE_NUM_TOKENS = [2**i for i in range(8) + ] + [i for i in range(256, 3073, 256)] + + class TorchCompileConfig(StrictBaseModel): """Configuration for torch.compile.""" enable_fullgraph: bool = Field( @@ -4996,13 +5008,12 @@ class TorchCompileConfig(StrictBaseModel): enable_piecewise_cuda_graph: bool = Field( default=False, - description="Enable piecewise CUDA graph in torch.compile.") + description="Deprecated. Use prefill_cuda_graph_backend='piecewise' " + "instead.") capture_num_tokens: Optional[List[PositiveInt]] = Field( default=None, - description= - "List of num of tokens to capture the piecewise CUDA graph for. If not provided, the number of tokens will be the same as cuda_graph_config.batch_sizes." - ) + description="Deprecated. Use prefill_capture_num_tokens instead.") @field_validator('capture_num_tokens') @classmethod @@ -5017,17 +5028,10 @@ def validate_capture_num_tokens(cls, v): "When torch compile is enabled, userbuffers is enabled by default.") max_num_streams: PositiveInt = Field( - default=1, + default=3, description= "The maximum number of CUDA streams to use for torch.compile.") - @model_validator(mode='after') - def set_default_capture_num_tokens(self) -> 'TorchCompileConfig': - if self.enable_piecewise_cuda_graph and self.capture_num_tokens is None: - self.capture_num_tokens = [2**i for i in range(8) - ] + [i for i in range(256, 3073, 256)] - return self - class TorchLlmArgs(BaseLlmArgs): # PyTorch backend specific configurations @@ -5221,6 +5225,20 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: torch_compile_config: Optional[TorchCompileConfig] = Field( default=None, description="Torch compile config.", status="prototype") + prefill_cuda_graph_backend: PrefillCudaGraphBackend = Field( + default=PrefillCudaGraphBackend.DISABLED, + description="CUDA graph implementation used for prefill requests. " + "Defaults to disabled.", + status="prototype", + telemetry=TelemetryField.categorical("disabled", "piecewise", + "breakable")) + + prefill_capture_num_tokens: Optional[List[int]] = Field( + default=None, + description= + "Token-count buckets captured by the selected prefill CUDA graph implementation.", + status="prototype") + enable_autotuner: bool = Field( default=True, description= @@ -5517,6 +5535,63 @@ def validate_encode_only_torch_compile_config(self) -> 'TorchLlmArgs': "graphs or disable enable_piecewise_cuda_graph.") return self + @model_validator(mode="after") + def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': + """Normalize legacy piecewise CUDA graph options into prefill fields.""" + backend_is_explicit = "prefill_cuda_graph_backend" in self.model_fields_set + buckets_are_explicit = "prefill_capture_num_tokens" in self.model_fields_set + compile_config = self.torch_compile_config + legacy_buckets_are_explicit = (compile_config is not None + and "capture_num_tokens" + in compile_config.model_fields_set) + + if compile_config is not None and compile_config.enable_piecewise_cuda_graph: + if (backend_is_explicit and self.prefill_cuda_graph_backend + != PrefillCudaGraphBackend.PIECEWISE): + raise ValueError( + "torch_compile_config.enable_piecewise_cuda_graph conflicts " + "with prefill_cuda_graph_backend") + logger.warning( + "TorchCompileConfig.enable_piecewise_cuda_graph is deprecated; " + "use prefill_cuda_graph_backend='piecewise' instead.") + self.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE + + legacy_buckets = (compile_config.capture_num_tokens + if compile_config is not None else None) + if legacy_buckets_are_explicit: + logger.warning( + "TorchCompileConfig.capture_num_tokens is deprecated; use " + "prefill_capture_num_tokens instead.") + if (legacy_buckets is not None and buckets_are_explicit + and self.prefill_capture_num_tokens is not None + and sorted(set(legacy_buckets)) != sorted( + set(self.prefill_capture_num_tokens))): + raise ValueError( + "torch_compile_config.capture_num_tokens conflicts with " + "prefill_capture_num_tokens") + if not buckets_are_explicit and legacy_buckets is not None: + self.prefill_capture_num_tokens = list(legacy_buckets) + + if self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED: + if self.prefill_capture_num_tokens is None: + self.prefill_capture_num_tokens = list( + _DEFAULT_PREFILL_CAPTURE_NUM_TOKENS) + if self.encode_only: + raise ValueError( + "encode_only does not support prefill CUDA graphs") + + if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE: + if self.torch_compile_config is None: + self.torch_compile_config = TorchCompileConfig() + elif (self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.BREAKABLE + and self.torch_compile_config is not None): + raise ValueError( + "breakable prefill CUDA graph does not support torch_compile_config" + ) + + return self + @model_validator(mode="after") def validate_speculative_config(self): if self.speculative_config: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e2ce02373430..89c76b57182f 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1204,6 +1204,24 @@ "kind": "value", "path": "pp_partition" }, + { + "allowed_values": [], + "annotation": "Optional[List[int]]", + "converter": "", + "kind": "value", + "path": "prefill_capture_num_tokens" + }, + { + "allowed_values": [ + "disabled", + "piecewise", + "breakable" + ], + "annotation": "", + "converter": "allowlist", + "kind": "categorical", + "path": "prefill_cuda_graph_backend" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 28e122fd7c45..de4578d584ad 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -33,9 +33,10 @@ DFlashDecodingConfig, DSparkDecodingConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MambaStateConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, - NGramDecodingConfig, PARDDecodingConfig, RocketSparseAttentionConfig, - SADecodingConfig, SamplingParams, SchedulerConfig, - SkipSoftmaxAttentionConfig, SAEnhancerConfig, TorchCompileConfig) + NGramDecodingConfig, PARDDecodingConfig, PrefillCudaGraphBackend, + RocketSparseAttentionConfig, SADecodingConfig, SamplingParams, + SchedulerConfig, SkipSoftmaxAttentionConfig, SAEnhancerConfig, + TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -3561,6 +3562,66 @@ def test_nvfp4_multi_gpus_piecewise_cuda_graph(self, tp_size, pp_size, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip_less_mpi_world_size(8) + @skip_pre_blackwell + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size,mtp_nextn,attention_dp,max_batch_size,moe_backend,fp8kv,chunked_prefill", + [ + (8, 1, 8, 0, True, 24, "CUTLASS", False, False), + ], + ids=["baseline"]) + def test_nvfp4_multi_gpus_breakable_cuda_graph(self, tp_size, pp_size, + ep_size, mtp_nextn, + attention_dp, max_batch_size, + moe_backend, fp8kv, + chunked_prefill): + sm_version = get_sm_version() + if moe_backend == "TRTLLM" and sm_version in (120, 121): + pytest.skip(f"{moe_backend} backend does not support SM 120 or 121") + + moe_config = MoeConfig(backend=moe_backend, max_num_tokens=16384) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7) + if fp8kv: + kv_cache_config.dtype = "fp8" + kv_cache_config.enable_block_reuse = True + + pytorch_config = dict( + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, + max_batch_size=max_batch_size, + ), + moe_config=moe_config, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + prefill_capture_num_tokens=[2048, 8192], + ) + + mtp_config = None + if mtp_nextn > 0: + mtp_config = MTPDecodingConfig(max_draft_len=mtp_nextn) + + llm_kwargs = dict( + max_batch_size=max_batch_size, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) + if chunked_prefill: + llm_kwargs.update( + enable_chunked_prefill=True, + max_num_tokens=8192, + ) + + with LLM(f"{llm_models_root()}/DeepSeek-V3.2-Exp-FP4-v2", + **pytorch_config, **llm_kwargs) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.skip_less_mpi_world_size(8) @skip_pre_blackwell @pytest.mark.parametrize( @@ -3920,6 +3981,86 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): eplb_config, mtp_nextn=mtp_nextn) + @pytest.mark.skip_less_mpi_world_size(8) + @pytest.mark.threadleak(enabled=False) + def test_mixed_breakable_cuda_graph(self): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + base_prompt_ids = tokenizer.encode( + "TensorRT-LLM accelerates reliable large language model inference " + "with efficient attention, parallelism, and CUDA graphs. ", + add_special_tokens=False, + ) + assert base_prompt_ids + + def make_prompt(prompt_length): + return (base_prompt_ids * + ((prompt_length + len(base_prompt_ids) - 1) // + len(base_prompt_ids)))[:prompt_length] + + generation_prompt = make_prompt(64) + context_prompt = make_prompt(129) + sampling_params = SamplingParams( + max_tokens=8, + min_tokens=8, + seed=42, + temperature=0, + ignore_eos=True, + detokenize=False, + add_special_tokens=False, + ) + common_llm_kwargs = dict( + tensor_parallel_size=8, + moe_expert_parallel_size=8, + moe_config=MoeConfig(backend="TRTLLM"), + enable_attention_dp=True, + max_batch_size=8, + max_num_tokens=1024, + max_seq_len=2048, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + dtype="fp8", + free_gpu_memory_fraction=0.6, + ), + cuda_graph_config=CudaGraphConfig( + batch_sizes=[1, 2, 4, 6, 8], + enable_padding=True, + ), + ) + + def run(backend): + with LLM( + self.MODEL_PATH, + **common_llm_kwargs, + prefill_cuda_graph_backend=backend, + prefill_capture_num_tokens=[128, 256, 512, 1024], + ) as llm: + generation_request = llm.generate_async( + generation_prompt, + sampling_params=sampling_params, + streaming=True, + ) + next(generation_request) + assert not generation_request.finished + + # Admit a context request while the first request is decoding. + context_request = llm.generate_async( + context_prompt, + sampling_params=sampling_params, + streaming=False, + ) + generation_output = generation_request.result() + context_output = context_request.result() + return [ + generation_output.outputs[0].token_ids, + context_output.outputs[0].token_ids, + ] + + eager_token_ids = run(PrefillCudaGraphBackend.DISABLED) + breakable_token_ids = run(PrefillCudaGraphBackend.BREAKABLE) + assert breakable_token_ids == eager_token_ids + _DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT = ( "Solve the problem carefully. End your response with a final line exactly " @@ -6157,6 +6298,46 @@ def test_bf16(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_blackwell + @pytest.mark.threadleak(enabled=False) + def test_bf16_breakable_prefill_cuda_graph(self): + model_path = f"{llm_models_root()}/Qwen3.5-4B" + prompts = [ + [[17] * 128], + [[17] * 129], + # The second request is admitted while the first is decoding, + # exercising BCG replay for a mixed context/decode batch. + [[17] * 64, [23] * 65], + [[31] * 256], + ] + sampling_params = SamplingParams(max_tokens=4) + + def run(backend): + results = [] + with LLM( + model_path, + trust_remote_code=True, + max_seq_len=1024, + max_num_tokens=512, + max_batch_size=4, + disable_overlap_scheduler=True, + kv_cache_config=self.kv_cache_config, + cuda_graph_config=CudaGraphConfig(enable_padding=True, + max_batch_size=4), + prefill_cuda_graph_backend=backend, + prefill_capture_num_tokens=[128, 256, 512], + ) as llm: + for batch in prompts: + results.append([ + output.outputs[0].token_ids for output in llm.generate( + batch, sampling_params=sampling_params) + ]) + return results + + eager_results = run(PrefillCudaGraphBackend.DISABLED) + breakable_results = run(PrefillCudaGraphBackend.BREAKABLE) + assert breakable_results == eager_results + @skip_pre_hopper def test_fp8(self): model_path = f"{llm_models_root()}/Qwen3.5-4B-FP8" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index bcacc3e90bbc..16023ee556ea 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -57,6 +57,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a16_mxfp4[latency-TRTLLM] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention[target_sparsity_0.9-fp8kv=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=True] + - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16_breakable_prefill_cuda_graph - accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[TRTLLM] - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 77ca7e7f50bf..f6fa9b3e5fd7 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -182,6 +182,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4_mtp_index_share[tp_size=8-ep_size=8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_mixed_breakable_cuda_graph TIMEOUT (120) ISOLATION - examples/test_deepseek_v4_pro.py::test_short_token_boundary_smoke TIMEOUT (120) - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) @@ -254,6 +255,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp3] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_pp4_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[baseline_fp8kv] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[baseline] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8_attn_dp] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[ep8] TIMEOUT (60) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py index 7325b3cecf25..02a2ee09e901 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -271,7 +271,11 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): # Call the deepseek_v4 output projection (mla_rope_inplace modifies attn_out_latent # in-place, so clone before passing to preserve original for reference) - output = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) + output = mla._deepseek_v4_o_proj( + attn_out_latent.clone(), + position_ids, + enable_dsv4_epilogue_fusion=False, + ) # Calculate reference output if dtype_str == "bf16": diff --git a/tests/unittest/_torch/compilation/test_remove_copy_pass.py b/tests/unittest/_torch/compilation/test_remove_copy_pass.py index ef78351e9453..8827c09e69f3 100644 --- a/tests/unittest/_torch/compilation/test_remove_copy_pass.py +++ b/tests/unittest/_torch/compilation/test_remove_copy_pass.py @@ -98,7 +98,7 @@ def test_remove_copy_for_mutates_tensor_list( graph.lint() -def test_remove_copy_for_mutates_args_restores_optional_none() -> None: +def test_remove_copy_for_mla_restores_final_output_mutation() -> None: graph = Graph() hidden_states = graph.placeholder("hidden_states") output = graph.placeholder("output") @@ -111,11 +111,8 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: "position_ids": None, "layer_idx": "0", "latent_cache_gen": None, - "enable_dsv4_epilogue_fusion": False, "_all_bases": (output,), "_output_base_index": 0, - "_dsv4_output_base_index": None, - "_dsv4_output_sf_base_index": None, }, ) mutated_output = graph.call_function(getitem, args=(functionalized, 1)) @@ -127,49 +124,5 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: inplace_nodes = [node for node in graph.nodes if node.target == inplace_func] assert len(inplace_nodes) == 1 assert inplace_nodes[0].kwargs["output"] is output - assert inplace_nodes[0].kwargs["dsv4_output"] is None - assert inplace_nodes[0].kwargs["dsv4_output_sf"] is None assert clone.args[0] is output graph.lint() - - -def test_remove_copy_for_mutates_args_rejects_getitem_for_optional_none( - monkeypatch: pytest.MonkeyPatch, -) -> None: - graph = Graph() - hidden_states = graph.placeholder("hidden_states") - output = graph.placeholder("output") - inplace_func = torch.ops.trtllm.mla_custom_op_inplace.default - functionalized = graph.call_function( - auto_functionalized_v2, - args=(inplace_func,), - kwargs={ - "hidden_states": hidden_states, - "position_ids": None, - "layer_idx": "0", - "latent_cache_gen": None, - "enable_dsv4_epilogue_fusion": False, - "_all_bases": (output,), - "_output_base_index": 0, - "_dsv4_output_base_index": None, - "_dsv4_output_sf_base_index": None, - }, - ) - optional_output = graph.call_function(getitem, args=(functionalized, 2)) - clone = graph.call_function(torch.ops.aten.clone.default, args=(optional_output,)) - graph.output(clone) - - monkeypatch.setattr( - remove_copy_pass, - "inplace_info", - lambda: {inplace_func: {1: "output", 2: "dsv4_output"}}, - ) - - with pytest.raises( - AssertionError, - match=( - "getitem user for optional output 'dsv4_output' has no " - "base tensor -- graph is malformed" - ), - ): - remove_copy_pass.remove_copy_for_mutates_args(graph) diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py new file mode 100644 index 000000000000..1c342aa61875 --- /dev/null +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -0,0 +1,290 @@ +# Adapted from SGLang's breakable CUDA graph tests. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, + break_graph, + eager_on_graph, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph import _copy_output +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import ( + BreakableCUDAGraphRunner, + BreakableCUDAGraphRunnerState, +) +from tensorrt_llm._torch.utils import make_weak_ref + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _capture(body): + graph = BreakableCUDAGraph() + with BreakableCUDAGraphCapture(graph, stream=torch.cuda.Stream()): + body() + return graph + + +def test_no_break_capture_and_repeated_replay(): + x = torch.zeros(4, device="cuda") + output = torch.zeros_like(x) + graph = _capture(lambda: output.copy_(x + 1)) + + assert graph.num_segments == 1 + assert graph.num_breaks == 0 + for value in (5, 11): + x.fill_(value) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, value + 1)) + + +def test_single_and_multiple_breakpoints(): + @eager_on_graph + def add_one(value): + return value + 1 + + @eager_on_graph + def double(value): + return value * 2 + + x = torch.zeros(4, device="cuda") + output = torch.zeros_like(x) + + def body(): + value = add_one(x + 1) + value = double(value + 1) + output.copy_(value) + + graph = _capture(body) + assert graph.num_segments == 3 + assert graph.num_breaks == 2 + + x.fill_(5) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 16)) + + +def test_outside_capture(): + @eager_on_graph + def outside(value): + return value + 2 + + value = torch.tensor([1.0, 2.0], device="cuda") + torch.testing.assert_close(outside(value), value + 2) + + +def test_eager_on_graph_during_torch_compile(): + @eager_on_graph + def add_one(value): + return value + 1 + + compiled_add_one = torch.compile(add_one, backend="eager", fullgraph=True) + value = torch.ones(4, device="cuda") + + torch.testing.assert_close(compiled_add_one(value), value + 1) + + +def test_make_weak_ref_supports_value_types_and_rejects_objects(): + unsupported = object() + with pytest.raises(TypeError, match="Invalid type"): + make_weak_ref(unsupported) + + value = {"nested": (None, "value", [1, 2.0, True])} + assert make_weak_ref(value) == value + + with pytest.raises(TypeError, match="Invalid type"): + make_weak_ref({unsupported: "value"}) + + +def test_break_graph_inserts_empty_breakpoint(): + x = torch.zeros(4, device="cuda") + output = torch.zeros_like(x) + + def body(): + value = x + 1 + break_graph() + output.copy_(value + 2) + + graph = _capture(body) + assert graph.num_segments == 2 + assert graph.num_breaks == 1 + x.fill_(10) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 13)) + + +def test_output_writeback_for_tensor_dict_and_object(): + class Output: + def __init__(self, tensor, label): + self.tensor = tensor + self.label = label + + tensor = torch.zeros(4, device="cuda") + assert _copy_output(tensor, torch.full_like(tensor, 3)) is tensor + torch.testing.assert_close(tensor, torch.full_like(tensor, 3)) + + output_dict = {"value": torch.zeros(4, device="cuda")} + assert _copy_output(output_dict, {"value": torch.ones(4, device="cuda")}) is output_dict + torch.testing.assert_close(output_dict["value"], torch.ones(4, device="cuda")) + + output_object = Output(torch.zeros(4, device="cuda"), "old") + assert ( + _copy_output(output_object, Output(torch.full((4,), 2.0, device="cuda"), "new")) + is output_object + ) + torch.testing.assert_close(output_object.tensor, torch.full_like(output_object.tensor, 2)) + assert output_object.label == "new" + + +def test_side_stream_is_joined_before_segment_end(): + x = torch.ones(4, device="cuda") + output = torch.zeros_like(x) + side_stream = torch.cuda.Stream() + + def body(): + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + output.copy_((x + 1) * 2) + + graph = _capture(body) + x.fill_(3) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 8)) + + +class _Body(nn.Module): + def __init__(self): + super().__init__() + self.forward_calls = 0 + + def forward(self, value): + self.forward_calls += 1 + return value + 1 + + +class _LogitsProcessor(nn.Module): + def __init__(self): + super().__init__() + self.forward_calls = 0 + + def forward(self, value): + self.forward_calls += 1 + return value * 2 + + +def test_runner_warmup_capture_execute_and_shared_output(): + body = _Body().cuda() + logits_processor = _LogitsProcessor().cuda() + runner = BreakableCUDAGraphRunner(body) + counters = {"outer": 0} + inputs = {} + + def engine_forward(): + counters["outer"] += 1 + if runner.is_capturing: + return runner.capture_model_body( + lambda: {"logits": logits_processor(body(inputs["value"]))} + ) + return {"logits": logits_processor(body(inputs["value"]))} + + inputs["value"] = torch.zeros((8, 4), device="cuda") + runner.capture(8, engine_forward) + first_shared_output = runner._shared_output + inputs["value"] = torch.zeros((4, 4), device="cuda") + runner.capture(4, engine_forward) + + assert runner.state == BreakableCUDAGraphRunnerState.IDLE + assert counters == {"outer": 6} + assert body.forward_calls == 6 + assert logits_processor.forward_calls == 6 + assert runner._shared_output is first_shared_output + + original_forward = body.forward + inputs["value"].fill_(3) + result = runner.execute(4, engine_forward) + torch.cuda.synchronize() + torch.testing.assert_close(result["logits"], torch.full((4, 4), 8.0, device="cuda")) + assert counters == {"outer": 7} + assert body.forward_calls == 6 + assert logits_processor.forward_calls == 7 + assert body.forward == original_forward + + +def test_runner_first_bucket_segments_share_one_memory_pool(): + class BreakableBody(nn.Module): + @staticmethod + @eager_on_graph + def eager_add_one(value): + return value + 1 + + @staticmethod + @eager_on_graph + def eager_double(value): + return value * 2 + + def forward(self, value): + value = self.eager_add_one(value + 1) + return self.eager_double(value + 1) + + body = BreakableBody().cuda() + runner = BreakableCUDAGraphRunner(body) + inputs = {"value": torch.zeros((8, 4), device="cuda")} + + def engine_forward(): + if runner.is_capturing: + return runner.capture_model_body(lambda: body(inputs["value"])) + return body(inputs["value"]) + + runner.capture(8, engine_forward) + + graph = runner._graphs[8] + assert graph.num_segments == 3 + assert runner._memory_pool is not None + assert all(segment.pool() == runner._memory_pool for segment in graph._segments) + + +def test_runner_graph_miss_nested_execute_and_exception_recovery(): + body = _Body().cuda() + runner = BreakableCUDAGraphRunner(body) + with pytest.raises(KeyError, match="No BCG captured"): + runner.execute(4, lambda: None) + + runner._graphs[4] = object() + runner._outputs[4] = torch.zeros(1, device="cuda") + original_forward = body.forward + + def nested(): + return runner.execute(4, lambda: None) + + with pytest.raises(RuntimeError, match="while runner is replay"): + runner.execute(4, nested) + assert runner.state == BreakableCUDAGraphRunnerState.IDLE + assert body.forward == original_forward + + def fail(): + raise ValueError("expected") + + with pytest.raises(ValueError, match="expected"): + runner.execute(4, fail) + assert runner.state == BreakableCUDAGraphRunnerState.IDLE + assert body.forward == original_forward + + +def test_runner_warmup_exception_restores_idle_state(): + runner = BreakableCUDAGraphRunner(_Body().cuda()) + + def fail(): + raise ValueError("expected") + + with pytest.raises(ValueError, match="expected"): + runner.warmup(fail) + assert runner.state == BreakableCUDAGraphRunnerState.IDLE diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 884fb8e2eebb..2b8b48b8cbf4 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -235,10 +235,15 @@ def _make_forward_only_engine( ) engine.spec_metadata = spec_metadata engine._set_up_spec_metadata = Mock(return_value=spec_metadata) - engine._prepare_inputs = Mock(return_value=({"prepared": True}, None)) + prepared_inputs = { + "prepared": True, + "input_ids": torch.zeros(2, dtype=torch.int32), + } + engine._prepare_inputs = Mock(return_value=(prepared_inputs, None)) outputs = {"logits": object()} engine._forward_step = Mock(return_value=outputs) engine._execute_logit_post_processors = Mock() + engine.breakable_cuda_graph_runner = None runner = Mock() runner.enabled = runner_enabled @@ -674,7 +679,8 @@ def test_forward_commits_candidate_only_on_graph_hit(self) -> None: prepare_args = engine._prepare_inputs.call_args.args self.assertIs(prepare_args[0], graph_batch) self.assertEqual(prepare_args[-1], frozenset({1})) - runner.replay.assert_called_once_with(key, {"prepared": True}) + prepared_inputs = engine._prepare_inputs.return_value[0] + runner.replay.assert_called_once_with(key, prepared_inputs) engine._forward_step.assert_not_called() engine._execute_logit_post_processors.assert_called_once_with( batch, outputs) @@ -748,7 +754,8 @@ def test_zero_runtime_draft_speculation_commits_graph_candidate( self.assertEqual( semantic_attn_metadata.update_spec_dec_param.call_args. kwargs["num_contexts"], 1) - runner.replay.assert_called_once_with(key, {"prepared": True}) + prepared_inputs = engine._prepare_inputs.return_value[0] + runner.replay.assert_called_once_with(key, prepared_inputs) def test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager( self) -> None: @@ -829,7 +836,8 @@ def test_forward_allows_guided_context_logits_on_graph_hit(self) -> None: prepare_args = engine._prepare_inputs.call_args.args self.assertIs(prepare_args[0], graph_batch) self.assertEqual(prepare_args[-1], frozenset({context.py_request_id})) - runner.replay.assert_called_once_with(key, {"prepared": True}) + prepared_inputs = engine._prepare_inputs.return_value[0] + runner.replay.assert_called_once_with(key, prepared_inputs) def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: engine, runner, resource_manager, _, _ = _make_forward_only_engine(None) diff --git a/tests/unittest/_torch/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py index 908069a02560..450279376da5 100644 --- a/tests/unittest/_torch/modules/test_mla_registry.py +++ b/tests/unittest/_torch/modules/test_mla_registry.py @@ -13,14 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import Mock, patch +import pytest import torch from torch import nn from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.mla import MLA +from tensorrt_llm._torch.modules.mla import MLA, create_mla_outputs_impl from tensorrt_llm.functional import PositionEmbeddingType @@ -80,3 +82,187 @@ def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None: assert registry["0"]() is target_mla assert registry["0_0"]() is draft_mla assert registry["0_1"]() is next_mla + + +def test_dsv4_epilogue_fusion_returns_final_output_inside_breakable_graph() -> None: + metadata = SimpleNamespace(num_contexts=1, num_generations=1, num_tokens=5) + mla_layer = Mock(spec=MLA) + mla_layer._should_use_dsv4_epilogue_fusion.return_value = True + mla_layer.create_output.return_value = torch.empty(8, 4, 2) + hidden_states = torch.empty(8, 8) + + with ( + patch( + "tensorrt_llm._torch.modules.mla._extract_mla_extra_attrs", + return_value=(metadata, mla_layer), + ), + patch( + "tensorrt_llm._torch.modules.mla.is_in_breakable_cuda_graph", + return_value=True, + ), + ): + output = create_mla_outputs_impl(hidden_states, "0") + + assert output is mla_layer.create_output.return_value + mla_layer.create_output.assert_called_once_with( + hidden_states, + 1, + enable_dsv4_epilogue_fusion=True, + ) + mla_layer._create_dsv4_epilogue_buffers.assert_not_called() + + +def test_create_mla_outputs_custom_op_returns_tensor() -> None: + schema = torch.ops.trtllm.create_mla_outputs.default._schema + assert [str(return_value.type) for return_value in schema.returns] == ["Tensor"] + + +def test_mla_custom_op_marks_only_final_output_mutable() -> None: + schema = torch.ops.trtllm.mla_custom_op_inplace.default._schema + mutated_args = [ + arg.name + for arg in schema.arguments + if arg.alias_info is not None and arg.alias_info.is_write + ] + assert mutated_args == ["output"] + + +def test_dsv4_epilogue_fusion_supports_mixed_batch() -> None: + mla_layer = SimpleNamespace( + _disable_dsv4_epilogue_fusion=False, + is_deepseek_v4=True, + mapping=SimpleNamespace( + has_cp_helix=lambda: False, + enable_attention_dp=True, + ), + num_heads=128, + num_heads_tp=128, + mqa=SimpleNamespace( + sparse_params=object(), + has_fp8_kv_cache=True, + ), + o_a_proj=SimpleNamespace(dtype=torch.float8_e4m3fn), + kv_lora_rank=448, + qk_rope_head_dim=64, + qk_head_dim=512, + v_head_dim=512, + n_local_groups=8, + inverse_rotary_emb=SimpleNamespace(is_neox=False), + ) + + with patch("tensorrt_llm._torch.modules.mla.is_sm_100f", return_value=True): + assert MLA._should_use_dsv4_epilogue_fusion(mla_layer, 1, 1) + + +def test_dsv4_fusion_create_output_uses_bucket_token_count() -> None: + mla_layer = SimpleNamespace( + n_local_groups=4, + o_lora_rank=3, + dtype=torch.bfloat16, + ) + hidden_states = torch.empty(8, 16) + + output = MLA.create_output( + mla_layer, + hidden_states, + num_contexts=1, + enable_dsv4_epilogue_fusion=True, + ) + + assert output.shape == (8, 4, 3) + assert output.dtype == torch.bfloat16 + + +def test_dsv4_fusion_o_proj_only_flattens_lora_output() -> None: + projected = torch.randn(7, 5) + mla_layer = SimpleNamespace( + n_local_groups=4, + o_lora_rank=3, + o_b_proj=Mock(return_value=projected), + ) + lora_o = torch.randn(7, 4, 3) + + output = MLA._deepseek_v4_o_proj( + mla_layer, + lora_o, + enable_dsv4_epilogue_fusion=True, + ) + + assert output is projected + mla_layer.o_b_proj.assert_called_once() + torch.testing.assert_close(mla_layer.o_b_proj.call_args.args[0], lora_o.flatten(1)) + + +def test_dsv4_epilogue_buffers_use_real_token_count() -> None: + mla_layer = SimpleNamespace( + n_local_groups=4, + num_heads_tp=128, + v_head_dim=512, + ) + q = torch.empty(8, 16) + + fp8_o, output_sf = MLA._create_dsv4_epilogue_buffers(mla_layer, q, num_tokens=5) + + assert fp8_o.shape == (4, 5, 32 * 512) + assert output_sf.shape == (4, 32 * 4, 8) + + +@pytest.mark.parametrize( + "num_context_tokens,num_generation_tokens,bucket_tokens", + [(5, 0, 8), (0, 3, 4), (5, 3, 12)], +) +def test_dsv4_epilogue_bmm_writes_only_phase_ranges( + num_context_tokens: int, + num_generation_tokens: int, + bucket_tokens: int, +) -> None: + groups = 2 + rank = 3 + output = torch.full((bucket_tokens, groups, rank), -1.0) + mla_layer = SimpleNamespace( + o_a_proj=torch.empty(0), + o_a_proj_scale=torch.empty(0), + ) + + def fake_bmm(_attn_fp8, _weight, attn_scale, _weight_scale, phase_output): + phase_output.fill_(attn_scale.item()) + + with patch.object( + torch.ops.trtllm, + "cute_dsl_fp8_bmm_blackwell", + side_effect=fake_bmm, + ) as bmm: + context_epilogue = None + if num_context_tokens: + context_epilogue = ( + torch.empty(groups, num_context_tokens, 4), + torch.tensor(11.0), + ) + generation_epilogue = None + if num_generation_tokens: + generation_epilogue = ( + torch.empty(groups, num_generation_tokens, 4), + torch.tensor(22.0), + ) + MLA._run_dsv4_o_lora_bmms( + mla_layer, + output, + num_context_tokens, + num_context_tokens + num_generation_tokens, + context_epilogue, + generation_epilogue, + ) + + assert bmm.call_count == bool(num_context_tokens) + bool(num_generation_tokens) + if num_context_tokens: + torch.testing.assert_close( + output[:num_context_tokens], torch.full_like(output[:num_context_tokens], 11.0) + ) + if num_generation_tokens: + generation_end = num_context_tokens + num_generation_tokens + torch.testing.assert_close( + output[num_context_tokens:generation_end], + torch.full_like(output[num_context_tokens:generation_end], 22.0), + ) + real_tokens = num_context_tokens + num_generation_tokens + torch.testing.assert_close(output[real_tokens:], torch.full_like(output[real_tokens:], -1.0)) diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index c9b9a42388a7..30f3727f688f 100644 --- a/tests/unittest/api_stability/api_stability_core.py +++ b/tests/unittest/api_stability/api_stability_core.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + # autoflake: skip_file import copy import inspect @@ -29,7 +32,7 @@ from tensorrt_llm.llmapi import (CalibConfig, CompletionOutput, GuidedDecodingParams, QuantConfig, RequestOutput, SamplingParams) -from tensorrt_llm.llmapi.llm_args import SamplerType +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend, SamplerType from tensorrt_llm.llmapi.llm_utils import LlmArgs from tensorrt_llm.logger import Singleton from tensorrt_llm.sampling_params import LogprobMode diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 5cac19e246df..c0ebc5cc9bd6 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -219,6 +219,14 @@ methods: annotation: Optional[tensorrt_llm.llmapi.llm_args.TorchCompileConfig] default: null status: prototype + prefill_cuda_graph_backend: + annotation: tensorrt_llm.llmapi.llm_args.PrefillCudaGraphBackend + default: disabled + status: prototype + prefill_capture_num_tokens: + annotation: Optional[List[int]] + default: null + status: prototype enable_autotuner: annotation: bool default: True diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index c5bdfb6a6a9e..ce5891685770 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -47,7 +47,8 @@ MambaStateConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, - PeftCacheConfig, PybindMirror, + PeftCacheConfig, + PrefillCudaGraphBackend, PybindMirror, RayPlacementConfig, SkipSoftmaxAttentionConfig, SleepConfig, SpeculativeConfig, @@ -1788,13 +1789,13 @@ class TestPiecewiseCudaGraphCaptureDefaults: Three invariants are exercised: - 1. `TorchCompileConfig.capture_num_tokens` defaults to a fixed - powers-of-2 + 256-stride list when `enable_piecewise_cuda_graph` - is True (and stays `None` otherwise). The fixed list keeps the - capture set small to bound startup time and CUDA graph memory; - the model-engine filter (invariants 2 and 3) clamps out-of-range - entries to the reachable ceiling and never invents sizes beyond - this list. + 1. `TorchLlmArgs.prefill_capture_num_tokens` defaults to a fixed + powers-of-2 + 256-stride list when a prefill CUDA graph backend is + enabled. The deprecated `TorchCompileConfig.capture_num_tokens` stays + `None` unless explicitly set. The fixed list keeps the capture set small + to bound startup time and CUDA graph memory; the model-engine filter + (invariants 2 and 3) clamps out-of-range entries to the reachable ceiling + and never invents sizes beyond this list. 2. `_filter_piecewise_capture_num_tokens` caps the candidate list at `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` -- the largest forward-pass `num_tokens` the warmup builder can @@ -1810,18 +1811,127 @@ class TestPiecewiseCudaGraphCaptureDefaults: _EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS = [2**i for i in range(8)] + list( range(256, 3073, 256)) - def test_torch_compile_config_capture_num_tokens_default_when_piecewise_enabled( - self): - """Default capture set is the powers-of-2 + 256-stride list. + def test_prefill_capture_num_tokens_uses_plain_int_list(self): + annotation = TorchLlmArgs.model_fields[ + "prefill_capture_num_tokens"].annotation + list_annotation = get_args(annotation)[0] + assert get_origin(list_annotation) is list + assert get_args(list_annotation) == (int, ) - Keeps the capture set bounded (~20 entries) so server startup - time and CUDA graph memory stay predictable. The model engine - further filters and appends the reachable ceiling, so - out-of-range entries (e.g. > max_seq_len-1) are never recorded - and gap ISLs still get a graph. - """ + def test_breakable_uses_default_capture_buckets(self): + args = TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE) + assert args.prefill_capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert args.torch_compile_config is None + + def test_piecewise_new_config_enables_default_torch_compile(self): + args = TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + prefill_capture_num_tokens=[512, 128, 512]) + assert args.torch_compile_config == TorchCompileConfig() + assert args.prefill_capture_num_tokens == [512, 128, 512] + + def test_legacy_piecewise_config_maps_to_new_fields(self): + args = TorchLlmArgs(model=llama_model_path, + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True, + capture_num_tokens=[128, 256])) + assert args.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE + assert args.prefill_capture_num_tokens == [256, 128] + + def test_explicit_new_buckets_with_legacy_piecewise_enable(self): + args = TorchLlmArgs(model=llama_model_path, + prefill_capture_num_tokens=[128, 256], + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True)) + assert args.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE + assert args.prefill_capture_num_tokens == [128, 256] + + def test_explicit_legacy_and_new_config_conflicts(self): + with pytest.raises(ValueError, match="conflicts"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True)) + + with pytest.raises(ValueError, match="conflicts"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + prefill_capture_num_tokens=[128], + torch_compile_config=TorchCompileConfig( + enable_piecewise_cuda_graph=True, capture_num_tokens=[256])) + + def test_breakable_rejects_explicit_torch_compile(self): + with pytest.raises(ValueError, match="does not support"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + torch_compile_config=TorchCompileConfig()) + + def test_prefill_filter_sorts_dedupes_and_drops_nonpositive(self): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + _filter_prefill_capture_num_tokens + + kept, unrecordable = _filter_prefill_capture_num_tokens( + [256, 0, -1, 128, 256], + max_num_tokens=512, + max_batch_size=1, + max_seq_len=513, + ) + assert kept == [128, 256] + assert unrecordable == [] + + @pytest.mark.parametrize("backend", [ + PrefillCudaGraphBackend.PIECEWISE, + PrefillCudaGraphBackend.BREAKABLE, + ]) + def test_piecewise_and_breakable_use_identical_padding(self, backend): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + PyTorchModelEngine + + engine = object.__new__(PyTorchModelEngine) + engine.enable_attention_dp = False + engine.prefill_cuda_graph_backend = backend + engine._prefill_cuda_graph_num_tokens = [128, 256, 512] + assert engine._get_padding_params(129, 1, None) == (256, True, None) + + def test_attention_dp_prefill_graph_uses_all_rank_decision(self): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + PyTorchModelEngine + + class FakeDist: + + def __init__(self, decisions): + self.decisions = decisions + + def tp_allgather(self, value): + del value + return self.decisions + + engine = object.__new__(PyTorchModelEngine) + engine.enable_attention_dp = True + engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.BREAKABLE + engine._prefill_cuda_graph_num_tokens = [128, 256, 512] + engine._get_all_rank_ctx_requests = lambda _: [0, 1, 0, 0] + + all_rank_num_tokens = [1, 129, 1, 1] + engine.dist = FakeDist([True, True, True, True]) + assert engine._get_padding_params(1, 0, + all_rank_num_tokens) == (256, True, + [256] * 4) + + engine.dist = FakeDist([True, False, True, True]) + assert engine._get_padding_params( + 1, 0, all_rank_num_tokens) == (1, False, all_rank_num_tokens) + + def test_torch_compile_config_does_not_populate_legacy_capture_buckets( + self): config = TorchCompileConfig(enable_piecewise_cuda_graph=True) - assert config.capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert config.capture_num_tokens is None def test_torch_compile_config_capture_num_tokens_stays_none_when_piecewise_disabled( self): @@ -1842,12 +1952,8 @@ def test_torch_compile_config_capture_num_tokens_user_override_preserved( # `validate_capture_num_tokens` dedupes and reverse-sorts. assert config.capture_num_tokens == sorted(set(user_list), reverse=True) - def test_torch_llm_args_capture_num_tokens_default_when_piecewise_enabled( + def test_torch_llm_args_prefill_buckets_default_when_piecewise_enabled( self): - """Same default applies when reached through `TorchLlmArgs` construction. - - This is the path real users hit via `trtllm-serve` YAML. - """ args = TorchLlmArgs( model=llama_model_path, max_batch_size=1, @@ -1859,7 +1965,8 @@ def test_torch_llm_args_capture_num_tokens_default_when_piecewise_enabled( torch_compile_config=TorchCompileConfig( enable_piecewise_cuda_graph=True), ) - assert args.torch_compile_config.capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert args.prefill_capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS + assert args.torch_compile_config.capture_num_tokens is None def test_piecewise_filter_never_invents_far_ceiling(self): """A ceiling far above the largest candidate is NOT added.