feat(models): add Qwen3-Next-80B-A3B support

benchmarks/bench_decode_moe.py

@@ -115,6 +115,11 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         help="seconds to wait for the spawned server to become ready",
     )
     p.add_argument("--json", dest="json_out", default=None, help="append the result rows here")
+    p.add_argument(
+        "--extra-args",
+        default="",
+        help="extra flags passed verbatim to the server (e.g. '--kv-reserve-tokens 2048 --max-seq-len-override 4096')",
+    )
     return p.parse_args(argv)
 
 
@@ -193,6 +198,9 @@ def serve_cmd(args: argparse.Namespace, backend: str, port: int) -> list[str]:
         cmd += ["--moe-cache-rate", str(args.cache_rate)]
     else:
         cmd.append("--moe-cache-auto")
+    if getattr(args, "extra_args", ""):
+        import shlex as _shlex
+        cmd += _shlex.split(args.extra_args)
     return cmd
 
 

docs/models.md

@@ -12,6 +12,7 @@ for them; other checkpoints of the same architectures work too.
 | Qwen3.6 / Qwen3.5 MoE | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-35B-A3B-FP8)), [nvidia/Qwen3.6-35B-A3B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-35B-A3B-NVFP4), [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8)) |
 | Qwen3.6 dense | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-27B-FP8)), [nvidia/Qwen3.6-27B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-27B-NVFP4) |
 | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) |
+| Qwen3-Next | [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct) ([-FP8](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct-FP8)), [nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4](https://huggingface.co/nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4) |
 | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) |
 | Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. |
 | MiniMax-M2.5 | [nvidia/MiniMax-M2.5-NVFP4](https://huggingface.co/nvidia/MiniMax-M2.5-NVFP4) |

python/freetoken/engine/config.py

@@ -89,7 +89,7 @@ def hf_config(self):
     def model_config(self) -> ModelConfig:
         spec = get_model_spec(self.hf_config.architectures[0])
         parse_config = _load_attr(spec.module, spec.parse_config)
-        return parse_config(self.hf_config)
+        return parse_config(self.hf_config, self.model_path)
 
     @property
     def max_seq_len(self) -> int:

python/freetoken/layers/moe.py

@@ -38,6 +38,7 @@ def __init__(
         apply_router_weight_on_input: bool = False,
         allocate_experts: bool = True,
         weight_format: str = "bf16",
+        scoring: str = "softmax",
     ):
         super().__init__()
 
@@ -53,6 +54,8 @@ def __init__(
         self.activation = activation
         self.apply_router_weight_on_input = apply_router_weight_on_input
         self.weight_format = weight_format
+        # Router scoring function: "softmax" (Qwen3.5/3.6) or "sigmoid" (Qwen3-Next).
+        self.scoring = scoring
         intermediate_size_per_partition = div_even(intermediate_size, tp_size)
         if allocate_experts:
             self._alloc_resident_experts(intermediate_size_per_partition)
@@ -176,6 +179,7 @@ def forward(
                 gating_output=router_logits,
                 topk=self.top_k,
                 renormalize=self.renormalize,
+                scoring=self.scoring,
             )
             return self._maybe_all_reduce(
                 self._resident_gemm(hidden_states, topk_weights, topk_ids)
@@ -190,6 +194,7 @@ def forward(
             renormalize=self.renormalize,
             activation=self.activation,
             apply_router_weight_on_input=self.apply_router_weight_on_input,
+            **({"scoring": self.scoring} if self.scoring != "softmax" else {}),
         )
         return self._maybe_all_reduce(final_hidden_states)
 
@@ -261,6 +266,7 @@ def decode_forward(
             gating_output=router_logits,
             topk=self.top_k,
             renormalize=self.renormalize,
+            scoring=self.scoring,
         )
         return self._decode_routed(hidden_states, topk_weights, topk_ids)
 
@@ -274,6 +280,7 @@ def prefill_forward(
             gating_output=router_logits,
             topk=self.top_k,
             renormalize=self.renormalize,
+            scoring=self.scoring,
         )
         return self._prefill_routed(hidden_states, topk_weights, topk_ids)
 
@@ -631,6 +638,9 @@ def make_moe_layer(
     else:
         kwargs["weight_format"] = weight_format
     layer = layer_cls(**kwargs)
+    # Router scoring (Qwen3-Next uses sigmoid): set post-construction so every layer
+    # class (resident + offload) takes it without signature changes.
+    layer.scoring = getattr(config, "moe_scoring_func", "softmax")
     for name, value in (extra_attrs or {}).items():
         setattr(layer, name, value)
     return layer

python/freetoken/models/config.py

@@ -158,6 +158,9 @@ class LinearGatedDeltaGroupConfig(BaseAttentionGroupConfig):
     value_head_dim: int
     conv_kernel_dim: int
     output_gate: bool
+    # Checkpoint ships the GDN input projection pre-fused as bf16 ``in_proj_qkvz`` +
+    # ``in_proj_ba`` (Qwen3-Next modelopt NVFP4) instead of the four unfused parts.
+    in_proj_split: bool = False
 
 
 @dataclass(frozen=True)
@@ -215,6 +218,9 @@ class ModelConfig:
     model_type: str
     architectures: list[str]
     moe_backend: str = "fused"
+    # Router scoring function: "softmax" (default, Qwen3.5/3.6) or "sigmoid"
+    # (Qwen3-Next ``scoring_func``). Read by make_moe_layer -> fused_topk.
+    moe_scoring_func: str = "softmax"
     # ----- optional, model-specific extensions (default keeps other models intact) -----
     moe_enabled: bool = False
     # Weight quantization of the MoE experts only. "none" keeps the default BF16

python/freetoken/models/deepseek_v4/config.py

@@ -20,7 +20,7 @@
 from .args import load_args
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     model_path = getattr(hf_config, "_name_or_path", None) or getattr(
         hf_config, "name_or_path", None
     )

python/freetoken/models/gemma4/config.py

@@ -97,7 +97,7 @@ def _attn_geometry(cfg: Any, layer_type: str, *, is_full: bool) -> tuple[int, in
     )
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     cfg, top_architectures, top_cfg = _text_config(hf_config)
     rope_params = cfg.rope_parameters
     swa_type, full_type = "sliding_attention", "full_attention"

python/freetoken/models/gemma4/gguf.py

@@ -48,7 +48,7 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int:
     return full_head_dim // 4
 
 
-def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig:
+def parse_gguf_config(shim: "GgufConfigShim", model_path: str | None = None) -> ModelConfig:
     m = shim.metadata
 
     def g(key: str):

python/freetoken/models/glm4_moe/config.py

@@ -16,7 +16,7 @@ def _rope_params(hf_config: Any) -> dict:
     return params
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     """Parse a HuggingFace ``Glm4MoeConfig`` (GLM-4.5/4.6/4.7) into FreeToken's
     :class:`ModelConfig`.
 

python/freetoken/models/glm_moe_dsa/config.py

@@ -50,7 +50,7 @@ def _dsa_on(args, num_layers: int) -> bool:
     )
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     args = load_args(hf_config)
     # Latent-KV MLA: the paged pool stores a single "head" = ckv (kv_lora_rank) | kpe
     # (qk_rope_head_dim); the model absorbs kv_b into Q/O and attends with flashinfer MLA.

python/freetoken/models/gpt_oss/config.py

@@ -38,7 +38,7 @@ def _rope_config(hf_config: Any, head_dim: int) -> RotaryConfig:
     )
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     head_dim = getattr(hf_config, "head_dim", None) or (
         hf_config.hidden_size // hf_config.num_attention_heads
     )

python/freetoken/models/llama/config.py

@@ -5,7 +5,7 @@
 from freetoken.models.config import ModelConfig, RotaryConfig
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     num_kv_heads = getattr(hf_config, "num_key_value_heads", hf_config.num_attention_heads)
     head_dim = (
         getattr(hf_config, "head_dim", None)

python/freetoken/models/minimax_m2/config.py

@@ -18,7 +18,7 @@ def _rope_params(hf_config: Any) -> dict:
     return params
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     head_dim = (
         getattr(hf_config, "head_dim", None)
         or hf_config.hidden_size // hf_config.num_attention_heads

python/freetoken/models/minimax_m3/config.py

@@ -58,7 +58,7 @@ def _log_mode_once(key: str, message: str) -> None:
     init_logger(__name__).info(message)
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     text = _text_config(hf_config)
 
     num_layers = text.num_hidden_layers

python/freetoken/models/mistral/config.py

@@ -5,7 +5,7 @@
 from freetoken.models.config import ModelConfig, RotaryConfig
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     num_kv_heads = getattr(hf_config, "num_key_value_heads", hf_config.num_attention_heads)
     head_dim = (
         getattr(hf_config, "head_dim", None)

python/freetoken/models/muse_glimmer/config.py

@@ -35,7 +35,7 @@ def _group_rope_theta(text: Any, layer_ids: tuple[int, ...], default: float) ->
     return thetas.pop()
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     text = _text_config(hf_config)
 
     head_dim = getattr(text, "head_dim", None) or text.hidden_size // text.num_attention_heads

python/freetoken/models/qwen2/config.py

@@ -5,7 +5,7 @@
 from freetoken.models.config import ModelConfig, RotaryConfig
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     num_kv_heads = getattr(hf_config, "num_key_value_heads", hf_config.num_attention_heads)
     head_dim = (
         getattr(hf_config, "head_dim", None)

python/freetoken/models/qwen3/config.py

@@ -5,7 +5,7 @@
 from freetoken.models.config import ModelConfig, RotaryConfig
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     num_kv_heads = getattr(hf_config, "num_key_value_heads", hf_config.num_attention_heads)
     head_dim = (
         getattr(hf_config, "head_dim", None)

python/freetoken/models/qwen3_5_moe/config.py

@@ -9,6 +9,7 @@
     RotaryConfig,
     detect_compressed_tensors_nvfp4,
 )
+from freetoken.models.loader import iter_weight_files
 
 
 def _quant_accessor(hf_config: Any):
@@ -20,6 +21,38 @@ def _quant_accessor(hf_config: Any):
     return quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d))
 
 
+def _is_ct_storage(model_path: str | None) -> bool:
+    """Whether the checkpoint stores compressed-tensors packed FP4 (``weight_packed`` +
+    ``weight_global_scale``). Modelopt-quantized NVFP4 checkpoints (nvidia Qwen3-Next-*-NVFP4)
+    advertise a compressed-tensors style config (``config_groups``/``ignore``) but store
+    modelopt two-level scales (``weight_scale``/``weight_scale_2``) instead -- probe the
+    first shard. Defaults to True (legacy ct path) when undeterminable."""
+    if model_path is None:
+        return True
+    import safetensors
+    for file in iter_weight_files(model_path):
+        with safetensors.safe_open(file, framework="pt", device="cpu") as f:
+            return any(k.endswith((".weight_packed", ".weight_global_scale")) for k in f.keys())
+    return True
+
+
+def _gdn_split_layout(model_path: str | None) -> bool:
+    """Whether the checkpoint ships the GDN input projection pre-fused as bf16
+    ``in_proj_qkvz`` + ``in_proj_ba`` (Qwen3-Next modelopt NVFP4 layout) instead of the
+    four unfused ``in_proj_{qkv,z,b,a}`` parts (Qwen3.5 ct-NVFP4 / bf16 layout)."""
+    if model_path is None:
+        return False
+    import safetensors
+    for file in iter_weight_files(model_path):
+        with safetensors.safe_open(file, framework="pt", device="cpu") as f:
+            keys = list(f.keys())
+            if any(k.endswith(".linear_attn.in_proj_qkvz.weight") for k in keys):
+                return True
+            if any(k.endswith(".linear_attn.in_proj_qkv.weight") for k in keys):
+                return False
+    return False
+
+
 def _fp8_block_quant(hf_config: Any) -> tuple[str, tuple[int, int] | None]:
     """Detect DeepSeek-V3-style 128x128 block-fp8 from HF ``quantization_config``.
 
@@ -60,6 +93,14 @@ def _expert_quant(hf_config: Any) -> str:
                     return "nvfp4"
                 if "fp8" in expert_algo:
                     return "fp8"
+    # compressed-tensors style config (config_groups + ignore): every Linear target with
+    # fp4 group weights is NVFP4. nvidia Qwen3-Next-*-NVFP4 advertises this config format
+    # while storing modelopt two-level scales; its routed experts are Linear, not ignored.
+    groups = get("config_groups") or {}
+    for spec in (groups.values() if isinstance(groups, dict) else []):
+        w = (spec or {}).get("weights") or {}
+        if w and str(w.get("type", "")).lower() == "float" and w.get("num_bits") == 4:
+            return "nvfp4"
     return "none"
 
 
@@ -136,7 +177,38 @@ def _layer_types(text: Any) -> list[str]:
     ]
 
 
-def parse_config(hf_config: Any) -> ModelConfig:
+def _shared_expert_quant(hf_config: Any, model_path: str | None) -> str:
+    """Whether the checkpoint stores the MoE ``shared_expert`` MLP as packed NVFP4.
+
+    The nvidia/modelopt ``NVFP4`` MoE checkpoints (e.g. Qwen3.5-122B-A10B-NVFP4) ship the
+    routed experts as packed FP4 but leave ``shared_expert`` bf16 -- the model's ``ignore``
+    list explicitly excludes ``model.language_model.layers.*.mlp.shared_expert*`` from
+    quantization. The earlier unconditional ``dense_quant = "nvfp4"`` assumption broke
+    ``load_state_dict`` on those checkpoints: the NVFP4 dense kernels try to pop a
+    ``weight_scale`` that the bf16 shared_expert does not carry (KeyError).
+
+    Detection mirrors weight.py's ``.weight_scale_2`` probe: NVFP4 linears always carry a
+    paired ``weight_scale_2`` alongside ``.weight``. We scan the shards in order and, for
+    the first one carrying any shared_expert gate weight, check for that suffix. Returns
+    ``"nvfp4"`` when the shared_expert is quantized, ``"none"`` when it is bf16, and
+    ``"nvfp4"`` (the legacy assumption) when ``model_path`` is None so existing callers
+    behave unchanged.
+    """
+    if model_path is None:
+        return "nvfp4"
+    needle = ".mlp.shared_expert.gate_proj.weight_scale_2"
+    import safetensors
+    for file in iter_weight_files(model_path):
+        with safetensors.safe_open(file, framework="pt", device="cpu") as f:
+            keys = list(f.keys())
+            if any("mlp.shared_expert.gate_proj.weight" in k for k in keys):
+                if any(k.endswith(needle) for k in keys):
+                    return "nvfp4"
+                return "none"
+    return "none"
+
+
+def parse_config(hf_config: Any, model_path: str | None = None) -> ModelConfig:
     text = getattr(hf_config, "text_config", hf_config)
 
     head_dim = (
@@ -176,13 +248,16 @@ def parse_config(hf_config: Any) -> ModelConfig:
     # NVFP4. The lm_head is detected separately (only the mixed checkpoint quantizes it).
     # MoE-NVFP4 keeps the shared_expert dense MLP native FP4 (expert_quant=="nvfp4"); a dense
     # (non-MoE) modelopt checkpoint instead tags the bare .mlp.{gate,up,down}_proj as NVFP4.
-    dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config)
+    dense_quant = _shared_expert_quant(hf_config, model_path) if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config)
     lm_head_quant = _lm_head_quant(hf_config)
 
     # compressed-tensors NVFP4 (dense Qwen3.6-27B): the attention (q/k/v/o, GDN out_proj) AND
     # the dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms stay bf16. Wire the shared
     # W4A16 kernels (attn_quant=="nvfp4" routes the attention/GDN linears through them too).
-    if _compressed_tensors_nvfp4(hf_config):
+    # Gate on the actual storage format: modelopt-quantized NVFP4 checkpoints (Qwen3-Next)
+    # share the compressed-tensors config style but keep q/k/v bf16 and store two-level
+    # scales, so they take the modelopt paths (attn dequant-at-load, _shared_expert_quant).
+    if _compressed_tensors_nvfp4(hf_config) and _is_ct_storage(model_path):
         attn_quant = "nvfp4"
         dense_quant = "nvfp4"
         lm_head_quant = "none"
@@ -219,6 +294,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         value_head_dim=text.linear_value_head_dim,
         conv_kernel_dim=text.linear_conv_kernel_dim,
         output_gate=True,
+        in_proj_split=_gdn_split_layout(model_path),
     )
     # Order groups by their first layer id for deterministic iteration.
     groups = tuple(
@@ -245,6 +321,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         moe_intermediate_size=getattr(text, "moe_intermediate_size", 0),
         shared_expert_intermediate_size=getattr(text, "shared_expert_intermediate_size", 0),
         norm_topk_prob=bool(getattr(text, "norm_topk_prob", False)),
+        moe_scoring_func=str(getattr(text, "scoring_func", None) or "softmax"),
         moe_enabled=moe_enabled,
         use_qk_norm=True,
         model_type=getattr(hf_config, "model_type", "qwen3_5_moe"),