fix(moe): preserve upstream quantized methods

Author: CherryLemon

tests/unit_tests/ops/test_fused_moe_layer.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2026 BAAI. All rights reserved.
+
+from unittest.mock import MagicMock, patch
+
+from vllm_fl.ops.fused_moe import layer
+
+
+def _runner_with_quant_method(quant_method):
+    runner = MagicMock()
+    runner._quant_method = quant_method
+    runner.moe_config = MagicMock()
+    return runner
+
+
+def test_fused_moe_fl_replaces_unquantized_method():
+    quant_method = MagicMock(spec=layer.UnquantizedFusedMoEMethod)
+    runner = _runner_with_quant_method(quant_method)
+    replacement = MagicMock()
+
+    with (
+        patch.object(layer, "_OrigFusedMoE", return_value=runner),
+        patch.object(
+            layer,
+            "UnquantizedFusedMoEMethodFL",
+            return_value=replacement,
+        ) as replacement_cls,
+        patch.object(layer, "replace_router_with_fl") as replace_router,
+    ):
+        result = layer.FusedMoEFL(test_arg=True)
+
+    assert result is runner
+    replacement_cls.assert_called_once_with(runner.moe_config)
+    runner._replace_quant_method.assert_called_once_with(replacement)
+    replace_router.assert_called_once_with()
+
+
+def test_fused_moe_fl_preserves_quantized_method():
+    quant_method = object()
+    runner = _runner_with_quant_method(quant_method)
+
+    with (
+        patch.object(layer, "_OrigFusedMoE", return_value=runner),
+        patch.object(layer, "UnquantizedFusedMoEMethodFL") as replacement_cls,
+        patch.object(layer, "replace_router_with_fl") as replace_router,
+        patch.object(layer.logger, "info_once") as info_once,
+    ):
+        result = layer.FusedMoEFL()
+
+    assert result is runner
+    replacement_cls.assert_not_called()
+    runner._replace_quant_method.assert_not_called()
+    replace_router.assert_called_once_with()
+    info_once.assert_called_once_with(
+        "Preserving upstream quantized MoE method %s in FusedMoEFL.",
+        "object",
+    )

vllm_fl/ops/fused_moe/layer.py

@@ -12,11 +12,15 @@
 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
     UnquantizedFusedMoEMethod,
 )
+from vllm.logger import init_logger
 
 from vllm_fl.ops.fused_moe.router import replace_router_with_fl
 from .fused_moe_utils import select_unquantized_moe_backend_oot
 
 
+logger = init_logger(__name__)
+
+
 class UnquantizedFusedMoEMethodFL(UnquantizedFusedMoEMethod):
     """OOT replacement for UnquantizedFusedMoEMethod that routes computation
     through flaggems operators."""
@@ -54,9 +58,17 @@ def FusedMoEFL(*args, **kwargs) -> MoERunner:
     #    _fused_moe_pkg.FusedMoE with FusedMoEFL.
     runner: MoERunner = _OrigFusedMoE(*args, **kwargs)
 
-    # 2. Replace quant_method with FL version.
-    fl_quant_method = UnquantizedFusedMoEMethodFL(runner.moe_config)
-    runner._replace_quant_method(fl_quant_method)
+    # 2. Replace only an upstream unquantized method with the FL version.
+    # Quantized methods own their weight/activation scaling metadata and must
+    # remain attached to the runner.
+    if isinstance(runner._quant_method, UnquantizedFusedMoEMethod):
+        fl_quant_method = UnquantizedFusedMoEMethodFL(runner.moe_config)
+        runner._replace_quant_method(fl_quant_method)
+    else:
+        logger.info_once(
+            "Preserving upstream quantized MoE method %s in FusedMoEFL.",
+            type(runner._quant_method).__name__,
+        )
 
     # 3. Replace router _compute_routing with FL version via monkey-patch.
     #    replace_router_with_fl() patches the class method so the router