fix(kunlunxin): fix 35B-A3B startup and restore NPU code

- fused_moe_utils.py: remove broken is_cuda() guard in apply() so XPU
  takes the fast path (fused_experts_impl) with proper return statement
- worker.py: restore NPU compatibility code (empty_cache patch,
  profile_run skip, _dummy_run skip) removed during PR268 merge

Co-Authored-By: Claude <noreply@anthropic.com>

Author: luoyc123

vllm_fl/ops/fused_moe/fused_moe_utils.py

@@ -12,44 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+from enum import Enum
 from typing import Any
 
 import torch
 
 import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config.kernel import MoEBackend
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
 )
-from vllm.model_executor.layers.fused_moe.fused_moe import (
-    TritonExperts,
-    try_get_optimal_moe_config,
-)
-from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
-    UnquantizedMoeBackend,
-    backend_to_kernel_cls,
-    map_unquantized_backend,
-)
-from vllm.model_executor.layers.fused_moe.utils import (
-    _resize_cache,
-    moe_kernel_quantize_input,
-)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+from vllm.model_executor.layers.fused_moe.oracle.unquantized import UnquantizedMoeBackend, map_unquantized_backend, backend_to_kernel_cls
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts, try_get_optimal_moe_config
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache, moe_kernel_quantize_input
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     FlashinferMoeBackend,
     get_flashinfer_moe_backend,
 )
-from vllm.platforms import current_platform
-from vllm.triton_utils import tl
-
+from vllm.triton_utils import tl, triton
 from vllm_fl.dispatch import CachedOp
 from vllm_fl.ops.fused_moe.activation import apply_moe_activation
 from vllm_fl.utils import use_flaggems
 # Kunlunxin: always use TritonExpertsFL regardless of flaggems setting
 # (avoids _moe_C arg mismatch when flaggems is off)
 from vllm_fl.dispatch.config.utils import get_platform_name
+
 _moe_align_block_size = CachedOp("moe_align_block_size")
 _invoke_fused_moe_triton_kernel = CachedOp("invoke_fused_moe_triton_kernel")
 _moe_sum = CachedOp("moe_sum")
@@ -102,10 +95,8 @@ def _move_to_back(
         ]
     return _AVAILABLE_BACKENDS
 
-
 ## Adopt from select_unquantized_moe_backend
-def select_unquantized_moe_backend_oot(
-    moe_config: FusedMoEConfig,
+def select_unquantized_moe_backend_oot(moe_config: FusedMoEConfig,
 ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
     """
     Select the primary Unquantized MoE backend.
@@ -118,16 +109,14 @@ def select_unquantized_moe_backend_oot(
 
     if current_platform.is_tpu():
         return UnquantizedMoeBackend.TPU, None
-
-    if current_platform.is_out_of_tree() and use_flaggems():
+    if current_platform.is_out_of_tree() and use_flaggems(): 
         return UnquantizedMoeBackend.TRITON, TritonExpertsFL
     elif current_platform.is_out_of_tree():
         # Kunlunxin uses TritonExpertsFL unconditionally because its
         # fused_experts_impl is patched to xtorch_ops, not flaggems.
         if get_platform_name() == "kunlunxin":
             return UnquantizedMoeBackend.TRITON, TritonExpertsFL
         return UnquantizedMoeBackend.TRITON, TritonExperts
-
     if moe_config.is_lora_enabled:
         return UnquantizedMoeBackend.TRITON, backend_to_kernel_cls(
             UnquantizedMoeBackend.TRITON
@@ -257,7 +246,6 @@ def _return_or_raise(
         "No Unquantized MoE backend supports the deployment configuration."
     )
 
-
 def _prepare_expert_assignment(
     topk_ids: torch.Tensor,
     config: dict[str, Any],
@@ -306,7 +294,6 @@ def _prepare_expert_assignment(
         ignore_invalid_experts=ignore_invalid_experts,
     )
 
-
 class TritonExpertsFL(TritonExperts):
     def apply(
         self,
@@ -326,53 +313,11 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
-        # Dynamic W8A8 is handled by TritonW8A8Experts so vLLM owns
-        # activation quantization. Do not allow it to fall back into the
-        # FlagGems contract, which expects floating-point input here.
-        if self.quant_config.use_int8_w8a8:
-            raise RuntimeError(
-                "W8A8 MoE must use TritonW8A8Experts, not TritonExpertsFL"
-            )
-
-        # Fast path (no LoRA, NVIDIA only): let FlagGems own both expert GEMMs
-        # for unquantized and W8A16 inputs.
-        if self._lora_context is None and current_platform.is_cuda():
-            import flag_gems
-
-            output.copy_(
-                flag_gems.fused_experts_impl(
-                    hidden_states,
-                    w1,
-                    w2,
-                    topk_weights,
-                    topk_ids,
-                    inplace=False,
-                    activation=activation.value,
-                    apply_router_weight_on_input=apply_router_weight_on_input,
-                    use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
-                    use_int8_w8a8=False,
-                    use_int8_w8a16=self.quant_config.use_int8_w8a16,
-                    use_int4_w4a16=self.quant_config.use_int4_w4a16,
-                    per_channel_quant=self.per_act_token_quant,
-                    global_num_experts=global_num_experts,
-                    expert_map=expert_map,
-                    w1_scale=self.w1_scale,
-                    w2_scale=self.w2_scale,
-                    a1_scale=a1q_scale,
-                    a2_scale=a2_scale,
-                    block_shape=self.block_shape,
-                    w1_bias=self.w1_bias,
-                    w2_bias=self.w2_bias,
-                )
-            )
         # Fast path (no LoRA): single fused call via patched fused_experts_impl.
         if self._lora_context is None:
             from vllm_fl.ops.fused_moe.fused_moe import fused_experts_impl as _fused_experts_impl
-        # Fast path (no LoRA, NVIDIA only): single fused FlagGems call.
-        if self._lora_context is None and current_platform.is_cuda():
-            import flag_gems
 
-            output.copy_(flag_gems.fused_experts_impl(
+            output.copy_(_fused_experts_impl(
                 hidden_states,
                 w1,
                 w2,

vllm_fl/worker/worker.py

@@ -259,35 +259,13 @@ def __init__(
                     path=fl_envs.FLAGGEMS_ENABLE_OPLIST_PATH,
                 )
             elif blacklist:
-                # On XPU (Kunlunxin), enable() / config_filter() triggers a Triton
-                # JIT _scatter_jit_function crash (err_code -714) during profile_run.
-                # Work around by converting the blacklist to a whitelist complement
-                # and calling only_enable() instead, which uses a different internal
-                # registration path that works correctly on this platform.
-                is_xpu = current_platform.vendor_name == "kunlunxin"
-                if is_xpu:
-                    all_op_names = [
-                        item[0] for item in flag_gems._FULL_CONFIG if len(item) >= 2 
-                    ]
-                    complement = [n for n in all_op_names if n not in blacklist]
-                    logger.info(
-                        f"[FlagGems] Blacklist mode (XPU): excluding {blacklist}, "
-                        f"enabling {len(complement)}/{len(all_op_names)} ATen ops"
-                    )
-                    flag_gems.only_enable(
-                        include=complement,
-                        record=True,
-                        once=True,
-                        path=fl_envs.FLAGGEMS_ENABLE_OPLIST_PATH,
-                    )
-                else:
-                    logger.info(f"[FlagGems] Disable the following ops: {blacklist}")
-                    flag_gems.enable(
-                        unused=blacklist,
-                        record=True,
-                        once=True,
-                        path=fl_envs.FLAGGEMS_ENABLE_OPLIST_PATH,
-                    )
+                logger.info(f"[FlagGems] Disable the following ops: {blacklist}")
+                flag_gems.enable(
+                    unused=blacklist,
+                    record=should_record,
+                    once=True,
+                    path=fl_envs.FLAGGEMS_ENABLE_OPLIST_PATH,
+                )
             else:
                 logger.info("[FlagGems] Enable all ops")
                 flag_gems.enable(