feat(gguf): serve gpt-oss gguf checkpoints on gpt_oss

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:<path>`: 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) |

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",

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})"

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,20 +125,27 @@ 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:
             h = mul_mat_id(gate_up, gate_up_type, x)
             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,28 +155,37 @@ 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)
 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)

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",

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:

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",
+]