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:
@@ -13,6 +13,8 @@ unsafe extern "C" {
|
||||
fn launch_mul_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
|
||||
fn launch_mul_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
|
||||
fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
|
||||
fn launch_gpt_oss_glu_bf16(gate_up: *const c_void, out: *mut c_void, n_elements: i32,
|
||||
alpha: f32, limit: f32, stream: *mut c_void);
|
||||
}
|
||||
|
||||
fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
|
||||
@@ -97,3 +99,31 @@ pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// gpt-oss fused GLU activation (BF16 only).
|
||||
/// Input: gate_up [rows, 2*D] with interleaved columns (gate=even, up=odd).
|
||||
/// Output: [rows, D]
|
||||
/// Computes: gate.clamp(max=limit) * sigmoid(gate * alpha) * (up.clamp(-limit,limit) + 1)
|
||||
pub fn gpt_oss_glu(gate_up: &Tensor, alpha: f32, limit: f32) -> Tensor {
|
||||
assert!(gate_up.is_contiguous());
|
||||
assert!(matches!(gate_up.device(), Device::Cuda(_)));
|
||||
assert_eq!(gate_up.dtype(), DType::BF16, "gpt_oss_glu requires BF16");
|
||||
assert_eq!(gate_up.ndim(), 2);
|
||||
let rows = gate_up.shape()[0];
|
||||
let cols = gate_up.shape()[1];
|
||||
assert_eq!(cols % 2, 0);
|
||||
let d = cols / 2;
|
||||
let out = Tensor::empty(&[rows, d], gate_up.dtype(), gate_up.device());
|
||||
let n_elements = (rows * d) as i32;
|
||||
unsafe {
|
||||
launch_gpt_oss_glu_bf16(
|
||||
gate_up.data_ptr() as *const c_void,
|
||||
out.data_ptr() as *mut c_void,
|
||||
n_elements,
|
||||
alpha,
|
||||
limit,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user