Upload folder using huggingface_hub

pyproject.toml

@@ -0,0 +1,11 @@
+[project]
+name = "vllm-solon-moe"
+version = "1.0.0"
+description = "vLLM serving plugin for Solon-MoE (per-layer MoE over native Gemma4)"
+requires-python = ">=3.10"
+
+[project.entry-points."vllm.general_plugins"]
+solon_moe = "vllm_solon_moe:register"
+
+[tool.setuptools]
+py-modules = ["vllm_solon_moe"]

vllm_solon_moe.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""
+vllm_solon_moe.py -- Solon-MoE support for vLLM via a thin shim over the
+native Gemma4 implementation.
+
+Background: Solon-MoE checkpoints use the *native* transformers Gemma4 MoE
+block (packed gate_up_proj [E, 2*ffn, hidden], router.proj +
+per_expert_scale + scale, post_feedforward_layernorm_2 as the lambda scale),
+applied selectively to a subset of layers listed in ``config.moe_layers``.
+vLLM's gemma4.py already implements the full MoE path but only supports a
+*global* ``enable_moe_block`` flag. This shim makes the flag per-layer --
+the exact same trick as the 57-line HF-side shim.
+
+Usage (pick one):
+
+  A. vLLM plugin entry point (recommended for a package):
+       [project.entry-points."vllm.general_plugins"]
+       solon_moe = "vllm_solon_moe:register"
+
+  B. Ad-hoc, no packaging:
+       VLLM_PLUGINS="" python -c "
+       import vllm_solon_moe; vllm_solon_moe.register()
+       from vllm import LLM; ..."
+     or put this file on PYTHONPATH and add to VLLM_PLUGINS mechanism.
+
+Serving-time lambda override (no checkpoint copy needed):
+       SOLON_LAMBDA=0.10 vllm serve /path/to/solon-moe --trust-remote-code
+  The checkpoint stores the *training* lambda (0.15); inference is validated
+  at 0.10. If SOLON_LAMBDA is unset the checkpoint value is used as-is.
+
+ADAPT notes (vLLM version drift):
+  - Constructor signatures of Gemma4DecoderLayer / Gemma4ForCausalLM vary
+    across vLLM versions (vllm_config vs individual args). The shim passes
+    through *args/**kwargs untouched, so it survives most drifts; if the
+    decoder-layer class is named differently, fix DECODER_CLS_NAME below.
+  - Verify once that vLLM's RMSNorm here is plain ``x*w`` (matching
+    transformers Gemma4), not the legacy Gemma ``x*(1+w)``:
+      grep -n "from.*layernorm import" .../models/gemma4.py
+    If it imports GemmaRMSNorm for these norms, lambda semantics break --
+    that must be patched before serving.
+"""
+import os
+import re
+
+from vllm.model_executor.models import gemma4 as _g4
+
+DECODER_CLS_NAME = "Gemma4DecoderLayer"   # ADAPT if named differently
+_LAYER_RE = re.compile(r"layers\.(\d+)")
+
+
+def _make_selective_layer(base_cls):
+    class SolonMoEDecoderLayer(base_cls):
+        def __init__(self, *args, **kwargs):
+            # Locate config and prefix among the args regardless of signature.
+            config = kwargs.get("config")
+            prefix = kwargs.get("prefix", "")
+            if config is None:
+                for a in args:
+                    if hasattr(a, "num_hidden_layers"):
+                        config = a
+                        break
+            if not prefix:
+                for a in args:
+                    if isinstance(a, str) and "layers." in a:
+                        prefix = a
+                        break
+            m = _LAYER_RE.search(prefix or "")
+            layer_idx = int(m.group(1)) if m else -1
+
+            moe_layers = getattr(config, "moe_layers", None)
+            if moe_layers is None:
+                want = bool(getattr(config, "enable_moe_block", False))
+            else:
+                want = layer_idx in set(moe_layers)
+
+            saved = getattr(config, "enable_moe_block", False)
+            saved2 = getattr(config, "use_second_mlp_block", False)
+            config.enable_moe_block = bool(want)
+            if hasattr(config, "use_second_mlp_block"):
+                config.use_second_mlp_block = False   # single source of truth
+            try:
+                super().__init__(*args, **kwargs)
+            finally:
+                config.enable_moe_block = saved
+                if hasattr(config, "use_second_mlp_block"):
+                    config.use_second_mlp_block = saved2
+
+    return SolonMoEDecoderLayer
+
+
+class SolonMoEForCausalLM(_g4.Gemma4ForCausalLM):
+    """Gemma4ForCausalLM with per-layer MoE selection + lambda override."""
+
+    def __init__(self, *, vllm_config, prefix: str = ""):
+        # NOTE: signature must be explicit new-style (*, vllm_config, prefix) —
+        # vLLM introspects it to choose the construction path. A generic
+        # (*args, **kwargs) passthrough makes vLLM fall back to old-style
+        # argument guessing, which then omits vllm_config. (Observed 2026-08.)
+        base = getattr(_g4, DECODER_CLS_NAME)
+        patched = _make_selective_layer(base)
+        setattr(_g4, DECODER_CLS_NAME, patched)
+        try:
+            super().__init__(vllm_config=vllm_config, prefix=prefix)
+        finally:
+            setattr(_g4, DECODER_CLS_NAME, base)
+
+    def load_weights(self, *args, **kwargs):
+        out = super().load_weights(*args, **kwargs)
+        self._apply_lambda_override()
+        return out
+
+    def _apply_lambda_override(self):
+        lam = os.environ.get("SOLON_LAMBDA")
+        if lam is None:
+            return
+        lam = float(lam)
+        n = 0
+        for name, module in self.named_modules():
+            if name.endswith("post_feedforward_layernorm_2") and hasattr(module, "weight"):
+                module.weight.data.fill_(lam)
+                n += 1
+        print(f"[solon_moe] SOLON_LAMBDA={lam} applied to {n} layer(s)"
+              + ("" if n else "  -- WARNING: no lambda norms found (naming drift?)"))
+
+
+def register():
+    from vllm import ModelRegistry
+    if "SolonMoEForCausalLM" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model(
+            "SolonMoEForCausalLM",
+            "vllm_solon_moe:SolonMoEForCausalLM",
+        )