From e64767f0f728ed25fd06182d984e56d9a89ffbca Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Sun, 1 Nov 2026 12:58:55 -0800 Subject: [PATCH 1/3] [TRTLLM-17146][feat] Add the Muse Spark Open decoder fast paths Fused kernels for the decoder's hot ops, each with the reference composition kept as the fallback for the inputs it does not support: - musespark_open_kernels: Triton kernels for the post-expert norm and combine, the gated residual join (alone and with the next prenorm), the routing weights, the full top-8 routing over 256 experts, the FP32 router projection, and the attention output norm and gate. - musespark_open_boundary: a CuTe DSL kernel that splits a residual join and the next prenorm across a CTA cluster (GB300, width 8192, 1 to 32 tokens). - musespark_open_finalize_ar: the expert finalize fused with a one-shot all-reduce over CUDA IPC (GB300, BF16 TP4 with EP4 on one node, 1 to 128 tokens). Construction is agreed across the TP group and falls back to the unfused finalize and NCCL. - The pre-expert norm fused with the experts' NVFP4 input quantization, and small batches pool on the auxiliary stream while the router runs. The four-rank tests spawn their ranks with MPIPoolExecutor; the single-GPU tests are listed in l0_gb300. Signed-off-by: Michal Guzek --- .../_torch/models/modeling_musespark_open.py | 436 +++++++- .../_torch/models/musespark_open_boundary.py | 327 ++++++ .../models/musespark_open_finalize_ar.py | 979 ++++++++++++++++++ .../_torch/models/musespark_open_kernels.py | 598 +++++++++++ .../test_lists/test-db/l0_gb300.yml | 11 + .../modeling/test_modeling_musespark_open.py | 18 +- .../test_musespark_open_boundary_fusion.py | 287 +++++ .../test_musespark_open_cluster_boundary.py | 201 ++++ .../test_musespark_open_finalize_allreduce.py | 384 +++++++ .../test_musespark_open_finalize_dispatch.py | 124 +++ .../test_musespark_open_finalize_lifecycle.py | 423 ++++++++ .../test_musespark_open_finalize_local.py | 288 ++++++ .../test_musespark_open_finalize_stress.py | 243 +++++ .../test_musespark_open_full_routing.py | 180 ++++ .../test_musespark_open_fusion_switches.py | 124 +++ .../test_musespark_open_moe_finalize.py | 222 ++++ ...test_musespark_open_moe_prepare_overlap.py | 76 ++ .../modeling/test_musespark_open_norm_gate.py | 76 ++ .../test_musespark_open_pre_expert_quant.py | 136 +++ .../test_musespark_open_router_projection.py | 102 ++ 20 files changed, 5180 insertions(+), 55 deletions(-) create mode 100644 tensorrt_llm/_torch/models/musespark_open_boundary.py create mode 100644 tensorrt_llm/_torch/models/musespark_open_finalize_ar.py create mode 100644 tensorrt_llm/_torch/models/musespark_open_kernels.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_boundary_fusion.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_finalize_allreduce.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_finalize_dispatch.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_finalize_lifecycle.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_finalize_local.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_finalize_stress.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_full_routing.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_fusion_switches.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_moe_finalize.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_moe_prepare_overlap.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_norm_gate.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_pre_expert_quant.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_router_projection.py diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open.py b/tensorrt_llm/_torch/models/modeling_musespark_open.py index 5338c1a919a7..0ddf1303599f 100644 --- a/tensorrt_llm/_torch/models/modeling_musespark_open.py +++ b/tensorrt_llm/_torch/models/modeling_musespark_open.py @@ -26,12 +26,20 @@ The residual stream is FP32. Each branch reads a BF16 cast of it, and its output rejoins through a norm-preserving gated residual add. + +On CUDA the routing, the router projection, the attention output norm and gate, +the expert finalize and the residual joins run as fused kernels +(`musespark_open_kernels`), mostly for small batches. On GB300 a residual join +and the next layer's prenorm can share one CTA-cluster kernel +(`musespark_open_boundary`), and with TP4 on one node the expert finalize runs +fused with the all-reduce (`musespark_open_finalize_ar`). Each falls back to +the reference composition for the inputs it does not support. """ import contextlib import copy import os -from typing import Any, Dict, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple import torch from torch import nn @@ -39,6 +47,8 @@ from tensorrt_llm._torch.configs.musespark_open import SLIDING_ATTENTION, MuseSparkOpenConfig from tensorrt_llm._utils import nvtx_range from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping from ...models.modeling_utils import QuantConfig from ..attention.backends import AttentionMetadata @@ -49,6 +59,7 @@ ) from ..attention.qk_norm_attention import QKNormRoPEAttention from ..distributed import AllReduce +from ..kernels.fused_ops.rmsnorm_fp4_quant import rmsnorm_fp4_quant_available from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding @@ -56,10 +67,15 @@ from ..modules.logits_processor import LogitsProcessor, scale_and_soft_cap_logits from ..modules.rms_norm import RMSNorm from ..moe.fused_moe import MiniMaxM2MoeRoutingMethod, create_moe -from ..utils import AuxStreamType +from ..moe.fused_moe.moe_op_backend import TRTLLMOpBackend +from ..utils import AuxStreamType, Fp4QuantizedTensor +from . import musespark_open_kernels as fused_kernels from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel +if TYPE_CHECKING: + from .musespark_open_finalize_ar import FinalizeIpcWorkspace + # Perception subtrees in the checkpoint. The text decoder never builds these, # and their tensors are hundreds of GiB, so the loader skips them unread. PERCEPTION_KEY_PREFIXES = ("vision_model.", "modalities.", "audio_model.") @@ -173,8 +189,21 @@ def post_expert_norm_combine( """Weight-first RMSNorm each selected expert output, then combine, in FP32. The MoE kernel's own finalize cannot do this: RMSNorm is nonlinear, so - normalizing after the top-k sum is a different function. + normalizing after the top-k sum is a different function. On CUDA this is + one fused kernel, which sums the squares and the slots in a different + order from the composition below. """ + if fused_kernels.post_expert_norm_combine_supported(gemm2_output, num_tokens): + return fused_kernels.post_expert_norm_combine( + gemm2_output, + expanded_idx_to_permuted_idx, + routing_weights, + norm_weight, + eps, + num_tokens=num_tokens, + top_k=top_k, + hidden_size=hidden_size, + ) normed, _ = post_expert_norm_rows( gemm2_output, expanded_idx_to_permuted_idx, @@ -195,6 +224,8 @@ def _routing_weights( output_dtype: torch.dtype, ) -> Tuple[torch.Tensor, torch.Tensor]: """Gather and renormalize the unbiased scores of the selected experts.""" + if fused_kernels.routing_weights_supported(scores, topk_idx, output_dtype): + return fused_kernels.routing_weights(scores, topk_idx, eps) weights = scores.gather(1, topk_idx) weights = weights / (weights.sum(dim=-1, keepdim=True) + eps) return topk_idx.to(torch.int32), weights.to(output_dtype) @@ -225,7 +256,14 @@ def requires_separated_routing(self) -> bool: return True def apply(self, router_logits: torch.Tensor, input_ids: Optional[torch.Tensor] = None): - scores, scores_with_bias = self.get_scores(router_logits, self.e_score_correction_bias) + bias = self.e_score_correction_bias + if fused_kernels.full_routing_supported(router_logits, bias, self.top_k, self.output_dtype): + # One kernel for the whole step. Its picks are sorted, with ties + # going to the lowest expert id, where `torch.topk(sorted=False)` + # leaves the order unspecified; the finalize sums the experts in + # pick order. + return fused_kernels.full_routing(router_logits, bias, self.renorm_eps) + scores, scores_with_bias = self.get_scores(router_logits, bias) _, topk_idx = torch.topk(scores_with_bias, k=self.top_k, dim=-1, sorted=False) # The weights come from the unbiased scores; the bias only steers which # experts are picked. @@ -356,13 +394,66 @@ def add_after_rmsnorm( `gain` selects the norm: `1 + w` for the weight-first FFN norm, `None` for the weightless attention one. The gain product and the normalized branch are each rounded to the activation dtype, as the - modules they stand for do. + modules they stand for do. On CUDA this is one kernel. """ alpha, beta = self.coefficients() + if fused_kernels.gated_residual_supported(residual, branch, gain): + return fused_kernels.gated_residual_after_rmsnorm( + residual, branch, gain, alpha, beta, eps + ) scaled = (gain * branch if gain is not None else branch).to(torch.float32) normed = (scaled * torch.rsqrt(scaled.pow(2).mean(-1, keepdim=True) + eps)).to(branch.dtype) return alpha * residual + beta * normed.to(torch.float32) + def add_after_rmsnorm_and_norm( + self, + residual: torch.Tensor, + branch: torch.Tensor, + eps: float, + next_norm: RMSNorm, + gain: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """`add_after_rmsnorm`, then `next_norm` of the result cast to the activation dtype. + + Returns `(residual, next_input)`: the updated FP32 residual stream and + the input of the branch that reads it next. On CUDA both come from one + kernel; on GB300, for 1 to 32 tokens of the 8192-wide stream, from a + kernel that splits each token across a CTA cluster. + """ + if not ( + next_norm.use_gemma + and fused_kernels.gated_residual_and_norm_supported( + residual, branch, gain, next_norm.weight + ) + ): + updated = self.add_after_rmsnorm(residual, branch, eps, gain) + return updated, next_norm(updated.to(branch.dtype)) + alpha, beta = self.coefficients() + if ( + not torch.compiler.is_compiling() + and residual.shape[1] == 8192 + and residual.shape[0] <= 32 + ): + # Imported here: the module loads CuTe DSL, which only this path needs. + from . import musespark_open_boundary as boundary + + if boundary.cluster_boundary_supported( + residual, branch, gain, alpha, beta, next_norm.weight + ): + return boundary.cluster_boundary( + residual, + branch, + gain, + alpha, + beta, + next_norm.weight, + eps, + next_norm.variance_epsilon, + ) + return fused_kernels.gated_residual_after_rmsnorm_and_norm( + residual, branch, gain, alpha, beta, eps, next_norm.weight, next_norm.variance_epsilon + ) + class MuseSparkOpenLogitsProcessor(LogitsProcessor): """LM head epilogue: metap multiplier then a tanh soft cap, both FP32.""" @@ -489,8 +580,12 @@ def apply_output_gate(self, attention_output: torch.Tensor, gate: torch.Tensor) Order matters: the checkpoint normalizes each head of the attention output before multiplying by `sigmoid(o_gate)`, and only then - projects. + projects. On CUDA, up to 128 tokens take one kernel for both. """ + if fused_kernels.attention_norm_gate_supported(attention_output, gate, self.head_dim): + return fused_kernels.attention_norm_gate( + attention_output, gate, self.head_dim, self.o_head_norm.variance_epsilon + ) normed = self.o_head_norm(attention_output.reshape(-1, self.head_dim)).reshape( attention_output.shape ) @@ -528,6 +623,15 @@ class MuseSparkOpenMoE(nn.Module): different model. """ + # The auxiliary stream the expert input is prepared on (see + # prepare_moe_inputs), and the fused finalize + all-reduce workspace, which + # `MuseSparkOpenModel` sets when it can serve this layer. Without them the + # MoE takes the unfused paths. + _prepare_stream: Optional[torch.cuda.Stream] = None + finalize_allreduce_workspace: Optional["FinalizeIpcWorkspace"] = None + # Whether this forward's finalize also reduces across the TP group. + _reduce_in_finalize = False + def __init__( self, model_config: ModelConfig[MuseSparkOpenConfig], @@ -595,6 +699,10 @@ def __init__( dtype=config.torch_dtype, use_gemma=True, ) + # `1 + w` of the pre-expert norm, for the norm fused with the experts' + # NVFP4 input quantization (see prepare_expert_input). Non-persistent, + # filled by cache_derived_state. + self.register_buffer("_pre_expert_quant_gain", None, persistent=False) self.post_expert_norm_eps = config.rms_post_norm_eps # Weight-first post-expert norm, applied per selected expert output. self.post_expert_norm_weight = nn.Parameter( @@ -663,11 +771,140 @@ def __init__( layer_idx, ("router", "pool", "experts", "finalize", "all_reduce", "unpool"), ) + if aux_stream is not None: + self._prepare_stream = aux_stream + self._prepare_fork_event = torch.cuda.Event() + self._prepare_ready_event = torch.cuda.Event() + + def cache_derived_state(self) -> None: + """Cache the pre-expert norm's `1 + w`; the loader calls this after the weights land.""" + weight = self.pre_expert_norm.weight + if weight is None or weight.is_meta: + return + self._pre_expert_quant_gain = (weight.float() + 1).to(weight.dtype) + + def post_load_weights(self) -> None: + self.cache_derived_state() + + def _apply(self, *args, **kwargs): + # As in MuseSparkOpenRMSNormWeightFirst: a cached `1 + w` would not + # follow a dtype move of the weight, so drop it on any move. + out = super()._apply(*args, **kwargs) + self._pre_expert_quant_gain = None + return out def router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - """FP32 router logits, scaled by the metap multiplier before routing.""" + """FP32 router logits, scaled by the metap multiplier before routing. + + On CUDA, up to 128 tokens take a kernel that keeps the FP32 weights and + accumulation and folds in the input cast and the scale. + """ + weight = getattr(self.router, "weight", None) + if ( + weight is not None + and not torch.compiler.is_compiling() + and fused_kernels.router_projection_supported(hidden_states, weight) + ): + return fused_kernels.router_projection( + hidden_states, + weight, + self.metap_multiplier, + overlapped=self._prepare_stream is not None and self.experts.comm is None, + ) return self.router(hidden_states.to(torch.float32)) * self.metap_multiplier + def prepare_expert_input(self, pooled: torch.Tensor) -> torch.Tensor | Fp4QuantizedTensor: + """The pre-expert norm, fused with the experts' NVFP4 input quantization where they take it. + + The fused kernel quantizes the FP32 norm result without the reference's + intermediate BF16 round. It serves the native TRTLLM-Gen NVFP4 experts + with a single static input scale and no communication; everything else + takes the norm alone and leaves quantization to the experts. + """ + experts = self.experts + backend = getattr(experts, "backend", experts) + if ( + not pooled.is_cuda + or pooled.dtype != torch.bfloat16 + or torch.compiler.is_compiling() + or getattr(experts, "comm", None) is not None + or not isinstance(getattr(backend, "op_backend", None), TRTLLMOpBackend) + or not getattr(backend, "has_nvfp4", False) + or getattr(backend, "fc31_act_scale", None) is not None + or pooled.shape[0] == 0 + or pooled.shape[0] > getattr(experts, "moe_max_num_tokens", pooled.shape[0]) + or not rmsnorm_fp4_quant_available() + ): + return self.pre_expert_norm(pooled) + scale = getattr(backend, "fc31_input_scale", None) + if scale is None or scale.numel() != 1: + return self.pre_expert_norm(pooled) + gain = self._pre_expert_quant_gain + if gain is None: + weight = self.pre_expert_norm.weight + gain = (weight.float() + 1).to(weight.dtype) + from flashinfer.norm import rmsnorm_fp4quant + + fp4, sf = rmsnorm_fp4quant( + pooled, + gain, + eps=self.pre_expert_norm.variance_epsilon, + global_scale=scale.reshape(1), + block_size=16, + scale_format="e4m3", + # The MoE swizzles the scales itself, after any dispatch. + is_sf_swizzled_layout=False, + ) + return Fp4QuantizedTensor(fp4.view(torch.uint8), sf.view(torch.uint8), is_sf_swizzled=False) + + def prepare_moe_inputs( + self, hidden_states: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor | Fp4QuantizedTensor, Optional[torch.cuda.Event]]: + """Router logits and the expert input. + + Returns `(router_logits, expert_input, input_ready_event)`. Up to 128 + tokens pool on the auxiliary stream while the router and the routing run + on the current one; the returned event marks the expert input as + written, and the MoE joins it before its first read. + """ + if not ( + self._prepare_stream is not None + and hidden_states.is_cuda + and 0 < hidden_states.shape[0] <= 128 + and not torch.compiler.is_compiling() + and self.experts.comm is None + ): + with self.nvtx("router"): + router_logits = self.router_logits(hidden_states) + with self.nvtx("pool"): + expert_input = self.prepare_expert_input(self.pooler(hidden_states)) + return router_logits, expert_input, None + self._prepare_fork_event.record() + with torch.cuda.stream(self._prepare_stream), self.nvtx("pool"): + self._prepare_fork_event.wait() + expert_input = self.prepare_expert_input(self.pooler(hidden_states)) + self._prepare_ready_event.record() + # No record_stream, which under graph capture would hold every layer's + # blocks until the capture ends: the current stream reads the expert + # input only after the join, and the next use of the auxiliary stream + # forks from the current one again. + with self.nvtx("router"): + router_logits = self.router_logits(hidden_states) + return router_logits, expert_input, self._prepare_ready_event + + def _finalize_allreduce_applies(self, num_tokens: int) -> bool: + """Whether this forward's finalize also reduces across the TP group. + + The workspace takes at most 128 tokens per call, and its calls must + follow one stream on each rank, so a batch the MoE would split into + chunks stays unfused. + """ + return ( + self.finalize_allreduce_workspace is not None + and 1 <= num_tokens <= min(128, self.experts.moe_max_num_tokens) + and not torch.compiler.is_compiling() + ) + def finalize( self, *, @@ -684,10 +921,23 @@ def finalize( Each output row depends on one `(token, expert)` row at a time and non-local slots contribute exact zeros, so combining the per-rank results equals finalizing after the combine. + + When forward chose the fused finalize + all-reduce, this returns the + BF16 sum over the TP group instead of this rank's FP32 partial. """ # The post-expert norm is the same for every expert. del token_selected_slots with self.nvtx("finalize"): + if self._reduce_in_finalize: + return self.finalize_allreduce_workspace.finalize( + gemm2_output, + expanded_idx_to_permuted_idx, + # TRTLLM-Gen dispatches BF16 weights; the kernel reads FP32. + routing_weights.to(torch.float32), + self.post_expert_norm_weight, + self.post_expert_norm_eps, + num_tokens, + ) if self.mapping.moe_tp_size > 1: # Backend row permutations are rank-local, so reduce in canonical # (token, selected-expert) order rather than in GEMM scratch @@ -717,33 +967,41 @@ def finalize( def forward( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata ) -> torch.Tensor: - with self.nvtx("router"): - router_logits = self.router_logits(hidden_states) - with self.nvtx("pool"): - pooled = self.pre_expert_norm(self.pooler(hidden_states)) - with self.nvtx("experts"): - # With the finalize registered, the MoE returns this rank's tokens - # already normalized, weighted and summed over the local experts. - combined = self.experts( - pooled, - router_logits, - do_finalize=False, - output_dtype=hidden_states.dtype, - all_rank_num_tokens=attn_metadata.all_rank_num_tokens - if attn_metadata is not None - else None, - use_dp_padding=False, - ) + router_logits, expert_input, input_ready_event = self.prepare_moe_inputs(hidden_states) + # Chosen here rather than in finalize, which cannot tell the reduce + # below that it already ran. + reduce_in_finalize = self._finalize_allreduce_applies(hidden_states.shape[0]) + self._reduce_in_finalize = reduce_in_finalize + try: + with self.nvtx("experts"): + # With the finalize registered, the MoE returns this rank's + # tokens already normalized, weighted and summed over the local + # experts. + combined = self.experts( + expert_input, + router_logits, + do_finalize=False, + output_dtype=hidden_states.dtype, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens + if attn_metadata is not None + else None, + use_dp_padding=False, + input_ready_event=input_ready_event, + ) + finally: + self._reduce_in_finalize = False # Cast before the reduce: with expert parallelism the partial sums cross # the wire at the activation width, and they are cast to it for the # unpooler either way. combined = combined.to(hidden_states.dtype) # Expert parallelism leaves a partial sum on each rank; reduce it over - # the TP group. Not under attention DP, where the ranks hold different - # tokens and the MoE's own combine has already reduced the partials, and - # not under pure MoE TP, where finalize has already reduced each expert. + # the TP group. Not when the fused finalize already did, not under + # attention DP, where the ranks hold different tokens and the MoE's own + # combine has already reduced the partials, and not under pure MoE TP, + # where finalize has already reduced each expert. if ( - not self.mapping.enable_attention_dp + not reduce_in_finalize + and not self.mapping.enable_attention_dp and self.mapping.has_tp() and self.mapping.moe_tp_size == 1 ): @@ -806,12 +1064,24 @@ def forward( position_ids: torch.IntTensor, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, + precomputed_attn_input: Optional[torch.Tensor] = None, + next_layer_norm: Optional[RMSNorm] = None, **kwargs, - ) -> torch.Tensor: + ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]: + """Run the layer on the FP32 residual stream `hidden_states`. + + Each residual join computes the next branch's prenorm in the same + kernel. Given `next_layer_norm`, the last join also normalizes for the + next layer, and the layer returns `(hidden_states, next_attn_input)`; + that layer takes the input as `precomputed_attn_input`. + """ # hidden_states is the FP32 residual stream; each branch reads a cast of # it to the activation dtype and writes back through the gated add. with self.nvtx("pre_qkv_proj_norm"): - attn_input = self.pre_qkv_proj_norm(hidden_states.to(self.dtype)) + if precomputed_attn_input is not None: + attn_input = precomputed_attn_input + else: + attn_input = self.pre_qkv_proj_norm(hidden_states.to(self.dtype)) # The attention projection and its TP reduction stay in this range. with self.nvtx("attn"): attn_out = self.self_attn( @@ -821,28 +1091,89 @@ def forward( **kwargs, ) with self.nvtx("norm_add_attn"): - hidden_states = self.add_attn.add_after_rmsnorm( - hidden_states, attn_out, self.post_attn_norm.variance_epsilon + hidden_states, ffn_input = self.add_attn.add_after_rmsnorm_and_norm( + hidden_states, attn_out, self.post_attn_norm.variance_epsilon, self.pre_ffn_norm ) with self.nvtx("moe"): - ffn_out = self.feed_forward( - self.pre_ffn_norm(hidden_states.to(self.dtype)), attn_metadata - ) + ffn_out = self.feed_forward(ffn_input, attn_metadata) with self.nvtx("norm_add_ffn"): - hidden_states = self.add_ffn.add_after_rmsnorm( - hidden_states, - ffn_out, - self.post_ffn_norm.variance_epsilon, - self.post_ffn_norm.gain(), - ) + next_attn_input = None + if next_layer_norm is None: + hidden_states = self.add_ffn.add_after_rmsnorm( + hidden_states, + ffn_out, + self.post_ffn_norm.variance_epsilon, + self.post_ffn_norm.gain(), + ) + else: + hidden_states, next_attn_input = self.add_ffn.add_after_rmsnorm_and_norm( + hidden_states, + ffn_out, + self.post_ffn_norm.variance_epsilon, + next_layer_norm, + self.post_ffn_norm.gain(), + ) # Speculative drafters that read target hidden states capture the FP32 # residual stream. spec_metadata = kwargs.get("spec_metadata") if spec_metadata is not None: spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, None) + if next_attn_input is not None: + return hidden_states, next_attn_input return hidden_states +def _finalize_allreduce_supported( + mapping: Mapping, config: MuseSparkOpenConfig, layers: nn.ModuleList +) -> bool: + """Whether the fused finalize + all-reduce can serve this deployment. + + It needs BF16 TP4 on one GB300 node with expert parallelism over the whole + TP group and the checkpoint's expert shape. Every rank reaches the same + answer, so either all of them construct the workspace, which is + collective, or none does. + """ + return ( + mapping.tp_size == 4 + and len({rank // mapping.gpus_per_node for rank in mapping.tp_group}) == 1 + and mapping.moe_ep_size == mapping.tp_size + and not mapping.enable_attention_dp + and config.torch_dtype == torch.bfloat16 + and torch.cuda.is_available() + and torch.cuda.get_device_capability() == (10, 3) + and all( + layer.feed_forward.expert_hidden_size == 4096 + and layer.feed_forward.top_k == 8 + and getattr(layer.feed_forward.experts, "comm", None) is None + for layer in layers + ) + ) + + +def _build_finalize_allreduce_workspace(mapping: Mapping) -> Optional["FinalizeIpcWorkspace"]: + """Construct the workspace, or return None to keep the unfused finalize and all-reduce. + + Construction probes the TP group and raises when the GPUs cannot map each + other's memory or the ranks disagree. That is a property of the + deployment, so it degrades to the unfused path instead of failing model + construction. + """ + from .musespark_open_finalize_ar import HAS_CUTEDSL, FinalizeIpcWorkspace + + if not HAS_CUTEDSL: + return None + try: + workspace = FinalizeIpcWorkspace(mapping) + except (ValueError, RuntimeError) as exc: + logger.warning( + f"MuseSparkOpen finalize+all-reduce fusion unavailable ({exc}); " + "using the unfused finalize and AllReduce path." + ) + return None + logger.info("MuseSparkOpen: the expert finalize runs fused with the TP4 all-reduce.") + return workspace + + class MuseSparkOpenModel(DecoderModel): def __init__(self, model_config: ModelConfig[MuseSparkOpenConfig]): super().__init__(model_config) @@ -870,6 +1201,14 @@ def __init__(self, model_config: ModelConfig[MuseSparkOpenConfig]): for layer_idx in range(config.num_hidden_layers) ] ) + # One fused finalize + all-reduce workspace serves every layer. + self.finalize_allreduce_workspace = None + if _finalize_allreduce_supported(model_config.mapping, config, self.layers): + self.finalize_allreduce_workspace = _build_finalize_allreduce_workspace( + model_config.mapping + ) + for layer in self.layers: + layer.feed_forward.finalize_allreduce_workspace = self.finalize_allreduce_workspace # The final norm uses the stored weight as is (gain centered at 0). self.norm = RMSNorm( hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype @@ -893,13 +1232,24 @@ def forward( embeds = embeds * self.m_emb # The FP32 residual stream, entered once. hidden_states = embeds.to(torch.float32) - for decoder_layer in self.layers: - hidden_states = decoder_layer( + # Each layer but the last also returns the next layer's attention input. + attn_input = None + for index, decoder_layer in enumerate(self.layers): + next_norm = ( + self.layers[index + 1].pre_qkv_proj_norm if index + 1 < len(self.layers) else None + ) + result = decoder_layer( position_ids=position_ids, hidden_states=hidden_states, attn_metadata=attn_metadata, + precomputed_attn_input=attn_input, + next_layer_norm=next_norm, **kwargs, ) + if next_norm is None: + hidden_states = result + else: + hidden_states, attn_input = result return self.norm(hidden_states.to(self.dtype)) diff --git a/tensorrt_llm/_torch/models/musespark_open_boundary.py b/tensorrt_llm/_torch/models/musespark_open_boundary.py new file mode 100644 index 000000000000..9b52d1625e82 --- /dev/null +++ b/tensorrt_llm/_torch/models/musespark_open_boundary.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A residual join and the next RMSNorm in one GB300 kernel, split across a CTA cluster. + +For 1 to 32 tokens of the 8192-wide residual stream, each token is split +across a cluster of CTAs that reduce both variances through distributed shared +memory. The kernel keeps the reference's BF16 rounding of the normalized branch +and of the residual fed to the next norm; its reduction order can still move +single values across a BF16 rounding boundary. Each reduction has its own +shared-memory inbox and barrier, and CUDA graph replays read live inputs. +""" + +import functools + +import torch + +# Unlike the finalize + all-reduce kernel, this one needs no minimum DSL release. +try: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + import cutlass.utils as utils + from cutlass import Float32, Int32 + from cutlass._mlir.dialects import llvm + from cutlass.cute.runtime import make_ptr + from cutlass.cutlass_dsl import T, dsl_user_op +except ImportError: + _HAS_CUTE_DSL = False +else: + _HAS_CUTE_DSL = True + +if _HAS_CUTE_DSL: + + @dsl_user_op + def _float_bits(value, *, loc=None, ip=None): + """The int32 with the same bit pattern as an FP32 value.""" + bits = llvm.inline_asm( + T.i32(), + [Float32(value).ir_value(loc=loc, ip=ip)], + "mov.b32 $0, $1;", + "=r,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return Int32(bits) + + class _ClusterBoundaryKernel: + def __init__( + self, + m: int, + cluster: int, + threads: int, + gain: bool, + post_eps: float, + pre_eps: float, + ): + self.m, self.cluster, self.threads, self.gain = m, cluster, threads, gain + self.post_eps, self.pre_eps = float(post_eps), float(pre_eps) + self.iters = 8192 // (cluster * threads * 8) + self.warps = threads // 32 + self.slots = cluster * self.warps + + @cute.jit + def __call__( + self, + branch: cute.Pointer, + residual: cute.Pointer, + gain: cute.Pointer, + alpha: cute.Pointer, + beta: cute.Pointer, + weight: cute.Pointer, + output: cute.Pointer, + next_output: cute.Pointer, + stream: cuda.CUstream, + ): + self.kernel(branch, residual, gain, alpha, beta, weight, output, next_output).launch( + grid=[self.m * self.cluster, 1, 1], + block=[self.threads, 1, 1], + cluster=[self.cluster, 1, 1], + smem=self.slots * 4 * 4 * 2 + 16, + stream=stream, + use_pdl=True, + ) + + @cute.jit + def load(self, ptr: cute.Pointer, dtype: cutlass.Constexpr): + layout = cute.make_layout((8, self.iters), stride=(1, self.threads * self.cluster * 8)) + reg = cute.make_rmem_tensor(cute.make_layout((8, self.iters)), dtype) + atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=128) + cute.copy(atom, cute.make_tensor(ptr.align(16), layout), reg) + return reg + + @cute.jit + def reduce( + self, value: Float32, partials: cute.Tensor, bars: cute.Tensor, phase: cutlass.Constexpr + ): + tid = cute.arch.thread_idx()[0] + lane, warp = tid % 32, tid // 32 + rank = cute.arch.block_idx()[0] % self.cluster + value = cute.arch.warp_reduction_sum(value) + if cutlass.const_expr(self.cluster > 1): + if lane < self.cluster: + # Lane c sends this warp's partial to CTA c as one 16-byte + # message that completes 16 bytes of that CTA's transaction. + slot = partials.iterator + (phase * self.slots + rank * self.warps + warp) * 4 + cute.arch.store_async_dsmem( + slot.align(16), + (_float_bits(value), Int32(0), Int32(0), Int32(0)), + bars.iterator + phase, + lane, + ) + cute.arch.mbarrier_wait(bars.iterator + phase, 0) + else: + if lane == 0: + partials[(phase * self.slots + warp) * 4] = value + cute.arch.barrier() + total = Float32(0) + for slot in range(lane, self.slots, 32): + total += partials[(phase * self.slots + slot) * 4] + return cute.arch.warp_reduction_sum(total) + + @cute.kernel + def kernel( + self, + branch: cute.Pointer, + residual: cute.Pointer, + gain: cute.Pointer, + alpha: cute.Pointer, + beta: cute.Pointer, + weight: cute.Pointer, + output: cute.Pointer, + next_output: cute.Pointer, + ): + tid = cute.arch.thread_idx()[0] + rank = cute.arch.block_idx()[0] % self.cluster + row = cute.arch.block_idx()[0] // self.cluster + offset = cute.assume((rank * self.threads + tid) * 8, divby=8) + smem = utils.SmemAllocator() + partials = smem.allocate_tensor( + Float32, cute.make_layout((self.slots * 4 * 2,)), byte_alignment=16 + ) + bars = smem.allocate_tensor(cutlass.Uint64, cute.make_layout((2,)), byte_alignment=8) + if cutlass.const_expr(self.cluster > 1): + if tid < 2: + cute.arch.mbarrier_init(bars.iterator + tid, 1) + cute.arch.mbarrier_arrive_and_expect_tx(bars.iterator + tid, self.slots * 16) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + # The previous kernel may still be writing any input, including + # coefficients computed per call, so wait before the first load. + cute.arch.griddepcontrol_wait() + a = self.load(alpha + offset, Float32) + b = self.load(beta + offset, Float32) + w = self.load(weight + offset, cutlass.BFloat16) + g = self.load(gain + offset, cutlass.BFloat16) + x = self.load(branch + row * 8192 + offset, cutlass.BFloat16) + r = self.load(residual + row * 8192 + offset, Float32) + scaled = cute.make_rmem_tensor(cute.make_layout((8, self.iters)), Float32) + updated = cute.make_rmem_tensor(cute.make_layout((8, self.iters)), Float32) + normed = cute.make_rmem_tensor(cute.make_layout((8, self.iters)), cutlass.BFloat16) + total = Float32(0) + for p in cutlass.range_constexpr(self.iters): + for v in cutlass.range_constexpr(8): + value = Float32(x[v, p]) + if cutlass.const_expr(self.gain): + value = Float32(cutlass.BFloat16(value * Float32(g[v, p]))) + scaled[v, p] = value + total += value * value + total = self.reduce(total, partials, bars, 0) + inv = cute.math.rsqrt(total / Float32(8192) + Float32(self.post_eps), fastmath=True) + total = Float32(0) + for p in cutlass.range_constexpr(self.iters): + for v in cutlass.range_constexpr(8): + value = Float32(cutlass.BFloat16(scaled[v, p] * inv)) + u = a[v, p] * r[v, p] + b[v, p] * value + updated[v, p] = u + rounded = Float32(cutlass.BFloat16(u)) + scaled[v, p] = rounded + total += rounded * rounded + total = self.reduce(total, partials, bars, 1) + inv = cute.math.rsqrt(total / Float32(8192) + Float32(self.pre_eps), fastmath=True) + for p in cutlass.range_constexpr(self.iters): + for v in cutlass.range_constexpr(8): + normed[v, p] = cutlass.BFloat16( + scaled[v, p] * inv * (Float32(w[v, p]) + Float32(1)) + ) + layout = cute.make_layout((8, self.iters), stride=(1, self.threads * self.cluster * 8)) + atom32 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + atom16 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=128 + ) + cute.copy( + atom32, updated, cute.make_tensor((output + row * 8192 + offset).align(16), layout) + ) + cute.copy( + atom16, + normed, + cute.make_tensor((next_output + row * 8192 + offset).align(16), layout), + ) + if cutlass.const_expr(self.cluster > 1): + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + cute.arch.griddepcontrol_launch_dependents() + + # Argument dtypes, in the kernel's argument order. + _DTYPES = ( + cutlass.BFloat16, + Float32, + cutlass.BFloat16, + Float32, + Float32, + cutlass.BFloat16, + Float32, + cutlass.BFloat16, + ) + + @functools.lru_cache(maxsize=128) + def _compile_boundary( + tokens: int, has_gain: bool, post_eps: float, pre_eps: float, device_index: int + ): + cluster = 16 if tokens == 1 else (8 if tokens <= 8 else 4) + threads = min(256, 8192 // (cluster * 8)) + with torch.cuda.device(device_index): + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + _ClusterBoundaryKernel(tokens, cluster, threads, has_gain, post_eps, pre_eps), + *[ + make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=16) + for dtype in _DTYPES + ], + stream, + options="--enable-tvm-ffi", + ) + + +@functools.lru_cache(maxsize=16) +def cluster_boundary_available(device: torch.device) -> bool: + """Whether the kernel runs on `device`: CuTe DSL and a GB300 (SM103) GPU.""" + return ( + _HAS_CUTE_DSL + and device.type == "cuda" + and torch.cuda.get_device_capability(device) == (10, 3) + ) + + +def cluster_boundary_supported( + residual: torch.Tensor, + branch: torch.Tensor, + gain: torch.Tensor | None, + alpha: torch.Tensor, + beta: torch.Tensor, + next_weight: torch.Tensor, +) -> bool: + """Check shape, dtype and vector-copy alignment before dispatch.""" + if residual.ndim != 2 or residual.shape[1] != 8192 or not 1 <= residual.shape[0] <= 32: + return False + if not cluster_boundary_available(residual.device): + return False + if ( + residual.dtype != torch.float32 + or branch.dtype != torch.bfloat16 + or branch.shape != residual.shape + ): + return False + gain_arg = next_weight if gain is None else gain + vectors = (gain_arg, alpha, beta, next_weight) + if any(t.shape != (8192,) for t in vectors): + return False + if any( + t.device != residual.device or not t.is_contiguous() or t.data_ptr() % 16 + for t in (residual, branch, *vectors) + ): + return False + return ( + gain_arg.dtype == torch.bfloat16 + and next_weight.dtype == torch.bfloat16 + and alpha.dtype == torch.float32 + and beta.dtype == torch.float32 + ) + + +def cluster_boundary( + residual: torch.Tensor, + branch: torch.Tensor, + gain: torch.Tensor | None, + alpha: torch.Tensor, + beta: torch.Tensor, + next_weight: torch.Tensor, + post_eps: float, + pre_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """`(alpha * residual + beta * rmsnorm(gain * branch), next Gemma RMSNorm of it)`. + + Raises `ValueError` for inputs `cluster_boundary_supported` rejects; the + caller keeps the fallback for them. + """ + if not cluster_boundary_supported(residual, branch, gain, alpha, beta, next_weight): + raise ValueError("Unsupported cluster boundary inputs, device, or alignment") + gain_arg = next_weight if gain is None else gain + updated, normalized = torch.empty_like(residual), torch.empty_like(branch) + compiled = _compile_boundary( + residual.shape[0], gain is not None, post_eps, pre_eps, residual.device.index + ) + args = (branch, residual, gain_arg, alpha, beta, next_weight, updated, normalized) + # Resolve the stream per call: a CUDA graph captures on a different stream + # from the eager call that filled the compilation cache. + stream = cuda.CUstream(torch.cuda.current_stream(residual.device).cuda_stream) + compiled(*[make_ptr(dtype, tensor.data_ptr()) for dtype, tensor in zip(_DTYPES, args)], stream) + return updated, normalized diff --git a/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py b/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py new file mode 100644 index 000000000000..ae9b6515100f --- /dev/null +++ b/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py @@ -0,0 +1,979 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fused expert finalize and tensor-parallel all-reduce for Muse Spark Open. + +`FinalizeIpcWorkspace.finalize` normalizes every selected expert row on its own +(weight-first RMSNorm with zero-centered gains), combines the rows of a token with +its routing weights, rounds that per-rank partial to BF16, and sums the partials +of the four ranks of a single-node TP group in FP32. One kernel launch does all +of it, and every rank ends up with the same BF16 result. + +Communication is one-shot over CUDA IPC. Each rank stores its BF16 partial into +its own slab of every peer's receive buffer and then reads the slabs addressed to +it. A 32-bit word that still holds the sentinel `0x80000000` has not arrived yet; +producers turn BF16 negative zero into positive zero, so published data never +looks like the sentinel. The reader writes the sentinel back as soon as it has a +value. + +Two receive buffers alternate between calls. Before a peer can write call +`n + 2` into a buffer, it must have read this rank's call `n + 1`, which this +rank only publishes after call `n` (including its resets) completed. A buffer is +therefore never rewritten while it still holds unread data, and a value left over +from an earlier call can never satisfy a later read. + +Each rank owns four int32 control words: the current buffer index and an arrival +counter. Every CTA reads the index once and then arrives on the counter; the last +arrival flips the index. `atom.inc` wraps the counter to zero in the same step, +so no counter grows without bound and graph replays advance the state on device. + +Workspace per rank: `2 x 4 x 128 x 4096` BF16 receive slots (8 MiB) plus the +control words. Construction is collective over the TP group, compiles and loads +the kernel, and either succeeds on every rank or raises the same exception type +on every rank. Nothing compiles afterwards. + +Release never frees memory that a peer may still map. CUDA leaves freeing an IPC +export undefined while an importer still has it open, so `close()` is +collective: once the owner has quiesced every rank, each rank synchronizes its +device, then closes its mappings of peer memory, and only after the whole group +confirmed both steps does any rank free its own buffer. Dropping the workspace +without `close()` stays local: it waits for this rank's device and closes its +mappings, but leaves its own buffer allocated until the process exits. +""" + +import functools +import importlib.metadata +import math +import os +import socket +import sys +import zlib + +import numpy as np +import torch + +from tensorrt_llm._ipc_utils import _ipc_mem_handle_from_bytes, _ipc_mem_handle_to_bytes +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +# Problem shape fixed by the supported model configuration. +_HIDDEN = 4096 +_TOP_K = 8 +_MAX_TOKENS = 128 +_RANKS = 4 +# BF16 values moved per 16-byte access. +_VECTOR = 8 +_SLAB_BYTES = _MAX_TOKENS * _HIDDEN * 2 +_BUFFERS = 2 +_SENTINEL = 0x80000000 +_CONTROL_WORDS = 4 +# Bump when the device protocol changes so mixed builds refuse to pair up. +_PROTOCOL_VERSION = 1 +# CuTe DSL release this kernel was validated with. +_VALIDATED_DSL = "4.8.0" +_IPC_ALIGNMENT = 2 << 20 + +# Construction status codes exchanged across the TP group. +_OK = 0 +_UNSUPPORTED = 1 +_FAILED = 2 + + +def _dsl_version_supported() -> bool: + from packaging.version import InvalidVersion, Version + + try: + installed = Version(importlib.metadata.version("nvidia-cutlass-dsl")) + except (importlib.metadata.PackageNotFoundError, InvalidVersion): + return False + return installed >= Version(_VALIDATED_DSL) + + +try: + import cuda.bindings.driver as cuda + import cuda.bindings.runtime as cudart + import cutlass + import cutlass.cute as cute + import cutlass.utils as cutlass_utils + from cutlass import BFloat16, Float32, Int32, Int64, Uint32 + from cutlass._mlir import ir + from cutlass._mlir.dialects import llvm + from cutlass.cute.runtime import make_ptr + from cutlass.cutlass_dsl import T, dsl_user_op +except ImportError: + HAS_CUTEDSL = False +else: + HAS_CUTEDSL = _dsl_version_supported() + + +def _pdl_enabled() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +if HAS_CUTEDSL: + _FOUR_WORDS = "!llvm.struct<(i32,i32,i32,i32)>" + + def _asm(result_type, operands, text, constraints, loc, ip, side_effects=True): + return llvm.inline_asm( + result_type, + operands, + text, + constraints, + has_side_effects=side_effects, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + def _u32(value, loc, ip): + return Uint32(value).ir_value(loc=loc, ip=ip) + + def _i64(value, loc, ip): + return Int64(value).ir_value(loc=loc, ip=ip) + + def _f32(value, loc, ip): + return Float32(value).ir_value(loc=loc, ip=ip) + + def _split_words(result, loc, ip): + return tuple( + Uint32(llvm.extractvalue(T.i32(), result, [i], loc=loc, ip=ip)) for i in range(4) + ) + + # Every global access below is a side-effecting asm statement, so none of + # them can move above griddepcontrol.wait or be merged across poll attempts. + + @dsl_user_op + def _load_input_x4(address, enabled, *, loc=None, ip=None): + """16-byte input load; yields zeros without touching memory when disabled.""" + result = _asm( + ir.Type.parse(_FOUR_WORDS), + [_i64(address, loc, ip), _u32(enabled, loc, ip)], + "{\n.reg .pred p;\n.reg .b32 a, b, c, d;\n" + "mov.b32 a, 0;\nmov.b32 b, 0;\nmov.b32 c, 0;\nmov.b32 d, 0;\n" + "setp.ne.u32 p, $5, 0;\n" + "@p ld.global.nc.v4.b32 {a, b, c, d}, [$4];\n" + "mov.b32 $0, a;\nmov.b32 $1, b;\nmov.b32 $2, c;\nmov.b32 $3, d;\n}", + "=r,=r,=r,=r,l,r", + loc, + ip, + ) + return _split_words(result, loc, ip) + + @dsl_user_op + def _load_input_s32(address, *, loc=None, ip=None): + result = _asm( + T.i32(), [_i64(address, loc, ip)], "ld.global.nc.b32 $0, [$1];", "=r,l", loc, ip + ) + return Int32(result) + + @dsl_user_op + def _load_input_f32(address, *, loc=None, ip=None): + result = _asm( + T.f32(), [_i64(address, loc, ip)], "ld.global.nc.f32 $0, [$1];", "=f,l", loc, ip + ) + return Float32(result) + + @dsl_user_op + def _load_peer_x4(address, *, loc=None, ip=None): + """Relaxed system-scope 16-byte load, as four independent 32-bit words.""" + result = _asm( + ir.Type.parse(_FOUR_WORDS), + [_i64(address, loc, ip)], + "ld.relaxed.sys.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + ) + return _split_words(result, loc, ip) + + @dsl_user_op + def _store_peer_x4(address, w0, w1, w2, w3, *, loc=None, ip=None): + """Relaxed system-scope 16-byte store, as four independent 32-bit words.""" + _asm( + None, + [_i64(address, loc, ip)] + [_u32(w, loc, ip) for w in (w0, w1, w2, w3)], + "st.relaxed.sys.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc, + ip, + ) + + @dsl_user_op + def _store_x4(address, w0, w1, w2, w3, *, loc=None, ip=None): + _asm( + None, + [_i64(address, loc, ip)] + [_u32(w, loc, ip) for w in (w0, w1, w2, w3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc, + ip, + ) + + @dsl_user_op + def _load_control(address, *, loc=None, ip=None): + result = _asm( + T.i32(), + [_i64(address, loc, ip)], + "ld.relaxed.gpu.global.b32 $0, [$1];", + "=r,l", + loc, + ip, + ) + return Uint32(result) + + @dsl_user_op + def _store_control(address, value, *, loc=None, ip=None): + _asm( + None, + [_i64(address, loc, ip), _u32(value, loc, ip)], + "st.relaxed.gpu.global.b32 [$0], $1;", + "l,r", + loc, + ip, + ) + + @dsl_user_op + def _arrive(address, last, *, loc=None, ip=None): + """Arrive on the counter; it wraps to zero when the old value equals `last`.""" + result = _asm( + T.i32(), + [_i64(address, loc, ip), _u32(last, loc, ip)], + "atom.acq_rel.gpu.global.inc.u32 $0, [$1], $2;", + "=r,l,r", + loc, + ip, + ) + return Uint32(result) + + @dsl_user_op + def _unpack_bf16x2(word, *, loc=None, ip=None): + """Return the FP32 values of the low and high BF16 halves of `word`.""" + result = _asm( + ir.Type.parse("!llvm.struct<(f32,f32)>"), + [_u32(word, loc, ip)], + "{\n.reg .b32 lo, hi;\nshl.b32 lo, $2, 16;\nand.b32 hi, $2, 0xffff0000;\n" + "mov.b32 $0, lo;\nmov.b32 $1, hi;\n}", + "=f,=f,r", + loc, + ip, + side_effects=False, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + @dsl_user_op + def _pack_bf16x2(low, high, *, loc=None, ip=None): + """Round to nearest-even BF16 and pack; `low` lands in the low half.""" + result = _asm( + T.i32(), + [_f32(high, loc, ip), _f32(low, loc, ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc, + ip, + side_effects=False, + ) + return Uint32(result) + + @dsl_user_op + def _pack_publishable_bf16x2(low, high, *, loc=None, ip=None): + """Like `_pack_bf16x2`, but negative zeros become positive zeros. + + The sentinel word is a +0 low half next to a -0 high half, so a payload + without negative zeros can never be mistaken for it. + """ + result = _asm( + T.i32(), + [_f32(high, loc, ip), _f32(low, loc, ip)], + "{\n.reg .b32 w, m;\n.reg .pred p;\n" + "cvt.rn.bf16x2.f32 w, $1, $2;\n" + "and.b32 m, w, 0xffff;\nsetp.eq.u32 p, m, 0x8000;\n@p and.b32 w, w, 0xffff0000;\n" + "and.b32 m, w, 0xffff0000;\nsetp.eq.u32 p, m, 0x80000000;\n@p and.b32 w, w, 0xffff;\n" + "mov.b32 $0, w;\n}", + "=r,f,f", + loc, + ip, + side_effects=False, + ) + return Uint32(result) + + @dsl_user_op + def _zero_unless(value, enabled, *, loc=None, ip=None): + result = _asm( + T.f32(), + [_f32(value, loc, ip), _u32(enabled, loc, ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $2, 0;\nselp.f32 $0, $1, 0f00000000, p;\n}", + "=f,f,r", + loc, + ip, + side_effects=False, + ) + return Float32(result) + + class _FinalizeAllReduceKernel: + """One-shot finalize + all-reduce across the four ranks. + + CTA `b` handles token `b`; each thread owns one 16-byte column vector of + the row for all `_TOP_K` slots. + """ + + def __init__(self, pdl: bool): + self.threads = _HIDDEN // _VECTOR + self.warps = self.threads // 32 + self.pdl = pdl + + @cute.jit + def __call__( + self, + gemm: cute.Pointer, + indices: cute.Pointer, + weights: cute.Pointer, + norm: cute.Pointer, + out: cute.Pointer, + control: cute.Pointer, + receive: cute.Pointer, + peer0: cute.Pointer, + peer1: cute.Pointer, + peer2: cute.Pointer, + peer3: cute.Pointer, + rows: Int32, + row_stride: Int32, + tokens: Int32, + rank: Int32, + eps: Float32, + dry: Int32, + stream: cuda.CUstream, + ): + self.kernel( + gemm, + indices, + weights, + norm, + out, + control, + receive, + peer0, + peer1, + peer2, + peer3, + rows, + row_stride, + tokens, + rank, + eps, + dry, + ).launch( + grid=[tokens, 1, 1], + block=[self.threads, 1, 1], + smem=128 + 32 * self.warps, + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel( + self, + gemm: cute.Pointer, + indices: cute.Pointer, + weights: cute.Pointer, + norm: cute.Pointer, + out: cute.Pointer, + control: cute.Pointer, + receive: cute.Pointer, + peer0: cute.Pointer, + peer1: cute.Pointer, + peer2: cute.Pointer, + peer3: cute.Pointer, + rows: Int32, + row_stride: Int32, + tokens: Int32, + rank: Int32, + eps: Float32, + dry: Int32, + ): + peers = (peer0, peer1, peer2, peer3) + tid = cute.arch.thread_idx()[0] + token = cute.arch.block_idx()[0] + blocks = cute.arch.grid_dim()[0] + column = tid * _VECTOR + lane = tid % 32 + warp = tid // 32 + + smem = cutlass_utils.SmemAllocator() + snapshot = smem.allocate_tensor(Int32, cute.make_layout((4,)), byte_alignment=16) + warp_sums = smem.allocate_tensor( + Float32, cute.make_layout((self.warps * _TOP_K,)), byte_alignment=16 + ) + sums = smem.allocate_tensor(Float32, cute.make_layout((_TOP_K,)), byte_alignment=16) + + # A dry launch only makes the driver load the module. + if dry == 0: + if cutlass.const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + + # Selected rows and routing weights; invalid slots read nothing. + selection = Int64(token) * _TOP_K + gemm_address = gemm.toint() + row_index = [] + row_weight = [] + row_words = [] + for k in cutlass.range_constexpr(_TOP_K): + index = _load_input_s32(indices.toint() + (selection + k) * 4) + row_index.append(index) + row_weight.append(_load_input_f32(weights.toint() + (selection + k) * 4)) + valid = Uint32((index >= 0) & (index < rows)) + address = gemm_address + (Int64(index) * Int64(row_stride) + Int64(column)) * 2 + row_words.append(_load_input_x4(address, valid)) + norm_words = _load_input_x4(norm.toint() + Int64(column) * 2, Uint32(1)) + + # Read the buffer index once, then arrive; the last arrival flips it. + # Issued after the input loads so both round trips overlap; the + # reduction barrier below makes the snapshot visible to every thread. + control_address = control.toint() + if tid == 0: + current = _load_control(control_address) + snapshot[0] = Int32(current) + previous = _arrive(control_address + 4, Uint32(blocks - 1)) + if previous == Uint32(blocks - 1): + _store_control(control_address, current ^ Uint32(1)) + + gains = [] + for w in cutlass.range_constexpr(4): + low, high = _unpack_bf16x2(norm_words[w]) + gains.append(low + Float32(1.0)) + gains.append(high + Float32(1.0)) + + # Sum of squares of the scaled row, per slot. + partials = [] + for k in cutlass.range_constexpr(_TOP_K): + square_sum = Float32(0.0) + for w in cutlass.range_constexpr(4): + low, high = _unpack_bf16x2(row_words[k][w]) + scaled_low = low * gains[2 * w] + scaled_high = high * gains[2 * w + 1] + square_sum = ( + square_sum + scaled_low * scaled_low + scaled_high * scaled_high + ) + partials.append(cute.arch.warp_reduction_sum(square_sum)) + if lane == 0: + for k in cutlass.range_constexpr(_TOP_K): + warp_sums[warp * _TOP_K + k] = partials[k] + cute.arch.barrier() + if tid < _TOP_K: + total = Float32(0.0) + for w in cutlass.range_constexpr(self.warps): + total = total + warp_sums[w * _TOP_K + tid] + sums[tid] = total + cute.arch.barrier() + + # Normalize each slot, weight it, and accumulate in FP32. + result = [Float32(0.0)] * _VECTOR + for k in cutlass.range_constexpr(_TOP_K): + valid = Uint32((row_index[k] >= 0) & (row_index[k] < rows)) + scale = row_weight[k] * cute.math.rsqrt(sums[k] / Float32(_HIDDEN) + eps) + scale = _zero_unless(scale, valid) + for w in cutlass.range_constexpr(4): + low, high = _unpack_bf16x2(row_words[k][w]) + result[2 * w] = result[2 * w] + scale * (low * gains[2 * w]) + result[2 * w + 1] = result[2 * w + 1] + scale * (high * gains[2 * w + 1]) + own = [] + for w in cutlass.range_constexpr(4): + own.append(_pack_publishable_bf16x2(result[2 * w], result[2 * w + 1])) + + # Publish this rank's partial into its slab of every peer's buffer. + buffer = Int64(snapshot[0]) + slot = (Int64(token) * _HIDDEN + Int64(column)) * 2 + own_slab = (buffer * _RANKS + Int64(rank)) * _SLAB_BYTES + for p in cutlass.range_constexpr(_RANKS): + if rank != p: + _store_peer_x4( + peers[p].toint() + own_slab + slot, own[0], own[1], own[2], own[3] + ) + + # Consume every rank's partial in rank order and reset its slot. + total = [Float32(0.0)] * _VECTOR + receive_address = receive.toint() + for p in cutlass.range_constexpr(_RANKS): + w0, w1, w2, w3 = own[0], own[1], own[2], own[3] + if rank != p: + address = receive_address + (buffer * _RANKS + p) * _SLAB_BYTES + slot + w0, w1, w2, w3 = _load_peer_x4(address) + while ( + (w0 == _SENTINEL) + | (w1 == _SENTINEL) + | (w2 == _SENTINEL) + | (w3 == _SENTINEL) + ): + w0, w1, w2, w3 = _load_peer_x4(address) + _store_x4(address, _SENTINEL, _SENTINEL, _SENTINEL, _SENTINEL) + words = (w0, w1, w2, w3) + for w in cutlass.range_constexpr(4): + low, high = _unpack_bf16x2(words[w]) + total[2 * w] = total[2 * w] + low + total[2 * w + 1] = total[2 * w + 1] + high + + packed = [] + for w in cutlass.range_constexpr(4): + packed.append(_pack_bf16x2(total[2 * w], total[2 * w + 1])) + _store_x4(out.toint() + slot, packed[0], packed[1], packed[2], packed[3]) + + # Dependents may start only after the last access to the buffers. + if cutlass.const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + + @functools.cache + def _compiled_kernel(device_index: int, pdl: bool): + def placeholder(dtype, alignment): + return make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=alignment) + + with torch.cuda.device(device_index): + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + _FinalizeAllReduceKernel(pdl), + placeholder(BFloat16, 16), + placeholder(Int32, 4), + placeholder(Float32, 4), + placeholder(BFloat16, 16), + placeholder(BFloat16, 16), + placeholder(Uint32, 16), + placeholder(Uint32, 16), + *[placeholder(Uint32, 16) for _ in range(_RANKS)], + Int32(0), + Int32(_HIDDEN), + Int32(1), + Int32(0), + Float32(1e-6), + Int32(1), + stream, + options="--enable-tvm-ffi", + ) + + +def _fingerprint() -> int: + dsl = importlib.metadata.version("nvidia-cutlass-dsl") if HAS_CUTEDSL else "none" + text = f"{_PROTOCOL_VERSION}|{dsl}|{_RANKS}|{_HIDDEN}|{_TOP_K}|{_MAX_TOKENS}" + return zlib.crc32(text.encode()) + + +def _describe(error: Exception) -> str: + return f"{type(error).__name__}: {error}" + + +def _host_id() -> int: + return zlib.crc32(socket.gethostname().encode()) + + +def _device_id(device: torch.device) -> int: + return zlib.crc32(str(torch.cuda.get_device_properties(device).uuid).encode()) + + +class FinalizeIpcWorkspace: + """Persistent state of the fused finalize + all-reduce for one TP4 group on one node. + + One instance can be shared by every layer of a model; calls must be issued + in the same order on every rank, ordered on one stream (or explicitly) per + rank. `close()` releases it on every rank; see its documentation for what + dropping it does instead. + """ + + # What the destructor sees if __init__ failed before setting up any state. + _closed = True + + def __init__(self, mapping: Mapping): + if not isinstance(mapping, Mapping): + raise TypeError(f"mapping must be a Mapping, got {type(mapping).__name__}") + self._ranks = mapping.tp_size + self._rank = mapping.tp_rank + self._device = torch.device("cuda", torch.cuda.current_device()) + self._export = 0 + self._imports = [] + self._control = None + self._kernel = None + self._peer_args = () + self._receive = 0 + self._closed = False + + from tensorrt_llm._torch.distributed.communicator import Distributed + + self._dist = Distributed.get(mapping) + self._agree("setup", *self._guarded(self._prepare, mapping)) + self._allocate_ipc() + self._agree("initialization", *self._guarded(self._initialize)) + + # Construction stages. Each stage turns any local failure into a status + # and agrees on it with the whole group, so that every rank either goes on + # or raises the same exception type; a rank never waits for peers that + # already gave up. + + @staticmethod + def _guarded(stage, *args): + """Run this rank's part of a stage; an exception of any type becomes a failure status. + + A rank that raised past the following agreement would leave its peers waiting in it. + TVM-FFI, for one, reports launch failures as ValueError, TypeError and other types. + """ + try: + return stage(*args) + except Exception as error: # noqa: BLE001 - every failure must reach the agreement + return _FAILED, _describe(error) + + def _prepare(self, mapping: Mapping): + if not HAS_CUTEDSL: + return _UNSUPPORTED, "CuTe DSL is unavailable or older than the validated release" + nodes = len({rank // mapping.gpus_per_node for rank in mapping.tp_group}) + if self._ranks != _RANKS or nodes != 1: + return _UNSUPPORTED, "the fusion needs TP4 on one node" + if torch.cuda.get_device_capability(self._device) != (10, 3): + return _UNSUPPORTED, "the kernel is validated on SM103 only" + try: + self._kernel = _compiled_kernel(self._device.index, _pdl_enabled()) + # Compilation can fail in many ways; every one must reach the agreement. + except Exception as error: # noqa: BLE001 + return _FAILED, f"kernel compilation failed: {error!r}" + return _OK, "" + + def _agree(self, stage: str, status: int, detail: str) -> None: + identity = [status, _fingerprint(), self._rank, _host_id(), _device_id(self._device)] + table = self._dist.tp_allgather_int64(identity) + problems = [] + if status != _OK: + problems.append(f"rank {self._rank}: {detail}") + statuses = table[:, 0] + if (table[:, 1] != table[0, 1]).any(): + statuses = np.maximum(statuses, _UNSUPPORTED) + problems.append("ranks run different kernel builds or configurations") + if sorted(table[:, 2].tolist()) != list(range(self._ranks)): + statuses = np.maximum(statuses, _UNSUPPORTED) + problems.append("TP ranks are not distinct") + hosts = len(set(table[:, 3].tolist())) + if hosts != 1: + statuses = np.maximum(statuses, _UNSUPPORTED) + problems.append(f"the TP group spans {hosts} hosts") + devices = {(int(host), int(device)) for host, device in table[:, 3:5]} + if len(devices) != self._ranks: + statuses = np.maximum(statuses, _UNSUPPORTED) + problems.append("two ranks share one GPU") + worst = int(statuses.max()) + if worst == _OK: + return + message = f"MuseSparkOpen finalize+all-reduce {stage} failed on rank(s) " + ", ".join( + str(i) for i, s in enumerate(statuses) if s != _OK + ) + if problems: + message += ": " + "; ".join(problems) + self._abandon(ValueError if worst == _UNSUPPORTED else RuntimeError, message) + + def _allocate_ipc(self) -> None: + size = math.ceil(_BUFFERS * self._ranks * _SLAB_BYTES / _IPC_ALIGNMENT) * _IPC_ALIGNMENT + status, detail, handle = _OK, "", bytes(64) + try: + error, pointer = cudart.cudaMalloc(size) + if error != cudart.cudaError_t.cudaSuccess: + status, detail = _FAILED, f"cudaMalloc failed: {error}" + else: + self._export = int(pointer) + error, ipc_handle = cudart.cudaIpcGetMemHandle(pointer) + if error != cudart.cudaError_t.cudaSuccess: + status, detail = _FAILED, f"cudaIpcGetMemHandle failed: {error}" + else: + handle = _ipc_mem_handle_to_bytes(ipc_handle) + except Exception as error: # noqa: BLE001 - every failure must reach the agreement + status, detail, handle = _FAILED, _describe(error), bytes(64) + words = np.frombuffer(handle.ljust(64, b"\0")[:64], dtype=np.int64) + table = self._dist.tp_allgather_int64([status, *words.tolist()]) + failed = [i for i, s in enumerate(table[:, 0]) if s != _OK] + if failed: + self._abandon( + RuntimeError, + f"MuseSparkOpen finalize+all-reduce allocation failed on rank(s) {failed}" + + (f": {detail}" if detail else ""), + ) + + status, detail = _OK, "" + peers = [] + try: + for peer in range(self._ranks): + if peer == self._rank: + peers.append(self._export) + continue + data = np.ascontiguousarray(table[peer, 1:]).tobytes() + error, pointer = cudart.cudaIpcOpenMemHandle( + _ipc_mem_handle_from_bytes(data), cudart.cudaIpcMemLazyEnablePeerAccess + ) + if error != cudart.cudaError_t.cudaSuccess: + status, detail = _FAILED, f"cudaIpcOpenMemHandle(rank {peer}) failed: {error}" + break + self._imports.append(int(pointer)) + peers.append(int(pointer)) + except Exception as error: # noqa: BLE001 - every failure must reach the agreement + status, detail = _FAILED, _describe(error) + table = self._dist.tp_allgather_int64([status]) + failed = [i for i, s in enumerate(table[:, 0]) if s != _OK] + if failed: + self._abandon( + RuntimeError, + f"MuseSparkOpen finalize+all-reduce peer mapping failed on rank(s) {failed}" + + (f": {detail}" if detail else ""), + ) + self._receive = self._export + self._peer_args = tuple(peers) + + def _initialize(self): + # Runs through _guarded: any exception becomes this rank's failure status. + with torch.cuda.device(self._device): + size = _BUFFERS * self._ranks * _SLAB_BYTES + (error,) = cuda.cuMemsetD32(self._receive, _SENTINEL, size // 4) + if error != cuda.CUresult.CUDA_SUCCESS: + return _FAILED, f"sentinel fill failed: {error}" + self._control = torch.zeros(_CONTROL_WORDS, dtype=torch.int32, device=self._device) + # Pointer arguments that stay fixed for the workspace's lifetime. + self._fixed = (self._control.data_ptr(), self._receive, *self._peer_args) + # A dry launch loads the module now, not inside a CUDA graph capture. + dummy = torch.zeros(_HIDDEN, dtype=torch.bfloat16, device=self._device).data_ptr() + self._launch(dummy, dummy, dummy, dummy, dummy, 0, _HIDDEN, 1, 1.0, 1) + torch.cuda.synchronize(self._device) + return _OK, "" + + def _launch(self, gemm, indices, weights, norm, out, rows, row_stride, tokens, eps, dry=0): + # TVM-FFI accepts device addresses as plain ints for pointer parameters. + self._kernel( + gemm, + indices, + weights, + norm, + out, + *self._fixed, + rows, + row_stride, + tokens, + self._rank, + eps, + dry, + cuda.CUstream(torch.cuda.current_stream(self._device).cuda_stream), + ) + + def finalize( + self, + gemm2_output: torch.Tensor, + indices: torch.Tensor, + routing_weights: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + num_tokens: int, + ) -> torch.Tensor: + """Enqueue the fused finalize + all-reduce on the current stream. + + Returns a new `[num_tokens, 4096]` BF16 tensor holding the full group + sum. Every rank must make the matching call with the same `num_tokens` + and `eps`. Requirements, checked before anything is launched: + + * `gemm2_output`: CUDA BF16 `[R, W]`, `W >= 4096`, unit column stride, + row stride a multiple of 8 and a 16-byte aligned base when `R > 0`. + Only rows addressed by valid slots and their first 4096 columns are + read. + * `indices`: contiguous int32 with `8 * num_tokens` entries; a slot is + valid when `0 <= index < R`. + * `routing_weights`: contiguous FP32 with `8 * num_tokens` entries. + * `norm_weight`: contiguous, 16-byte aligned BF16 with 4096 entries, + stored as offsets from one. + * `eps`: positive and finite. `num_tokens`: in `[1, 128]`. + + Raises `RuntimeError` once the workspace is closed. + """ + if self._closed: + raise RuntimeError("this MuseSparkOpen finalize workspace is closed") + rows, row_stride = self._validate( + gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens + ) + out = torch.empty((num_tokens, _HIDDEN), dtype=torch.bfloat16, device=self._device) + self._launch( + gemm2_output.data_ptr() if rows > 0 else 0, + indices.data_ptr(), + routing_weights.data_ptr(), + norm_weight.data_ptr(), + out.data_ptr(), + rows, + row_stride, + num_tokens, + float(eps), + ) + return out + + def _validate(self, gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens): + if isinstance(num_tokens, bool) or not isinstance(num_tokens, int): + raise ValueError(f"num_tokens must be an int, got {type(num_tokens).__name__}") + if not 1 <= num_tokens <= _MAX_TOKENS: + raise ValueError(f"num_tokens must be in [1, {_MAX_TOKENS}], got {num_tokens}") + if not (isinstance(eps, (int, float)) and math.isfinite(eps) and eps > 0): + raise ValueError(f"eps must be positive and finite, got {eps!r}") + if ( + gemm2_output.device != self._device + or gemm2_output.dtype != torch.bfloat16 + or gemm2_output.dim() != 2 + or gemm2_output.shape[1] < _HIDDEN + ): + raise ValueError( + "gemm2_output must be a [R, >=4096] BF16 tensor on the workspace device" + ) + rows = gemm2_output.shape[0] + row_stride = _HIDDEN + if rows > 0: + if gemm2_output.stride(1) != 1 or gemm2_output.data_ptr() % 16: + raise ValueError("gemm2_output needs unit column stride and a 16-byte aligned base") + if rows > 1: + row_stride = gemm2_output.stride(0) + if row_stride % _VECTOR: + raise ValueError("gemm2_output row stride must be a multiple of 8 elements") + expected = num_tokens * _TOP_K + for name, tensor, dtype in ( + ("indices", indices, torch.int32), + ("routing_weights", routing_weights, torch.float32), + ): + if ( + tensor.device != self._device + or tensor.dtype != dtype + or not tensor.is_contiguous() + or tensor.numel() != expected + ): + raise ValueError( + f"{name} must be a contiguous {dtype} tensor with {expected} entries" + ) + if ( + norm_weight.device != self._device + or norm_weight.dtype != torch.bfloat16 + or not norm_weight.is_contiguous() + or norm_weight.numel() != _HIDDEN + or norm_weight.data_ptr() % 16 + ): + raise ValueError( + "norm_weight must be a contiguous, 16-byte aligned BF16 tensor of 4096" + ) + return rows, row_stride + + # Release. + + def close(self) -> None: + """Release the workspace on every rank of the TP group. Collective. + + The owner first stops, on every rank, all `finalize` calls and all + replays of CUDA graphs that captured one. Then every rank calls + `close()` while the TP group's MPI communicator is still alive: + + 1. each rank synchronizes its device, and the group confirms that all did; + 2. each rank closes its mappings of peer memory, and the group confirms + that all did, so no rank maps any peer's buffer any more; + 3. each rank frees its own buffer, and the group confirms that all did. + + If any rank fails a step, every rank raises `RuntimeError` and nothing + a peer may still map is freed; it stays allocated until the process + exits. The workspace is unusable afterwards, and calling `close()` + again does nothing. + + Dropping the workspace without `close()` involves no other rank: it + waits for this rank's device and then closes this rank's mappings of + peer memory. It never frees this rank's own buffer, which peers may + still map; that stays allocated until the process exits. + """ + if self._closed: + return + self._closed = True + problem = self._release_together() + if problem: + raise RuntimeError(f"MuseSparkOpen finalize workspace close failed: {problem}") + + def _abandon(self, error_type: type, message: str) -> None: + """Undo a failed construction on every rank, then raise the agreed error.""" + problem = self._release_together() + self._closed = True + if problem: + message += f" (cleanup: {problem})" + raise error_type(message) + + def _release_together(self) -> str: + """Release this rank's resources in an order that is safe for every peer. + + Collective. Returns an empty string on success, or what failed. + """ + held = bool(self._export or self._imports) + failed = self._failed_ranks(self._sync() if held else True) + if failed: + self._retain() + return f"device synchronization failed on rank(s) {failed}" + failed = self._failed_ranks(self._close_imports()) + if failed: + self._retain() + return f"unmapping peer memory failed on rank(s) {failed}" + # No rank maps a peer's buffer any more, so every rank may free its own. + failed = self._failed_ranks(self._free_export()) + if failed: + return f"freeing workspace memory failed on rank(s) {failed}" + return "" + + def _failed_ranks(self, ok: bool) -> list[int]: + table = self._dist.tp_allgather_int64([int(ok)]) + return [rank for rank, value in enumerate(table[:, 0]) if not value] + + def _sync(self) -> bool: + try: + torch.cuda.synchronize(self._device) + except RuntimeError as error: + logger.warning(f"MuseSparkOpen finalize workspace release: {error}") + return False + return True + + def _close_imports(self) -> bool: + ok = True + for pointer in self._imports: + (error,) = cudart.cudaIpcCloseMemHandle(pointer) + if error != cudart.cudaError_t.cudaSuccess: + logger.warning( + f"MuseSparkOpen finalize workspace release: unmapping failed: {error}" + ) + ok = False + self._imports = [] + return ok + + def _free_export(self) -> bool: + if not self._export: + return True + (error,) = cudart.cudaFree(self._export) + if error != cudart.cudaError_t.cudaSuccess: + logger.warning(f"MuseSparkOpen finalize workspace release: cudaFree failed: {error}") + return False + self._export = 0 + return True + + def _retain(self) -> None: + """Keep everything a peer may still map allocated until the process exits.""" + self._export, self._imports = 0, [] + + def __del__(self): + # Local only, so it frees nothing that a peer may still map (see close()). + if self._closed or sys.is_finalizing(): + return + self._closed = True + if not (self._export or self._imports): + return + if not self._sync(): + self._retain() + return + self._close_imports() + if self._export: + logger.debug( + "MuseSparkOpen finalize workspace dropped without close(); its IPC buffer " + "stays allocated until the process exits" + ) + self._export = 0 diff --git a/tensorrt_llm/_torch/models/musespark_open_kernels.py b/tensorrt_llm/_torch/models/musespark_open_kernels.py new file mode 100644 index 000000000000..223c181c36e5 --- /dev/null +++ b/tensorrt_llm/_torch/models/musespark_open_kernels.py @@ -0,0 +1,598 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Triton kernels for the Muse Spark Open decoder's fused paths. + +Each kernel stands in for a short chain of elementwise and reduction ops in +`modeling_musespark_open`. The `*_supported` check next to each launcher states +the inputs the kernel takes; the caller keeps the reference composition for +everything else. The kernels round to the activation dtype wherever the +reference does, so they differ from it only in reduction order. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +# Warps per program, measured on GB300 at the checkpoint's widths with 1, 32 +# and 256 tokens: each is the fastest, or within noise of it, at every count. +_FINALIZE_NUM_WARPS = 8 +_RESIDUAL_NUM_WARPS = 16 + + +@triton.jit +def _post_expert_norm_combine_kernel( + out_ptr, # [num_tokens, HIDDEN] FP32 + gemm2_ptr, # [rows, >= HIDDEN], row-major + index_ptr, # [num_tokens * TOP_K]; -1 marks a slot whose expert is not local + weights_ptr, # [num_tokens * TOP_K] + norm_weight_ptr, # [HIDDEN], stored as an offset from 1 + gemm2_row_stride, + gemm2_num_rows, + eps, + TOP_K: tl.constexpr, + HIDDEN: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """One program per token: weight-first RMSNorm of each selected row, then the weighted sum. + + A slot whose index is negative or past the end of `gemm2` is never read: + the MoE kernel leaves those rows unwritten, and a NaN there must not reach + the sum. + """ + token = tl.program_id(0) + hoff = tl.arange(0, BLOCK_H) + koff = tl.arange(0, BLOCK_K) + in_row = hoff < HIDDEN + in_slot = koff < TOP_K + + gain = tl.load(norm_weight_ptr + hoff, mask=in_row, other=0.0).to(tl.float32) + 1.0 + index = tl.load(index_ptr + token * TOP_K + koff, mask=in_slot, other=-1) + local = in_slot & (index >= 0) & (index < gemm2_num_rows) + safe = tl.where(local, index, 0).to(tl.int64) + + rows = tl.load( + gemm2_ptr + safe[:, None] * gemm2_row_stride + hoff[None, :], + mask=local[:, None] & in_row[None, :], + other=0.0, + ).to(tl.float32) + scaled = rows * gain[None, :] + variance = tl.sum(scaled * scaled, axis=1) / HIDDEN + normed = scaled * tl.rsqrt(variance + eps)[:, None] + weights = tl.load(weights_ptr + token * TOP_K + koff, mask=in_slot, other=0.0).to(tl.float32) + acc = tl.sum(tl.where(local[:, None], normed * weights[:, None], 0.0), axis=0) + tl.store(out_ptr + token * HIDDEN + hoff, acc, mask=in_row) + + +def post_expert_norm_combine_supported(gemm2_output: torch.Tensor, num_tokens: int) -> bool: + """Whether `post_expert_norm_combine` takes these inputs: CUDA, row-major, some tokens.""" + return gemm2_output.is_cuda and num_tokens > 0 and gemm2_output.stride(-1) == 1 + + +def post_expert_norm_combine( + gemm2_output: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + routing_weights: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + *, + num_tokens: int, + top_k: int, + hidden_size: int, +) -> torch.Tensor: + """The per-expert post-norm and the top-k combine, in FP32, as one kernel.""" + out = torch.empty((num_tokens, hidden_size), dtype=torch.float32, device=gemm2_output.device) + _post_expert_norm_combine_kernel[(num_tokens,)]( + out, + gemm2_output, + expanded_idx_to_permuted_idx.reshape(-1), + routing_weights.reshape(-1), + norm_weight, + gemm2_output.stride(0), + gemm2_output.shape[0], + eps, + TOP_K=top_k, + HIDDEN=hidden_size, + BLOCK_K=triton.next_power_of_2(top_k), + BLOCK_H=triton.next_power_of_2(hidden_size), + num_warps=_FINALIZE_NUM_WARPS, + ) + return out + + +@triton.jit +def _gated_residual_after_rmsnorm_kernel( + out_ptr, # [T, HIDDEN] FP32, the updated residual stream + residual_ptr, # [T, HIDDEN] FP32 + branch_ptr, # [T, HIDDEN], activation dtype + gain_ptr, # [HIDDEN] `1 + w`; read only with HAS_GAIN + alpha_ptr, # [HIDDEN] FP32 + beta_ptr, # [HIDDEN] FP32 + eps, + next_out_ptr, # [T, HIDDEN], activation dtype; written only with FUSE_NEXT_NORM + next_weight_ptr, # [HIDDEN]; read only with FUSE_NEXT_NORM + next_eps, + HIDDEN: tl.constexpr, + BLOCK_H: tl.constexpr, + HAS_GAIN: tl.constexpr, + FUSE_NEXT_NORM: tl.constexpr, +): + """`alpha * residual + beta * rmsnorm(gain * branch)` for one token. + + The gain product and the normalized branch are each rounded to the + activation dtype, as in the reference. With `FUSE_NEXT_NORM`, the updated + residual is also rounded to the activation dtype and passed through the + next Gemma RMSNorm, which applies `1 + w` after normalizing. + """ + token = tl.program_id(0) + offs = tl.arange(0, BLOCK_H) + mask = offs < HIDDEN + base = token * HIDDEN + offs + + branch = tl.load(branch_ptr + base, mask=mask, other=0.0) + scaled = branch.to(tl.float32) + if HAS_GAIN: + gain = tl.load(gain_ptr + offs, mask=mask, other=0.0).to(tl.float32) + scaled = (scaled * gain).to(branch.dtype).to(tl.float32) + variance = tl.sum(scaled * scaled, axis=0) / HIDDEN + normed = (scaled * tl.rsqrt(variance + eps)).to(branch.dtype).to(tl.float32) + + residual = tl.load(residual_ptr + base, mask=mask, other=0.0).to(tl.float32) + alpha = tl.load(alpha_ptr + offs, mask=mask, other=0.0) + beta = tl.load(beta_ptr + offs, mask=mask, other=0.0) + updated = alpha * residual + beta * normed + tl.store(out_ptr + base, updated, mask=mask) + if FUSE_NEXT_NORM: + next_input = updated.to(branch.dtype).to(tl.float32) + next_variance = tl.sum(next_input * next_input, axis=0) / HIDDEN + next_weight = tl.load(next_weight_ptr + offs, mask=mask, other=0.0).to(tl.float32) + next_output = next_input * tl.rsqrt(next_variance + next_eps) + next_output = next_output * (next_weight + 1.0) + tl.store(next_out_ptr + base, next_output, mask=mask) + + +def gated_residual_supported( + residual: torch.Tensor, branch: torch.Tensor, gain: Optional[torch.Tensor] +) -> bool: + """Whether `gated_residual_after_rmsnorm` takes these inputs.""" + return ( + residual.is_cuda + and residual.ndim == 2 + and residual.shape == branch.shape + and residual.shape[0] > 0 + and residual.is_contiguous() + and branch.is_contiguous() + # A gain of another dtype makes the reference promote `gain * branch` + # and skip the round to the activation dtype that the kernel takes. + and (gain is None or (gain.dtype == branch.dtype and gain.is_contiguous())) + ) + + +def gated_residual_after_rmsnorm( + residual: torch.Tensor, + branch: torch.Tensor, + gain: Optional[torch.Tensor], + alpha: torch.Tensor, + beta: torch.Tensor, + eps: float, +) -> torch.Tensor: + """One residual join: the FP32 `alpha * residual + beta * rmsnorm(gain * branch)`.""" + out = torch.empty_like(residual, dtype=torch.float32) + hidden_size = residual.shape[-1] + _gated_residual_after_rmsnorm_kernel[(residual.shape[0],)]( + out, + residual, + branch, + # Unread without a gain, but a pointer argument still needs a tensor. + alpha if gain is None else gain, + alpha, + beta, + eps, + out, + alpha, + eps, + HIDDEN=hidden_size, + BLOCK_H=triton.next_power_of_2(hidden_size), + HAS_GAIN=gain is not None, + FUSE_NEXT_NORM=False, + num_warps=_RESIDUAL_NUM_WARPS, + ) + return out + + +def gated_residual_and_norm_supported( + residual: torch.Tensor, + branch: torch.Tensor, + gain: Optional[torch.Tensor], + next_weight: torch.Tensor, +) -> bool: + """Whether `gated_residual_after_rmsnorm_and_norm` takes these inputs.""" + return ( + residual.dtype == torch.float32 + and branch.dtype == torch.bfloat16 + and gated_residual_supported(residual, branch, gain) + and next_weight.dtype == branch.dtype + and next_weight.is_contiguous() + ) + + +def gated_residual_after_rmsnorm_and_norm( + residual: torch.Tensor, + branch: torch.Tensor, + gain: Optional[torch.Tensor], + alpha: torch.Tensor, + beta: torch.Tensor, + eps: float, + next_weight: torch.Tensor, + next_eps: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + """A residual join and the next Gemma RMSNorm; returns `(residual, next_input)`.""" + updated = torch.empty_like(residual) + next_input = torch.empty_like(branch) + hidden_size = residual.shape[-1] + _gated_residual_after_rmsnorm_kernel[(residual.shape[0],)]( + updated, + residual, + branch, + alpha if gain is None else gain, + alpha, + beta, + eps, + next_input, + next_weight, + next_eps, + HIDDEN=hidden_size, + BLOCK_H=triton.next_power_of_2(hidden_size), + HAS_GAIN=gain is not None, + FUSE_NEXT_NORM=True, + num_warps=_RESIDUAL_NUM_WARPS, + ) + return updated, next_input + + +@triton.jit +def _routing_weights_kernel( + scores_ptr, + indices_ptr, + selected_ptr, + weights_ptr, + score_stride, + eps, + TOP_K: tl.constexpr, + BLOCK_K: tl.constexpr, +): + row = tl.program_id(0) + slots = tl.arange(0, BLOCK_K) + mask = slots < TOP_K + indices = tl.load(indices_ptr + row * TOP_K + slots, mask=mask, other=0) + scores = tl.load(scores_ptr + row * score_stride + indices, mask=mask, other=0.0) + denominator = tl.sum(scores, axis=0) + eps + weights = tl.div_rn(scores, denominator) + tl.store(selected_ptr + row * TOP_K + slots, indices.to(tl.int32), mask=mask) + tl.store(weights_ptr + row * TOP_K + slots, weights, mask=mask) + + +def routing_weights_supported( + scores: torch.Tensor, topk_idx: torch.Tensor, output_dtype: torch.dtype +) -> bool: + """Whether `routing_weights` takes these inputs: FP32 scores and int64 top-8 ids on CUDA.""" + return ( + scores.is_cuda + and scores.dtype == torch.float32 + and scores.ndim == 2 + and scores.is_contiguous() + and scores.shape[0] > 0 + and topk_idx.is_cuda + and topk_idx.dtype == torch.int64 + and topk_idx.ndim == 2 + and topk_idx.is_contiguous() + and topk_idx.shape == (scores.shape[0], 8) + and output_dtype == torch.float32 + ) + + +def routing_weights( + scores: torch.Tensor, topk_idx: torch.Tensor, eps: float +) -> Tuple[torch.Tensor, torch.Tensor]: + """The int32 top-8 ids and their scores, renormalized to sum to one.""" + selected = torch.empty_like(topk_idx, dtype=torch.int32) + weights = torch.empty_like(topk_idx, dtype=torch.float32) + _routing_weights_kernel[(scores.shape[0],)]( + scores, + topk_idx, + selected, + weights, + scores.stride(0), + eps, + TOP_K=8, + BLOCK_K=8, + num_warps=1, + ) + return selected, weights + + +@triton.jit +def _router_projection_kernel( + hidden_states, + weight, + output, + scale, + EXPERTS: tl.constexpr, + HIDDEN: tl.constexpr, + ROWS: tl.constexpr, +): + token = tl.program_id(0) + expert = tl.program_id(1) * ROWS + tl.arange(0, ROWS) + offsets = tl.arange(0, HIDDEN) + x = tl.load(hidden_states + token * HIDDEN + offsets).to(tl.float32) + w = tl.load( + weight + expert[:, None] * HIDDEN + offsets[None, :], + expert[:, None] < EXPERTS, + other=0.0, + ) + logits = tl.sum(w * x[None, :], axis=1) * scale + tl.store(output + token * EXPERTS + expert, logits, expert < EXPERTS) + + +@triton.jit +def _router_tensorcore_kernel( + hidden_states, + weight, + partials, + TOKENS: tl.constexpr, +): + # Sixteen tokens share each weight tile. The 8192-wide dot product is split + # into sixteen 512-wide chunks whose FP32 partials the reduce kernel sums; + # TF32x3 keeps the FP32 weights' precision. + tokens = tl.program_id(0) * 16 + tl.arange(0, 16) + experts = tl.program_id(1) * 32 + tl.arange(0, 32) + split = tl.program_id(2) + acc = tl.zeros((16, 32), tl.float32) + for block in range(4): + hidden = split * 512 + block * 128 + tl.arange(0, 128) + x = tl.load( + hidden_states + tokens[:, None] * 8192 + hidden[None, :], + tokens[:, None] < TOKENS, + other=0.0, + ).to(tl.float32) + w = tl.load(weight + experts[None, :] * 8192 + hidden[:, None]) + acc = tl.dot(x, w, acc, input_precision="tf32x3") + tl.store( + partials + split * TOKENS * 256 + tokens[:, None] * 256 + experts[None, :], + acc, + tokens[:, None] < TOKENS, + ) + + +@triton.jit +def _router_tensorcore_reduce_kernel(partials, output, scale, SIZE: tl.constexpr): + offsets = tl.program_id(0) * 128 + tl.arange(0, 128) + splits = tl.arange(0, 16) + values = tl.load( + partials + splits[:, None] * SIZE + offsets[None, :], + offsets[None, :] < SIZE, + other=0.0, + ) + tl.store(output + offsets, tl.sum(values, axis=0) * scale, offsets < SIZE) + + +def router_projection_supported(hidden_states: torch.Tensor, weight: torch.Tensor) -> bool: + """Whether `router_projection` takes these inputs: up to 128 BF16 tokens of width 8192.""" + return ( + hidden_states.is_cuda + and hidden_states.dtype == torch.bfloat16 + and hidden_states.ndim == 2 + and 0 < hidden_states.shape[0] <= 128 + and hidden_states.shape[1] == 8192 + and hidden_states.is_contiguous() + and weight.device == hidden_states.device + and weight.dtype == torch.float32 + and weight.shape == (256, 8192) + and weight.is_contiguous() + ) + + +def _router_launch(tokens: int, overlapped: bool) -> Tuple[int, int]: + """(experts per program, warps) of the CUDA-core router, measured on GB300. + + `overlapped` is the shape measured while the pooler runs on another stream. + """ + if overlapped: + if tokens in (3, 5): + return 1, 16 + if 4 <= tokens <= 11: + return 2, 16 + if tokens in (13, 15): + return 2, 8 + if tokens == 14: + return 4, 8 + return (1, 16) if tokens <= 2 else (2, 4) + + +def router_projection( + hidden_states: torch.Tensor, weight: torch.Tensor, scale: float, *, overlapped: bool +) -> torch.Tensor: + """`scale * hidden_states @ weight.T` with FP32 weights and accumulation. + + Fewer than 16 tokens use CUDA-core dot products; 16 or more share weight + tiles on the tensor cores. + """ + tokens = hidden_states.shape[0] + output = torch.empty((tokens, 256), device=hidden_states.device, dtype=torch.float32) + if tokens >= 16: + partials = torch.empty((16, tokens, 256), device=hidden_states.device, dtype=torch.float32) + _router_tensorcore_kernel[(triton.cdiv(tokens, 16), 8, 16)]( + hidden_states, weight, partials, TOKENS=tokens, num_warps=4 + ) + _router_tensorcore_reduce_kernel[(triton.cdiv(tokens * 256, 128),)]( + partials, output, scale, SIZE=tokens * 256, num_warps=4 + ) + return output + rows, warps = _router_launch(tokens, overlapped) + _router_projection_kernel[(tokens, 256 // rows)]( + hidden_states, + weight, + output, + scale, + EXPERTS=256, + HIDDEN=8192, + ROWS=rows, + num_warps=warps, + enable_fp_fusion=False, + ) + return output + + +@triton.jit +def _full_routing_kernel(logits, bias, selected, weights, eps, EXPERTS: tl.constexpr): + """Top-8 of `sigmoid(logits) + bias`, weighted by the renormalized unbiased scores. + + The picks come out in descending order and exact ties go to the lowest + expert id, so the order in which the finalize sums the experts is fixed. + """ + # The `255 - expert` tie-break key needs at most 256 experts. + tl.static_assert(EXPERTS <= 256) + row = tl.program_id(0) + expert = tl.arange(0, EXPERTS) + scores = tl.sigmoid(tl.load(logits + row * EXPERTS + expert)) + ranked = scores + tl.load(bias + expert) + # Both signed zeros must sort as one value before the id breaks the tie. + ranked = tl.where(ranked == 0.0, 0.0, ranked) + bits = ranked.to(tl.uint32, bitcast=True) + ordered = tl.where((bits >> 31) != 0, ~bits, bits ^ 0x80000000).to(tl.uint32) + keys = (ordered.to(tl.uint64) << 32) | (255 - expert).to(tl.uint64) + sorted_keys = tl.sort(keys, descending=True) + ids = 255 - (sorted_keys & 0xFFFFFFFF).to(tl.int32) + picked = tl.gather(scores, ids, axis=0) + picked = tl.where(expert < 8, picked, 0.0) + normalized = tl.div_rn(picked, tl.sum(picked, axis=0) + eps) + tl.store(selected + row * 8 + expert, ids, expert < 8) + tl.store(weights + row * 8 + expert, normalized, expert < 8) + + +def full_routing_supported( + router_logits: torch.Tensor, + bias: Optional[torch.Tensor], + top_k: int, + output_dtype: torch.dtype, +) -> bool: + """Whether `full_routing` takes these inputs: up to 128 tokens over 256 experts, top-8.""" + return ( + router_logits.is_cuda + and router_logits.dtype == torch.float32 + and router_logits.ndim == 2 + and router_logits.is_contiguous() + and 0 < router_logits.shape[0] <= 128 + and router_logits.shape[1] == 256 + and top_k == 8 + and output_dtype == torch.float32 + and bias is not None + and bias.is_cuda + and bias.dtype == torch.float32 + and bias.shape == (256,) + and bias.is_contiguous() + ) + + +def full_routing( + router_logits: torch.Tensor, bias: torch.Tensor, eps: float +) -> Tuple[torch.Tensor, torch.Tensor]: + """The whole routing step for 256 experts and top-8; returns int32 ids and FP32 weights.""" + tokens = router_logits.shape[0] + selected = torch.empty((tokens, 8), device=router_logits.device, dtype=torch.int32) + weights = torch.empty((tokens, 8), device=router_logits.device, dtype=torch.float32) + _full_routing_kernel[(tokens,)]( + router_logits, bias, selected, weights, eps, EXPERTS=256, num_warps=4 + ) + return selected, weights + + +@triton.jit +def _attention_norm_gate_kernel( + out, + attention, + gate, + attention_stride, + gate_token_stride, + gate_head_stride, + eps, + HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + """Weightless RMSNorm of one head of the attention output, times `sigmoid(gate)`. + + The norm and the sigmoid are each rounded to BF16 before the product, as + the reference's separate ops do. + """ + row = tl.program_id(0) + token = row // HEADS + head = row % HEADS + offsets = tl.arange(0, HEAD_DIM) + x = tl.load(attention + token * attention_stride + head * HEAD_DIM + offsets).to(tl.float32) + variance = tl.sum(x * x, axis=0) / HEAD_DIM + normed = (x * tl.rsqrt(variance + eps)).to(tl.bfloat16) + g = tl.load(gate + token * gate_token_stride + head * gate_head_stride + offsets).to(tl.float32) + sigmoid = tl.sigmoid(g).to(tl.bfloat16) + result = normed.to(tl.float32) * sigmoid.to(tl.float32) + tl.store(out + row * HEAD_DIM + offsets, result) + + +def attention_norm_gate_supported( + attention_output: torch.Tensor, gate: torch.Tensor, head_dim: int +) -> bool: + """Whether `attention_norm_gate` takes these inputs: up to 128 BF16 tokens, 64-wide heads.""" + return ( + attention_output.is_cuda + and attention_output.dtype == torch.bfloat16 + and attention_output.ndim == 2 + and attention_output.stride(-1) == 1 + and 0 < attention_output.shape[0] <= 128 + and head_dim == 64 + and attention_output.shape[1] % head_dim == 0 + and gate.is_cuda + and gate.dtype == torch.bfloat16 + and gate.ndim in (2, 3) + and gate.stride(-1) == 1 + and gate.shape[0] == attention_output.shape[0] + and gate.numel() == attention_output.numel() + and (gate.ndim == 2 or gate.shape[-1] == head_dim) + ) + + +def attention_norm_gate( + attention_output: torch.Tensor, gate: torch.Tensor, head_dim: int, eps: float +) -> torch.Tensor: + """Per-head weightless RMSNorm of `[tokens, heads * head_dim]`, gated by `sigmoid(gate)`. + + `gate` is `[tokens, heads * head_dim]` or `[tokens, heads, head_dim]`, with + any token and head strides. + """ + output = torch.empty_like(attention_output, memory_format=torch.contiguous_format) + heads = attention_output.shape[1] // head_dim + _attention_norm_gate_kernel[(attention_output.shape[0] * heads,)]( + output, + attention_output, + gate, + attention_output.stride(0), + gate.stride(0), + gate.stride(1) if gate.ndim == 3 else head_dim, + eps, + HEADS=heads, + HEAD_DIM=head_dim, + num_warps=4, + ) + return output diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index 1de8159db276..601c6d9b2271 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -96,6 +96,17 @@ l0_gb300: - unittest/_torch/modules/test_mxfp8_draft_head.py - unittest/_torch/modules/test_q35_mtp_glue.py - unittest/_torch/pyexecutor/test_steady_spec_prepare.py + # Muse Spark Open decoder fast paths on one GPU: the fused Triton kernels and the + # GB300 cluster boundary kernel, against their reference compositions. + - unittest/_torch/modeling/test_musespark_open_router_projection.py + - unittest/_torch/modeling/test_musespark_open_full_routing.py + - unittest/_torch/modeling/test_musespark_open_norm_gate.py + - unittest/_torch/modeling/test_musespark_open_moe_finalize.py + - unittest/_torch/modeling/test_musespark_open_pre_expert_quant.py + - unittest/_torch/modeling/test_musespark_open_moe_prepare_overlap.py + - unittest/_torch/modeling/test_musespark_open_boundary_fusion.py + - unittest/_torch/modeling/test_musespark_open_cluster_boundary.py + - unittest/_torch/modeling/test_musespark_open_finalize_local.py - condition: ranges: system_gpu_count: diff --git a/tests/unittest/_torch/modeling/test_modeling_musespark_open.py b/tests/unittest/_torch/modeling/test_modeling_musespark_open.py index 7c387b68c02b..30ae5b459980 100644 --- a/tests/unittest/_torch/modeling/test_modeling_musespark_open.py +++ b/tests/unittest/_torch/modeling/test_modeling_musespark_open.py @@ -802,35 +802,29 @@ def test_missing_in_range_keys_reports_a_dropped_tensor(): # Residual join # # A decoder layer joins each branch as `gated_add(residual, rmsnorm(branch))` -# over the FP32 residual stream. +# over the FP32 residual stream. On CUDA the join is a fused kernel, checked +# against this composition in test_musespark_open_boundary_fusion.py. # -------------------------------------------------------------------------- -@pytest.mark.parametrize("device", ["cpu", DEVICE]) -def test_residual_join_is_the_gated_add_of_the_normed_branch(device): - """The join equals the gated add of the normalized branch, to the last bit. +def test_residual_join_is_the_gated_add_of_the_normed_branch(): + """Off CUDA, the join equals the gated add of the normalized branch, to the last bit. Checked on both sides of the layer: the weight-first FFN norm and the weightless attention norm. """ hidden, eps, tokens = 256, 1e-6, 8 gen = torch.Generator().manual_seed(3) - residual = torch.randn(tokens, hidden, generator=gen, dtype=torch.float32).to(device) - branch = ( - torch.randn(tokens, hidden, generator=gen, dtype=torch.float32) - .to(torch.bfloat16) - .to(device) - ) + residual = torch.randn(tokens, hidden, generator=gen, dtype=torch.float32) + branch = torch.randn(tokens, hidden, generator=gen, dtype=torch.float32).to(torch.bfloat16) gate = MuseSparkOpenGatedResidual(hidden, torch.bfloat16) with torch.no_grad(): gate.gate.normal_(0, 0.5, generator=gen) - gate = gate.to(device) gate.cache_derived_state() norm = MuseSparkOpenRMSNormWeightFirst(hidden, eps, torch.bfloat16) with torch.no_grad(): norm.weight.normal_(0, 0.3, generator=gen) - norm = norm.to(device) norm.cache_derived_state() def weightless_norm(x): diff --git a/tests/unittest/_torch/modeling/test_musespark_open_boundary_fusion.py b/tests/unittest/_torch/modeling/test_musespark_open_boundary_fusion.py new file mode 100644 index 000000000000..77b028a6426e --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_boundary_fusion.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused residual joins, alone and with the next prenorm: rounding, graphs, layer carry.""" + +import contextlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_musespark_open import ( + MuseSparkOpenDecoderLayer, + MuseSparkOpenGatedResidual, + MuseSparkOpenModel, + MuseSparkOpenRMSNormWeightFirst, +) +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _case(tokens, hidden): + torch.manual_seed(719) + residual = torch.randn(tokens, hidden, device="cuda", dtype=torch.float32) + branch = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + gate = MuseSparkOpenGatedResidual(hidden, torch.bfloat16).cuda() + norm = RMSNorm(hidden_size=hidden, eps=1e-5, dtype=torch.bfloat16, use_gemma=True).cuda() + with torch.no_grad(): + gate.gate.normal_(0, 0.5) + norm.weight.normal_(0, 0.02) + gate.cache_derived_state() + return residual, branch, gate, norm + + +@pytest.mark.parametrize("with_gain", [False, True]) +@pytest.mark.parametrize("tokens", [1, 2, 16, 128, 8192]) +@pytest.mark.parametrize("hidden", [256, 8192]) +def test_boundary_fusion_preserves_residual_and_prenorm(tokens, hidden, with_gain): + with torch.inference_mode(): + residual, branch, gate, norm = _case(tokens, hidden) + gain = ( + (torch.randn(hidden, device="cuda", dtype=torch.bfloat16) * 0.02 + 1) + if with_gain + else None + ) + expected_residual = gate.add_after_rmsnorm(residual, branch, 1e-5, gain) + expected_norm = norm(expected_residual.to(torch.bfloat16)) + actual_residual, actual_norm = gate.add_after_rmsnorm_and_norm( + residual, branch, 1e-5, norm, gain + ) + torch.testing.assert_close(actual_residual, expected_residual, rtol=2e-5, atol=2e-5) + torch.testing.assert_close(actual_norm, expected_norm, rtol=0.016, atol=0.001) + assert actual_residual.dtype == torch.float32 + assert actual_norm.dtype == torch.bfloat16 + relative_l2 = ( + actual_norm.float() - expected_norm.float() + ).norm() / expected_norm.float().norm() + assert relative_l2 < 1e-3 + + +# The plain residual join: `gated_add(residual, rmsnorm(branch))`, checked +# against the composition it replaces, including the real `RMSNorm` module on +# the attention side. + + +def _join_case(tokens, hidden, dtype, seed=0): + gen = torch.Generator(device="cuda").manual_seed(seed) + residual = torch.randn(tokens, hidden, generator=gen, dtype=torch.float32, device="cuda") + branch = torch.randn(tokens, hidden, generator=gen, dtype=torch.float32, device="cuda").to( + dtype + ) + gate = MuseSparkOpenGatedResidual(hidden, dtype).cuda() + with torch.no_grad(): + gate.gate.normal_(0, 0.5, generator=gen) + gate.cache_derived_state() + return residual, branch, gate + + +def _weight_first_norm(hidden, eps, dtype): + norm = MuseSparkOpenRMSNormWeightFirst(hidden, eps, dtype).cuda() + with torch.no_grad(): + norm.weight.normal_(0, 0.3) + norm.cache_derived_state() + return norm + + +def _weightless_norm(eps): + def norm(x): + x = x.to(torch.float32) + return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)).to(torch.bfloat16) + + return norm + + +@pytest.mark.parametrize("tokens", [1, 5, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fused_join_matches_the_weightless_norm_path(tokens, dtype): + """Attention side: `add_attn(h, post_attn_norm(y))` with a real weightless `RMSNorm`.""" + hidden, eps = 256, 1e-6 + residual, branch, gate = _join_case(tokens, hidden, dtype) + reference_norm = RMSNorm(hidden_size=hidden, eps=eps, dtype=dtype, has_weights=False).cuda() + + fused = gate.add_after_rmsnorm(residual, branch, eps) + reference = gate(residual, reference_norm(branch)) + + assert fused.dtype is torch.float32 + assert (fused - reference).abs().max() < 1e-4 * reference.abs().max() + torch.testing.assert_close(fused, reference, rtol=2e-5, atol=2e-5) + + +@pytest.mark.parametrize("tokens", [1, 5, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fused_join_matches_the_weight_first_norm_path(tokens, dtype): + """FFN side, with the cached `1 + w` passed as the gain, as the decoder layer does.""" + hidden, eps = 256, 1e-6 + residual, branch, gate = _join_case(tokens, hidden, dtype) + norm = _weight_first_norm(hidden, eps, dtype) + + fused = gate.add_after_rmsnorm(residual, branch, eps, norm.gain()) + reference = gate(residual, norm(branch)) + + assert (fused - reference).abs().max() < 1e-4 * reference.abs().max() + torch.testing.assert_close(fused, reference, rtol=2e-5, atol=2e-5) + # Dropping the gain is a different function, so a kernel ignoring it fails here. + assert not torch.allclose(fused, gate.add_after_rmsnorm(residual, branch, eps), atol=1e-2) + + +def test_a_gain_of_another_dtype_takes_the_composition(): + """The composition promotes `gain * branch` across dtypes and skips a round the kernel takes. + + Bit-exact equality with the composition shows the kernel did not run: it + reduces the variance in another order. + """ + hidden, eps, tokens = 128, 1e-6, 4 + residual, branch, gate = _join_case(tokens, hidden, torch.float16) + gain = torch.empty(hidden, dtype=torch.float32, device="cuda").normal_(0, 0.3) + 1.0 + + got = gate.add_after_rmsnorm(residual, branch, eps, gain) + + scaled = (gain * branch).to(torch.float32) + normed = (scaled * torch.rsqrt(scaled.pow(2).mean(-1, keepdim=True) + eps)).to(branch.dtype) + torch.testing.assert_close(got, gate(residual, normed), atol=0, rtol=0) + + +def test_the_fused_join_gives_the_same_answer_with_or_without_the_gain_cache(): + hidden, eps, tokens = 128, 1e-6, 4 + residual, branch, gate = _join_case(tokens, hidden, torch.bfloat16) + norm = MuseSparkOpenRMSNormWeightFirst(hidden, eps, torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.normal_(0, 0.3) + + assert norm._weight_plus_one is None + uncached = gate.add_after_rmsnorm(residual, branch, eps, norm.gain()) + norm.cache_derived_state() + cached = gate.add_after_rmsnorm(residual, branch, eps, norm.gain()) + + torch.testing.assert_close(cached, uncached, atol=0, rtol=0) + + +def test_the_fused_join_replays_from_a_cuda_graph(): + hidden, eps, tokens = 256, 1e-6, 8 + residual, branch, gate = _join_case(tokens, hidden, torch.bfloat16) + norm = _weight_first_norm(hidden, eps, torch.bfloat16) + + gate.add_after_rmsnorm(residual, branch, eps, norm.gain()) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = gate.add_after_rmsnorm(residual, branch, eps, norm.gain()) + + fresh_residual, fresh_branch, _ = _join_case(tokens, hidden, torch.bfloat16, seed=7) + residual.copy_(fresh_residual) + branch.copy_(fresh_branch) + graph.replay() + torch.cuda.synchronize() + + assert torch.isfinite(static_out).all() + torch.testing.assert_close(static_out, gate(residual, norm(branch)), rtol=2e-5, atol=2e-5) + + +@pytest.mark.parametrize("tokens", [1, 4]) +def test_the_fused_join_differs_from_the_composition_only_by_reduction_order(tokens): + """At the checkpoint's width, within a bound only the variance's reduction order can explain. + + The kernel keeps both of the composition's rounds to the activation dtype + (the gain product and the normalized branch); one that dropped either would + land further from the composition than this bound. On larger inputs the + other reduction order eventually moves a value across a BF16 rounding + boundary, which the tolerance tests above allow for. + """ + hidden, eps = 8192, 1e-6 + residual, branch, gate = _join_case(tokens, hidden, torch.bfloat16, seed=3) + norm = _weight_first_norm(hidden, eps, torch.bfloat16) + + for gain, reference_norm in ((norm.gain(), norm), (None, _weightless_norm(eps))): + fused = gate.add_after_rmsnorm(residual, branch, eps, gain) + unfused = gate(residual, reference_norm(branch)) + assert torch.isfinite(fused).all() + assert (fused - unfused).abs().max() < 1e-6 * unfused.abs().max() + + +def test_boundary_fusion_graph_replay(): + with torch.inference_mode(): + residual, branch, gate, norm = _case(2, 8192) + gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm) + residual.normal_() + branch.normal_() + expected = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm) + graph.replay() + torch.cuda.synchronize() + for got, want in zip(actual, expected): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +def _stub_layer(layer_idx, hidden): + layer = MuseSparkOpenDecoderLayer.__new__(MuseSparkOpenDecoderLayer) + torch.nn.Module.__init__(layer) + layer.layer_idx = layer_idx + layer.dtype = torch.bfloat16 + layer.nvtx = lambda name: contextlib.nullcontext() + layer.pre_qkv_proj_norm = RMSNorm( + hidden_size=hidden, eps=1e-5, dtype=torch.bfloat16, use_gemma=True + ).cuda() + layer.pre_ffn_norm = RMSNorm( + hidden_size=hidden, eps=1e-5, dtype=torch.bfloat16, use_gemma=True + ).cuda() + layer.post_attn_norm = RMSNorm( + hidden_size=hidden, eps=1e-5, dtype=torch.bfloat16, has_weights=False + ).cuda() + layer.post_ffn_norm = MuseSparkOpenRMSNormWeightFirst(hidden, 1e-5, torch.bfloat16).cuda() + layer.add_attn = MuseSparkOpenGatedResidual(hidden, torch.bfloat16).cuda() + layer.add_ffn = MuseSparkOpenGatedResidual(hidden, torch.bfloat16).cuda() + layer.self_attn = lambda hidden_states, **kwargs: hidden_states * 0.7 + layer.feed_forward = lambda x, metadata: x * 0.9 + for param in layer.parameters(): + param.normal_(0, 0.03) + return layer + + +@pytest.mark.parametrize("tokens", [1, 2, 16]) +def test_three_layer_boundary_carry_and_graph(tokens): + """Each layer hands the next its attention input; the result matches layer by layer.""" + with torch.inference_mode(): + torch.manual_seed(951) + hidden = 256 + model = MuseSparkOpenModel.__new__(MuseSparkOpenModel) + torch.nn.Module.__init__(model) + model.dtype = torch.bfloat16 + model.m_emb = 1.0 + model.embed_norm = RMSNorm( + hidden_size=hidden, eps=1e-5, dtype=torch.bfloat16, has_weights=False + ).cuda() + model.norm = RMSNorm(hidden_size=hidden, eps=1e-5, dtype=torch.bfloat16).cuda() + model.layers = torch.nn.ModuleList([_stub_layer(index, hidden) for index in range(3)]) + for module in model.modules(): + if hasattr(module, "cache_derived_state"): + module.cache_derived_state() + x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + residual = model.embed_norm(x).to(torch.float32) + for layer in model.layers: + residual = layer(position_ids=None, hidden_states=residual, attn_metadata=None) + expected = model.norm(residual.to(torch.bfloat16)) + + captures = [] + metadata = SimpleNamespace( + maybe_capture_hidden_states=lambda idx, state, residual: captures.append( + (idx, state, residual) + ) + ) + got = model(attn_metadata=None, inputs_embeds=x, spec_metadata=metadata) + # Drafters capture every layer's FP32 residual stream. + assert [entry[0] for entry in captures] == [0, 1, 2] + assert all( + state.dtype == torch.float32 and residual is None for _, state, residual in captures + ) + torch.testing.assert_close(got, expected, rtol=0.04, atol=0.004) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + got = model(attn_metadata=None, inputs_embeds=x) + x.normal_() + expected = model(attn_metadata=None, inputs_embeds=x) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(got, expected, rtol=0, atol=0) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py b/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py new file mode 100644 index 000000000000..0e57f08dd020 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cluster boundary rounding, live graph inputs, and bounded dispatch.""" + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenGatedResidual +from tensorrt_llm._torch.models.musespark_open_boundary import cluster_boundary_available +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not cluster_boundary_available(torch.device("cuda", 0)), + reason="GB300 and CuTe DSL required", +) + + +@pytest.mark.parametrize("tokens", [1, 2, 3, 4, 8, 16, 17, 32]) +@pytest.mark.parametrize("with_gain", [False, True]) +@pytest.mark.parametrize("eps", [1e-6, 1e-5]) +@torch.inference_mode() +def test_cluster_boundary_rounding_and_mutable_graph( + tokens: int, with_gain: bool, eps: float, monkeypatch: pytest.MonkeyPatch +) -> None: + torch.manual_seed(1294 + tokens) + residual = torch.randn(tokens, 8192, device="cuda", dtype=torch.float32) + branch = torch.randn(tokens, 8192, device="cuda", dtype=torch.bfloat16) + gate = MuseSparkOpenGatedResidual(8192, torch.bfloat16).cuda() + gate.gate.normal_(0, 0.5) + gate.cache_derived_state() + norm = RMSNorm(hidden_size=8192, eps=eps, dtype=torch.bfloat16, use_gemma=True).cuda() + norm.weight.normal_(0, 0.02) + gain = torch.empty(8192, device="cuda", dtype=torch.bfloat16) if with_gain else None + if gain is not None: + gain.normal_(1, 0.1) + + # Compile on the eager stream, then capture on the graph stream. + gate.add_after_rmsnorm_and_norm(residual, branch, eps, norm, gain) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = gate.add_after_rmsnorm_and_norm(residual, branch, eps, norm, gain) + for scale in (1e-3, 1.0, 1e3): + residual.normal_().mul_(scale) + branch.normal_().mul_(scale) + norm.weight.normal_(0, 0.02) + if gain is not None: + gain.normal_(1, 0.1) + graph.replay() + with monkeypatch.context() as reference: + reference.setattr( + "tensorrt_llm._torch.models.musespark_open_boundary.cluster_boundary_supported", + lambda *args: False, + ) + expected = gate.add_after_rmsnorm_and_norm(residual, branch, eps, norm, gain) + torch.cuda.synchronize() + assert actual[0].dtype == torch.float32 and actual[1].dtype == torch.bfloat16 + for got in actual: + assert torch.isfinite(got).all() + # Different variance trees can cross individual BF16 thresholds. + # The tight aggregate bound rejects omitting a BF16 rounding stage. + residual_error = (actual[0] - expected[0]).norm() / expected[0].norm().clamp_min(1e-12) + normalized_error = (actual[1].float() - expected[1].float()).norm() / expected[ + 1 + ].float().norm().clamp_min(1e-12) + assert residual_error < 2e-4 + assert normalized_error < 1e-3 + eager = gate.add_after_rmsnorm_and_norm(residual, branch, eps, norm, gain) + for got, want in zip(actual, eager): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +@pytest.mark.parametrize("tokens,hidden", [(33, 8192), (2, 4096)]) +@torch.inference_mode() +def test_cluster_boundary_unsupported_shapes_keep_fallback( + tokens: int, hidden: int, monkeypatch: pytest.MonkeyPatch +) -> None: + from tensorrt_llm._torch.models import musespark_open_boundary + + def unexpected(*args, **kwargs): + raise AssertionError("Unsupported shape reached the cluster kernel") + + monkeypatch.setattr(musespark_open_boundary, "cluster_boundary", unexpected) + residual = torch.randn(tokens, hidden, device="cuda", dtype=torch.float32) + branch = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + gate = MuseSparkOpenGatedResidual(hidden, torch.bfloat16).cuda() + norm = RMSNorm(hidden_size=hidden, eps=1e-5, dtype=torch.bfloat16, use_gemma=True).cuda() + actual = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm) + expected = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm) + for got, want in zip(actual, expected): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +@torch.inference_mode() +def test_cluster_boundary_unaligned_contiguous_branch_keeps_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.models import musespark_open_boundary + + def unexpected(*args, **kwargs): + raise AssertionError("Unaligned input reached the vectorized cluster kernel") + + monkeypatch.setattr(musespark_open_boundary, "cluster_boundary", unexpected) + residual = torch.randn(2, 8192, device="cuda", dtype=torch.float32) + branch = torch.randn(2 * 8192 + 1, device="cuda", dtype=torch.bfloat16)[1:].view(2, 8192) + assert branch.is_contiguous() and branch.data_ptr() % 16 != 0 + gate = MuseSparkOpenGatedResidual(8192, torch.bfloat16).cuda() + norm = RMSNorm(hidden_size=8192, eps=1e-5, dtype=torch.bfloat16, use_gemma=True).cuda() + actual = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm) + expected = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm) + for got, want in zip(actual, expected): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +@torch.inference_mode() +def test_cluster_boundary_rejects_empty_inputs() -> None: + from tensorrt_llm._torch.models.musespark_open_boundary import cluster_boundary_supported + + residual = torch.empty(0, 8192, device="cuda", dtype=torch.float32) + branch = torch.empty_like(residual, dtype=torch.bfloat16) + coefficient = torch.ones(8192, device="cuda", dtype=torch.float32) + weight = torch.ones(8192, device="cuda", dtype=torch.bfloat16) + assert not cluster_boundary_supported(residual, branch, None, coefficient, coefficient, weight) + + +@pytest.mark.parametrize("device_index", [0, 1]) +@pytest.mark.parametrize("cache_enabled", [False, True]) +@pytest.mark.parametrize("tokens", [2, 16]) +@torch.inference_mode() +def test_cluster_boundary_side_stream_and_gate_cache( + device_index: int, cache_enabled: bool, tokens: int, monkeypatch: pytest.MonkeyPatch +) -> None: + if torch.cuda.device_count() <= device_index: + pytest.skip("Second GPU unavailable") + with torch.cuda.device(device_index): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + residual = torch.randn(tokens, 8192, device="cuda", dtype=torch.float32) + branch = torch.randn(tokens, 8192, device="cuda", dtype=torch.bfloat16) + gain = torch.ones(8192, device="cuda", dtype=torch.bfloat16) + gate = MuseSparkOpenGatedResidual(8192, torch.bfloat16).cuda() + gate.gate.normal_(0, 0.5) + if cache_enabled: + # Uncached, the coefficients are recomputed from the live gate. + gate.cache_derived_state() + norm = RMSNorm(hidden_size=8192, eps=1e-5, dtype=torch.bfloat16, use_gemma=True).cuda() + norm.weight.normal_(0, 0.02) + gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm, gain) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm, gain) + for _ in range(8): + residual.normal_() + branch.normal_() + gain.normal_(1, 0.1) + norm.weight.normal_(0, 0.02) + if not cache_enabled: + gate.gate.normal_(0, 0.5) + graph.replay() + with monkeypatch.context() as reference: + reference.setattr( + "tensorrt_llm._torch.models.musespark_open_boundary.cluster_boundary_supported", + lambda *args: False, + ) + expected = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm, gain) + stream.synchronize() + for got, want, tolerance in zip(actual, expected, (2e-4, 1e-3)): + assert got.device.index == device_index + assert torch.isfinite(got).all() + error = (got.float() - want.float()).norm() / want.float().norm() + assert error < tolerance + + +def test_cluster_boundary_does_not_import_the_finalize_module() -> None: + import ast + import inspect + + from tensorrt_llm._torch.models import musespark_open_boundary + + tree = ast.parse(inspect.getsource(musespark_open_boundary)) + imported = set() + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + imported.add(node.module or "") + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + assert not any("finalize_ar" in name for name in imported) + + +def test_cluster_boundary_availability_ignores_finalize_fusion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Disabling the finalize fusion must not disable the boundary kernel.""" + from tensorrt_llm._torch.models import musespark_open_boundary, musespark_open_finalize_ar + + monkeypatch.setattr(musespark_open_finalize_ar, "HAS_CUTEDSL", False) + musespark_open_boundary.cluster_boundary_available.cache_clear() + try: + assert musespark_open_boundary.cluster_boundary_available(torch.device("cuda", 0)) + finally: + musespark_open_boundary.cluster_boundary_available.cache_clear() diff --git a/tests/unittest/_torch/modeling/test_musespark_open_finalize_allreduce.py b/tests/unittest/_torch/modeling/test_musespark_open_finalize_allreduce.py new file mode 100644 index 000000000000..758925155471 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_finalize_allreduce.py @@ -0,0 +1,384 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused finalize + all-reduce parity with the unfused finalize and NCCL, on four GB300 ranks. + +Each test runs its body on four MPI workers spawned by the test process. The +reference on every rank is the unfused finalize of the same inputs followed by +an NCCL all-reduce; every rank must also end with the same bits. +""" + +import contextlib +import pickle +import sys + +import cloudpickle +import pytest +import torch +from mpi4py import MPI + +import tensorrt_llm +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.models.modeling_musespark_open import ( + MuseSparkOpenMoE, + post_expert_norm_combine, +) +from tensorrt_llm._torch.models.musespark_open_finalize_ar import HAS_CUTEDSL, FinalizeIpcWorkspace +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL) + +RANKS = 4 + +pytestmark = [ + # The worker pool is reused across tests, and NCCL starts threads that outlive each test. + pytest.mark.threadleak(enabled=False), + pytest.mark.skipif( + torch.cuda.device_count() < RANKS + or not HAS_CUTEDSL + or torch.cuda.get_device_capability() != (10, 3), + reason="four GB300 GPUs and CuTe DSL required", + ), + pytest.mark.parametrize("mpi_pool_executor", [RANKS], indirect=True), +] + + +def _run_rank(body, *args): + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + mapping = Mapping(world_size=RANKS, tp_size=RANKS, rank=rank, gpus_per_node=RANKS) + with torch.inference_mode(): + return body(rank, mapping, *args) + + +def _on_every_rank(executor, body, *args) -> list: + """`body(rank, mapping, *args)` on each worker; the per-rank results in rank order.""" + return list(executor.map(_run_rank, [body] * RANKS, *[[arg] * RANKS for arg in args])) + + +def _reference(nccl, gemm, indices, weights, norm, eps, tokens): + partial = post_expert_norm_combine( + gemm, indices, weights, norm, eps, num_tokens=tokens, top_k=8, hidden_size=4096 + ) + return nccl(partial.to(torch.bfloat16)) + + +def _within_gates(actual: torch.Tensor, expected: torch.Tensor) -> bool: + actual, expected = actual.float(), expected.float() + if not torch.allclose(actual, expected, rtol=0.04, atol=0.02): + return False + # Every expert can be masked in a one-token case, which makes both sides zero. + return bool((actual - expected).norm() <= 0.01 * expected.norm()) + + +def _same_on_all_ranks(tensor: torch.Tensor) -> bool: + gathered = MPI.COMM_WORLD.allgather(tensor.view(torch.int16).cpu().numpy()) + return all((peer == gathered[0]).all() for peer in gathered) + + +def _all_ranks(ok: bool) -> bool: + """Combine per-rank verdicts, so a failure on one rank never strands the others.""" + return all(MPI.COMM_WORLD.allgather(bool(ok))) + + +def _inputs(tokens: int, rank: int, seed: int, rows=None, width=4096): + """Routing shared by all ranks; every fourth slot is local to `rank`.""" + shared = torch.Generator(device="cuda").manual_seed(seed) + weights = torch.rand(tokens, 8, device="cuda", generator=shared) + weights /= weights.sum(-1, keepdim=True) + norm = (torch.randn(4096, device="cuda", generator=shared) * 0.03).to(torch.bfloat16) + local = torch.Generator(device="cuda").manual_seed(seed * 1009 + rank) + rows = max(64, tokens * 8) if rows is None else rows + gemm = torch.randn(rows, width, device="cuda", generator=local).to(torch.bfloat16) + indices = torch.arange(tokens * 8, device="cuda", dtype=torch.int32).reshape(tokens, 8) + indices = torch.where(indices % RANKS == rank, indices, -1).contiguous() + return gemm, indices, weights, norm + + +def _graph_parity(rank, mapping, tokens, case): + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + last = rank == RANKS - 1 + gemm, indices, weights, norm = _inputs(tokens, rank, 1270 + tokens) + source = gemm.clone() + if case == "upper": + indices[0, rank] = 10000 + if case == "empty_rank" and last: + indices.fill_(-1) + gemm.fill_(float("nan")) + if case == "zero_rows" and last: + indices.zero_() + gemm = torch.empty(0, 4096, device="cuda", dtype=torch.bfloat16) + + for _ in range(3): + workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(5): + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + ok = True + for repeat in range(3): + if not (case in ("empty_rank", "zero_rows") and last): + gemm.copy_(source * (-1 if repeat % 2 else 1) + 0.1 * repeat) + expected = _reference(nccl, gemm, indices, weights, norm, 1e-8, tokens) + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph.replay() + torch.cuda.synchronize() + ok &= _within_gates(actual, expected) + ok &= _same_on_all_ranks(actual) + del graph + return _all_ranks(ok) + + +@pytest.mark.parametrize("tokens", [1, 2, 8, 16, 32, 80, 115, 122, 123, 128]) +@pytest.mark.parametrize("case", ["normal", "upper", "empty_rank", "zero_rows"]) +def test_finalize_allreduce_graph_parity(mpi_pool_executor, tokens: int, case: str) -> None: + """Replays of a captured graph track mutated inputs, including a rank with no work.""" + assert all(_on_every_rank(mpi_pool_executor, _graph_parity, tokens, case)) + + +def _compiles_nothing(rank, mapping): + import cutlass.cute as cute + + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + + def no_compilation(*args, **kwargs): + raise AssertionError("a kernel was compiled after workspace construction") + + ok = True + with pytest.MonkeyPatch.context() as patch: + patch.setattr(cute, "compile", no_compilation) + for step, tokens in enumerate( + (1, 2, 7, 8, 9, 16, 31, 32, 33, 64, 65, 100, 127, 128, 1, 128) + ): + rows = (0, 1, 3, max(64, tokens * 8))[(step + rank) % 4] + gemm, indices, weights, norm = _inputs(tokens, rank, 300 + step, rows=rows) + eps = (1e-8, 1e-6, 1e-5)[step % 3] + actual = workspace.finalize(gemm, indices, weights, norm, eps, tokens) + ok &= _within_gates(actual, _reference(nccl, gemm, indices, weights, norm, eps, tokens)) + ok &= _same_on_all_ranks(actual) + # Global capture mode turns any compilation or module load during capture into an error. + gemm, indices, weights, norm = _inputs(16, rank, 999) + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, capture_error_mode="global"): + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, 16) + graph.replay() + expected = _reference(nccl, gemm, indices, weights, norm, 1e-8, 16) + torch.cuda.synchronize() + ok &= _within_gates(actual, expected) + ok &= _same_on_all_ranks(actual) + del graph + return _all_ranks(ok) + + +def test_finalize_compiles_nothing_after_construction(mpi_pool_executor) -> None: + """Token counts, row counts, eps and routing changes reuse the constructed kernel.""" + assert all(_on_every_rank(mpi_pool_executor, _compiles_nothing)) + + +def _replays_track_inputs(rank, mapping, tokens): + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + calls = [_inputs(tokens, rank, 500 + i) for i in range(5)] + for gemm, indices, weights, norm in calls: + workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = [workspace.finalize(g, i, w, n, 1e-8, tokens) for g, i, w, n in calls] + ok = True + for repeat in range(4): + for index, (gemm, indices, weights, norm) in enumerate(calls): + fresh = _inputs(tokens, rank, 700 + 10 * repeat + index) + gemm.copy_(fresh[0]) + indices.copy_(torch.roll(fresh[1], shifts=repeat, dims=1)) + weights.copy_(fresh[2]) + norm.copy_(fresh[3]) + expected = [_reference(nccl, g, i, w, n, 1e-8, tokens) for g, i, w, n in calls] + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph.replay() + torch.cuda.synchronize() + for actual, want in zip(outputs, expected): + ok &= _within_gates(actual, want) + ok &= _same_on_all_ranks(actual) + del graph + return _all_ranks(ok) + + +@pytest.mark.parametrize("tokens", [1, 9, 64, 128]) +def test_finalize_graph_replays_track_every_input(mpi_pool_executor, tokens: int) -> None: + """Each captured call reads its own inputs; stale data would surface as a mismatch.""" + assert all(_on_every_rank(mpi_pool_executor, _replays_track_inputs, tokens)) + + +def _contract_case(rank, mapping, case, tokens): + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + gemm, indices, weights, norm = _inputs(tokens, rank, 1100 + tokens) + eps_values = (1e-8,) + if case == "negative_zero_partials": + # Alternating-sign rows times subnormal weights round to BF16 (+0, -0) + # pairs: the sentinel bit pattern, unless negative zeros are canonicalized. + signs = torch.where(torch.arange(4096, device="cuda") % 2 == 0, 1.0, -1.0) + gemm = signs.expand(gemm.shape[0], -1).to(torch.bfloat16).contiguous() + weights = torch.full_like(weights, 1e-41) + elif case == "all_invalid": + indices = torch.full_like(indices, -1) + elif case == "duplicate_rows": + indices[:, 1] = indices[:, 0] + elif case == "extreme_indices": + rows = gemm.shape[0] + extremes = torch.tensor([rows - 1, rows, -1, -(2**31), 2**31 - 1], dtype=torch.int32) + indices[:, :5] = extremes.to(indices.device) + elif case == "padded_rows_with_nan": + wide = torch.full( + (gemm.shape[0], 4096 + 64), float("nan"), device="cuda", dtype=torch.bfloat16 + ) + wide[:, :4096] = gemm + gemm = wide[:, :4096] + elif case == "eps_sweep": + gemm = (gemm.float() * 1e-2).to(torch.bfloat16) + eps_values = (1e-6, 1e-3, 0.1, 1.0) + elif case == "single_contributor" and rank != 0: + indices = torch.full_like(indices, -1) + ok = True + for eps in eps_values: + actual = workspace.finalize(gemm, indices, weights, norm, eps, tokens) + ok &= _within_gates(actual, _reference(nccl, gemm, indices, weights, norm, eps, tokens)) + ok &= _same_on_all_ranks(actual) + if case == "all_invalid": + ok &= bool((actual.float() == 0).all()) + return _all_ranks(ok) + + +@pytest.mark.parametrize( + "case", + [ + "negative_zero_partials", + "all_invalid", + "duplicate_rows", + "extreme_indices", + "padded_rows_with_nan", + "eps_sweep", + "single_contributor", + ], +) +@pytest.mark.parametrize("tokens", [1, 33, 128]) +def test_finalize_contract_cases(mpi_pool_executor, case: str, tokens: int) -> None: + assert all(_on_every_rank(mpi_pool_executor, _contract_case, case, tokens)) + + +def _malformed_calls(rank, mapping): + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + tokens = 4 + gemm, indices, weights, norm = _inputs(tokens, rank, 77) + padded = torch.zeros(gemm.numel() + 8, device="cuda", dtype=torch.bfloat16) + misaligned = padded[1 : 1 + gemm.numel()].view(gemm.shape) + odd_stride = torch.zeros(gemm.shape[0], 4100, device="cuda", dtype=torch.bfloat16)[:, :4096] + malformed = [ + (gemm.float(), indices, weights, norm, 1e-8, tokens), + (gemm[:, :2048], indices, weights, norm, 1e-8, tokens), + (gemm.t(), indices, weights, norm, 1e-8, tokens), + (misaligned, indices, weights, norm, 1e-8, tokens), + (odd_stride, indices, weights, norm, 1e-8, tokens), + (gemm, indices.long(), weights, norm, 1e-8, tokens), + (gemm, indices[:, :4].contiguous(), weights, norm, 1e-8, tokens), + (gemm, indices, weights.half(), norm, 1e-8, tokens), + (gemm, indices, weights, norm[:2048], 1e-8, tokens), + (gemm, indices, weights, norm, 0.0, tokens), + (gemm, indices, weights, norm, float("nan"), tokens), + (gemm, indices, weights, norm, 1e-8, 0), + (gemm, indices, weights, norm, 1e-8, 129), + ] + rejected = [] + for call in malformed: + try: + workspace.finalize(*call) + except ValueError: + rejected.append(True) + else: + rejected.append(False) + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + expected = _reference(nccl, gemm, indices, weights, norm, 1e-8, tokens) + ok = all(rejected) and _within_gates(actual, expected) + ok &= _same_on_all_ranks(actual) + return _all_ranks(ok) + + +def test_finalize_rejects_malformed_calls_before_launch(mpi_pool_executor) -> None: + """Every rank rejects the same malformed call; the next valid call still works.""" + assert all(_on_every_rank(mpi_pool_executor, _malformed_calls)) + + +def _fp32_accumulation(rank, mapping): + workspace = FinalizeIpcWorkspace(mapping) + tokens = 3 + rows = torch.ones(8, 4096, dtype=torch.bfloat16, device="cuda") + norm = torch.zeros(4096, dtype=torch.bfloat16, device="cuda") + ok = True + # Rank partials: rank 0 contributes 1, every other rank 2^-8. + indices = torch.full((tokens, 8), -1, dtype=torch.int32, device="cuda") + indices[:, 0] = 0 + weights = torch.zeros(tokens, 8, device="cuda") + weights[:, 0] = 1.0 if rank == 0 else 2.0**-8 + actual = workspace.finalize(rows, indices, weights, norm, 1e-8, tokens) + expected = float(torch.tensor(1.0 + (RANKS - 1) * 2.0**-8).to(torch.bfloat16)) + ok &= bool((actual.float() == expected).all()) + # Slot partials on rank 0: eight slots on one row with weights [1, 2^-9 x 7] give 1.015625. + indices = torch.zeros(tokens, 8, dtype=torch.int32, device="cuda") + if rank != 0: + indices.fill_(-1) + weights = torch.full((tokens, 8), 2.0**-9, device="cuda") + weights[:, 0] = 1.0 + actual = workspace.finalize(rows, indices, weights, norm, 1e-8, tokens) + ok &= bool((actual.float() == 1.015625).all()) + return _all_ranks(ok) + + +def test_finalize_accumulates_in_fp32(mpi_pool_executor) -> None: + """Exact sums that BF16 accumulation would round away; tolerance gates cannot see this.""" + assert all(_on_every_rank(mpi_pool_executor, _fp32_accumulation)) + + +def _model_hook(rank, mapping, tokens): + """The MoE's registered finalize, with forward having chosen the fused path.""" + moe = MuseSparkOpenMoE.__new__(MuseSparkOpenMoE) + torch.nn.Module.__init__(moe) + moe.mapping = mapping + moe.nvtx = lambda name: contextlib.nullcontext() + moe.finalize_allreduce_workspace = FinalizeIpcWorkspace(mapping) + moe._reduce_in_finalize = True + gemm, indices, weights, norm = _inputs(tokens, rank, 4400 + tokens) + moe.post_expert_norm_weight = norm + moe.post_expert_norm_eps = 1e-6 + # TRTLLM-Gen dispatches BF16 routing weights. + weights = weights.to(torch.bfloat16) + actual = moe.finalize( + gemm2_output=gemm, + expanded_idx_to_permuted_idx=indices, + token_selected_slots=None, + routing_weights=weights, + num_tokens=tokens, + ) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + expected = _reference(nccl, gemm, indices, weights, norm, 1e-6, tokens) + ok = actual.dtype == torch.bfloat16 and actual.shape == (tokens, 4096) + ok &= _within_gates(actual, expected) + ok &= _same_on_all_ranks(actual) + return _all_ranks(ok) + + +@pytest.mark.parametrize("tokens", [1, 128]) +def test_the_moe_finalize_hook_reduces_across_the_group(mpi_pool_executor, tokens: int) -> None: + """The hook the MoE scheduler calls returns the group sum when forward chose the fusion.""" + assert all(_on_every_rank(mpi_pool_executor, _model_hook, tokens)) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_finalize_dispatch.py b/tests/unittest/_torch/modeling/test_musespark_open_finalize_dispatch.py new file mode 100644 index 000000000000..4fb5de60612b --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_finalize_dispatch.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Which finalize the MoE forward takes, and that the TP reduce runs exactly once. + +The experts are a stub that calls the registered finalize the way the MoE +scheduler does, so these drive the real `forward` and `finalize` on the host. +""" + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_musespark_open import ( + MuseSparkOpenMoE, + post_expert_norm_combine, +) + +pytestmark = pytest.mark.cpu_only + +WIDTH = 16 + + +def _moe(tokens, workspace, *, moe_max_num_tokens=8192): + module = MuseSparkOpenMoE.__new__(MuseSparkOpenMoE) + torch.nn.Module.__init__(module) + module.top_k = 8 + module.expert_hidden_size = WIDTH + module.post_expert_norm_weight = torch.zeros(WIDTH, dtype=torch.bfloat16) + module.post_expert_norm_eps = 1e-6 + module.mapping = SimpleNamespace(enable_attention_dp=False, has_tp=lambda: True, moe_tp_size=1) + module.nvtx = lambda name: nullcontext() + module.finalize_allreduce_workspace = workspace + module._reduce_in_finalize = False + module.all_reduce = Mock(side_effect=lambda partial: partial + 100) + module.unpooler = lambda combined: combined + ready = object() + module.prepare_moe_inputs = lambda hidden: (torch.zeros(tokens, 256), "expert input", ready) + + # What the scheduler hands the registered finalize: TRTLLM-Gen dispatches + # BF16 routing weights. + gemm2 = torch.randn(tokens * 8, WIDTH, dtype=torch.bfloat16) + index = torch.arange(tokens * 8, dtype=torch.int32) + routing = torch.full((tokens, 8), 1 / 8, dtype=torch.bfloat16) + seen = {} + + def experts(expert_input, router_logits, **kwargs): + seen.update(kwargs) + seen["reduce_in_finalize"] = module._reduce_in_finalize + return module.finalize( + gemm2_output=gemm2, + expanded_idx_to_permuted_idx=index, + token_selected_slots=None, + routing_weights=routing, + num_tokens=tokens, + ) + + module.experts = Mock(side_effect=experts, comm=None, moe_max_num_tokens=moe_max_num_tokens) + unfused = post_expert_norm_combine( + gemm2, + index, + routing, + module.post_expert_norm_weight, + module.post_expert_norm_eps, + num_tokens=tokens, + top_k=8, + hidden_size=WIDTH, + ) + return module, seen, ready, unfused + + +def _workspace(tokens): + return SimpleNamespace( + finalize=Mock(return_value=torch.full((tokens, WIDTH), 3.0, dtype=torch.bfloat16)) + ) + + +@pytest.mark.parametrize( + "tokens,with_workspace,compiling,moe_max_num_tokens,expect_fused", + [ + (1, True, False, 8192, True), + (128, True, False, 8192, True), + (129, True, False, 8192, False), + (8, False, False, 8192, False), + (8, True, True, 8192, False), + # A batch the MoE would split into chunks. + (8, True, False, 4, False), + ], +) +def test_the_forward_picks_one_finalize_and_reduces_once( + tokens, with_workspace, compiling, moe_max_num_tokens, expect_fused +): + workspace = _workspace(tokens) if with_workspace else None + module, seen, ready, unfused = _moe(tokens, workspace, moe_max_num_tokens=moe_max_num_tokens) + hidden = torch.ones(tokens, 32, dtype=torch.bfloat16) + + with patch("torch.compiler.is_compiling", return_value=compiling): + result = module.forward(hidden, None) + + assert seen["input_ready_event"] is ready + assert seen["do_finalize"] is False + assert seen["reduce_in_finalize"] is expect_fused + assert module._reduce_in_finalize is False + if expect_fused: + workspace.finalize.assert_called_once() + routing_weights = workspace.finalize.call_args.args[2] + assert routing_weights.dtype == torch.float32 + module.all_reduce.assert_not_called() + torch.testing.assert_close(result, workspace.finalize.return_value, rtol=0, atol=0) + else: + if workspace is not None: + workspace.finalize.assert_not_called() + module.all_reduce.assert_called_once() + torch.testing.assert_close(result, unfused.to(torch.bfloat16) + 100, rtol=0, atol=0) + + +def test_a_failed_expert_call_does_not_leave_the_fused_choice_behind(): + module, _, _, _ = _moe(4, _workspace(4)) + module.experts.side_effect = RuntimeError("injected") + with pytest.raises(RuntimeError, match="injected"): + module.forward(torch.ones(4, 32, dtype=torch.bfloat16), None) + assert module._reduce_in_finalize is False diff --git a/tests/unittest/_torch/modeling/test_musespark_open_finalize_lifecycle.py b/tests/unittest/_torch/modeling/test_musespark_open_finalize_lifecycle.py new file mode 100644 index 000000000000..9bf68308fdcd --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_finalize_lifecycle.py @@ -0,0 +1,423 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Construction, failure and teardown of the fused finalize + all-reduce workspace, on four ranks. + +The four ranks share CUDA IPC buffers, and no rank may free its buffer while a +peer still maps it. The tests record every unmap and free on the node-wide +monotonic clock, delay one rank's unmaps, and check that every unmap finished +before any free started. Each test runs its body on four MPI workers spawned by +the test process. +""" + +import gc +import pickle +import sys +import time + +import cloudpickle +import pytest +import torch +from mpi4py import MPI + +import tensorrt_llm +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.models import modeling_musespark_open, musespark_open_finalize_ar +from tensorrt_llm._torch.models.modeling_musespark_open import post_expert_norm_combine +from tensorrt_llm._torch.models.musespark_open_finalize_ar import HAS_CUTEDSL, FinalizeIpcWorkspace +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL) + +RANKS = 4 +_LEAK_TOLERANCE = 32 << 20 +_UNMAP_DELAY_S = 0.5 + +pytestmark = [ + # The worker pool is reused across tests, and NCCL starts threads that outlive each test. + pytest.mark.threadleak(enabled=False), + pytest.mark.skipif( + torch.cuda.device_count() < RANKS + or not HAS_CUTEDSL + or torch.cuda.get_device_capability() != (10, 3), + reason="four GB300 GPUs and CuTe DSL required", + ), + pytest.mark.parametrize("mpi_pool_executor", [RANKS], indirect=True), +] + + +def _run_rank(body, *args): + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + mapping = Mapping(world_size=RANKS, tp_size=RANKS, rank=rank, gpus_per_node=RANKS) + with torch.inference_mode(): + return body(rank, mapping, *args) + + +def _on_every_rank(executor, body, *args) -> list: + """`body(rank, mapping, *args)` on each worker; the per-rank results in rank order.""" + return list(executor.map(_run_rank, [body] * RANKS, *[[arg] * RANKS for arg in args])) + + +def _check_call(workspace, nccl, rank, tokens, seed) -> bool: + shared = torch.Generator(device="cuda").manual_seed(seed) + weights = torch.rand(tokens, 8, device="cuda", generator=shared) + weights /= weights.sum(-1, keepdim=True) + norm = (torch.randn(4096, device="cuda", generator=shared) * 0.03).to(torch.bfloat16) + local = torch.Generator(device="cuda").manual_seed(seed * 1009 + rank) + gemm = torch.randn(max(64, tokens * 8), 4096, device="cuda", generator=local).to(torch.bfloat16) + indices = torch.arange(tokens * 8, device="cuda", dtype=torch.int32).reshape(tokens, 8) + indices = torch.where(indices % RANKS == rank, indices, -1).contiguous() + expected = nccl( + post_expert_norm_combine( + gemm, indices, weights, norm, 1e-8, num_tokens=tokens, top_k=8, hidden_size=4096 + ).to(torch.bfloat16) + ) + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + torch.cuda.synchronize() + actual, expected = actual.float(), expected.float() + return bool( + torch.allclose(actual, expected, rtol=0.04, atol=0.02) + and (actual - expected).norm() <= 0.01 * expected.norm() + ) + + +def _all_ranks(ok: bool) -> bool: + return all(MPI.COMM_WORLD.allgather(bool(ok))) + + +class _ReleaseLog: + """Records this rank's IPC unmaps and buffer frees; can delay or fail its unmaps.""" + + def __init__(self, patch, *, unmap_delay: float = 0.0, fail_unmap: bool = False): + runtime = musespark_open_finalize_ar.cudart + unmap_peer, free, malloc = ( + runtime.cudaIpcCloseMemHandle, + runtime.cudaFree, + runtime.cudaMalloc, + ) + self.unmaps, self.frees, self.buffers = [], [], set() + + def recorded_unmap(pointer): + if unmap_delay and not self.unmaps: + time.sleep(unmap_delay) + if fail_unmap: + result = (runtime.cudaError_t.cudaErrorInvalidValue,) + else: + result = unmap_peer(pointer) + self.unmaps.append(time.monotonic_ns()) + return result + + def recorded_free(pointer): + if int(pointer) in self.buffers: + self.frees.append(time.monotonic_ns()) + return free(pointer) + + def recorded_malloc(size): + error, pointer = malloc(size) + if error == runtime.cudaError_t.cudaSuccess: + self.buffers.add(int(pointer)) + return error, pointer + + patch.setattr(runtime, "cudaIpcCloseMemHandle", recorded_unmap) + patch.setattr(runtime, "cudaFree", recorded_free) + patch.setattr(runtime, "cudaMalloc", recorded_malloc) + + def watch(self, workspace) -> None: + if workspace._export: + self.buffers.add(workspace._export) + + def clear(self) -> None: + self.unmaps.clear() + self.frees.clear() + + +def _release_order(log: _ReleaseLog) -> tuple[bool, int, int]: + """Whether every rank's unmaps ended before any rank's frees began (the ranks share one + node clock), plus the group's unmap and free counts.""" + logs = MPI.COMM_WORLD.allgather((log.unmaps, log.frees)) + unmaps = [stamp for done, _ in logs for stamp in done] + frees = [stamp for _, done in logs for stamp in done] + ordered = not unmaps or not frees or max(unmaps) < min(frees) + return ordered, len(unmaps), len(frees) + + +def _free_device_memory() -> int: + gc.collect() + torch.cuda.synchronize() + torch.cuda.empty_cache() + return torch.cuda.mem_get_info()[0] + + +def _use_and_close_workspace(mapping, nccl, rank, cycle, log) -> bool: + """One full lifetime: eager calls, a captured graph, quiescence, then close().""" + ok = True + workspace = FinalizeIpcWorkspace(mapping) + log.watch(workspace) + for step in range(3): + ok &= _check_call(workspace, nccl, rank, 1 + 40 * step, 100 * cycle + step) + gemm = torch.randn(64, 4096, device="cuda").to(torch.bfloat16) + indices = torch.full((8, 8), -1, dtype=torch.int32, device="cuda") + weights = torch.ones(8, 8, device="cuda") + norm = torch.zeros(4096, dtype=torch.bfloat16, device="cuda") + workspace.finalize(gemm, indices, weights, norm, 1e-8, 8) + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = workspace.finalize(gemm, indices, weights, norm, 1e-8, 8) + graph.replay() + graph.replay() + torch.cuda.synchronize() + ok &= bool((output.float() == 0).all()) + del graph, output + MPI.COMM_WORLD.Barrier() + log.clear() + workspace.close() + ordered, unmaps, frees = _release_order(log) + ok &= ordered + ok &= (unmaps, frees) == (RANKS * (RANKS - 1), RANKS) + del workspace + gc.collect() + MPI.COMM_WORLD.Barrier() + return ok + + +def _teardown_and_recreate(rank, mapping): + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + # The first lifetime also pays one-time costs (NCCL buffers, kernel modules). + with pytest.MonkeyPatch.context() as patch: + ok = _use_and_close_workspace(mapping, nccl, rank, 0, _ReleaseLog(patch)) + MPI.COMM_WORLD.Barrier() + baseline = _free_device_memory() + for cycle in range(1, 11): + with pytest.MonkeyPatch.context() as patch: + delay = _UNMAP_DELAY_S if rank == cycle % RANKS else 0.0 + log = _ReleaseLog(patch, unmap_delay=delay) + ok &= _use_and_close_workspace(mapping, nccl, rank, cycle, log) + leaked = baseline - _free_device_memory() + return _all_ranks(ok and leaked < _LEAK_TOLERANCE) + + +def test_finalize_workspace_teardown_and_recreate(mpi_pool_executor) -> None: + """Quiesce, close and rebuild ten times after eager and graph use, with one rank slow + to unmap each time; no buffer is freed early and nothing leaks.""" + assert all(_on_every_rank(mpi_pool_executor, _teardown_and_recreate)) + + +def _inject(patch, stage): + module = musespark_open_finalize_ar + if stage == "fingerprint": + patch.setattr(module, "_fingerprint", lambda *args: 12345) + elif stage == "compile": + + def broken_compile(*args): + raise TypeError("injected compilation failure") + + patch.setattr(module, "_compiled_kernel", broken_compile) + elif stage == "allocate": + patch.setattr( + module.cudart, + "cudaMalloc", + lambda size: (module.cudart.cudaError_t.cudaErrorMemoryAllocation, 0), + ) + elif stage == "map_peers": + patch.setattr( + module.cudart, + "cudaIpcOpenMemHandle", + lambda handle, flags: (module.cudart.cudaError_t.cudaErrorInvalidValue, 0), + ) + elif stage == "setup_raises": + + def broken_prepare(self, mapping): + raise KeyError("injected setup failure") + + patch.setattr(FinalizeIpcWorkspace, "_prepare", broken_prepare) + elif stage == "handle_raises": + + def broken_export(handle): + raise ValueError("injected handle export failure") + + patch.setattr(module, "_ipc_mem_handle_to_bytes", broken_export) + elif stage == "open_raises": + + def broken_import(data): + raise TypeError("injected handle import failure") + + patch.setattr(module, "_ipc_mem_handle_from_bytes", broken_import) + elif stage == "launch_raises": + # TVM-FFI reports launch failures as ValueError, TypeError and other types. + def broken_launch(self, *args, **kwargs): + raise ValueError("injected launch failure") + + patch.setattr(FinalizeIpcWorkspace, "_launch", broken_launch) + else: + patch.setattr( + module.cuda, + "cuMemsetD32", + lambda *args: (module.cuda.CUresult.CUDA_ERROR_INVALID_VALUE,), + ) + + +def _construction_failure(rank, mapping, stage): + failing = rank == RANKS - 1 + with pytest.MonkeyPatch.context() as patch: + log = _ReleaseLog(patch, unmap_delay=_UNMAP_DELAY_S if rank == 0 else 0.0) + if failing: + _inject(patch, stage) + expected = ValueError if stage == "fingerprint" else RuntimeError + MPI.COMM_WORLD.Barrier() + baseline = _free_device_memory() + raised = None + try: + FinalizeIpcWorkspace(mapping) + except (ValueError, RuntimeError) as error: + raised = type(error) + ordered, unmaps, frees = _release_order(log) + # Group-wide IPC events the release must produce; the failing rank maps nothing at + # "map_peers" and allocates nothing at "allocate". + events = { + "allocate": (0, RANKS - 1), + "map_peers": ((RANKS - 1) ** 2, RANKS), + "initialize": (RANKS * (RANKS - 1), RANKS), + "handle_raises": (0, RANKS), + "open_raises": ((RANKS - 1) ** 2, RANKS), + "launch_raises": (RANKS * (RANKS - 1), RANKS), + } + ordered &= (unmaps, frees) == events.get(stage, (0, 0)) + leaked = baseline - _free_device_memory() + MPI.COMM_WORLD.Barrier() + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + works = _check_call(workspace, nccl, rank, 16, 4242) + workspace.close() + return _all_ranks(raised is expected and ordered and leaked < _LEAK_TOLERANCE and works) + + +@pytest.mark.parametrize( + "stage", + [ + "fingerprint", + "compile", + "allocate", + "map_peers", + "initialize", + "setup_raises", + "handle_raises", + "open_raises", + "launch_raises", + ], +) +def test_finalize_construction_failure_is_agreed(mpi_pool_executor, stage: str) -> None: + """A failure on one rank, reported as an error code or raised as any exception type, + makes every rank raise the same type and release in group order, with rank 0 slow to + unmap; nothing leaks and a retry then works.""" + assert all(_on_every_rank(mpi_pool_executor, _construction_failure, stage)) + + +def _model_falls_back(rank, mapping): + warnings = [] + with pytest.MonkeyPatch.context() as patch: + patch.setattr(modeling_musespark_open.logger, "warning", warnings.append) + if rank == RANKS - 1: + _inject(patch, "allocate") + workspace = modeling_musespark_open._build_finalize_allreduce_workspace(mapping) + ok = workspace is None and len(warnings) == 1 and "allocation failed" in warnings[0] + return _all_ranks(ok) + + +def test_a_failure_on_one_rank_leaves_every_rank_on_the_unfused_path(mpi_pool_executor) -> None: + """The model's construction degrades on every rank alike instead of failing or hanging.""" + assert all(_on_every_rank(mpi_pool_executor, _model_falls_back)) + + +def _zero_call(tokens: int = 4): + gemm = torch.zeros(64, 4096, device="cuda", dtype=torch.bfloat16) + indices = torch.full((tokens, 8), -1, dtype=torch.int32, device="cuda") + weights = torch.zeros(tokens, 8, device="cuda") + norm = torch.zeros(4096, dtype=torch.bfloat16, device="cuda") + return gemm, indices, weights, norm, 1e-8, tokens + + +def _failed_release(rank, mapping, failure): + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + workspace = FinalizeIpcWorkspace(mapping) + ok = _check_call(workspace, nccl, rank, 8, 77) + failing = rank == RANKS - 1 + message = "" + with pytest.MonkeyPatch.context() as patch: + log = _ReleaseLog(patch, fail_unmap=failure == "unmap" and failing) + log.watch(workspace) + if failing and failure == "sync": + patch.setattr(workspace, "_sync", lambda: False) + MPI.COMM_WORLD.Barrier() + try: + workspace.close() + except RuntimeError as error: + message = str(error) + try: + workspace.finalize(*_zero_call()) + rejected = False + except RuntimeError: + rejected = True + retry = FinalizeIpcWorkspace(mapping) + works = _check_call(retry, nccl, rank, 16, 4243) + retry.close() + ok &= f"rank(s) [{RANKS - 1}]" in message and not log.frees and rejected and works + return _all_ranks(ok) + + +@pytest.mark.parametrize("failure", ["sync", "unmap"]) +def test_finalize_failed_release_frees_nothing(mpi_pool_executor, failure: str) -> None: + """If one rank fails to synchronize or to unmap, every rank raises from close(), no IPC + buffer is freed, and the group can still build a working workspace.""" + assert all(_on_every_rank(mpi_pool_executor, _failed_release, failure)) + + +def _drop_after_enqueue(rank, mapping): + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + other = FinalizeIpcWorkspace(mapping) + ok = True + with pytest.MonkeyPatch.context() as patch: + log = _ReleaseLog(patch) + for repeat in range(3): + workspace = FinalizeIpcWorkspace(mapping) + log.watch(workspace) + lagging = repeat % RANKS + tokens = 64 + shared = torch.Generator(device="cuda").manual_seed(repeat) + weights = torch.rand(tokens, 8, device="cuda", generator=shared) + weights /= weights.sum(-1, keepdim=True) + norm = torch.zeros(4096, dtype=torch.bfloat16, device="cuda") + gemm = torch.randn(512, 4096, device="cuda").to(torch.bfloat16) + indices = torch.arange(tokens * 8, device="cuda", dtype=torch.int32).reshape(tokens, 8) + indices = torch.where(indices % RANKS == rank, indices, -1).contiguous() + expected = nccl( + post_expert_norm_combine( + gemm, indices, weights, norm, 1e-8, num_tokens=tokens, top_k=8, hidden_size=4096 + ).to(torch.bfloat16) + ) + if rank == lagging: + torch.cuda._sleep(20_000_000) + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + if rank == lagging: + del workspace + gc.collect() + torch.cuda.synchronize() + actual, expected = actual.float(), expected.float() + ok &= bool(torch.allclose(actual, expected, rtol=0.04, atol=0.02)) + MPI.COMM_WORLD.Barrier() + if rank != lagging: + del workspace + gc.collect() + MPI.COMM_WORLD.Barrier() + ok &= _check_call(other, nccl, rank, 24, 500 + repeat) + other.close() + return _all_ranks(ok and not log.frees) + + +def test_finalize_drop_right_after_enqueue(mpi_pool_executor) -> None: + """Dropping a workspace without close(), right after a call, waits for the call, frees no + IPC buffer, and leaves another live workspace working.""" + assert all(_on_every_rank(mpi_pool_executor, _drop_after_enqueue)) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_finalize_local.py b/tests/unittest/_torch/modeling/test_musespark_open_finalize_local.py new file mode 100644 index 000000000000..e18895e5113c --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_finalize_local.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Single-process contracts of the finalize+all-reduce module: imports, input checks, release order. + +These need no peers. Argument validation runs on a workspace whose launcher is +replaced by a recorder, so a rejected call is proven to launch nothing. Release +runs on a workspace holding fake buffers, with the TP group and the CUDA calls +replaced by recorders, so the order of every step can be checked exactly. +""" + +import subprocess +import sys + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.models import musespark_open_finalize_ar +from tensorrt_llm._torch.models.musespark_open_finalize_ar import FinalizeIpcWorkspace + +_BLOCKED_DSL = """ +import sys +for name in ("cutlass", "cutlass.cute", "cutlass.utils", "cutlass._mlir", "cutlass.cutlass_dsl"): + sys.modules[name] = None +from tensorrt_llm._torch.models import musespark_open_boundary, musespark_open_finalize_ar +from tensorrt_llm._torch.models.modeling_musespark_open import _build_finalize_allreduce_workspace +print(musespark_open_finalize_ar.HAS_CUTEDSL, musespark_open_boundary._HAS_CUTE_DSL, + _build_finalize_allreduce_workspace(object())) +""" + + +def test_modules_import_without_cute_dsl() -> None: + """Both kernels report unavailable and the model keeps the unfused path.""" + result = subprocess.run( + [sys.executable, "-c", _BLOCKED_DSL], capture_output=True, text=True, timeout=600 + ) + assert result.returncode == 0, result.stderr[-2000:] + assert result.stdout.split()[-3:] == ["False", "False", "None"] + + +@pytest.fixture +def recorded(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + workspace = FinalizeIpcWorkspace.__new__(FinalizeIpcWorkspace) + workspace._device = torch.device("cuda", torch.cuda.current_device()) + workspace._ranks, workspace._rank = 4, 0 + workspace._export, workspace._imports = 0, [] + workspace._kernel = object() + workspace._closed = False + launches = [] + workspace._launch = lambda *args, **kwargs: launches.append((args, kwargs)) + return workspace, launches + + +def _valid(tokens=4, rows=64, width=4096): + gemm = torch.zeros(rows, width, device="cuda", dtype=torch.bfloat16) + indices = torch.full((tokens, 8), -1, device="cuda", dtype=torch.int32) + weights = torch.zeros(tokens, 8, device="cuda") + norm = torch.zeros(4096, device="cuda", dtype=torch.bfloat16) + return [gemm, indices, weights, norm, 1e-8, tokens] + + +def _malformed(): + """Each entry replaces one argument of a valid call.""" + gemm = torch.zeros(64, 4096, device="cuda", dtype=torch.bfloat16) + return { + "fp32 rows": (0, gemm.float()), + "narrow rows": (0, gemm[:, :2048]), + "transposed rows": (0, torch.zeros(4096, 64, device="cuda", dtype=torch.bfloat16).t()), + "misaligned rows": ( + 0, + torch.zeros(64 * 4096 + 8, device="cuda", dtype=torch.bfloat16)[1 : 1 + 64 * 4096].view( + 64, 4096 + ), + ), + "row stride not a multiple of 8": ( + 0, + torch.zeros(64, 4100, device="cuda", dtype=torch.bfloat16)[:, :4096], + ), + "cpu rows": (0, gemm.cpu()), + "int64 indices": (1, torch.full((4, 8), -1, device="cuda", dtype=torch.int64)), + "short indices": (1, torch.full((4, 4), -1, device="cuda", dtype=torch.int32)), + "strided indices": (1, torch.full((4, 16), -1, device="cuda", dtype=torch.int32)[:, ::2]), + "fp16 weights": (2, torch.zeros(4, 8, device="cuda", dtype=torch.float16)), + "short weights": (2, torch.zeros(4, 4, device="cuda")), + "short norm": (3, torch.zeros(2048, device="cuda", dtype=torch.bfloat16)), + "fp32 norm": (3, torch.zeros(4096, device="cuda")), + "misaligned norm": (3, torch.zeros(4104, device="cuda", dtype=torch.bfloat16)[1:4097]), + "zero eps": (4, 0.0), + "negative eps": (4, -1e-6), + "nan eps": (4, float("nan")), + "inf eps": (4, float("inf")), + "no tokens": (5, 0), + "too many tokens": (5, 129), + "float token count": (5, 4.0), + } + + +_MALFORMED_CASES = [ + "fp32 rows", + "narrow rows", + "transposed rows", + "misaligned rows", + "row stride not a multiple of 8", + "cpu rows", + "int64 indices", + "short indices", + "strided indices", + "fp16 weights", + "short weights", + "short norm", + "fp32 norm", + "misaligned norm", + "zero eps", + "negative eps", + "nan eps", + "inf eps", + "no tokens", + "too many tokens", + "float token count", +] + + +@pytest.mark.parametrize("case", _MALFORMED_CASES) +def test_malformed_call_raises_before_launch(recorded, case: str) -> None: + workspace, launches = recorded + position, value = _malformed()[case] + arguments = _valid() + arguments[position] = value + with pytest.raises(ValueError): + workspace.finalize(*arguments) + assert launches == [] + + +@pytest.mark.parametrize( + "rows, width, stride", + [(0, 4096, None), (1, 4096, None), (1, 4100, 4100), (64, 4096, None), (64, 4160, 4160)], +) +def test_legal_per_rank_differences_launch(recorded, rows: int, width: int, stride) -> None: + """Row counts and padded or arbitrary strides may differ between ranks.""" + workspace, launches = recorded + arguments = _valid(rows=rows, width=width) + if stride is not None: + arguments[0] = arguments[0][:, :4096] + output = workspace.finalize(*arguments) + assert len(launches) == 1 + assert output.shape == (4, 4096) and output.dtype == torch.bfloat16 + args, _ = launches[0] + gemm_pointer, row_count = args[0], args[5] + assert row_count == rows + assert gemm_pointer == (0 if rows == 0 else arguments[0].data_ptr()) + + +def test_construction_stage_reports_any_exception_as_failure() -> None: + """A rank must reach the agreement whatever its stage raised, or its peers wait forever.""" + + def broken_stage(argument): + raise KeyError(argument) + + status, detail = FinalizeIpcWorkspace._guarded(broken_stage, "injected") + assert status == musespark_open_finalize_ar._FAILED + assert detail.startswith("KeyError") + + +def test_closed_workspace_rejects_calls(recorded) -> None: + workspace, launches = recorded + workspace._closed = True + with pytest.raises(RuntimeError, match="closed"): + workspace.finalize(*_valid()) + assert launches == [] + + +class _FakeGroup: + """Four-rank TP group whose peers echo this rank, except one peer that can fail one step.""" + + def __init__(self, events): + self.events = events + self.failing_step = None + self.steps = 0 + + def tp_allgather_int64(self, values): + self.steps += 1 + self.events.append("agree") + table = np.array([list(values)] * 4, dtype=np.int64) + if self.steps == self.failing_step: + table[2, 0] = 0 + return table + + +_UNMAPS = ["unmap 0x1000", "unmap 0x2000", "unmap 0x3000"] +_ORDERED_RELEASE = ["sync", "agree", *_UNMAPS, "agree", "free 0xe000", "agree"] + + +@pytest.fixture +def releasable(monkeypatch: pytest.MonkeyPatch): + """A workspace holding a fake IPC buffer and three fake peer mappings.""" + if not torch.cuda.is_available() or not hasattr(musespark_open_finalize_ar, "cudart"): + pytest.skip("CUDA and cuda-python required") + runtime = musespark_open_finalize_ar.cudart + events = [] + + def unmap(pointer): + events.append(f"unmap {pointer:#x}") + return (runtime.cudaError_t.cudaSuccess,) + + def free(pointer): + events.append(f"free {pointer:#x}") + return (runtime.cudaError_t.cudaSuccess,) + + def sync(): + events.append("sync") + return True + + monkeypatch.setattr(runtime, "cudaIpcCloseMemHandle", unmap) + monkeypatch.setattr(runtime, "cudaFree", free) + workspace = FinalizeIpcWorkspace.__new__(FinalizeIpcWorkspace) + workspace._device = torch.device("cuda", torch.cuda.current_device()) + workspace._export, workspace._imports = 0xE000, [0x1000, 0x2000, 0x3000] + workspace._dist = _FakeGroup(events) + workspace._sync = sync + workspace._closed = False + yield workspace, events + # The fake pointers must never reach CUDA once the recorders are gone. + workspace._closed = True + + +def test_close_frees_only_after_the_group_unmapped(releasable) -> None: + workspace, events = releasable + workspace.close() + assert events == _ORDERED_RELEASE + workspace.close() + assert events == _ORDERED_RELEASE, "a second close must do nothing" + + +def test_failed_construction_releases_in_the_same_order(releasable) -> None: + workspace, events = releasable + with pytest.raises(ValueError, match="injected"): + workspace._abandon(ValueError, "injected") + assert events == _ORDERED_RELEASE + + +@pytest.mark.parametrize( + "failing_step, released", + [(1, ["sync", "agree"]), (2, ["sync", "agree", *_UNMAPS, "agree"])], + ids=["peer_sync_fails", "peer_unmap_fails"], +) +def test_peer_failure_frees_nothing(releasable, failing_step: int, released) -> None: + workspace, events = releasable + workspace._dist.failing_step = failing_step + with pytest.raises(RuntimeError, match=r"rank\(s\) \[2\]"): + workspace.close() + assert events == released + + +def test_local_unmap_failure_frees_nothing(releasable, monkeypatch: pytest.MonkeyPatch) -> None: + workspace, events = releasable + runtime = musespark_open_finalize_ar.cudart + monkeypatch.setattr( + runtime, + "cudaIpcCloseMemHandle", + lambda pointer: (runtime.cudaError_t.cudaErrorInvalidValue,), + ) + with pytest.raises(RuntimeError, match="unmapping"): + workspace.close() + assert "free 0xe000" not in events + + +def test_local_sync_failure_releases_nothing(releasable) -> None: + workspace, events = releasable + workspace._sync = lambda: False + with pytest.raises(RuntimeError, match="synchronization"): + workspace.close() + assert events == ["agree"] + + +def test_drop_unmaps_peers_but_frees_nothing(releasable) -> None: + """Without close() the destructor involves no other rank and keeps the shared buffer.""" + workspace, events = releasable + workspace.__del__() + assert events == ["sync", *_UNMAPS] + + +def test_drop_after_failed_sync_releases_nothing(releasable) -> None: + workspace, events = releasable + workspace._sync = lambda: False + workspace.__del__() + assert events == [] diff --git a/tests/unittest/_torch/modeling/test_musespark_open_finalize_stress.py b/tests/unittest/_torch/modeling/test_musespark_open_finalize_stress.py new file mode 100644 index 000000000000..d0606a4ce23d --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_finalize_stress.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Long and skewed call sequences for the fused finalize + all-reduce workspace, on four ranks. + +Every call carries its own inputs, so a read of stale or not-yet-written data +shows up as a mismatch against the unfused finalize and NCCL. Each test runs +its body on four MPI workers spawned by the test process. +""" + +import pickle +import random +import sys +import time + +import cloudpickle +import pytest +import torch +from mpi4py import MPI + +import tensorrt_llm +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.models.modeling_musespark_open import post_expert_norm_combine +from tensorrt_llm._torch.models.musespark_open_finalize_ar import HAS_CUTEDSL, FinalizeIpcWorkspace +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL) + +RANKS = 4 + +pytestmark = [ + # The worker pool is reused across tests, and NCCL starts threads that outlive each test. + pytest.mark.threadleak(enabled=False), + pytest.mark.skipif( + torch.cuda.device_count() < RANKS + or not HAS_CUTEDSL + or torch.cuda.get_device_capability() != (10, 3), + reason="four GB300 GPUs and CuTe DSL required", + ), + pytest.mark.parametrize("mpi_pool_executor", [RANKS], indirect=True), +] + + +def _run_rank(body, *args): + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + mapping = Mapping(world_size=RANKS, tp_size=RANKS, rank=rank, gpus_per_node=RANKS) + with torch.inference_mode(): + return body(rank, mapping, *args) + + +def _on_every_rank(executor, body, *args) -> list: + """`body(rank, mapping, *args)` on each worker; the per-rank results in rank order.""" + return list(executor.map(_run_rank, [body] * RANKS, *[[arg] * RANKS for arg in args])) + + +def _reference(nccl, gemm, indices, weights, norm, tokens): + partial = post_expert_norm_combine( + gemm, indices, weights, norm, 1e-8, num_tokens=tokens, top_k=8, hidden_size=4096 + ) + return nccl(partial.to(torch.bfloat16)) + + +def _within_gates(actual: torch.Tensor, expected: torch.Tensor) -> bool: + actual, expected = actual.float(), expected.float() + if not torch.allclose(actual, expected, rtol=0.04, atol=0.02): + return False + return bool((actual - expected).norm() <= 0.01 * expected.norm()) + + +def _all_ranks(ok: bool) -> bool: + return all(MPI.COMM_WORLD.allgather(bool(ok))) + + +def _inputs(tokens: int, rank: int, seed: int, rows=None): + shared = torch.Generator(device="cuda").manual_seed(seed) + weights = torch.rand(tokens, 8, device="cuda", generator=shared) + weights /= weights.sum(-1, keepdim=True) + norm = (torch.randn(4096, device="cuda", generator=shared) * 0.03).to(torch.bfloat16) + local = torch.Generator(device="cuda").manual_seed(seed * 1009 + rank) + rows = max(64, tokens * 8) if rows is None else rows + gemm = torch.randn(rows, 4096, device="cuda", generator=local).to(torch.bfloat16) + indices = torch.arange(tokens * 8, device="cuda", dtype=torch.int32).reshape(tokens, 8) + indices = torch.where(indices % RANKS == rank, indices, -1).contiguous() + return gemm, indices, weights, norm + + +def _refill(call, tokens, rank, seed): + """Overwrite a captured call's inputs in place with a fresh payload.""" + fresh = _inputs(tokens, rank, seed, rows=call[0].shape[0]) + for target, source in zip(call, fresh): + target.copy_(source) + + +def _capture(workspace, calls): + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = [workspace.finalize(g, i, w, n, 1e-8, t) for (g, i, w, n), t in calls] + return graph, outputs + + +def _random_sequences(rank, mapping): + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + schedule = random.Random(2026) # identical on every rank + graphs = [] + for count in (5, 7, 1): + tokens = [schedule.randint(1, 128) for _ in range(count)] + calls = [(_inputs(t, rank, 5000 + 100 * len(graphs) + i), t) for i, t in enumerate(tokens)] + for (gemm, indices, weights, norm), t in calls: + workspace.finalize(gemm, indices, weights, norm, 1e-8, t) + graphs.append((calls, *_capture(workspace, calls))) + # Start with abrupt 128 <-> 1 switches and ramps, then random traffic. + plan = [128, 1] * 8 + list(range(1, 129, 9)) + list(range(128, 0, -11)) + ok = True + for step in range(1500): + local = random.Random(step * 31 + rank) # per-rank row counts may differ + if step < len(plan) or schedule.random() < 0.7: + tokens = plan[step] if step < len(plan) else schedule.randint(1, 128) + rows = local.choice((0, 1, 5, None)) + gemm, indices, weights, norm = _inputs(tokens, rank, 10_000 + step, rows) + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + ok &= _within_gates(actual, _reference(nccl, gemm, indices, weights, norm, tokens)) + else: + calls, graph, outputs = graphs[schedule.randrange(len(graphs))] + for index, (call, tokens) in enumerate(calls): + _refill(call, tokens, rank, 20_000 + 10 * step + index) + expected = [_reference(nccl, *call, tokens) for call, tokens in calls] + graph.replay() + for actual, want in zip(outputs, expected): + ok &= _within_gates(actual, want) + torch.cuda.synchronize() + del graphs + return _all_ranks(ok) + + +def test_finalize_random_call_sequences(mpi_pool_executor) -> None: + """Eager calls and three graphs (5, 7 and 1 calls) interleaved in a seeded order.""" + assert all(_on_every_rank(mpi_pool_executor, _random_sequences)) + + +def _delayed_ranks(rank, mapping, pattern): + workspace = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + schedule = random.Random(77) + ok = True + for step in range(200): + tokens = schedule.randint(1, 128) + slow = { + "always_slow": 0, + "rotating": step % RANKS, + "random": schedule.randrange(RANKS), + }[pattern] + cycles = int(schedule.uniform(0, 4e6)) + host_pause = schedule.random() < 0.1 + gemm, indices, weights, norm = _inputs(tokens, rank, 30_000 + step) + expected = _reference(nccl, gemm, indices, weights, norm, tokens) + if rank == slow: + torch.cuda._sleep(cycles) + if host_pause: + time.sleep(0.002) + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + ok &= _within_gates(actual, expected) + # The same skew inside a captured graph: the slow rank sleeps before every call. + calls = [(_inputs(t, rank, 40_000 + t), t) for t in (3, 128, 17)] + for (gemm, indices, weights, norm), t in calls: + workspace.finalize(gemm, indices, weights, norm, 1e-8, t) + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = [] + for (gemm, indices, weights, norm), t in calls: + if rank == 0: + torch.cuda._sleep(1_000_000) + outputs.append(workspace.finalize(gemm, indices, weights, norm, 1e-8, t)) + for repeat in range(30): + for index, (call, tokens) in enumerate(calls): + _refill(call, tokens, rank, 50_000 + 10 * repeat + index) + expected = [_reference(nccl, *call, tokens) for call, tokens in calls] + graph.replay() + for actual, want in zip(outputs, expected): + ok &= _within_gates(actual, want) + torch.cuda.synchronize() + del graph + return _all_ranks(ok) + + +@pytest.mark.parametrize("pattern", ["always_slow", "rotating", "random"]) +def test_finalize_with_delayed_ranks(mpi_pool_executor, pattern: str) -> None: + """A lagging rank must never let the others reuse a buffer it has not drained.""" + assert all(_on_every_rank(mpi_pool_executor, _delayed_ranks, pattern)) + + +def _two_workspaces(rank, mapping): + first = FinalizeIpcWorkspace(mapping) + second = FinalizeIpcWorkspace(mapping) + nccl = AllReduce(mapping, strategy=AllReduceStrategy.NCCL).cuda() + schedule = random.Random(5) + ok = True + for step in range(120): + workspace = first if step % 3 else second + tokens = schedule.randint(1, 128) + gemm, indices, weights, norm = _inputs(tokens, rank, 60_000 + step) + actual = workspace.finalize(gemm, indices, weights, norm, 1e-8, tokens) + ok &= _within_gates(actual, _reference(nccl, gemm, indices, weights, norm, tokens)) + calls = [(_inputs(t, rank, 70_000 + t), t) for t in (5, 64, 128)] + owners = (first, second, first) + for owner, ((gemm, indices, weights, norm), t) in zip(owners, calls): + owner.finalize(gemm, indices, weights, norm, 1e-8, t) + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = [ + owner.finalize(g, i, w, n, 1e-8, t) for owner, ((g, i, w, n), t) in zip(owners, calls) + ] + for repeat in range(10): + for index, (call, tokens) in enumerate(calls): + _refill(call, tokens, rank, 80_000 + 10 * repeat + index) + expected = [_reference(nccl, *call, tokens) for call, tokens in calls] + graph.replay() + for actual, want in zip(outputs, expected): + ok &= _within_gates(actual, want) + torch.cuda.synchronize() + del graph + # Retiring one workspace must not disturb the other. + MPI.COMM_WORLD.Barrier() + del second + for step in range(20): + tokens = schedule.randint(1, 128) + gemm, indices, weights, norm = _inputs(tokens, rank, 90_000 + step) + actual = first.finalize(gemm, indices, weights, norm, 1e-8, tokens) + ok &= _within_gates(actual, _reference(nccl, gemm, indices, weights, norm, tokens)) + return _all_ranks(ok) + + +def test_finalize_two_workspaces_interleave(mpi_pool_executor) -> None: + """Independent workspaces share no state, eagerly or inside one graph.""" + assert all(_on_every_rank(mpi_pool_executor, _two_workspaces)) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_full_routing.py b/tests/unittest/_torch/modeling/test_musespark_open_full_routing.py new file mode 100644 index 000000000000..8a46e1ec441e --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_full_routing.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused routing: selection, bias semantics, tie policy, fallbacks and graph replay.""" + +from unittest.mock import patch + +import pytest +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.models import musespark_open_kernels +from tensorrt_llm._torch.models.modeling_musespark_open import ( + MuseSparkOpenRoutingMethod, + _routing_weights, +) + +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _router(bias): + return MuseSparkOpenRoutingMethod( + top_k=8, num_experts=256, callable_e_score_correction_bias=lambda: bias, renorm_eps=1e-6 + ) + + +@cuda_only +@pytest.mark.parametrize("tokens", [1, 2, 16, 128]) +@pytest.mark.parametrize("tied", [False, True]) +def test_full_routing_selection_and_weights(tokens, tied): + torch.manual_seed(910) + logits = torch.randn(tokens, 256, device="cuda", dtype=torch.float32) + bias = torch.randn(256, device="cuda", dtype=torch.float32) * 0.03 + if tied: + logits.zero_() + bias.zero_() + router = _router(bias) + selected, weights = router.apply(logits) + scores = logits.sigmoid() + if tied: + # Exact ties go to the lowest expert ids. + expected = torch.arange(8, device="cuda").expand(tokens, -1) + else: + expected = torch.topk(scores + bias, 8, sorted=True).indices + torch.testing.assert_close(selected.long(), expected, rtol=0, atol=0) + raw = scores.gather(1, expected) + torch.testing.assert_close( + weights, raw / (raw.sum(-1, keepdim=True) + 1e-6), rtol=1e-5, atol=1e-6 + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + got = router.apply(logits) + logits.normal_() + want = router.apply(logits) + graph.replay() + torch.cuda.synchronize() + for actual, reference in zip(got, want): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +def test_full_routing_cpu_fallback(): + logits = torch.randn(2, 256) + bias = torch.randn(256) + router = _router(bias) + ids, weights = router.apply(logits) + scores = logits.sigmoid() + expected = torch.topk(scores + bias, 8, sorted=False).indices + torch.testing.assert_close(ids.long(), expected) + raw = scores.gather(1, expected) + torch.testing.assert_close(weights, raw / (raw.sum(-1, keepdim=True) + 1e-6)) + + +@cuda_only +@pytest.mark.parametrize("tokens", [129, 512]) +def test_batches_past_128_tokens_take_the_reference_selection(tokens): + """Past 128 tokens selection is `torch.topk`; only the weights use a kernel.""" + torch.manual_seed(913) + logits = torch.randn(tokens, 256, device="cuda") + bias = torch.randn(256, device="cuda") * 0.03 + ids, weights = _router(bias).apply(logits) + scores = logits.sigmoid() + expected = torch.topk(scores + bias, 8, sorted=False).indices + torch.testing.assert_close(ids.long(), expected, rtol=0, atol=0) + raw = scores.gather(1, expected) + torch.testing.assert_close( + weights, raw / (raw.sum(-1, keepdim=True) + 1e-6), rtol=1e-6, atol=1e-7 + ) + + +@cuda_only +@pytest.mark.parametrize("tokens", [1, 3, 256, 8192]) +@pytest.mark.parametrize("case", ["random", "tiny", "zero", "ties"]) +def test_the_routing_weights_kernel_matches_the_composition_exactly(tokens, case): + """The weights kernel serves selections the fused routing does not make (over 128 tokens).""" + torch.manual_seed(0) + eps = 1e-15 + logits = torch.randn(tokens, 256, dtype=torch.float32, device="cuda") + if case == "tiny": + logits.fill_(-60.0) + elif case == "ties": + logits.zero_() + scores = torch.sigmoid(logits) + if case == "zero": + scores.zero_() + bias = torch.linspace(-0.2, 0.2, 256, device="cuda") + if case == "ties": + bias.zero_() + _, indices = torch.topk(scores + bias, k=8, dim=-1, sorted=False) + expected = scores.gather(1, indices) + expected = expected / (expected.sum(-1, keepdim=True) + eps) + + kernel = musespark_open_kernels._routing_weights_kernel + with patch.object(kernel, "run", wraps=kernel.run) as run: + selected, weights = _routing_weights(scores, indices, eps, torch.float32) + assert run.call_count == 1, "the composition answered instead of the kernel" + + assert selected.dtype == torch.int32 and weights.dtype == torch.float32 + assert selected.shape == weights.shape == (tokens, 8) + torch.testing.assert_close(selected, indices.to(torch.int32), rtol=0, atol=0) + torch.testing.assert_close(weights, expected, rtol=0, atol=0) + + +@triton.jit +def _reference_argmax_routing_kernel(logits, bias, selected, weights, eps, EXPERTS: tl.constexpr): + row = tl.program_id(0) + expert = tl.arange(0, EXPERTS) + slots = tl.arange(0, 8) + scores = tl.sigmoid(tl.load(logits + row * EXPERTS + expert)) + ranked = scores + tl.load(bias + expert) + ids = tl.full((8,), 0, tl.int32) + values = tl.full((8,), 0.0, tl.float32) + for k in tl.static_range(8): + best = tl.argmax(ranked, axis=0, tie_break_left=True) + value = tl.sum(tl.where(expert == best, scores, 0.0), axis=0) + ids = tl.where(slots == k, best, ids) + values = tl.where(slots == k, value, values) + ranked = tl.where(expert == best, -float("inf"), ranked) + values = tl.div_rn(values, tl.sum(values, axis=0) + eps) + tl.store(selected + row * 8 + slots, ids) + tl.store(weights + row * 8 + slots, values) + + +@cuda_only +@pytest.mark.parametrize("tokens", [1, 2, 16, 128]) +@torch.inference_mode() +def test_sorted_routing_matches_an_argmax_reference_on_mutable_inputs(tokens): + """Repeated argmax with left tie-breaking states the selection order independently.""" + logits = torch.randn(tokens, 256, device="cuda") + bias = torch.randn(256, device="cuda") + router = _router(bias) + router.apply(logits) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = router.apply(logits) + for seed in range(32): + torch.manual_seed(seed) + logits.normal_() + bias.normal_(0, 0.01) + if seed == 0: + logits.zero_() + bias.fill_(-0.5) + if seed == 1: + logits.mul_(1e-6) + bias.zero_() + if seed == 2: + bias.normal_(-1, 0.01) + if seed == 3: + logits.fill_(100) + bias.zero_() + if seed == 4: + logits.floor_() + bias.zero_() + graph.replay() + expected = (torch.empty_like(actual[0]), torch.empty_like(actual[1])) + _reference_argmax_routing_kernel[(tokens,)]( + logits, bias, *expected, 1e-6, EXPERTS=256, num_warps=4 + ) + torch.cuda.synchronize() + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1], expected[1], rtol=1e-6, atol=1e-7) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_fusion_switches.py b/tests/unittest/_torch/modeling/test_musespark_open_fusion_switches.py new file mode 100644 index 000000000000..872a0862cd95 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_fusion_switches.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""When the model builds the fused finalize + all-reduce workspace, and how it degrades. + +Host-side only: the workspace constructor and the device queries are replaced, +so these pin the selection without a GPU or IPC. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.models import modeling_musespark_open as modeling +from tensorrt_llm._torch.models import musespark_open_finalize_ar as finalize_ar + +pytestmark = pytest.mark.cpu_only + +MAPPING = object() + + +@pytest.fixture +def warnings_seen(monkeypatch): + seen = [] + monkeypatch.setattr(modeling.logger, "warning", seen.append) + return seen + + +@pytest.mark.parametrize("error", [ValueError("tp mismatch"), RuntimeError("peer not accessible")]) +def test_a_failed_probe_falls_back_with_a_warning(monkeypatch, warnings_seen, error): + """A group whose GPUs cannot map each other's memory must not fail model construction.""" + + def raise_probe_failure(mapping): + assert mapping is MAPPING + raise error + + monkeypatch.setattr(finalize_ar, "HAS_CUTEDSL", True) + monkeypatch.setattr(finalize_ar, "FinalizeIpcWorkspace", raise_probe_failure) + assert modeling._build_finalize_allreduce_workspace(MAPPING) is None + assert len(warnings_seen) == 1 + assert str(error) in warnings_seen[0] + + +def test_the_workspace_is_built_when_the_probe_succeeds(monkeypatch, warnings_seen): + workspace = object() + monkeypatch.setattr(finalize_ar, "HAS_CUTEDSL", True) + monkeypatch.setattr(finalize_ar, "FinalizeIpcWorkspace", lambda mapping: workspace) + assert modeling._build_finalize_allreduce_workspace(MAPPING) is workspace + assert warnings_seen == [] + + +def test_without_cute_dsl_nothing_is_probed(monkeypatch, warnings_seen): + def must_not_construct(mapping): + raise AssertionError("FinalizeIpcWorkspace must not be constructed without CuTe DSL") + + monkeypatch.setattr(finalize_ar, "HAS_CUTEDSL", False) + monkeypatch.setattr(finalize_ar, "FinalizeIpcWorkspace", must_not_construct) + assert modeling._build_finalize_allreduce_workspace(MAPPING) is None + assert warnings_seen == [] + + +def _mapping(**changes): + values = dict( + tp_size=4, gpus_per_node=4, tp_group=[0, 1, 2, 3], moe_ep_size=4, enable_attention_dp=False + ) + values.update(changes) + return SimpleNamespace(**values) + + +def _layers(expert_hidden_size=4096, top_k=8, comm=None): + moe = SimpleNamespace( + expert_hidden_size=expert_hidden_size, top_k=top_k, experts=SimpleNamespace(comm=comm) + ) + return [SimpleNamespace(feed_forward=moe) for _ in range(3)] + + +@pytest.fixture +def device(monkeypatch): + """A GB300 unless a test says otherwise.""" + state = SimpleNamespace(available=True, capability=(10, 3)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: state.available) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args: state.capability) + return state + + +BF16 = SimpleNamespace(torch_dtype=torch.bfloat16) + + +def test_bf16_tp4_ep4_on_one_gb300_node_is_supported(device): + assert modeling._finalize_allreduce_supported(_mapping(), BF16, _layers()) + + +@pytest.mark.parametrize( + "mapping", + [ + _mapping(tp_size=8, tp_group=list(range(8)), moe_ep_size=8), + _mapping(tp_size=2, tp_group=[0, 1], moe_ep_size=2), + # Four ranks over two nodes. + _mapping(tp_group=[0, 1, 4, 5]), + _mapping(moe_ep_size=1), + _mapping(enable_attention_dp=True), + ], + ids=["tp8", "tp2", "two_nodes", "moe_tp", "attention_dp"], +) +def test_other_parallel_layouts_are_not(device, mapping): + assert not modeling._finalize_allreduce_supported(mapping, BF16, _layers()) + + +@pytest.mark.parametrize( + "layers", + [_layers(expert_hidden_size=2048), _layers(top_k=4), _layers(comm=object())], + ids=["expert_width", "top_k", "moe_comm"], +) +def test_other_expert_shapes_and_communication_are_not(device, layers): + assert not modeling._finalize_allreduce_supported(_mapping(), BF16, layers) + + +def test_other_dtypes_and_devices_are_not(device): + fp16 = SimpleNamespace(torch_dtype=torch.float16) + assert not modeling._finalize_allreduce_supported(_mapping(), fp16, _layers()) + device.capability = (10, 0) + assert not modeling._finalize_allreduce_supported(_mapping(), BF16, _layers()) + device.available = False + assert not modeling._finalize_allreduce_supported(_mapping(), BF16, _layers()) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_moe_finalize.py b/tests/unittest/_torch/modeling/test_musespark_open_moe_finalize.py new file mode 100644 index 000000000000..b9b02c94e45e --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_moe_finalize.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused per-expert post-norm combine against the reference composition. + +On CUDA `post_expert_norm_combine` is one kernel. These check it against the +op-by-op composition at the checkpoint's width and at a narrow one, and pin what +only the kernel can get wrong: a row behind an invalid slot must never be read, +because the MoE kernel leaves it unwritten and a NaN there would reach the sum, +and the kernel must replay from a CUDA graph. +""" + +import pytest +import torch + +from tensorrt_llm._torch.models import musespark_open_kernels +from tensorrt_llm._torch.models.modeling_musespark_open import ( + post_expert_norm_combine, + post_expert_norm_rows, +) + +EPS = 1e-6 +# The checkpoint's expert width and top-k. +HIDDEN = 4096 +TOP_K = 8 + +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference_combine(gemm2, index, routing, weight, eps, *, num_tokens, top_k, hidden_size): + """The op-by-op composition the fused kernel stands for.""" + rows, _ = post_expert_norm_rows( + gemm2, index, weight, eps, num_tokens=num_tokens, top_k=top_k, hidden_size=hidden_size + ) + return (rows * routing.reshape(num_tokens, top_k, 1).to(torch.float32)).sum(dim=1) + + +def _case( + num_tokens, + *, + hidden=HIDDEN, + top_k=TOP_K, + dtype=torch.bfloat16, + invalid=(), + pad=0, + seed=0, + device="cuda", +): + """A finalize input with `invalid` flat slots set to -1 and NaN behind them. + + NaN is what the MoE kernel can leave in the rows of experts that are not + local to this rank: it never writes them. + """ + gen = torch.Generator(device=device).manual_seed(seed) + slots = num_tokens * top_k + gemm2 = torch.randn(slots, hidden + pad, generator=gen, device=device, dtype=torch.float32) + index = torch.arange(slots, dtype=torch.int32, device=device) + for slot in invalid: + index[slot] = -1 + gemm2[slot].fill_(float("nan")) + if pad: + gemm2[:, hidden:].fill_(float("nan")) + routing = torch.rand(num_tokens, top_k, generator=gen, device=device, dtype=torch.float32) + weight = torch.randn(hidden, generator=gen, device=device, dtype=torch.float32) * 0.1 + return gemm2.to(dtype), index, routing, weight.to(dtype) + + +@cuda_only +@pytest.mark.parametrize("num_tokens", [1, 7, 256]) +def test_fused_matches_the_reference_at_the_checkpoint_width(num_tokens): + """The kernel sums in another order, far below the BF16 grain the result is cast to.""" + gemm2, index, routing, weight = _case(num_tokens, invalid=(0, 3)) + kwargs = dict(num_tokens=num_tokens, top_k=TOP_K, hidden_size=HIDDEN) + + fused = post_expert_norm_combine(gemm2, index, routing, weight, EPS, **kwargs) + reference = _reference_combine(gemm2, index, routing, weight, EPS, **kwargs) + + assert torch.isfinite(fused).all() + max_abs = (fused - reference).abs().max().item() + assert max_abs < 1e-4 * reference.abs().max().item() + cosine = torch.nn.functional.cosine_similarity( + fused.reshape(1, -1).double(), reference.reshape(1, -1).double() + ).item() + assert cosine > 1 - 1e-9 + torch.testing.assert_close(fused, reference, rtol=2e-5, atol=2e-5) + + +@cuda_only +def test_fused_never_reads_a_row_behind_an_invalid_slot(): + """A rank owning none of the selected experts contributes exact zeros, not NaN.""" + num_tokens = 16 + gemm2, index, routing, weight = _case(num_tokens, invalid=tuple(range(num_tokens * TOP_K))) + + combined = post_expert_norm_combine( + gemm2, index, routing, weight, EPS, num_tokens=num_tokens, top_k=TOP_K, hidden_size=HIDDEN + ) + + assert torch.equal(combined, torch.zeros_like(combined)) + + +@cuda_only +def test_fused_selects_out_a_permuted_index_past_the_end_of_gemm2(): + """An index past the end of `gemm2` is selected out exactly like a -1 slot.""" + num_tokens, slots = 4, 4 * TOP_K + gemm2, index, routing, weight = _case(num_tokens) + kwargs = dict(num_tokens=num_tokens, top_k=TOP_K, hidden_size=HIDDEN) + + out_of_range = index.clone() + out_of_range[1] = slots + out_of_range[5] = slots + 1024 + as_invalid = index.clone() + as_invalid[1] = -1 + as_invalid[5] = -1 + + got = post_expert_norm_combine(gemm2, out_of_range, routing, weight, EPS, **kwargs) + want = post_expert_norm_combine(gemm2, as_invalid, routing, weight, EPS, **kwargs) + + assert torch.isfinite(got).all() + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=cuda_only)]) +def test_the_padding_columns_of_a_wider_gemm2_row_are_ignored(device): + """The MoE kernel's rows may be padded, and the padding can hold anything.""" + num_tokens = 4 + gemm2, index, routing, weight = _case(num_tokens, pad=64, device=device) + kwargs = dict(num_tokens=num_tokens, top_k=TOP_K, hidden_size=HIDDEN) + + combined = post_expert_norm_combine(gemm2, index, routing, weight, EPS, **kwargs) + + assert combined.shape == (num_tokens, HIDDEN) + assert torch.isfinite(combined).all() + torch.testing.assert_close( + combined, + _reference_combine(gemm2[:, :HIDDEN], index, routing, weight, EPS, **kwargs), + rtol=2e-5, + atol=2e-5, + ) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=cuda_only)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_the_combine_accumulates_in_fp32_whatever_the_expert_output_dtype(device, dtype): + num_tokens = 8 + gemm2, index, routing, weight = _case( + num_tokens, hidden=256, top_k=4, dtype=dtype, device=device + ) + kwargs = dict(num_tokens=num_tokens, top_k=4, hidden_size=256) + + combined = post_expert_norm_combine(gemm2, index, routing.to(dtype), weight, EPS, **kwargs) + + assert combined.dtype is torch.float32 + torch.testing.assert_close( + combined, + _reference_combine(gemm2, index, routing.to(dtype), weight, EPS, **kwargs), + rtol=2e-5, + atol=2e-5, + ) + + +@cuda_only +def test_fused_masks_the_tail_of_a_narrow_row(): + """Eight live lanes of a block padded to the next power of two, and two slots of top-k.""" + num_tokens, hidden, top_k = 4, 8, 2 + gemm2, index, routing, weight = _case( + num_tokens, hidden=hidden, top_k=top_k, dtype=torch.float32, invalid=(3,) + ) + kwargs = dict(num_tokens=num_tokens, top_k=top_k, hidden_size=hidden) + + fused = post_expert_norm_combine(gemm2, index, routing, weight, EPS, **kwargs) + + assert torch.isfinite(fused).all() + torch.testing.assert_close( + fused, + _reference_combine(gemm2, index, routing, weight, EPS, **kwargs), + rtol=2e-5, + atol=2e-5, + ) + + +@cuda_only +def test_fused_finalize_replays_from_a_cuda_graph(): + """Warm up, capture, then replay against fresh inputs.""" + num_tokens = 32 + gemm2, index, routing, weight = _case(num_tokens, invalid=(1, 2)) + kwargs = dict(num_tokens=num_tokens, top_k=TOP_K, hidden_size=HIDDEN) + + post_expert_norm_combine(gemm2, index, routing, weight, EPS, **kwargs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = post_expert_norm_combine(gemm2, index, routing, weight, EPS, **kwargs) + + fresh_gemm2, _, fresh_routing, _ = _case(num_tokens, invalid=(1, 2), seed=5) + gemm2.copy_(fresh_gemm2) + routing.copy_(fresh_routing) + graph.replay() + torch.cuda.synchronize() + + expected = _reference_combine(gemm2, index, routing, weight, EPS, **kwargs) + assert torch.isfinite(static_out).all() + torch.testing.assert_close(static_out, expected, rtol=2e-5, atol=2e-5) + + +@cuda_only +@pytest.mark.parametrize("layout", ["no_tokens", "column_major"]) +def test_unsupported_inputs_take_the_composition(layout, monkeypatch): + def unexpected(*args, **kwargs): + raise AssertionError("an unsupported input reached the fused combine") + + monkeypatch.setattr(musespark_open_kernels, "post_expert_norm_combine", unexpected) + num_tokens = 0 if layout == "no_tokens" else 4 + gemm2, index, routing, weight = _case(num_tokens, dtype=torch.float32) + if layout == "column_major": + gemm2 = gemm2.t().contiguous().t() + kwargs = dict(num_tokens=num_tokens, top_k=TOP_K, hidden_size=HIDDEN) + + combined = post_expert_norm_combine(gemm2, index, routing, weight, EPS, **kwargs) + + assert combined.shape == (num_tokens, HIDDEN) + torch.testing.assert_close( + combined, _reference_combine(gemm2, index, routing, weight, EPS, **kwargs) + ) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_moe_prepare_overlap.py b/tests/unittest/_torch/modeling/test_musespark_open_moe_prepare_overlap.py new file mode 100644 index 000000000000..cf4e788ec9f5 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_moe_prepare_overlap.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small batches pool on the auxiliary stream while the router runs, eagerly and in graphs.""" + +import contextlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenMoE + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _module(*, stream=True, comm=None): + mod = MuseSparkOpenMoE.__new__(MuseSparkOpenMoE) + torch.nn.Module.__init__(mod) + mod._prepare_stream = torch.cuda.Stream() if stream else None + mod._prepare_fork_event = torch.cuda.Event() + mod._prepare_ready_event = torch.cuda.Event() + mod.nvtx = lambda name: contextlib.nullcontext() + mod.experts = SimpleNamespace(comm=comm) + router_weight = torch.randn(256, 256, device="cuda") + mod.router_logits = lambda x: x.float() @ router_weight + mod.pooler = torch.nn.Linear(256, 128, bias=False, dtype=torch.bfloat16, device="cuda") + mod.prepare_expert_input = lambda x: x * 0.5 + return mod + + +def _serial(mod, x): + return mod.router_logits(x), mod.prepare_expert_input(mod.pooler(x)) + + +def _prepare(mod, x): + router_logits, expert_input, ready = mod.prepare_moe_inputs(x) + if ready is not None: + # The MoE joins the event before its first read of the expert input. + ready.wait() + return router_logits, expert_input, ready + + +@pytest.mark.parametrize("tokens", [1, 2, 8, 16, 80, 128, 129]) +def test_prepare_overlap_matches_serial_and_graph_replay(tokens): + torch.manual_seed(923) + with torch.inference_mode(): + mod = _module() + x = torch.randn(tokens, 256, device="cuda", dtype=torch.bfloat16) + expected = _serial(mod, x) + router_logits, expert_input, ready = _prepare(mod, x) + assert (ready is not None) == (tokens <= 128) + torch.cuda.synchronize() + for got, want in zip((router_logits, expert_input), expected): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = _prepare(mod, x)[:2] + for _ in range(3): + x.normal_() + expected = _serial(mod, x) + graph.replay() + torch.cuda.synchronize() + for got, want in zip(captured, expected): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +@pytest.mark.parametrize("reason", ["no_stream", "moe_comm", "no_tokens"]) +def test_the_overlap_stays_off_where_it_does_not_apply(reason): + with torch.inference_mode(): + mod = _module(stream=reason != "no_stream", comm=object() if reason == "moe_comm" else None) + x = torch.randn(0 if reason == "no_tokens" else 4, 256, device="cuda", dtype=torch.bfloat16) + router_logits, expert_input, ready = mod.prepare_moe_inputs(x) + assert ready is None + for got, want in zip((router_logits, expert_input), _serial(mod, x)): + torch.testing.assert_close(got, want, rtol=0, atol=0) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_norm_gate.py b/tests/unittest/_torch/modeling/test_musespark_open_norm_gate.py new file mode 100644 index 000000000000..2227d9a54ccf --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_norm_gate.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused attention output norm and gate: rounding, strided gates, fallbacks and graph replay.""" + +import pytest +import torch + +from tensorrt_llm._torch.models import musespark_open_kernels +from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenAttention +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _attention(): + mod = MuseSparkOpenAttention.__new__(MuseSparkOpenAttention) + torch.nn.Module.__init__(mod) + mod.head_dim = 64 + mod._fuse_qk_norm_rope_gate = False + mod.o_head_norm = RMSNorm( + hidden_size=64, eps=1e-5, dtype=torch.bfloat16, has_weights=False + ).cuda() + return mod + + +def _reference(mod, x, gate): + return mod.o_head_norm(x.reshape(-1, 64)).reshape_as(x) * gate.reshape_as(x).sigmoid() + + +def _inputs(tokens, strided): + x = torch.randn(tokens, 2048, device="cuda", dtype=torch.bfloat16) + if strided: + # The gate half of a per-head interleaved projection. + gate = torch.randn(tokens, 32, 128, device="cuda", dtype=torch.bfloat16)[:, :, 64:] + else: + gate = torch.randn_like(x) + return x, gate + + +@pytest.mark.parametrize("tokens", [1, 2, 16, 128]) +@pytest.mark.parametrize("strided", [False, True]) +@torch.inference_mode() +def test_norm_gate_and_graph_replay(tokens, strided): + torch.manual_seed(950) + mod = _attention() + x, gate = _inputs(tokens, strided) + got = mod.apply_output_gate(x, gate) + expected = _reference(mod, x, gate) + # The fused norm reduces in another order, which can move single values + # across a BF16 rounding boundary. + torch.testing.assert_close(got, expected, rtol=0.016, atol=0.002) + assert (got.float() - expected.float()).norm() / expected.float().norm() < 0.003 + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + got = mod.apply_output_gate(x, gate) + x.normal_() + gate.normal_() + eager = mod.apply_output_gate(x, gate) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(got, eager, rtol=0, atol=0) + reference = _reference(mod, x, gate) + torch.testing.assert_close(got, reference, rtol=0.016, atol=0.002) + + +@pytest.mark.parametrize("tokens", [129, 512]) +@torch.inference_mode() +def test_unsupported_batches_use_the_reference(tokens, monkeypatch): + def unexpected(*args, **kwargs): + raise AssertionError("an unsupported batch reached the fused norm and gate") + + monkeypatch.setattr(musespark_open_kernels, "attention_norm_gate", unexpected) + torch.manual_seed(951) + mod = _attention() + x, gate = _inputs(tokens, strided=True) + torch.testing.assert_close(mod.apply_output_gate(x, gate), _reference(mod, x, gate)) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_pre_expert_quant.py b/tests/unittest/_torch/modeling/test_musespark_open_pre_expert_quant.py new file mode 100644 index 000000000000..c2384dcbf21b --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_pre_expert_quant.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The pre-expert norm fused with the experts' NVFP4 input quantization: gating, layout, graphs.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.kernels.fused_ops.rmsnorm_fp4_quant import rmsnorm_fp4_quant_available +from tensorrt_llm._torch.models import modeling_musespark_open +from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenMoE +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.moe.fused_moe.moe_op_backend import TRTLLMOpBackend +from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.fp4_block_scale import TRTLLMGenNvfp4Base +from tensorrt_llm._torch.utils import Fp4QuantizedTensor + +fused_available = pytest.mark.skipif( + not torch.cuda.is_available() or not rmsnorm_fp4_quant_available(), + reason="CUDA and FlashInfer's fused norm and quantization required", +) + + +def _module(device, comm=None): + module = MuseSparkOpenMoE.__new__(MuseSparkOpenMoE) + torch.nn.Module.__init__(module) + module.register_buffer("_pre_expert_quant_gain", None, persistent=False) + module.pre_expert_norm = RMSNorm( + hidden_size=4096, eps=1e-5, dtype=torch.bfloat16, device=device, use_gemma=True + ) + backend = SimpleNamespace( + op_backend=TRTLLMOpBackend(), + has_nvfp4=True, + fc31_input_scale=torch.tensor([448.0], dtype=torch.float32, device=device), + ) + module.experts = SimpleNamespace(backend=backend, comm=comm) + return module + + +@pytest.mark.parametrize("comm", [None, object()]) +def test_cpu_inputs_take_the_norm(comm): + module = _module("cpu", comm) + x = torch.randn(2, 4096, dtype=torch.float32) + torch.testing.assert_close(module.prepare_expert_input(x), module.pre_expert_norm(x)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("reason", ["comm", "awq", "chunks", "backend", "quant", "scale"]) +def test_unsupported_consumers_take_the_norm(reason): + module = _module("cuda") + backend = module.experts.backend + if reason == "comm": + module.experts.comm = object() + elif reason == "awq": + backend.fc31_act_scale = torch.ones(4096, device="cuda") + elif reason == "backend": + backend.op_backend = object() + elif reason == "quant": + backend.has_nvfp4 = False + elif reason == "scale": + backend.fc31_input_scale = torch.ones(2, device="cuda") + else: + module.experts.moe_max_num_tokens = 1 + x = torch.randn(2, 4096, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close(module.prepare_expert_input(x), module.pre_expert_norm(x)) + + +@fused_available +def test_without_the_fused_kernel_the_norm_is_kept(monkeypatch): + with torch.inference_mode(): + module = _module("cuda") + x = torch.randn(2, 4096, device="cuda", dtype=torch.bfloat16) + assert isinstance(module.prepare_expert_input(x), Fp4QuantizedTensor) + monkeypatch.setattr(modeling_musespark_open, "rmsnorm_fp4_quant_available", lambda: False) + torch.testing.assert_close(module.prepare_expert_input(x), module.pre_expert_norm(x)) + + +@fused_available +def test_the_moe_takes_the_layout_and_graphs_replay_it(): + torch.manual_seed(812) + with torch.inference_mode(): + module = _module("cuda") + x = torch.randn(2, 4096, device="cuda", dtype=torch.bfloat16) + output = module.prepare_expert_input(x) + assert isinstance(output, Fp4QuantizedTensor) + assert not output.is_sf_swizzled + # The quantized-input branch reads nothing off the module, so the + # namespace stands in for the NVFP4 experts. + packed, scales = TRTLLMGenNvfp4Base.quantize_input(module.experts.backend, output) + assert packed.data_ptr() == output.fp4_tensor.data_ptr() + assert scales.data_ptr() == output.scaling_factor.data_ptr() + assert packed.shape == (2, 2048) + assert scales.shape == (2, 256) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = module.prepare_expert_input(x) + x.normal_() + expected = module.prepare_expert_input(x) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured.fp4_tensor, expected.fp4_tensor, rtol=0, atol=0) + torch.testing.assert_close(captured.scaling_factor, expected.scaling_factor, rtol=0, atol=0) + + +@fused_available +def test_the_cached_gain_matches_the_one_computed_per_call(): + torch.manual_seed(813) + with torch.inference_mode(): + module = _module("cuda") + module.pre_expert_norm.weight.normal_(0, 0.1) + x = torch.randn(4, 4096, device="cuda", dtype=torch.bfloat16) + uncached = module.prepare_expert_input(x) + module.cache_derived_state() + weight = module.pre_expert_norm.weight + torch.testing.assert_close( + module._pre_expert_quant_gain, (weight.float() + 1).to(weight.dtype), rtol=0, atol=0 + ) + cached = module.prepare_expert_input(x) + torch.testing.assert_close(cached.fp4_tensor, uncached.fp4_tensor, rtol=0, atol=0) + torch.testing.assert_close(cached.scaling_factor, uncached.scaling_factor, rtol=0, atol=0) + + +@pytest.mark.parametrize("stage", ["post_load_weights", "cache_derived_state"]) +def test_the_gain_cache_follows_the_weight(stage): + """Filled by either loader walk, skipped while meta, dropped by a dtype move.""" + with torch.device("meta"): + module = _module("meta") + getattr(module, stage)() + assert module._pre_expert_quant_gain is None + + module = _module("cpu") + getattr(module, stage)() + assert module._pre_expert_quant_gain is not None + module.to(torch.float32) + assert module._pre_expert_quant_gain is None diff --git a/tests/unittest/_torch/modeling/test_musespark_open_router_projection.py b/tests/unittest/_torch/modeling/test_musespark_open_router_projection.py new file mode 100644 index 000000000000..682f9fe52019 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_router_projection.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small-batch FP32 router projection, including mutable CUDA-graph inputs.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenMoE + + +def _module(device, hidden=8192, experts=256): + module = MuseSparkOpenMoE.__new__(MuseSparkOpenMoE) + torch.nn.Module.__init__(module) + module.router = torch.nn.Linear(hidden, experts, bias=False, device=device, dtype=torch.float32) + module.metap_multiplier = 0.375 + module._prepare_stream = None + module.experts = SimpleNamespace(comm=None) + return module + + +def _fp64_reference(module, x): + return (x.double() @ module.router.weight.double().t() * module.metap_multiplier).float() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tokens", [1, 2, 3, 8, 15, 16, 17, 32, 48, 80, 128]) +@torch.no_grad() +def test_router_matches_fp64_and_replays_changed_inputs(tokens): + """Both the CUDA-core (under 16 tokens) and the tensor-core kernel.""" + torch.manual_seed(1263 + tokens) + module = _module("cuda") + x = torch.randn(tokens, 8192, device="cuda", dtype=torch.bfloat16) + calls = [] + hook = module.router.register_forward_hook(lambda *_: calls.append(True)) + module.router_logits(x) + hook.remove() + assert calls == [], "a supported batch fell back to the Linear" + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = module.router_logits(x) + for factor in (1.0, -0.5, 2.0): + x.normal_().mul_(factor) + # The replay must read the live weights too, not a cast or a cache. + module.router.weight.mul_(-0.75) + graph.replay() + torch.testing.assert_close(output, _fp64_reference(module, x), rtol=2e-5, atol=2e-6) + assert output.dtype == torch.float32 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tokens", range(1, 17)) +@torch.no_grad() +def test_overlap_tuned_router_reads_mutable_graph_inputs(tokens: int) -> None: + """The launch shapes used while the pooler runs on the auxiliary stream.""" + torch.manual_seed(1291 + tokens) + module = _module("cuda") + module._prepare_stream = torch.cuda.Stream() + x = torch.randn(tokens, 8192, device="cuda", dtype=torch.bfloat16) + module.router_logits(x) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = module.router_logits(x) + for factor in (1.0, -0.5, 2.0): + x.normal_().mul_(factor) + module.router.weight.mul_(-0.75) + graph.replay() + torch.testing.assert_close(output, _fp64_reference(module, x), rtol=2e-5, atol=2e-6) + + +@torch.no_grad() +def test_cpu_fallback_calls_linear(): + module = _module("cpu", hidden=16, experts=8) + x = torch.randn(2, 16, dtype=torch.bfloat16) + calls = [] + hook = module.router.register_forward_hook(lambda *_: calls.append(True)) + actual = module.router_logits(x) + hook.remove() + assert calls == [True] + expected = torch.nn.functional.linear(x.float(), module.router.weight) * module.metap_multiplier + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tokens,noncontiguous", [(0, False), (129, False), (16, True)]) +@torch.no_grad() +def test_unsupported_gpu_inputs_use_linear(tokens, noncontiguous): + module = _module("cuda") + width = 16384 if noncontiguous else 8192 + x = torch.randn(tokens, width, device="cuda", dtype=torch.bfloat16) + if noncontiguous: + x = x[:, ::2] + calls = [] + hook = module.router.register_forward_hook(lambda *_: calls.append(True)) + actual = module.router_logits(x) + hook.remove() + assert calls == [True] + expected = torch.nn.functional.linear(x.float(), module.router.weight) * module.metap_multiplier + torch.testing.assert_close(actual, expected, rtol=0, atol=0) From 46c9aa13eadb7cc22e7773b16b1c97b8a23d489b Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 2 Nov 2026 13:01:01 -0800 Subject: [PATCH 2/3] [TRTLLM-17146][fix] Pass the router GEMM's token count at run time; check the cluster kernel runs in its tests - The tensor-core router kernels take the token count as a run-time argument instead of a constexpr, so token counts from 16 to 128 no longer compile one kernel each. - The cluster boundary's rounding/graph and side-stream tests record the calls that reach the cluster kernel, so they fail if the model falls back to the Triton path for their inputs. Signed-off-by: Michal Guzek --- .../_torch/models/musespark_open_kernels.py | 20 ++++++++-------- .../test_musespark_open_cluster_boundary.py | 23 +++++++++++++++++++ 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/models/musespark_open_kernels.py b/tensorrt_llm/_torch/models/musespark_open_kernels.py index 223c181c36e5..6dd3aee8d32d 100644 --- a/tensorrt_llm/_torch/models/musespark_open_kernels.py +++ b/tensorrt_llm/_torch/models/musespark_open_kernels.py @@ -353,7 +353,7 @@ def _router_tensorcore_kernel( hidden_states, weight, partials, - TOKENS: tl.constexpr, + num_tokens, ): # Sixteen tokens share each weight tile. The 8192-wide dot product is split # into sixteen 512-wide chunks whose FP32 partials the reduce kernel sums; @@ -366,28 +366,28 @@ def _router_tensorcore_kernel( hidden = split * 512 + block * 128 + tl.arange(0, 128) x = tl.load( hidden_states + tokens[:, None] * 8192 + hidden[None, :], - tokens[:, None] < TOKENS, + tokens[:, None] < num_tokens, other=0.0, ).to(tl.float32) w = tl.load(weight + experts[None, :] * 8192 + hidden[:, None]) acc = tl.dot(x, w, acc, input_precision="tf32x3") tl.store( - partials + split * TOKENS * 256 + tokens[:, None] * 256 + experts[None, :], + partials + split * num_tokens * 256 + tokens[:, None] * 256 + experts[None, :], acc, - tokens[:, None] < TOKENS, + tokens[:, None] < num_tokens, ) @triton.jit -def _router_tensorcore_reduce_kernel(partials, output, scale, SIZE: tl.constexpr): +def _router_tensorcore_reduce_kernel(partials, output, scale, size): offsets = tl.program_id(0) * 128 + tl.arange(0, 128) splits = tl.arange(0, 16) values = tl.load( - partials + splits[:, None] * SIZE + offsets[None, :], - offsets[None, :] < SIZE, + partials + splits[:, None] * size + offsets[None, :], + offsets[None, :] < size, other=0.0, ) - tl.store(output + offsets, tl.sum(values, axis=0) * scale, offsets < SIZE) + tl.store(output + offsets, tl.sum(values, axis=0) * scale, offsets < size) def router_projection_supported(hidden_states: torch.Tensor, weight: torch.Tensor) -> bool: @@ -436,10 +436,10 @@ def router_projection( if tokens >= 16: partials = torch.empty((16, tokens, 256), device=hidden_states.device, dtype=torch.float32) _router_tensorcore_kernel[(triton.cdiv(tokens, 16), 8, 16)]( - hidden_states, weight, partials, TOKENS=tokens, num_warps=4 + hidden_states, weight, partials, num_tokens=tokens, num_warps=4 ) _router_tensorcore_reduce_kernel[(triton.cdiv(tokens * 256, 128),)]( - partials, output, scale, SIZE=tokens * 256, num_warps=4 + partials, output, scale, size=tokens * 256, num_warps=4 ) return output rows, warps = _router_launch(tokens, overlapped) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py b/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py index 0e57f08dd020..dbcd7efcdf22 100644 --- a/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py +++ b/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py @@ -15,6 +15,25 @@ ) +def _record_cluster_launches(monkeypatch: pytest.MonkeyPatch) -> list: + """Record the token count of every call that reaches the cluster kernel. + + The model falls back silently for inputs the kernel rejects, so a test + that compares against the fallback must show the kernel actually ran. + """ + from tensorrt_llm._torch.models import musespark_open_boundary + + launches = [] + cluster_boundary = musespark_open_boundary.cluster_boundary + + def recorded(residual, *args): + launches.append(residual.shape[0]) + return cluster_boundary(residual, *args) + + monkeypatch.setattr(musespark_open_boundary, "cluster_boundary", recorded) + return launches + + @pytest.mark.parametrize("tokens", [1, 2, 3, 4, 8, 16, 17, 32]) @pytest.mark.parametrize("with_gain", [False, True]) @pytest.mark.parametrize("eps", [1e-6, 1e-5]) @@ -34,11 +53,13 @@ def test_cluster_boundary_rounding_and_mutable_graph( if gain is not None: gain.normal_(1, 0.1) + launches = _record_cluster_launches(monkeypatch) # Compile on the eager stream, then capture on the graph stream. gate.add_after_rmsnorm_and_norm(residual, branch, eps, norm, gain) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): actual = gate.add_after_rmsnorm_and_norm(residual, branch, eps, norm, gain) + assert launches == [tokens, tokens] for scale in (1e-3, 1.0, 1e3): residual.normal_().mul_(scale) branch.normal_().mul_(scale) @@ -144,10 +165,12 @@ def test_cluster_boundary_side_stream_and_gate_cache( gate.cache_derived_state() norm = RMSNorm(hidden_size=8192, eps=1e-5, dtype=torch.bfloat16, use_gemma=True).cuda() norm.weight.normal_(0, 0.02) + launches = _record_cluster_launches(monkeypatch) gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm, gain) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=stream): actual = gate.add_after_rmsnorm_and_norm(residual, branch, 1e-5, norm, gain) + assert launches == [tokens, tokens] for _ in range(8): residual.normal_() branch.normal_() From 78e99e0cfefa0407d919d33c59daf9a4a597d1bb Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 2 Nov 2026 13:01:02 -0800 Subject: [PATCH 3/3] [TRTLLM-17146][fix] Compile the cluster boundary once per cluster size; list the fast-path tests on B200 and GB300 multi-GPU - The cluster boundary kernel takes the token count at launch, so only the three cluster sizes (with the gain and the two epsilons) compile, instead of one compilation per token count from 1 to 32. - l0_gb300_multi_gpus: the three four-rank finalize + all-reduce files. - l0_b200: the six Triton fast-path files, which need only CUDA. - Document that the workspace construction is the constructor's first TP collective, and that the model never calls close(). Signed-off-by: Michal Guzek --- .../_torch/models/modeling_musespark_open.py | 5 +++ .../_torch/models/musespark_open_boundary.py | 31 +++++++++++++------ .../models/musespark_open_finalize_ar.py | 4 +++ .../test_lists/test-db/l0_b200.yml | 11 +++++-- .../test-db/l0_gb300_multi_gpus.yml | 5 +++ 5 files changed, 45 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open.py b/tensorrt_llm/_torch/models/modeling_musespark_open.py index 0ddf1303599f..ef6903934ac6 100644 --- a/tensorrt_llm/_torch/models/modeling_musespark_open.py +++ b/tensorrt_llm/_torch/models/modeling_musespark_open.py @@ -1157,6 +1157,11 @@ def _build_finalize_allreduce_workspace(mapping: Mapping) -> Optional["FinalizeI other's memory or the ranks disagree. That is a property of the deployment, so it degrades to the unfused path instead of failing model construction. + + It is also the first TP collective in the model's constructor: every rank + must reach it. A rank that fails earlier in the constructor leaves its + peers blocked here until the launcher stops them, instead of failing with + the same error. """ from .musespark_open_finalize_ar import HAS_CUTEDSL, FinalizeIpcWorkspace diff --git a/tensorrt_llm/_torch/models/musespark_open_boundary.py b/tensorrt_llm/_torch/models/musespark_open_boundary.py index 9b52d1625e82..2aa235a05b9e 100644 --- a/tensorrt_llm/_torch/models/musespark_open_boundary.py +++ b/tensorrt_llm/_torch/models/musespark_open_boundary.py @@ -61,14 +61,13 @@ def _float_bits(value, *, loc=None, ip=None): class _ClusterBoundaryKernel: def __init__( self, - m: int, cluster: int, threads: int, gain: bool, post_eps: float, pre_eps: float, ): - self.m, self.cluster, self.threads, self.gain = m, cluster, threads, gain + self.cluster, self.threads, self.gain = cluster, threads, gain self.post_eps, self.pre_eps = float(post_eps), float(pre_eps) self.iters = 8192 // (cluster * threads * 8) self.warps = threads // 32 @@ -85,10 +84,11 @@ def __call__( weight: cute.Pointer, output: cute.Pointer, next_output: cute.Pointer, + tokens: Int32, stream: cuda.CUstream, ): self.kernel(branch, residual, gain, alpha, beta, weight, output, next_output).launch( - grid=[self.m * self.cluster, 1, 1], + grid=[tokens * self.cluster, 1, 1], block=[self.threads, 1, 1], cluster=[self.cluster, 1, 1], smem=self.slots * 4 * 4 * 2 + 16, @@ -232,20 +232,28 @@ def kernel( cutlass.BFloat16, ) - @functools.lru_cache(maxsize=128) + def _cluster_size(tokens: int) -> int: + """CTAs per token: fewer tokens spread each one over more CTAs.""" + return 16 if tokens == 1 else (8 if tokens <= 8 else 4) + + # The token count is a launch argument, so only the three cluster sizes, + # the gain and the two epsilons select a compilation. The executor's warmup + # runs each CUDA graph batch size eagerly before capturing it, so a capture + # never compiles. + @functools.lru_cache(maxsize=32) def _compile_boundary( - tokens: int, has_gain: bool, post_eps: float, pre_eps: float, device_index: int + cluster: int, has_gain: bool, post_eps: float, pre_eps: float, device_index: int ): - cluster = 16 if tokens == 1 else (8 if tokens <= 8 else 4) threads = min(256, 8192 // (cluster * 8)) with torch.cuda.device(device_index): stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) return cute.compile( - _ClusterBoundaryKernel(tokens, cluster, threads, has_gain, post_eps, pre_eps), + _ClusterBoundaryKernel(cluster, threads, has_gain, post_eps, pre_eps), *[ make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=16) for dtype in _DTYPES ], + Int32(1), stream, options="--enable-tvm-ffi", ) @@ -316,12 +324,17 @@ def cluster_boundary( raise ValueError("Unsupported cluster boundary inputs, device, or alignment") gain_arg = next_weight if gain is None else gain updated, normalized = torch.empty_like(residual), torch.empty_like(branch) + tokens = residual.shape[0] compiled = _compile_boundary( - residual.shape[0], gain is not None, post_eps, pre_eps, residual.device.index + _cluster_size(tokens), gain is not None, post_eps, pre_eps, residual.device.index ) args = (branch, residual, gain_arg, alpha, beta, next_weight, updated, normalized) # Resolve the stream per call: a CUDA graph captures on a different stream # from the eager call that filled the compilation cache. stream = cuda.CUstream(torch.cuda.current_stream(residual.device).cuda_stream) - compiled(*[make_ptr(dtype, tensor.data_ptr()) for dtype, tensor in zip(_DTYPES, args)], stream) + compiled( + *[make_ptr(dtype, tensor.data_ptr()) for dtype, tensor in zip(_DTYPES, args)], + tokens, + stream, + ) return updated, normalized diff --git a/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py b/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py index ae9b6515100f..be9cd5dc397e 100644 --- a/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py +++ b/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py @@ -886,6 +886,10 @@ def close(self) -> None: waits for this rank's device and then closes this rank's mappings of peer memory. It never frees this rank's own buffer, which peers may still map; that stays allocated until the process exits. + + The model never calls `close()`: the workspace lives as long as the + process serving the model, so a process that builds the model several + times keeps one 8 MiB buffer per build until it exits. """ if self._closed: return diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index c78b989b906a..595f6b382906 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -341,11 +341,18 @@ l0_b200: - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_e4b_dummy - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_multimodal_26b_dummy - # Muse Spark Open: the decoder's GPU module tests and the tower load; the - # remaining Muse files are cpu_only and run in l0_cpu. + # Muse Spark Open: the decoder's GPU module tests, the tower load and the + # decoder's Triton fast paths, which need only CUDA. The GB300 kernels' tests + # run in l0_gb300 and l0_gb300_multi_gpus, the cpu_only files in l0_cpu. - unittest/_torch/modeling/test_modeling_musespark_open.py - unittest/_torch/modeling/test_musespark_open_moe_tp.py - unittest/_torch/modeling/test_musespark_open_tower_load.py + - unittest/_torch/modeling/test_musespark_open_router_projection.py + - unittest/_torch/modeling/test_musespark_open_full_routing.py + - unittest/_torch/modeling/test_musespark_open_norm_gate.py + - unittest/_torch/modeling/test_musespark_open_moe_finalize.py + - unittest/_torch/modeling/test_musespark_open_boundary_fusion.py + - unittest/_torch/modeling/test_musespark_open_moe_prepare_overlap.py - unittest/tools/test_layer_wise_benchmarks.py::test_deepseek_r1_ctx_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_nemotron_gen_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_qwen3_next_gen_tep[1] diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 97d9ebd7ac45..34dd0dedce03 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -69,6 +69,11 @@ l0_gb300_multi_gpus: - unittest/_torch/multi_gpu/test_ll_exchange.py - unittest/_torch/modeling -k "modeling_llama" - unittest/_torch/modeling -k "modeling_gpt_oss" + # Muse Spark Open's fused expert finalize + TP4 all-reduce over CUDA IPC; each + # test spawns its four ranks with MPIPoolExecutor. + - unittest/_torch/modeling/test_musespark_open_finalize_allreduce.py + - unittest/_torch/modeling/test_musespark_open_finalize_lifecycle.py + - unittest/_torch/modeling/test_musespark_open_finalize_stress.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py # KDA runtime state-pool and mixed-batch core coverage.