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