From 087e15f473bab530d39b7bf71b0284bfb5f0af58 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Sun, 23 Aug 2026 07:30:52 +0000
Subject: [PATCH] Fix DeepSeek MoE routed experts losing linear_class
 quantization

_DeepSeekMoEFFN.__init__ threaded linear_class into _SharedExpertMLP
but not into the MoELayer that builds the routed-expert dense-loop
fallback. A quantized config therefore quantized attention, dense FFN,
and the shared expert, but silently left every routed MoE expert's
gate/up/down projections as plain float MatMul -- losing quantization
and breaking the fuse_dense_moe_to_qmoe post-hoc rewrite, which only
pattern-matches a quantized MatMulNBits dense-fallback shape.

Fix: pass linear_class through to MoELayer's routed-expert
construction (one line). DeepSeek-V4's MoE FFN is unaffected -- it
recomputes its own quantized class per-expert independently of
MoELayer's linear_class and was already correct.

Added test_deepseek_moe_ffn_linear_class_reaches_routed_experts,
which isolates and asserts on module.moe.experts vs
module.shared_experts separately (not just an aggregate creation
count) so the regression is caught precisely. Verified: fails on
pre-fix code, passes post-fix. Full suite (3249 passed, 9 skipped)
and ruff check/format clean.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/models/deepseek.py      |  9 +++++-
 src/mobius/models/deepseek_test.py | 51 ++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py
index 329810eb..430f6aef 100644
--- a/src/mobius/models/deepseek.py
+++ b/src/mobius/models/deepseek.py
@@ -315,7 +315,14 @@ def __init__(
         super().__init__()
         assert config.num_local_experts is not None
         assert config.moe_intermediate_size is not None
-        self.moe = MoELayer(config, gate=gate)
+        # ``linear_class`` must reach the routed-expert dense-loop fallback
+        # too (not just the shared expert below): otherwise a quantized
+        # config quantizes every other linear in the model (attention,
+        # dense FFN, shared expert) but silently leaves the routed MoE
+        # experts as plain float `MatMul`, which both loses quantization
+        # and breaks the `fuse_dense_moe_to_qmoe` post-hoc rewrite (it
+        # only matches a quantized `MatMulNBits` dense-fallback pattern).
+        self.moe = MoELayer(config, gate=gate, linear_class=linear_class)
         # Shared expert uses moe_intermediate_size * n_shared_experts
         n_shared = config.n_shared_experts or 1
         shared_intermediate = config.moe_intermediate_size * n_shared
diff --git a/src/mobius/models/deepseek_test.py b/src/mobius/models/deepseek_test.py
index 8cedd3e7..6bcfa131 100644
--- a/src/mobius/models/deepseek_test.py
+++ b/src/mobius/models/deepseek_test.py
@@ -97,6 +97,57 @@ def test_deepseek_dense_moe_fallback_matches_numpy_reference():
     assert all(node.domain != "com.microsoft" for node in graph)
 
 
+def test_deepseek_moe_ffn_linear_class_reaches_routed_experts():
+    """``linear_class`` must reach the routed dense-loop experts, not just the shared expert.
+
+    Regression test for a bug where ``_DeepSeekMoEFFN`` constructed its
+    ``MoELayer`` without forwarding ``linear_class``, so a quantized config
+    quantized attention/dense-FFN/shared-expert linears but silently left the
+    routed MoE experts as plain float ``MatMul`` -- losing quantization and
+    breaking the ``fuse_dense_moe_to_qmoe`` post-hoc rewrite, which only
+    matches a quantized ``MatMulNBits`` dense-fallback pattern.
+    """
+    from mobius.components._common import Linear
+
+    created: list[Linear] = []
+
+    class TrackingLinear(Linear):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            created.append(self)
+
+    config = make_config(
+        hidden_size=4,
+        intermediate_size=8,
+        moe_intermediate_size=3,
+        num_local_experts=4,
+        num_experts_per_tok=2,
+        n_group=2,
+        topk_group=1,
+        n_shared_experts=1,
+        scoring_func="sigmoid",
+        topk_method="noaux_tc",
+    )
+    module = _DeepSeekMoEFFN(config, DeepSeekMoEGate(config), linear_class=TrackingLinear)
+
+    # The dense loop-over-experts fallback must have been built with
+    # TrackingLinear for every routed expert's projections. Check
+    # `module.moe.experts` directly (not the global `created` list) so this
+    # assertion can't pass merely because the shared expert below picked up
+    # the class -- that path already worked before this fix.
+    assert module.moe.experts is not None
+    routed_linears = [m for m in module.moe.experts.modules() if isinstance(m, Linear)]
+    assert len(routed_linears) > 0
+    assert all(isinstance(m, TrackingLinear) for m in routed_linears)
+    # The shared expert must still be quantized too (it already worked
+    # before this fix; guard against a future regression there as well).
+    shared_linears = [m for m in module.shared_experts.modules() if isinstance(m, Linear)]
+    assert len(shared_linears) > 0
+    assert all(isinstance(m, TrackingLinear) for m in shared_linears)
+    # Sanity: the tracking list saw both groups (routed + shared).
+    assert len(created) == len(routed_linears) + len(shared_linears)
+
+
 def test_noaux_tc_supports_single_expert_groups():
     config = make_config(
         hidden_size=4,