From a2de146fb6c2b2bb4cf8e2947b56bee251c0e1d1 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Mon, 13 Jul 2026 20:24:41 +0800 Subject: [PATCH] Add Qwen3.6 MoE inference support --- README.md | 16 +- crates/xserv-kernels/build.rs | 1 + crates/xserv-kernels/src/activation.rs | 42 + crates/xserv-kernels/src/attention.rs | 12 +- crates/xserv-kernels/src/deltanet.rs | 278 +++++ crates/xserv-kernels/src/lib.rs | 12 +- crates/xserv-kernels/src/moe.rs | 61 + crates/xserv-kernels/src/rope.rs | 52 + crates/xserv-model/src/bin/xserv-chat.rs | 60 +- crates/xserv-model/src/bin/xserv-cli.rs | 18 +- crates/xserv-model/src/config.rs | 180 ++- crates/xserv-model/src/decode_graph.rs | 4 +- crates/xserv-model/src/lib.rs | 10 +- crates/xserv-model/src/loader.rs | 60 +- crates/xserv-model/src/qwen35_moe.rs | 1447 ++++++++++++++++++++++ crates/xserv-model/src/sampling.rs | 64 +- crates/xserv-server/src/api.rs | 165 ++- crates/xserv-server/src/engine.rs | 21 +- crates/xserv-server/src/main.rs | 44 +- crates/xserv-server/src/pp_engine.rs | 146 ++- crates/xserv-server/src/tp_engine.rs | 37 +- csrc/activation/activations.cu | 83 ++ csrc/attention/deltanet.cu | 349 ++++++ csrc/attention/flash_attention.cu | 20 +- csrc/attention/paged_attention.cu | 4 +- csrc/embedding/rope.cu | 54 + csrc/moe/moe_sparse.cu | 62 + 27 files changed, 3153 insertions(+), 149 deletions(-) create mode 100644 crates/xserv-kernels/src/deltanet.rs create mode 100644 crates/xserv-model/src/qwen35_moe.rs create mode 100644 csrc/attention/deltanet.cu diff --git a/README.md b/README.md index 82c7317..0d4d233 100644 --- a/README.md +++ b/README.md @@ -3,13 +3,15 @@ > 从零用 **Rust + CUDA** 构建的 LLM 推理引擎，目标是吃透 LLM Serving 全栈技术。 xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架，自己实现了张量抽象、CUDA kernel、 -分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。支持 **Qwen3-8B**（BF16） -和 **gpt-oss-20b**（MoE，BF16/FP8/MXFP4 量化），多卡 TP/PP，并提供一套与 **llama.cpp** +分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。支持 **Qwen3-8B**、 +**Qwen3.6-35B-A3B**（BF16）和 **gpt-oss-20b**（MoE，BF16/FP8/MXFP4 量化），多卡 +TP/PP，并提供一套与 **llama.cpp** 对比正确性和性能的标准 benchmark。 ## 现状一览 -- **模型**：GPT-2（124M）、Qwen3-8B（BF16）、gpt-oss-20b（32 专家 top-4 MoE，harmony 格式） +- **模型**：GPT-2（124M）、Qwen3-8B（BF16）、Qwen3.6-35B-A3B（256 专家 top-8 MoE，BF16）、 + gpt-oss-20b（32 专家 top-4 MoE，harmony 格式） - **性能**（RTX 5090，贪心，单流）： - Qwen3-8B BF16 单卡：约 56 tok/s（HF transformers 的 1.4×） - gpt-oss-20b FP8 稀疏 MoE + CUDA Graph decode：**TPOT 5.8ms（~172 tok/s， @@ -110,6 +112,14 @@ curl http://localhost:8080/v1/chat/completions \ 其它端点：`GET /health`、`GET /v1/models`。 +Qwen3.6-35B-A3B 当前走 8 卡流水线并行（BF16，暂不支持单卡或 TP）： + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + ./target/release/xserv-server /path/to/qwen3.6-35b-a3b \ + --pp 8 --max-batch 1 --max-seq-len 4096 +``` + ### 2. 命令行推理 ```bash diff --git a/crates/xserv-kernels/build.rs b/crates/xserv-kernels/build.rs index b1806cf..7ddc879 100644 --- a/crates/xserv-kernels/build.rs +++ b/crates/xserv-kernels/build.rs @@ -30,6 +30,7 @@ fn main() { .file("../../csrc/attention/flash_attention.cu") .file("../../csrc/attention/paged_attention.cu") .file("../../csrc/attention/reshape_and_cache.cu") + .file("../../csrc/attention/deltanet.cu") .file("../../csrc/moe/moe_kernels.cu") .file("../../csrc/moe/moe_sparse.cu") .file("../../csrc/quantization/dequant_fp8.cu") diff --git a/crates/xserv-kernels/src/activation.rs b/crates/xserv-kernels/src/activation.rs index 45adc1b..af35c40 100644 --- a/crates/xserv-kernels/src/activation.rs +++ b/crates/xserv-kernels/src/activation.rs @@ -6,6 +6,10 @@ unsafe extern "C" { fn launch_gelu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void); fn launch_silu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void); fn launch_silu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void); + fn launch_sigmoid_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void); + fn launch_sigmoid_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void); + fn launch_softplus_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void); + fn launch_softplus_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void); fn launch_scale_f32( x: *const c_void, out: *mut c_void, @@ -48,6 +52,14 @@ unsafe extern "C" { n: i32, stream: *mut c_void, ); + fn launch_row_scale_bf16( + x: *const c_void, + scale: *const c_void, + out: *mut c_void, + rows: i32, + cols: i32, + stream: *mut c_void, + ); fn launch_silu_mul_bf16( gate: *const c_void, up: *const c_void, @@ -151,6 +163,12 @@ pub fn gelu(x: &Tensor) -> Tensor { pub fn silu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_silu_f32, launch_silu_bf16) } +pub fn sigmoid(x: &Tensor) -> Tensor { + dispatch_unary(x, launch_sigmoid_f32, launch_sigmoid_bf16) +} +pub fn softplus(x: &Tensor) -> Tensor { + dispatch_unary(x, launch_softplus_f32, launch_softplus_bf16) +} pub fn scale(x: &Tensor, scale_val: f32) -> Tensor { assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_))); @@ -190,6 +208,30 @@ pub fn mul(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16) } +pub fn row_scale_bf16(x: &Tensor, scale: &Tensor) -> Tensor { + assert_eq!(x.ndim(), 2); + assert_eq!(scale.ndim(), 2); + assert_eq!(x.dtype(), DType::BF16); + assert_eq!(scale.dtype(), DType::BF16); + assert!(x.is_contiguous() && scale.is_contiguous()); + assert!(matches!(x.device(), Device::Cuda(_))); + let rows = x.shape()[0]; + let cols = x.shape()[1]; + assert_eq!(scale.shape(), &[rows, 1]); + let out = Tensor::empty(&[rows, cols], DType::BF16, x.device()); + unsafe { + launch_row_scale_bf16( + x.data_ptr() as _, + scale.data_ptr() as _, + out.data_ptr() as *mut c_void, + rows as i32, + cols as i32, + xserv_cuda::current_stream_raw(), + ); + } + out +} + /// Row-broadcast bias add: out[r, c] = x[r, c] + bias[c] (BF16 only). pub fn bias_add_2d(x: &Tensor, bias: &Tensor) -> Tensor { assert_eq!(x.ndim(), 2); diff --git a/crates/xserv-kernels/src/attention.rs b/crates/xserv-kernels/src/attention.rs index 9da96a2..5540d62 100644 --- a/crates/xserv-kernels/src/attention.rs +++ b/crates/xserv-kernels/src/attention.rs @@ -412,8 +412,8 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens "num_q_heads must be divisible by num_kv_heads" ); assert!( - head_dim <= 128, - "flash_attention supports head_dim up to 128" + head_dim <= 256, + "flash_attention supports head_dim up to 256" ); // Dispatch to specialized decode kernel for single-token generation @@ -479,7 +479,7 @@ pub fn flash_attention_sinks( assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]); assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]); assert!(num_q_heads % num_kv_heads == 0); - assert!(head_dim <= 128); + assert!(head_dim <= 256); assert_eq!( sinks.shape()[0], num_q_heads, @@ -548,7 +548,7 @@ pub fn paged_decode_attention( num_q_heads % num_kv_heads == 0, "GQA: num_q_heads must be divisible by num_kv_heads" ); - assert!(head_dim <= 128); + assert!(head_dim <= 256); let scale = 1.0 / (head_dim as f32).sqrt(); let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device()); @@ -601,7 +601,7 @@ pub fn paged_decode_attention_tree( assert_eq!(q.shape()[2], 1); assert_eq!(q.dtype(), DType::BF16); assert!(num_q_heads % num_kv_heads == 0); - assert!(head_dim <= 128); + assert!(head_dim <= 256); let scale = 1.0 / (head_dim as f32).sqrt(); let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device()); @@ -653,7 +653,7 @@ pub fn paged_decode_attention_sinks( assert_eq!(q.shape()[2], 1); assert_eq!(q.dtype(), DType::BF16); assert!(num_q_heads % num_kv_heads == 0); - assert!(head_dim <= 128); + assert!(head_dim <= 256); let scale = 1.0 / (head_dim as f32).sqrt(); let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device()); diff --git a/crates/xserv-kernels/src/deltanet.rs b/crates/xserv-kernels/src/deltanet.rs new file mode 100644 index 0000000..f118178 --- /dev/null +++ b/crates/xserv-kernels/src/deltanet.rs @@ -0,0 +1,278 @@ +use std::ffi::c_void; + +use xserv_tensor::{DType, Device, Tensor}; + +unsafe extern "C" { + fn launch_depthwise_causal_conv1d_silu_bf16( + x: *const c_void, + w: *const c_void, + out: *mut c_void, + seq_len: i32, + channels: i32, + kernel: i32, + stream: *mut c_void, + ); + fn launch_deltanet_ar_bf16( + q: *const c_void, + k: *const c_void, + v: *const c_void, + beta: *const c_void, + alpha_softplus: *const c_void, + a_log: *const c_void, + state: *mut c_void, + out: *mut c_void, + key_heads: i32, + value_heads: i32, + head_dim: i32, + stream: *mut c_void, + ); + fn launch_depthwise_causal_conv1d_stateful_silu_bf16( + x: *const c_void, + w: *const c_void, + state: *mut c_void, + out: *mut c_void, + seq_len: i32, + channels: i32, + kernel: i32, + stream: *mut c_void, + ); + fn launch_l2_normalize_rows_bf16( + x: *const c_void, + out: *mut c_void, + rows: i32, + cols: i32, + eps: f32, + stream: *mut c_void, + ); + fn launch_deltanet_recurrent_bf16_f32( + q: *const c_void, + k: *const c_void, + v: *const c_void, + beta: *const c_void, + alpha_softplus: *const c_void, + a_log: *const c_void, + state: *mut c_void, + out: *mut c_void, + seq_len: i32, + key_heads: i32, + value_heads: i32, + head_dim: i32, + stream: *mut c_void, + ); +} + +fn conv_kernel_shape(w: &Tensor, channels: usize) -> usize { + match w.ndim() { + 2 => { + assert_eq!(w.shape()[0], channels); + w.shape()[1] + } + 3 => { + assert_eq!(w.shape()[0], channels); + assert_eq!(w.shape()[1], 1); + w.shape()[2] + } + _ => panic!("conv weight must be [C,K] or [C,1,K]"), + } +} + +/// Depthwise causal Conv1D + SiLU for Qwen3.5/3.6 DeltaNet. +/// x: [seq_len, channels] BF16, w: [channels, 1, kernel] or [channels, kernel] BF16. +pub fn depthwise_causal_conv1d_silu_bf16(x: &Tensor, w: &Tensor) -> Tensor { + assert_eq!(x.ndim(), 2); + assert!(x.is_contiguous() && w.is_contiguous()); + assert!(matches!(x.device(), Device::Cuda(_))); + assert_eq!(x.dtype(), DType::BF16); + assert_eq!(w.dtype(), DType::BF16); + let seq_len = x.shape()[0]; + let channels = x.shape()[1]; + let kernel = conv_kernel_shape(w, channels); + let out = Tensor::empty(&[seq_len, channels], DType::BF16, x.device()); + unsafe { + launch_depthwise_causal_conv1d_silu_bf16( + x.data_ptr() as *const c_void, + w.data_ptr() as *const c_void, + out.data_ptr() as *mut c_void, + seq_len as i32, + channels as i32, + kernel as i32, + xserv_cuda::current_stream_raw(), + ); + } + out +} + +/// Stateful depthwise Conv1D + SiLU. `state` is BF16 `[channels, kernel-1]` +/// and is updated in-place with the newest raw projection values. +pub fn depthwise_causal_conv1d_stateful_silu_bf16( + x: &Tensor, + w: &Tensor, + state: &mut Tensor, +) -> Tensor { + assert_eq!(x.ndim(), 2); + assert!(x.is_contiguous() && w.is_contiguous() && state.is_contiguous()); + assert!(matches!(x.device(), Device::Cuda(_))); + assert_eq!(x.dtype(), DType::BF16); + assert_eq!(w.dtype(), DType::BF16); + assert_eq!(state.dtype(), DType::BF16); + let seq_len = x.shape()[0]; + let channels = x.shape()[1]; + let kernel = conv_kernel_shape(w, channels); + assert_eq!(state.shape(), &[channels, kernel - 1]); + let out = Tensor::empty(&[seq_len, channels], DType::BF16, x.device()); + unsafe { + launch_depthwise_causal_conv1d_stateful_silu_bf16( + x.data_ptr() as *const c_void, + w.data_ptr() as *const c_void, + state.data_ptr() as *mut c_void, + out.data_ptr() as *mut c_void, + seq_len as i32, + channels as i32, + kernel as i32, + xserv_cuda::current_stream_raw(), + ); + } + out +} + +/// Row-wise L2 normalization with FP32 accumulation. +pub fn l2_normalize_bf16(x: &Tensor, eps: f32) -> Tensor { + assert_eq!(x.ndim(), 2); + assert!(x.is_contiguous()); + assert!(matches!(x.device(), Device::Cuda(_))); + assert_eq!(x.dtype(), DType::BF16); + let rows = x.shape()[0]; + let cols = x.shape()[1]; + let out = Tensor::empty(x.shape(), DType::BF16, x.device()); + unsafe { + launch_l2_normalize_rows_bf16( + x.data_ptr() as *const c_void, + out.data_ptr() as *mut c_void, + rows as i32, + cols as i32, + eps, + xserv_cuda::current_stream_raw(), + ); + } + out +} + +/// Stateful Qwen3.5/3.6 Gated DeltaNet for prefill or decode. +/// q/k: `[seq,key_heads,head_dim]`, v: `[seq,value_heads,head_dim]`. +/// beta/alpha: `[seq,value_heads]`; FP32 state is updated in-place. +pub fn deltanet_recurrent_bf16_f32( + q: &Tensor, + k: &Tensor, + v: &Tensor, + beta: &Tensor, + alpha_softplus: &Tensor, + a_log: &Tensor, + state: &mut Tensor, +) -> Tensor { + assert_eq!(q.ndim(), 3); + assert_eq!(k.ndim(), 3); + assert_eq!(v.ndim(), 3); + assert_eq!(beta.ndim(), 2); + assert_eq!(alpha_softplus.ndim(), 2); + assert_eq!(a_log.ndim(), 1); + assert_eq!(state.ndim(), 3); + assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous()); + assert!(beta.is_contiguous() && alpha_softplus.is_contiguous() && a_log.is_contiguous()); + assert!(state.is_contiguous()); + assert_eq!(q.dtype(), DType::BF16); + assert_eq!(k.dtype(), DType::BF16); + assert_eq!(v.dtype(), DType::BF16); + assert_eq!(beta.dtype(), DType::BF16); + assert_eq!(alpha_softplus.dtype(), DType::BF16); + assert_eq!(a_log.dtype(), DType::BF16); + assert_eq!(state.dtype(), DType::F32); + let seq_len = q.shape()[0]; + let key_heads = q.shape()[1]; + let head_dim = q.shape()[2]; + let value_heads = v.shape()[1]; + assert_eq!(k.shape(), &[seq_len, key_heads, head_dim]); + assert_eq!(v.shape(), &[seq_len, value_heads, head_dim]); + assert_eq!(beta.shape(), &[seq_len, value_heads]); + assert_eq!(alpha_softplus.shape(), &[seq_len, value_heads]); + assert_eq!(a_log.shape(), &[value_heads]); + assert_eq!(state.shape(), &[value_heads, head_dim, head_dim]); + assert_eq!(value_heads % key_heads, 0); + let out = Tensor::empty( + &[seq_len, value_heads, head_dim], + DType::BF16, + q.device(), + ); + unsafe { + launch_deltanet_recurrent_bf16_f32( + q.data_ptr() as *const c_void, + k.data_ptr() as *const c_void, + v.data_ptr() as *const c_void, + beta.data_ptr() as *const c_void, + alpha_softplus.data_ptr() as *const c_void, + a_log.data_ptr() as *const c_void, + state.data_ptr() as *mut c_void, + out.data_ptr() as *mut c_void, + seq_len as i32, + key_heads as i32, + value_heads as i32, + head_dim as i32, + xserv_cuda::current_stream_raw(), + ); + } + out +} + +/// Single-token autoregressive DeltaNet state update for Qwen3.5/3.6. +/// q/k: [key_heads, head_dim], v: [value_heads, head_dim], beta/alpha_softplus/a_log: [value_heads]. +/// state: [value_heads, head_dim, head_dim] updated in-place. Returns [1, value_heads * head_dim]. +pub fn deltanet_ar_bf16( + q: &Tensor, + k: &Tensor, + v: &Tensor, + beta: &Tensor, + alpha_softplus: &Tensor, + a_log: &Tensor, + state: &mut Tensor, +) -> Tensor { + assert_eq!(q.ndim(), 2); + assert_eq!(k.ndim(), 2); + assert_eq!(v.ndim(), 2); + assert_eq!(beta.ndim(), 1); + assert_eq!(alpha_softplus.ndim(), 1); + assert_eq!(a_log.ndim(), 1); + assert_eq!(state.ndim(), 3); + assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous()); + assert!(beta.is_contiguous() && alpha_softplus.is_contiguous() && a_log.is_contiguous()); + assert!(state.is_contiguous()); + assert!(matches!(q.device(), Device::Cuda(_))); + assert_eq!(q.dtype(), DType::BF16); + assert_eq!(k.dtype(), DType::BF16); + assert_eq!(v.dtype(), DType::BF16); + let key_heads = q.shape()[0]; + let head_dim = q.shape()[1]; + let value_heads = v.shape()[0]; + assert_eq!(k.shape(), &[key_heads, head_dim]); + assert_eq!(v.shape()[1], head_dim); + assert_eq!(beta.shape(), &[value_heads]); + assert_eq!(alpha_softplus.shape(), &[value_heads]); + assert_eq!(a_log.shape(), &[value_heads]); + assert_eq!(state.shape(), &[value_heads, head_dim, head_dim]); + let out = Tensor::empty(&[1, value_heads * head_dim], DType::BF16, q.device()); + unsafe { + launch_deltanet_ar_bf16( + q.data_ptr() as *const c_void, + k.data_ptr() as *const c_void, + v.data_ptr() as *const c_void, + beta.data_ptr() as *const c_void, + alpha_softplus.data_ptr() as *const c_void, + a_log.data_ptr() as *const c_void, + state.data_ptr() as *mut c_void, + out.data_ptr() as *mut c_void, + key_heads as i32, + value_heads as i32, + head_dim as i32, + xserv_cuda::current_stream_raw(), + ); + } + out +} diff --git a/crates/xserv-kernels/src/lib.rs b/crates/xserv-kernels/src/lib.rs index 2cd8122..a2c1173 100644 --- a/crates/xserv-kernels/src/lib.rs +++ b/crates/xserv-kernels/src/lib.rs @@ -1,6 +1,7 @@ pub mod activation; pub mod argmax; pub mod attention; +pub mod deltanet; pub mod dispatch; pub mod embedding; pub mod gemm; @@ -12,18 +13,25 @@ pub mod rope; pub mod softmax; pub mod transpose; -pub use activation::{add, bias_add_2d, gelu, gpt_oss_glu, mul, scale, silu, silu_mul}; +pub use activation::{ + add, bias_add_2d, gelu, gpt_oss_glu, mul, row_scale_bf16, scale, sigmoid, silu, silu_mul, + softplus, +}; pub use argmax::{argmax_bf16_single, argmax_bf16_to_host}; pub use attention::{ attention, copy_kv_position, decode_attention, flash_attention, flash_attention_sinks, paged_decode_attention, paged_decode_attention_sinks, paged_decode_attention_tree, reshape_and_cache_batched_bf16, reshape_and_cache_bf16, }; +pub use deltanet::{ + deltanet_ar_bf16, deltanet_recurrent_bf16_f32, depthwise_causal_conv1d_silu_bf16, + depthwise_causal_conv1d_stateful_silu_bf16, l2_normalize_bf16, +}; pub use embedding::{embedding, embedding_device_ids}; pub use gemm::{GemmBackend, batched_matmul, matmul, matmul_batched_gemv}; pub use layernorm::layernorm; pub use rmsnorm::{add_rmsnorm, rmsnorm}; -pub use rope::{RopeCache, rope_inplace, rope_inplace_device_pos}; +pub use rope::{RopeCache, partial_rope_bf16, rope_inplace, rope_inplace_device_pos}; pub use softmax::softmax; pub use transpose::{ merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, diff --git a/crates/xserv-kernels/src/moe.rs b/crates/xserv-kernels/src/moe.rs index e04060e..123b765 100644 --- a/crates/xserv-kernels/src/moe.rs +++ b/crates/xserv-kernels/src/moe.rs @@ -42,6 +42,21 @@ unsafe extern "C" { stream: *mut c_void, ); + fn launch_moe_sparse_gemv_bf16_bf16( + x: *const c_void, + w: *const c_void, + topk_ids: *const c_void, + y: *mut c_void, + num_tokens: i32, + n: i32, + k: i32, + top_k: i32, + expert_start: i32, + local_experts: i32, + x_per_slot: i32, + stream: *mut c_void, + ); + fn launch_moe_sparse_gemv_fp8_bf16( x: *const c_void, w: *const c_void, @@ -244,6 +259,52 @@ pub fn moe_weighted_sum( out } +/// Sparse MoE GEMV (BF16 weights): compute only routed experts. +/// x: [num_tokens, K] or [num_tokens * top_k, K], w_t: [local_experts, N, K]. +pub fn moe_sparse_gemv_bf16( + x: &Tensor, + w_t: &Tensor, + topk_ids: &Tensor, + top_k: usize, + expert_start: usize, + local_experts: usize, + x_per_slot: bool, +) -> Tensor { + assert_eq!(x.ndim(), 2); + assert_eq!(w_t.ndim(), 3); + assert_eq!(x.dtype(), DType::BF16); + assert_eq!(w_t.dtype(), DType::BF16); + assert!(x.is_contiguous() && w_t.is_contiguous()); + let local = w_t.shape()[0]; + let n = w_t.shape()[1]; + let k = w_t.shape()[2]; + assert_eq!(local, local_experts); + assert_eq!(x.shape()[1], k); + let num_tokens = if x_per_slot { + x.shape()[0] / top_k + } else { + x.shape()[0] + }; + let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device()); + unsafe { + launch_moe_sparse_gemv_bf16_bf16( + x.data_ptr() as *const c_void, + w_t.data_ptr() as *const c_void, + topk_ids.data_ptr() as *const c_void, + y.data_ptr() as *mut c_void, + num_tokens as i32, + n as i32, + k as i32, + top_k as i32, + expert_start as i32, + local_experts as i32, + if x_per_slot { 1 } else { 0 }, + xserv_cuda::current_stream_raw(), + ); + } + y +} + /// Sparse MoE GEMV (FP8 W8A16): compute only the routed experts. /// /// x: [num_tokens, K] BF16 (x_per_slot=false, gate_up) or diff --git a/crates/xserv-kernels/src/rope.rs b/crates/xserv-kernels/src/rope.rs index 552a59e..729a184 100644 --- a/crates/xserv-kernels/src/rope.rs +++ b/crates/xserv-kernels/src/rope.rs @@ -23,6 +23,17 @@ unsafe extern "C" { head_dim: i32, stream: *mut c_void, ); + fn launch_partial_rope_bf16( + x: *const c_void, + out: *mut c_void, + positions: *const c_void, + num_tokens: i32, + num_heads: i32, + head_dim: i32, + n_rot: i32, + theta: f32, + stream: *mut c_void, + ); fn launch_compute_rope_cache( cos_cache: *mut c_void, sin_cache: *mut c_void, @@ -171,6 +182,47 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) { rope_inplace_device_pos(x, cache, pos_gpu.as_ptr() as *const c_void); } +/// Apply partial RoPE and return a new tensor. +/// x: [num_tokens, num_heads, head_dim] BF16 on GPU. Only the first `n_rot` +/// dimensions are rotated; the tail is copied unchanged. +pub fn partial_rope_bf16(x: &Tensor, positions: &[u32], n_rot: usize, theta: f32) -> Tensor { + assert_eq!(x.ndim(), 3); + assert!(x.is_contiguous()); + assert!(matches!(x.device(), Device::Cuda(_))); + assert_eq!(x.dtype(), DType::BF16); + let num_tokens = x.shape()[0]; + let num_heads = x.shape()[1]; + let head_dim = x.shape()[2]; + assert_eq!(positions.len(), num_tokens); + assert!(n_rot <= head_dim && n_rot % 2 == 0); + + let pos_bytes = unsafe { + std::slice::from_raw_parts( + positions.as_ptr() as *const u8, + num_tokens * std::mem::size_of::(), + ) + }; + let mut pos_gpu = + xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions"); + pos_gpu.copy_from_host(pos_bytes).unwrap(); + let out = Tensor::empty(x.shape(), DType::BF16, x.device()); + + unsafe { + launch_partial_rope_bf16( + x.data_ptr() as *const c_void, + out.data_ptr() as *mut c_void, + pos_gpu.as_ptr() as *const c_void, + num_tokens as i32, + num_heads as i32, + head_dim as i32, + n_rot as i32, + theta, + xserv_cuda::current_stream_raw(), + ); + } + out +} + /// RoPE in-place with positions already on the GPU (u32, [num_tokens]). /// Used by the CUDA-graph decode path, where the position lives in a /// persistent device buffer updated outside the captured region. diff --git a/crates/xserv-model/src/bin/xserv-chat.rs b/crates/xserv-model/src/bin/xserv-chat.rs index c02f9e4..4b5a7e5 100644 --- a/crates/xserv-model/src/bin/xserv-chat.rs +++ b/crates/xserv-model/src/bin/xserv-chat.rs @@ -5,8 +5,8 @@ use std::sync::{Arc, mpsc}; use std::thread; use xserv_model::{ - BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, SamplingParams, - loader, sample, sample_greedy_penalized, + BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, ModelFamily, PagedKVCache, Qwen3, + SamplingParams, loader, sample, sample_greedy_penalized, }; use xserv_tensor::{DType, Device}; use xserv_tokenizer::Tokenizer; @@ -93,24 +93,26 @@ fn tp_worker_loop( rank as u32, )); let weights = loader::load_model_dir(&model_dir, Device::Cpu); - let model = if config.is_moe() { - ChatModel::GptOss(GptOss::from_weights_tp( + let model = match config.family() { + ModelFamily::GptOss => ChatModel::GptOss(GptOss::from_weights_tp( config.clone(), weights, rank, world, rank as u32, Some(tp), - )) - } else { - ChatModel::Qwen3(Qwen3::from_weights_tp( + )), + ModelFamily::Qwen35Moe => { + panic!("Qwen3.5/3.6 MoE chat inference is not implemented yet") + } + _ => ChatModel::Qwen3(Qwen3::from_weights_tp( config.clone(), weights, rank, world, rank as u32, Some(tp), - )) + )), }; let local_kv = config.num_kv_heads() / world; let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE; @@ -323,20 +325,27 @@ fn main() { ); let config = ModelConfig::from_file(&opts.model_dir.join("config.json")); - let model_type = config.model_type.as_deref().unwrap_or("unknown"); - let is_moe = config.is_moe(); + let model_type = config.model_type_str(); + let model_family = config.family(); + let is_gpt_oss = model_family == ModelFamily::GptOss; let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1); eprintln!( "Model: {model_type}{}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}", - if is_moe { " (MoE)" } else { "" }, + if config.is_gpt_oss() { " (MoE)" } else { "" }, config.num_layers(), config.hidden(), config.num_heads(), config.num_kv_heads(), - config.vocab_size, + config.vocab_size(), max_seq_len ); + if model_family == ModelFamily::Qwen35Moe { + eprintln!( + "Qwen3.5/3.6 MoE is recognized but chat inference is not implemented yet; it requires DeltaNet/linear-attention and Qwen MoE support." + ); + std::process::exit(1); + } let world = opts.tp; if world > 1 { @@ -374,7 +383,7 @@ fn main() { let tp = Arc::new(xserv_distributed::TpContext::init(0, world, id, 0)); let weights = loader::load_model_dir(&opts.model_dir, Device::Cpu); eprintln!("Loaded {} tensors", weights.len()); - let m = if is_moe { + let m = if is_gpt_oss { ChatModel::GptOss(GptOss::from_weights_tp( config.clone(), weights, @@ -412,7 +421,7 @@ fn main() { eprintln!("Loading weights..."); let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0)); eprintln!("Loaded {} tensors", weights.len()); - let m = if is_moe { + let m = if is_gpt_oss { ChatModel::GptOss(GptOss::from_weights(config.clone(), weights)) } else { ChatModel::Qwen3(Qwen3::from_weights(config.clone(), weights)) @@ -465,7 +474,7 @@ fn main() { _ => {} } - if is_moe { + if is_gpt_oss { // Harmony multi-turn: re-render the whole conversation (prior // analysis dropped) and re-prefill into a freshly cleared slot. let prompt = @@ -495,7 +504,7 @@ fn main() { max_new_tokens, use_color, &tp_handle, - is_moe, + is_gpt_oss, opts.enable_thinking, ); moe_history.push((input.to_string(), answer)); @@ -540,7 +549,7 @@ fn main() { max_new_tokens, use_color, &tp_handle, - is_moe, + is_gpt_oss, opts.enable_thinking, ); match finish { @@ -672,6 +681,7 @@ fn parse_args() -> CliOptions { temperature, top_k, top_p, + ..SamplingParams::default() }, system_prompt, enable_thinking, @@ -846,25 +856,25 @@ fn generate_with_paged_cache( max_tokens: usize, use_color: bool, tp: &Option, - is_moe: bool, + is_gpt_oss: bool, enable_thinking: bool, ) -> (Finish, String) { - let harmony_end_id = if is_moe { + let harmony_end_id = if is_gpt_oss { tokenizer.special_token_id("<|end|>") } else { None }; - let harmony_channel_id = if is_moe { + let harmony_channel_id = if is_gpt_oss { tokenizer.special_token_id("<|channel|>") } else { None }; - let harmony_message_id = if is_moe { + let harmony_message_id = if is_gpt_oss { tokenizer.special_token_id("<|message|>") } else { None }; - let harmony_special: Vec = if is_moe { + let harmony_special: Vec = if is_gpt_oss { [ "<|channel|>", "<|start|>", @@ -888,7 +898,7 @@ fn generate_with_paged_cache( InAnalysis, InFinal, } - let mut hstate = if is_moe { + let mut hstate = if is_gpt_oss { HarmonyState::InFinal } else { HarmonyState::Normal @@ -931,7 +941,7 @@ fn generate_with_paged_cache( let mut next = pick(&logits, sampling, &history); let mut decode_buffer = Vec::new(); let mut in_thinking = false; - let show_thinking = is_moe && enable_thinking; + let show_thinking = is_gpt_oss && enable_thinking; // Visible answer tokens, returned for multi-turn history. For moe this is // the final-channel content only (analysis is suppressed/gray); for Qwen3 // it is everything printed. The caller decodes these into the assistant @@ -1046,7 +1056,7 @@ fn generate_with_paged_cache( next = pick(&logits, sampling, &history); continue; } - if is_moe && hstate != HarmonyState::InFinal { + if is_gpt_oss && hstate != HarmonyState::InFinal { // Between harmony messages (after a channel's <|end|>, before the // next <|channel|>): the model emits a role header like "assistant". // That's structural, not user-visible content — suppress it. Only diff --git a/crates/xserv-model/src/bin/xserv-cli.rs b/crates/xserv-model/src/bin/xserv-cli.rs index f5fe4f9..59e5db7 100644 --- a/crates/xserv-model/src/bin/xserv-cli.rs +++ b/crates/xserv-model/src/bin/xserv-cli.rs @@ -1,7 +1,7 @@ use std::io::{self, Write}; use std::path::PathBuf; use xserv_model::{ - BLOCK_SIZE, KVCache, ModelConfig, PagedKVCache, SamplingParams, loader, sample, + BLOCK_SIZE, KVCache, ModelConfig, ModelFamily, PagedKVCache, SamplingParams, loader, sample, sample_greedy_penalized, }; use xserv_tensor::{DType, Device}; @@ -43,6 +43,7 @@ fn main() { temperature: flag(&args, "--temperature", 0.0f32), top_k: flag(&args, "--top-k", 0usize), top_p: flag(&args, "--top-p", 1.0f32), + ..SamplingParams::default() }; let rep_penalty = flag(&args, "--rep-penalty", 1.0f32); let rep_window = flag(&args, "--rep-window", 512usize); @@ -56,22 +57,29 @@ fn main() { ); let config = ModelConfig::from_file(&model_dir.join("config.json")); - let model_type = config.model_type.as_deref().unwrap_or("unknown"); + let model_type = config.model_type_str(); + let model_family = config.family(); eprintln!( "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}", config.num_layers(), config.hidden(), config.num_heads(), config.num_kv_heads(), - config.vocab_size + config.vocab_size() ); + if model_family == ModelFamily::Qwen35Moe { + eprintln!( + "Qwen3.5/3.6 MoE is recognized but CLI inference is not implemented yet; it requires DeltaNet/linear-attention and Qwen MoE support." + ); + std::process::exit(1); + } eprintln!("Loading weights..."); let weights = loader::load_model_dir(&model_dir, Device::Cuda(0)); eprintln!("Loaded {} tensors", weights.len()); - let is_qwen3 = model_type.contains("qwen"); - let is_gpt_oss = model_type.contains("gpt_oss"); + let is_qwen3 = model_family == ModelFamily::Qwen3; + let is_gpt_oss = model_family == ModelFamily::GptOss; let dtype = if is_qwen3 || is_gpt_oss { DType::BF16 } else { diff --git a/crates/xserv-model/src/config.rs b/crates/xserv-model/src/config.rs index 9d5b24b..8010363 100644 --- a/crates/xserv-model/src/config.rs +++ b/crates/xserv-model/src/config.rs @@ -1,6 +1,27 @@ use serde::Deserialize; use std::path::Path; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ModelFamily { + Gpt2, + Qwen3, + GptOss, + Qwen35Moe, + Unknown, +} + +impl ModelFamily { + pub fn as_str(self) -> &'static str { + match self { + ModelFamily::Gpt2 => "gpt2", + ModelFamily::Qwen3 => "qwen3", + ModelFamily::GptOss => "gpt_oss", + ModelFamily::Qwen35Moe => "qwen3_5_moe", + ModelFamily::Unknown => "unknown", + } + } +} + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { pub rope_type: Option, @@ -10,10 +31,21 @@ pub struct RopeScaling { pub beta_slow: Option, } +#[derive(Debug, Clone, Deserialize)] +pub struct RopeParameters { + pub rope_type: Option, + pub rope_theta: Option, + pub partial_rotary_factor: Option, + pub mrope_interleaved: Option, + pub mrope_section: Option>, +} + #[derive(Debug, Clone, Deserialize)] pub struct ModelConfig { pub architectures: Option>, pub model_type: Option, + #[serde(default)] + pub text_config: Option>, // Modern HF naming #[serde(default)] @@ -26,6 +58,7 @@ pub struct ModelConfig { pub num_key_value_heads: Option, #[serde(default)] pub num_hidden_layers: Option, + #[serde(default)] pub vocab_size: usize, #[serde(default)] pub max_position_embeddings: Option, @@ -56,14 +89,22 @@ pub struct ModelConfig { #[serde(default)] pub tie_word_embeddings: Option, - // MoE (gpt-oss) + // MoE (gpt-oss / Qwen MoE) #[serde(default)] pub num_local_experts: Option, #[serde(default)] + pub num_experts: Option, + #[serde(default)] pub num_experts_per_tok: Option, #[serde(default)] + pub moe_intermediate_size: Option, + #[serde(default)] + pub shared_expert_intermediate_size: Option, + #[serde(default)] pub layer_types: Option>, #[serde(default)] + pub full_attention_interval: Option, + #[serde(default)] pub sliding_window: Option, #[serde(default)] pub attention_bias: Option, @@ -72,11 +113,29 @@ pub struct ModelConfig { #[serde(default)] pub rope_scaling: Option, #[serde(default)] + pub rope_parameters: Option, + #[serde(default)] pub swiglu_limit: Option, #[serde(default)] pub geglu_alpha: Option, #[serde(default)] pub hidden_act: Option, + + // Qwen3.5/3.6 linear attention / DeltaNet metadata + #[serde(default)] + pub linear_conv_kernel_dim: Option, + #[serde(default)] + pub linear_key_head_dim: Option, + #[serde(default)] + pub linear_num_key_heads: Option, + #[serde(default)] + pub linear_num_value_heads: Option, + #[serde(default)] + pub linear_value_head_dim: Option, + #[serde(default)] + pub partial_rotary_factor: Option, + #[serde(default)] + pub mtp_num_hidden_layers: Option, } impl ModelConfig { @@ -87,80 +146,153 @@ impl ModelConfig { .unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display())) } + pub fn text(&self) -> &Self { + self.text_config.as_deref().unwrap_or(self) + } + + pub fn model_type_str(&self) -> &str { + self.text() + .model_type + .as_deref() + .or(self.model_type.as_deref()) + .unwrap_or("unknown") + } + + pub fn family(&self) -> ModelFamily { + let top_type = self.model_type.as_deref().unwrap_or(""); + let text_type = self.text().model_type.as_deref().unwrap_or(top_type); + let has_arch = |needle: &str| { + self.architectures + .as_ref() + .map(|arch| arch.iter().any(|a| a.contains(needle))) + .unwrap_or(false) + }; + + if top_type == "qwen3_5_moe" || text_type == "qwen3_5_moe_text" || has_arch("Qwen3_5Moe") { + ModelFamily::Qwen35Moe + } else if text_type == "gpt_oss" || top_type == "gpt_oss" { + ModelFamily::GptOss + } else if text_type.contains("qwen") || top_type.contains("qwen") { + ModelFamily::Qwen3 + } else if text_type.contains("gpt2") || top_type.contains("gpt2") { + ModelFamily::Gpt2 + } else { + ModelFamily::Unknown + } + } + pub fn hidden(&self) -> usize { - self.hidden_size - .or(self.n_embd) + let c = self.text(); + c.hidden_size + .or(c.n_embd) .expect("hidden_size or n_embd required") } pub fn num_heads(&self) -> usize { - self.num_attention_heads - .or(self.n_head) + let c = self.text(); + c.num_attention_heads + .or(c.n_head) .expect("num_attention_heads or n_head required") } pub fn num_layers(&self) -> usize { - self.num_hidden_layers - .or(self.n_layer) + let c = self.text(); + c.num_hidden_layers + .or(c.n_layer) .expect("num_hidden_layers or n_layer required") } pub fn max_seq_len(&self) -> usize { - self.max_position_embeddings - .or(self.n_positions) - .unwrap_or(2048) + let c = self.text(); + c.max_position_embeddings.or(c.n_positions).unwrap_or(2048) } pub fn ffn_hidden(&self) -> usize { - self.intermediate_size - .or(self.n_inner) - .unwrap_or(self.hidden() * 4) + let c = self.text(); + c.intermediate_size + .or(c.moe_intermediate_size) + .or(c.n_inner) + .unwrap_or_else(|| self.hidden() * 4) + } + + pub fn vocab_size(&self) -> usize { + self.text().vocab_size } pub fn num_kv_heads(&self) -> usize { - self.num_key_value_heads.unwrap_or(self.num_heads()) + self.text().num_key_value_heads.unwrap_or(self.num_heads()) } pub fn head_dim(&self) -> usize { - self.explicit_head_dim + self.text() + .explicit_head_dim .unwrap_or_else(|| self.hidden() / self.num_heads()) } pub fn ln_eps(&self) -> f32 { - self.layer_norm_eps - .or(self.layer_norm_epsilon) + let c = self.text(); + c.layer_norm_eps + .or(c.layer_norm_epsilon) + .or(c.rms_norm_eps) .unwrap_or(1e-5) as f32 } + pub fn rope_theta_value(&self) -> Option { + let c = self.text(); + c.rope_theta + .or_else(|| c.rope_parameters.as_ref().and_then(|rp| rp.rope_theta)) + } + pub fn tied_embeddings(&self) -> bool { - self.tie_word_embeddings.unwrap_or(true) + self.text().tie_word_embeddings.unwrap_or(true) } pub fn num_experts(&self) -> usize { - self.num_local_experts.unwrap_or(0) + self.text() + .num_local_experts + .or(self.text().num_experts) + .unwrap_or(0) } pub fn experts_per_token(&self) -> usize { - self.num_experts_per_tok.unwrap_or(1) + self.text().num_experts_per_tok.unwrap_or(1) } pub fn is_moe(&self) -> bool { - self.num_local_experts.unwrap_or(0) > 1 + self.num_experts() > 1 + } + + pub fn is_gpt_oss(&self) -> bool { + self.family() == ModelFamily::GptOss + } + + pub fn is_qwen35_moe(&self) -> bool { + self.family() == ModelFamily::Qwen35Moe } pub fn is_sliding_layer(&self, layer_idx: usize) -> bool { - self.layer_types + self.text() + .layer_types .as_ref() .and_then(|lt| lt.get(layer_idx)) .map(|t| t == "sliding_attention") .unwrap_or(false) } + pub fn is_linear_attention_layer(&self, layer_idx: usize) -> bool { + self.text() + .layer_types + .as_ref() + .and_then(|lt| lt.get(layer_idx)) + .map(|t| t == "linear_attention") + .unwrap_or(false) + } + pub fn window_size(&self) -> usize { - self.sliding_window.unwrap_or(0) + self.text().sliding_window.unwrap_or(0) } pub fn geglu_alpha(&self) -> f32 { - self.geglu_alpha.unwrap_or(1.702) as f32 + self.text().geglu_alpha.unwrap_or(1.702) as f32 } } diff --git a/crates/xserv-model/src/decode_graph.rs b/crates/xserv-model/src/decode_graph.rs index 74bad8f..b04044c 100644 --- a/crates/xserv-model/src/decode_graph.rs +++ b/crates/xserv-model/src/decode_graph.rs @@ -98,9 +98,9 @@ impl DecodeGraphState { let num_kv_heads = config.num_kv_heads(); let head_dim = config.head_dim(); let intermediate = config.ffn_hidden(); - let vocab_size = config.vocab_size; + let vocab_size = config.vocab_size(); let num_layers = config.num_layers(); - let eps = config.rms_norm_eps.unwrap_or(1e-6) as f32; + let eps = config.ln_eps(); let es = 2usize; // BF16 = 2 bytes let stream = CudaStream::new().expect("create CUDA stream for graph"); diff --git a/crates/xserv-model/src/lib.rs b/crates/xserv-model/src/lib.rs index 8af2412..492809d 100644 --- a/crates/xserv-model/src/lib.rs +++ b/crates/xserv-model/src/lib.rs @@ -8,10 +8,11 @@ pub mod kv_cache; pub mod loader; pub mod paged_kv_cache; pub mod qwen3; +pub mod qwen35_moe; pub mod qwen3_graph; pub mod sampling; -pub use config::ModelConfig; +pub use config::{ModelConfig, ModelFamily}; pub use decode_graph::{DecodeGraphState, LayerWeightPtrs}; pub use gpt_oss::GptOss; pub use gpt_oss_graph::{GptOssDecodeGraph, GraphedGptOssDecoder}; @@ -19,7 +20,12 @@ pub use gpt2::{GPT2, KVCache}; pub use kv_cache::GpuKVCache; pub use paged_kv_cache::{BLOCK_SIZE, BlockAllocator, Location, PagedKVCache}; pub use qwen3::Qwen3; -pub use sampling::{SamplingParams, sample, sample_greedy_penalized}; +pub use qwen35_moe::{ + Qwen35AttentionMeta, Qwen35FullAttentionOutput, Qwen35FullAttentionWeights, + Qwen35LinearAttentionWeights, Qwen35LinearProjectionOutput, Qwen35Moe, Qwen35MoeSpec, + Qwen35RecurrentCache, Qwen35TensorMeta, Qwen35TensorSpec, +}; +pub use sampling::{SamplingParams, sample, sample_greedy_penalized, sample_with_history}; /// Initialize GPU kernel hooks. Called automatically by model constructors, /// but safe to call multiple times (idempotent via OnceLock). diff --git a/crates/xserv-model/src/loader.rs b/crates/xserv-model/src/loader.rs index b5ddd97..47a2aca 100644 --- a/crates/xserv-model/src/loader.rs +++ b/crates/xserv-model/src/loader.rs @@ -1,10 +1,20 @@ use half::{bf16, f16}; use safetensors::SafeTensors; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::path::Path; use xserv_tensor::{DType, Device, Tensor}; pub fn load_safetensors(path: &Path, device: Device) -> HashMap { + load_safetensors_filtered(path, device, |_| true) +} + +/// Load only tensors accepted by `keep`. The safetensors file is still read as +/// one byte buffer, but unneeded tensor payloads are not copied into Tensors. +pub fn load_safetensors_filtered( + path: &Path, + device: Device, + keep: impl Fn(&str) -> bool, +) -> HashMap { let data = std::fs::read(path).unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display())); let st = SafeTensors::deserialize(&data) @@ -13,6 +23,9 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap let mut tensors = HashMap::new(); for (name, view) in st.tensors() { + if !keep(&name) { + continue; + } let shape: Vec = view.shape().to_vec(); let raw_bytes = view.data(); let dtype = match view.dtype() { @@ -36,27 +49,64 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap /// Load from a directory containing model.safetensors (or sharded files) + config.json. pub fn load_model_dir(dir: &Path, device: Device) -> HashMap { + load_model_dir_filtered(dir, device, |_| true) +} + +/// Load a filtered subset of a model directory. For indexed sharded models, +/// consult `model.safetensors.index.json` first so shards containing no wanted +/// tensors are never read. This is critical for pipeline parallelism on large +/// models: each stage should read only its layer range, not the full checkpoint. +pub fn load_model_dir_filtered( + dir: &Path, + device: Device, + keep: impl Fn(&str) -> bool, +) -> HashMap { let single = dir.join("model.safetensors"); if single.exists() { - return load_safetensors(&single, device); + return load_safetensors_filtered(&single, device, keep); } + let index_path = dir.join("model.safetensors.index.json"); + let wanted_shards: Option> = if index_path.exists() { + let text = std::fs::read_to_string(&index_path) + .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display())); + let index: serde_json::Value = serde_json::from_str(&text) + .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display())); + let map = index + .get("weight_map") + .and_then(|v| v.as_object()) + .unwrap_or_else(|| panic!("{} has no weight_map", index_path.display())); + Some( + map.iter() + .filter(|(name, _)| keep(name)) + .filter_map(|(_, shard)| shard.as_str().map(str::to_string)) + .collect(), + ) + } else { + None + }; + // Try sharded: model-00001-of-NNNNN.safetensors let mut all_tensors = HashMap::new(); let mut entries: Vec<_> = std::fs::read_dir(dir) .unwrap() .filter_map(|e| e.ok()) .filter(|e| { - e.path() + let is_safetensors = e + .path() .file_name() .map(|f| f.to_string_lossy().ends_with(".safetensors")) - .unwrap_or(false) + .unwrap_or(false); + let is_wanted = wanted_shards.as_ref().is_none_or(|wanted| { + wanted.contains(&e.file_name().to_string_lossy().to_string()) + }); + is_safetensors && is_wanted }) .collect(); entries.sort_by_key(|e| e.file_name()); for entry in entries { - let tensors = load_safetensors(&entry.path(), device); + let tensors = load_safetensors_filtered(&entry.path(), device, &keep); all_tensors.extend(tensors); } diff --git a/crates/xserv-model/src/qwen35_moe.rs b/crates/xserv-model/src/qwen35_moe.rs new file mode 100644 index 0000000..4959ef2 --- /dev/null +++ b/crates/xserv-model/src/qwen35_moe.rs @@ -0,0 +1,1447 @@ +use std::collections::{BTreeMap, HashMap}; + +use half::bf16; +use xserv_kernels::{ + GemmBackend, add, attention, bias_add_2d, deltanet_ar_bf16, + deltanet_recurrent_bf16_f32, depthwise_causal_conv1d_silu_bf16, + depthwise_causal_conv1d_stateful_silu_bf16, embedding, flash_attention, + l2_normalize_bf16, matmul, merge_heads_gpu, mul, paged_decode_attention, + partial_rope_bf16, repeat_kv_gpu, rmsnorm, row_scale_bf16, sigmoid, silu, silu_mul, + softplus, +}; +use xserv_tensor::{DType, Device, Tensor}; + +use crate::config::ModelConfig; +use crate::paged_kv_cache::PagedKVCache; + +/// Qwen3.5/3.6 stores every RMSNorm weight except the Gated DeltaNet output +/// norm as a zero-centered delta. The reference implementation applies +/// `weight + 1` before the RMSNorm multiplication. +fn qwen35_norm_with_unit_offset(weight: Tensor) -> Tensor { + assert_eq!(weight.device(), Device::Cpu); + assert_eq!(weight.dtype(), DType::BF16); + let shape = weight.shape().to_vec(); + let values = weight + .as_slice::() + .iter() + .map(|value| bf16::from_f32(value.to_f32() + 1.0)) + .collect::>(); + Tensor::from_slice(&values, &shape) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Qwen35LayerKind { + LinearAttention, + FullAttention, +} + +#[derive(Debug, Clone)] +pub struct Qwen35MoeSpec { + pub layer_kinds: Vec, + pub hidden: usize, + pub vocab: usize, + pub num_heads: usize, + pub num_kv_heads: usize, + pub head_dim: usize, + pub linear_key_heads: usize, + pub linear_value_heads: usize, + pub linear_key_dim: usize, + pub linear_value_dim: usize, + pub linear_conv_kernel: usize, + pub num_experts: usize, + pub experts_per_token: usize, + pub moe_intermediate: usize, + pub shared_intermediate: usize, +} + +#[derive(Debug, Clone)] +pub struct Qwen35TensorSpec { + pub name: String, + pub dtype: &'static str, + pub shape: Vec, +} + +#[derive(Debug, Clone)] +pub struct Qwen35TensorMeta { + pub name: String, + pub dtype: String, + pub shape: Vec, +} + +#[derive(Debug, Clone)] +pub struct Qwen35LinearAttentionMeta { + pub in_proj_qkv: Qwen35TensorMeta, + pub in_proj_z: Qwen35TensorMeta, + pub in_proj_a: Qwen35TensorMeta, + pub in_proj_b: Qwen35TensorMeta, + pub conv1d: Qwen35TensorMeta, + pub a_log: Qwen35TensorMeta, + pub dt_bias: Qwen35TensorMeta, + pub norm: Qwen35TensorMeta, + pub out_proj: Qwen35TensorMeta, +} + +#[derive(Debug, Clone)] +pub struct Qwen35FullAttentionMeta { + pub q_proj: Qwen35TensorMeta, + pub k_proj: Qwen35TensorMeta, + pub v_proj: Qwen35TensorMeta, + pub o_proj: Qwen35TensorMeta, + pub q_norm: Qwen35TensorMeta, + pub k_norm: Qwen35TensorMeta, +} + +#[derive(Debug, Clone)] +pub struct Qwen35MoeLayerMeta { + pub gate: Qwen35TensorMeta, + pub experts_gate_up: Qwen35TensorMeta, + pub experts_down: Qwen35TensorMeta, + pub shared_gate: Qwen35TensorMeta, + pub shared_up: Qwen35TensorMeta, + pub shared_down: Qwen35TensorMeta, + pub shared_expert_gate: Qwen35TensorMeta, +} + +#[derive(Debug, Clone)] +pub struct Qwen35LayerMeta { + pub input_norm: Qwen35TensorMeta, + pub post_attention_norm: Qwen35TensorMeta, + pub attention: Qwen35AttentionMeta, + pub moe: Qwen35MoeLayerMeta, +} + +#[derive(Debug, Clone)] +pub enum Qwen35AttentionMeta { + Linear(Qwen35LinearAttentionMeta), + Full(Qwen35FullAttentionMeta), +} + +#[derive(Debug, Clone)] +pub struct Qwen35ModelMeta { + pub embed_tokens: Qwen35TensorMeta, + pub norm: Qwen35TensorMeta, + pub lm_head: Qwen35TensorMeta, + pub layers: Vec, +} + +impl Qwen35TensorMeta { + pub fn matches_spec(&self, spec: &Qwen35TensorSpec) -> bool { + self.dtype == spec.dtype && self.shape == spec.shape + } +} + +impl Qwen35TensorSpec { + fn new(name: impl Into, dtype: &'static str, shape: impl Into>) -> Self { + Self { + name: name.into(), + dtype, + shape: shape.into(), + } + } +} + +impl Qwen35MoeSpec { + pub fn from_config(config: &ModelConfig) -> Self { + let text = config.text(); + let layer_kinds = (0..config.num_layers()) + .map(|i| { + if config.is_linear_attention_layer(i) { + Qwen35LayerKind::LinearAttention + } else { + Qwen35LayerKind::FullAttention + } + }) + .collect(); + Self { + layer_kinds, + hidden: config.hidden(), + vocab: config.vocab_size(), + num_heads: config.num_heads(), + num_kv_heads: config.num_kv_heads(), + head_dim: config.head_dim(), + linear_key_heads: text.linear_num_key_heads.unwrap_or(16), + linear_value_heads: text.linear_num_value_heads.unwrap_or(32), + linear_key_dim: text.linear_key_head_dim.unwrap_or(128), + linear_value_dim: text.linear_value_head_dim.unwrap_or(128), + linear_conv_kernel: text.linear_conv_kernel_dim.unwrap_or(4), + num_experts: config.num_experts(), + experts_per_token: config.experts_per_token(), + moe_intermediate: config.ffn_hidden(), + shared_intermediate: text + .shared_expert_intermediate_size + .unwrap_or_else(|| config.ffn_hidden()), + } + } + + pub fn linear_layers(&self) -> usize { + self.layer_kinds + .iter() + .filter(|&&kind| kind == Qwen35LayerKind::LinearAttention) + .count() + } + + pub fn full_attention_layers(&self) -> usize { + self.layer_kinds.len() - self.linear_layers() + } + + pub fn linear_key_width(&self) -> usize { + self.linear_key_heads * self.linear_key_dim + } + + pub fn linear_value_width(&self) -> usize { + self.linear_value_heads * self.linear_value_dim + } + + pub fn linear_qkv_width(&self) -> usize { + 2 * self.linear_key_width() + self.linear_value_width() + } + + pub fn linear_conv_width(&self) -> usize { + self.linear_qkv_width() + } + + pub fn full_q_width(&self) -> usize { + 2 * self.num_heads * self.head_dim + } + + pub fn full_query_width(&self) -> usize { + self.num_heads * self.head_dim + } + + pub fn full_gate_width(&self) -> usize { + self.num_heads * self.head_dim + } + + pub fn full_kv_width(&self) -> usize { + self.num_kv_heads * self.head_dim + } + + pub fn full_output_width(&self) -> usize { + self.num_heads * self.head_dim + } + + pub fn describe_full_attention_path(&self) -> String { + format!( + "q_proj=[{},{}] -> query={} + gate={}, k/v=[{},{}], o_proj=[{},{}], qk_norm={}, gate=sigmoid(gate) before o_proj", + self.full_q_width(), + self.hidden, + self.full_query_width(), + self.full_gate_width(), + self.full_kv_width(), + self.hidden, + self.hidden, + self.full_output_width(), + self.head_dim, + ) + } + + pub fn expected_tensor_specs(&self) -> Vec { + let mut specs = Vec::new(); + specs.push(Qwen35TensorSpec::new( + "model.language_model.embed_tokens.weight", + "BF16", + [self.vocab, self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + "model.language_model.norm.weight", + "BF16", + [self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + "lm_head.weight", + "BF16", + [self.vocab, self.hidden], + )); + + for (i, kind) in self.layer_kinds.iter().copied().enumerate() { + let p = format!("model.language_model.layers.{i}"); + self.push_common_layer_specs(&mut specs, &p); + match kind { + Qwen35LayerKind::LinearAttention => { + self.push_linear_attention_specs(&mut specs, &p) + } + Qwen35LayerKind::FullAttention => self.push_full_attention_specs(&mut specs, &p), + } + } + specs + } + + pub fn build_meta( + &self, + tensors: &BTreeMap, + ) -> Result { + let specs = self.expected_tensor_specs(); + for spec in &specs { + let meta = tensors + .get(&spec.name) + .ok_or_else(|| format!("missing tensor {}", spec.name))?; + if !meta.matches_spec(spec) { + return Err(format!( + "tensor {} expected {} {:?}, got {} {:?}", + spec.name, spec.dtype, spec.shape, meta.dtype, meta.shape + )); + } + } + + let take = |name: &str| -> Qwen35TensorMeta { + tensors + .get(name) + .unwrap_or_else(|| panic!("validated tensor missing: {name}")) + .clone() + }; + let mut layers = Vec::with_capacity(self.layer_kinds.len()); + for (i, kind) in self.layer_kinds.iter().copied().enumerate() { + let p = format!("model.language_model.layers.{i}"); + let attention = match kind { + Qwen35LayerKind::LinearAttention => { + Qwen35AttentionMeta::Linear(Qwen35LinearAttentionMeta { + in_proj_qkv: take(&format!("{p}.linear_attn.in_proj_qkv.weight")), + in_proj_z: take(&format!("{p}.linear_attn.in_proj_z.weight")), + in_proj_a: take(&format!("{p}.linear_attn.in_proj_a.weight")), + in_proj_b: take(&format!("{p}.linear_attn.in_proj_b.weight")), + conv1d: take(&format!("{p}.linear_attn.conv1d.weight")), + a_log: take(&format!("{p}.linear_attn.A_log")), + dt_bias: take(&format!("{p}.linear_attn.dt_bias")), + norm: take(&format!("{p}.linear_attn.norm.weight")), + out_proj: take(&format!("{p}.linear_attn.out_proj.weight")), + }) + } + Qwen35LayerKind::FullAttention => { + Qwen35AttentionMeta::Full(Qwen35FullAttentionMeta { + q_proj: take(&format!("{p}.self_attn.q_proj.weight")), + k_proj: take(&format!("{p}.self_attn.k_proj.weight")), + v_proj: take(&format!("{p}.self_attn.v_proj.weight")), + o_proj: take(&format!("{p}.self_attn.o_proj.weight")), + q_norm: take(&format!("{p}.self_attn.q_norm.weight")), + k_norm: take(&format!("{p}.self_attn.k_norm.weight")), + }) + } + }; + layers.push(Qwen35LayerMeta { + input_norm: take(&format!("{p}.input_layernorm.weight")), + post_attention_norm: take(&format!("{p}.post_attention_layernorm.weight")), + attention, + moe: Qwen35MoeLayerMeta { + gate: take(&format!("{p}.mlp.gate.weight")), + experts_gate_up: take(&format!("{p}.mlp.experts.gate_up_proj")), + experts_down: take(&format!("{p}.mlp.experts.down_proj")), + shared_gate: take(&format!("{p}.mlp.shared_expert.gate_proj.weight")), + shared_up: take(&format!("{p}.mlp.shared_expert.up_proj.weight")), + shared_down: take(&format!("{p}.mlp.shared_expert.down_proj.weight")), + shared_expert_gate: take(&format!("{p}.mlp.shared_expert_gate.weight")), + }, + }); + } + Ok(Qwen35ModelMeta { + embed_tokens: take("model.language_model.embed_tokens.weight"), + norm: take("model.language_model.norm.weight"), + lm_head: take("lm_head.weight"), + layers, + }) + } + + fn push_common_layer_specs(&self, specs: &mut Vec, p: &str) { + specs.push(Qwen35TensorSpec::new( + format!("{p}.input_layernorm.weight"), + "BF16", + [self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.post_attention_layernorm.weight"), + "BF16", + [self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.mlp.gate.weight"), + "BF16", + [self.num_experts, self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.mlp.experts.gate_up_proj"), + "BF16", + [self.num_experts, 2 * self.moe_intermediate, self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.mlp.experts.down_proj"), + "BF16", + [self.num_experts, self.hidden, self.moe_intermediate], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.mlp.shared_expert.gate_proj.weight"), + "BF16", + [self.shared_intermediate, self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.mlp.shared_expert.up_proj.weight"), + "BF16", + [self.shared_intermediate, self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.mlp.shared_expert.down_proj.weight"), + "BF16", + [self.hidden, self.shared_intermediate], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.mlp.shared_expert_gate.weight"), + "BF16", + [1, self.hidden], + )); + } + + fn push_linear_attention_specs(&self, specs: &mut Vec, p: &str) { + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.in_proj_qkv.weight"), + "BF16", + [self.linear_qkv_width(), self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.in_proj_z.weight"), + "BF16", + [self.linear_value_width(), self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.in_proj_a.weight"), + "BF16", + [self.linear_value_heads, self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.in_proj_b.weight"), + "BF16", + [self.linear_value_heads, self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.conv1d.weight"), + "BF16", + [self.linear_conv_width(), 1, self.linear_conv_kernel], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.A_log"), + "BF16", + [self.linear_value_heads], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.dt_bias"), + "BF16", + [self.linear_value_heads], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.norm.weight"), + "BF16", + [self.linear_value_dim], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.linear_attn.out_proj.weight"), + "BF16", + [self.hidden, self.linear_value_width()], + )); + } + + fn push_full_attention_specs(&self, specs: &mut Vec, p: &str) { + specs.push(Qwen35TensorSpec::new( + format!("{p}.self_attn.q_proj.weight"), + "BF16", + [self.full_q_width(), self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.self_attn.k_proj.weight"), + "BF16", + [self.full_kv_width(), self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.self_attn.v_proj.weight"), + "BF16", + [self.full_kv_width(), self.hidden], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.self_attn.o_proj.weight"), + "BF16", + [self.hidden, self.num_heads * self.head_dim], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.self_attn.q_norm.weight"), + "BF16", + [self.head_dim], + )); + specs.push(Qwen35TensorSpec::new( + format!("{p}.self_attn.k_norm.weight"), + "BF16", + [self.head_dim], + )); + } +} + +pub struct Qwen35FullAttentionWeights { + pub q_w_t: Tensor, + pub k_w_t: Tensor, + pub v_w_t: Tensor, + pub o_w_t: Tensor, + pub q_norm: Tensor, + pub k_norm: Tensor, +} + +pub struct Qwen35FullAttentionOutput { + pub q_full: Tensor, + pub query: Tensor, + pub gate: Tensor, + pub k: Tensor, + pub v: Tensor, + pub q_normed: Tensor, + pub k_normed: Tensor, + pub q_rope: Tensor, + pub k_rope: Tensor, + pub attn_out: Tensor, + pub attn_merged: Tensor, + pub gate_sigmoid: Tensor, + pub gated_attn: Tensor, + pub projected: Tensor, +} + +pub struct Qwen35LinearAttentionWeights { + pub in_proj_qkv_t: Tensor, + pub in_proj_z_t: Tensor, + pub in_proj_a_t: Tensor, + pub in_proj_b_t: Tensor, + pub conv1d: Tensor, + pub a_log: Tensor, + pub dt_bias: Tensor, + pub norm: Tensor, + pub out_proj_t: Tensor, +} + +pub struct Qwen35LinearProjectionOutput { + pub qkv: Tensor, + pub z: Tensor, + pub a: Tensor, + pub b: Tensor, + pub beta: Tensor, + pub alpha_softplus: Tensor, + pub q: Tensor, + pub k: Tensor, + pub v: Tensor, + pub conv: Tensor, + pub q_conv: Tensor, + pub k_conv: Tensor, + pub v_conv: Tensor, + pub recurrent_out: Option, + pub norm_gated: Option, + pub projected: Option, +} + +pub struct Qwen35Moe { + pub config: ModelConfig, + pub spec: Qwen35MoeSpec, + embed_tokens: Tensor, + layers: Vec, + norm: Tensor, + lm_head_t: Tensor, + layer_start: usize, + is_first_stage: bool, + is_last_stage: bool, +} + +struct Qwen35Block { + input_norm: Tensor, + post_norm: Tensor, + attention: Qwen35AttentionWeights, + moe: Qwen35MoeWeights, +} + +enum Qwen35AttentionWeights { + Linear(Qwen35LinearAttentionWeights), + Full(Qwen35FullAttentionWeights), +} + +struct Qwen35MoeWeights { + router_w_t: Tensor, + expert_gate_up: Tensor, + expert_down: Tensor, + shared_gate_w_t: Tensor, + shared_up_w_t: Tensor, + shared_down_w_t: Tensor, + shared_router_w_t: Tensor, +} + +struct Qwen35LinearState { + conv: Tensor, + recurrent: Tensor, +} + +/// Per-stage recurrent state for the serial PP engine. The outer dimension is +/// sequence slot and the inner dimension is local layer index; full-attention +/// layers have no recurrent entry. +pub struct Qwen35RecurrentCache { + slots: Vec>>, + layer_kinds: Vec, + conv_channels: usize, + conv_history: usize, + value_heads: usize, + value_dim: usize, + device: Device, +} + +impl Qwen35RecurrentCache { + fn make_layer_state(&self, kind: Qwen35LayerKind) -> Option { + if kind == Qwen35LayerKind::FullAttention { + return None; + } + Some(Qwen35LinearState { + conv: Tensor::zeros( + &[self.conv_channels, self.conv_history], + DType::BF16, + self.device, + ), + recurrent: Tensor::zeros( + &[self.value_heads, self.value_dim, self.value_dim], + DType::F32, + self.device, + ), + }) + } + + pub fn reset_slot(&mut self, slot: usize) { + assert!(slot < self.slots.len()); + let kinds = self.layer_kinds.clone(); + self.slots[slot] = kinds + .into_iter() + .map(|kind| self.make_layer_state(kind)) + .collect(); + } + + fn linear_mut(&mut self, slot: usize, layer: usize) -> &mut Qwen35LinearState { + self.slots[slot][layer] + .as_mut() + .expect("linear layer missing recurrent state") + } +} + +impl Qwen35Moe { + pub fn take_linear_attention_projection_weights( + weights: &mut HashMap, + layer: usize, + device: xserv_tensor::Device, + ) -> Qwen35LinearAttentionWeights { + let p = format!("model.language_model.layers.{layer}.linear_attn"); + let mut take = |name: &str| -> Tensor { + weights + .remove(name) + .unwrap_or_else(|| panic!("missing tensor {name}")) + .to_device(device) + }; + Qwen35LinearAttentionWeights { + in_proj_qkv_t: take(&format!("{p}.in_proj_qkv.weight")) + .transpose(0, 1) + .contiguous(), + in_proj_z_t: take(&format!("{p}.in_proj_z.weight")) + .transpose(0, 1) + .contiguous(), + in_proj_a_t: take(&format!("{p}.in_proj_a.weight")) + .transpose(0, 1) + .contiguous(), + in_proj_b_t: take(&format!("{p}.in_proj_b.weight")) + .transpose(0, 1) + .contiguous(), + conv1d: take(&format!("{p}.conv1d.weight")).contiguous(), + a_log: take(&format!("{p}.A_log")), + dt_bias: take(&format!("{p}.dt_bias")), + norm: take(&format!("{p}.norm.weight")), + out_proj_t: take(&format!("{p}.out_proj.weight")) + .transpose(0, 1) + .contiguous(), + } + } + + pub fn project_linear_attention( + spec: &Qwen35MoeSpec, + input: &Tensor, + weights: &Qwen35LinearAttentionWeights, + ) -> Qwen35LinearProjectionOutput { + let seq_len = input.shape()[0]; + assert_eq!(input.shape(), &[seq_len, spec.hidden]); + let qkv = matmul(input, &weights.in_proj_qkv_t, GemmBackend::CuBlas); + let z = matmul(input, &weights.in_proj_z_t, GemmBackend::CuBlas); + let a = matmul(input, &weights.in_proj_a_t, GemmBackend::CuBlas); + let b = matmul(input, &weights.in_proj_b_t, GemmBackend::CuBlas); + let beta = sigmoid(&b); + let alpha_softplus = softplus(&bias_add_2d(&a, &weights.dt_bias)); + let key_width = spec.linear_key_width(); + let value_width = spec.linear_value_width(); + let q = qkv.narrow(1, 0, key_width).contiguous(); + let k = qkv.narrow(1, key_width, key_width).contiguous(); + let v = qkv.narrow(1, 2 * key_width, value_width).contiguous(); + let conv = depthwise_causal_conv1d_silu_bf16(&qkv, &weights.conv1d); + let q_conv = conv.narrow(1, 0, key_width).contiguous(); + let k_conv = conv.narrow(1, key_width, key_width).contiguous(); + let v_conv = conv.narrow(1, 2 * key_width, value_width).contiguous(); + Qwen35LinearProjectionOutput { + qkv, + z, + a, + b, + beta, + alpha_softplus, + q, + k, + v, + conv, + q_conv, + k_conv, + v_conv, + recurrent_out: None, + norm_gated: None, + projected: None, + } + } + + pub fn linear_attention_recurrent_step( + spec: &Qwen35MoeSpec, + projected: &Qwen35LinearProjectionOutput, + weights: &Qwen35LinearAttentionWeights, + state: &mut Tensor, + ) -> Tensor { + assert_eq!(projected.q_conv.shape()[0], 1, "AR step requires one token"); + let q = projected + .q_conv + .reshape(&[spec.linear_key_heads, spec.linear_key_dim]); + let k = projected + .k_conv + .reshape(&[spec.linear_key_heads, spec.linear_key_dim]); + let v = projected + .v_conv + .reshape(&[spec.linear_value_heads, spec.linear_value_dim]); + let beta = projected.beta.reshape(&[spec.linear_value_heads]); + let alpha_softplus = projected.alpha_softplus.reshape(&[spec.linear_value_heads]); + deltanet_ar_bf16(&q, &k, &v, &beta, &alpha_softplus, &weights.a_log, state) + } + + pub fn finish_linear_attention( + spec: &Qwen35MoeSpec, + recurrent_out: &Tensor, + z: &Tensor, + weights: &Qwen35LinearAttentionWeights, + eps: f32, + ) -> (Tensor, Tensor) { + assert_eq!(recurrent_out.shape(), &[1, spec.linear_value_width()]); + assert_eq!(z.shape(), &[1, spec.linear_value_width()]); + let normed = rmsnorm( + &recurrent_out.reshape(&[spec.linear_value_heads, spec.linear_value_dim]), + &weights.norm, + eps, + ) + .reshape(&[1, spec.linear_value_width()]); + let gated = mul(&normed, &silu(z)); + let projected = matmul(&gated, &weights.out_proj_t, GemmBackend::CuBlas); + (gated, projected) + } + + pub fn take_full_attention_weights( + weights: &mut HashMap, + layer: usize, + device: xserv_tensor::Device, + ) -> Qwen35FullAttentionWeights { + let p = format!("model.language_model.layers.{layer}.self_attn"); + let mut take = |name: &str, unit_offset: bool| -> Tensor { + let weight = weights + .remove(name) + .unwrap_or_else(|| panic!("missing tensor {name}")); + let weight = if unit_offset { + qwen35_norm_with_unit_offset(weight) + } else { + weight + }; + weight.to_device(device) + }; + Qwen35FullAttentionWeights { + q_w_t: take(&format!("{p}.q_proj.weight"), false) + .transpose(0, 1) + .contiguous(), + k_w_t: take(&format!("{p}.k_proj.weight"), false) + .transpose(0, 1) + .contiguous(), + v_w_t: take(&format!("{p}.v_proj.weight"), false) + .transpose(0, 1) + .contiguous(), + o_w_t: take(&format!("{p}.o_proj.weight"), false) + .transpose(0, 1) + .contiguous(), + q_norm: take(&format!("{p}.q_norm.weight"), true), + k_norm: take(&format!("{p}.k_norm.weight"), true), + } + } + + pub fn forward_full_attention_layer( + spec: &Qwen35MoeSpec, + input: &Tensor, + weights: &Qwen35FullAttentionWeights, + positions: &[u32], + n_rot: usize, + rope_theta: f32, + eps: f32, + ) -> Qwen35FullAttentionOutput { + let seq_len = input.shape()[0]; + assert_eq!(input.shape(), &[seq_len, spec.hidden]); + assert_eq!(positions.len(), seq_len); + let q_full = matmul(input, &weights.q_w_t, GemmBackend::CuBlas); + let k = matmul(input, &weights.k_w_t, GemmBackend::CuBlas); + let v = matmul(input, &weights.v_w_t, GemmBackend::CuBlas); + // q_proj rows are laid out [head, {query, gate}, head_dim]. + let q_and_gate = q_full.reshape(&[seq_len, spec.num_heads, 2, spec.head_dim]); + let query = q_and_gate + .narrow(2, 0, 1) + .squeeze(2) + .contiguous() + .reshape(&[seq_len, spec.full_query_width()]); + let gate = q_and_gate + .narrow(2, 1, 1) + .squeeze(2) + .contiguous() + .reshape(&[seq_len, spec.full_gate_width()]); + let q_flat = query.reshape(&[seq_len * spec.num_heads, spec.head_dim]); + let k_flat = k.reshape(&[seq_len * spec.num_kv_heads, spec.head_dim]); + let q_normed = rmsnorm(&q_flat, &weights.q_norm, eps); + let k_normed = rmsnorm(&k_flat, &weights.k_norm, eps); + let q_rope = partial_rope_bf16( + &q_normed.reshape(&[seq_len, spec.num_heads, spec.head_dim]), + positions, + n_rot, + rope_theta, + ) + .reshape(&[seq_len * spec.num_heads, spec.head_dim]); + let k_rope = partial_rope_bf16( + &k_normed.reshape(&[seq_len, spec.num_kv_heads, spec.head_dim]), + positions, + n_rot, + rope_theta, + ) + .reshape(&[seq_len * spec.num_kv_heads, spec.head_dim]); + let q_attn = q_rope + .reshape(&[seq_len, spec.num_heads, spec.head_dim]) + .transpose(0, 1) + .unsqueeze(0) + .contiguous(); + let k_attn = k_rope + .reshape(&[seq_len, spec.num_kv_heads, spec.head_dim]) + .transpose(0, 1) + .unsqueeze(0) + .contiguous(); + let v_attn = v + .reshape(&[seq_len, spec.num_kv_heads, spec.head_dim]) + .transpose(0, 1) + .unsqueeze(0) + .contiguous(); + let n_rep = spec.num_heads / spec.num_kv_heads; + let k_full = repeat_kv_gpu(&k_attn, n_rep); + let v_full = repeat_kv_gpu(&v_attn, n_rep); + let attn_out = attention(&q_attn, &k_full, &v_full, true); + let attn_merged = merge_heads_gpu(&attn_out, seq_len, spec.num_heads, spec.head_dim); + let gate_sigmoid = sigmoid(&gate); + let gated_attn = mul(&attn_merged, &gate_sigmoid); + let projected = matmul(&gated_attn, &weights.o_w_t, GemmBackend::CuBlas); + + Qwen35FullAttentionOutput { + q_full, + query, + gate, + k, + v, + q_normed, + k_normed, + q_rope, + k_rope, + attn_out, + attn_merged, + gate_sigmoid, + gated_attn, + projected, + } + } + + pub fn from_weights_tp( + config: ModelConfig, + weights: HashMap, + _rank: usize, + _world: usize, + _device: u32, + _tp: Option>, + ) -> Self { + let spec = Qwen35MoeSpec::from_config(&config); + let language_tensors = weights + .keys() + .filter(|k| k.starts_with("model.language_model.")) + .count(); + let visual_tensors = weights + .keys() + .filter(|k| k.starts_with("model.visual.")) + .count(); + panic!( + "Qwen3.5/3.6 MoE is recognized but inference is not implemented yet: \ + layers={}, linear_layers={}, full_attention_layers={}, hidden={}, heads={}/{}, \ + experts={}, top_k={}, moe_intermediate={}, shared_intermediate={}, \ + loaded_tensors={} (language={}, visual={}). \ + Required next components: Gated DeltaNet linear attention \ + (qkv={}, z={}, a/b={}, conv=[{},1,{}], state head_dim={}), \ + Qwen full-attention output gating + IMRoPE (q_width={}, kv_width={}), \ + and routed+shared Qwen MoE.", + spec.layer_kinds.len(), + spec.linear_layers(), + spec.full_attention_layers(), + spec.hidden, + spec.num_heads, + spec.num_kv_heads, + spec.num_experts, + spec.experts_per_token, + spec.moe_intermediate, + spec.shared_intermediate, + weights.len(), + language_tensors, + visual_tensors, + spec.linear_qkv_width(), + spec.linear_value_width(), + spec.linear_value_heads, + spec.linear_conv_width(), + spec.linear_conv_kernel, + spec.linear_value_dim, + spec.full_q_width(), + spec.full_kv_width(), + ); + } +} + +impl Qwen35Moe { + /// Return whether a checkpoint tensor is needed by one pipeline stage. + /// Vision and MTP tensors are intentionally excluded from text generation. + pub fn weight_belongs_to_pp_stage( + config: &ModelConfig, + name: &str, + stage: usize, + num_stages: usize, + ) -> bool { + let num_layers = config.num_layers(); + assert_eq!(num_layers % num_stages, 0); + let per_stage = num_layers / num_stages; + let lo = stage * per_stage; + let hi = lo + per_stage; + if stage == 0 && name == "model.language_model.embed_tokens.weight" { + return true; + } + if stage + 1 == num_stages + && (name == "model.language_model.norm.weight" || name == "lm_head.weight") + { + return true; + } + let Some(rest) = name.strip_prefix("model.language_model.layers.") else { + return false; + }; + let Some(layer) = rest.split('.').next().and_then(|s| s.parse::().ok()) else { + return false; + }; + (lo..hi).contains(&layer) + } + + /// Pipeline-parallel text-model load. Every stage owns a contiguous layer + /// range; only stage 0 has embeddings and only the last stage has the head. + pub fn from_weights_pp( + config: ModelConfig, + mut w: HashMap, + stage: usize, + num_stages: usize, + device: u32, + ) -> Self { + crate::init_kernels(); + let spec = Qwen35MoeSpec::from_config(&config); + let num_layers = config.num_layers(); + assert!(num_stages >= 1); + assert_eq!(num_layers % num_stages, 0); + let per_stage = num_layers / num_stages; + let lo = stage * per_stage; + let hi = lo + per_stage; + let is_first_stage = stage == 0; + let is_last_stage = stage + 1 == num_stages; + let dev = Device::Cuda(device); + + let take = |w: &mut HashMap, name: &str| -> Tensor { + w.remove(name) + .unwrap_or_else(|| panic!("missing weight: {name}")) + }; + let repl = |t: Tensor| -> Tensor { t.to_device(dev) }; + let wt = |t: Tensor| -> Tensor { t.to_device(dev).transpose(0, 1).contiguous() }; + let placeholder = || Tensor::from_slice(&[bf16::ZERO], &[1, 1]).to_device(dev); + + let embed_tokens = if is_first_stage { + repl(take(&mut w, "model.language_model.embed_tokens.weight")) + } else { + placeholder() + }; + let norm = if is_last_stage { + repl(qwen35_norm_with_unit_offset(take( + &mut w, + "model.language_model.norm.weight", + ))) + } else { + placeholder() + }; + let lm_head_t = if is_last_stage { + wt(take(&mut w, "lm_head.weight")) + } else { + placeholder() + }; + + eprintln!( + "[qwen36-pp] stage {stage}/{num_stages}: layers [{lo}, {hi}) {}{}", + if is_first_stage { "+embed " } else { "" }, + if is_last_stage { "+norm+lm_head" } else { "" }, + ); + let mut layers = Vec::with_capacity(per_stage); + for i in lo..hi { + let p = format!("model.language_model.layers.{i}"); + let attention = match spec.layer_kinds[i] { + Qwen35LayerKind::LinearAttention => Qwen35AttentionWeights::Linear( + Self::take_linear_attention_projection_weights(&mut w, i, dev), + ), + Qwen35LayerKind::FullAttention => Qwen35AttentionWeights::Full( + Self::take_full_attention_weights(&mut w, i, dev), + ), + }; + let moe_p = format!("{p}.mlp"); + let moe = Qwen35MoeWeights { + router_w_t: wt(take(&mut w, &format!("{moe_p}.gate.weight"))), + expert_gate_up: repl(take( + &mut w, + &format!("{moe_p}.experts.gate_up_proj"), + )) + .contiguous(), + expert_down: repl(take(&mut w, &format!("{moe_p}.experts.down_proj"))) + .contiguous(), + shared_gate_w_t: wt(take( + &mut w, + &format!("{moe_p}.shared_expert.gate_proj.weight"), + )), + shared_up_w_t: wt(take( + &mut w, + &format!("{moe_p}.shared_expert.up_proj.weight"), + )), + shared_down_w_t: wt(take( + &mut w, + &format!("{moe_p}.shared_expert.down_proj.weight"), + )), + shared_router_w_t: wt(take( + &mut w, + &format!("{moe_p}.shared_expert_gate.weight"), + )), + }; + layers.push(Qwen35Block { + input_norm: repl(qwen35_norm_with_unit_offset(take( + &mut w, + &format!("{p}.input_layernorm.weight"), + ))), + post_norm: repl(qwen35_norm_with_unit_offset(take( + &mut w, + &format!("{p}.post_attention_layernorm.weight"), + ))), + attention, + moe, + }); + xserv_cuda::allocator::cached_trim(); + } + + if !w.is_empty() { + let mut unused: Vec<_> = w.keys().cloned().collect(); + unused.sort(); + eprintln!( + "[qwen36-pp] stage {stage}: {} filtered weight(s) were not consumed: {:?}", + unused.len(), + unused + ); + } + + Self { + config, + spec, + embed_tokens, + layers, + norm, + lm_head_t, + layer_start: lo, + is_first_stage, + is_last_stage, + } + } + + pub fn new_recurrent_cache(&self, max_seqs: usize) -> Qwen35RecurrentCache { + let layer_kinds = self.spec.layer_kinds + [self.layer_start..self.layer_start + self.layers.len()] + .to_vec(); + let mut cache = Qwen35RecurrentCache { + slots: (0..max_seqs).map(|_| Vec::new()).collect(), + layer_kinds, + conv_channels: self.spec.linear_conv_width(), + conv_history: self.spec.linear_conv_kernel - 1, + value_heads: self.spec.linear_value_heads, + value_dim: self.spec.linear_value_dim, + device: self.layers.first().map(|l| l.input_norm.device()).unwrap_or(Device::Cpu), + }; + for slot in 0..max_seqs { + cache.reset_slot(slot); + } + cache + } + + pub fn embed(&self, token_ids: &[u32]) -> Tensor { + debug_assert!(self.is_first_stage); + embedding(&self.embed_tokens, token_ids) + } + + pub fn head(&self, x: &Tensor) -> Tensor { + debug_assert!(self.is_last_stage); + let x = rmsnorm(x, &self.norm, self.config.ln_eps()); + matmul(&x, &self.lm_head_t, GemmBackend::CuBlas) + } + + pub fn forward_layers_prefill( + &self, + mut x: Tensor, + slot: usize, + paged_cache: &mut PagedKVCache, + recurrent_cache: &mut Qwen35RecurrentCache, + ) -> Tensor { + let n_tokens = x.shape()[0]; + let pos_offset = paged_cache.seq_len(slot); + paged_cache.ensure_capacity(slot, pos_offset + n_tokens); + // gather_kv_contiguous reads seq_len, so expose the new logical length + // before the per-layer prefill loop, matching the existing Qwen3 path. + paged_cache.advance_seq_len(slot, n_tokens); + let positions: Vec = (pos_offset..pos_offset + n_tokens) + .map(|p| p as u32) + .collect(); + + for (local_idx, layer) in self.layers.iter().enumerate() { + let residual = x.clone(); + let normed = rmsnorm(&x, &layer.input_norm, self.config.ln_eps()); + let attn = match &layer.attention { + Qwen35AttentionWeights::Linear(weights) => { + let state = recurrent_cache.linear_mut(slot, local_idx); + self.forward_linear_attention(&normed, weights, state) + } + Qwen35AttentionWeights::Full(weights) => self.forward_full_attention_prefill( + &normed, + weights, + &positions, + slot, + local_idx, + paged_cache, + ), + }; + let attn_residual = add(&residual, &attn); + let moe_in = rmsnorm(&attn_residual, &layer.post_norm, self.config.ln_eps()); + let moe_out = self.forward_moe(&moe_in, &layer.moe); + x = add(&attn_residual, &moe_out); + } + x + } + + /// Serial PP v1 decodes one sequence/token at a time. + pub fn forward_layers_decode( + &self, + mut x: Tensor, + slot: usize, + paged_cache: &mut PagedKVCache, + recurrent_cache: &mut Qwen35RecurrentCache, + ) -> Tensor { + assert_eq!(x.shape(), &[1, self.spec.hidden]); + let position = paged_cache.seq_len(slot); + let kv_lens = [(position + 1) as i32]; + paged_cache.ensure_capacity(slot, position + 1); + paged_cache.sync_active_batch_with_lens(&[slot], &kv_lens); + + for (local_idx, layer) in self.layers.iter().enumerate() { + let residual = x.clone(); + let normed = rmsnorm(&x, &layer.input_norm, self.config.ln_eps()); + let attn = match &layer.attention { + Qwen35AttentionWeights::Linear(weights) => { + let state = recurrent_cache.linear_mut(slot, local_idx); + self.forward_linear_attention(&normed, weights, state) + } + Qwen35AttentionWeights::Full(weights) => self.forward_full_attention_decode( + &normed, + weights, + position as u32, + slot, + local_idx, + paged_cache, + ), + }; + let attn_residual = add(&residual, &attn); + let moe_in = rmsnorm(&attn_residual, &layer.post_norm, self.config.ln_eps()); + let moe_out = self.forward_moe(&moe_in, &layer.moe); + x = add(&attn_residual, &moe_out); + } + paged_cache.advance_seq_len(slot, 1); + x + } + + fn forward_linear_attention( + &self, + input: &Tensor, + weights: &Qwen35LinearAttentionWeights, + state: &mut Qwen35LinearState, + ) -> Tensor { + let n_tokens = input.shape()[0]; + let qkv = matmul(input, &weights.in_proj_qkv_t, GemmBackend::CuBlas); + let z = matmul(input, &weights.in_proj_z_t, GemmBackend::CuBlas); + let a = matmul(input, &weights.in_proj_a_t, GemmBackend::CuBlas); + let b = matmul(input, &weights.in_proj_b_t, GemmBackend::CuBlas); + let beta = sigmoid(&b); + let alpha = softplus(&bias_add_2d(&a, &weights.dt_bias)); + let conv = depthwise_causal_conv1d_stateful_silu_bf16( + &qkv, + &weights.conv1d, + &mut state.conv, + ); + let key_width = self.spec.linear_key_width(); + let value_width = self.spec.linear_value_width(); + let q = conv.narrow(1, 0, key_width).contiguous(); + let k = conv.narrow(1, key_width, key_width).contiguous(); + let v = conv + .narrow(1, 2 * key_width, value_width) + .contiguous() + .reshape(&[ + n_tokens, + self.spec.linear_value_heads, + self.spec.linear_value_dim, + ]); + let q = l2_normalize_bf16( + &q.reshape(&[ + n_tokens * self.spec.linear_key_heads, + self.spec.linear_key_dim, + ]), + self.config.ln_eps(), + ) + .reshape(&[ + n_tokens, + self.spec.linear_key_heads, + self.spec.linear_key_dim, + ]); + let k = l2_normalize_bf16( + &k.reshape(&[ + n_tokens * self.spec.linear_key_heads, + self.spec.linear_key_dim, + ]), + self.config.ln_eps(), + ) + .reshape(&[ + n_tokens, + self.spec.linear_key_heads, + self.spec.linear_key_dim, + ]); + let recurrent = deltanet_recurrent_bf16_f32( + &q, + &k, + &v, + &beta, + &alpha, + &weights.a_log, + &mut state.recurrent, + ); + let normed = rmsnorm( + &recurrent.reshape(&[ + n_tokens * self.spec.linear_value_heads, + self.spec.linear_value_dim, + ]), + &weights.norm, + self.config.ln_eps(), + ) + .reshape(&[n_tokens, value_width]); + let gated = mul(&normed, &silu(&z)); + matmul(&gated, &weights.out_proj_t, GemmBackend::CuBlas) + } + + fn project_full_attention( + &self, + input: &Tensor, + weights: &Qwen35FullAttentionWeights, + positions: &[u32], + ) -> (Tensor, Tensor, Tensor, Tensor) { + let n_tokens = input.shape()[0]; + let q_full = matmul(input, &weights.q_w_t, GemmBackend::CuBlas); + // q_proj rows are interleaved per head: [query, gate]. + let q_and_gate = q_full.reshape(&[ + n_tokens, + self.spec.num_heads, + 2, + self.spec.head_dim, + ]); + let query = q_and_gate + .narrow(2, 0, 1) + .squeeze(2) + .contiguous(); + let gate = q_and_gate + .narrow(2, 1, 1) + .squeeze(2) + .contiguous() + .reshape(&[n_tokens, self.spec.full_gate_width()]); + let k = matmul(input, &weights.k_w_t, GemmBackend::CuBlas); + let v = matmul(input, &weights.v_w_t, GemmBackend::CuBlas); + let q = rmsnorm( + &query.reshape(&[n_tokens * self.spec.num_heads, self.spec.head_dim]), + &weights.q_norm, + self.config.ln_eps(), + ) + .reshape(&[n_tokens, self.spec.num_heads, self.spec.head_dim]); + let k = rmsnorm( + &k.reshape(&[ + n_tokens * self.spec.num_kv_heads, + self.spec.head_dim, + ]), + &weights.k_norm, + self.config.ln_eps(), + ) + .reshape(&[n_tokens, self.spec.num_kv_heads, self.spec.head_dim]); + let n_rot = (self.spec.head_dim as f64 + * self.config.text().partial_rotary_factor.unwrap_or(0.25)) + as usize; + let theta = self.config.rope_theta_value().unwrap_or(10_000_000.0) as f32; + // For text-only requests all three MRoPE position axes are identical, + // so interleaved MRoPE reduces to this partial RoPE operation. + let q = partial_rope_bf16(&q, positions, n_rot, theta) + .transpose(0, 1) + .contiguous() + .unsqueeze(0); + let k = partial_rope_bf16(&k, positions, n_rot, theta) + .transpose(0, 1) + .contiguous() + .unsqueeze(0); + let v = v + .reshape(&[n_tokens, self.spec.num_kv_heads, self.spec.head_dim]) + .transpose(0, 1) + .contiguous() + .unsqueeze(0); + (gate, q, k, v) + } + + fn finish_full_attention( + &self, + gate: &Tensor, + attn: &Tensor, + weights: &Qwen35FullAttentionWeights, + ) -> Tensor { + let n_tokens = gate.shape()[0]; + let merged = attn + .squeeze(0) + .transpose(0, 1) + .contiguous() + .reshape(&[n_tokens, self.spec.full_output_width()]); + let gated = mul(&merged, &sigmoid(gate)); + matmul(&gated, &weights.o_w_t, GemmBackend::CuBlas) + } + + fn forward_full_attention_prefill( + &self, + input: &Tensor, + weights: &Qwen35FullAttentionWeights, + positions: &[u32], + slot: usize, + layer: usize, + paged_cache: &mut PagedKVCache, + ) -> Tensor { + let n_tokens = input.shape()[0]; + let start = paged_cache.seq_len(slot) - n_tokens; + let (gate, q, k, v) = self.project_full_attention(input, weights, positions); + paged_cache.append_tokens(slot, layer, &k, &v, n_tokens, start); + let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer); + let attn = flash_attention(&q, &k_full, &v_full, true); + self.finish_full_attention(&gate, &attn, weights) + } + + fn forward_full_attention_decode( + &self, + input: &Tensor, + weights: &Qwen35FullAttentionWeights, + position: u32, + slot: usize, + layer: usize, + paged_cache: &mut PagedKVCache, + ) -> Tensor { + let (gate, q, k, v) = self.project_full_attention(input, weights, &[position]); + // append_tokens expects [1, kv_heads, tokens, dim], which project_full + // returns for the one-token decode case. + paged_cache.append_tokens(slot, layer, &k, &v, 1, position as usize); + let attn = paged_decode_attention( + &q, + paged_cache.k_pool(layer).as_ptr() as *const std::ffi::c_void, + paged_cache.v_pool(layer).as_ptr() as *const std::ffi::c_void, + paged_cache.block_table_gpu().as_ptr() as *const i32, + paged_cache.context_lens_gpu().as_ptr() as *const i32, + 1, + self.spec.num_heads, + self.spec.num_kv_heads, + self.spec.head_dim, + paged_cache.max_blocks_per_seq(), + ); + self.finish_full_attention(&gate, &attn, weights) + } + + fn forward_moe(&self, input: &Tensor, weights: &Qwen35MoeWeights) -> Tensor { + let n_tokens = input.shape()[0]; + let router_logits = matmul(input, &weights.router_w_t, GemmBackend::CuBlas); + let (topk_ids, topk_weights) = xserv_kernels::moe::moe_topk_softmax( + &router_logits, + self.spec.num_experts, + self.spec.experts_per_token, + ); + + let shared_gate = matmul(input, &weights.shared_gate_w_t, GemmBackend::CuBlas); + let shared_up = matmul(input, &weights.shared_up_w_t, GemmBackend::CuBlas); + let shared_hidden = silu_mul(&shared_gate, &shared_up); + let shared_down = matmul( + &shared_hidden, + &weights.shared_down_w_t, + GemmBackend::CuBlas, + ); + let shared_scale = sigmoid(&matmul( + input, + &weights.shared_router_w_t, + GemmBackend::CuBlas, + )); + let shared_out = row_scale_bf16(&shared_down, &shared_scale); + + let gate_up = xserv_kernels::moe::moe_sparse_gemv_bf16( + input, + &weights.expert_gate_up, + &topk_ids, + self.spec.experts_per_token, + 0, + self.spec.num_experts, + false, + ); + let gate = gate_up + .narrow(2, 0, self.spec.moe_intermediate) + .contiguous(); + let up = gate_up + .narrow( + 2, + self.spec.moe_intermediate, + self.spec.moe_intermediate, + ) + .contiguous(); + let routed_hidden = silu_mul(&gate, &up); + let down = xserv_kernels::moe::moe_sparse_gemv_bf16( + &routed_hidden.reshape(&[ + n_tokens * self.spec.experts_per_token, + self.spec.moe_intermediate, + ]), + &weights.expert_down, + &topk_ids, + self.spec.experts_per_token, + 0, + self.spec.num_experts, + true, + ); + let routed_out = xserv_kernels::moe::moe_weighted_sum_sparse( + &down, + &topk_ids, + &topk_weights, + 0, + self.spec.num_experts, + ); + add(&routed_out, &shared_out) + } +} diff --git a/crates/xserv-model/src/sampling.rs b/crates/xserv-model/src/sampling.rs index 2751b34..11f00f6 100644 --- a/crates/xserv-model/src/sampling.rs +++ b/crates/xserv-model/src/sampling.rs @@ -1,5 +1,6 @@ use half::bf16; use rand::Rng; +use std::collections::HashSet; use xserv_tensor::{DType, Device, Tensor}; #[derive(Clone)] @@ -7,6 +8,8 @@ pub struct SamplingParams { pub temperature: f32, pub top_k: usize, pub top_p: f32, + pub presence_penalty: f32, + pub repetition_penalty: f32, } impl Default for SamplingParams { @@ -15,6 +18,8 @@ impl Default for SamplingParams { temperature: 0.0, top_k: 0, top_p: 1.0, + presence_penalty: 0.0, + repetition_penalty: 1.0, } } } @@ -22,12 +27,21 @@ impl Default for SamplingParams { /// Sample a token from logits with shape [seq_len, vocab_size]. /// Uses the last position's logits. Handles both F32 and BF16 dtypes. pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 { + sample_with_history(logits, params, &[]) +} + +/// Sample while applying penalties to tokens already present in the request. +/// Presence penalty follows the OpenAI convention (subtract once per distinct +/// token); repetition penalty follows the HF convention. +pub fn sample_with_history(logits: &Tensor, params: &SamplingParams, history: &[u32]) -> u32 { assert_eq!(logits.ndim(), 2); // Greedy fast path: GPU argmax + 4-byte D2H instead of copying the whole // [seq, vocab] logits to the host and scanning it (~201k bf16/token). // NaN logits lose every `>` comparison in the kernel, matching the // NaN-safe host argmax below. if params.temperature == 0.0 + && params.presence_penalty == 0.0 + && params.repetition_penalty == 1.0 && logits.dtype() == DType::BF16 && matches!(logits.device(), Device::Cuda(_)) && logits.is_contiguous() @@ -55,11 +69,6 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 { _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()), }; - // Greedy - if params.temperature == 0.0 { - return argmax(&last_row); - } - // NaN-safe: sampling path uses partial_cmp().unwrap() in top-k/top-p // sorts and softmax; a single NaN logit would panic the engine thread. // Replace NaN with -inf (equivalent to masking) instead. @@ -74,6 +83,29 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 { eprintln!("[sampling] WARNING: NaN logits encountered in sample()"); } + if !history.is_empty() && (params.presence_penalty != 0.0 || params.repetition_penalty != 1.0) { + let seen: HashSet = history.iter().copied().collect(); + for id in seen { + let i = id as usize; + if i >= last_row.len() { + continue; + } + if params.repetition_penalty != 1.0 { + let value = last_row[i]; + last_row[i] = if value > 0.0 { + value / params.repetition_penalty + } else { + value * params.repetition_penalty + }; + } + last_row[i] -= params.presence_penalty; + } + } + + if params.temperature == 0.0 { + return argmax(&last_row); + } + // Apply temperature let mut logits_f32: Vec = last_row.iter().map(|v| v / params.temperature).collect(); @@ -195,3 +227,25 @@ fn argmax(data: &[f32]) -> u32 { } best_i as u32 } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn greedy_sampling_applies_history_penalties() { + let logits = Tensor::from_slice(&[10.0f32, 9.0, 0.0], &[1, 3]); + + let presence = SamplingParams { + presence_penalty: 2.0, + ..SamplingParams::default() + }; + assert_eq!(sample_with_history(&logits, &presence, &[0]), 1); + + let repetition = SamplingParams { + repetition_penalty: 2.0, + ..SamplingParams::default() + }; + assert_eq!(sample_with_history(&logits, &repetition, &[0]), 1); + } +} diff --git a/crates/xserv-server/src/api.rs b/crates/xserv-server/src/api.rs index c7a7d70..162ee6e 100644 --- a/crates/xserv-server/src/api.rs +++ b/crates/xserv-server/src/api.rs @@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize}; use std::convert::Infallible; use std::path::Path; use std::sync::Arc; +use std::sync::atomic::Ordering; use tokio_stream::StreamExt; use tokio_stream::wrappers::ReceiverStream; use uuid::Uuid; @@ -25,11 +26,35 @@ pub struct ChatRequest { #[serde(default)] pub stream: Option, #[serde(default)] + pub stream_options: Option, + #[serde(default)] pub temperature: Option, #[serde(default)] pub top_k: Option, #[serde(default)] pub top_p: Option, + #[serde(default)] + pub presence_penalty: Option, + #[serde(default)] + pub repetition_penalty: Option, + #[serde(default)] + pub chat_template_kwargs: Option, + #[serde(default)] + pub enable_thinking: Option, +} + +#[derive(Deserialize)] +pub struct StreamOptions { + #[serde(default)] + pub include_usage: bool, +} + +#[derive(Deserialize, Default)] +pub struct ChatTemplateKwargs { + #[serde(default)] + pub enable_thinking: Option, + #[serde(default)] + pub preserve_thinking: Option, } #[derive(Deserialize, Serialize, Clone)] @@ -102,12 +127,17 @@ impl ChatTemplate { } } - pub fn render(&self, messages: &[Message]) -> String { + pub fn render( + &self, + messages: &[Message], + enable_thinking: Option, + preserve_thinking: Option, + ) -> String { if self.source.is_empty() { return build_prompt_hardcoded(messages, &self.model_type); } - match self.render_jinja(messages) { + match self.render_jinja(messages, enable_thinking, preserve_thinking) { Ok(prompt) => prompt, Err(e) => { eprintln!("[chat-template] Jinja render error: {e}, falling back to hardcoded"); @@ -116,7 +146,12 @@ impl ChatTemplate { } } - fn render_jinja(&self, messages: &[Message]) -> Result { + fn render_jinja( + &self, + messages: &[Message], + enable_thinking: Option, + preserve_thinking: Option, + ) -> Result { let mut env = minijinja::Environment::new(); // Register custom functions the template may call. @@ -127,6 +162,42 @@ impl ChatTemplate { env.add_filter("startswith", |s: String, prefix: String| -> bool { s.starts_with(&prefix) }); + env.set_unknown_method_callback(|_, value, method, args| { + use minijinja::value::from_args; + use minijinja::{Error, ErrorKind, Value}; + + let Some(value) = value.as_str() else { + return Err(Error::from(ErrorKind::UnknownMethod)); + }; + match method { + "startswith" => { + let (prefix,): (String,) = from_args(args)?; + Ok(Value::from(value.starts_with(&prefix))) + } + "endswith" => { + let (suffix,): (String,) = from_args(args)?; + Ok(Value::from(value.ends_with(&suffix))) + } + "split" => { + let (separator,): (String,) = from_args(args)?; + Ok(Value::from_serialize( + value + .split(&separator) + .map(str::to_owned) + .collect::>(), + )) + } + "rstrip" => { + let (chars,): (String,) = from_args(args)?; + Ok(Value::from(value.trim_end_matches(|c| chars.contains(c)))) + } + "lstrip" => { + let (chars,): (String,) = from_args(args)?; + Ok(Value::from(value.trim_start_matches(|c| chars.contains(c)))) + } + _ => Err(Error::from(ErrorKind::UnknownMethod)), + } + }); env.add_template("chat", &self.source)?; let tmpl = env.get_template("chat")?; @@ -136,6 +207,8 @@ impl ChatTemplate { add_generation_prompt => true, bos_token => "", eos_token => "", + enable_thinking => enable_thinking, + preserve_thinking => preserve_thinking, }; tmpl.render(ctx) @@ -256,8 +329,12 @@ fn build_prompt_gpt_oss(messages: &[Message]) -> String { // HTTP handlers // --------------------------------------------------------------------------- -pub async fn health() -> &'static str { - "ok" +pub async fn health(Extension(state): Extension>) -> Response { + if state.engine_ready.load(Ordering::Acquire) { + (StatusCode::OK, "ok").into_response() + } else { + (StatusCode::SERVICE_UNAVAILABLE, "loading").into_response() + } } pub async fn list_models(Extension(state): Extension>) -> Json { @@ -291,7 +368,10 @@ async fn chat_non_stream(state: Arc, req: ChatRequest) -> Response { return response; } - let prompt = state.chat_template.render(&req.messages); + let (enable_thinking, preserve_thinking) = template_options(&req); + let prompt = state + .chat_template + .render(&req.messages, enable_thinking, preserve_thinking); let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt); let prompt_token_count = prompt_tokens.len(); @@ -363,18 +443,26 @@ fn chat_stream(state: Arc, req: ChatRequest) -> Response { return response; } - let prompt = state.chat_template.render(&req.messages); + let (enable_thinking, preserve_thinking) = template_options(&req); + let prompt = state + .chat_template + .render(&req.messages, enable_thinking, preserve_thinking); let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt); + let prompt_token_count = prompt_tokens.len(); + let include_usage = req + .stream_options + .as_ref() + .is_some_and(|options| options.include_usage); let max_seq_len = state.max_seq_len; - if prompt_tokens.len() >= max_seq_len { + if prompt_token_count >= max_seq_len { return bad_request(format!( "prompt is {} tokens, exceeds max_seq_len {}", - prompt_tokens.len(), + prompt_token_count, max_seq_len )); } - let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len()); + let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count); let (engine_tx, engine_rx) = tokio::sync::mpsc::channel::(64); let gen_req = GenerateRequest { @@ -393,10 +481,12 @@ fn chat_stream(state: Arc, req: ChatRequest) -> Response { tokio::spawn(async move { let mut engine_stream = ReceiverStream::new(engine_rx); let mut first = true; + let mut completion_token_count = 0usize; while let Some(event) = engine_stream.next().await { match event { GenerateEvent::Token { text, .. } => { + completion_token_count += 1; if first { // First chunk: role announcement let chunk = @@ -422,6 +512,16 @@ fn chat_stream(state: Arc, req: ChatRequest) -> Response { let fr = normalize_finish_reason(&finish_reason); let chunk = make_chunk(&id, &model_name, created, None, None, fr); let _ = sse_tx.send(Ok(Event::default().data(chunk))).await; + if include_usage { + let usage = make_usage_chunk( + &id, + &model_name, + created, + prompt_token_count, + completion_token_count, + ); + let _ = sse_tx.send(Ok(Event::default().data(usage))).await; + } let _ = sse_tx .send(Ok(Event::default().data("[DONE]".to_string()))) .await; @@ -464,6 +564,18 @@ fn validate_request(req: &ChatRequest, model_name: &str) -> Option { return Some(bad_request("top_k must be <= 1_000_000")); } } + if let Some(p) = req.presence_penalty { + if !p.is_finite() || !(-2.0..=2.0).contains(&p) { + return Some(bad_request("presence_penalty must be in [-2, 2]")); + } + } + if let Some(p) = req.repetition_penalty { + if !p.is_finite() || p <= 0.0 { + return Some(bad_request( + "repetition_penalty must be a finite value greater than 0", + )); + } + } None } @@ -546,6 +658,28 @@ fn make_chunk( .to_string() } +fn make_usage_chunk( + id: &str, + model: &str, + created: u64, + prompt_tokens: usize, + completion_tokens: usize, +) -> String { + serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + }) + .to_string() +} + fn unix_timestamp() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -558,9 +692,20 @@ fn sampling_params(req: &ChatRequest) -> SamplingParams { temperature: req.temperature.unwrap_or(0.0), top_k: req.top_k.unwrap_or(0), top_p: req.top_p.unwrap_or(1.0), + presence_penalty: req.presence_penalty.unwrap_or(0.0), + repetition_penalty: req.repetition_penalty.unwrap_or(1.0), } } +fn template_options(req: &ChatRequest) -> (Option, Option) { + let kwargs = req.chat_template_kwargs.as_ref(); + ( + req.enable_thinking + .or_else(|| kwargs.and_then(|value| value.enable_thinking)), + kwargs.and_then(|value| value.preserve_thinking), + ) +} + /// Map engine finish_reason strings to OpenAI-standard values. Any engine-internal /// code (e.g. "error" from tp/pp client-stall) collapses to None so SDK clients see /// a clean null instead of an unknown value. diff --git a/crates/xserv-server/src/engine.rs b/crates/xserv-server/src/engine.rs index 253b072..7fa5dab 100644 --- a/crates/xserv-server/src/engine.rs +++ b/crates/xserv-server/src/engine.rs @@ -4,7 +4,9 @@ use std::sync::Once; use std::sync::mpsc; use std::time::Instant; use xserv_model::loader; -use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample}; +use xserv_model::{ + BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample_with_history, +}; use xserv_tensor::{DType, Device}; use xserv_tokenizer::Tokenizer; @@ -232,7 +234,7 @@ impl Engine { slot, &mut self.paged_cache, ); - let next = sample(&logits, &seq.sampling); + let next = sample_with_history(&logits, &seq.sampling, &seq.prompt_tokens); seq.generated_tokens.push(next); seq.prefilled = true; emit_token(&self.tokenizer, seq, next); @@ -310,7 +312,11 @@ impl Engine { // (~1.2 MB for B=4, Qwen3 vocab=152K). let all_greedy = decode_indices .iter() - .all(|&i| running[i].sampling.temperature == 0.0); + .all(|&i| { + running[i].sampling.temperature == 0.0 + && running[i].sampling.presence_penalty == 0.0 + && running[i].sampling.repetition_penalty == 1.0 + }); if all_greedy { let next_ids = xserv_kernels::argmax_bf16_to_host(&logits); for (j, &i) in decode_indices.iter().enumerate() { @@ -329,7 +335,10 @@ impl Engine { for (j, &i) in decode_indices.iter().enumerate() { let row_start = j * vocab_size; let row_logits = &data[row_start..row_start + vocab_size]; - let next = if running[i].sampling.temperature == 0.0 { + let unpenalized_greedy = running[i].sampling.temperature == 0.0 + && running[i].sampling.presence_penalty == 0.0 + && running[i].sampling.repetition_penalty == 1.0; + let next = if unpenalized_greedy { row_logits .iter() .enumerate() @@ -339,7 +348,9 @@ impl Engine { } else { let row_tensor = xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]); - sample(&row_tensor, &running[i].sampling) + let mut history = running[i].prompt_tokens.clone(); + history.extend_from_slice(&running[i].generated_tokens); + sample_with_history(&row_tensor, &running[i].sampling, &history) }; running[i].generated_tokens.push(next); emit_token(&self.tokenizer, &mut running[i], next); diff --git a/crates/xserv-server/src/main.rs b/crates/xserv-server/src/main.rs index 2610a4c..aff2904 100644 --- a/crates/xserv-server/src/main.rs +++ b/crates/xserv-server/src/main.rs @@ -10,8 +10,9 @@ use axum::{ }; use engine::GenerateRequest; use std::path::PathBuf; +use std::sync::atomic::AtomicBool; use std::sync::{Arc, Mutex, mpsc}; -use xserv_model::ModelConfig; +use xserv_model::{ModelConfig, ModelFamily}; pub struct AppState { pub model_name: String, @@ -19,6 +20,7 @@ pub struct AppState { pub engine_sender: Mutex>, pub engine_tokenizer: Mutex, pub max_seq_len: usize, + pub engine_ready: Arc, } #[tokio::main] @@ -77,12 +79,18 @@ async fn main() { std::process::exit(1); } let model_config = ModelConfig::from_file(&model_dir.join("config.json")); - // gpt-oss is only implemented in the TP engine; route it there even at - // tp=1 (single-rank world) so quantized models can serve on one GPU. - let is_gpt_oss = model_config.model_type.as_deref() == Some("gpt_oss"); - if pp > 1 && is_gpt_oss { + let model_family = model_config.family(); + if model_family == ModelFamily::Qwen35Moe && pp <= 1 { eprintln!( - "gpt-oss is not supported by the pipeline-parallel engine (Qwen3 only); use --tp instead" + "Qwen3.5/3.6 MoE model '{}' currently requires pipeline parallelism; use --pp 8 on dash5", + model_config.model_type_str() + ); + std::process::exit(1); + } + if pp > 1 && !matches!(model_family, ModelFamily::Qwen3 | ModelFamily::Qwen35Moe) { + eprintln!( + "{} is not supported by the pipeline-parallel engine; use --tp instead", + model_family.as_str() ); std::process::exit(1); } @@ -109,27 +117,42 @@ async fn main() { // behind, instead of letting them pile up in RAM. try_send in the API // handler surfaces this as 503 to the client. let (tx, rx) = mpsc::sync_channel::(256); + let engine_ready = Arc::new(AtomicBool::new(false)); let model_dir_clone = model_dir.clone(); + let engine_ready_clone = Arc::clone(&engine_ready); std::thread::spawn(move || { if pp > 1 { // Pipeline-parallel path: stage-0 coordinator + worker stage threads. - pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx); - } else if tp <= 1 && !is_gpt_oss { + pp_engine::run_pp( + &model_dir_clone, + pp, + max_seq_len, + rx, + engine_ready_clone, + ); + } else if tp <= 1 && model_family == ModelFamily::Qwen3 { let mut engine = engine::Engine::load_with_swap( &model_dir_clone, max_batch, max_seq_len, swap_space_gb, ); + engine_ready_clone.store(true, std::sync::atomic::Ordering::Release); engine.run(rx); } else { // Tensor-parallel path: rank-0 coordinator + worker rank threads. - tp_engine::run_tp(&model_dir_clone, tp, max_seq_len, rx); + tp_engine::run_tp( + &model_dir_clone, + tp, + max_seq_len, + rx, + engine_ready_clone, + ); } }); - let model_type = model_config.model_type.clone().unwrap_or_default(); + let model_type = model_config.model_type_str().to_string(); let chat_template = api::ChatTemplate::load(&model_dir, &model_type); let state = Arc::new(AppState { model_name, @@ -137,6 +160,7 @@ async fn main() { engine_sender: Mutex::new(tx), engine_tokenizer: Mutex::new(tokenizer), max_seq_len, + engine_ready, }); let app = Router::new() diff --git a/crates/xserv-server/src/pp_engine.rs b/crates/xserv-server/src/pp_engine.rs index 42615cb..9a7b23e 100644 --- a/crates/xserv-server/src/pp_engine.rs +++ b/crates/xserv-server/src/pp_engine.rs @@ -16,14 +16,18 @@ use std::ffi::c_void; use std::path::{Path, PathBuf}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; use std::thread; use half::bf16; use xserv_distributed::{PpContext, UniqueId}; use xserv_model::loader; -use xserv_model::sampling::SamplingParams; -use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, sample}; +use xserv_model::sampling::{SamplingParams, sample_with_history}; +use xserv_model::{ + BLOCK_SIZE, ModelConfig, ModelFamily, PagedKVCache, Qwen3, Qwen35Moe, + Qwen35RecurrentCache, +}; use xserv_tensor::{DType, Device, Tensor}; use xserv_tokenizer::Tokenizer; @@ -39,7 +43,7 @@ enum PpCommand { /// Receive `[n_tokens, hidden]` from the previous stage, run this stage's /// layers; if last stage, sample with `sampling` and return the token. Prefill { - n_tokens: usize, + tokens: Vec, slot: usize, sampling: SamplingParams, }, @@ -52,13 +56,69 @@ enum PpCommand { } struct StageCtx { - model: Qwen3, + model: PpModel, + qwen35_recurrent: Option, cache: PagedKVCache, pp: Arc, hidden: usize, device: u32, } +enum PpModel { + Qwen3(Qwen3), + Qwen35(Qwen35Moe), +} + +impl StageCtx { + fn reset_slot(&mut self, slot: usize) { + if let Some(cache) = &mut self.qwen35_recurrent { + cache.reset_slot(slot); + } + } + + fn embed(&self, tokens: &[u32]) -> Tensor { + match &self.model { + PpModel::Qwen3(model) => model.embed(tokens), + PpModel::Qwen35(model) => model.embed(tokens), + } + } + + fn head(&self, x: &Tensor) -> Tensor { + match &self.model { + PpModel::Qwen3(model) => model.head(x), + PpModel::Qwen35(model) => model.head(x), + } + } + + fn forward_prefill(&mut self, x: Tensor, slot: usize) -> Tensor { + match &self.model { + PpModel::Qwen3(model) => model.forward_layers_prefill(x, slot, &mut self.cache), + PpModel::Qwen35(model) => model.forward_layers_prefill( + x, + slot, + &mut self.cache, + self.qwen35_recurrent + .as_mut() + .expect("Qwen3.6 recurrent cache"), + ), + } + } + + fn forward_decode(&mut self, x: Tensor, slot: usize) -> Tensor { + match &self.model { + PpModel::Qwen3(model) => model.forward_layers_decode(x, &[slot], &mut self.cache), + PpModel::Qwen35(model) => model.forward_layers_decode( + x, + slot, + &mut self.cache, + self.qwen35_recurrent + .as_mut() + .expect("Qwen3.6 recurrent cache"), + ), + } + } +} + /// Build this stage: NCCL init, load + slice weights, size a per-stage KV pool /// for THIS stage's layers only (so per-GPU KV is ~1/P). fn build_stage( @@ -71,14 +131,40 @@ fn build_stage( id: UniqueId, ) -> StageCtx { let pp = Arc::new(PpContext::init(stage, world, id, device)); - let weights = loader::load_model_dir(model_dir, Device::Cpu); - let model = Qwen3::from_weights_pp(config.clone(), weights, stage, world, device); + let model = match config.family() { + ModelFamily::Qwen35Moe => { + let weights = loader::load_model_dir_filtered(model_dir, Device::Cpu, |name| { + Qwen35Moe::weight_belongs_to_pp_stage(config, name, stage, world) + }); + PpModel::Qwen35(Qwen35Moe::from_weights_pp( + config.clone(), + weights, + stage, + world, + device, + )) + } + _ => { + let weights = loader::load_model_dir(model_dir, Device::Cpu); + PpModel::Qwen3(Qwen3::from_weights_pp( + config.clone(), + weights, + stage, + world, + device, + )) + } + }; // The KV cache only needs this stage's layers; build it from a config clone // whose layer count is the per-stage count (heads are NOT split under PP). let per_stage = config.num_layers() / world; let mut stage_config = config.clone(); - stage_config.num_hidden_layers = Some(per_stage); + if let Some(text) = stage_config.text_config.as_mut() { + text.num_hidden_layers = Some(per_stage); + } else { + stage_config.num_hidden_layers = Some(per_stage); + } let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE); let total_blocks = max_blocks_per_seq + 8; // v1 serial: one active sequence @@ -91,8 +177,13 @@ fn build_stage( DType::BF16, device, ); + let qwen35_recurrent = match &model { + PpModel::Qwen35(model) => Some(model.new_recurrent_cache(4)), + PpModel::Qwen3(_) => None, + }; StageCtx { model, + qwen35_recurrent, cache, pp, hidden: config.hidden(), @@ -138,40 +229,51 @@ fn worker_loop( max_seq_len, id, ); + let _ = ack_tx.send(()); let is_last = stage == world - 1; + let mut token_history = vec![Vec::::new(); 4]; let prev = stage - 1; let next = stage + 1; while let Ok(cmd) = cmd_rx.recv() { match cmd { PpCommand::Register(slot) => { + token_history[slot].clear(); + sc.reset_slot(slot); let _ = sc.cache.register_sequence(slot); let _ = ack_tx.send(()); } PpCommand::Free(slot) => { + token_history[slot].clear(); sc.cache.free_sequence(slot); let _ = ack_tx.send(()); } PpCommand::Prefill { - n_tokens, + tokens, slot, sampling, } => { + let n_tokens = tokens.len(); + token_history[slot] = tokens; let x = recv_hidden(&sc, n_tokens, prev); - let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache); + let x = sc.forward_prefill(x, slot); if is_last { - let logits = sc.model.head(&x); - let _ = token_tx.send(sample(&logits, &sampling)); + let logits = sc.head(&x); + let token = sample_with_history(&logits, &sampling, &token_history[slot]); + token_history[slot].push(token); + let _ = token_tx.send(token); } else { send_hidden(&sc, &x, next); } } PpCommand::Decode { slot, sampling } => { let x = recv_hidden(&sc, 1, prev); - let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache); + let x = sc.forward_decode(x, slot); if is_last { - let logits = sc.model.head(&x); - let _ = token_tx.send(sample(&logits, &sampling)); + let logits = sc.head(&x); + let token = sample_with_history(&logits, &sampling, &token_history[slot]); + token_history[slot].push(token); + let _ = token_tx.send(token); } else { send_hidden(&sc, &x, next); } @@ -191,6 +293,7 @@ pub fn run_pp( world: usize, max_seq_len: usize, rx: mpsc::Receiver, + ready: Arc, ) { assert!(world >= 2, "run_pp requires world >= 2"); let config = ModelConfig::from_file(&model_dir.join("config.json")); @@ -231,6 +334,10 @@ pub fn run_pp( // Stage 0 (this thread): coordinator + embedding + first layers. let mut sc = build_stage(model_dir, &config, 0, world, 0, max_seq_len, id); + for _ in 1..world { + ack_rx.recv().expect("PP worker exited during model load"); + } + ready.store(true, Ordering::Release); eprintln!("[pp-engine] ready (pp={world}, max_seq_len={max_seq_len})"); let n_workers = world - 1; @@ -249,6 +356,7 @@ pub fn run_pp( let slot = 0usize; while let Ok(req) = rx.recv() { broadcast(&cmd_txs, PpCommand::Register(slot)); + sc.reset_slot(slot); sc.cache.register_sequence(slot).expect("register slot"); wait_acks(&ack_rx); @@ -256,13 +364,13 @@ pub fn run_pp( broadcast( &cmd_txs, PpCommand::Prefill { - n_tokens: req.prompt_tokens.len(), + tokens: req.prompt_tokens.clone(), slot, sampling: req.sampling.clone(), }, ); - let x = sc.model.embed(&req.prompt_tokens); - let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache); + let x = sc.embed(&req.prompt_tokens); + let x = sc.forward_prefill(x, slot); send_hidden(&sc, &x, next_peer); let mut next = token_rx.recv().expect("prefill token"); @@ -287,8 +395,8 @@ pub fn run_pp( sampling: req.sampling.clone(), }, ); - let x = sc.model.embed(&[next]); - let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache); + let x = sc.embed(&[next]); + let x = sc.forward_decode(x, slot); send_hidden(&sc, &x, next_peer); next = token_rx.recv().expect("decode token"); generated += 1; diff --git a/crates/xserv-server/src/tp_engine.rs b/crates/xserv-server/src/tp_engine.rs index 29c3f73..bed534c 100644 --- a/crates/xserv-server/src/tp_engine.rs +++ b/crates/xserv-server/src/tp_engine.rs @@ -14,14 +14,15 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; use std::thread; use xserv_distributed::{TpContext, UniqueId}; use xserv_model::loader; use xserv_model::{ - BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, sample, - sample_greedy_penalized, + BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, ModelFamily, PagedKVCache, Qwen3, + sample_greedy_penalized, sample_with_history, }; use xserv_tensor::{DType, Device, Tensor}; use xserv_tokenizer::Tokenizer; @@ -105,24 +106,26 @@ fn build_rank( tp: Option>, ) -> RankCtx { let weights = loader::load_model_dir(model_dir, Device::Cpu); - let model = if config.is_moe() { - TpModel::GptOss(GptOss::from_weights_tp( + let model = match config.family() { + ModelFamily::GptOss => TpModel::GptOss(GptOss::from_weights_tp( config.clone(), weights, rank, world, device, tp, - )) - } else { - TpModel::Qwen3(Qwen3::from_weights_tp( + )), + ModelFamily::Qwen35Moe => { + panic!("Qwen3.5/3.6 MoE reached TP engine before inference support was implemented") + } + _ => TpModel::Qwen3(Qwen3::from_weights_tp( config.clone(), weights, rank, world, device, tp, - )) + )), }; let local_kv = config.num_kv_heads() / world; let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE; @@ -164,6 +167,7 @@ fn worker_loop( max_seq_len, Some(tp), ); + let _ = ack_tx.send(()); while let Ok(cmd) = cmd_rx.recv() { match cmd { TpCommand::Register(slot) => { @@ -196,6 +200,7 @@ pub fn run_tp( world: usize, max_seq_len: usize, rx: mpsc::Receiver, + ready: Arc, ) { // world=1 is a valid single-rank configuration (gpt-oss has no // single-GPU engine path; NCCL init and all_reduce no-op at world=1). @@ -235,6 +240,10 @@ pub fn run_tp( // Rank 0 (this thread). let tp = Arc::new(TpContext::init(0, world, id, 0)); let mut rc = build_rank(model_dir, &config, 0, world, 0, max_seq_len, Some(tp)); + for _ in 1..world { + ack_rx.recv().expect("TP worker exited during model load"); + } + ready.store(true, Ordering::Release); eprintln!("[tp-engine] ready (tp={world}, max_seq_len={max_seq_len})"); // Optional repetition penalty to break greedy repetition loops (reasoning @@ -254,7 +263,7 @@ pub fn run_tp( let start = history.len().saturating_sub(rep_window); sample_greedy_penalized(logits, &history[start..], rep_penalty) } else { - sample(logits, sp) + sample_with_history(logits, sp, history) } }; @@ -288,9 +297,9 @@ pub fn run_tp( .model .forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache); wait_acks(&ack_rx); - let mut gen_ids: Vec = Vec::new(); - let mut next = pick(&logits, &req.sampling, &gen_ids); - gen_ids.push(next); + let mut token_history = req.prompt_tokens.clone(); + let mut next = pick(&logits, &req.sampling, &token_history); + token_history.push(next); let mut decode_buf: Vec = Vec::new(); let mut generated = 1usize; @@ -317,8 +326,8 @@ pub fn run_tp( ); let logits = rank_decode(&mut rc, &[next], &[pos], &[slot]); wait_acks(&ack_rx); - next = pick(&logits, &req.sampling, &gen_ids); - gen_ids.push(next); + next = pick(&logits, &req.sampling, &token_history); + token_history.push(next); generated += 1; stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf); }; diff --git a/csrc/activation/activations.cu b/csrc/activation/activations.cu index fc9c672..92523f3 100644 --- a/csrc/activation/activations.cu +++ b/csrc/activation/activations.cu @@ -36,6 +36,35 @@ __global__ void silu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) { if (idx < n) out[idx] = __float2bfloat16(silu_f(__bfloat162float(x[idx]))); } +__global__ void sigmoid_f32(const float* x, float* out, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) out[idx] = 1.0f / (1.0f + expf(-x[idx])); +} + +__global__ void sigmoid_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + float v = __bfloat162float(x[idx]); + out[idx] = __float2bfloat16(1.0f / (1.0f + expf(-v))); + } +} + +__global__ void softplus_f32(const float* x, float* out, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + float v = x[idx]; + out[idx] = log1pf(expf(-fabsf(v))) + fmaxf(v, 0.0f); + } +} + +__global__ void softplus_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + float v = __bfloat162float(x[idx]); + out[idx] = __float2bfloat16(log1pf(expf(-fabsf(v))) + fmaxf(v, 0.0f)); + } +} + __global__ void scale_f32_kernel(const float* x, float* out, float scale, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) out[idx] = x[idx] * scale; @@ -108,6 +137,21 @@ __global__ void mul_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b, if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) * __bfloat162float(b[idx])); } +__global__ void row_scale_bf16_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ scale, + __nv_bfloat16* __restrict__ out, + int rows, + int cols +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = rows * cols; + if (idx >= total) return; + int row = idx / cols; + float v = __bfloat162float(x[idx]) * __bfloat162float(scale[row]); + out[idx] = __float2bfloat16(v); +} + extern "C" { void launch_gelu_f32(const void* x, void* out, int n, void* stream) { @@ -140,6 +184,36 @@ void launch_silu_bf16(const void* x, void* out, int n, void* stream) { CUDA_CHECK_LAST_ERROR(); } +void launch_sigmoid_f32(const void* x, void* out, int n, void* stream) { + int block = 256; + int grid = (n + block - 1) / block; + sigmoid_f32<<>>((const float*)x, (float*)out, n); + CUDA_CHECK_LAST_ERROR(); +} + +void launch_sigmoid_bf16(const void* x, void* out, int n, void* stream) { + int block = 256; + int grid = (n + block - 1) / block; + sigmoid_bf16<<>>( + (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n); + CUDA_CHECK_LAST_ERROR(); +} + +void launch_softplus_f32(const void* x, void* out, int n, void* stream) { + int block = 256; + int grid = (n + block - 1) / block; + softplus_f32<<>>((const float*)x, (float*)out, n); + CUDA_CHECK_LAST_ERROR(); +} + +void launch_softplus_bf16(const void* x, void* out, int n, void* stream) { + int block = 256; + int grid = (n + block - 1) / block; + softplus_bf16<<>>( + (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n); + CUDA_CHECK_LAST_ERROR(); +} + void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream) { int block = 256; int grid = (n + block - 1) / block; @@ -193,6 +267,15 @@ void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* strea CUDA_CHECK_LAST_ERROR(); } +void launch_row_scale_bf16(const void* x, const void* scale, void* out, int rows, int cols, void* stream) { + int n = rows * cols; + int block = 256; + int grid = (n + block - 1) / block; + row_scale_bf16_kernel<<>>( + (const __nv_bfloat16*)x, (const __nv_bfloat16*)scale, (__nv_bfloat16*)out, rows, cols); + CUDA_CHECK_LAST_ERROR(); +} + void launch_silu_mul_bf16(const void* gate, const void* up, void* out, int n, void* stream) { int block = 256; int grid = (n + block - 1) / block; diff --git a/csrc/attention/deltanet.cu b/csrc/attention/deltanet.cu new file mode 100644 index 0000000..7448192 --- /dev/null +++ b/csrc/attention/deltanet.cu @@ -0,0 +1,349 @@ +#include +#include +#include "../common.cuh" + +__global__ void depthwise_causal_conv1d_silu_bf16_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ out, + int seq_len, + int channels, + int kernel +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = seq_len * channels; + if (idx >= total) return; + + int c = idx % channels; + int t = idx / channels; + float acc = 0.0f; + for (int k = 0; k < kernel; ++k) { + int src_t = t - (kernel - 1 - k); + if (src_t >= 0) { + float xv = __bfloat162float(x[src_t * channels + c]); + float wv = __bfloat162float(w[c * kernel + k]); + acc += xv * wv; + } + } + float y = acc / (1.0f + expf(-acc)); + out[idx] = __float2bfloat16(y); +} + +__global__ void deltanet_ar_bf16_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + const __nv_bfloat16* __restrict__ beta, + const __nv_bfloat16* __restrict__ alpha_softplus, + const __nv_bfloat16* __restrict__ a_log, + __nv_bfloat16* __restrict__ state, + __nv_bfloat16* __restrict__ out, + int key_heads, + int value_heads, + int head_dim +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = value_heads * head_dim; + if (idx >= total) return; + + int h = idx / head_dim; + int j = idx % head_dim; + int key_h = h * key_heads / value_heads; + + float beta_h = __bfloat162float(beta[h]); + float gate = expf(__bfloat162float(alpha_softplus[h]) * __bfloat162float(a_log[h])); + + float sk = 0.0f; + for (int i = 0; i < head_dim; ++i) { + float s = __bfloat162float(state[(h * head_dim + i) * head_dim + j]); + float ki = __bfloat162float(k[key_h * head_dim + i]); + s *= gate; + state[(h * head_dim + i) * head_dim + j] = __float2bfloat16(s); + sk += s * ki; + } + + float vj = __bfloat162float(v[h * head_dim + j]); + float d = (vj - sk) * beta_h; + + float out_j = 0.0f; + for (int i = 0; i < head_dim; ++i) { + float ki = __bfloat162float(k[key_h * head_dim + i]); + float qi = __bfloat162float(q[key_h * head_dim + i]) / sqrtf((float)head_dim); + int sidx = (h * head_dim + i) * head_dim + j; + float s = __bfloat162float(state[sidx]) + ki * d; + state[sidx] = __float2bfloat16(s); + out_j += s * qi; + } + out[h * head_dim + j] = __float2bfloat16(out_j); +} + +// Stateful depthwise causal convolution for hybrid recurrent attention. +// `state` stores the preceding `kernel - 1` inputs in chronological order. +// One thread owns a channel, so the state update is race-free for any seq_len. +__global__ void depthwise_causal_conv1d_stateful_silu_bf16_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ state, + __nv_bfloat16* __restrict__ out, + int seq_len, + int channels, + int kernel +) { + int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= channels) return; + + const int history = kernel - 1; + for (int t = 0; t < seq_len; ++t) { + float acc = 0.0f; + for (int k = 0; k < kernel; ++k) { + int src_t = t - history + k; + float xv; + if (src_t < 0) { + xv = __bfloat162float(state[c * history + history + src_t]); + } else { + xv = __bfloat162float(x[(long long)src_t * channels + c]); + } + acc += xv * __bfloat162float(w[c * kernel + k]); + } + out[(long long)t * channels + c] = + __float2bfloat16(acc / (1.0f + expf(-acc))); + } + + // Keep the last `history` raw projection values for the next call. + for (int h = 0; h < history; ++h) { + int src_t = seq_len - history + h; + __nv_bfloat16 value; + if (src_t < 0) { + value = state[c * history + history + src_t]; + } else { + value = x[(long long)src_t * channels + c]; + } + state[c * history + h] = value; + } +} + +// L2-normalize each row using FP32 accumulation. Qwen3.5/3.6 applies this to +// convolved Q and K before the delta rule (this is not RMSNorm). +__global__ void l2_normalize_rows_bf16_kernel( + const __nv_bfloat16* __restrict__ x, + __nv_bfloat16* __restrict__ out, + int rows, + int cols, + float eps +) { + int row = blockIdx.x; + if (row >= rows) return; + + float local = 0.0f; + for (int i = threadIdx.x; i < cols; i += blockDim.x) { + float v = __bfloat162float(x[(long long)row * cols + i]); + local += v * v; + } + __shared__ float sums[256]; + sums[threadIdx.x] = local; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) sums[threadIdx.x] += sums[threadIdx.x + stride]; + __syncthreads(); + } + // Match torch.nn.functional.normalize / llama.cpp: clamp the norm by eps, + // rather than adding eps to every non-zero norm. + float inv = rsqrtf(fmaxf(sums[0], eps * eps)); + for (int i = threadIdx.x; i < cols; i += blockDim.x) { + out[(long long)row * cols + i] = + __float2bfloat16(__bfloat162float(x[(long long)row * cols + i]) * inv); + } +} + +// Recurrent Gated DeltaNet over one or more tokens. HF stores V heads grouped +// by K head, so V head h uses K/Q head floor(h / (value_heads/key_heads)). +// The recurrent matrix is FP32 as required by `mamba_ssm_dtype=float32`. +__global__ void deltanet_recurrent_bf16_f32_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + const __nv_bfloat16* __restrict__ beta, + const __nv_bfloat16* __restrict__ alpha_softplus, + const __nv_bfloat16* __restrict__ a_log, + float* __restrict__ state, + __nv_bfloat16* __restrict__ out, + int seq_len, + int key_heads, + int value_heads, + int head_dim +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = value_heads * head_dim; + if (idx >= total) return; + + int h = idx / head_dim; + int j = idx % head_dim; + int key_h = h * key_heads / value_heads; + float a = __bfloat162float(a_log[h]); + float neg_a = -expf(a); + float q_scale = rsqrtf((float)head_dim); + + for (int t = 0; t < seq_len; ++t) { + const __nv_bfloat16* q_t = q + ((long long)t * key_heads + key_h) * head_dim; + const __nv_bfloat16* k_t = k + ((long long)t * key_heads + key_h) * head_dim; + const __nv_bfloat16* v_t = v + ((long long)t * value_heads + h) * head_dim; + float b = __bfloat162float(beta[(long long)t * value_heads + h]); + float alpha = __bfloat162float(alpha_softplus[(long long)t * value_heads + h]); + float decay = expf(neg_a * alpha); + + float sk = 0.0f; + for (int i = 0; i < head_dim; ++i) { + long long sidx = ((long long)h * head_dim + i) * head_dim + j; + float s = state[sidx] * decay; + state[sidx] = s; + sk += s * __bfloat162float(k_t[i]); + } + + float d = (__bfloat162float(v_t[j]) - sk) * b; + float out_j = 0.0f; + for (int i = 0; i < head_dim; ++i) { + long long sidx = ((long long)h * head_dim + i) * head_dim + j; + float s = state[sidx] + __bfloat162float(k_t[i]) * d; + state[sidx] = s; + out_j += s * (__bfloat162float(q_t[i]) * q_scale); + } + out[((long long)t * value_heads + h) * head_dim + j] = __float2bfloat16(out_j); + } +} + +extern "C" { + +void launch_depthwise_causal_conv1d_silu_bf16( + const void* x, + const void* w, + void* out, + int seq_len, + int channels, + int kernel, + void* stream +) { + int total = seq_len * channels; + int block = 256; + int grid = (total + block - 1) / block; + depthwise_causal_conv1d_silu_bf16_kernel<<>>( + (const __nv_bfloat16*)x, + (const __nv_bfloat16*)w, + (__nv_bfloat16*)out, + seq_len, + channels, + kernel + ); + CUDA_CHECK_LAST_ERROR(); +} + +void launch_deltanet_ar_bf16( + const void* q, + const void* k, + const void* v, + const void* beta, + const void* alpha_softplus, + const void* a_log, + void* state, + void* out, + int key_heads, + int value_heads, + int head_dim, + void* stream +) { + int total = value_heads * head_dim; + int block = 128; + int grid = (total + block - 1) / block; + deltanet_ar_bf16_kernel<<>>( + (const __nv_bfloat16*)q, + (const __nv_bfloat16*)k, + (const __nv_bfloat16*)v, + (const __nv_bfloat16*)beta, + (const __nv_bfloat16*)alpha_softplus, + (const __nv_bfloat16*)a_log, + (__nv_bfloat16*)state, + (__nv_bfloat16*)out, + key_heads, + value_heads, + head_dim + ); + CUDA_CHECK_LAST_ERROR(); +} + +void launch_depthwise_causal_conv1d_stateful_silu_bf16( + const void* x, + const void* w, + void* state, + void* out, + int seq_len, + int channels, + int kernel, + void* stream +) { + int block = 256; + int grid = (channels + block - 1) / block; + depthwise_causal_conv1d_stateful_silu_bf16_kernel<<>>( + (const __nv_bfloat16*)x, + (const __nv_bfloat16*)w, + (__nv_bfloat16*)state, + (__nv_bfloat16*)out, + seq_len, + channels, + kernel + ); + CUDA_CHECK_LAST_ERROR(); +} + +void launch_l2_normalize_rows_bf16( + const void* x, + void* out, + int rows, + int cols, + float eps, + void* stream +) { + l2_normalize_rows_bf16_kernel<<>>( + (const __nv_bfloat16*)x, + (__nv_bfloat16*)out, + rows, + cols, + eps + ); + CUDA_CHECK_LAST_ERROR(); +} + +void launch_deltanet_recurrent_bf16_f32( + const void* q, + const void* k, + const void* v, + const void* beta, + const void* alpha_softplus, + const void* a_log, + void* state, + void* out, + int seq_len, + int key_heads, + int value_heads, + int head_dim, + void* stream +) { + int total = value_heads * head_dim; + int block = 128; + int grid = (total + block - 1) / block; + deltanet_recurrent_bf16_f32_kernel<<>>( + (const __nv_bfloat16*)q, + (const __nv_bfloat16*)k, + (const __nv_bfloat16*)v, + (const __nv_bfloat16*)beta, + (const __nv_bfloat16*)alpha_softplus, + (const __nv_bfloat16*)a_log, + (float*)state, + (__nv_bfloat16*)out, + seq_len, + key_heads, + value_heads, + head_dim + ); + CUDA_CHECK_LAST_ERROR(); +} + +} diff --git a/csrc/attention/flash_attention.cu b/csrc/attention/flash_attention.cu index 1042007..a0e1666 100644 --- a/csrc/attention/flash_attention.cu +++ b/csrc/attention/flash_attention.cu @@ -13,13 +13,13 @@ // O [batch, num_q_heads, q_len, head_dim] // // Shared memory (BF16): -// smem_q[BR][head_dim] — 64 * 128 * 2 = 16 KB (loaded once per Q tile) -// smem_kv[BC][head_dim] — 64 * 128 * 2 = 16 KB (alternates K and V) -// Total: 32 KB (fits in default 48 KB shared memory) +// Use 32-row tiles so head_dim=256 still fits in the default 48 KB shared +// memory limit: 2 * 32 * 256 * 2 = 32 KB. -#define BR 64 -#define BC 64 +#define BR 32 +#define BC 32 #define THREADS_PER_BLOCK 128 +#define FLASH_HEAD_DIM_MAX 256 __global__ void flash_attention_bf16_kernel( const __nv_bfloat16* __restrict__ Q, @@ -73,8 +73,8 @@ __global__ void flash_attention_bf16_kernel( // Thread t (0 <= t < q_tile_rows) owns Q row t bool owns_row = (tid < q_tile_rows); - // Per-thread FP32 accumulators (head_dim up to 128) - float O_acc[128]; + // Per-thread FP32 accumulators (head_dim up to 256) + float O_acc[FLASH_HEAD_DIM_MAX]; float m_val = -INFINITY; float l_val = 0.0f; if (owns_row) { @@ -250,7 +250,7 @@ __global__ void flash_attention_sinks_bf16_kernel( bool owns_row = (tid < q_tile_rows); - float O_acc[128]; + float O_acc[FLASH_HEAD_DIM_MAX]; float m_val = -INFINITY; float l_val = 0.0f; if (owns_row) { @@ -384,7 +384,7 @@ __global__ void flash_attention_sinks_bf16_kernel( // ============================================================ #define DECODE_THREADS 256 -#define HEAD_DIM_MAX 128 +#define HEAD_DIM_MAX 256 __global__ void decode_attention_bf16_kernel( const __nv_bfloat16* __restrict__ Q, @@ -413,7 +413,7 @@ __global__ void decode_attention_bf16_kernel( const __nv_bfloat16* V_base = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim; __nv_bfloat16* O_ptr = O + ((long long)batch_idx * num_q_heads + q_head) * head_dim; - // Load Q vector into registers (head_dim <= 128) + // Load Q vector into registers (head_dim <= 256) float q_reg[HEAD_DIM_MAX]; for (int d = 0; d < head_dim; d++) { q_reg[d] = __bfloat162float(Q_ptr[d]); diff --git a/csrc/attention/paged_attention.cu b/csrc/attention/paged_attention.cu index 0859485..e9a73f7 100644 --- a/csrc/attention/paged_attention.cu +++ b/csrc/attention/paged_attention.cu @@ -21,11 +21,11 @@ // active batch) so we just index by blockIdx.x_seq. // context_lens [batch] int32 — number of valid tokens per sequence. // -// One CUDA block: 256 threads, head_dim <= 128. +// One CUDA block: 256 threads, head_dim <= 256. #define PAGED_BLOCK_SIZE 16 #define PAGED_THREADS 256 -#define PAGED_HEAD_DIM_MAX 128 +#define PAGED_HEAD_DIM_MAX 256 __global__ void paged_decode_attention_bf16_kernel( const __nv_bfloat16* __restrict__ Q, diff --git a/csrc/embedding/rope.cu b/csrc/embedding/rope.cu index 2d2c4cc..5cddfed 100644 --- a/csrc/embedding/rope.cu +++ b/csrc/embedding/rope.cu @@ -69,6 +69,42 @@ __global__ void rope_bf16( x[base + pair_idx + half_dim] = __float2bfloat16(x1 * cos_val + x0 * sin_val); } +__global__ void partial_rope_bf16( + const __nv_bfloat16* __restrict__ x, + __nv_bfloat16* __restrict__ out, + const int* __restrict__ positions, + int num_heads, int head_dim, int n_rot, float theta +) { + int token_idx = blockIdx.x; + int head_idx = blockIdx.y; + int pair_idx = threadIdx.x; + int half_rot = n_rot / 2; + if (pair_idx >= half_rot) return; + + int pos = positions[token_idx]; + float freq = 1.0f / powf(theta, (float)(2 * pair_idx) / (float)n_rot); + float angle = (float)pos * freq; + float sin_val, cos_val; + sincosf(angle, &sin_val, &cos_val); + + int base = (token_idx * num_heads + head_idx) * head_dim; + float x0 = __bfloat162float(x[base + pair_idx]); + float x1 = __bfloat162float(x[base + pair_idx + half_rot]); + + out[base + pair_idx] = __float2bfloat16(x0 * cos_val - x1 * sin_val); + out[base + pair_idx + half_rot] = __float2bfloat16(x1 * cos_val + x0 * sin_val); +} + +__global__ void copy_partial_rope_tail_bf16( + const __nv_bfloat16* __restrict__ x, + __nv_bfloat16* __restrict__ out, + int total +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) return; + out[idx] = x[idx]; +} + // Precompute cos/sin cache on GPU __global__ void compute_rope_cache( float* __restrict__ cos_cache, // [max_seq_len, half_dim] @@ -109,6 +145,24 @@ void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache, CUDA_CHECK_LAST_ERROR(); } +void launch_partial_rope_bf16(const void* x, void* out, const void* positions, + int num_tokens, int num_heads, int head_dim, + int n_rot, float theta, void* stream) { + int total = num_tokens * num_heads * head_dim; + int block_copy = 256; + int grid_copy = (total + block_copy - 1) / block_copy; + copy_partial_rope_tail_bf16<<>>( + (const __nv_bfloat16*)x, (__nv_bfloat16*)out, total); + CUDA_CHECK_LAST_ERROR(); + + dim3 grid(num_tokens, num_heads); + int block = n_rot / 2; + partial_rope_bf16<<>>( + (const __nv_bfloat16*)x, (__nv_bfloat16*)out, (const int*)positions, + num_heads, head_dim, n_rot, theta); + CUDA_CHECK_LAST_ERROR(); +} + void launch_compute_rope_cache(void* cos_cache, void* sin_cache, int max_seq_len, int half_dim, float theta, void* stream) { diff --git a/csrc/moe/moe_sparse.cu b/csrc/moe/moe_sparse.cu index bed86a3..69c120d 100644 --- a/csrc/moe/moe_sparse.cu +++ b/csrc/moe/moe_sparse.cu @@ -34,6 +34,51 @@ #define SPARSE_TILE_N 8 // output columns per block (= warps per block) +// Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16). +// Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV +// loses nothing to tensor cores and skips activation quantization. +__global__ void moe_sparse_gemv_bf16_bf16_kernel( + const __nv_bfloat16* __restrict__ x, // [T, K] or [T*top_k, K] + const __nv_bfloat16* __restrict__ w, // [local_experts, N, K] + const int* __restrict__ topk_ids, // [T, top_k] global expert ids + __nv_bfloat16* __restrict__ y, // [T, top_k, N] + int N, int K, int top_k, + int expert_start, int local_experts, + int x_per_slot +) { + int token = blockIdx.z; + int slot = blockIdx.y; + int eid = topk_ids[token * top_k + slot]; + int lid = eid - expert_start; + if (lid < 0 || lid >= local_experts) return; + + extern __shared__ float xs[]; + const __nv_bfloat16* xrow = + x + (long long)(x_per_slot ? token * top_k + slot : token) * K; + for (int i = threadIdx.x; i < K; i += blockDim.x) { + xs[i] = __bfloat162float(xrow[i]); + } + __syncthreads(); + + int n = blockIdx.x * SPARSE_TILE_N + (threadIdx.x >> 5); + if (n >= N) return; + int lane = threadIdx.x & 31; + + const __nv_bfloat16* wrow = w + ((long long)lid * N + n) * K; + float acc = 0.0f; + for (int i = lane; i < K; i += 32) { + acc += xs[i] * __bfloat162float(wrow[i]); + } + + #pragma unroll + for (int o = 16; o > 0; o >>= 1) { + acc += __shfl_down_sync(0xffffffffu, acc, o); + } + if (lane == 0) { + y[((long long)token * top_k + slot) * N + n] = __float2bfloat16(acc); + } +} + // Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16). // Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV // loses nothing to tensor cores and skips activation quantization. @@ -194,6 +239,23 @@ __global__ void moe_weighted_sum_sparse_bf16_kernel( extern "C" { +void launch_moe_sparse_gemv_bf16_bf16( + const void* x, const void* w, const void* topk_ids, void* y, + int num_tokens, int N, int K, int top_k, + int expert_start, int local_experts, int x_per_slot, + void* stream +) { + dim3 grid((N + SPARSE_TILE_N - 1) / SPARSE_TILE_N, top_k, num_tokens); + int block = SPARSE_TILE_N * 32; + size_t smem = (size_t)K * sizeof(float); + moe_sparse_gemv_bf16_bf16_kernel<<>>( + (const __nv_bfloat16*)x, (const __nv_bfloat16*)w, (const int*)topk_ids, + (__nv_bfloat16*)y, + N, K, top_k, expert_start, local_experts, x_per_slot + ); + CUDA_CHECK_LAST_ERROR(); +} + void launch_moe_sparse_gemv_fp8_bf16( const void* x, const void* w, const void* w_scales, const void* bias, const void* topk_ids, void* y,