refactor(musa): remove direct flaggems calls from vendor backend

vllm_fl/dispatch/backends/vendor/musa/musa.py

@@ -142,40 +142,3 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) ->
 
         return AttentionBackendEnum.TRITON_ATTN.get_path()
 
-    def topk_softmax(
-        self,
-        topk_weights: torch.Tensor,
-        topk_indices: torch.Tensor,
-        token_expert_indices: torch.Tensor,
-        gating_output: torch.Tensor,
-        renormalize: bool = False,
-        e_score_correction_bias: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """MoE top-k softmax via FlagGems (no _moe_C CUDA extension needed).
-
-        ``torch.ops._moe_C.topk_softmax`` is a CUDA-only custom op that is
-        not available on MUSA.  FlagGems ships a portable triton implementation
-        that runs on MUSA devices.
-
-        Args:
-            topk_weights: Output tensor for top-k weights (modified in-place).
-            topk_indices: Output tensor for top-k expert indices (modified in-place).
-            token_expert_indices: Token-to-expert index mapping.
-            gating_output: Raw logits from the router gate.
-            renormalize: Whether to renormalize top-k weights to sum to 1.
-            e_score_correction_bias: Optional per-expert bias for score correction.
-
-        Returns:
-            Tuple of (topk_weights, topk_indices).
-        """
-        from vllm_fl.dispatch.backends.flaggems.impl.fused_moe import (
-            topk_softmax_flaggems,
-        )
-
-        return topk_softmax_flaggems(
-            topk_weights,
-            topk_indices,
-            token_expert_indices,
-            gating_output,
-            renormalize,
-        )

vllm_fl/dispatch/backends/vendor/musa/patch.py

@@ -158,9 +158,12 @@ def patch_cuda_stream_for_musa():
             return
 
         # aux_stream helper used by MoE shared-expert overlap.
-        # MUSA does not support torch.cuda.Stream() — return None so that
-        # SharedExperts falls back to the synchronous (non-overlapped) path,
-        # which is guarded by ``if self._stream is not None``.
+        # torch_musa currently does not support parallel CUDA streams for
+        # compute overlap (torch.cuda.Stream() raises or returns a non-functional
+        # stream on MUSA devices).  Return None so that SharedExperts falls back
+        # to the synchronous (non-overlapped) path, which is guarded by
+        # ``if self._stream is not None``.
+        # TODO: Enable parallel stream overlap once torch_musa supports it.
         try:
             import vllm.utils.torch_utils as _tu
 
@@ -300,7 +303,7 @@ def _make_launcher_musa(self):
 
 
 def patch_moe_topk_softmax_for_musa():
-    """Patch MoE top-k softmax for MUSA via FlagGems.
+    """Patch MoE top-k softmax for MUSA via the dispatch system.
 
     ``torch.ops._moe_C.topk_softmax`` is a CUDA-only extension not available
     on MUSA.  vllm 0.24.0 calls it through two entry points:
@@ -317,16 +320,18 @@ def patch_moe_topk_softmax_for_musa():
     this point since they only depend on torch and vllm internals that are
     already loaded by the time apply_musa_patches() is invoked in the worker.
 
+    The actual implementation is resolved through the dispatch system's
+    op_backends configuration (musa.yaml: topk_softmax -> [flagos, reference]),
+    rather than directly importing from the flaggems backend.
+
     TODO: remove once MUSA ships a compiled _moe_C extension.
     """
     try:
-        from vllm_fl.dispatch.backends.flaggems.impl.fused_moe import (
-            topk_softmax_flaggems,
-        )
+        from vllm_fl.dispatch import call_op
     except Exception as exc:
         logger.warning(
             "patch_moe_topk_softmax_for_musa: cannot import "
-            "topk_softmax_flaggems — MoE models will fail on MUSA: %s", exc)
+            "dispatch call_op — MoE models will fail on MUSA: %s", exc)
         return
 
     def _topk_softmax_musa(
@@ -337,10 +342,13 @@ def _topk_softmax_musa(
         renormalize=False,
         e_score_correction_bias=None,
     ):
-        # topk_softmax_flaggems modifies topk_weights/topk_ids in-place.
+        # Delegate to the dispatch system which resolves the best available
+        # implementation based on musa.yaml op_backends configuration.
+        # The dispatch call modifies topk_weights/topk_ids in-place.
         # Must return (topk_weights, topk_ids) to match vllm_topk_softmax
         # signature which callers unpack as: topk_weights, topk_ids = topk_func(...)
-        topk_softmax_flaggems(
+        call_op(
+            "topk_softmax",
             topk_weights,
             topk_ids,
             token_expert_indices,
@@ -368,7 +376,7 @@ def _topk_softmax_musa(
 
     if patched:
         logger.info(
-            "Patched vllm_topk_softmax for MUSA (FlagGems) in: %s",
+            "Patched vllm_topk_softmax for MUSA (via dispatch) in: %s",
             ", ".join(patched))
 
 

vllm_fl/dispatch/backends/vendor/musa/register_ops.py

@@ -76,15 +76,6 @@ def register_builtins(registry) -> None:
             vendor="musa",
             priority=BackendPriority.VENDOR,
         ),
-        # MoE top-k softmax (no _moe_C CUDA extension on MUSA, use FlagGems)
-        OpImpl(
-            op_name="topk_softmax",
-            impl_id="vendor.musa",
-            kind=BackendImplKind.VENDOR,
-            fn=_bind_is_available(backend.topk_softmax, is_avail),
-            vendor="musa",
-            priority=BackendPriority.VENDOR,
-        ),
     ]
 
     registry.register_many(impls)

vllm_fl/dispatch/config/musa.yaml

@@ -45,6 +45,11 @@ op_backends:
     - flagos
     - vendor:musa
     - reference
+  # MoE top-k softmax: MUSA has no _moe_C CUDA extension,
+  # use FlagGems triton implementation with reference as fallback
+  topk_softmax:
+    - flagos
+    - reference
 
 # FlagOS operator blacklist
 # Musa is CUDA-compatible, so most FlagGems ops should work