From ba29216d3674cb9b7028d67e7a25b8bf7b8612c9 Mon Sep 17 00:00:00 2001 From: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:34:51 -0700 Subject: [PATCH 1/4] [None][fix] Make Mamba2 mixer tolerate prefill CUDA graph token padding Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> (cherry picked from commit e9b4ad793665d0d226e2214749f1582aacf9c411) --- .../_torch/modules/mamba/mamba2_mixer.py | 26 ++++++++++++------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index d88c0ed23bc4..411cc28640f9 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -318,7 +318,6 @@ def forward( num_prefill_tokens = attn_metadata.num_ctx_tokens num_decode_tokens = attn_metadata.num_tokens - num_prefill_tokens num_actual_tokens = attn_metadata.num_tokens - seqlen_split_size = [num_prefill_tokens, num_decode_tokens] batch_split_size = [num_prefills, num_decodes] state_indices = mamba_metadata.state_indices[:num_prefills + @@ -339,7 +338,12 @@ def forward( # Split z and dt with views. z = zxbcdt[:, :self.tp_d_inner] dt = zxbcdt[:, self.tp_d_inner + self.tp_conv_dim:] - dt_p, dt_d = torch.split(dt, seqlen_split_size, dim=0) + # Slice instead of exact-sum split: under piecewise CUDA graphs the + # token dim is padded to the capture bucket, so hidden_states can + # carry more rows than num_actual_tokens; the pad tail belongs to + # neither the prefill nor the decode segment. + dt_p = dt[:num_prefill_tokens] + dt_d = dt[num_prefill_tokens:num_actual_tokens] # Decode path uses regular view since no transpose is needed. xbc_d = zxbcdt[num_prefill_tokens:num_actual_tokens, @@ -352,11 +356,12 @@ def forward( dtype=zxbcdt.dtype, device=zxbcdt.device, ) - preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split( - preallocated_ssm_out, - [num_prefill_tokens, num_decode_tokens], - dim=0, - ) + # Zero the pad tail (torch.empty) so the full-length gated norm below + # sees defined values in pad rows; no-op slice when not padded. + preallocated_ssm_out[num_actual_tokens:].zero_() + preallocated_ssm_out_p = preallocated_ssm_out[:num_prefill_tokens] + preallocated_ssm_out_d = preallocated_ssm_out[ + num_prefill_tokens:num_actual_tokens] if num_prefills > 0: @@ -764,14 +769,17 @@ def convert_dt(): ) # norm - hidden_states = self.norm(preallocated_ssm_out, z[:num_actual_tokens]) + # Full padded length through norm/out_proj so the residual stream + # keeps a consistent row count (pad rows are zeros; row-wise norm + # keeps them finite). The caller trims real tokens via gather_ids. + hidden_states = self.norm(preallocated_ssm_out, z) # out_proj out = self.out_proj(hidden_states, lora_params=lora_params, layer_idx=self.layer_idx) - return out[:num_actual_tokens] + return out # We want to cache the largest indexing vector we'd ever need and mask it, vs From 58945e76baa81a60056eca5f9c7eda3ea5ac5d61 Mon Sep 17 00:00:00 2001 From: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> Date: Sat, 22 Aug 2026 01:30:21 -0700 Subject: [PATCH 2/4] [None][fix] Add Mamba2 piecewise boundary op so the traced graph stays composition-free Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> (cherry picked from commit 60a719bdb7675581e980a176f4ff97b150a1e59b) --- .../_torch/compilation/piecewise_optimizer.py | 1 + tensorrt_llm/_torch/compilation/utils.py | 15 ++ .../_torch/modules/mamba/mamba2_mixer.py | 252 +++++++++++++++--- 3 files changed, 229 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py index 53ee6d35edf7..5ade1813c7c1 100644 --- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py +++ b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py @@ -26,6 +26,7 @@ def _piecewise_boundary_ops(): "mla_custom_op_inplace", "mla_dsa_attn_inplace", "gdn_custom_op_inplace", + "mamba2_custom_op_inplace", "minimax_m3_attn_custom_op_inplace", ] return [ diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py index a03e84014109..b3e9d189e047 100644 --- a/tensorrt_llm/_torch/compilation/utils.py +++ b/tensorrt_llm/_torch/compilation/utils.py @@ -201,6 +201,21 @@ def inplace_info(): "gdn_custom_op_inplace": { 1: "output" }, + # Registered lazily: the op only exists once mamba2_mixer is imported + # (Mamba2/NemotronH family). Void boundary op mutating ssm_out: + # auto_functionalized returns (None, ssm_out), hence index 1. + "mamba2_custom_op_inplace": { + 1: "ssm_out" + }, + # Registered lazily: the op only exists once mamba2_mixer is + # imported (Mamba2/NemotronH family). Void op mutating (state, out): + # auto_functionalized returns (None, state, out), hence indices 1/2. + # Without this entry the pass leaves the functionalization clone of + # the full per-layer SSM state cache in every decode graph. + "flashinfer_selective_state_update": { + 1: "state", + 2: "out" + }, "minimax_m3_attn_custom_op_inplace": { 1: "output" }, diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 411cc28640f9..8ab742c721ab 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -15,6 +15,8 @@ import functools import os +import weakref +from typing import Optional import torch from einops import rearrange, repeat @@ -31,6 +33,7 @@ from ...model_config import ModelConfig from ...peft.lora.layer import LoraLayer, LoraModuleType from ...speculative import SpecMetadata +from ...utils import get_model_extra_attrs, is_torch_compiling from ..linear import Linear, TensorParallelMode from .causal_conv1d import causal_conv1d_fn, causal_conv1d_update from .causal_conv1d_triton import \ @@ -46,6 +49,107 @@ from .ssd_combined import mamba_chunk_scan_combined +def _extract_mamba2_extra_attrs(layer_idx: str): + extra_attrs = get_model_extra_attrs() + assert extra_attrs is not None, "Model extra attrs is not set" + + metadata_ref = extra_attrs.get("attention_metadata", None) + assert metadata_ref is not None, "Attention metadata is not set" + metadata = metadata_ref() + assert isinstance(metadata, AttentionMetadata) + + mamba2_layers = extra_attrs.get("mamba2_layers", None) + assert mamba2_layers is not None, "Mamba2 layer is not registered" + layer_ref = mamba2_layers.get(layer_idx, None) + assert layer_ref is not None, \ + f"Cannot find Mamba2 layer for layer {layer_idx}" + mamba_layer = layer_ref() + assert isinstance(mamba_layer, Mamba2Mixer) + + return metadata, mamba_layer, extra_attrs.get("spec_metadata", None) + + +@torch.library.custom_op("trtllm::mamba2_custom_op_inplace", + mutates_args=("ssm_out", )) +def mamba2_custom_op_inplace(zxbcdt: torch.Tensor, layer_idx: str, + ssm_out: torch.Tensor) -> None: + # Piecewise boundary op (mirror of trtllm::gdn_custom_op_inplace): the + # whole conv+SSM core runs eagerly inside one opaque node, so the traced + # graph never reads batch-composition ints (num_ctx_tokens/num_decodes + # would otherwise specialize a dynamo variant per composition, which + # exhausts the recompile limit and hard-fails under fullgraph=True) and + # every surrounding piece stays uniform in the num_tokens dim. + attn_metadata, mamba_layer, spec_metadata = _extract_mamba2_extra_attrs( + layer_idx) + mamba_layer.forward_core(zxbcdt, attn_metadata, + attn_metadata.mamba_metadata, spec_metadata, + ssm_out) + + +@torch.library.custom_op("trtllm::flashinfer_selective_state_update", + mutates_args=("state", "out"), + device_types="cuda") +def _flashinfer_selective_state_update_op( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + out: torch.Tensor, + dt_bias: Optional[torch.Tensor] = None, + dt_softplus: bool = False, + state_batch_indices: Optional[torch.Tensor] = None, + rand_seed: Optional[torch.Tensor] = None, + philox_rounds: int = 10) -> None: + # Opaque wrapper for the torch.compile path: flashinfer resolves its JIT + # module inside the Python call (torch.cuda.device_count() and friends), + # which dynamo cannot trace; the resulting graph break sits inside the + # decoder-layer loop, so dynamo skips the whole forward frame and the + # piecewise backend then fails on child frames that carry no + # input_ids/inputs_embeds placeholder. + # Contract notes: the op is void (remove_copy_for_mutates_args indexes + # every getitem user into the inplace_info map, so a real return value + # would KeyError at index 0), and `out` must be a plain intermediate, + # not a view of the caller's preallocated buffer (mutable view args are + # not handled by the rewrite). + kwargs = {} + if rand_seed is not None: + kwargs["rand_seed"] = rand_seed + kwargs["philox_rounds"] = philox_rounds + selective_state_update_fi(state, + x, + dt, + A, + B, + C, + D, + z=None, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + out=out, + **kwargs) + + +@_flashinfer_selective_state_update_op.register_fake +def _(state, + x, + dt, + A, + B, + C, + D, + out, + dt_bias=None, + dt_softplus=False, + state_batch_indices=None, + rand_seed=None, + philox_rounds=10) -> None: + return None + + class Mamba2Mixer(nn.Module): def __init__( @@ -72,6 +176,22 @@ def __init__( ): super().__init__() + # Register into the model's extra attrs so the mamba2 boundary custom + # op can recover this module (and the live metadata) from just a + # layer-idx string while the surrounding forward is traced. + self.layer_idx_str = str(layer_idx) + self.register_to_config = False + if config is not None: + if "mamba2_layers" not in config.extra_attrs: + config.extra_attrs["mamba2_layers"] = {} + suffix = 0 + while self.layer_idx_str in config.extra_attrs["mamba2_layers"]: + self.layer_idx_str = str(layer_idx) + f"_{suffix}" + suffix += 1 + config.extra_attrs["mamba2_layers"][self.layer_idx_str] = \ + weakref.ref(self) + self.register_to_config = True + config = config or ModelConfig() if config.mapping.enable_attention_dp: @@ -312,6 +432,58 @@ def forward( **kwargs, ) -> torch.Tensor: + # in_proj (LoRA is applied internally by Linear layer) + zxbcdt = self.in_proj(hidden_states, + lora_params=lora_params, + layer_idx=self.layer_idx) + z = zxbcdt[:, :self.tp_d_inner] + + # Preallocated output buffer shared by the prefill and decode + # segments; forward_core fills the real rows and zeroes the pad tail. + # Allocated here, not in forward_core, on every path: the boundary op + # is void and mutates this buffer in place (mutates_args + inplace_info), + # which is what makes it a piecewise / BCG boundary. The eager path is + # unified on purpose; its only extra cost is zero_() of an empty pad + # slice (a no-op). Do not move the allocation into forward_core. + preallocated_ssm_out = torch.empty( + [zxbcdt.shape[0], (self.tp_nheads * self.head_dim)], + dtype=zxbcdt.dtype, + device=zxbcdt.device, + ) + + if self.register_to_config and is_torch_compiling(): + # Route the conv+SSM core through the opaque boundary op (see its + # comment) so the traced graph stays free of batch-composition + # ints and uniform in the num_tokens dim. + torch.ops.trtllm.mamba2_custom_op_inplace(zxbcdt, + self.layer_idx_str, + preallocated_ssm_out) + else: + self.forward_core(zxbcdt, attn_metadata, mamba_metadata, + spec_metadata, preallocated_ssm_out) + + # norm + # Full padded length through norm/out_proj so the residual stream + # keeps a consistent row count (pad rows are zeros; row-wise norm + # keeps them finite). The caller trims real tokens via gather_ids. + hidden_states = self.norm(preallocated_ssm_out, z) + + # out_proj + out = self.out_proj(hidden_states, + lora_params=lora_params, + layer_idx=self.layer_idx) + + return out + + def forward_core( + self, + zxbcdt: torch.Tensor, + attn_metadata: AttentionMetadata, + mamba_metadata: Mamba2Metadata, + spec_metadata: SpecMetadata | None, + preallocated_ssm_out: torch.Tensor, + ) -> None: + # calculate split size num_prefills = attn_metadata.num_contexts num_decodes = attn_metadata.seq_lens.shape[0] - num_prefills @@ -330,13 +502,7 @@ def forward( state_indices_p, state_indices_d = torch.split(state_indices, batch_split_size) - # in_proj (LoRA is applied internally by Linear layer) - zxbcdt = self.in_proj(hidden_states, - lora_params=lora_params, - layer_idx=self.layer_idx) - - # Split z and dt with views. - z = zxbcdt[:, :self.tp_d_inner] + # Split dt with views (z is consumed by the gated norm in forward). dt = zxbcdt[:, self.tp_d_inner + self.tp_conv_dim:] # Slice instead of exact-sum split: under piecewise CUDA graphs the # token dim is padded to the capture bucket, so hidden_states can @@ -349,15 +515,9 @@ def forward( xbc_d = zxbcdt[num_prefill_tokens:num_actual_tokens, self.tp_d_inner:self.tp_d_inner + self.tp_conv_dim] - # Preallocate output tensor to avoid memcpy cost for merging prefill - # and decode outputs - preallocated_ssm_out = torch.empty( - [zxbcdt.shape[0], (self.tp_nheads * self.head_dim)], - dtype=zxbcdt.dtype, - device=zxbcdt.device, - ) - # Zero the pad tail (torch.empty) so the full-length gated norm below - # sees defined values in pad rows; no-op slice when not padded. + # Zero the pad tail (the buffer comes from torch.empty in forward) so + # the full-length gated norm sees defined values in pad rows; no-op + # slice when not padded. preallocated_ssm_out[num_actual_tokens:].zero_() preallocated_ssm_out_p = preallocated_ssm_out[:num_prefill_tokens] preallocated_ssm_out_d = preallocated_ssm_out[ @@ -757,29 +917,43 @@ def convert_dt(): ssu_kwargs['rand_seed'] = rand_seed[:1] ssu_kwargs['philox_rounds'] = self._philox_rounds - self.selective_state_update_func( - ssm_states, - x_d, - dt_d, - A, - B_d, - C_d, - D, - **ssu_kwargs, - ) - - # norm - # Full padded length through norm/out_proj so the residual stream - # keeps a consistent row count (pad rows are zeros; row-wise norm - # keeps them finite). The caller trims real tokens via gather_ids. - hidden_states = self.norm(preallocated_ssm_out, z) - - # out_proj - out = self.out_proj(hidden_states, - lora_params=lora_params, - layer_idx=self.layer_idx) - - return out + # Only while dynamo itself traces forward_core, i.e. a mixer + # built without a model config. Registered layers run this + # body eagerly inside the boundary op, where the plain call is + # the right one; the engine-level is_torch_compiling() flag + # stays set for the whole run and is not the test here. + if self._use_flashinfer and torch.compiler.is_compiling(): + # Route through the opaque custom op (see its comment); + # `out` is a fresh intermediate copied into the + # preallocated view afterwards so the op itself has no + # mutable view args. + ssu_out = torch.empty_like(x_d) + torch.ops.trtllm.flashinfer_selective_state_update( + ssm_states, + x_d, + dt_d, + A, + B_d, + C_d, + D, + ssu_out, + dt_bias=dt_bias, + dt_softplus=self.delta_softplus, + state_batch_indices=state_indices_d, + rand_seed=ssu_kwargs.get('rand_seed'), + philox_rounds=self._philox_rounds) + ssu_kwargs['out'].copy_(ssu_out) + else: + self.selective_state_update_func( + ssm_states, + x_d, + dt_d, + A, + B_d, + C_d, + D, + **ssu_kwargs, + ) # We want to cache the largest indexing vector we'd ever need and mask it, vs From 1c4aeb63823d97f5397465a0a248111de203435a Mon Sep 17 00:00:00 2001 From: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> Date: Sat, 22 Aug 2026 05:23:48 -0700 Subject: [PATCH 3/4] [None][fix] Route Mamba2 boundary op through eager_on_graph for breakable CUDA graphs Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> (cherry picked from commit 88bf5470c2cb461606ef4b41c2edc7af0d138eac) --- .../_torch/modules/mamba/mamba2_mixer.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 8ab742c721ab..6a001d8f1242 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -32,6 +32,8 @@ from ...attention.backends import AttentionMetadata from ...model_config import ModelConfig from ...peft.lora.layer import LoraLayer, LoraModuleType +from ...pyexecutor.breakable_cuda_graph import (eager_on_graph, + is_in_breakable_cuda_graph) from ...speculative import SpecMetadata from ...utils import get_model_extra_attrs, is_torch_compiling from ..linear import Linear, TensorParallelMode @@ -86,6 +88,9 @@ def mamba2_custom_op_inplace(zxbcdt: torch.Tensor, layer_idx: str, ssm_out) +maybe_bcg_mamba2_custom_op_inplace = eager_on_graph(mamba2_custom_op_inplace) + + @torch.library.custom_op("trtllm::flashinfer_selective_state_update", mutates_args=("state", "out"), device_types="cuda") @@ -451,13 +456,16 @@ def forward( device=zxbcdt.device, ) - 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): # Route the conv+SSM core through the opaque boundary op (see its - # comment) so the traced graph stays free of batch-composition - # ints and uniform in the num_tokens dim. - torch.ops.trtllm.mamba2_custom_op_inplace(zxbcdt, - self.layer_idx_str, - preallocated_ssm_out) + # comment): under torch.compile the traced graph stays free of + # batch-composition ints; under breakable CUDA graph capture the + # op is the eager bridge between captured segments. + maybe_bcg_mamba2_custom_op_inplace(zxbcdt, self.layer_idx_str, + preallocated_ssm_out) else: self.forward_core(zxbcdt, attn_metadata, mamba_metadata, spec_metadata, preallocated_ssm_out) From d508688eba53f90cb6a651f9f02877d1ff77e866 Mon Sep 17 00:00:00 2001 From: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> Date: Tue, 8 Sep 2026 00:54:07 -0700 Subject: [PATCH 4/4] [None][test] Add Nemotron-H breakable prefill CUDA graph parity test (MPI + Ray) Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> --- .../test_lists/test-db/l0_dgx_h100.yml | 2 + .../test_lists/test-db/l0_h100.yml | 3 + .../modeling/test_modeling_nemotron_h.py | 190 +++++++++++++++++- 3 files changed, 193 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 960bb97b3326..912fff1bce17 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -38,6 +38,7 @@ l0_dgx_h100: - disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16] # ------------- Model specific tests --------------- - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8[tp2_ep1] + - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph[tp2] TIMEOUT (90) - condition: ranges: system_gpu_count: @@ -267,6 +268,7 @@ l0_dgx_h100: - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part3" - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part4" - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" + - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph -k "tp2" TIMEOUT (90) - unittest/llmapi/test_async_llm.py -m "gpu2" - examples/test_ray.py::test_llm_inference_distributed_ray[tp2] - examples/test_ray.py::test_llm_inference_distributed_ray[pp2] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index eb898fc3a2b6..6d97cfaea0db 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -83,6 +83,8 @@ l0_h100: # reuse dense-8B greedy/logprob equality; accuracy stays on GSM8K/MMLU). - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_cuda_graph_overlap_scheduler - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_chunked_prefill + # Breakable prefill CUDA graphs (BCG) on the hybrid Mamba path; mpi_ray_parity (Ray id runs in the ray stage below). + - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph[tp1] TIMEOUT (90) - unittest/_torch/modeling/test_multimodal_encoder_graph.py # Qwen3.5-MoE-VL is hybrid (Mamba SSM + attention); FlashInfer's # chunk_gated_delta_rule GDN prefill kernel is sm90+ only, so this @@ -245,6 +247,7 @@ l0_h100: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logprobs[False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logprobs[True-TinyLlama-1.1B-Chat-v1.0] - unittest/_torch/executor/test_overlap_scheduler.py + - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph -k "tp1" TIMEOUT (90) - unittest/executor/test_shim_ray.py - unittest/_torch/ray_orchestrator/single_gpu/test_llm_sleep.py - unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py -m "part0" diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index d39ae9d40596..eedd53bd0c58 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -1,12 +1,14 @@ import pytest import torch from utils.llm_data import llm_models_root -from utils.util import skip_fp8_pre_ada, skip_gpu_memory_less_than +from utils.util import (skip_fp8_pre_ada, skip_gpu_memory_less_than, + skip_single_gpu) from tensorrt_llm import LLM from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.llmapi.llm import RequestOutput -from tensorrt_llm.llmapi.llm_args import CudaGraphConfig, LoadFormat +from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, LoadFormat, + PrefillCudaGraphBackend) from tensorrt_llm.sampling_params import SamplingParams @@ -216,3 +218,187 @@ def test_nemotron_h_chunked_prefill(): ) > 0, f"Prompt {i}: chunked prefill produced empty output" assert len(output.outputs[0].text ) > 0, f"Prompt {i}: chunked prefill produced empty text" + + +# Breakable prefill CUDA graphs (BCG) on the hybrid path. Under BCG the Mamba2 +# layers run through the eager_on_graph boundary op (mamba2_mixer.py), so this +# covers the NemotronH-specific bridge rather than the generic BCG runner +# (which unittest/_torch/executor/test_breakable_cuda_graph.py exercises). +_BCG_MAX_NUM_TOKENS = 256 +_BCG_CAPTURE_NUM_TOKENS = [64, 128, 256] +# Probability-ratio bound shared with ray_orchestrator/multi_gpu/ +# test_accuracy_with_allreduce_strategy.py::compare_logprobs (e^-2.30 ~ 0.1x). +_BCG_LOGPROB_TOLERANCE = 2.30 +# Context batches: an exact capture bucket, one token past a bucket (padded +# replay), two context requests of unequal length in one batch, and a prompt +# longer than max_num_tokens (chunked prefill: 256 + 44). +_BCG_CONTEXT_BATCHES = [ + [[17] * 128], + [[17] * 129], + [[17] * 64, [23] * 65], + [[31] * 300], +] +# The mixed-batch section adds the streaming decode request and the context +# request admitted while it decodes. +_BCG_MIXED_BATCH_REQUESTS = 2 +# Decode length of that streaming request: with EOS ignored this is a ~1 s +# generation window, orders of magnitude longer than the client-side gap +# before the second request is submitted. +_BCG_STREAM_MAX_TOKENS = 64 +# One (token_ids, logprobs) entry per request from each backend run. +_BCG_NUM_REQUESTS = (sum(len(batch) for batch in _BCG_CONTEXT_BATCHES) + + _BCG_MIXED_BATCH_REQUESTS) + + +def _first_step_logprobs(output) -> torch.Tensor: + """Log-probabilities of the first generated position for one request.""" + return torch.log_softmax( + output.outputs[0].generation_logits[0].float().cpu(), dim=-1) + + +def _assert_mixed_batch_overlap(decoding_out, admitted_out) -> None: + """Prove that `admitted` was prefilled while `decoding` was still decoding. + + Both requests return per-request iteration metrics taken from the engine's + iteration counter. With the overlap scheduler disabled, a batch capacity of + 4 and 65 + 1 tokens inside the token budget, the scheduler runs every active + generation request each iteration, so admitted's first iteration (its + context chunk) falling inside decoding's generation window means that + iteration carried a context chunk and a decode token together. + """ + dec = decoding_out.outputs[0].request_perf_metrics + adm = admitted_out.outputs[0].request_perf_metrics + assert dec is not None and adm is not None, "request_perf_metrics missing" + assert None not in (dec.first_iter, dec.last_iter, adm.first_iter), ( + f"iteration metrics not populated: decoding {dec.first_iter}.." + f"{dec.last_iter}, admitted {adm.first_iter}") + assert dec.first_iter < adm.first_iter <= dec.last_iter, ( + f"no mixed batch: decoding ran iterations [{dec.first_iter}, " + f"{dec.last_iter}] but admitted was prefilled at iteration " + f"{adm.first_iter}") + + +def _run_nemotron_h_prefill_backend(backend: PrefillCudaGraphBackend, + tp_size: int): + """Run the fixed prompt schedule with one prefill CUDA graph backend. + + Returns one (token_ids, first_step_logprobs) pair per request, in a + deterministic order shared by both backends. + """ + # ignore_eos on every compared request: the two backends may legitimately + # pick different greedy tokens after the first step (see the test + # docstring), and if one of them is EOS the arms would stop at different + # lengths for a difference the test otherwise accepts. + sampling_params = SamplingParams(max_tokens=4, + temperature=0.0, + ignore_eos=True, + return_generation_logits=True) + per_request = [] + with LLM( + model=f"{llm_models_root(check=True)}/{_NANO_30B_BF16}", + tensor_parallel_size=tp_size, + # Pin NCCL so the allreduce path under BCG is the plain one on + # every platform (AUTO may pick MNNVL on multi-node NVLink). + allreduce_strategy="NCCL", + max_batch_size=4, + max_num_tokens=_BCG_MAX_NUM_TOKENS, + enable_chunked_prefill=True, + disable_overlap_scheduler=True, + gather_generation_logits=True, + kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"), + cuda_graph_config=CudaGraphConfig(enable_padding=True, + max_batch_size=4), + prefill_cuda_graph_backend=backend, + prefill_capture_num_tokens=_BCG_CAPTURE_NUM_TOKENS, + ) as llm: + # The validator must not have downgraded the requested backend. + assert llm.args.prefill_cuda_graph_backend == backend + for batch in _BCG_CONTEXT_BATCHES: + for output in llm.generate(batch, sampling_params=sampling_params): + per_request.append((list(output.outputs[0].token_ids), + _first_step_logprobs(output))) + + # Mixed batch: admit a context request while another request is + # decoding, so BCG replays a batch that carries both a context chunk + # and decode tokens. The overlap is proved afterwards from the two + # requests' iteration metrics rather than assumed. + decoding = llm.generate_async([17] * 128, + sampling_params=SamplingParams( + max_tokens=_BCG_STREAM_MAX_TOKENS, + temperature=0.0, + ignore_eos=True, + return_generation_logits=True, + return_perf_metrics=True), + streaming=True) + # In streaming mode every response carries only the newest step's + # logits, so the first-step distribution must be read from the first + # streamed response, before the stream moves on. + next(decoding) + decoding_first_step_logprobs = _first_step_logprobs(decoding) + admitted = llm.generate_async([23] * 65, + sampling_params=SamplingParams( + max_tokens=4, + temperature=0.0, + ignore_eos=True, + return_generation_logits=True, + return_perf_metrics=True), + streaming=False) + decoding_out = decoding.result() + admitted_out = admitted.result() + _assert_mixed_batch_overlap(decoding_out, admitted_out) + per_request.append((list(decoding_out.outputs[0].token_ids), + decoding_first_step_logprobs)) + per_request.append((list(admitted_out.outputs[0].token_ids), + _first_step_logprobs(admitted_out))) + return per_request + + +@skip_gpu_memory_less_than((2 * 30 + 1) * 2**30) +@pytest.mark.mpi_ray_parity +@pytest.mark.parametrize("tp_size", [ + 1, + pytest.param(2, marks=skip_single_gpu), +], + ids=lambda n: f"tp{n}") +def test_nemotron_h_breakable_prefill_cuda_graph(tp_size): + """Real-weight Nano: breakable prefill CUDA graphs vs eager prefill. + + Same context / chunked-prefill / mixed-admission schedule with + prefill_cuda_graph_backend DISABLED and BREAKABLE. The first generated + token is the direct product of the (captured) prefill, so its + distribution is compared per request within the repo's accepted + probability-ratio bound and BCG's greedy pick must be one of eager's top-2. + Full-sequence greedy equality is reported but not required: Nano-30B-A3B + is MoE and flips greedy tokens under padding-induced numeric drift (see + the note above test_nemotron_h_sanity). NCCL allreduce is pinned: tp1 has + no allreduce at all, tp2 (skipped below 2 GPUs) runs the NCCL allreduce inside the captured + segments. Runs under the MPI executor and, with --run-ray, under the Ray + executor (mpi_ray_parity). The mixed-admission overlap is proved from the + two requests' iteration metrics, not assumed. + """ + eager = _run_nemotron_h_prefill_backend(PrefillCudaGraphBackend.DISABLED, + tp_size) + bcg = _run_nemotron_h_prefill_backend(PrefillCudaGraphBackend.BREAKABLE, + tp_size) + assert len(eager) == len(bcg) == _BCG_NUM_REQUESTS + + identical = 0 + worst_diff = 0.0 + for i, ((eager_ids, eager_lp), (bcg_ids, + bcg_lp)) in enumerate(zip(eager, bcg)): + assert len(bcg_ids) == len(eager_ids) > 0, ( + f"request {i}: BCG produced {len(bcg_ids)} tokens, " + f"eager {len(eager_ids)}") + eager_top1 = eager_ids[0] + diff = (eager_lp[eager_top1] - bcg_lp[eager_top1]).abs().item() + worst_diff = max(worst_diff, diff) + assert diff < _BCG_LOGPROB_TOLERANCE, ( + f"request {i}: BCG log-prob of eager's first token differs by " + f"{diff:.3f} nats (bound {_BCG_LOGPROB_TOLERANCE})") + eager_top2 = torch.topk(eager_lp, 2).indices.tolist() + assert bcg_ids[0] in eager_top2, ( + f"request {i}: BCG first token {bcg_ids[0]} not in eager top-2 " + f"{eager_top2}") + identical += int(bcg_ids == eager_ids) + print(f"BCG vs eager: {identical}/{len(eager)} sequences identical, " + f"worst first-token log-prob diff {worst_diff:.3f} nats")