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/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/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/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/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 6696f65dd..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 @@ -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,