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

@@ -33,6 +33,18 @@ unsafe extern "C" {
head_dim: i32, max_blocks_per_seq: i32,
scale: f32, stream: *mut c_void,
);
fn launch_paged_decode_attention_sinks_bf16(
q: *const c_void,
k_cache: *const c_void,
v_cache: *const c_void,
o: *mut c_void,
block_tables: *const i32,
context_lens: *const i32,
sinks: *const c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
head_dim: i32, max_blocks_per_seq: i32,
scale: f32, window_size: i32, stream: *mut c_void,
);
fn launch_reshape_and_cache_bf16(
k_src: *const c_void, v_src: *const c_void,
k_pool: *mut c_void, v_pool: *mut c_void,
@@ -337,3 +349,58 @@ pub fn paged_decode_attention(
output
}
/// Paged decode attention with attention sinks and optional sliding window.
///
/// sinks_ptr: pointer to [num_q_heads] BF16 on GPU (or null for no sinks)
/// window_size: 0 = full attention, >0 = sliding window
#[allow(clippy::too_many_arguments)]
pub fn paged_decode_attention_sinks(
q: &Tensor,
k_cache_ptr: *const c_void,
v_cache_ptr: *const c_void,
block_tables_ptr: *const i32,
context_lens_ptr: *const i32,
sinks_ptr: *const c_void,
batch: usize,
num_q_heads: usize,
num_kv_heads: usize,
head_dim: usize,
max_blocks_per_seq: usize,
window_size: usize,
) -> Tensor {
assert_eq!(q.ndim(), 4);
assert_eq!(q.shape()[2], 1);
assert_eq!(q.dtype(), DType::BF16);
assert!(num_q_heads % num_kv_heads == 0);
assert!(head_dim <= 128);
let scale = 1.0 / (head_dim as f32).sqrt();
let output = Tensor::empty(
&[batch, num_q_heads, 1, head_dim],
DType::BF16,
q.device(),
);
unsafe {
launch_paged_decode_attention_sinks_bf16(
q.data_ptr() as *const c_void,
k_cache_ptr,
v_cache_ptr,
output.data_ptr() as *mut c_void,
block_tables_ptr,
context_lens_ptr,
sinks_ptr,
batch as i32,
num_q_heads as i32,
num_kv_heads as i32,
head_dim as i32,
max_blocks_per_seq as i32,
scale,
window_size as i32,
std::ptr::null_mut(),
);
}
output
}