phase19: MoE support — gpt-oss-20b end-to-end inference with TP=2

Add Mixture-of-Experts support for the gpt-oss-20b model (20.9B params,
32 experts × top-4 routing). Key additions:

- ModelConfig: MoE fields (num_local_experts, layer_types, sliding_window,
  attention_bias, explicit head_dim, rope_scaling, swiglu_limit)
- YaRN RoPE: RopeCache::new_yarn() with correct frequency interpolation
  and attention_scaling = 0.1*ln(factor)+1
- Custom GLU kernel: gpt_oss_glu_bf16 (clamped sigmoid gate activation)
- Paged attention with sinks + sliding window kernel variant
- GptOss model struct with expert-parallel TP (split 32 experts across ranks)
- bench-gpt-oss binary for TP inference benchmarking

Verified on dash5 with 2x RTX 5090: 63.6 tok/s decode, ~160ms TTFT.
Model generates topically-coherent output (needs chat template for quality).

Known issues:
- Custom GEMV kernel produces NaN with small N (workaround: pad to M=2)
- Prefill doesn't use attention sinks (uses standard flash attention)
- Output quality requires chat template formatting

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-30 15:18:01 +08:00
parent 46bfb59f30
commit 9ad91a4a92
12 changed files with 1390 additions and 44 deletions

View File

@@ -1,6 +1,15 @@
use serde::Deserialize;
use std::path::Path;
#[derive(Debug, Clone, Deserialize)]
pub struct RopeScaling {
pub rope_type: Option<String>,
pub factor: Option<f64>,
pub original_max_position_embeddings: Option<usize>,
pub beta_fast: Option<f64>,
pub beta_slow: Option<f64>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ModelConfig {
pub architectures: Option<Vec<String>>,
@@ -46,6 +55,24 @@ pub struct ModelConfig {
pub rope_theta: Option<f64>,
#[serde(default)]
pub tie_word_embeddings: Option<bool>,
// MoE (gpt-oss)
#[serde(default)]
pub num_local_experts: Option<usize>,
#[serde(default)]
pub num_experts_per_tok: Option<usize>,
#[serde(default)]
pub layer_types: Option<Vec<String>>,
#[serde(default)]
pub sliding_window: Option<usize>,
#[serde(default)]
pub attention_bias: Option<bool>,
#[serde(default, rename = "head_dim")]
pub explicit_head_dim: Option<usize>,
#[serde(default)]
pub rope_scaling: Option<RopeScaling>,
#[serde(default)]
pub swiglu_limit: Option<f64>,
}
impl ModelConfig {
@@ -81,7 +108,7 @@ impl ModelConfig {
}
pub fn head_dim(&self) -> usize {
self.hidden() / self.num_heads()
self.explicit_head_dim.unwrap_or_else(|| self.hidden() / self.num_heads())
}
pub fn ln_eps(&self) -> f32 {
@@ -93,4 +120,28 @@ impl ModelConfig {
pub fn tied_embeddings(&self) -> bool {
self.tie_word_embeddings.unwrap_or(true)
}
pub fn num_experts(&self) -> usize {
self.num_local_experts.unwrap_or(0)
}
pub fn experts_per_token(&self) -> usize {
self.num_experts_per_tok.unwrap_or(1)
}
pub fn is_moe(&self) -> bool {
self.num_local_experts.unwrap_or(0) > 1
}
pub fn is_sliding_layer(&self, layer_idx: usize) -> bool {
self.layer_types
.as_ref()
.and_then(|lt| lt.get(layer_idx))
.map(|t| t == "sliding_attention")
.unwrap_or(false)
}
pub fn window_size(&self) -> usize {
self.sliding_window.unwrap_or(0)
}
}