diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open.py b/tensorrt_llm/_torch/models/modeling_musespark_open.py index 5338c1a919a7..ef6903934ac6 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,94 @@ 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. + + 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 + + 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 +1206,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 +1237,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..2aa235a05b9e --- /dev/null +++ b/tensorrt_llm/_torch/models/musespark_open_boundary.py @@ -0,0 +1,340 @@ +# 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, + cluster: int, + threads: int, + gain: bool, + post_eps: float, + pre_eps: float, + ): + 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 + 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, + tokens: Int32, + stream: cuda.CUstream, + ): + self.kernel(branch, residual, gain, alpha, beta, weight, output, next_output).launch( + grid=[tokens * 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, + ) + + 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( + cluster: int, has_gain: bool, post_eps: float, pre_eps: float, device_index: int + ): + 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(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", + ) + + +@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) + tokens = residual.shape[0] + compiled = _compile_boundary( + _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)], + 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 new file mode 100644 index 000000000000..be9cd5dc397e --- /dev/null +++ b/tensorrt_llm/_torch/models/musespark_open_finalize_ar.py @@ -0,0 +1,983 @@ +# 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. + + 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 + 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..6dd3aee8d32d --- /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, + 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; + # 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] < 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 * num_tokens * 256 + tokens[:, None] * 256 + experts[None, :], + acc, + tokens[:, None] < num_tokens, + ) + + +@triton.jit +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, + 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, num_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_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.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/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. 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..dbcd7efcdf22 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_cluster_boundary.py @@ -0,0 +1,224 @@ +# 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", +) + + +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]) +@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) + + 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) + 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) + 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_() + 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)