diff --git a/docs/cli.md b/docs/cli.md index feda4f7a34..68822084ee 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -90,7 +90,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does, and [perf | `--moe-cache-split` | equal | How the pools share the slots: `equal` (in proportion to their layer counts), or `profile:`: every pool keeps `top_k` slots and the rest go to the most used experts in the profile's per-layer expert counts. With `offload`, a pool smaller than the expert count lowers `--max-running-requests` to its slots / top_k | | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts | | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | -| `--moe-cpu-layers` | all on GPU | With `offload`/`hybrid`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for hosts that cap CUDA pinned memory (Windows/WSL): it pins layers from the middle out until CUDA refuses one and OS-locks the rest (`FREETOKEN_PIN_BUDGET_GB` and FTW checkpoints lock from a size estimate instead), see [models.md](models.md#hosts-that-cap-pinned-memory-wsl); every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4, GGUF on x86 CPUs) | +| `--moe-cpu-layers` | all on GPU | With `offload`/`hybrid`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for hosts that cap CUDA pinned memory (Windows/WSL): it pins layers from the middle out until CUDA refuses one and OS-locks the rest (`FREETOKEN_PIN_BUDGET_GB` and FTW checkpoints lock from a size estimate instead), see [models.md](models.md#hosts-that-cap-pinned-memory-wsl); every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4, GGUF on x86 CPUs, except gpt-oss GGUF files whose per-expert biases decode on the GPU only) | | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU | | `--moe-locked-fetch` | auto | With `hybrid` and locked layers: share of a locked layer's misses fetched per step into its GPU slots through pinned staging. Auto takes p / (2p + c) from the `ft bench bw` row of the layer's kind (at most 1 per step without a row); 0 = locked layers decode on CPU only | | `--moe-locked-max-fetch` | 8 | Most experts a locked layer fetches per step; sizes the staging (that many rows of the widest expert, per bank) | diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index bc7b43707a..c5eec307d9 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -236,6 +236,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int top_k=4, moe_intermediate_size=2880, expert_formats=("mxfp4_triton",), + arch_aliases=("GptOssGGUFForCausalLM",), # the GGUF row builds the same classes ), AotModel( name="openai/gpt-oss-20b", @@ -245,6 +246,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int top_k=4, moe_intermediate_size=2880, expert_formats=("mxfp4_triton",), + arch_aliases=("GptOssGGUFForCausalLM",), ), AotModel( name="zai-org/GLM-4.7", diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py index b17a91359b..999e5c7edc 100644 --- a/python/freetoken/layers/quantization/moe/base.py +++ b/python/freetoken/layers/quantization/moe/base.py @@ -159,12 +159,12 @@ def slot_limit(self, cfg: MoEConfig) -> int | None: """Most GPU cache slots the kernel can address for ``cfg``; None for no limit.""" return self.max_slots - def _common_reject(self, cfg: MoEConfig, *, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool) -> str | None: + def _common_reject(self, cfg: MoEConfig, *, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool, bias_ok: bool = False) -> str | None: if not tp_ok and cfg.tp_size > 1: return "TP > 1 is not supported for this expert format" if not cpu_ok and cfg.decode_target != "gpu": return "has no CPU executor format; decode must run on the GPU" - if cfg.has_bias: + if cfg.has_bias and not bias_ok: return "kernel has no bias epilogue" if plain_silu_only and not cfg.plain_silu: return f"kernel is plain-silu only, experts use ({cfg.activation}, alpha={cfg.alpha}, beta={cfg.beta}, limit={cfg.limit})" diff --git a/python/freetoken/layers/quantization/moe/gguf.py b/python/freetoken/layers/quantization/moe/gguf.py index 01e754d644..4172ea4c74 100644 --- a/python/freetoken/layers/quantization/moe/gguf.py +++ b/python/freetoken/layers/quantization/moe/gguf.py @@ -36,8 +36,9 @@ def expert_types(cfg: MoEConfig) -> tuple[int, int]: def _bank_spec(ggml_type: int, rows: int, k: int) -> BankSpec: + # F16 / F32 expert banks are loaded and run as bf16, the fused MoE's dtype dtype = float_dtype(ggml_type) - return BankSpec((rows, k), dtype) if dtype is not None else BankSpec((rows, row_bytes(ggml_type, k)), torch.uint8) + return BankSpec((rows, k), torch.bfloat16) if dtype is not None else BankSpec((rows, row_bytes(ggml_type, k)), torch.uint8) _DECODE_CONFIG = {"BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "num_warps": 8} @@ -52,7 +53,7 @@ def float_mul_mat_id(w: torch.Tensor, x: torch.Tensor, ids: torch.Tensor, topk_w t, top_k = ids.shape per_route = x.dim() == 3 - # f16 experts take f16 activations, as llama.cpp's f16 path does + # the float banks run bf16, so the activations cast to the banks' dtype a = x.reshape(-1, x.shape[-1]).to(w.dtype).contiguous() out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device) if config is None: @@ -75,7 +76,9 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: reason = gated_epilogue_reason(cfg) if reason: return f"float GGUF experts run on the bf16 fused MoE: {reason}" - return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False) + if cfg.has_bias and cfg.decode_target != "gpu": + return "gpt-oss per-expert biases have no CPU executor format; decode must run on the GPU" + return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False, bias_ok=True) def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: gate_up, down = expert_types(cfg) @@ -84,6 +87,11 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: if cfg.scheme.weight.scale is not None: # each expert's NVFP4 weight_scale_2 as (gate, up, down, 1): a 16-byte row, which the slot cache's copies need banks["global"] = BankSpec((4,), torch.float32) + if cfg.has_bias: + # gpt-oss per-expert biases: bank rows so they follow the expert through the slot cache (decode) + # and the expert-ordered prefill / resident views, as llama.cpp's build_moe_ffn applies them + banks["gate_up_bias"] = BankSpec((2 * i,), torch.bfloat16) + banks["down_bias"] = BankSpec((h,), torch.bfloat16) return banks def pack(self, pieces, cfg: MoEConfig, out): @@ -117,13 +125,18 @@ def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor return ggml.mul_mat_id(w, ggml_type, a, ids) return float_mul_mat_id(w, a, ids, topk_weights, config) - # llama.cpp's build_moe_ffn: each NVFP4 GEMM result times the expert's weight_scale_2, gate and up before the glu, down before the routing weight + # llama.cpp's build_moe_ffn: each NVFP4 GEMM result times the expert's weight_scale_2 (before + # its bias), gate and up before the glu, down before the routing weight scales = view.tensors.get("global") gate_up_scales = scales if scales is not None and gate_up_type == NVFP4 else None down_scales = scales if scales is not None and down_type == NVFP4 else None + # the per-expert biases (gpt-oss) add around the gated activation, as in llama.cpp's build_moe_ffn; + # the view indexes them by slot id on the decode path and by expert id on prefill / resident + gate_up_bias = view.tensors.get("gate_up_bias") + down_bias = view.tensors.get("down_bias") glu = fused_glu(layer) small = x.shape[0] <= min(ggml.mmvq_max_tokens(t) for t in (gate_up_type, down_type) if is_quantized(t)) - if glu is not None and small and is_quantized(gate_up_type): + if glu is not None and small and is_quantized(gate_up_type) and gate_up_bias is None: gu = (None, None) if gate_up_scales is None else (gate_up_scales[:, 0], gate_up_scales[:, 1]) inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu, *gu) else: @@ -131,6 +144,8 @@ def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor t, k, two_i = h.shape if gate_up_scales is not None: h.view(t, k, 2, two_i // 2).mul_(gate_up_scales[ids, :2].unsqueeze(-1)) + if gate_up_bias is not None: + h.add_(gate_up_bias[ids]) inter = torch.empty(t, k, two_i // 2, dtype=h.dtype, device=h.device) gated_act_and_mul(layer.activation, h.view(-1, two_i), inter.view(-1, two_i // 2), alpha=float(layer.alpha), limit=limit_or_inf(layer)) del h # free the gate_up output before the down GEMM allocates its own @@ -140,18 +155,24 @@ def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor y = mul_mat_id(view.tensors["down"], down_type, inter) if down_scales is not None: topk_weights = topk_weights * down_scales[ids, 2] + if down_bias is not None: + y.add_(down_bias[ids]) return ggml.moe_weighted_sum(y, topk_weights) @staticmethod def _apply_float(layer, x, topk_weights, topk_ids, view: ExpertView, ggml_type: int, *, is_prefill: bool): - """Float experts: their banks in their own dtype, through the bf16 fused MoE (the unquantized kernel's path).""" + """Float experts: their bf16 banks through the bf16 fused MoE (the unquantized kernel's path).""" from freetoken.moe.fused import fused_experts_decode_impl, fused_experts_impl gate_up, down = view.tensors["gate_up"], view.tensors["down"] if gate_up.dtype != x.dtype: raise NotImplementedError(f"GGUF {type_name(ggml_type)} experts need {gate_up.dtype} activations, got {x.dtype}") impl = fused_experts_impl if is_prefill or is_resident(layer) else fused_experts_decode_impl - return impl(x, gate_up, down, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, float(layer.alpha), limit_or_inf(layer)) + return impl( + x, gate_up, down, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, + float(layer.alpha), limit_or_inf(layer), + w1_bias=view.tensors.get("gate_up_bias"), w2_bias=view.tensors.get("down_bias"), + ) @register_method(QuantKind.GGUF, LayerKind.MOE) @@ -159,9 +180,12 @@ class GGUFMoEMethod(MoEMethod): candidates = (GgmlMoEKernel,) def cpu_spec(self) -> tuple[str | None, dict[str, int]]: + if self.cfg.has_bias: + raise NotImplementedError("GGUF experts with per-expert biases (gpt-oss) do not run on the cpu executor") gate_up, down = expert_types(self.cfg) - # BF16 experts keep the executor's bf16 kernels, which use AVX512-BF16 where the CPU has it; ggml-cpu is built for AVX2 only - if float_dtype(gate_up) is torch.bfloat16 and float_dtype(down) is torch.bfloat16: + # every float type (BF16 / F16 / F32) loads and runs as bf16 banks: the executor's bf16 + # kernels, which use AVX512-BF16 where the CPU has it (ggml-cpu is built for AVX2 only) + if float_dtype(gate_up) is not None and float_dtype(down) is not None: return "bf16", {} # an NVFP4 layer of a file with weight scales reads the layout's global bank too scaled = self.cfg.scheme.weight.scale is not None and NVFP4 in (gate_up, down) diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 1f234b1c4a..845ffedfe9 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -23,6 +23,7 @@ "flint": "FlintGGUFForCausalLM", "deepseek4": "DeepseekV4GGUFForCausalLM", "gemma4": "Gemma4GGUFForCausalLM", + "gpt-oss": "GptOssGGUFForCausalLM", "glm4moe": "Glm4MoeGGUFForCausalLM", # glm5-next: the llama.cpp spelling; glm5next: unsloth's first release "glm5-next": "Glm5NextGGUFForCausalLM", diff --git a/python/freetoken/models/gguf/experts.py b/python/freetoken/models/gguf/experts.py index 0efb5982fb..6f1b8c6185 100644 --- a/python/freetoken/models/gguf/experts.py +++ b/python/freetoken/models/gguf/experts.py @@ -1,8 +1,9 @@ """Routed-expert banks straight from a GGUF's ``blk.N.ffn_{gate,up,down}_exps`` tensors, in their packed ggml rows: pinned host banks for the offload cache, or device banks for resident experts. -Each layer keeps its own ggml types, so each layer gets its own bank shapes: ``gate_up`` [E, 2I, row_bytes] (expert e's gate rows, then its up rows) and ``down`` [E, H, row_bytes]; float experts keep their dtype ([E, 2I, H] and [E, H, I]). +Each layer keeps its own ggml types, so each layer gets its own bank shapes: ``gate_up`` [E, 2I, row_bytes] (expert e's gate rows, then its up rows) and ``down`` [E, H, row_bytes]; float experts (BF16 / F16 / F32) load and run bf16 ([E, 2I, H] and [E, H, I]). A GGUF that stores ``ffn_gate_up_exps`` (gemma4) already holds the gate_up rows in that order. -When NVFP4 experts carry per-expert ``.scale`` tensors, every layer also gets ``global`` [E, 4] fp32: the (gate, up, down, 1) weight_scale_2, 1 where a tensor has none.""" +When NVFP4 experts carry per-expert ``.scale`` tensors, every layer also gets ``global`` [E, 4] fp32: the (gate, up, down, 1) weight_scale_2, 1 where a tensor has none. +gpt-oss files carry per-expert gate / up / down biases; every layer then also gets ``gate_up_bias`` [E, 2I] and ``down_bias`` [E, H] bf16 rows, so the bias follows the expert through the slot cache like its weights do.""" from __future__ import annotations @@ -16,17 +17,17 @@ from freetoken.utils import init_logger from .reader import GgufTensorInfo, gguf_tensor_index -from .tensors import weight_scale +from .tensors import float_tensor, weight_scale logger = init_logger(__name__) def _expert_rows(info: GgufTensorInfo) -> torch.Tensor: - """[E, rows, row_bytes] packed ggml rows, or [E, rows, K] for a float tensor.""" + """[E, rows, row_bytes] packed ggml rows, or [E, rows, K] bf16 for an F16 / BF16 / F32 tensor (the fused MoE runs bf16).""" k, rows, experts = info.ne raw = torch.from_numpy(info.data()).view(experts, rows, info.nbytes // (experts * rows)) dtype = float_dtype(info.ggml_type) - return raw if dtype is None else raw.view(dtype) + return raw if dtype is None else raw.view(dtype).to(torch.bfloat16) def _drop_page_cache(info: GgufTensorInfo) -> None: @@ -69,7 +70,10 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int index = gguf_tensor_index(model_path) scaled = has_expert_scales({name: info.ggml_type for name, info in index.items()}) + # gpt-oss: a bias tensor per expert and role; other families ship none + biased = f"blk.{first_layer}.ffn_gate_exps.bias" in index roles = ("gate_up", "down", "global") if scaled else ("gate_up", "down") + roles += ("gate_up_bias", "down_bias") if biased else () sources: dict[str, list[torch.Tensor]] = {role: [] for role in roles} tails: dict[str, list[int]] = {role: [] for role in roles} tags = [] @@ -90,6 +94,9 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int tails["down"].append(mmq_tail_bytes(down.ggml_type, down.ne[0])) if scaled: tails["global"].append(0) + if biased: + tails["gate_up_bias"].append(0) + tails["down_bias"].append(0) rows = [_expert_rows(info) for info in parts] d = _expert_rows(down) gate_up_shape = (rows[0].shape[0], sum(r.shape[1] for r in rows), rows[0].shape[2]) @@ -108,9 +115,30 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int banks["global"] = torch.empty(experts, 4, dtype=torch.float32, device=device) if scaled: banks["global"].copy_(_expert_scales(parts, down, experts)) + if biased: + missing = [f"{blk}.ffn_{p}_exps.bias" for p in ("gate", "up", "down") if f"{blk}.ffn_{p}_exps.bias" not in index] + if missing: + raise NotImplementedError(f"{model_path}: {', '.join(missing)}: the per-expert biases ship on every layer of a gpt-oss file") + # gpt-oss per-expert biases: gate|up as one [E, 2I] row block, down as [E, H] + gate_b, up_b, down_b = ( + float_tensor(index[f"{blk}.ffn_{p}_exps.bias"]).to(torch.bfloat16) for p in ("gate", "up", "down") + ) + gate_up_bias, down_bias = torch.cat([gate_b, up_b], dim=1), down_b + for role, rows_bias in (("gate_up_bias", gate_up_bias), ("down_bias", down_bias)): + if rows_bias.shape[0] != experts: + raise NotImplementedError(f"{blk}.{role}: {tuple(rows_bias.shape)} for {experts} experts") + if device is None: + host[role] = HostBank(tuple(rows_bias.shape), torch.bfloat16) + banks[role] = host[role].tensor + else: + banks[role] = torch.empty_like(rows_bias, device=device) + banks[role].copy_(rows_bias) if dummy: banks["gate_up"].zero_() banks["down"].zero_() + if biased: + banks["gate_up_bias"].zero_() + banks["down_bias"].zero_() else: start = 0 for r in rows: diff --git a/python/freetoken/models/gpt_oss/__init__.py b/python/freetoken/models/gpt_oss/__init__.py index 45ef7d37c2..dbf9ed0421 100644 --- a/python/freetoken/models/gpt_oss/__init__.py +++ b/python/freetoken/models/gpt_oss/__init__.py @@ -1,5 +1,14 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import GptOssForCausalLM from .weight import iter_expert_pieces, iter_weights -__all__ = ["GptOssForCausalLM", "parse_config", "iter_weights", "iter_expert_pieces"] +__all__ = [ + "GptOssForCausalLM", + "parse_config", + "parse_gguf_config", + "iter_gguf_weights", + "iter_weights", + "iter_expert_pieces", + "to_gguf", +] diff --git a/python/freetoken/models/gpt_oss/gguf.py b/python/freetoken/models/gpt_oss/gguf.py new file mode 100644 index 0000000000..9cd21d2ac6 --- /dev/null +++ b/python/freetoken/models/gpt_oss/gguf.py @@ -0,0 +1,198 @@ +"""gpt-oss GGUFs: the llama.cpp ``gpt-oss`` arch served by the family's own HF classes with every tensor in its packed ggml type. + +The converter (conversion/gpt_oss.py) writes the attention projections fused as ``attn_qkv`` (``--fuse-qkv``; separate ``attn_q``/``attn_k``/``attn_v`` without), the attention sinks as ``attn_sinks.weight``, and a per-expert bias for each of gate / up / down - the only GGUF family whose experts carry biases, so they load as two dense tensors the ggml kernel adds around the gated activation. The MXFP4 releases (ggml-org) repack the routed experts into ggml MXFP4, scales ahead of the blocks in llama.cpp's nibble order, which the ggml kernels read as packed rows. + +The GGUF carries no per-layer attention types in any release: llama.cpp derives the alternating ``sliding_attention``/``full_attention`` layout from ``attention.sliding_window`` (its ``pattern`` KV, which no release writes, is honored when present), and so does this loader. The swiglu's alpha (1.702) and limit (7) are family constants llama.cpp also hardcodes; the GGUF stores neither. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +# gpt-oss swiglu constants, as llama.cpp's swiglu_oai op and every HF config ship them +_ACT_ALPHA = 1.702 +_ACT_LIMIT = 7.0 + +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", + "mlp.gate": "ffn_gate_inp", + "mlp.experts.0.gate_proj": "ffn_gate_exps", + "mlp.experts.0.up_proj": "ffn_up_exps", + "mlp.experts.0.down_proj": "ffn_down_exps", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + index = gguf_tensor_index(shim.model_path) + arch = shim.model_type # the GGUF architecture: gpt-oss + + def kv(key: str): + if (name := f"{arch}.{key}") not in md: + raise KeyError(f"missing GGUF metadata key {name}") + return md[name] + + check_scales(shim.model_path) + # llama.cpp treats the attention-output, router and per-expert biases as required tensors + # (openai-moe.cpp); without this check a GGUF that lost them would serve silently bias-free + missing = [ + f"blk.0.{name}.bias" + for name in ("attn_output", "ffn_gate_inp", "ffn_gate_exps", "ffn_up_exps", "ffn_down_exps") + if f"blk.0.{name}.bias" not in index + ] + if "blk.0.attn_qkv.bias" not in index and "blk.0.attn_q.bias" not in index: + missing.append("blk.0.attn_{qkv|q,k,v}.bias") + if missing: + raise NotImplementedError( + f"{shim.model_path}: {', '.join(missing)}: every gpt-oss release ships these (reconvert with a current llama.cpp)" + ) + scaling = md.get(f"{arch}.rope.scaling.type", "none") + if scaling == "yarn": + # every release's rope_scaling; the betas default to llama.cpp's when a file omits them + rope_scaling = { + "rope_type": "yarn", + "factor": float(kv("rope.scaling.factor")), + "original_max_position_embeddings": int(kv("rope.scaling.original_context_length")), + "beta_fast": float(md.get(f"{arch}.rope.scaling.yarn_beta_fast", 32.0)), + "beta_slow": float(md.get(f"{arch}.rope.scaling.yarn_beta_slow", 1.0)), + } + elif scaling == "none": + rope_scaling = None + else: + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + num_layers = int(kv("block_count")) + # llama.cpp's load_swa_pattern (openai-moe.cpp, default 2): the per-layer pattern array + # verbatim, else the scalar's il % p < p - 1 layout, layer 0 sliding + pattern = md.get(f"{arch}.attention.sliding_window.pattern") + if pattern is None: + layer_types = ["sliding_attention" if i % 2 < 1 else "full_attention" for i in range(num_layers)] + elif isinstance(pattern, list): + if len(pattern) != num_layers: + raise NotImplementedError( + f"{shim.model_path}: attention.sliding_window.pattern has {len(pattern)} entries for {num_layers} layers" + ) + layer_types = ["sliding_attention" if bool(p) else "full_attention" for p in pattern] + else: + p = int(pattern) + layer_types = ["sliding_attention" if p == 0 or i % p < p - 1 else "full_attention" for i in range(num_layers)] + return SimpleNamespace( + model_type="gpt_oss", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=num_layers, + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(kv("attention.key_length")), + rope_theta=float(kv("rope.freq_base")), + rope_scaling=rope_scaling, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("expert_feed_forward_length")), + num_local_experts=int(kv("expert_count")), + num_experts_per_tok=int(kv("expert_used_count")), + layer_types=layer_types, + sliding_window=int(kv("attention.sliding_window")), + attention_bias=True, + swiglu_limit=_ACT_LIMIT, + hidden_act_alpha=_ACT_ALPHA, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return replace(parse_config(_hf_config(shim)), expert_quant="gguf") + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + config = parse_gguf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = config.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def bias(name: str) -> torch.Tensor: + return float_tensor(index[f"{name}.bias"]) + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not config.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(config.num_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.post_attention_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + if f"{blk}.attn_qkv.bias" in index: + yield f"{sa}.qkv_proj.bias", bias(f"{blk}.attn_qkv") + elif f"{blk}.attn_q.bias" in index: + yield f"{sa}.qkv_proj.bias", torch.cat([bias(f"{blk}.attn_{p}") for p in ("q", "k", "v")]) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + if f"{blk}.attn_output.bias" in index: + yield f"{sa}.o_proj.bias", bias(f"{blk}.attn_output") + yield f"{sa}.sinks", float_tensor(t(f"{blk}.attn_sinks")) + mlp = f"{pre}.mlp" + yield f"{mlp}.router.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + if f"{blk}.ffn_gate_inp.bias" in index: + yield f"{mlp}.router.bias", bias(f"{blk}.ffn_gate_inp") + # the per-expert gate/up/down biases ride the expert banks (models/gguf/experts.py), + # following each expert through the slot cache like its weights do + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/gpt_oss/model.py b/python/freetoken/models/gpt_oss/model.py index 8398fb8be3..99de80b88c 100644 --- a/python/freetoken/models/gpt_oss/model.py +++ b/python/freetoken/models/gpt_oss/model.py @@ -51,6 +51,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/gpt_oss/moe.py b/python/freetoken/models/gpt_oss/moe.py index 3b58946b4d..140b3fa37c 100644 --- a/python/freetoken/models/gpt_oss/moe.py +++ b/python/freetoken/models/gpt_oss/moe.py @@ -41,7 +41,8 @@ def __init__(self, config: ModelConfig, layer_id: int | None = None, *, prefix: config.num_experts, has_bias=config.has_router_bias, ) - if config.moe_weight_format != "mxfp4": + gguf = config.expert_quant == "gguf" + if config.moe_weight_format != "mxfp4" and not gguf: raise ValueError( f"gpt-oss supports only mxfp4 expert weights, got " f"moe_weight_format={config.moe_weight_format!r}" @@ -49,10 +50,11 @@ def __init__(self, config: ModelConfig, layer_id: int | None = None, *, prefix: self.experts = make_moe_layer( config, layer_id=layer_id, - activation="gpt_oss_swiglu", + # the GGUF banks store gate|up split: swigluoai is the uninterleaved spelling of gpt_oss_swiglu + activation="swigluoai" if gguf else "gpt_oss_swiglu", alpha=config.hidden_act_alpha, limit=config.swiglu_limit, - interleaved=True, + interleaved=not gguf, has_bias=True, resident_cls=GptOssMoELayer, offload_cls=GptOssOffloadMoELayer, diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 3ed67c99bf..b19bc9c977 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -451,6 +451,15 @@ class ModelSpec: "GptOssForCausalLM", packed_modules_mapping=_DENSE_PACKED, ), + # llama.cpp gpt-oss GGUF (gpt-oss-20b / 120b): the same model classes with every + # tensor in its packed ggml type + "GptOssGGUFForCausalLM": ModelSpec( + "freetoken.models.gpt_oss", + "GptOssForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED, + ), "Glm4MoeForCausalLM": ModelSpec( "freetoken.models.glm4_moe", "Glm4MoeForCausalLM", diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index 1f275749d7..898bece64c 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -466,7 +466,7 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None return banks, {}, (H, I) def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_type: int, scaled: bool = False) -> tuple[dict, dict, tuple[int, int]]: - """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float rows keep their dtype. + """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float types never reach here (cpu_spec sends them to the bf16 format, whose banks load bf16). ``scaled`` adds global [E, 4] fp32, each expert's (gate, up, down, 1) NVFP4 weight_scale_2.""" from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes, type_name diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index 5c34fb33ca..02f9bd99d2 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -210,6 +210,8 @@ def fused_experts_impl( apply_router_weight_on_input: bool = False, act_alpha: float = 1.0, act_limit: float = float("inf"), + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, ) -> torch.Tensor: """Returns ``hidden_states`` itself, overwritten with the routed output. A caller that still needs the input afterwards (a shared expert, a residual) must read it BEFORE this @@ -225,6 +227,7 @@ def fused_experts_impl( assert w1.is_contiguous(), "Expert weights1 must be contiguous" assert w2.is_contiguous(), "Expert weights2 must be contiguous" assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] + assert w1_bias is None or not apply_router_weight_on_input, "per-expert biases need the unweighted first GEMM" num_tokens, _ = hidden_states.shape E, N, _ = w1.shape M = num_tokens @@ -285,6 +288,9 @@ def fused_experts_impl( config, compute_type=compute_type, ) + if w1_bias is not None: + # gpt-oss: the expert's gate/up bias is added before the gated activation + intermediate_cache1 += w1_bias[curr_topk_ids] gated_act_and_mul(activation, intermediate_cache1.view(-1, N), intermediate_cache2, alpha=act_alpha, limit=act_limit) fused_moe_kernel_triton( intermediate_cache2, @@ -305,6 +311,9 @@ def fused_experts_impl( intermediate_cache3, out_hidden_states[begin_token_idx:end_token_idx], ) + if w2_bias is not None: + # the down bias is weighted like the expert's output: sum_k w_k * bias_k + out_hidden_states[begin_token_idx:end_token_idx] += (curr_topk_weights.unsqueeze(-1) * w2_bias[curr_topk_ids]).sum(1) return out_hidden_states @@ -318,6 +327,8 @@ def fused_experts_decode_impl( apply_router_weight_on_input: bool = False, act_alpha: float = 1.0, act_limit: float = float("inf"), + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, ) -> torch.Tensor: from freetoken.kernel import fused_moe_decode_kernel_triton, moe_sum_reduce_triton from freetoken.layers import gated_act_and_mul @@ -336,6 +347,7 @@ def fused_experts_decode_impl( assert w1.dtype == hidden_states.dtype and w2.dtype == hidden_states.dtype assert topk_weights.dtype == torch.float32 assert topk_ids.dtype == torch.int32 + assert w1_bias is None or not apply_router_weight_on_input, "per-expert biases need the unweighted first GEMM" M, _ = hidden_states.shape top_k = topk_ids.shape[1] gate_up_dim = w1.shape[1] @@ -363,6 +375,8 @@ def fused_experts_decode_impl( compute_type=hidden_states.dtype, ) + if w1_bias is not None: + intermediate_cache1 += w1_bias[topk_ids] intermediate_cache2 = torch.empty( (M * top_k, intermediate_size), device=hidden_states.device, @@ -389,4 +403,6 @@ def fused_experts_decode_impl( out_hidden_states = torch.empty_like(hidden_states) moe_sum_reduce_triton(intermediate_cache3, out_hidden_states) + if w2_bias is not None: + out_hidden_states += (topk_weights.unsqueeze(-1) * w2_bias[topk_ids]).sum(1) return out_hidden_states diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 6610b128cd..4925c47050 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -46,7 +46,7 @@ # Half the host/cache footprint of bf16; the grouped GEMM (kernel/triton/fp8_blockscale_moe) # reads the routed fp8 rows directly and dequantizes in the K-loop (no bf16 materialization). "fp8_block": ("gate_up", "gate_up_scale", "down", "down_scale"), - # GGUF experts in their own ggml type per layer: gate_up [L*E, 2I, row_bytes] (gate rows, then up rows), down [L*E, H, row_bytes]; float experts keep their dtype + # GGUF experts in their own ggml type per layer: gate_up [L*E, 2I, row_bytes] (gate rows, then up rows), down [L*E, H, row_bytes]; float experts load and run bf16 "gguf": ("gate_up", "down"), # GGUF experts whose NVFP4 tensors carry per-expert scales: global [L*E, 4] fp32 (gate, up, down, 1) weight_scale_2 "gguf_scaled": ("gate_up", "down", "global"), @@ -483,10 +483,17 @@ def set_bank_sources( # loaders and FTW files may still name the banks the old way (gate_up_packed, ...) by_role = {canonical_role(name): per_layer for name, per_layer in sources.items()} - if set(by_role) != {canonical_role(n) for n in self.bank_schema}: - raise AssertionError( - f"banks {sorted(sources)} do not match the {self.quant_format!r} schema {self.bank_schema}" - ) + schema = {canonical_role(n) for n in self.bank_schema} + if set(by_role) != schema: + # gguf banks carry the file's own layout, so a file may hold roles the static + # schema lacks: the gpt-oss per-expert biases ride their own bank rows + extra = set(by_role) - schema + if self.quant_format in ("gguf", "gguf_scaled") and not (schema - set(by_role)) and extra <= {"gate_up_bias", "down_bias"}: + self.bank_schema = tuple(self.bank_schema) + tuple(sorted(extra)) + else: + raise AssertionError( + f"banks {sorted(sources)} do not match the {self.quant_format!r} schema {self.bank_schema}" + ) sources = {name: by_role[canonical_role(name)] for name in self.bank_schema} residency = layer_residency or [HostResidency.PINNED.value] * self.num_layers assert len(residency) == self.num_layers, (len(residency), self.num_layers) diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py index da34dba20d..fcfbf62fc5 100644 --- a/tests/moe/test_fused_moe.py +++ b/tests/moe/test_fused_moe.py @@ -2,12 +2,15 @@ import torch -def _activation_and_mul(gate_up: torch.Tensor, activation: str) -> torch.Tensor: +def _activation_and_mul(gate_up: torch.Tensor, activation: str, alpha: float = 1.0, limit: float = float("inf")) -> torch.Tensor: gate, up = gate_up.chunk(2, dim=-1) if activation == "silu": return torch.nn.functional.silu(gate) * up if activation == "gelu": return torch.nn.functional.gelu(gate) * up + if activation == "swigluoai": + gate = gate.clamp(max=limit) + return gate / (1.0 + torch.exp(-alpha * gate)) * (up.clamp(-limit, limit) + 1.0) raise AssertionError(f"unsupported test activation {activation}") @@ -19,6 +22,10 @@ def _reference_fused_experts_decode( topk_ids: torch.Tensor, activation: str = "silu", apply_router_weight_on_input: bool = False, + act_alpha: float = 1.0, + act_limit: float = float("inf"), + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, ) -> torch.Tensor: output = torch.zeros_like(hidden_states, dtype=torch.float32) hidden_fp32 = hidden_states.float() @@ -29,10 +36,14 @@ def _reference_fused_experts_decode( expert_id = int(topk_ids[token_idx, route_idx]) route_weight = topk_weights[token_idx, route_idx].float() gate_up = torch.matmul(w1_fp32[expert_id], hidden_fp32[token_idx]) + if w1_bias is not None: + gate_up = gate_up + w1_bias[expert_id].float() if apply_router_weight_on_input: gate_up = gate_up * route_weight - activated = _activation_and_mul(gate_up, activation) + activated = _activation_and_mul(gate_up, activation, alpha=act_alpha, limit=act_limit) contribution = torch.matmul(w2_fp32[expert_id], activated) + if w2_bias is not None: + contribution = contribution + w2_bias[expert_id].float() if not apply_router_weight_on_input: contribution = contribution * route_weight output[token_idx] += contribution @@ -156,6 +167,75 @@ def test_fused_experts_decode_matches_grouped_impl(dtype, apply_router_weight_on torch.testing.assert_close(output, expected, rtol=2e-2, atol=2e-2) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("impl_name", ["decode", "grouped"]) +def test_fused_experts_with_per_expert_biases_match_reference(impl_name): + """gpt-oss's expert banks carry per-expert biases: w1 before the gated activation, w2 before the routing weight.""" + from freetoken.moe.fused import fused_experts_decode_impl, fused_experts_impl + + device = torch.device("cuda") + batch_size = 4 + num_experts = 37 + hidden_size = 32 + intermediate_size = 24 + top_k = 4 + torch.manual_seed(104) + + hidden_states = 0.25 * torch.randn(batch_size, hidden_size, device=device, dtype=torch.bfloat16) + gate_up = torch.randn( + num_experts, + intermediate_size * 2, + hidden_size, + device=device, + dtype=torch.bfloat16, + ) * 0.25 + down = torch.randn( + num_experts, + hidden_size, + intermediate_size, + device=device, + dtype=torch.bfloat16, + ) * 0.25 + w1_bias = 0.5 * torch.randn(num_experts, intermediate_size * 2, device=device, dtype=torch.bfloat16) + w2_bias = 0.5 * torch.randn(num_experts, hidden_size, device=device, dtype=torch.bfloat16) + weights = torch.rand(batch_size, top_k, device=device, dtype=torch.float32) + topk_weights = weights / weights.sum(dim=-1, keepdim=True) + topk_ids = torch.tensor( + [[31, 4, 18, 0], [29, 7, 35, 12], [6, 21, 1, 33], [16, 3, 28, 9]], + device=device, + dtype=torch.int32, + ) + + impl = fused_experts_decode_impl if impl_name == "decode" else fused_experts_impl + output = impl( + hidden_states.clone(), + gate_up, + down, + topk_weights, + topk_ids, + activation="swigluoai", + act_alpha=1.702, + act_limit=7.0, + w1_bias=w1_bias, + w2_bias=w2_bias, + ) + expected = _reference_fused_experts_decode( + hidden_states, + gate_up, + down, + topk_weights, + topk_ids, + activation="swigluoai", + act_alpha=1.702, + act_limit=7.0, + w1_bias=w1_bias, + w2_bias=w2_bias, + ) + torch.cuda.synchronize() + + torch.testing.assert_close(output, expected, rtol=5e-2, atol=5e-2) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_fused_experts_grouped_impl_is_cuda_graph_capturable(): from freetoken.moe.fused import fused_experts_impl