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>

crates/xserv-kernels/src/activation.rs

@@ -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
+}

crates/xserv-kernels/src/attention.rs

@@ -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
+}

crates/xserv-kernels/src/lib.rs

@@ -10,10 +10,10 @@ pub mod rope;
 pub mod softmax;
 pub mod transpose;
 
-pub use activation::{add, gelu, mul, scale, silu, silu_mul};
+pub use activation::{add, gelu, gpt_oss_glu, mul, scale, silu, silu_mul};
 pub use argmax::{argmax_bf16_single, argmax_bf16_to_host};
 pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
-pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention, reshape_and_cache_bf16, reshape_and_cache_batched_bf16};
+pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention, paged_decode_attention_sinks, reshape_and_cache_bf16, reshape_and_cache_batched_bf16};
 pub use embedding::embedding;
 pub use gemm::{batched_matmul, matmul, GemmBackend};
 pub use layernorm::layernorm;

crates/xserv-kernels/src/rope.rs

@@ -37,6 +37,81 @@ impl RopeCache {
 
         Self { cos, sin, max_seq_len, half_dim }
     }
+
+    /// YaRN (Yet another RoPE extensioN) RoPE cache. Applies frequency-dependent
+    /// interpolation so the model can extrapolate beyond its training context.
+    pub fn new_yarn(
+        max_seq_len: usize,
+        head_dim: usize,
+        theta: f64,
+        factor: f64,
+        original_max_pos: usize,
+        beta_fast: f64,
+        beta_slow: f64,
+    ) -> Self {
+        let half_dim = head_dim / 2;
+        let dim = head_dim as f64;
+
+        // find_correction_dim: inverse formula to find dimension from number of rotations
+        let find_correction_dim = |num_rotations: f64| -> f64 {
+            dim * (original_max_pos as f64 / (num_rotations * 2.0 * std::f64::consts::PI)).ln()
+                / (2.0 * theta.ln())
+        };
+
+        let low_raw = find_correction_dim(beta_fast);
+        let high_raw = find_correction_dim(beta_slow);
+        // config has truncate=false, so use raw values (no floor/ceil)
+        let low = low_raw.max(0.0);
+        let high = high_raw.min((half_dim - 1) as f64);
+
+        // Compute inv_freq with YaRN interpolation
+        let mut inv_freq = vec![0.0f64; half_dim];
+        for i in 0..half_dim {
+            let pos_freq = theta.powf((2 * i) as f64 / dim);
+            let inv_freq_extrapolation = 1.0 / pos_freq;  // original
+            let inv_freq_interpolation = 1.0 / (factor * pos_freq);  // scaled
+
+            // Linear ramp: 0 where we keep original, 1 where we interpolate
+            let ramp = if (high - low).abs() < 0.001 {
+                0.5
+            } else {
+                ((i as f64 - low) / (high - low)).clamp(0.0, 1.0)
+            };
+            let extrapolation_factor = 1.0 - ramp;
+
+            inv_freq[i] = inv_freq_interpolation * (1.0 - extrapolation_factor)
+                + inv_freq_extrapolation * extrapolation_factor;
+        }
+
+        // Attention scaling factor for YaRN: 0.1 * ln(factor) + 1.0
+        let attn_factor = 0.1 * factor.ln() + 1.0;
+
+        // Build cos/sin cache on CPU then upload
+        let total = max_seq_len * half_dim;
+        let mut cos_host = vec![0.0f32; total];
+        let mut sin_host = vec![0.0f32; total];
+        for pos in 0..max_seq_len {
+            for i in 0..half_dim {
+                let angle = pos as f64 * inv_freq[i];
+                cos_host[pos * half_dim + i] = (angle.cos() * attn_factor) as f32;
+                sin_host[pos * half_dim + i] = (angle.sin() * attn_factor) as f32;
+            }
+        }
+
+        let nbytes = total * std::mem::size_of::<f32>();
+        let mut cos = GpuBuffer::alloc(nbytes).expect("alloc yarn cos_cache");
+        let mut sin = GpuBuffer::alloc(nbytes).expect("alloc yarn sin_cache");
+        let cos_bytes = unsafe {
+            std::slice::from_raw_parts(cos_host.as_ptr() as *const u8, nbytes)
+        };
+        let sin_bytes = unsafe {
+            std::slice::from_raw_parts(sin_host.as_ptr() as *const u8, nbytes)
+        };
+        cos.copy_from_host(cos_bytes).unwrap();
+        sin.copy_from_host(sin_bytes).unwrap();
+
+        Self { cos, sin, max_seq_len, half_dim }
+    }
 }
 
 /// Apply RoPE in-place to x.