From 8030fd04406b2f10122355baefd3c71e7bfa12a1 Mon Sep 17 00:00:00 2001 From: Ahmed Hassanin Date: Sun, 30 Aug 2026 16:04:15 -0400 Subject: [PATCH 1/2] feat(models): support poolside/Laguna-S-2.1-NVFP4 Laguna S 2.1 is a 48-layer, 117.6B/8.5B-active MoE with several traits no existing model module covers: * Per-layer attention width: 48 heads on the 12 full-attention layers, 72 on the 36 sliding ones (window 512), with KV fixed at 8 heads. * Two RoPE configs nested under `rope_parameters` by layer type -- YaRN (theta 500k, partial_rotary_factor 0.5) for full, default (theta 10k, full rotary) for sliding. There is no top-level `rope_theta`. * Softplus output gating applied per head before o_proj, and unconditional QK-RMSNorm with no config flag. * Sigmoid router with a selection-only correction bias, `norm_topk_prob`, a routed scaling factor of 2.5, and one always-on shared expert. Layer 0 is dense. * Mixed-precision experts: layers 1-39 ship compressed-tensors NVFP4 (`weight_packed` / `weight_scale` / `weight_global_scale`, group 16), while 40-47 ship plain bf16. `ModelConfig.num_qo_heads` is set to the maximum per-layer head count rather than the nominal one: the Triton backend sizes its decode scratch and CUDA-graph capture buffers from that single value and never resizes them on replay, so the nominal 48 would let the 72-head layers write out of bounds under graph replay. The bf16 expert layers are quantized to NVFP4 at conversion so one bank layout and one kernel path covers all 47 MoE layers. `Nvfp4ExpertSourceSpec` gains three optional hooks to keep this on the shared loader (with its completion tracking, pin-after-fill and FTW streaming) instead of a private one: `kind_map` for compressed-tensors key names, `global_transform` for its quant-side scale convention, and `synthesize_layer` for bank layers absent from the checkpoint. Existing specs leave all three None and are unaffected. Verified against the checkpoint and its reference implementation: converted banks dequantize bit-exactly for the packed layers; the synthesized layers land at 9.51% relative RMS against a 9.41% floor measured with ideal fp32 scales; RoPE matches `modeling_laguna.py` to ~5e-07 for both layer types; router expert ids are identical and weights match once the reference's bf16 cast is applied. Reasoning and tool-call parsers resolve automatically: Laguna's think tags are ``/`` with the opening tag pre-filled by the chat template, and its tool envelope is GLM-4.7's grammar without the inter-tag newlines. --- python/freetoken/kernel/aot_models.py | 16 ++ python/freetoken/models/config.py | 9 + python/freetoken/models/laguna/__init__.py | 11 + python/freetoken/models/laguna/attention.py | 143 ++++++++++ python/freetoken/models/laguna/config.py | 238 ++++++++++++++++ python/freetoken/models/laguna/model.py | 91 ++++++ python/freetoken/models/laguna/moe.py | 94 ++++++ python/freetoken/models/laguna/weight.py | 300 ++++++++++++++++++++ python/freetoken/models/nvfp4_banks.py | 107 ++++++- python/freetoken/models/register.py | 12 + python/freetoken/server/args.py | 15 +- python/freetoken/server/generation.py | 6 + python/freetoken/server/reasoning_parser.py | 1 + 13 files changed, 1035 insertions(+), 8 deletions(-) create mode 100644 python/freetoken/models/laguna/__init__.py create mode 100644 python/freetoken/models/laguna/attention.py create mode 100644 python/freetoken/models/laguna/config.py create mode 100644 python/freetoken/models/laguna/model.py create mode 100644 python/freetoken/models/laguna/moe.py create mode 100644 python/freetoken/models/laguna/weight.py diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index c9c2fb98e..7faeca5e0 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -339,6 +339,22 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int kv_groups=((2, 128),), # sliding and full (NoPE) layers share the same geometry aliases=("RedHatAI/Muse-Glimmer-30B-NVFP4",), ), + AotModel( + # 1:3 full/SWA hybrid (window 512), but both groups are 8 kv heads x 128 + # head_dim, so they share one store variant (gpt-oss precedent). Only the + # QUERY width differs per layer (48 full / 72 SWA), which the store kernel + # never sees. Routed experts keep the native "nvfp4" 6-bank layout: I=1024 + # is narrow enough that select_nvfp4_backend resolves to the Triton + # inline-dequant kernels, so the marlin/b12x repacks never run. + name="poolside/Laguna-S-2.1-NVFP4", + architecture="LagunaForCausalLM", + hidden_size=3072, + kv_groups=((8, 128),), + top_k=10, + moe_intermediate_size=1024, + expert_formats=("nvfp4",), + arch_aliases=("LagunaForConditionalGeneration",), + ), AotModel( name="meta-llama/Llama-3.1-8B-Instruct", architecture="LlamaForCausalLM", diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 229cce812..ad5519b90 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -299,6 +299,15 @@ class ModelConfig: moe_weight_format: str | None = None swiglu_limit: float | None = None hidden_act_alpha: float = 1.702 + # Per-layer head count (Laguna S 2.1: 48 for full, 72 for SWA layers). Generic: + # hybrid models whose attention width varies by layer type set this and the model + # module sizes its projections from it (``num_qo_heads`` stays the max, so the + # backends' per-layer scratch and CUDA-graph buffers fit the widest layer). + num_attention_heads_per_layer: tuple[int, ...] | None = None + # Laguna (laguna) payload (LagunaArgs): the attention output-gating mode ("per-head" + # vs per-element) and its per-layer types. Opaque to model-agnostic engine code; + # None for every other model. + laguna_args: Any | None = None # Full DeepseekV4Args payload for the DSV4-specific machinery (MLA sparse attention, # CSA/HCA compressors, Lightning Indexer, manifold-constrained Hyper-Connections, # hash routing). Opaque to model-agnostic engine code; None for non-DSV4 models. diff --git a/python/freetoken/models/laguna/__init__.py b/python/freetoken/models/laguna/__init__.py new file mode 100644 index 000000000..a5960bbc9 --- /dev/null +++ b/python/freetoken/models/laguna/__init__.py @@ -0,0 +1,11 @@ +from .config import LagunaArgs, parse_config +from .model import LagunaForCausalLM +from .weight import iter_weights, setup_offload_expert_banks + +__all__ = [ + "LagunaArgs", + "LagunaForCausalLM", + "parse_config", + "iter_weights", + "setup_offload_expert_banks", +] diff --git a/python/freetoken/models/laguna/attention.py b/python/freetoken/models/laguna/attention.py new file mode 100644 index 000000000..dfdec7017 --- /dev/null +++ b/python/freetoken/models/laguna/attention.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F +from freetoken.attention import AttentionSpec +from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info +from freetoken.layers import BaseOP, LinearReplicated +from freetoken.layers.norm import RMSNorm +from freetoken.layers.rotary import get_rope +from freetoken.models.config import FullAttentionGroupConfig, SWAAttentionGroupConfig +from freetoken.utils import nvtx_annotate + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +class LagunaAttention(BaseOP): + """Laguna attention with per-layer head count, QK RMSNorm, and per-head gating.""" + + def __init__(self, config: ModelConfig, layer_id: int): + self.layer_id = layer_id + group = config.attention_group_for_layer(layer_id) + self.is_swa = isinstance(group, SWAAttentionGroupConfig) + if not isinstance(group, (FullAttentionGroupConfig, SWAAttentionGroupConfig)): + raise ValueError(f"LagunaAttention does not support {group.kind!r} layers") + + # Per-layer head count (SWA 72 vs full 48 for S 2.1). Falls back to global. + if getattr(config, "num_attention_heads_per_layer", None) is not None: + heads = config.num_attention_heads_per_layer[layer_id] # type: ignore[index] + else: + heads = config.num_qo_heads + self.num_qo_heads_global = int(heads) + self.head_dim = group.head_dim + self.num_kv_heads_global = group.num_kv_heads + + # TP partition. Every projection below is sized from the *_global head counts + # (Laguna keeps q/k/v/g/o split, and the per-layer head count varies), so this + # model is TP=1 only; assert rather than silently replicate into wrong shapes. + tp_size = get_tp_info().size + assert tp_size == 1, ( + "Laguna does not support tensor parallelism: its per-layer head counts " + f"(48 full / 72 SWA) are not sharded by this module (tp_size={tp_size})" + ) + + self.q_dim = self.num_qo_heads_global * self.head_dim + self.kv_dim = self.num_kv_heads_global * self.head_dim + # Laguna stores split q/k/v/g/o as separate linears (see configuration_laguna.py + # base_model_tp_plan). We keep them split to avoid a fusion layer. + hidden = config.hidden_size + self.q_proj = LinearReplicated(hidden, self.q_dim, has_bias=False) + self.k_proj = LinearReplicated(hidden, self.kv_dim, has_bias=False) + self.v_proj = LinearReplicated(hidden, self.kv_dim, has_bias=False) + self.o_proj = LinearReplicated(self.num_qo_heads_global * self.head_dim, hidden, has_bias=False) + + # Per-head gating (config.json gating "per-head"). + args = config.laguna_args + gating = args.gating if args is not None else "per-head" + self.gating_enabled = bool(gating) + self.gate_per_head = gating == "per-head" + if self.gating_enabled: + if self.gate_per_head: + g_out = self.num_qo_heads_global + else: + # per-element + g_out = self.num_qo_heads_global * self.head_dim + self.g_proj = LinearReplicated(hidden, g_out, has_bias=False) + else: + self.g_proj = None # type: ignore[assignment] + + # QK RMSNorm (LagunaRMSNorm in HF -> vanilla RMSNorm with scale). + eps = config.rms_norm_eps + self.q_norm = RMSNorm(self.head_dim, eps=eps) + self.k_norm = RMSNorm(self.head_dim, eps=eps) + + # Rope per group (full=YARN, swa=default). + rotary_config = group.rotary_config + self.rotary = get_rope( + head_dim=self.head_dim, + rotary_dim=rotary_config.rotary_dim, + max_position=rotary_config.max_position, + base=rotary_config.base, + rope_scaling=( + tuple(rotary_config.scaling.items()) + if rotary_config.scaling + else None + ), + ) + self.attn_spec = AttentionSpec( + sliding_window=group.sliding_window if self.is_swa else None, + sm_scale=config.attn_sm_scale, + ) + + @nvtx_annotate("LAGUNA_MHA") + def forward(self, x: torch.Tensor) -> torch.Tensor: + ctx = get_global_ctx() + T = x.shape[0] + + # Split projections. + q = self.q_proj.forward(x) + k = self.k_proj.forward(x) + v = self.v_proj.forward(x) + + # Per-head QK RMSNorm, applied BEFORE RoPE (modeling_laguna.py). Normalize in + # place on a [-1, head_dim] view of the projection output -- the flat and + # [T, heads, head_dim] views share storage, so RoPE below sees the normed values. + self.q_norm.forward_inplace(q.view(-1, self.head_dim)) + self.k_norm.forward_inplace(k.view(-1, self.head_dim)) + + # RoPE rotates the first ``rotary_dim`` dims of each head (partial rotary: 64 of + # 128 on full layers, 128 on SWA). The kernel takes [T, heads*head_dim] and + # derives the head count from head_size, so it handles both widths unchanged. + pos = ctx.batch.positions.reshape(-1) + if pos.device != q.device or pos.dtype != torch.long: + pos = pos.to(device=q.device, dtype=torch.long) + q, k = self.rotary.forward(pos, q, k) + + o = ctx.attn_backend.forward( + q.view(T, self.num_qo_heads_global, self.head_dim).contiguous(), + k.contiguous(), + v.contiguous(), + self.layer_id, + ctx.batch, + attn_spec=self.attn_spec, + ) + o = o.reshape(T, self.num_qo_heads_global * self.head_dim) + + # Softplus output gating, applied BEFORE o_proj (modeling_laguna.py:452-461). + if self.gating_enabled: + gate = F.softplus(self.g_proj.forward(x).float()).to(o.dtype) + if self.gate_per_head: + # [T, heads] broadcast across head_dim + o = (o.view(T, self.num_qo_heads_global, self.head_dim) * gate.unsqueeze(-1)).reshape( + T, self.num_qo_heads_global * self.head_dim + ) + else: + o = o * gate + return self.o_proj.forward(o) + + +__all__ = ["LagunaAttention"] diff --git a/python/freetoken/models/laguna/config.py b/python/freetoken/models/laguna/config.py new file mode 100644 index 000000000..0e0df24a9 --- /dev/null +++ b/python/freetoken/models/laguna/config.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from freetoken.models.config import ( + FullAttentionGroupConfig, + ModelConfig, + RotaryConfig, + SWAAttentionGroupConfig, + detect_compressed_tensors_nvfp4, +) + + +@dataclass(frozen=True) +class LagunaArgs: + """Laguna-specific attention output gating (``modeling_laguna.py`` ``g_proj``). + + ``gating`` is ``"per-head"`` (one gate per head, broadcast over head_dim), + ``True``/``"per-element"`` (one gate per channel) or ``False`` (no gating). + """ + + gating: str | bool + gating_types: tuple[str, ...] | None = None + + +def parse_config(hf_config: Any) -> ModelConfig: + # Laguna is not a multimodal wrapper; but keep text_config indirection + # for parity with gemma4/qwen4_exp in case a future variant wraps. + text = getattr(hf_config, "text_config", None) + if text is not None: + cfg = text + top_architectures = getattr(hf_config, "architectures", None) + top_cfg = hf_config + else: + cfg = hf_config + top_architectures = getattr(cfg, "architectures", None) + top_cfg = cfg + + head_dim = int(getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads) + num_kv_heads = int(getattr(cfg, "num_key_value_heads", cfg.num_attention_heads)) + max_position = int(getattr(cfg, "max_position_embeddings", 4096)) + rms_norm_eps = float(getattr(cfg, "rms_norm_eps", 1e-6)) + hidden_act = str(getattr(cfg, "hidden_act", "silu")) + vocab_size = int(cfg.vocab_size) + hidden_size = int(cfg.hidden_size) + intermediate_size = int(getattr(cfg, "intermediate_size", 0)) + num_layers = int(cfg.num_hidden_layers) + num_attention_heads = int(cfg.num_attention_heads) + num_attention_heads_per_layer = getattr(cfg, "num_attention_heads_per_layer", None) + if num_attention_heads_per_layer is not None: + num_attention_heads_per_layer = tuple(int(x) for x in num_attention_heads_per_layer) + assert len(num_attention_heads_per_layer) == num_layers, ( + f"num_attention_heads_per_layer len {len(num_attention_heads_per_layer)} != num_layers {num_layers}" + ) + # ModelConfig.num_qo_heads sizes shared, layer-agnostic state: the Triton + # backend's decode scratch and its CUDA-graph capture buffers + # (attention/triton.py num_q_heads -> init_capture_graph). Laguna S 2.1's + # config says 48 but its 36 SWA layers run 72 heads, and the captured buffer + # is never resized on replay -- so it must be the MAX over layers, not the + # nominal value. Per-layer projections read num_attention_heads_per_layer. + num_qo_heads = max(num_attention_heads_per_layer) + else: + num_qo_heads = num_attention_heads + tie_word_embeddings = bool(getattr(cfg, "tie_word_embeddings", False)) + sliding_window = getattr(cfg, "sliding_window", None) + # Layer types: expected 1:3 full:swa pattern ("full_attention"/"sliding_attention"). + layer_types = getattr(cfg, "layer_types", None) + if layer_types is None: + layer_types = ["full_attention"] * num_layers + layer_types = list(layer_types) + assert len(layer_types) == num_layers + + # Rope: Laguna stores two dicts under rope_parameters. + rope_parameters = getattr(cfg, "rope_parameters", None) or {} + if not isinstance(rope_parameters, dict): + rope_parameters = {} + full_rope = rope_parameters.get("full_attention") + swa_rope = rope_parameters.get("sliding_attention") + # Fall back to outer rope_parameters when per-type not present (BF16 future etc.). + if full_rope is None: + full_rope = {k: v for k, v in rope_parameters.items() if k not in ("full_attention", "sliding_attention")} + if not full_rope: + full_rope = {"rope_type": "default", "rope_theta": 500000.0} + if swa_rope is None: + # Laguna without separate SWA rope (unlikely) — reuse full. + swa_rope = full_rope + + # Full rope is YARN for S 2.1. + full_rope_type = str(full_rope.get("rope_type", "default")) + full_rope_theta = float(full_rope.get("rope_theta", 500000.0)) + full_partial = float(full_rope.get("partial_rotary_factor", 0.5)) + swa_rope_type = str(swa_rope.get("rope_type", "default")) + swa_rope_theta = float(swa_rope.get("rope_theta", 10000.0)) + swa_partial = float(swa_rope.get("partial_rotary_factor", 1.0)) + + # Build RotaryConfigs. Full uses YARN scaling dict, SWA uses default. + # For YARN we carry the full dict (factor/beta/attention_factor...) but + # only the scalar keys are hashable for get_rope cache key. + def _rope_scaling(rope_dict: dict[str, Any], rope_type: str) -> dict[str, Any] | None: + if rope_type == "default": + return None + # YARN needs the standard params; filter to scalar non-list values. + # layers/rotary.py consumes rope_type/yarn params explicitly. + scaling: dict[str, Any] = {"rope_type": rope_type} + for k in ("factor", "beta_fast", "beta_slow", "original_max_position_embeddings", + "attention_factor", "truncate", "mscale", "mscale_all_dim"): + if k in rope_dict: + scaling[k] = rope_dict[k] + # Also carry rope_theta via base field, not here. + return scaling + + full_rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=int(head_dim * full_partial), + max_position=max_position, + base=full_rope_theta, + scaling=_rope_scaling(full_rope, full_rope_type), + ) + swa_rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=int(head_dim * swa_partial), + max_position=max_position, + base=swa_rope_theta, + scaling=_rope_scaling(swa_rope, swa_rope_type), + ) + + # Split layer ids. + full_ids = tuple(i for i, t in enumerate(layer_types) if t == "full_attention") + swa_ids = tuple(i for i, t in enumerate(layer_types) if t != "full_attention") + # Canonical SWA type is "sliding_attention" but treat any non-full as SWA. + if not full_ids: + # Degenerate (all SWA) — still build two groups but full empty. + full_ids = () + if not swa_ids: + swa_ids = () + + # Sliding window required when SWA layers exist. + sw = int(sliding_window) if sliding_window is not None else 0 + if swa_ids and sw <= 0: + # Default rescue: Laguna S 2.1 ships 512; keep a clear error if missing. + raise ValueError("Laguna config has sliding_attention layers but no sliding_window") + + # Architecture bookkeeping. + architectures = ( + top_architectures + or getattr(cfg, "architectures", None) + or ["LagunaForCausalLM"] + ) + + # MoE. + num_experts = int(getattr(cfg, "num_experts", 0) or 0) + num_experts_per_tok = int(getattr(cfg, "num_experts_per_tok", 0) or 0) + moe_intermediate_size = int(getattr(cfg, "moe_intermediate_size", 0) or 0) + shared_expert_intermediate_size = int(getattr(cfg, "shared_expert_intermediate_size", 0) or 0) + norm_topk_prob = bool(getattr(cfg, "norm_topk_prob", True)) + # Laguna names it moe_routed_scaling_factor (2.5 for S 2.1); it multiplies the + # routed-expert sum before the shared expert is added (modeling_laguna.py:255). + # The repo convention folds it into topk_weights instead, which is equivalent. + routed_scaling_factor = float(getattr(cfg, "moe_routed_scaling_factor", 1.0) or 1.0) + moe_enabled = num_experts > 0 and num_experts_per_tok > 0 + # decoder_sparse_step + mlp_only_layers gating is handled in model.py (layer 0 dense). + + # Quant: Laguna NVFP4 on routed experts (compressed-tensors W4A16). Only layers + # 1-39 are packed; 40-47 ship bf16 experts (the checkpoint's `ignore` list) and are + # quantized to NVFP4 at conversion so all 47 layers share one bank layout. + try: + is_nvfp4 = detect_compressed_tensors_nvfp4(top_cfg) + except ValueError: + # An unsupported 4-bit float scheme (e.g. MXFP4 group_size 32) must surface + # clearly instead of silently falling back to bf16. + raise + expert_quant = "nvfp4" if is_nvfp4 else "none" + + gating = getattr(cfg, "gating", "per-head") + gating_types = getattr(cfg, "gating_types", None) + if gating_types is not None: + gating_types = tuple(str(x) for x in gating_types) + + # Laguna S 2.1: only layer 0 is dense (mlp_only_layers [0]). + # Map to first_k_dense_replace so ModelConfig.num_moe_layers == 47. + mlp_only_layers = getattr(cfg, "mlp_only_layers", None) + if mlp_only_layers is not None and len(mlp_only_layers) > 0 and mlp_only_layers == [0]: + first_k_dense_replace = 1 + else: + first_k_dense_replace = int(getattr(cfg, "first_k_dense_replace", 0) or 0) + + return ModelConfig( + num_layers=num_layers, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + hidden_size=hidden_size, + vocab_size=vocab_size, + intermediate_size=intermediate_size, + rms_norm_eps=rms_norm_eps, + tie_word_embeddings=tie_word_embeddings, + rotary_config=full_rotary, + hidden_act=hidden_act, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + moe_intermediate_size=moe_intermediate_size, + shared_expert_intermediate_size=shared_expert_intermediate_size, + norm_topk_prob=norm_topk_prob, + routed_scaling_factor=routed_scaling_factor, + model_type=getattr(cfg, "model_type", "laguna"), + architectures=list(architectures), + moe_enabled=moe_enabled, + first_k_dense_replace=first_k_dense_replace, + expert_quant=expert_quant, + attn_quant="none", + dense_quant="none", + lm_head_quant="none", + use_qk_norm=True, + attn_sm_scale=None, # 1/sqrt(head_dim) default + attention_groups=( + FullAttentionGroupConfig( + name="full", + layer_ids=full_ids, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rotary_config=full_rotary, + ), + SWAAttentionGroupConfig( + name="swa", + layer_ids=swa_ids, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rotary_config=swa_rotary, + sliding_window=sw, + ), + ), + num_attention_heads_per_layer=num_attention_heads_per_layer, + laguna_args=LagunaArgs(gating=gating, gating_types=gating_types), + ) + + +__all__ = ["LagunaArgs", "parse_config"] diff --git a/python/freetoken/models/laguna/model.py b/python/freetoken/models/laguna/model.py new file mode 100644 index 000000000..e2c82b6c5 --- /dev/null +++ b/python/freetoken/models/laguna/model.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNormFused, VocabParallelEmbedding +from freetoken.models.blocks import BaseLLMModel, GatedMLP +from freetoken.utils import nvtx_annotate + +from .attention import LagunaAttention +from .moe import LagunaSparseMoeBlock + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +class LagunaDecoderLayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self._layer_id = layer_id + self.self_attn = LagunaAttention(config, layer_id) + # Determine dense vs MoE. Laguna S 2.1: mlp_only_layers [0] → layer 0 dense, + # rest MoE sparse. No decoder_sparse_step concept beyond that. + # We also respect generic first_k_dense_replace if set. + is_dense = False + # Check mlp_only_layers style: we didn't store it on config yet, but + # config.intermediate_size path for dense can be used: if layer 0 or + # first_k_dense_replace covers this layer. + # For Laguna we know mlp_only_layers=[0]; encode it via a model-config + # field would be ideal, but we can inline the check here via num_layers/type. + # Since we don't have the list on ModelConfig, treat layer 0 as dense when + # model_type is laguna and num_experts>0. + if config.model_type == "laguna": + mlp_only = {0} + if layer_id in mlp_only: + is_dense = True + if getattr(config, "first_k_dense_replace", 0) > layer_id: + is_dense = True + if not config.moe_enabled: + is_dense = True + if is_dense: + # Dense GatedMLP uses intermediate_size (12288). + self.mlp = GatedMLP(config) + else: + self.mlp = LagunaSparseMoeBlock(config, layer_id) + self.input_layernorm = RMSNormFused(size=config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNormFused(size=config.hidden_size, eps=config.rms_norm_eps) + + @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") + def forward(self, x: torch.Tensor, residual: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor]: + x, residual = self.input_layernorm.forward(x, residual) + x = self.self_attn.forward(x) + x, residual = self.post_attention_layernorm.forward(x, residual) + x = self.mlp.forward(x) + return x, residual + + +class LagunaModel(BaseOP): + def __init__(self, config: ModelConfig): + self.embed_tokens = VocabParallelEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + ) + self.layers = OPList([LagunaDecoderLayer(config, lid) for lid in range(config.num_layers)]) + self.norm = RMSNormFused(size=config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + x = self.embed_tokens.forward(input_ids) + residual: torch.Tensor | None = None + for layer in self.layers.op_list: + x, residual = layer.forward(x, residual) + return self.norm.forward(x, residual)[0] + + +class LagunaForCausalLM(BaseLLMModel): + def __init__(self, config: ModelConfig): + self.model = LagunaModel(config) + self.lm_head = ParallelLMHead( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + ) + super().__init__() + + def forward(self) -> torch.Tensor: + output = self.model.forward(get_global_ctx().batch.input_ids) + return self.lm_head.forward(output) + + +__all__ = ["LagunaForCausalLM", "LagunaModel", "LagunaDecoderLayer"] diff --git a/python/freetoken/models/laguna/moe.py b/python/freetoken/models/laguna/moe.py new file mode 100644 index 000000000..8a00178f6 --- /dev/null +++ b/python/freetoken/models/laguna/moe.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from freetoken.layers import BaseOP, make_moe_layer + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +class LagunaGate(BaseOP): + """Gate with bias for HF-faithful keys mlp.gate.weight + mlp.gate.e_score_correction_bias.""" + + def __init__(self, hidden_size: int, num_experts: int): + self.weight = torch.empty(num_experts, hidden_size) + self.e_score_correction_bias = torch.empty(num_experts) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.linear(x, self.weight) + + +class LagunaSparseMoeBlock(BaseOP): + """Laguna MoE: sigmoid-routed experts + shared expert. + + Uses the offload/resident seam ``make_moe_layer(...).routed_forward`` so + the existing MoE kernels are reused unchanged. Keys are + ``mlp.gate.weight`` + ``mlp.gate.e_score_correction_bias`` (HF-faithful). + """ + + def __init__(self, config: ModelConfig, layer_id: int): + self.layer_id = int(layer_id) + first_dense = int(getattr(config, "first_k_dense_replace", 0)) + offload_id = layer_id - first_dense + if offload_id < 0: + offload_id = 0 + self._offload_id = offload_id + hidden = int(config.hidden_size) + self.gate = LagunaGate(hidden, int(config.num_experts)) + self.top_k = int(config.num_experts_per_tok) + self.norm_topk_prob = bool(config.norm_topk_prob) + self.routed_scaling_factor = config.routed_scaling_factor + self.experts = make_moe_layer( + config, + layer_id=offload_id, + # ``_route`` below already renormalized, and ``routed_forward`` bypasses the + # layer's internal router entirely -- so the flag must not re-apply it. + renormalize=False, + ) + shared_inter = int(config.shared_expert_intermediate_size or config.moe_intermediate_size) + self.shared_expert = _SharedExpert(config.hidden_size, shared_inter, config.hidden_act, config.rms_norm_eps) + + def _route(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Sigmoid router (modeling_laguna.py LagunaTopKRouter.forward). + + The bias shifts SELECTION only; the returned weights are gathered from the + unbiased scores, then renormalized and scaled -- the DeepSeek/GLM ``noaux_tc`` + convention (cf. ``glm4_moe/moe.py``). + """ + logits = self.gate.forward(hidden).float() + scores = torch.sigmoid(logits) + scores_for_selection = scores + self.gate.e_score_correction_bias.float() + _, topk_ids = torch.topk(scores_for_selection, self.top_k, dim=-1) + topk_weights = scores.gather(-1, topk_ids) + if self.norm_topk_prob: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + topk_weights = topk_weights * self.routed_scaling_factor + return topk_weights.to(torch.float32).contiguous(), topk_ids.to(torch.int32).contiguous() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden = hidden_states.view(-1, hidden_states.shape[-1]) + topk_weights, topk_ids = self._route(hidden) + routed = self.experts.routed_forward(hidden, topk_weights, topk_ids.clone()) + shared = self.shared_expert.forward(hidden) + out = routed + shared + return out.view(hidden_states.shape) + + +class _SharedExpert(BaseOP): + def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str, eps: float): + from freetoken.layers import LinearColParallelMerged, LinearRowParallel, silu_and_mul, gelu_and_mul, gelu_tanh_and_mul + + self.gate_up_proj = LinearColParallelMerged( + hidden_size, [intermediate_size, intermediate_size], has_bias=False + ) + act_map = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} + self.act_fn = act_map.get(hidden_act, silu_and_mul) + self.down_proj = LinearRowParallel(intermediate_size, hidden_size, has_bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] + return self.down_proj.forward(self.act_fn(self.gate_up_proj.forward(x))) + + +__all__ = ["LagunaGate", "LagunaSparseMoeBlock"] diff --git a/python/freetoken/models/laguna/weight.py b/python/freetoken/models/laguna/weight.py new file mode 100644 index 000000000..aab41085d --- /dev/null +++ b/python/freetoken/models/laguna/weight.py @@ -0,0 +1,300 @@ +"""Laguna weight loading. + +Dense tensors (attention q/k/v/o/g, QK norms, router gate, shared experts, layer 0's +dense MLP, embeddings, lm_head, norms) are all bf16 and stream through ``iter_weights``. + +Routed experts go to the offload cache as native NVFP4 banks. Two checkpoint quirks +drive ``setup_offload_expert_banks``: + +* **compressed-tensors naming.** Laguna ships ``weight_packed`` / ``weight_scale`` / + ``weight_global_scale`` (W4A16, group 16, ``tensor_group``), not ModelOpt's + ``weight`` / ``weight_scale`` / ``weight_scale_2``. The shared loader keys off the + ModelOpt names, so the regex maps ``kind`` onto them. +* **The global scale is quant-side.** ``weight_global_scale`` is ``448*6/amax``; the + dequant kernel wants its reciprocal (``fp4 * block_scale * global``), so it is + inverted on the way into the bank. ModelOpt's ``weight_scale_2`` is already + dequant-side and is stored verbatim -- cf. ``models/loader.py`` and the comment in + ``models/nvfp4_banks.py``. + +* **Mixed precision.** Only layers 1-39 are quantized; layers 40-47 ship bf16 experts + (the checkpoint's ``ignore`` list). Those are quantized to real NVFP4 here so all 47 + MoE layers share one bank layout and one kernel path. +""" + +from __future__ import annotations + +import re +from typing import Iterator + +import safetensors +import torch +from freetoken.distributed import get_tp_info +from freetoken.models.loader import iter_weight_files, shard_tensor +from freetoken.models.nvfp4_banks import ( + Nvfp4ExpertSourceSpec, + load_nvfp4_expert_source_banks, + load_nvfp4_expert_source_banks_parallel, +) +from freetoken.utils import cached_load_hf_config +from tqdm import tqdm + +from .config import parse_config + +_EXPERT_RE = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.\d+\.") + +# compressed-tensors NVFP4 expert keys. The named groups layer/expert/proj/kind are a +# hard contract with models/nvfp4_banks.py, which indexes matches by group name; ``kind`` +# is mapped onto the ModelOpt names the shared loader dispatches on. +_EXPERT_KEY_RE = re.compile( + r"^model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight_packed|weight_scale|weight_global_scale)$" +) + +_E2M1_MAX = 6.0 # largest finite E2M1 magnitude +_FP8_E4M3_MAX = 448.0 # largest finite e4m3 magnitude (block-scale dtype) +# E2M1 magnitudes by code (codes 8-15 are these negated; bit 3 is the sign). +# Mirrors kernel/triton/nvfp4_dequant.py's LUT. +_E2M1_MAGNITUDES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _is_expert(name: str) -> bool: + return _EXPERT_RE.search(name) is not None + + +def iter_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Dense (non-expert) weights. Routed experts always go through the offload banks. + + ``include_moe_experts`` is accepted for interface parity but never yields experts: + Laguna's are NVFP4 (and partly bf16), which the generic stacked-bf16 path cannot + represent. ``setup_offload_expert_banks`` owns them. + """ + config = parse_config(cached_load_hf_config(model_path)) + tp = get_tp_info() + + shared_buf: dict[str, dict[str, torch.Tensor]] = {} + dense_buf: dict[str, dict[str, torch.Tensor]] = {} + + for file in tqdm(iter_weight_files(model_path), desc="Loading weights", disable=not tp.is_primary()): + with safetensors.safe_open(file, framework="pt", device=str(device)) as f: + for raw_name in f.keys(): + # e_score_correction_bias sits under mlp.experts.* in the checkpoint but + # belongs to the router, so it is a dense tensor despite the prefix. + is_bias = raw_name.endswith("e_score_correction_bias") + if _is_expert(raw_name) and not is_bias: + continue + if not include_non_moe: + continue + + # HF-faithful key: modeling_laguna.py's _checkpoint_conversion_mapping + # moves the bias onto the router. + name = raw_name.replace( + "mlp.experts.e_score_correction_bias", "mlp.gate.e_score_correction_bias" + ) + + raw = f.get_tensor(raw_name) + tensor = shard_tensor(name, raw, rank=tp.rank, world_size=tp.size, num_kv_heads=config.num_kv_heads) + del raw + + # Fuse gate+up on the output-row axis (gate first), matching + # LinearColParallelMerged's [gate, up] layout. + if "shared_expert.gate_proj" in name or "shared_expert.up_proj" in name: + prefix = name.split(".shared_expert.")[0] + ".shared_expert" + slot = "gate" if "gate_proj" in name else "up" + slots = shared_buf.setdefault(prefix, {}) + slots[slot] = tensor + if "gate" in slots and "up" in slots: + merged = torch.cat([slots["gate"], slots["up"]], dim=0) + del shared_buf[prefix] + yield f"{prefix}.gate_up_proj.weight", merged + continue + if name.endswith(("mlp.gate_proj.weight", "mlp.up_proj.weight")): + # Dense MLP (layer 0). "mlp.gate.weight" (the router) does not match. + mlp_prefix = name[: name.index(".mlp.") + 4] + slot = "gate" if "gate_proj" in name else "up" + slots = dense_buf.setdefault(mlp_prefix, {}) + slots[slot] = tensor + if "gate" in slots and "up" in slots: + merged = torch.cat([slots["gate"], slots["up"]], dim=0) + del dense_buf[mlp_prefix] + yield f"{mlp_prefix}.gate_up_proj.weight", merged + continue + + yield name, tensor + + assert not shared_buf, f"Laguna: incomplete shared_expert merges: {list(shared_buf)}" + assert not dense_buf, f"Laguna: incomplete dense mlp merges: {list(dense_buf)}" + + +def _quant_bf16_to_nvfp4( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """BF16 ``[O, K]`` -> NVFP4 ``(packed [O, K//2], block_scale [O, K//16], global [O])``. + + Exact inverse of ``kernel/triton/nvfp4_dequant.py``: ``w = E2M1[code] * block_scale + * global``. Used only for Laguna's bf16 expert layers (40-47) so every MoE layer + ends up in one bank layout. + + Two-level scaling, matching the packed layers byte-for-byte in convention: the + global is **per-tensor** (verified against the checkpoint -- for layer 1 expert 0 + ``448*6/amax`` reproduces the stored ``weight_global_scale`` exactly), and each + 16-wide group's scale is stored fp8-e4m3. A per-row global would inflate the scale + on low-magnitude rows and push the block scales off the e4m3 grid (~2.4% error vs + ~0.2%). The returned global is already **inverted** (dequant-side), like the packed + path, and is broadcast to one entry per output row because that is the bank layout. + """ + O, K = weight.shape + assert K % 16 == 0, f"NVFP4 group size 16 does not divide K={K}" + w = weight.float().reshape(O, K // 16, 16) + + # Per-TENSOR global: maps the tensor's amax onto E2M1_MAX * FP8_MAX, so the largest + # weight is the largest code under the largest block scale. + amax = w.abs().amax().clamp(min=1e-12).item() + g_quant = (_FP8_E4M3_MAX * _E2M1_MAX) / amax # quant-side, as stored in the checkpoint + + # Per-group block scale, quantized through e4m3 and read back: the codes must be + # chosen against the scale the kernel will actually see, not the ideal one. + group_amax = w.abs().amax(dim=-1) # [O, groups] + block_fp8 = (group_amax * g_quant / _E2M1_MAX).clamp(max=_FP8_E4M3_MAX).to(torch.float8_e4m3fn) + + # Effective per-element step, guarding all-zero groups (scale 0 -> codes 0). + step = (block_fp8.float() / g_quant).clamp(min=1e-30) + normalized = w / step[:, :, None] + + # Round the magnitude to the nearest E2M1 value: a non-uniform table, so this is a + # nearest-neighbour search over 8 magnitudes, not a linear round. + mags = torch.tensor(_E2M1_MAGNITUDES, dtype=torch.float32, device=weight.device) + codes = (normalized.abs().unsqueeze(-1) - mags).abs().argmin(dim=-1).to(torch.uint8) + codes |= (normalized < 0).to(torch.uint8) << 3 # bit 3 is the sign + + # Pack two codes per byte, low nibble = lower K index (nvfp4_dequant.py stores + # 2*byte_off from the low nibble). + pairs = codes.reshape(O, K // 2, 2) + packed = pairs[..., 0] | (pairs[..., 1] << 4) + + g_dequant = torch.full((O,), 1.0 / g_quant, dtype=torch.float16, device=weight.device) + # The bank dtype is fp16, so a very small amax pushes 1/g into the subnormal range + # and eventually to zero (which would silently blank the layer). Real Laguna expert + # tensors sit at amax ~0.09-0.24 -> 1/g ~3e-5..9e-5, subnormal but exact to ~0.02%. + # Fail loudly rather than emit zeros if a future checkpoint is far smaller. + assert g_dequant[0].item() > 0.0, ( + f"NVFP4 global scale underflowed fp16 (amax={amax:.3e}); the expert bank dtype " + "cannot represent this tensor's scale" + ) + return packed.contiguous(), block_fp8.contiguous(), g_dequant + + +def _synthesize_bf16_layer( + reader, moe_layer_ids: list[int], E: int, I: int, bank_layer: int, banks: dict +) -> None: + """Quantize one bf16 expert layer straight into its banks (layers 40-47). + + Called by the shared loader for bank layers absent from the checkpoint. Reads and + quantizes one expert at a time, so peak extra memory is a single projection. + """ + lid = moe_layer_ids[bank_layer] + for eid in range(E): + for proj, row_off in (("gate_proj", 0), ("up_proj", I)): + base = f"model.layers.{lid}.mlp.experts.{eid}.{proj}" + packed, scale, glob = _quant_bf16_to_nvfp4(reader.get_tensor(f"{base}.weight")) + banks["gate_up_packed"][eid, row_off : row_off + I] = packed + banks["gate_up_scale"][eid, row_off : row_off + I] = scale + banks["gate_up_global"][eid, row_off : row_off + I] = glob + base = f"model.layers.{lid}.mlp.experts.{eid}.down_proj" + packed, scale, glob = _quant_bf16_to_nvfp4(reader.get_tensor(f"{base}.weight")) + banks["down_packed"][eid] = packed + banks["down_scale"][eid] = scale + banks["down_global"][eid] = glob + + +def _laguna_nvfp4_spec(model_path: str, model_config) -> Nvfp4ExpertSourceSpec: + from freetoken.models.loader import ShardReader + + first_dense = int(model_config.first_k_dense_replace) + num_layers = int(model_config.num_layers) + moe_layer_ids = list(range(first_dense, num_layers)) + assert len(moe_layer_ids) == model_config.num_moe_layers, ( + f"Laguna MoE layer count {len(moe_layer_ids)} != {model_config.num_moe_layers}" + ) + E = int(model_config.num_experts) + I = int(model_config.moe_intermediate_size) + + # Opened lazily: only touched if the checkpoint actually has bf16 expert layers. + reader = ShardReader(model_path, torch.device("cpu")) + + return Nvfp4ExpertSourceSpec( + key_pattern=_EXPERT_KEY_RE, + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + # Experts exist for layers [first_k_dense_replace, num_layers); banks pack by + # MoE-layer index so the leading dense layer leaves no hole. + layer_to_bank=lambda layer, config: ( + None + if layer < config.first_k_dense_replace or layer >= config.num_layers + else layer - config.first_k_dense_replace + ), + desc="Laguna NVFP4 experts", + kind_map={ + "weight_packed": "weight", + "weight_scale": "weight_scale", + "weight_global_scale": "weight_scale_2", + }, + # compressed-tensors stores the quant-side scale; the kernel wants 1/g. + global_transform=lambda t: 1.0 / t.float(), + synthesize_layer=lambda bank_layer, banks: _synthesize_bf16_layer( + reader, moe_layer_ids, E, I, bank_layer, banks + ), + ) + + +def setup_offload_expert_banks( + model_path: str, + model_config, + *, + device: torch.device, + dtype: torch.dtype, + dummy: bool = False, + parallel: bool = False, + workers: int = 8, + chunk: int = 8 << 20, + decode_target: str = "gpu", + layer_sink=None, +): + """Build Laguna's native NVFP4 expert banks (all 47 MoE layers). + + Overrides the generic ``nvfp4`` provider because the checkpoint is mixed precision: + layers 1-39 ship packed experts, 40-47 ship bf16 ones that are quantized here so a + single bank layout and kernel path covers every layer. Bank layout, layer-completion + tracking, pin-after-fill and FTW streaming are the shared machinery. + + The native (6-bank) layout is returned unconditionally: Laguna's ``I`` is 1024, and + ``select_nvfp4_backend`` keeps narrow-MoE models on the Triton inline-dequant + kernels, which read exactly this layout. + """ + from freetoken.models.loader import drop_page_cache + from freetoken.moe.expert_banks import ExpertBanks + + if dummy: + from freetoken.models.weight import dummy_nvfp4_expert_sources + + return ExpertBanks("nvfp4", dummy_nvfp4_expert_sources(model_config), streamed=False) + + spec = _laguna_nvfp4_spec(model_path, model_config) + primary = get_tp_info().is_primary() + loader = ( + load_nvfp4_expert_source_banks_parallel if parallel else load_nvfp4_expert_source_banks + ) + kwargs = {"workers": workers, "chunk": chunk} if parallel else {} + sources = loader( + model_path, model_config, spec, + drop_page_cache=drop_page_cache, primary=primary, layer_sink=layer_sink, **kwargs, + ) + return ExpertBanks("nvfp4", sources, streamed=layer_sink is not None) + + +__all__ = ["iter_weights", "setup_offload_expert_banks"] diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 0e3ab6a51..cf51a0c75 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -22,6 +22,31 @@ class Nvfp4ExpertSourceSpec: proj_to_role: dict[str, str] layer_to_bank: LayerToBank desc: str + # ``kind`` group -> ModelOpt kind ("weight" / "weight_scale" / "weight_scale_2"). + # compressed-tensors checkpoints (Laguna) ship weight_packed / weight_scale / + # weight_global_scale; normalizing here keeps the dispatch below in ModelOpt terms. + # None -> the pattern already yields ModelOpt names. + kind_map: dict[str, str] | None = None + # Applied to each global-scale tensor as it is read. compressed-tensors stores the + # QUANT-side scale (``448*6/amax``) but the dequant kernel wants its reciprocal + # (``fp4 * block_scale * global``); ModelOpt's ``weight_scale_2`` is already + # dequant-side, so the default is identity. Cf. models/loader.py. + global_transform: Callable[[torch.Tensor], torch.Tensor] | None = None + # Mixed-precision checkpoints: bank layers whose experts are absent from the + # checkpoint entirely (Laguna quantizes only layers 1-39; 40-47 ship bf16). Called + # as ``fill(bank_layer, {bank_name: tensor})`` for each missing layer and must write + # all E experts of all 6 banks; each counts as ``E*6`` placements so the completion + # tracker fires it exactly like a natively-packed layer. + synthesize_layer: Callable[[int, dict[str, torch.Tensor]], None] | None = None + + +def _normalize_kind(spec: Nvfp4ExpertSourceSpec, kind: str) -> str: + if spec.kind_map is None: + return kind + try: + return spec.kind_map[kind] + except KeyError: + raise ValueError(f"{spec.desc}: unmapped NVFP4 expert tensor kind {kind!r}") from None def _num_moe_layers(config) -> int: @@ -62,6 +87,44 @@ def _alloc_nvfp4_host_banks(num_layers: int, E: int, H: int, I: int): }, num_layers) +def _read_global(spec: Nvfp4ExpertSourceSpec, handle, name: str) -> torch.Tensor: + """Read one global-scale tensor, applying the spec's dequant-side transform.""" + tensor = handle.get_tensor(name) + if spec.global_transform is not None: + tensor = spec.global_transform(tensor) + return tensor.to(torch.float16) + + +def _missing_bank_layers( + spec: Nvfp4ExpertSourceSpec, num_layers: int, seen: set[int] +) -> list[int]: + """Bank layers with no packed experts in the checkpoint (mixed-precision models).""" + missing = sorted(set(range(num_layers)) - seen) + if missing and spec.synthesize_layer is None: + raise ValueError( + f"{spec.desc}: no NVFP4 expert tensors for bank layers {missing}; " + "the checkpoint is mixed precision but the spec has no synthesize_layer" + ) + return missing + + +def _synthesize_missing( + spec: Nvfp4ExpertSourceSpec, + missing_layers: list[int], + banks: dict[str, list[torch.Tensor]], + tracker, + per_layer: int, +) -> int: + """Fill layers absent from the checkpoint, noting them so the sink fires normally.""" + placed = 0 + for bank_layer in missing_layers: + spec.synthesize_layer(bank_layer, {name: per[bank_layer] for name, per in banks.items()}) + for _ in range(per_layer): + tracker.note(bank_layer) + placed += per_layer + return placed + + def load_nvfp4_expert_source_banks( model_path: str, config, @@ -101,6 +164,7 @@ def load_nvfp4_expert_source_banks( weight_shards: dict[str, list[tuple[str, re.Match[str], int]]] = collections.defaultdict(list) global_shards: dict[str, list[tuple[str, re.Match[str], int]]] = collections.defaultdict(list) + seen_layers: set[int] = set() for name, shard in weight_map.items(): match = spec.key_pattern.match(name) if match is None: @@ -112,7 +176,8 @@ def load_nvfp4_expert_source_banks( proj = match.group("proj") if proj not in spec.proj_to_role: raise ValueError(f"{spec.desc}: unknown NVFP4 expert projection {proj!r}") - kind = match.group("kind") + seen_layers.add(bank_layer) + kind = _normalize_kind(spec, match.group("kind")) if kind == "weight_scale_2": global_shards[shard].append((name, match, bank_layer)) elif kind in {"weight", "weight_scale"}: @@ -120,6 +185,8 @@ def load_nvfp4_expert_source_banks( else: raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") + missing_layers = _missing_bank_layers(spec, num_layers, seen_layers) + globals_map: dict[tuple[int, int, str], torch.Tensor] = {} for shard in sorted(global_shards): path = os.path.join(folder, shard) @@ -130,7 +197,7 @@ def load_nvfp4_expert_source_banks( int(match.group("expert")), match.group("proj"), ) - globals_map[key] = f.get_tensor(name).to(torch.float16) + globals_map[key] = _read_global(spec, f, name) drop_page_cache(path) _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill @@ -140,6 +207,14 @@ def load_nvfp4_expert_source_banks( down_packed = [b.tensor for b in _hb["down_packed"]] down_scale = [b.tensor for b in _hb["down_scale"]] down_global = [b.tensor for b in _hb["down_global"]] + _bank_tensors = { + "gate_up_packed": gate_up_packed, + "gate_up_scale": gate_up_scale, + "gate_up_global": gate_up_global, + "down_packed": down_packed, + "down_scale": down_scale, + "down_global": down_global, + } from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline @@ -154,7 +229,7 @@ def _load(sink) -> int: expert = int(match.group("expert")) proj = match.group("proj") role = spec.proj_to_role[proj] - kind = match.group("kind") + kind = _normalize_kind(spec, match.group("kind")) tensor = f.get_tensor(name) if kind == "weight": if role == "gate": @@ -181,6 +256,10 @@ def _load(sink) -> int: tracker.note(bank_layer_id) placed += 1 drop_page_cache(path) + if missing_layers: + placed += _synthesize_missing( + spec, missing_layers, _bank_tensors, tracker, E * 6 + ) return placed if layer_sink is not None: @@ -229,6 +308,7 @@ def load_nvfp4_expert_source_banks_parallel( weight_info: dict[str, tuple[re.Match[str], int]] = {} # name -> (match, bank_layer) global_names_by_shard: dict[str, list[str]] = collections.defaultdict(list) + seen_layers: set[int] = set() for name, shard in weight_map.items(): match = spec.key_pattern.match(name) if match is None: @@ -236,7 +316,8 @@ def load_nvfp4_expert_source_banks_parallel( bank_layer = _bank_layer(spec, int(match.group("layer")), config) if bank_layer is None: continue - kind = match.group("kind") + seen_layers.add(bank_layer) + kind = _normalize_kind(spec, match.group("kind")) if kind == "weight_scale_2": global_names_by_shard[shard].append(name) elif kind in {"weight", "weight_scale"}: @@ -244,6 +325,8 @@ def load_nvfp4_expert_source_banks_parallel( else: raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") + missing_layers = _missing_bank_layers(spec, num_layers, seen_layers) + # Pass 1: tiny per-tensor global scales (serial; data is scalar-per-expert). globals_map: dict[tuple[int, int, str], torch.Tensor] = {} for shard in sorted(global_names_by_shard): @@ -253,7 +336,7 @@ def load_nvfp4_expert_source_banks_parallel( for name in global_names_by_shard[shard]: m = spec.key_pattern.match(name) globals_map[(int(m.group("layer")), int(m.group("expert")), m.group("proj"))] = ( - f.get_tensor(name).to(torch.float16) + _read_global(spec, f, name) ) drop_page_cache(path) @@ -264,6 +347,14 @@ def load_nvfp4_expert_source_banks_parallel( down_packed = [b.tensor for b in _hb["down_packed"]] down_scale = [b.tensor for b in _hb["down_scale"]] down_global = [b.tensor for b in _hb["down_global"]] + _bank_tensors = { + "gate_up_packed": gate_up_packed, + "gate_up_scale": gate_up_scale, + "gate_up_global": gate_up_global, + "down_packed": down_packed, + "down_scale": down_scale, + "down_global": down_global, + } from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline @@ -279,7 +370,7 @@ def _load(sink) -> int: expert = int(match.group("expert")) proj = match.group("proj") role = spec.proj_to_role[proj] - kind = match.group("kind") + kind = _normalize_kind(spec, match.group("kind")) if kind == "weight": if role == "gate": gate_up_packed[bank_layer_id][expert, :I] = tensor @@ -300,6 +391,10 @@ def _load(sink) -> int: down_global[bank_layer_id][expert] = g tracker.note(bank_layer_id) placed += 1 + if missing_layers: + placed += _synthesize_missing( + spec, missing_layers, _bank_tensors, tracker, E * 6 + ) return placed if layer_sink is not None: diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index b94d8291b..5806268bb 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -130,6 +130,18 @@ class ModelSpec: "freetoken.models.glm_moe_dsa", "GlmMoeDsaForCausalLM", ), + # Laguna S/XS 2.1 (model_type laguna): 48-layer 1:3 full/SWA hybrid (window 512), + # per-head softplus gating, per-layer head count (48 vs 72), yarn full rope + + # default SWA rope, sigmoid router with bias, shared expert. NVFP4 routed experts + # (compressed-tensors W4A16) served from offload cache. + "LagunaForCausalLM": ModelSpec( + "freetoken.models.laguna", + "LagunaForCausalLM", + ), + "LagunaForConditionalGeneration": ModelSpec( + "freetoken.models.laguna", + "LagunaForCausalLM", + ), } diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 6696f65dd..110f1ba4f 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -163,6 +163,11 @@ def _infer_tool_call_parser(model_path: str) -> str: return "deepseekv32" if "glm" in marker: return "glm47" + # Laguna's tool envelope is GLM-4.7's byte-for-byte (name + + # / pairs), just without the newlines between tags, + # which Glm47Detector already tolerates. + if "laguna" in marker: + return "glm47" if "mistral" in marker: return "mistral" return "llama3" @@ -196,6 +201,12 @@ def _infer_reasoning_parser(model_path: str) -> str | None: return "qwen3" if "glm" in marker: return "glm" + # Laguna wraps its chain-of-thought in / (chat_template.jinja + # "laguna_glm_thinking_v8"), so the generic think-tag parser applies. Its + # template PRE-OPENS when enable_thinking (the default), leaving the + # model to emit only the closing tag -- see _make_reasoning_parser. + if "laguna" in marker: + return "laguna" # M3 first ("minimax" is a substring): tags + 3 thinking gears, # not M2's always-on implicit . if "minimax_m3" in marker or "minimax-m3" in marker or "minimaxm3" in marker: @@ -455,13 +466,13 @@ def _infer_reasoning_parser(model_path: str) -> str | None: type=str, default="auto", choices=[ - "auto", "off", "deepseekv32", "gpt_oss", "qwen3", "glm", + "auto", "off", "deepseekv32", "gpt_oss", "qwen3", "glm", "laguna", "minimax", "minimax_m3", "muse_glimmer", "gemma4", ], help=( "Reasoning parser that splits chain-of-thought into reasoning_content " "for OpenAI responses. 'auto' selects per model family (gpt-oss Harmony, " - " for qwen3/glm/minimax, for minimax-m3, ATEM to=self " + " for qwen3/glm/laguna/minimax, for minimax-m3, ATEM to=self " "channels for muse-glimmer, gemma thought, dsv4); 'off' disables it." ), ) diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a..299cb23ee 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -351,6 +351,12 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: # GLM's template honors enable_thinking (default on) even with tools; the # generic fallback would force thinking and mislabel disabled output as reasoning. force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False + elif parser_name == "laguna": + # Laguna's template pre-opens when enable_thinking (its default is true, + # and generation_config sets enable_thinking in default_chat_template_kwargs), so + # the model emits only the closing . With thinking off the template + # pre-closes instead, and the visible answer must stay content. + force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False elif parser_name == "gemma4": # Gemma4 defaults thinking off even when tools are present: its template injects an # empty thought channel before generation. Do not let Codex tool definitions make all diff --git a/python/freetoken/server/reasoning_parser.py b/python/freetoken/server/reasoning_parser.py index 6788675ca..ee167cc1e 100644 --- a/python/freetoken/server/reasoning_parser.py +++ b/python/freetoken/server/reasoning_parser.py @@ -878,6 +878,7 @@ class ReasoningParser: "gpt_oss": GptOssHarmonyReasoningParser, "qwen3": ThinkReasoningParser, "glm": ThinkReasoningParser, + "laguna": ThinkReasoningParser, "minimax": ThinkReasoningParser, "minimax_m3": MiniMaxM3ReasoningParser, "muse_glimmer": MuseGlimmerReasoningParser, From 52dd1601756980fef8cea25eaa3cd4619c929830 Mon Sep 17 00:00:00 2001 From: Ahmed Hassanin Date: Sun, 30 Aug 2026 16:05:26 -0400 Subject: [PATCH 2/2] fix(server): serve from a source checkout on Windows Five separate POSIX-only assumptions each stop `python -m freetoken` before it can answer a request on Windows. None affect POSIX behaviour. * ZMQ IPC. libzmq is built without the `ipc://` transport on Windows, so all five inter-process sockets fail to bind with "Protocol not supported". The addresses now come from one helper on SchedulerConfig that keeps `ipc://` wherever the ZMQ build supports it and otherwise uses a contiguous localhost TCP block. The base port is reserved once in the parent and travels to the workers in the config they are already handed, since separate processes cannot negotiate it at runtime. * Event loop. pyzmq's asyncio sockets need `add_reader`, which the Windows Proactor loop does not implement. The frontend's ZMQ listener died at startup and every request then hung until it timed out while the scheduler logged the prefill as normal. Setting the loop policy alone is not enough because uvicorn's loop factory hardcodes ProactorEventLoop on win32, so the server now runs on a selector loop we create and pass in via `loop="none"`. * `os.posix_fadvise`, `mmap.PROT_READ` and `mmap.madvise` are absent on Windows; the first raised AttributeError past the OSError guard, and the second is spelled `access=ACCESS_READ` there. The parallel O_DIRECT reader was already gated behind a platform check, so page-cache dropping simply becomes a no-op and host memory pressure during a large conversion is correspondingly higher. Checked by converting and serving a 93 GiB NVFP4 checkpoint from a source checkout: prompts answer in ~2s, thinking and tool calls parse, and a runtime cache rebuild followed by further generation works. --- python/freetoken/checkpoint/ftw.py | 7 +++- python/freetoken/models/loader.py | 9 +++- python/freetoken/moe/host_banks.py | 3 +- python/freetoken/scheduler/config.py | 60 +++++++++++++++++++++++++-- python/freetoken/server/api_server.py | 17 +++++++- python/freetoken/server/args.py | 4 +- 6 files changed, 91 insertions(+), 9 deletions(-) diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index a365e3d5b..abdd393fb 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/python/freetoken/checkpoint/ftw.py @@ -261,7 +261,12 @@ def _map(self, file: str) -> memoryview: if entry is None: fd = os.open(os.path.join(self.dir, file), os.O_RDONLY) try: - m = mmap.mmap(fd, 0, prot=mmap.PROT_READ) + # ``prot`` is POSIX-only; Windows spells a read-only mapping + # ``access=ACCESS_READ`` (and rejects ``prot`` outright). + if hasattr(mmap, "PROT_READ"): + m = mmap.mmap(fd, 0, prot=mmap.PROT_READ) + else: + m = mmap.mmap(fd, 0, access=mmap.ACCESS_READ) finally: os.close(fd) # the mapping keeps its own reference to the file try: diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 494193641..a9888331b 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -54,7 +54,14 @@ def iter_weight_files(model_path: str) -> list[str]: def drop_page_cache(path: str) -> None: - """drop a file's page cache: banks + full checkpoint cache don't both fit in host RAM (OOM).""" + """drop a file's page cache: banks + full checkpoint cache don't both fit in host RAM (OOM). + + POSIX-only (``posix_fadvise``); a no-op on Windows, which exposes no equivalent hint + for another process's cached pages. The cache there is left to the OS, so host RAM + pressure during a large conversion is higher. + """ + if not hasattr(os, "posix_fadvise"): + return try: fd = os.open(path, os.O_RDONLY) try: diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 436f52d2d..888b12a84 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -147,7 +147,8 @@ def release(self) -> None: For buffers that are done being read (the converter). No-op for born-pinned banks: registered pages cannot be dropped.""" if self._pinned: return - self._buf.madvise(mmap.MADV_DONTNEED) + if hasattr(self._buf, "madvise"): + self._buf.madvise(mmap.MADV_DONTNEED) def lock(self) -> None: """mlock the (now-filled) buffer: resident without CUDA pin quota, but no device address -- only the CPU executor can serve a locked layer. diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index b2bbe0aef..d34f9b608 100644 --- a/python/freetoken/scheduler/config.py +++ b/python/freetoken/scheduler/config.py @@ -1,5 +1,6 @@ from __future__ import annotations +import socket from dataclasses import dataclass, field from freetoken.engine import EngineConfig @@ -11,6 +12,43 @@ def _get_pid_suffix() -> str: return f".pid={os.getpid()}" +def _ipc_supported() -> bool: + """Whether ZMQ can bind ``ipc://`` transports. + + libzmq is built without the IPC transport on Windows (there is no AF_UNIX + equivalent it can use), so every ``ipc://`` bind raises "Protocol not + supported". Probing the build is cheaper and more honest than sniffing the + platform. + """ + try: + import zmq + except ImportError: # zmq missing entirely -> nothing to bind anyway + return False + return bool(zmq.has("ipc")) + + +def _reserve_port_base(count: int = 8) -> int: + """Reserve a contiguous, currently-free localhost port block for the TCP fallback. + + The workers are separate processes, so they cannot negotiate ports at runtime: + the parent picks a base once and every address derives from it deterministically, + and the resolved base travels to the children inside the (pickled) config. Bind + to port 0 to let the OS pick a free port, then step past the block we intend to + use so a second server on the same host lands elsewhere. + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", 0)) + base = probe.getsockname()[1] + # Keep the block inside the ephemeral range; wrap rather than overflow. + if base + count > 65535: + base -= count + return base + + +def _default_ipc_base_port() -> int: + return _reserve_port_base() if not _ipc_supported() else 0 + + @dataclass(frozen=True) class SchedulerConfig(EngineConfig): max_extend_tokens: int = 8192 @@ -21,18 +59,34 @@ class SchedulerConfig(EngineConfig): # networking config _unique_suffix: str = field(default_factory=_get_pid_suffix) + # Base of the localhost port block used when ipc:// is unavailable (Windows). + # 0 means "ipc:// works, no ports needed". Resolved once in the parent so the + # workers inherit the same addresses through the config they are handed. + _ipc_base_port: int = field(default_factory=_default_ipc_base_port) + + def _socket_addr(self, index: int) -> str: + """One inter-process socket address. ``index`` must be unique per socket. + + ``ipc://`` where the ZMQ build supports it (the POSIX path, unchanged); + otherwise a fixed offset into the reserved localhost port block. Loopback + TCP is visible to other local processes, unlike a filesystem socket, so + the block is bound to 127.0.0.1 only. + """ + if self._ipc_base_port == 0: + return f"ipc:///tmp/freetoken_{index}{self._unique_suffix}" + return f"tcp://127.0.0.1:{self._ipc_base_port + index}" @property def zmq_backend_addr(self) -> str: - return "ipc:///tmp/freetoken_0" + self._unique_suffix + return self._socket_addr(0) @property def zmq_detokenizer_addr(self) -> str: - return "ipc:///tmp/freetoken_1" + self._unique_suffix + return self._socket_addr(1) @property def zmq_scheduler_broadcast_addr(self) -> str: - return "ipc:///tmp/freetoken_2" + self._unique_suffix + return self._socket_addr(2) @property def max_forward_len(self) -> int: diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..2c6f6db63 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -930,6 +930,15 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s global _GLOBAL_STATE, _MODEL_SAMPLING + # pyzmq's asyncio sockets need ``add_reader``, which Windows' Proactor loop does not + # implement (it works only if tornado>=6.1 is importable). Without a selector loop the + # frontend's ZMQ listener dies at startup and every request hangs until it times out, + # while the scheduler happily logs the prefill. Setting the policy is not enough: + # uvicorn's own loop factory hardcodes ProactorEventLoop on win32, so the loop is also + # created explicitly below and handed to uvicorn via ``loop="none"``. No-op on POSIX. + if hasattr(asyncio, "WindowsSelectorEventLoopPolicy"): + asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) + if config.sampling_defaults == "model" and not config.use_dummy_weight: _MODEL_SAMPLING = load_generation_sampling(config.model_path) # Always surface the effective default sampling (model-recommended where available, @@ -1037,4 +1046,10 @@ def _on_meta(meta: dict) -> None: _serve_and_run_shell(host, port) return # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group. - uvicorn.run(app, host=host, port=port) + if hasattr(asyncio, "WindowsSelectorEventLoopPolicy"): + # ``loop="none"`` keeps uvicorn off its own factory (which would hand us a + # Proactor loop on win32) and runs the server on the selector loop we own. + server = uvicorn.Server(uvicorn.Config(app, host=host, port=port, loop="none")) + asyncio.run(server.serve()) + else: + uvicorn.run(app, host=host, port=port) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 110f1ba4f..4b578c5bc 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -50,13 +50,13 @@ def share_tokenizer(self) -> bool: @property def zmq_frontend_addr(self) -> str: - return "ipc:///tmp/freetoken_3" + self._unique_suffix + return self._socket_addr(3) @property def zmq_tokenizer_addr(self) -> str: if self.share_tokenizer: return self.zmq_detokenizer_addr - result = "ipc:///tmp/freetoken_4" + self._unique_suffix + result = self._socket_addr(4) assert result != self.zmq_detokenizer_addr return result