moe(wip): gpt-oss-20b groundwork — config fields, arch doc, MXFP4 tools

Phase 19 start. config.rs: explicit head_dim (gpt-oss=64) + MoE fields
(num_local_experts, num_experts_per_tok, swiglu_limit, sliding_window,
layer_types) with accessors; Qwen3/GPT-2 paths unchanged (fall back to
hidden/num_heads when head_dim absent).

docs/19-moe-gpt-oss.md: architecture + exact HF reference math (router
softmax-after-topk, interleaved clamped (up+1)*glu experts, attention
sinks, alternating sliding window, rotate_half RoPE theta=150000,
head_dim 64), verified tensor layout, MXFP4 dequant plan.
docs/MOE_PROGRESS.md: resume/handoff snapshot.

tools/mxfp4_probe.py: inspect safetensors + validate MXFP4 decode (done).
tools/gptoss_dequant.py: MXFP4 experts -> plain BF16 safetensors dir so
the existing loader reads it (no MXFP4 in Rust for the first pass).

Verified: llama.cpp (dash5, LLM_ARCH_OPENAI_MOE) runs the gpt-oss-20b
MXFP4 GGUF correctly (17*24 -> 408) = the correctness oracle. MXFP4 decode
validated in numpy. Model + GGUF staged on dash5.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-05-29 21:01:53 +08:00
parent 057a3c68a3
commit c7d0750c32
4 changed files with 260 additions and 1 deletions

View File

@@ -46,6 +46,28 @@ pub struct ModelConfig {
pub rope_theta: Option<f64>,
#[serde(default)]
pub tie_word_embeddings: Option<bool>,
// Explicit head_dim (gpt-oss: 64, which is NOT hidden/num_heads). When
// absent, head_dim() falls back to hidden/num_heads (Qwen3, GPT-2).
#[serde(default)]
pub head_dim: Option<usize>,
// MoE (gpt-oss). Absent for dense models.
#[serde(default)]
pub num_local_experts: Option<usize>,
#[serde(default)]
pub num_experts_per_tok: Option<usize>,
// gpt-oss clamped-SwiGLU limit (config: swiglu_limit, default 7.0).
#[serde(default)]
pub swiglu_limit: Option<f64>,
// Sliding-window attention (gpt-oss: 128 on alternating layers). The
// pattern is given by `layer_types` (e.g. "sliding_attention" /
// "full_attention" per layer); absent for dense models.
#[serde(default)]
pub sliding_window: Option<usize>,
#[serde(default)]
pub layer_types: Option<Vec<String>>,
}
impl ModelConfig {
@@ -81,7 +103,48 @@ impl ModelConfig {
}
pub fn head_dim(&self) -> usize {
self.hidden() / self.num_heads()
// gpt-oss sets head_dim explicitly (64 != 2880/64). Dense models omit it.
self.head_dim.unwrap_or_else(|| self.hidden() / self.num_heads())
}
// ----- MoE (gpt-oss) -----
/// True for MoE models (have an expert count in config).
pub fn is_moe(&self) -> bool {
self.num_local_experts.is_some()
}
pub fn num_experts(&self) -> usize {
self.num_local_experts.unwrap_or(0)
}
pub fn experts_per_tok(&self) -> usize {
self.num_experts_per_tok.unwrap_or(0)
}
/// Clamp bound for gpt-oss SwiGLU (config `swiglu_limit`, default 7.0).
pub fn swiglu_limit(&self) -> f32 {
self.swiglu_limit.unwrap_or(7.0) as f32
}
/// Whether layer `i` uses sliding-window attention. gpt-oss alternates per
/// `layer_types`; if that's absent but `sliding_window` is set, fall back to
/// the common "every other layer" pattern (even = sliding). Dense → false.
pub fn layer_uses_sliding_window(&self, layer: usize) -> bool {
if self.sliding_window.is_none() {
return false;
}
match &self.layer_types {
Some(types) => types
.get(layer)
.map(|t| t.contains("sliding"))
.unwrap_or(false),
None => layer % 2 == 0,
}
}
pub fn sliding_window(&self) -> Option<usize> {
self.sliding_window
}
pub fn ln_eps(&self) -> f32 {