diff --git a/benchmarks/bench_decode_moe.py b/benchmarks/bench_decode_moe.py index 566217927..2ffe5dc80 100644 --- a/benchmarks/bench_decode_moe.py +++ b/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 diff --git a/docs/models.md b/docs/models.md index e4850a124..44b022092 100644 --- a/docs/models.md +++ b/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) | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f39..1219ddac5 100644 --- a/python/freetoken/engine/config.py +++ b/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: diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index d68d8ded5..4d50c0b1c 100644 --- a/python/freetoken/layers/moe.py +++ b/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 diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1f8..489f1eba4 100644 --- a/python/freetoken/models/config.py +++ b/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 diff --git a/python/freetoken/models/deepseek_v4/config.py b/python/freetoken/models/deepseek_v4/config.py index 1005b8a3b..bea0a0831 100644 --- a/python/freetoken/models/deepseek_v4/config.py +++ b/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 ) diff --git a/python/freetoken/models/gemma4/config.py b/python/freetoken/models/gemma4/config.py index 057e0a553..578726c4b 100644 --- a/python/freetoken/models/gemma4/config.py +++ b/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" diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py index 437822b51..dfc8ee235 100644 --- a/python/freetoken/models/gemma4/gguf.py +++ b/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): diff --git a/python/freetoken/models/glm4_moe/config.py b/python/freetoken/models/glm4_moe/config.py index 395af4b10..c2d156fb8 100644 --- a/python/freetoken/models/glm4_moe/config.py +++ b/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`. diff --git a/python/freetoken/models/glm_moe_dsa/config.py b/python/freetoken/models/glm_moe_dsa/config.py index c0ba25c5f..2b51a7c78 100644 --- a/python/freetoken/models/glm_moe_dsa/config.py +++ b/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. diff --git a/python/freetoken/models/gpt_oss/config.py b/python/freetoken/models/gpt_oss/config.py index f3d555d46..e3ace36d0 100644 --- a/python/freetoken/models/gpt_oss/config.py +++ b/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 ) diff --git a/python/freetoken/models/llama/config.py b/python/freetoken/models/llama/config.py index c8cc633a2..f601b8a7f 100644 --- a/python/freetoken/models/llama/config.py +++ b/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) diff --git a/python/freetoken/models/minimax_m2/config.py b/python/freetoken/models/minimax_m2/config.py index 74e877a90..0ea7e53d4 100644 --- a/python/freetoken/models/minimax_m2/config.py +++ b/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 diff --git a/python/freetoken/models/minimax_m3/config.py b/python/freetoken/models/minimax_m3/config.py index 906870e30..5b358bc43 100644 --- a/python/freetoken/models/minimax_m3/config.py +++ b/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 diff --git a/python/freetoken/models/mistral/config.py b/python/freetoken/models/mistral/config.py index 8d4d428b7..347dcd6c9 100644 --- a/python/freetoken/models/mistral/config.py +++ b/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) diff --git a/python/freetoken/models/muse_glimmer/config.py b/python/freetoken/models/muse_glimmer/config.py index 4cc756cb6..41fd117c5 100644 --- a/python/freetoken/models/muse_glimmer/config.py +++ b/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 diff --git a/python/freetoken/models/qwen2/config.py b/python/freetoken/models/qwen2/config.py index d63f4e638..b86f24230 100644 --- a/python/freetoken/models/qwen2/config.py +++ b/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) diff --git a/python/freetoken/models/qwen3/config.py b/python/freetoken/models/qwen3/config.py index d730d62b7..3d1fce80d 100644 --- a/python/freetoken/models/qwen3/config.py +++ b/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) diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index dc16cfff8..71f1be27e 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/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"), diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..4450f5b50 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -53,7 +53,7 @@ class Qwen3_5GatedDeltaNet(BaseOP): def __init__( self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size, rms_norm_eps, layer_id, expert_quant: str = "none", - attn_quant: str = "none", + attn_quant: str = "none", in_proj_split: bool = False, ): self.layer_id = layer_id # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as @@ -77,13 +77,21 @@ def __init__( self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._fp8 = self._block_fp8 or self._pertensor_fp8 + # Modelopt NVFP4 checkpoints with the pre-fused split layout (Qwen3-Next) keep + # in_proj_qkvz/in_proj_ba bf16: same two-GEMM path as fp8, but a bf16 qkvz GEMM. + self._split = self._fp8 or in_proj_split self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._fp8: - ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged - self.in_proj_qkvz = ColMerged( - hidden_size, [self.conv_dim, self.value_dim], has_bias=False - ) + if self._split: + if self._fp8: + ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged + self.in_proj_qkvz = ColMerged( + hidden_size, [self.conv_dim, self.value_dim], has_bias=False + ) + else: + self.in_proj_qkvz = LinearColParallelMerged( + hidden_size, [self.conv_dim, self.value_dim], has_bias=False + ) self.in_proj_ba = LinearColParallelMerged( hidden_size, [num_v_heads, num_v_heads], has_bias=False ) @@ -161,7 +169,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._fp8: + if self._split: qkvz = self.in_proj_qkvz.forward(hidden_states) conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) ba = self.in_proj_ba.forward(hidden_states) diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24f..bcc9b868d 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -44,6 +44,7 @@ def __init__(self, config: ModelConfig, layer_id: int): layer_id=layer_id, expert_quant=config.expert_quant, attn_quant=config.attn_quant, + in_proj_split=g.in_proj_split, ) else: self.self_attn = Qwen3_5Attention(config, layer_id) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index b34140890..9b4e2717d 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -25,7 +25,7 @@ from freetoken.utils import cached_load_hf_config, download_hf_weight from tqdm import tqdm -from .config import _compressed_tensors_nvfp4, parse_config +from .config import _compressed_tensors_nvfp4, _is_ct_storage, parse_config # Expert weights are stored pre-fused per layer: experts.gate_up_proj / experts.down_proj. _PACKED_EXPERT_PATTERN = re.compile( @@ -37,8 +37,58 @@ # weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP # head's ``mtp.layers.N.mlp.experts.*`` tensors (served text-only, dropped). _NVFP4_EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") + +# Original (un-quantized) Qwen3-Next checkpoints store each routed expert as separate +# gate/up/down tensors per expert (``model.layers.N.mlp.experts.E..weight``). Under +# offload they are packed here into the whole-layer ``experts.gate_up_proj`` / +# ``experts.down_proj`` sources the banks loader expects (include_moe_experts=True); the +# dense pass drops them like the NVFP4 per-expert tensors. +_BF16_EXPERT_PART_RE = re.compile( + r"^model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\.weight$" +) + +class _Bf16ExpertPacker: + """Packs plain-bf16 per-expert parts (``...experts.E..weight``) into the + whole-layer stacked ``experts.gate_up_proj`` / ``experts.down_proj`` sources the + bank builder expects: fuse gate|up per expert, stack all ``num_experts`` experts + per layer. Non-expert names pass straight through, so a stream can be routed + through :meth:`feed` unconditionally.""" + + def __init__(self, num_experts: int) -> None: + self._num_experts = num_experts + self._parts: dict[tuple[int, int], dict[str, torch.Tensor]] = {} + self._layers: dict[int, dict[int, tuple[torch.Tensor, torch.Tensor]]] = {} + + def feed(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]]: + m = _BF16_EXPERT_PART_RE.match(name) + if m is None: + return [(name, tensor)] + layer, idx = int(m["layer"]), int(m["idx"]) + slots = self._parts.setdefault((layer, idx), {}) + slots[m["proj"]] = tensor + if len(slots) < 3: + return [] + del self._parts[(layer, idx)] + self._layers.setdefault(layer, {})[idx] = ( + torch.cat([slots["gate_proj"], slots["up_proj"]], dim=0), + slots["down_proj"], + ) + if len(self._layers[layer]) < self._num_experts: + return [] + done = self._layers.pop(layer) + ordered = [done[i] for i in range(self._num_experts)] + prefix = f"model.layers.{layer}.mlp.experts" + return [ + (prefix + ".gate_up_proj", torch.stack([g for g, _ in ordered], dim=0)), + (prefix + ".down_proj", torch.stack([d for _, d in ordered], dim=0)), + ] + + def leftovers(self) -> tuple[list, list]: + """Incomplete keys at end of stream: (partial experts, partial layers).""" + return sorted(self._parts), sorted(self._layers) _NVFP4_EXPERT_KEY_RE = re.compile( - r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"^model\.(?:language_model\.)?layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" ) _NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( @@ -170,6 +220,51 @@ def _try_fuse( return None +# Qwen3-Next modelopt checkpoints ship the pre-fused GDN input projections in HF's +# interleaved layout (Qwen3NextGatedDeltaNet.fix_query_key_value_ordering): per k-head +# group the raw rows are [q_g | k_g | v_pair | z_pair] (qkvz) and [b_pair | a_pair] (ba). +# The engine's GDN instead splits contiguous [q|k|v|z] / [b|a] blocks, so the rows must +# be de-interleaved at load time -- otherwise q/k/v/z (and b/a) are silently scrambled. +def _gdn_split_reorder(name: str, tensor: torch.Tensor, g) -> torch.Tensor: + nk, nv, hd = g.num_key_heads, g.num_value_heads, g.key_head_dim + ratio = nv // nk + tail = tensor.shape[1:] + rows = tensor.shape[0] + if ".linear_attn.in_proj_qkvz." in name: + grouped_rows = nk * (2 + 2 * ratio) * hd + if rows == grouped_rows: + # full weight: one row per logical output row + grouped = tensor.view(nk, 2 + 2 * ratio, hd, *tail) + take = lambda sl: grouped[:, sl].reshape(-1, *tail) + elif grouped_rows % rows == 0 and grouped_rows // rows == hd: + # per-128-output-row block scale (fp8 weight_scale_inv): each scale row + # covers exactly one head_dim-sized chunk of the interleaved layout + grouped = tensor.view(nk, 2 + 2 * ratio, *tail) + take = lambda sl: grouped[:, sl].reshape(-1, *tail) + else: + raise ValueError( + f"cannot de-interleave {name}: {rows} rows vs {grouped_rows} logical rows" + ) + return torch.cat( + [ + take(0), # q (all k-heads) + take(1), # k + take(slice(2, 2 + ratio)), # v (all v-heads) + take(slice(2 + ratio, None)), # z + ], + dim=0, + ) + if ".linear_attn.in_proj_ba." in name: + if rows != 2 * nv: + raise ValueError(f"cannot de-interleave {name}: {rows} rows vs {2 * nv}") + grouped = tensor.view(nk, 2, ratio, *tail) + return torch.cat( + [grouped[:, 0].reshape(-1, *tail), grouped[:, 1].reshape(-1, *tail)], + dim=0, + ) + return tensor + + def iter_weights( model_path: str, device: torch.device, @@ -178,8 +273,38 @@ def iter_weights( include_non_moe: bool, ) -> Iterator[tuple[str, torch.Tensor]]: hf_config = cached_load_hf_config(model_path) - config = parse_config(hf_config) - if _compressed_tensors_nvfp4(hf_config): + config = parse_config(hf_config, model_path) + try: + g = config.linear_attention_group() + except Exception: + g = None + # Only the pre-fused split layout (probed at parse time) needs the de-interleave; + # fused in_proj checkpoints (Qwen3.6) and unfused parts pass through unchanged. + reorder = g is not None and getattr(g, "in_proj_split", False) + for name, tensor in _iter_weights_flat( + model_path, device, + include_moe_experts=include_moe_experts, include_non_moe=include_non_moe, + ): + if reorder and ( + ".linear_attn.in_proj_qkvz." in name or ".linear_attn.in_proj_ba." in name + ): + tensor = _gdn_split_reorder(name, tensor, g) + yield name, tensor + + +def _iter_weights_flat( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + hf_config = cached_load_hf_config(model_path) + config = parse_config(hf_config, model_path) + # Gate on the storage format: modelopt-quantized NVFP4 checkpoints (Qwen3-Next) share + # the compressed-tensors config style but store modelopt scales and a pre-fused GDN + # in_proj -- they take the pure-NVFP4 path below. + if _compressed_tensors_nvfp4(hf_config) and _is_ct_storage(model_path): # Dense compressed-tensors NVFP4 (e.g. Qwen3.6-27B): attn (q/k/v/o, GDN out_proj) + # dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms bf16. yield from _iter_weights_compressed_tensors( @@ -221,6 +346,7 @@ def iter_weights( shared_buf: dict[str, dict[str, torch.Tensor]] = {} nvfp4_shared_buf: dict[str, dict[str, tuple]] = {} fuse_buf: dict[str, dict[int, torch.Tensor]] = {} + bf16_packer = _Bf16ExpertPacker(config.num_experts) for file in tqdm( iter_weight_files(model_path), @@ -233,7 +359,13 @@ def iter_weights( # Per-expert NVFP4 tensors go to the offload cache (load_nvfp4_expert_sources), # not the dense pass. bf16-base stacked experts (experts.gate_up_proj) have no # ``.mlp.experts..`` so they are unaffected and still hit _PACKED_EXPERT. - if _NVFP4_EXPERT_RE.search(raw_name): + if _NVFP4_EXPERT_RE.search(raw_name) and ( + not raw_name.endswith(".weight") + or raw_name.removesuffix(".weight") + ".weight_scale" in keyset + ): + # modelopt-quantized per-expert tensor (or its scale): handled by + # load_nvfp4_expert_source_banks, never this dense pass. A plain + # bf16 per-expert .weight (no scale sibling) falls through below. continue # Standalone modelopt scales are consumed with their .weight, never yielded. if raw_name.endswith(_SCALE_SUFFIXES): @@ -243,6 +375,12 @@ def iter_weights( if name is None: continue + if _BF16_EXPERT_PART_RE.match(name) is not None: + if not include_moe_experts: + continue # routed experts live in the offload banks, not here + yield from bf16_packer.feed(name, f.get_tensor(raw_name)) + continue + is_expert = _PACKED_EXPERT_PATTERN.match(name) is not None if is_expert and not include_moe_experts: continue @@ -290,6 +428,9 @@ def iter_weights( assert not shared_buf, f"Incomplete shared-expert merges: {list(shared_buf.keys())}" assert not nvfp4_shared_buf, f"Incomplete NVFP4 shared-expert merges: {list(nvfp4_shared_buf.keys())}" assert not fuse_buf, f"Incomplete projection fusions: {list(fuse_buf.keys())}" + parts, layers = bf16_packer.leftovers() + assert not parts, f"Incomplete per-expert parts: {parts[:5]}..." + assert not layers, f"Incomplete per-expert layers: {layers}" # ====================================================================================== @@ -667,8 +808,9 @@ def iter_weights_parallel( chunk: int = 8 << 20, ) -> Iterator[tuple[str, torch.Tensor]]: """experts-only parallel read via the common chunked multi-threaded O_DIRECT reader. - Qwen3.5 stores experts pre-fused/pre-stacked per layer (already ``[E, ...]``), so no - merge/stack -- just rename and yield; bank builder places by name as the serial path.""" + Pre-fused/pre-stacked Qwen3.5 layers pass through renamed; plain-bf16 per-expert + checkpoints (original Qwen3-Next) get the same fuse+stack packing as the serial path + via ``_Bf16ExpertPacker``, so both layouts land in the banks byte-identically.""" assert include_moe_experts and not include_non_moe, ( "qwen3_5_moe parallel reader is experts-only (used by load_moe_expert_sources)" ) @@ -677,14 +819,26 @@ def iter_weights_parallel( if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe weight loading currently supports TP=1 only") + config = parse_config(cached_load_hf_config(model_path)) + packer = _Bf16ExpertPacker(config.num_experts) + def _is_expert(raw_name: str) -> bool: name = _rename(raw_name) - return name is not None and _PACKED_EXPERT_PATTERN.match(name) is not None + return name is not None and ( + _PACKED_EXPERT_PATTERN.match(name) is not None + or _BF16_EXPERT_PART_RE.match(name) is not None + ) for raw_name, tensor in iter_expert_tensors_parallel( model_path, _is_expert, workers=workers, chunk=chunk ): - yield _rename(raw_name), tensor + # Pre-packed layers pass through the packer unchanged; plain bf16 per-expert + # parts get fused and stacked per layer exactly like the serial path. + yield from packer.feed(_rename(raw_name), tensor) + + parts, layers = packer.leftovers() + assert not parts, f"Incomplete per-expert parts: {parts[:5]}..." + assert not layers, f"Incomplete per-expert layers: {layers}" # ====================================================================================== @@ -715,10 +869,10 @@ def _is_expert(raw_name: str) -> bool: } _FP8_KIND_SUFFIXES = (".weight_scale_inv", ".weight") -# Routed-expert checkpoint key (per-expert, un-fused). ``mtp.layers...`` is excluded by the +# Routed-expert checkpoint key (per-expert, un-fused). The optional `language_model.` anchor # ``model.language_model.`` anchor, so the parallel reader only sees the real experts. _FP8_EXPERT_RE = re.compile( - r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"^model\.(?:language_model\.)?layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." r"(?Pgate|up|down)_proj\.(?Pweight|weight_scale_inv)$" ) @@ -790,7 +944,7 @@ def _iter_weights_fp8( # Resident (non-offload) experts: build the stacked fp8 banks once via the shared # parallel reader (pageable host -- the engine copies the per-layer slices to GPU), # then yield per-layer views into the Fp8ResidentMoE buffers. - config = parse_config(cached_load_hf_config(model_path)) + config = parse_config(cached_load_hf_config(model_path), model_path) if not config.is_moe: return # dense checkpoint: no routed experts to build as resident banks L, E, H, I, dense = _moe_dims(config) @@ -976,7 +1130,7 @@ def _load(sink) -> None: for li in tqdm(range(L), desc="Loading fp8 experts (serial)", disable=not primary): layer = dense + li for e in range(E): - p = f"model.language_model.layers.{layer}.mlp.experts.{e}" + p = f"model.layers.{layer}.mlp.experts.{e}" for proj in ("gate", "up", "down"): for kind in ("weight", "weight_scale_inv"): key = f"{p}.{proj}_proj.{kind}" @@ -1030,7 +1184,7 @@ def _load(sink) -> None: gu_rows = torch.empty(E, 2 * I, H, dtype=torch.bfloat16, device=device) dn_rows = torch.empty(E, H, I, dtype=torch.bfloat16, device=device) for e in range(E): - p = f"model.language_model.layers.{layer}.mlp.experts.{e}" + p = f"model.layers.{layer}.mlp.experts.{e}" gu_rows[e, :I] = _deq(f"{p}.gate_proj") gu_rows[e, I:] = _deq(f"{p}.up_proj") dn_rows[e] = _deq(f"{p}.down_proj") diff --git a/python/freetoken/models/qwen3_moe/config.py b/python/freetoken/models/qwen3_moe/config.py index fcec5420a..cd1666782 100644 --- a/python/freetoken/models/qwen3_moe/config.py +++ b/python/freetoken/models/qwen3_moe/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", diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca01..206f9f356 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -65,6 +65,14 @@ class ModelSpec: "freetoken.models.qwen3_5_moe", "Qwen3_5MoEForCausalLM", ), + # Qwen3-Next (model_type qwen3_next): Gated DeltaNet + MoE with FP8 block quantization. + # Same architecture family as Qwen3.5 MoE (linear_attn + full_attn interleaving, MoE + # experts) -- reuses the qwen3_5_moe package. Weights are flat (model.layers.*, no + # language_model prefix), block-fp8 (e4m3 + weight_scale_inv, 128x128 blocks). + "Qwen3NextForCausalLM": ModelSpec( + "freetoken.models.qwen3_5_moe", + "Qwen3_5MoEForCausalLM", + ), # Muse-Glimmer-30B (model_type muse_glimmer): multimodal wrapper config (text tower in # text_config, weights under model.language_model.); served text-only. Dense gated GQA # with a [SWA x3, full] pattern -- full layers are NoPE -- weightless qk norms, centered diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index fe7e417d7..7c077c592 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -18,13 +18,18 @@ def _torch_fused_topk( topk: int, renormalize: bool, num_token_non_padded: torch.Tensor | None, + scoring: str = "softmax", ) -> Tuple[torch.Tensor, torch.Tensor]: - """Pure-torch softmax router matching triton_kernels.topk (Windows fallback). + """Pure-torch router matching triton_kernels.topk (Windows fallback). - Softmax over all experts, select the top-k, and (when ``renormalize``) rescale the - selected weights to sum to 1 -- the standard fused-MoE routing convention. + Softmax (default) or sigmoid (Qwen3-Next ``scoring_func="sigmoid"``) over all + experts, select the top-k, and (when ``renormalize``) rescale the selected + weights to sum to 1 -- the standard fused-MoE routing convention. """ - probs = torch.softmax(gating_output.float(), dim=-1) + if scoring == "sigmoid": + probs = torch.sigmoid(gating_output.float()) + else: + probs = torch.softmax(gating_output.float(), dim=-1) topk_weights, topk_ids = torch.topk(probs, topk, dim=-1) if renormalize: topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) @@ -41,9 +46,15 @@ def fused_topk( topk: int, renormalize: bool, num_token_non_padded: torch.Tensor | None = None, + scoring: str = "softmax", ) -> Tuple[torch.Tensor, torch.Tensor]: assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + # Sigmoid scoring has no triton_kernels counterpart: always take the torch path + # (which also covers the non-power-of-2 topk case, e.g. Qwen3-Next's 10). + if scoring != "softmax": + return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded, scoring) + from freetoken.kernel.backend import is_triton_kernels_installed # triton_kernels ships no Windows wheel, and unlike flashinfer/sgl_kernel it is not one @@ -60,10 +71,16 @@ def fused_topk( "(numerically equivalent, slower). Expected on Windows (no wheel); on Linux " "install triton_kernels to restore the fused router." ) - return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded) + return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded, scoring) from triton_kernels.topk import topk as triton_kernels_topk + # triton_kernels' _topk_forward uses tl.arange(0, N_EXPTS_ACT) which requires + # topk to be a power of 2 (e.g. 8 for Qwen3.5). Models with non-power-of-2 topk + # (e.g. Qwen3-Next's 10) hit a Triton compilation error; fall back to pure-torch. + if topk & (topk - 1) != 0: + return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded, scoring) + logits = gating_output.float() softmax_first = not renormalize if softmax_first: @@ -434,12 +451,14 @@ def forward( renormalize: bool, activation: str = "silu", apply_router_weight_on_input: bool = False, + scoring: str = "softmax", ) -> torch.Tensor: topk_weights, topk_ids = fused_topk( hidden_states=hidden_states, gating_output=gating_output, topk=topk, renormalize=renormalize, + scoring=scoring, ) return fused_experts_impl( hidden_states, diff --git a/tests/models/test_qwen3_next_weights.py b/tests/models/test_qwen3_next_weights.py new file mode 100644 index 000000000..36fedbbd0 --- /dev/null +++ b/tests/models/test_qwen3_next_weights.py @@ -0,0 +1,116 @@ +"""Qwen3-Next GDN in_proj de-interleaving and sigmoid routing. + +Qwen3-Next modelopt checkpoints ship the pre-fused GDN input projections in HF's +interleaved layout (``Qwen3NextGatedDeltaNet.fix_query_key_value_ordering``): +per k-head group the raw rows are ``[q_g | k_g | v_pair | z_pair]`` (qkvz) and +``[b_pair | a_pair]`` (ba). The engine's GDN instead splits contiguous +``[q|k|v|z]`` / ``[b|a]`` blocks, so ``_gdn_split_reorder`` must de-interleave +at load time. Loading the interleaved layout as-is scrambles q/k/v/z silently -- +output is fluent garbage and decode speed looks normal -- so the permutation is +verified by round-trip here: build the interleaved layout from a known +contiguous reference and require the loader to recover it exactly, for both the +bf16 weights and the per-128-row fp8 scale blocks that alias head_dim. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.models.qwen3_5_moe.weight import _gdn_split_reorder +from freetoken.moe.fused import _torch_fused_topk + +# Small but non-degenerate GDN geometry: 4 k-head groups, 2 v-heads per group. +NK, NV, HD, HIDDEN = 4, 8, 8, 20 +RATIO = NV // NK + + +def _g(): + return SimpleNamespace(num_key_heads=NK, num_value_heads=NV, key_head_dim=HD) + + +def _interleave_qkvz(contig: torch.Tensor) -> torch.Tensor: + """Inverse of _gdn_split_reorder: HF's fix_query_key_value_ordering layout.""" + q, k, v, z = torch.split(contig, [NK * HD, NK * HD, NV * HD, NV * HD], dim=0) + + def head(t, h): + return t[h * HD : (h + 1) * HD] + + parts = [] + for g in range(NK): + parts += [head(q, g), head(k, g)] + parts += [head(v, h) for h in range(g * RATIO, (g + 1) * RATIO)] + parts += [head(z, h) for h in range(g * RATIO, (g + 1) * RATIO)] + return torch.cat(parts, dim=0) + + +def _interleave_ba(contig: torch.Tensor) -> torch.Tensor: + b, a = torch.split(contig, [NV, NV], dim=0) + parts = [] + for g in range(NK): + parts += [b[h : h + 1] for h in range(g * RATIO, (g + 1) * RATIO)] + parts += [a[h : h + 1] for h in range(g * RATIO, (g + 1) * RATIO)] + return torch.cat(parts, dim=0) + + +def test_in_proj_qkvz_round_trip(): + contig = torch.arange(NK * (2 + 2 * RATIO) * HD * HIDDEN, dtype=torch.float32) + contig = contig.view(-1, HIDDEN) + interleaved = _interleave_qkvz(contig) + name = "model.layers.0.linear_attn.in_proj_qkvz.weight" + out = _gdn_split_reorder(name, interleaved, _g()) + assert torch.equal(out, contig) + + +def test_in_proj_ba_round_trip(): + contig = torch.arange(2 * NV * HIDDEN, dtype=torch.float32).view(-1, HIDDEN) + interleaved = _interleave_ba(contig) + name = "model.layers.0.linear_attn.in_proj_ba.weight" + out = _gdn_split_reorder(name, interleaved, _g()) + assert torch.equal(out, contig) + + +def test_in_proj_qkvz_scale_blocks_round_trip(): + """The fp8 per-128-row ``weight_scale_inv`` aliases head_dim: each scale row + covers exactly one head_dim-sized chunk of the interleaved layout.""" + n_blocks = NK * (2 + 2 * RATIO) # one scale row per (group, slot) block + contig_scale = torch.arange(n_blocks, dtype=torch.float32).view(-1, 1) + q, k, v, z = torch.split(contig_scale, [NK, NK, NV, NV], dim=0) + parts = [] + for g in range(NK): + parts += [q[g : g + 1], k[g : g + 1]] + parts += [v[h : h + 1] for h in range(g * RATIO, (g + 1) * RATIO)] + parts += [z[h : h + 1] for h in range(g * RATIO, (g + 1) * RATIO)] + interleaved_scale = torch.cat(parts, dim=0) + name = "model.layers.0.linear_attn.in_proj_qkvz.weight_scale_inv" + out = _gdn_split_reorder(name, interleaved_scale, _g()) + assert torch.equal(out, contig_scale) + + +def test_other_names_pass_through(): + t = torch.arange(2 * NV * HIDDEN, dtype=torch.float32).view(-1, HIDDEN) + name = "model.layers.0.linear_attn.out_proj.weight" + out = _gdn_split_reorder(name, t, _g()) + assert out is t + + +def test_ambiguous_qkvz_rows_rejected(): + bad = torch.zeros(NK * (2 + 2 * RATIO) * HD + 1, HIDDEN) + name = "model.layers.0.linear_attn.in_proj_qkvz.weight" + with pytest.raises(ValueError, match="cannot de-interleave"): + _gdn_split_reorder(name, bad, _g()) + + +def test_sigmoid_router_matches_reference(): + """Qwen3-Next scores experts with sigmoid + top-10: verify the torch router + against a hand-rolled reference, including renormalization.""" + torch.manual_seed(0) + logits = torch.randn(5, 32) + weights, ids = _torch_fused_topk(logits, 10, True, None, "sigmoid") + probs = torch.sigmoid(logits.float()) + ref_weights, ref_ids = torch.topk(probs, 10, dim=-1) + ref_weights = ref_weights / ref_weights.sum(dim=-1, keepdim=True) + assert torch.equal(ids, ref_ids) + assert torch.allclose(weights, ref_weights)