[Model] Support hybrid MegaMoE/TRT-LLM execution for Muse Spark NVFP4

cmake/external_projects/deepgemm.cmake

@@ -35,9 +35,10 @@ else()
   # SFs for the FP8 einsum output (vllm-project/DeepGEMM#40), plus
   # fp4_fp4_mega_moe and fused post-expert RMSNorm for Muse Spark Open
   # (vllm-project/DeepGEMM#39), and the MX MegaMoE alignment fix
-  # (vllm-project/DeepGEMM#43).
+  # (vllm-project/DeepGEMM#43), and the TRT-LLM NVFP4 weight layout
+  # for Muse Spark MegaMoE (vllm-project/DeepGEMM#47).
   # Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API.
-  set(_DEEPGEMM_UPSTREAM_TAG "0820030756945b2daec208e8e8255268360f337d")
+  set(_DEEPGEMM_UPSTREAM_TAG "5cdb145704ec4672ecf4b61441599521ce635df6")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)

tests/kernels/moe/test_muse_spark_trtllm_moe.py

@@ -432,3 +432,98 @@ def test_muse_spark_moe_cudagraph(nvfp4, dist_init, workspace_init):
         ckpt, quant_config = _bf16_checkpoint(), None
     _, _, _, vllm_config, runner = build_moe(MuseSpark1xRMSNorm, ckpt, quant_config)
     assert_cudagraph_replays(runner, vllm_config, *_inputs(5), *_inputs(6))
+
+
+@torch.inference_mode()
+def test_megamoe_trtllm_weight_layout_matches_native(dist_init, workspace_init):
+    """Shared TRT storage must preserve native MegaMoE outputs and graph replay."""
+    from vllm.models.muse_spark.nvidia.mega_moe import _get_deep_gemm
+    from vllm.models.muse_spark.nvidia.ops.prepare_megamoe_nvfp4 import (
+        prepare_megamoe_nvfp4_inputs,
+    )
+
+    torch.manual_seed(137)
+    ckpt = _nvfp4_checkpoint()
+    norm = MuseSpark1xRMSNormWithInputScale(
+        hidden_size=K,
+        eps=1e-8,
+        has_weight=True,
+        is_post_norm=True,
+        zero_centered_gamma=True,
+    ).cuda()
+    norm.weight.copy_(torch.linspace(-0.5, 0.5, K, device="cuda"))
+    layers = []
+    for layout, threshold in (("native", 0), ("trtllm", 16)):
+        config = VllmConfig(
+            additional_config={"muse_spark_trtllm_token_threshold": threshold}
+        )
+        with set_current_vllm_config(config):
+            layer = MuseSparkMegaMoEExperts(
+                config,
+                num_experts=E,
+                num_local_experts=E,
+                experts_start_idx=0,
+                top_k=TOPK,
+                hidden_size=K,
+                intermediate_size=N,
+                post_expert_norm=norm,
+                num_max_tokens_per_rank=128,
+                prefix=f"shared_layout.{layout}",
+            ).cuda()
+            assert layer.weight_layout == layout
+            assert set(layer.load_weights(ckpt.items())) == set(ckpt)
+            layer.finalize_weights()
+            layers.append(layer)
+    native, shared = layers
+    for view, parameter in (
+        (shared._transformed_l1[0], shared.w13_weight),
+        (shared._transformed_l1[1], shared.w13_weight_scale),
+        (shared._transformed_l2[0], shared.w2_weight),
+        (shared._transformed_l2[1], shared.w2_weight_scale),
+    ):
+        assert view.data_ptr() == parameter.data_ptr()
+        assert view.untyped_storage().nbytes() == parameter.untyped_storage().nbytes()
+    shared.finalize_weights()  # Must not shuffle the already finalized weights again.
+
+    dg = _get_deep_gemm()
+    buffer = dg.SymmBuffer(
+        torch.distributed.group.WORLD, E, 128, TOPK, K, N, mma_type="fp4xfp4"
+    )
+    try:
+        for tokens in (1, 8, 32, 64):
+            x = torch.randn(tokens, K, device="cuda", dtype=torch.bfloat16) / 10
+            weights, ids = torch.randn(tokens, E, device="cuda").topk(TOPK)
+            weights = weights.softmax(-1)
+            ids = ids.int()
+            ids[::3, -1] = -1
+            buffer.topk_idx.fill_(-1)
+            buffer.topk_weights.zero_()
+
+            def stage(x=x, weights=weights, ids=ids, tokens=tokens):
+                prepare_megamoe_nvfp4_inputs(
+                    x,
+                    weights,
+                    ids,
+                    buffer.x[:tokens],
+                    buffer.x_sf[:tokens].view(torch.uint8),
+                    buffer.topk_idx[:tokens],
+                    buffer.topk_weights[:tokens],
+                    None,
+                )
+
+            stage()
+            expected = native.forward_prepared(x, buffer)
+            actual = shared.forward_prepared(x, buffer)
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = shared.forward_prepared(x, buffer)
+            for _ in range(3):
+                x.normal_(std=0.1)
+                ids.copy_(torch.randn(tokens, E, device="cuda").topk(TOPK).indices)
+                stage()
+                expected = native.forward_prepared(x, buffer)
+                graph.replay()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    finally:
+        buffer.destroy()

vllm/models/common/ops/sequence_parallel.py

@@ -23,13 +23,17 @@ def _custom_collective(name: str, x: torch.Tensor) -> torch.Tensor | None:
 
 
 @lru_cache
-def is_sequence_parallel_active(is_enabled: bool, num_tokens: int) -> bool:
+def is_sequence_parallel_active(
+    is_enabled: bool, num_tokens: int, min_tokens: int = 0
+) -> bool:
     """Whether a batch of ``num_tokens`` shards the stream along the sequence.
 
     Below one token per rank the shards are mostly padding, so the batch
-    stays replicated.
+    stays replicated. A model may raise this lower bound with ``min_tokens``.
     """
-    return is_enabled and num_tokens >= get_tensor_model_parallel_world_size()
+    return is_enabled and num_tokens >= max(
+        get_tensor_model_parallel_world_size(), min_tokens
+    )
 
 
 def sp_all_gather(x: torch.Tensor) -> torch.Tensor:

vllm/models/muse_spark/nvidia/hybrid_moe.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Token-threshold TRT-LLM execution over Muse MegaMoE's shared NVFP4 weights."""
+
+import functools
+
+import torch
+
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    UnfinalizedMoEOutput,
+    convert_flashinfer_moe_output,
+)
+from vllm.models.muse_spark.nvidia.mega_moe import MuseSparkMegaMoEExperts
+from vllm.models.muse_spark.nvidia.ops.mnnvl_fusions import (
+    maybe_create_moe_finalize_all_reduce,
+)
+from vllm.models.muse_spark.nvidia.ops.prepare_trtllm_nvfp4 import (
+    norm_quant_trtllm_nvfp4,
+)
+
+
+@functools.cache
+def _get_decode_workspace(hidden_size: int, top_k: int, max_tokens: int):
+    workspace = maybe_create_moe_finalize_all_reduce(
+        hidden_size, top_k, max_tokens, enable_pdl=True
+    )
+    if workspace is None:
+        raise RuntimeError("Muse hybrid MoE requires the fused EP finalize workspace")
+    return workspace
+
+
+class MuseSparkHybridMoEExperts(MuseSparkMegaMoEExperts):
+    """TRT-LLM for small replicated batches; MegaMoE for sequence shards."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._decode_workspace = None
+        if self.is_nvfp4:
+            if get_tensor_model_parallel_world_size() not in (4, 8):
+                raise ValueError(
+                    "Muse hybrid MoE currently supports single-node TEP4/8"
+                )
+            self._decode_workspace = _get_decode_workspace(
+                self.hidden_size, self.top_k, self.trtllm_token_threshold
+            )
+            norm = self.post_expert_norm
+            self._decode_workspace.compile(
+                row_stride=self.hidden_size,
+                weight_dtype=norm.weight.dtype,
+                eps=norm.variance_epsilon,
+                gain_center=norm.fused_gain_center,
+                weight_before_norm=norm.applies_weight_before_norm,
+            )
+
+    def finalize_weights(self) -> None:
+        if self._transformed_l1 is not None:
+            return
+        super().finalize_weights()
+        if self.is_nvfp4:
+            # Preserve MegaMoE's per-expert scales. TRT follows its existing
+            # static-activation recipe, with one maximum per local expert bank.
+            a13 = self.w13_input_scale.max().float().reshape(1)
+            a2 = self.w2_input_scale.max().float().reshape(1)
+            self._input_global_scale = a13.reciprocal()
+            self._g1_gate = (a13 * self.w13_weight_scale_2[:, 0]).contiguous()
+            self._g1_up = (a13 * self.w13_weight_scale_2[:, 1] / a2).contiguous()
+            self._g2_scale = (a2 * self.w2_weight_scale_2).contiguous()
+
+    def use_trtllm(self, num_tokens: int, sp_active: bool) -> bool:
+        return (
+            self.is_nvfp4
+            and not sp_active
+            and 0 < num_tokens <= self.trtllm_token_threshold
+        )
+
+    def prepare_decode_inputs(self, pooled, norm, buffer):
+        num_tokens = pooled.shape[0]
+        return norm_quant_trtllm_nvfp4(
+            pooled,
+            num_tokens,
+            norm.weight,
+            norm.variance_epsilon,
+            norm.fused_gain_center,
+            self._input_global_scale,
+            out=(
+                buffer.x[:num_tokens].view(torch.uint8),
+                buffer.x_sf[:num_tokens].view(torch.uint8),
+            ),
+        )
+
+    def forward_trtllm(self, hidden_states: torch.Tensor, buffer) -> torch.Tensor:
+        import flashinfer
+        from flashinfer.comm.mnnvl_cutedsl_moe_finalize import (
+            moe_finalize_rmsnorm_allreduce,
+        )
+
+        num_tokens = hidden_states.shape[0]
+        raw = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(
+            topk_ids=(
+                buffer.topk_idx[:num_tokens].to(torch.int32),
+                buffer.topk_weights[:num_tokens],
+            ),
+            routing_bias=None,
+            hidden_states=buffer.x[:num_tokens].view(torch.uint8),
+            hidden_states_scale=buffer.x_sf[:num_tokens].view(torch.float8_e4m3fn),
+            gemm1_weights=self.w13_weight,
+            gemm1_weights_scale=self.w13_weight_scale,
+            gemm1_bias=None,
+            gemm1_alpha=None,
+            gemm1_beta=None,
+            gemm1_clamp_limit=None,
+            gemm2_weights=self.w2_weight,
+            gemm2_weights_scale=self.w2_weight_scale,
+            gemm2_bias=None,
+            output1_scale_scalar=self._g1_up,
+            output1_scale_gate_scalar=self._g1_gate,
+            output2_scale_scalar=self._g2_scale,
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            n_group=0,
+            topk_group=0,
+            intermediate_size=self.intermediate_size,
+            local_expert_offset=self.experts_start_idx,
+            local_num_experts=self.num_local_experts,
+            routed_scaling_factor=None,
+            routing_method_type=7,
+            do_finalize=False,
+            tune_max_num_tokens=self.trtllm_token_threshold,
+        )
+        result = convert_flashinfer_moe_output(
+            raw, do_finalize=False, num_tokens=num_tokens, top_k=self.top_k
+        )
+        assert isinstance(result, UnfinalizedMoEOutput)
+        norm = self.post_expert_norm
+        y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
+        moe_finalize_rmsnorm_allreduce(
+            result.gemm2_permuted,
+            result.expert_weights,
+            result.expanded_idx_to_permuted_idx,
+            norm.weight,
+            y,
+            self._decode_workspace,
+            eps=norm.variance_epsilon,
+            gain_center=norm.fused_gain_center,
+            weight_before_norm=norm.applies_weight_before_norm,
+        )
+        return y