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..6d3ee0eceaf7 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,40 @@ 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 breakable backend is experimental. The integration coverage in this change +includes BF16 Qwen3.5 on one GPU and NVFP4 DeepSeek models on multiple GPUs, +with both context-only and mixed context/decode batches using the KV cache. + +The following restrictions are enforced: + +- `torch_compile_config`, LoRA, and multimodal models are rejected during + engine initialization. +- Speculative decoding is supported. +- Context-logit requests run eagerly instead of replaying a breakable CUDA + graph. + +Other model families, quantization modes, and parallel configurations are not +yet covered by this experimental backend's integration tests and should be +validated before use. + ## Tips for Piecewise CUDA Graph ### Piecewise CUDA Graph & Generation Only CUDA Graph @@ -59,9 +87,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/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 7d0ee04aff18..58625b057419 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -161,18 +161,65 @@ def create_sparse_attn_weights(self) -> None: # Fused epilogue buffer management and output projection. -def _validate_dsv4_epilogue_buffers( +def _create_dsv4_epilogue_buffers( self, + q: torch.Tensor, num_tokens: int, - dsv4_epilogue_output: tuple[torch.Tensor, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: - fp8_o, output_sf = dsv4_epilogue_output + if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: + raise ValueError( + "DSv4 fused epilogue requires num_heads_tp to be divisible by n_local_groups." + ) + heads_per_group = self.num_heads_tp // self.n_local_groups scale_buf_m = (num_tokens + 3) // 4 * 4 - if fp8_o.shape[1] != num_tokens or output_sf.shape[2] != scale_buf_m: - raise RuntimeError("Invalid DSv4 fused epilogue buffers for current token count.") + fp8_o = q.new_empty( + (self.n_local_groups, num_tokens, heads_per_group * self.v_head_dim), + dtype=torch.float8_e4m3fn, + ) + output_sf = q.new_empty( + ( + self.n_local_groups, + heads_per_group * (self.v_head_dim // 128), + scale_buf_m, + ), + dtype=torch.float32, + ) return fp8_o, output_sf +def _run_dsv4_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 prepare_sparse_attn_outputs( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata ) -> list[torch.Tensor]: @@ -183,9 +230,6 @@ def _should_use_dsv4_epilogue_fusion() -> bool: return False if num_contexts == 0 and num_generations == 0: return False - if num_contexts > 0 and num_generations > 0: - # The fused buffers do not carry token offsets for a mixed batch. - return False if self.mapping.has_cp_helix() or not is_sm_100f(): return False if not getattr(self.mapping, "enable_attention_dp", False): @@ -206,32 +250,15 @@ def _should_use_dsv4_epilogue_fusion() -> bool: return False return not self.inverse_rotary_emb.is_neox - def _create_dsv4_epilogue_buffers() -> tuple[torch.Tensor, torch.Tensor]: - if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: - raise ValueError( - "DSv4 fused epilogue requires num_heads_tp to be divisible by n_local_groups." - ) - heads_per_group = self.num_heads_tp // self.n_local_groups - num_tokens = attn_metadata.num_tokens - scale_buf_m = (num_tokens + 3) // 4 * 4 - fp8_o = hidden_states.new_empty( - (self.n_local_groups, num_tokens, heads_per_group * self.v_head_dim), - dtype=torch.float8_e4m3fn, - ) - output_sf = hidden_states.new_empty( - ( - self.n_local_groups, - heads_per_group * (self.v_head_dim // 128), - scale_buf_m, - ), - dtype=torch.float32, - ) - return fp8_o, output_sf - if _should_use_dsv4_epilogue_fusion(): - attn_output = [self.create_output(hidden_states[:0], attn_metadata.num_contexts)] - attn_output.extend(_create_dsv4_epilogue_buffers()) - return attn_output + num_tokens = hidden_states.shape[0] + return [ + torch.empty( + [num_tokens, self.n_local_groups, self.o_lora_rank], + device=hidden_states.device, + dtype=self.dtype, + ) + ] return [self.create_output(hidden_states, attn_metadata.num_contexts)] @@ -243,24 +270,12 @@ def project_sparse_attn_output( all_reduce_params: Optional["AllReduceParams"] = None, ) -> torch.Tensor: del attn_metadata, all_reduce_params - if len(attn_output) > 1: - attn_fp8, attn_scale = attn_output[1:] - num_tokens = attn_fp8.shape[1] - o_lora = torch.empty( - [num_tokens, self.n_local_groups, self.o_lora_rank], - device=attn_fp8.device, - dtype=self.dtype, - ) - torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( - attn_fp8, - self.o_a_proj, - attn_scale, - self.o_a_proj_scale, - o_lora.transpose(0, 1), - ) - return self.o_b_proj(o_lora.flatten(1)) - attn_output_tensor = attn_output[0] + # BCG/mixed-batch epilogue fusion runs o_a_proj at the end of attention, + # so this 3D tensor is O-LoRA output and only o_b_proj remains. + if attn_output_tensor.ndim == 3: + return self.o_b_proj(attn_output_tensor.flatten(1)) + assert position_ids is not None num_tokens = attn_output_tensor.shape[0] attn_output_tensor = attn_output_tensor.view(num_tokens, self.num_heads_tp, -1) @@ -347,11 +362,11 @@ def forward_generation_sparse_attn( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, - sparse_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + enable_dsv4_epilogue_fusion: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Run the DeepSeek-V4 generation absorption path.""" if get_sm_version() < 100: @@ -393,13 +408,11 @@ def forward_generation_sparse_attn( quant_q_buffer, ) - attention_output = output - output_sf = None + dsv4_output = output + o_lora_bmm_input_scale = None inverse_rope_cos_sin = None - if sparse_epilogue_output is not None: - attention_output, output_sf = _validate_dsv4_epilogue_buffers( - self, num_tokens, sparse_epilogue_output - ) + if enable_dsv4_epilogue_fusion: + dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers(self, q, num_tokens) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -411,8 +424,8 @@ def forward_generation_sparse_attn( attn_metadata, attention_input_type=AttentionInputType.generation_only, out_scale=self.out_scale, - output=attention_output, - output_sf=output_sf, + output=dsv4_output, + output_sf=o_lora_bmm_input_scale, latent_cache=latent_cache, q_pe=q_pe, sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), @@ -423,11 +436,13 @@ def forward_generation_sparse_attn( mla_bmm2_scale=mla_bmm2_scale, quant_q_buffer=quant_q_buffer, dsv4_inv_rope_cos_sin_cache=inverse_rope_cos_sin, - enable_dsv4_epilogue_fusion=sparse_epilogue_output is not None, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) - if sparse_epilogue_output is not None: - return attn_out_latent + if enable_dsv4_epilogue_fusion: + assert dsv4_output is not None and o_lora_bmm_input_scale is not None + return dsv4_output, o_lora_bmm_input_scale + assert output is not None if self.mapping.has_cp_helix(): raise RuntimeError( "DeepSeek-V4 + CP Helix is not supported because the post-process " @@ -445,11 +460,11 @@ def forward_context_sparse_attn( compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + output: Optional[torch.Tensor], latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, - sparse_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + enable_dsv4_epilogue_fusion: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Run the DeepSeek-V4 context absorption path.""" if get_sm_version() < 100: @@ -475,13 +490,11 @@ def forward_context_sparse_attn( quant_q_buffer = None quant_scale_qkv = None - attention_output = output - output_sf = None + dsv4_output = output + o_lora_bmm_input_scale = None inverse_rope_cos_sin = None - if sparse_epilogue_output is not None: - attention_output, output_sf = _validate_dsv4_epilogue_buffers( - self, num_tokens, sparse_epilogue_output - ) + if enable_dsv4_epilogue_fusion: + dsv4_output, o_lora_bmm_input_scale = _create_dsv4_epilogue_buffers(self, q, num_tokens) inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin attn_out_latent = self._attn_forward_gen( @@ -493,21 +506,24 @@ def forward_context_sparse_attn( attn_metadata, attention_input_type=AttentionInputType.context_only, out_scale=self.out_scale, - output=attention_output, - output_sf=output_sf, + output=dsv4_output, + output_sf=o_lora_bmm_input_scale, latent_cache=latent_cache, q_pe=q_pe, quant_q_buffer=quant_q_buffer, quant_scale_qkv=quant_scale_qkv, sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), dsv4_inv_rope_cos_sin_cache=inverse_rope_cos_sin, - enable_dsv4_epilogue_fusion=sparse_epilogue_output is not None, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) self._fused_quant_q_buffer = None self._fused_q_pe = None - if sparse_epilogue_output is not None: - return attn_out_latent + if enable_dsv4_epilogue_fusion: + assert dsv4_output is not None and o_lora_bmm_input_scale is not None + return dsv4_output, o_lora_bmm_input_scale + + assert output is not None if self.mapping.has_cp_helix(): raise RuntimeError( "DeepSeek-V4 + CP Helix is not supported because the post-process " @@ -532,14 +548,11 @@ def forward_sparse_attn( """Run DeepSeek-V4 MLA and write into the algorithm-defined output buffers.""" assert self.mha is None and self.mqa is not None, "DeepSeek-V4 is only supported in MQA mode" output = attn_output[0] - sparse_epilogue_output = (attn_output[1], attn_output[2]) if len(attn_output) > 1 else None + enable_dsv4_epilogue_fusion = output.ndim == 3 num_contexts = attn_metadata.num_contexts num_generations = attn_metadata.num_generations num_ctx_tokens = attn_metadata.num_ctx_tokens num_tokens = attn_metadata.num_tokens - if sparse_epilogue_output is not None and ((num_contexts > 0) == (num_generations > 0)): - raise RuntimeError("DSv4 epilogue fusion requires a context-only or generation-only batch.") - hidden_states = hidden_states[:num_tokens, ...] if position_ids is not None: position_ids = position_ids[..., :num_tokens] @@ -750,6 +763,8 @@ def _indexer_branch(): assert output is not None, "output must be provided" + context_o_lora_bmm_input = None + generation_o_lora_bmm_input = None if num_contexts > 0: q_ctx = q[:num_ctx_tokens, ...] topk_indices_ctx = topk_indices[:num_ctx_tokens, :] if topk_indices is not None else None @@ -761,17 +776,17 @@ def _indexer_branch(): assert ctx_position_ids is not None k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) - forward_context_sparse_attn( + context_o_lora_bmm_input = forward_context_sparse_attn( self, q_ctx, compressed_kv_ctx, k_pe_ctx, attn_metadata, - output[:num_ctx_tokens, :], + None if enable_dsv4_epilogue_fusion else output[:num_ctx_tokens, :], position_ids=ctx_position_ids, latent_cache=latent_cache_ctx, topk_indices=topk_indices_ctx, - sparse_epilogue_output=sparse_epilogue_output, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, ) if num_generations > 0: @@ -789,17 +804,31 @@ def _indexer_branch(): assert gen_position_ids is not None k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) - forward_generation_sparse_attn( + generation_o_lora_bmm_input = forward_generation_sparse_attn( self, q_gen, compressed_kv_gen, k_pe_gen, attn_metadata, - output[num_ctx_tokens:num_tokens, :], + None if enable_dsv4_epilogue_fusion else output[num_ctx_tokens:num_tokens, :], position_ids=gen_position_ids, latent_cache=latent_cache_gen, topk_indices=topk_indices_gen, - sparse_epilogue_output=sparse_epilogue_output, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, + ) + + if enable_dsv4_epilogue_fusion: + assert context_o_lora_bmm_input is None or isinstance(context_o_lora_bmm_input, tuple) + assert generation_o_lora_bmm_input is None or isinstance(generation_o_lora_bmm_input, tuple) + # The fused kernel output is group-first, which BCG cannot slice on + # dim 0. Write O-LoRA as token-first so replay can slice the bucket. + _run_dsv4_o_lora_bmms( + self, + output, + num_ctx_tokens, + num_tokens, + context_o_lora_bmm_input, + generation_o_lora_bmm_input, ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py index 67c912fba56a..5e4bc2efc4aa 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py @@ -7,6 +7,7 @@ import torch +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph from tensorrt_llm._torch.utils import Fp4QuantizedTensor from .module import _forward_dsa_attn, forward_dsa_proj @@ -146,3 +147,6 @@ def _mla_dsa_attn_inplace_fake( output: torch.Tensor, ) -> None: """Model the in-place output mutation during fake-tensor propagation.""" + + +maybe_bcg_mla_dsa_attn_inplace = eager_on_graph(mla_dsa_attn_inplace) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py index 5f07977a505d..5202067f14a6 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py @@ -198,7 +198,7 @@ def forward_sparse_attn_custom_op( ) q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] indexer_intermediates = proj_outputs[4:] - torch.ops.trtllm.mla_dsa_attn_inplace( + custom_ops.maybe_bcg_mla_dsa_attn_inplace( q, compressed_kv, k_pe, diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py index 73164f885660..53ee6d35edf7 100644 --- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py +++ b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py @@ -12,7 +12,7 @@ from tensorrt_llm.llmapi.utils import enable_llm_debug from ..utils import (get_model_extra_attrs, - get_per_request_piecewise_cuda_graph_flag, + get_per_request_prefill_cuda_graph_flag, get_piecewise_cuda_graph_flag, make_weak_ref, set_piecewise_running) from .multi_stream.auto_multi_stream import multi_stream_schedule @@ -202,7 +202,7 @@ def __call__(self, *args): if (runtime_num_of_token is None or runtime_num_of_token not in self.entries or not get_piecewise_cuda_graph_flag() - or not get_per_request_piecewise_cuda_graph_flag()): + or not get_per_request_prefill_cuda_graph_flag()): return self.default_callable(*args) if self.is_first_runner or self.is_last_runner: diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index e1978a577a3b..f9c2ed8dbe42 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -63,6 +63,7 @@ ) from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..utils import ( ActivationType, AuxStreamType, @@ -661,6 +662,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). @@ -1321,8 +1325,8 @@ def _forward_attention_core( output = q.new_empty( (q.shape[0], self.num_heads * self.head_dim), dtype=self.attn_activation_dtype ) - if self.register_to_config and is_torch_compiling(): - minimax_m3_attn_custom_op_inplace( + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): + maybe_bcg_minimax_m3_attn_custom_op_inplace( q, k, v, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 284eee0b0574..b6fea37271ba 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -23,6 +23,8 @@ cp_allgather, reducescatter) from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result +from ..pyexecutor.breakable_cuda_graph import (eager_on_graph, + is_in_breakable_cuda_graph) from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, is_torch_compiling) from .linear import (Linear, TensorParallelMode, WeightMode, @@ -117,6 +119,9 @@ def attn_custom_op_inplace( ) +maybe_bcg_attn_custom_op_inplace = eager_on_graph(attn_custom_op_inplace) + + def _helix_zero_kv_mask( attn_metadata: AttentionMetadata, num_tokens: int, @@ -963,20 +968,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 32aef7e61bcc..331bfe442976 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 6bdf977d131e..646509bb25da 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -44,6 +44,7 @@ from ..attention_backend.utils import create_attention from ..distributed import AllReduceParams from ..model_config import ModelConfig +from ..pyexecutor.breakable_cuda_graph import eager_on_graph from ..utils import ( AuxStreamType, Fp4QuantizedTensor, @@ -111,24 +112,27 @@ def _extract_mla_extra_attrs(layer_idx: str): return metadata, mla_layer -def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> list[torch.Tensor]: +def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) - return mla_layer._create_outputs(hidden_states, metadata) + outputs = mla_layer._create_outputs(hidden_states, metadata) + if len(outputs) != 1: + raise RuntimeError("MLA custom ops require exactly one output tensor.") + return outputs[0] @torch.library.custom_op("trtllm::create_mla_outputs", mutates_args=()) -def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> list[torch.Tensor]: +def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @create_mla_outputs.register_fake -def _create_mla_outputs_fake(hidden_states, layer_idx): +def _create_mla_outputs_fake(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @torch.library.custom_op( "trtllm::mla_custom_op_inplace", - mutates_args=("output", "sparse_output", "sparse_output_sf"), + mutates_args=("output",), ) def mla_custom_op_inplace( hidden_states: torch.Tensor, @@ -136,8 +140,6 @@ def mla_custom_op_inplace( layer_idx: str, output: torch.Tensor, latent_cache_gen: Optional[torch.Tensor], - sparse_output: Optional[torch.Tensor], - sparse_output_sf: Optional[torch.Tensor], hidden_states_fp4: Optional[torch.Tensor] = None, hidden_states_sf: Optional[torch.Tensor] = None, ) -> None: @@ -151,22 +153,18 @@ def mla_custom_op_inplace( scaling_factor=hidden_states_sf, unquantized_hidden_states=hidden_states, ) - attn_output = [output] - if sparse_output is not None: - attn_output.append(sparse_output) - if sparse_output_sf is not None: - if sparse_output is None: - raise RuntimeError("sparse_output_sf requires sparse_output") - attn_output.append(sparse_output_sf) mla_layer.forward_impl( position_ids, hidden_states, metadata, - attn_output=attn_output, + attn_output=[output], latent_cache_gen=latent_cache_gen, ) +maybe_bcg_mla_custom_op_inplace = eager_on_graph(mla_custom_op_inplace) + + def fp8_block_scaling_bmm_out( mat1: torch.Tensor, mat2_fp8: torch.Tensor, @@ -1723,36 +1721,26 @@ def _forward_custom_op( return output = attn_output[0] - sparse_output = None - sparse_output_sf = None - if len(attn_output) > 3: - raise RuntimeError("MLA output hooks may return at most two sparse output buffers.") - if len(attn_output) > 1: - sparse_output = attn_output[1] - if len(attn_output) > 2: - sparse_output_sf = attn_output[2] + if len(attn_output) != 1: + raise RuntimeError("MLA custom ops require exactly one output tensor.") if isinstance(hidden_states, Fp4QuantizedTensor): - torch.ops.trtllm.mla_custom_op_inplace( + maybe_bcg_mla_custom_op_inplace( hidden_states.unquantized_hidden_states, position_ids, self.layer_idx_str, output, latent_cache_gen, - sparse_output, - sparse_output_sf, 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, output, latent_cache_gen, - sparse_output, - sparse_output_sf, ) def _project_output( @@ -1816,9 +1804,9 @@ def forward( "unquantized_hidden_states view" ) output_hidden_states = hidden_states.unquantized_hidden_states - attn_output = torch.ops.trtllm.create_mla_outputs( - output_hidden_states, self.layer_idx_str - ) + attn_output = [ + torch.ops.trtllm.create_mla_outputs(output_hidden_states, self.layer_idx_str) + ] self._forward_custom_op( hidden_states, position_ids, diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__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..c70f02480927 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py @@ -0,0 +1,290 @@ +# 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): + logger.warning( + "Dropping a wait from the breakable CUDA graph capture stream " + "to a non-capturing stream; this dependency will not be " + "preserved during replay.", + stacklevel=2, + ) + return + _original_wait_stream(self, other) + forked.discard(other) + 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 = 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 f4406c7946d6..1e55429f06fa 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -35,6 +35,7 @@ create_input_processor_with_hash) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, + PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) from tensorrt_llm.logger import logger @@ -74,8 +75,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, @@ -216,6 +219,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, @@ -484,6 +491,7 @@ def __init__( setattr(self, "moe_load_balancer", moe_load_balancer) else: self.model = model + self._validate_breakable_cuda_graph_compatibility() pretrained_config = self.model.model_config.pretrained_config model_type = getattr(pretrained_config, "model_type", None) self._enable_dsv4_adp_dummy_fixes = should_enable_dsv4_adp_dummy_fixes( @@ -625,15 +633,14 @@ def __init__( and bool(self._cuda_graph_seq_lens)) self.torch_compile_config = self.llm_args.torch_compile_config + self.prefill_cuda_graph_backend = self.llm_args.prefill_cuda_graph_backend torch_compile_enabled = bool(self.torch_compile_config is not None) torch_compile_fullgraph = self.torch_compile_config.enable_fullgraph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_fullgraph'].default torch_compile_inductor_enabled = self.torch_compile_config.enable_inductor if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_inductor'].default - torch_compile_piecewise_cuda_graph = self.torch_compile_config.enable_piecewise_cuda_graph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ - 'enable_piecewise_cuda_graph'].default - torch_compile_piecewise_cuda_graph_num_tokens = self.torch_compile_config.capture_num_tokens if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ - 'capture_num_tokens'].default + torch_compile_piecewise_cuda_graph = (self.prefill_cuda_graph_backend == + PrefillCudaGraphBackend.PIECEWISE) torch_compile_enable_userbuffers = self.torch_compile_config.enable_userbuffers if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_userbuffers'].default torch_compile_max_num_streams = self.torch_compile_config.max_num_streams if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ @@ -642,14 +649,14 @@ def __init__( self._torch_compile_enabled = torch_compile_enabled self._torch_compile_piecewise_cuda_graph = torch_compile_piecewise_cuda_graph - piecewise_cuda_graph_num_tokens = ( - torch_compile_piecewise_cuda_graph_num_tokens - or cuda_graph_batch_sizes or []) + prefill_cuda_graph_num_tokens = self.llm_args.prefill_capture_num_tokens + if prefill_cuda_graph_num_tokens is None: + prefill_cuda_graph_num_tokens = cuda_graph_batch_sizes or [] num_extra_decoding_steps = self._get_num_extra_decoding_steps() - self._piecewise_cuda_graph_num_tokens, unrecordable = ( - _filter_piecewise_capture_num_tokens( - piecewise_cuda_graph_num_tokens, + self._prefill_cuda_graph_num_tokens, unrecordable = ( + _filter_prefill_capture_num_tokens( + prefill_cuda_graph_num_tokens, max_num_tokens=self.max_num_tokens, max_batch_size=self.batch_size, max_seq_len=self.max_seq_len, @@ -657,7 +664,7 @@ def __init__( )) if unrecordable: logger.warning( - f"Skipping piecewise CUDA graph capture for num_tokens=" + f"Skipping prefill CUDA graph capture for num_tokens=" f"{unrecordable}: exceeds reachable ceiling " f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)=" f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. " @@ -682,7 +689,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, @@ -926,6 +933,17 @@ def __init__( enable_encoder_decoder_mixed_cuda_graph), ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) + self.breakable_cuda_graph_runner = None + if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: + 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) # Initialize CUDA Graph LoRA manager if LoRA is enabled self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None @@ -2046,7 +2064,8 @@ def _get_graphs_to_capture( def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): """Warm up or capture CUDA graphs for the configured graph shapes.""" if not (self.cuda_graph_runner.enabled - or self._torch_compile_piecewise_cuda_graph): + or self.prefill_cuda_graph_backend + != PrefillCudaGraphBackend.DISABLED): return self._capture_generation_cuda_graphs(resource_manager) @@ -2054,7 +2073,7 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. if not self.cuda_graph_runner.is_warmup_only: - self._capture_piecewise_cuda_graphs(resource_manager) + self._capture_prefill_cuda_graphs(resource_manager) @torch.inference_mode() @with_warmup_flag @@ -2450,60 +2469,80 @@ def _capture_mixed_encoder_decoder_cuda_graphs( self.enable_spec_decode = saved_enable_spec_decode self.runtime_draft_len = saved_runtime_draft_len - def _capture_piecewise_cuda_graphs(self, resource_manager: ResourceManager): - """Captures piecewise CUDA graphs for context/prefill steps via torch.compile.""" - if not (self._torch_compile_piecewise_cuda_graph - and self._torch_compile_enabled): + def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): + """Capture configured CUDA graphs for context/prefill steps.""" + if (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.DISABLED + or (self.prefill_cuda_graph_backend + == PrefillCudaGraphBackend.PIECEWISE + and not self._torch_compile_enabled)): return - logger.info("Running piecewise CUDA graph warmup...") - piecewise_cuda_graph_num_tokens = sorted( - self._piecewise_cuda_graph_num_tokens, reverse=True) + logger.info("Running prefill CUDA graph warmup...") + prefill_cuda_graph_num_tokens = sorted( + self._prefill_cuda_graph_num_tokens, reverse=True) - with capture_piecewise_cuda_graph(True), self.no_cuda_graph(): - for num_tokens in piecewise_cuda_graph_num_tokens: + capture_context = (capture_piecewise_cuda_graph(True) + if self._torch_compile_piecewise_cuda_graph else + contextlib.nullcontext()) + with capture_context, self.no_cuda_graph(): + for num_tokens in prefill_cuda_graph_num_tokens: warmup_request = self._create_warmup_request( resource_manager, num_tokens, 0) with self._release_batch_context(warmup_request, 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 ### @@ -3090,6 +3129,23 @@ def is_multimodal(self) -> bool: return True return isinstance(self.input_processor, BaseMultimodalInputProcessor) + def _validate_breakable_cuda_graph_compatibility(self) -> None: + if self.llm_args.prefill_cuda_graph_backend != PrefillCudaGraphBackend.BREAKABLE: + return + + if isinstance(self.model, DecoderModelForCausalLM): + return + decoder_model = getattr(self.model, "llm", None) + if (self.llm_args.disable_mm_encoder + and isinstance(decoder_model, DecoderModelForCausalLM) + and getattr(self.model, "mm_encoder", None) is None): + return + if (isinstance(self.model, MultimodalModelMixin) or isinstance( + self.input_processor, BaseMultimodalInputProcessor)): + raise ValueError( + "breakable prefill CUDA graph does not support multimodal models" + ) + def _set_up_multimodal_encoder_attn_metadata(self) -> None: """Construct AttentionMetadata for any multimodal encoders inside the loaded model, using the engine's encoder runtime sizes @@ -3336,6 +3392,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() @@ -3519,45 +3578,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" @@ -3567,19 +3621,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 @@ -4016,7 +4069,7 @@ def _prepare_encoder_decoder_inputs_fast( attn_all_rank_num_tokens) = self._get_padding_params( total_num_tokens, scheduled_requests.num_context_requests, attn_all_rank_num_tokens) - set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + set_per_request_prefill_cuda_graph_flag(can_run_piecewise_cuda_graph) attn_metadata.padded_num_tokens = (padded_num_tokens if padded_num_tokens != total_num_tokens else None) @@ -4650,7 +4703,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 @@ -5782,9 +5835,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 @@ -6057,9 +6111,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: @@ -6544,6 +6598,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: @@ -7043,7 +7098,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, @@ -7167,14 +7221,34 @@ 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 not gather_context_logits + and breakable_runner.has_graph(num_tokens)) + if can_run_breakable_graph and not breakable_runner.is_warming_up: + outputs = breakable_runner.execute( + 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 cb62ec99f76b..2697b60f4077 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -170,8 +170,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") @@ -397,12 +397,14 @@ def piecewise_cuda_graph(enable: bool): set_piecewise_cuda_graph_flag(prev_enable) -def set_per_request_piecewise_cuda_graph_flag(enable: bool): - _global_attrs.per_request_piecewise_cuda_graph_flag = enable +def set_per_request_prefill_cuda_graph_flag(enable: bool): + """Set whether the current batch can use its prefill CUDA graph backend.""" + _global_attrs.per_request_prefill_cuda_graph_flag = enable -def get_per_request_piecewise_cuda_graph_flag() -> bool: - return getattr(_global_attrs, 'per_request_piecewise_cuda_graph_flag', True) +def get_per_request_prefill_cuda_graph_flag() -> bool: + """Return whether the current batch can use its prefill CUDA graph backend.""" + return getattr(_global_attrs, 'per_request_prefill_cuda_graph_flag', True) def create_lm_head_tp_mapping(mapping: Mapping, token_count: int) -> Mapping: diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 2ddc301eb01a..aff2e9736977 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -20,12 +20,12 @@ MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, NGramDecodingConfig, PARDDecodingConfig, - PrometheusMetricsConfig, ReorderRequestPolicyConfig, - RocketSparseAttentionConfig, SADecodingConfig, - SAEnhancerConfig, SaveHiddenStatesDecodingConfig, - SchedulerConfig, SkipSoftmaxAttentionConfig, - TorchCompileConfig, TorchLlmArgs, - TriAttentionKvCacheCompressionConfig, + PrefillCudaGraphBackend, PrometheusMetricsConfig, + ReorderRequestPolicyConfig, RocketSparseAttentionConfig, + SADecodingConfig, SAEnhancerConfig, + SaveHiddenStatesDecodingConfig, SchedulerConfig, + SkipSoftmaxAttentionConfig, TorchCompileConfig, + TorchLlmArgs, TriAttentionKvCacheCompressionConfig, UserProvidedDecodingConfig) from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder @@ -93,6 +93,7 @@ 'SkipSoftmaxAttentionConfig', 'TriAttentionKvCacheCompressionConfig', 'PrometheusMetricsConfig', + 'PrefillCudaGraphBackend', 'ThinkingBudgetLogitsProcessor', 'add_thinking_budget_logits_processor', 'MultimodalConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1afb2cc25994..131d1302c30b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5009,6 +5009,18 @@ class SamplerType(StrEnum): auto = "auto" +class PrefillCudaGraphBackend(StrEnum): + """CUDA graph implementation used for prefill requests.""" + + DISABLED = "disabled" + PIECEWISE = "piecewise" + BREAKABLE = "breakable" + + +_DEFAULT_PREFILL_CAPTURE_NUM_TOKENS = [2**i for i in range(8) + ] + [i for i in range(256, 3073, 256)] + + class TorchCompileConfig(StrictBaseModel): """Configuration for torch.compile.""" enable_fullgraph: bool = Field( @@ -5020,13 +5032,12 @@ class TorchCompileConfig(StrictBaseModel): enable_piecewise_cuda_graph: bool = Field( default=False, - description="Enable piecewise CUDA graph in torch.compile.") + description="Deprecated. Use prefill_cuda_graph_backend='piecewise' " + "instead.") capture_num_tokens: Optional[List[PositiveInt]] = Field( default=None, - description= - "List of num of tokens to capture the piecewise CUDA graph for. If not provided, the number of tokens will be the same as cuda_graph_config.batch_sizes." - ) + description="Deprecated. Use prefill_capture_num_tokens instead.") @field_validator('capture_num_tokens') @classmethod @@ -5041,17 +5052,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 @@ -5297,6 +5301,20 @@ def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': torch_compile_config: Optional[TorchCompileConfig] = Field( default=None, description="Torch compile config.", status="prototype") + prefill_cuda_graph_backend: PrefillCudaGraphBackend = Field( + default=PrefillCudaGraphBackend.DISABLED, + description="CUDA graph implementation used for prefill requests. " + "Defaults to disabled.", + status="prototype", + telemetry=TelemetryField.categorical("disabled", "piecewise", + "breakable")) + + prefill_capture_num_tokens: Optional[List[int]] = Field( + default=None, + description= + "Token-count buckets captured by the selected prefill CUDA graph implementation.", + status="prototype") + enable_autotuner: bool = Field( default=True, description= @@ -5593,6 +5611,65 @@ 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: + if self.torch_compile_config is not None: + raise ValueError( + "breakable prefill CUDA graph does not support torch_compile_config" + ) + if self.enable_lora or self.lora_config is not None: + raise ValueError( + "breakable prefill CUDA graph does not support LoRA") + + return self + @model_validator(mode="after") def validate_speculative_config(self): if self.speculative_config: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 38436f22514a..e16d33f7ea45 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1286,6 +1286,24 @@ "kind": "value", "path": "pp_partition" }, + { + "allowed_values": [], + "annotation": "Optional[List[int]]", + "converter": "", + "kind": "value", + "path": "prefill_capture_num_tokens" + }, + { + "allowed_values": [ + "disabled", + "piecewise", + "breakable" + ], + "annotation": "", + "converter": "allowlist", + "kind": "categorical", + "path": "prefill_cuda_graph_backend" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 11cca4e46f76..4578e6ad42b4 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 @@ -3596,6 +3597,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( @@ -3933,6 +3994,86 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): eplb_config, mtp_nextn=mtp_nextn) + @pytest.mark.skip_less_mpi_world_size(8) + @pytest.mark.threadleak(enabled=False) + def test_mixed_breakable_cuda_graph(self): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + base_prompt_ids = tokenizer.encode( + "TensorRT-LLM accelerates reliable large language model inference " + "with efficient attention, parallelism, and CUDA graphs. ", + add_special_tokens=False, + ) + assert base_prompt_ids + + def make_prompt(prompt_length): + return (base_prompt_ids * + ((prompt_length + len(base_prompt_ids) - 1) // + len(base_prompt_ids)))[:prompt_length] + + generation_prompt = make_prompt(64) + context_prompt = make_prompt(129) + sampling_params = SamplingParams( + max_tokens=8, + min_tokens=8, + seed=42, + temperature=0, + ignore_eos=True, + detokenize=False, + add_special_tokens=False, + ) + common_llm_kwargs = dict( + tensor_parallel_size=8, + moe_expert_parallel_size=8, + moe_config=MoeConfig(backend="TRTLLM"), + enable_attention_dp=True, + max_batch_size=8, + max_num_tokens=1024, + max_seq_len=2048, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + dtype="fp8", + free_gpu_memory_fraction=0.6, + ), + cuda_graph_config=CudaGraphConfig( + batch_sizes=[1, 2, 4, 6, 8], + enable_padding=True, + ), + ) + + def run(backend): + with LLM( + self.MODEL_PATH, + **common_llm_kwargs, + prefill_cuda_graph_backend=backend, + prefill_capture_num_tokens=[128, 256, 512, 1024], + ) as llm: + generation_request = llm.generate_async( + generation_prompt, + sampling_params=sampling_params, + streaming=True, + ) + next(generation_request) + assert not generation_request.finished + + # Admit a context request while the first request is decoding. + context_request = llm.generate_async( + context_prompt, + sampling_params=sampling_params, + streaming=False, + ) + generation_output = generation_request.result() + context_output = context_request.result() + return [ + generation_output.outputs[0].token_ids, + context_output.outputs[0].token_ids, + ] + + eager_token_ids = run(PrefillCudaGraphBackend.DISABLED) + breakable_token_ids = run(PrefillCudaGraphBackend.BREAKABLE) + assert breakable_token_ids == eager_token_ids + _DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT = ( "Solve the problem carefully. End your response with a final line exactly " @@ -6201,6 +6342,47 @@ 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, + disable_mm_encoder=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 ba99e5284330..634f8981f776 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -57,6 +57,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a16_mxfp4[latency-TRTLLM] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention[target_sparsity_0.9-fp8kv=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=True] + - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16_breakable_prefill_cuda_graph - accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[TRTLLM] - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 74d7936cab15..089c9cb70d04 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 196c06531831..887e5f66555c 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -276,9 +276,9 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): attn_out_latent = torch.randn(num_tokens, num_heads, qk_head_dim, dtype=dtype, device=device) position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) - # Call the deepseek_v4 output projection (mla_rope_inplace modifies attn_out_latent - # in-place, so clone before passing to preserve original for reference) - output = project_sparse_attn_output(mla, [attn_out_latent.clone()], position_ids) + # The non-fused MLA path stores attention output as a flattened 2D buffer. + # mla_rope_inplace modifies it in place, so preserve the 3D reference input. + output = project_sparse_attn_output(mla, [attn_out_latent.clone().flatten(1)], position_ids) # Calculate reference output if dtype_str == "bf16": diff --git a/tests/unittest/_torch/compilation/test_remove_copy_pass.py b/tests/unittest/_torch/compilation/test_remove_copy_pass.py index bf2794e65230..974906f485a2 100644 --- a/tests/unittest/_torch/compilation/test_remove_copy_pass.py +++ b/tests/unittest/_torch/compilation/test_remove_copy_pass.py @@ -138,7 +138,7 @@ def test_remove_copy_for_mutates_tensor_list( graph.lint() -def test_remove_copy_for_mutates_args_restores_optional_none() -> None: +def test_remove_copy_for_mla_restores_final_output_mutation() -> None: graph = Graph() hidden_states = graph.placeholder("hidden_states") output = graph.placeholder("output") @@ -153,8 +153,6 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: "latent_cache_gen": None, "_all_bases": (output,), "_output_base_index": 0, - "_sparse_output_base_index": None, - "_sparse_output_sf_base_index": None, }, ) mutated_output = graph.call_function(getitem, args=(functionalized, 1)) @@ -166,48 +164,5 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: inplace_nodes = [node for node in graph.nodes if node.target == inplace_func] assert len(inplace_nodes) == 1 assert inplace_nodes[0].kwargs["output"] is output - assert inplace_nodes[0].kwargs["sparse_output"] is None - assert inplace_nodes[0].kwargs["sparse_output_sf"] is None assert clone.args[0] is output graph.lint() - - -def test_remove_copy_for_mutates_args_rejects_getitem_for_optional_none( - monkeypatch: pytest.MonkeyPatch, -) -> None: - graph = Graph() - hidden_states = graph.placeholder("hidden_states") - output = graph.placeholder("output") - inplace_func = torch.ops.trtllm.mla_custom_op_inplace.default - functionalized = graph.call_function( - auto_functionalized_v2, - args=(inplace_func,), - kwargs={ - "hidden_states": hidden_states, - "position_ids": None, - "layer_idx": "0", - "latent_cache_gen": None, - "_all_bases": (output,), - "_output_base_index": 0, - "_sparse_output_base_index": None, - "_sparse_output_sf_base_index": None, - }, - ) - optional_output = graph.call_function(getitem, args=(functionalized, 2)) - clone = graph.call_function(torch.ops.aten.clone.default, args=(optional_output,)) - graph.output(clone) - - monkeypatch.setattr( - remove_copy_pass, - "inplace_info", - lambda: {inplace_func: {1: "output", 2: "sparse_output"}}, - ) - - with pytest.raises( - AssertionError, - match=( - "getitem user for optional output 'sparse_output' has no " - "base tensor -- graph is malformed" - ), - ): - remove_copy_pass.remove_copy_for_mutates_args(graph) diff --git a/tests/unittest/_torch/executor/test_breakable_cuda_graph.py b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py new file mode 100644 index 000000000000..7c38f9e328ad --- /dev/null +++ b/tests/unittest/_torch/executor/test_breakable_cuda_graph.py @@ -0,0 +1,343 @@ +# 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 logging +import weakref + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, + break_graph, + eager_on_graph, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph import _copy_output +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_eager_output_storage_survives_allocator_churn(): + eager_output_ptrs = [] + eager_output_refs = [] + + @eager_on_graph + def add_one(value): + output = value + 1 + eager_output_ptrs.append(output.data_ptr()) + eager_output_refs.append(weakref.ref(output)) + return output + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + x = torch.zeros(1024, device="cuda") + output = torch.zeros_like(x) + graph = BreakableCUDAGraph() + + with BreakableCUDAGraphCapture(graph, stream=stream): + output.copy_(add_one(x) * 2) + + captured_output_ptr = eager_output_ptrs[0] + assert eager_output_refs[0]() is not None + churn = [torch.full_like(x, value) for value in range(4)] + assert all(tensor.data_ptr() != captured_output_ptr for tensor in churn) + + x.fill_(5) + graph.replay() + + stream.synchronize() + torch.testing.assert_close(output, torch.full_like(output, 12)) + for value, tensor in enumerate(churn): + torch.testing.assert_close(tensor, torch.full_like(tensor, value)) + + +def test_outside_capture(): + @eager_on_graph + def outside(value): + 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)) + + +def test_dropped_wait_stream_logs_callsite(caplog): + side_stream = torch.cuda.Stream() + + def body(): + torch.cuda.current_stream().wait_stream(side_stream) + + with caplog.at_level( + logging.WARNING, + logger=("tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph"), + ): + _capture(body) + + record = next(record for record in caplog.records if "Dropping a wait" in record.message) + assert record.funcName == "body" + + +class _Body(nn.Module): + def __init__(self): + super().__init__() + 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 ed0f3bd4d00b..bb83bb0ee231 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -15,6 +15,7 @@ MultimodalEncoderMixin from tensorrt_llm._torch.models.modeling_multimodal_mixin import \ MultimodalModelMixin +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( @@ -25,6 +26,7 @@ PyTorchModelEngine, _build_request_multimodal_input, _make_single_token_context_graph_batch) from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchLlmArgs) @@ -41,7 +43,8 @@ from tensorrt_llm._torch.speculative.spec_sampler_base import \ SampleStateTensorsSpec from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.inputs.registry import BaseMultimodalDummyInputsBuilder +from tensorrt_llm.inputs.registry import (BaseMultimodalDummyInputsBuilder, + BaseMultimodalInputProcessor) from tensorrt_llm.llmapi import (CudaGraphConfig, SADecodingConfig, SamplingParams) from tensorrt_llm.mapping import CpType, Mapping @@ -238,10 +241,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 @@ -817,7 +825,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) @@ -891,7 +900,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: @@ -972,7 +982,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) @@ -1002,6 +1013,35 @@ def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: self.assertIs(context.py_multimodal_data, multimodal_data) self.assertIn("multimodal_embedding", multimodal_data) + def test_breakable_graph_falls_back_for_context_logits(self) -> None: + engine, _, resource_manager, _, outputs = \ + _make_forward_only_engine(None) + breakable_runner = Mock() + breakable_runner.is_capturing = False + breakable_runner.is_warming_up = False + breakable_runner.has_graph.return_value = True + breakable_runner.execute.return_value = outputs + engine.breakable_cuda_graph_runner = breakable_runner + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [_make_request_stub(1)] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock() + ), patch( + "tensorrt_llm._torch.pyexecutor.model_engine.get_per_request_prefill_cuda_graph_flag", + return_value=True): + actual_outputs = engine.forward( + batch, + resource_manager, + gather_context_logits=True, + ) + + self.assertIs(actual_outputs, outputs) + breakable_runner.execute.assert_not_called() + engine._forward_step.assert_called_once() + def test_generation_only_forward_does_not_call_new_selector(self) -> None: key = KeyType(batch_size=1, draft_len=0, is_first_draft=False) engine, runner, resource_manager, _, _ = _make_forward_only_engine(key) @@ -1139,6 +1179,44 @@ def test_encoder_cuda_graph_stages_and_restores_fixed_sequence_slots( (fixed_slot_output[511:512], fixed_slot_output[:400])) torch.testing.assert_close(restored_output, expected_output) + def test_breakable_rejects_multimodal_models(self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = DummyLegacyMultimodalIndexModel() + engine.input_processor = None + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + disable_mm_encoder=False) + + self.assertTrue(engine.is_multimodal) + with self.assertRaisesRegex(ValueError, "multimodal models"): + engine._validate_breakable_cuda_graph_compatibility() + + def test_breakable_allows_text_decoder_with_multimodal_processor( + self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = Mock(spec=DecoderModelForCausalLM) + engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + disable_mm_encoder=False) + + self.assertTrue(engine.is_multimodal) + engine._validate_breakable_cuda_graph_compatibility() + + def test_breakable_allows_multimodal_wrapper_in_text_only_mode( + self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = DummyLegacyMultimodalIndexModel() + engine.model.llm = Mock(spec=DecoderModelForCausalLM) + engine.model.mm_encoder = None + engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + disable_mm_encoder=True) + + self.assertTrue(engine.is_multimodal) + engine._validate_breakable_cuda_graph_compatibility() + def test_prepare_multimodal_indices_uses_mixin_token_ids(self) -> None: engine = object.__new__(PyTorchModelEngine) engine.model = DummyMultimodalIndexModel() diff --git a/tests/unittest/_torch/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py index 908069a02560..4e40bf419666 100644 --- a/tests/unittest/_torch/modules/test_mla_registry.py +++ b/tests/unittest/_torch/modules/test_mla_registry.py @@ -13,12 +13,20 @@ # 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.attention_backend.sparse.deepseek_v4.module import ( + _create_dsv4_epilogue_buffers, + _run_dsv4_o_lora_bmms, + prepare_sparse_attn_outputs, + project_sparse_attn_output, +) from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.mla import MLA from tensorrt_llm.functional import PositionEmbeddingType @@ -80,3 +88,167 @@ 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 _make_dsv4_epilogue_layer() -> SimpleNamespace: + return SimpleNamespace( + _disable_dsv4_epilogue_fusion=False, + mapping=SimpleNamespace( + has_cp_helix=lambda: False, + enable_attention_dp=True, + ), + num_heads=128, + num_heads_tp=128, + mqa=SimpleNamespace( + sparse_params=object(), + has_fp8_kv_cache=True, + ), + o_a_proj=SimpleNamespace(dtype=torch.float8_e4m3fn), + kv_lora_rank=448, + qk_rope_head_dim=64, + qk_head_dim=512, + v_head_dim=512, + n_local_groups=8, + o_lora_rank=3, + dtype=torch.bfloat16, + inverse_rotary_emb=SimpleNamespace(is_neox=False), + create_output=Mock(), + ) + + +def test_mla_custom_op_marks_only_final_output_mutable() -> None: + schema = torch.ops.trtllm.mla_custom_op_inplace.default._schema + mutated_args = [ + arg.name + for arg in schema.arguments + if arg.alias_info is not None and arg.alias_info.is_write + ] + assert mutated_args == ["output"] + + +def test_create_mla_outputs_custom_op_returns_tensor() -> None: + schema = torch.ops.trtllm.create_mla_outputs.default._schema + assert [str(return_value.type) for return_value in schema.returns] == ["Tensor"] + + +def test_dsv4_epilogue_fusion_supports_mixed_batch() -> None: + mla_layer = _make_dsv4_epilogue_layer() + metadata = SimpleNamespace(num_contexts=1, num_generations=1) + hidden_states = torch.empty(8, 16) + + with patch( + "tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module.is_sm_100f", + return_value=True, + ): + outputs = prepare_sparse_attn_outputs(mla_layer, hidden_states, metadata) + + assert len(outputs) == 1 + assert outputs[0].shape == (8, 8, 3) + assert outputs[0].dtype == torch.bfloat16 + mla_layer.create_output.assert_not_called() + + +def test_dsv4_fusion_create_output_uses_bucket_token_count() -> None: + mla_layer = _make_dsv4_epilogue_layer() + metadata = SimpleNamespace(num_contexts=1, num_generations=0) + hidden_states = torch.empty(8, 16) + + with patch( + "tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module.is_sm_100f", + return_value=True, + ): + output = prepare_sparse_attn_outputs(mla_layer, hidden_states, metadata)[0] + + assert output.shape == (8, 8, 3) + assert output.dtype == torch.bfloat16 + + +def test_dsv4_fusion_o_proj_only_flattens_lora_output() -> None: + projected = torch.randn(7, 5) + mla_layer = SimpleNamespace( + n_local_groups=4, + o_lora_rank=3, + o_b_proj=Mock(return_value=projected), + ) + lora_o = torch.randn(7, 4, 3) + + output = project_sparse_attn_output(mla_layer, [lora_o]) + + assert output is projected + mla_layer.o_b_proj.assert_called_once() + torch.testing.assert_close(mla_layer.o_b_proj.call_args.args[0], lora_o.flatten(1)) + + +def test_dsv4_epilogue_buffers_use_real_token_count() -> None: + mla_layer = SimpleNamespace( + n_local_groups=4, + num_heads_tp=128, + v_head_dim=512, + ) + q = torch.empty(8, 16) + + fp8_o, output_sf = _create_dsv4_epilogue_buffers(mla_layer, q, num_tokens=5) + + assert fp8_o.shape == (4, 5, 32 * 512) + assert output_sf.shape == (4, 32 * 4, 8) + + +@pytest.mark.parametrize( + "num_context_tokens,num_generation_tokens,bucket_tokens", + [(5, 0, 8), (0, 3, 4), (5, 3, 12)], +) +def test_dsv4_epilogue_bmm_writes_only_phase_ranges( + num_context_tokens: int, + num_generation_tokens: int, + bucket_tokens: int, +) -> None: + groups = 2 + rank = 3 + output = torch.full((bucket_tokens, groups, rank), -1.0) + mla_layer = SimpleNamespace( + o_a_proj=torch.empty(0), + o_a_proj_scale=torch.empty(0), + ) + + def fake_bmm(_attn_fp8, _weight, attn_scale, _weight_scale, phase_output): + phase_output.fill_(attn_scale.item()) + + with patch.object( + torch.ops.trtllm, + "cute_dsl_fp8_bmm_blackwell", + side_effect=fake_bmm, + ) as bmm: + context_epilogue = None + if num_context_tokens: + context_epilogue = ( + torch.empty(groups, num_context_tokens, 4), + torch.tensor(11.0), + ) + generation_epilogue = None + if num_generation_tokens: + generation_epilogue = ( + torch.empty(groups, num_generation_tokens, 4), + torch.tensor(22.0), + ) + _run_dsv4_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 409eafe66895..c10d9123c4d7 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -231,6 +231,14 @@ methods: annotation: Optional[tensorrt_llm.llmapi.llm_args.TorchCompileConfig] default: null status: prototype + prefill_cuda_graph_backend: + annotation: tensorrt_llm.llmapi.llm_args.PrefillCudaGraphBackend + default: disabled + status: prototype + prefill_capture_num_tokens: + annotation: Optional[List[int]] + default: null + status: prototype enable_autotuner: annotation: bool default: True diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2fa8772fd073..0a2aeb12c5f2 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -48,7 +48,8 @@ MambaStateConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, - PeftCacheConfig, PybindMirror, + PeftCacheConfig, + PrefillCudaGraphBackend, PybindMirror, RayPlacementConfig, SkipSoftmaxAttentionConfig, SleepConfig, SpeculativeConfig, @@ -1935,13 +1936,13 @@ class TestPiecewiseCudaGraphCaptureDefaults: Three invariants are exercised: - 1. `TorchCompileConfig.capture_num_tokens` defaults to a fixed - powers-of-2 + 256-stride list when `enable_piecewise_cuda_graph` - is True (and stays `None` otherwise). The fixed list keeps the - capture set small to bound startup time and CUDA graph memory; - the model-engine filter (invariants 2 and 3) clamps out-of-range - entries to the reachable ceiling and never invents sizes beyond - this list. + 1. `TorchLlmArgs.prefill_capture_num_tokens` defaults to a fixed + powers-of-2 + 256-stride list when a prefill CUDA graph backend is + enabled. The deprecated `TorchCompileConfig.capture_num_tokens` stays + `None` unless explicitly set. The fixed list keeps the capture set small + to bound startup time and CUDA graph memory; the model-engine filter + (invariants 2 and 3) clamps out-of-range entries to the reachable ceiling + and never invents sizes beyond this list. 2. `_filter_piecewise_capture_num_tokens` caps the candidate list at `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` -- the largest forward-pass `num_tokens` the warmup builder can @@ -1957,18 +1958,151 @@ 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_breakable_allows_speculative_decoding(self): + speculative_config = MTPDecodingConfig(max_draft_len=1) + args = TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + speculative_config=speculative_config) + assert args.speculative_config == speculative_config + + @pytest.mark.parametrize("lora_kwargs", [ + { + "enable_lora": True + }, + { + "lora_config": LoraConfig(lora_target_modules=["attn_q"]) + }, + ]) + def test_breakable_rejects_lora(self, lora_kwargs): + with pytest.raises(ValueError, match="LoRA"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + **lora_kwargs, + ) + + def test_prefill_filter_sorts_dedupes_and_drops_nonpositive(self): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + _filter_prefill_capture_num_tokens + + 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): @@ -1989,12 +2123,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, @@ -2006,7 +2136,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.