xserv/docs/19-moe-gpt-oss.md

Phase 19: MoE — gpt-oss-20b

目标：在 xserv 支持 MoE，用 openai/gpt-oss-20b 端到端跑通，并与 llama.cpp 在 AIME 2025 / GSM8K 上对比正确性与性能。MXFP4 expert 权重加载时反量化为 BF16；整模型 ~40GB 单卡放不下 → 复用 Phase 18 的 PP（PP=2 ~20GB/卡，PP=4 ~10GB/卡）。

实时进度与重启续作指南见 docs/MOE_PROGRESS.md。

1. 架构（config.json，已核对）

num_hidden_layers=24, hidden=2880, head_dim=64（≠hidden/heads）, n_heads=64, n_kv_heads=8（GQA n_rep=8）, expert intermediate=2880, num_local_experts=32, num_experts_per_tok=4, vocab=201088, max_pos=131072, rope_theta=150000, sliding_window=128（交替层，见 layer_types）, rms_norm_eps=1e-5, swiglu_limit=7.0, alpha=1.702, tie_embeddings=false。 量化：MXFP4，仅 expert MLP（gate_up/down 的 _blocks+_scales）； attn/router/embed/lm_head 为 BF16。

2. 参考数学（HF transformers modeling_gpt_oss.py，逐字核对）

RMSNorm — 标准（fp32 算 variance，eps=1e-5）。

Router（GptOssTopKRouter，softmax 在 topk 之后，含 bias）

logits  = x @ W_router^T + b_router          # [T, 32]
top_val, idx = topk(logits, k=4, dim=-1)     # [T, 4]
top_val = softmax(top_val, dim=-1)           # 仅对选中的 4 个归一化
scores  = zeros[T,32].scatter(1, idx, top_val)

Experts（GptOssExperts，fused gate_up，interleaved；clamped；(up+1)·glu）

alpha=1.702; limit=7.0
gate_up = x @ gate_up_proj[e] + gate_up_proj_bias[e]   # [.., 2*dim]
gate = gate_up[..., ::2]; up = gate_up[..., 1::2]      # 偶/奇 交错
gate = clamp(gate, max=limit)                          # 仅上界
up   = clamp(up, min=-limit, max=limit)
glu  = gate * sigmoid(gate * alpha)
h    = (up + 1) * glu                                   # 注意 (up+1)
y_e  = h @ down_proj[e] + down_proj_bias[e]
out  = Σ_{e∈top4} scores[t,e] * y_e

Attention（eager_attention_forward，带 sinks）

scaling = head_dim**-0.5 = 64**-0.5；q/k/v/o 都有 bias
RoPE(theta=150000) on q,k；repeat_kv(n_rep=8)
attn = (q @ k^T) * scaling + causal_mask        # 滑窗层叠加 banded(window=128)
sinks = module.sinks[head]                       # 每 head 一个标量
combined = cat([attn, sinks broadcast], dim=-1)  # 多一列
combined -= combined.max(-1, keepdim)            # 数值稳定
probs = softmax(combined, -1)
scores = probs[..., :-1]                          # 丢掉 sink 列 => 概率不归一到 1
o = (scores @ v) -> merge heads -> @Wo + bo

sinks 等价于 softmax 分母多了 exp(sink)——可学习的"不注意"通道。 交替 sliding window：config layer_types 标明哪些层 window=128，其余全注意力。

与 Qwen3 的新增点：MoE FFN、MXFP4 反量化、attention sinks（softmax 多一列再丢）、 交替 sliding window、q/k/v/o bias、head_dim=64、clamped (up+1)*glu、rope_theta=150000。

实测张量布局（layer 0，已用 tools/mxfp4_probe.py 核对）

self_attn.q_proj.weight [4096,2880] +bias[4096]   # 64 heads*64
self_attn.k_proj.weight [512,2880]  +bias[512]    # 8 kv*64
self_attn.v_proj.weight [512,2880]  +bias[512]
self_attn.o_proj.weight [2880,4096] +bias[2880]
self_attn.sinks [64]                               # 每 q-head 一个标量（BF16）
input_layernorm.weight [2880]; post_attention_layernorm.weight [2880]
mlp.router.weight [32,2880] +bias[32]
mlp.experts.gate_up_proj_blocks [32,5760,90,16] U8 + _scales [32,5760,90] U8 + _bias[32,5760] BF16
mlp.experts.down_proj_blocks    [32,2880,90,16] U8 + _scales [32,2880,90] U8 + _bias[32,2880] BF16
# 全局: model.embed_tokens.weight, model.norm.weight, lm_head.weight (BF16)

MXFP4 打包：[..., nblk=90, 16] U8，每 16 字节 = 32 个 FP4 码（低 nibble=偶 idx，高 nibble=奇 idx）， 每 block 一个 E8M0 scale；90*32 = 2880 = 输入(hidden)维。即 gate_up 每 expert 权重逻辑 shape [5760 out, 2880 in]（已转置存储：行=out，列=in，与 HF nn.Linear 一致 y=x·Wᵀ）。

RoPE（rotate_half，非 interleave）

dim = head_dim = 64; base = rope_theta = 150000
inv_freq = 1 / base^(arange(0,64,2)/64)        # 32 项
freqs = pos ⊗ inv_freq                          # [S, 32]；cos/sin = cos(freqs)/sin(freqs)  (不 doubling)
# 应用: x=[.., 64], first=x[:32], second=x[32:]
#   out_first  = first*cos - second*sin
#   out_second = second*cos + first*sin

⚠️ 与 Qwen3 的 RoPE kernel（interleave）不同 —— gptoss 走 rotate_half。需单独处理。

Decoder layer（pre-norm 残差，结构同 Qwen3）

h = x + attn(input_norm(x))            # attn 含 sinks/bias/滑窗
out = h + moe(post_norm(h))            # moe = router + top4 experts 加权和

最终：logits = lm_head(norm(h_last))。无 q_norm/k_norm（与 Qwen3 不同，gptoss 没有）。

3. MXFP4 反量化（expert 权重）

expert 张量名：model.layers.{i}.mlp.experts.gate_up_proj_blocks/_scales、 ...down_proj_blocks/_scales（bias 为 BF16）。MXFP4：每 32 元素一 block 共享一个 E8M0(8-bit 指数) scale，每元素 4-bit FP4(E2M1)。反量化 val = fp4_lut[code] * 2^(e8m0 - 127)。P19.1 先用 Python(numpy) 反量化并与 HF 一层 数值对照（block 方向 / LUT / gate_up interleave），再写进 Rust loader。

4. 路线（正确优先）

1. P19.1 Python 侦查 + MXFP4 反量化验证（不依赖 GPU）。
2. P19.2 config.rs 加 MoE 字段（Qwen3 路径不变）。
3. P19.3 gptoss.rs：dense（attn+sinks+bias+滑窗 / norm / lm_head）+ MoE FFN （正确优先：逐 token top-4 gather→clamped SwiGLU→加权和）；MXFP4 在 from_weights 反量化为 BF16。验收：prefill logits 与 HF BF16 容差内一致（top-1 一致）。
4. P19.4 接 PP（experts 随层切），--pp 端到端；PP=2/4 与 PP=1 等价。
5. P19.5 llama.cpp 对比（升级 submodule 到支持 gpt-oss 的版本 + 取/转 GGUF）， 跑 AIME 2025 + GSM8K，复用 tools/bench + summarize_fullq.py。

5. 风险

• MXFP4 格式细节必须逐字对 → Python 反量化兜底。
• attention sinks + 交替滑窗：现有 flash/paged kernel 未必支持 → 正确优先版本先走朴素 attention（显式 mask + sink 列）。
• llama.cpp pinned b9371 早于 gpt-oss（约 2025-08）→ 需升级 submodule，有连锁影响。
• 性能：MoE 正确优先版本（逐 expert gather/scatter）会慢；先对再快。
• 环境：huggingface.co 被墙，需经代理 + hf-mirror 下载（见 MOE_PROGRESS.md §2）。

6. 不在本阶段范围

GPU 原生 MXFP4 + 按需反量化 kernel（先全 BF16）；高性能 grouped-GEMM / expert parallel； TP×MoE；单卡运行（需 MXFP4-native）。