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:
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user