gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Qwen3.6 validation and benchmark coverage

Browse Source
This commit is contained in:
Gahow Wang
2026-07-13 20:24:57 +08:00
parent a2de146fb6
commit 0acaca34cb
12 changed files with 1657 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

280
crates/xserv-model/src/bin/check-qwen35-moe.rs Normal file
View File

@@ -0,0 +1,280 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use serde_json::Value;
use xserv_model::{ModelConfig, ModelFamily, Qwen35MoeSpec, Qwen35TensorMeta, Qwen35TensorSpec};
use xserv_tokenizer::Tokenizer;
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: check-qwen35-moe <metadata-model-dir> [safetensors-shard-dir]");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let shard_dir = args
.get(2)
.map(PathBuf::from)
.unwrap_or_else(|| model_dir.clone());
let config = ModelConfig::from_file(&model_dir.join("config.json"));
if config.family() != ModelFamily::Qwen35Moe {
eprintln!(
"expected Qwen3.5/3.6 MoE config, got {}",
config.family().as_str()
);
std::process::exit(2);
}
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
let index_path = model_dir.join("model.safetensors.index.json");
let index_text = std::fs::read_to_string(&index_path)
.unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
let index: Value = serde_json::from_str(&index_text)
.unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
let weight_map = index
.get("weight_map")
.and_then(|v| v.as_object())
.expect("model.safetensors.index.json must contain weight_map");
let keys: BTreeSet<String> = weight_map.keys().cloned().collect();
println!("family={}", config.family().as_str());
println!("model_type={}", config.model_type_str());
println!("layers={}", config.num_layers());
println!("hidden={}", config.hidden());
println!(
"heads={} kv_heads={} head_dim={}",
config.num_heads(),
config.num_kv_heads(),
config.head_dim()
);
println!(
"vocab_size={} tokenizer_vocab={}",
config.vocab_size(),
tokenizer.vocab_size()
);
println!("max_seq_len={}", config.max_seq_len());
let text = config.text();
println!(
"linear_dims key_heads={:?} value_heads={:?} key_dim={:?} value_dim={:?} conv_kernel={:?}",
text.linear_num_key_heads,
text.linear_num_value_heads,
text.linear_key_head_dim,
text.linear_value_head_dim,
text.linear_conv_kernel_dim
);
println!(
"rope partial={:?} params={:?}",
text.partial_rotary_factor,
text.rope_parameters
.as_ref()
.and_then(|rp| rp.mrope_section.clone())
);
println!("mtp_layers={:?}", text.mtp_num_hidden_layers);
println!(
"experts={} top_k={} moe_intermediate={}",
config.num_experts(),
config.experts_per_token(),
config.ffn_hidden()
);
println!("tensor_count={}", keys.len());
let shard_names: BTreeSet<&str> = weight_map.values().filter_map(|v| v.as_str()).collect();
println!("shard_count={}", shard_names.len());
if let Some(total_size) = index
.pointer("/metadata/total_size")
.and_then(|v| v.as_f64())
{
println!("total_size_gb={:.2}", total_size / 1e9);
}
let mut layer_types = BTreeMap::<String, usize>::new();
for i in 0..config.num_layers() {
let kind = if config.is_linear_attention_layer(i) {
"linear_attention"
} else {
"full_attention"
};
*layer_types.entry(kind.to_string()).or_default() += 1;
}
println!("layer_types={layer_types:?}");
let mut missing = Vec::new();
let spec = Qwen35MoeSpec::from_config(&config);
println!(
"full_attention_path={}",
spec.describe_full_attention_path()
);
let expected = spec.expected_tensor_specs();
for tensor in &expected {
if !keys.contains(&tensor.name) {
missing.push(tensor.name.clone());
}
}
let visual = keys
.iter()
.filter(|k| k.starts_with("model.visual."))
.count();
let language = keys
.iter()
.filter(|k| k.starts_with("model.language_model."))
.count();
println!("language_tensors={language}");
println!("visual_tensors={visual}");
if missing.is_empty() {
println!("schema_check=ok");
} else {
println!("schema_check=missing {}", missing.len());
for key in missing.iter().take(50) {
println!("missing={key}");
}
std::process::exit(3);
}
validate_headers(&shard_dir, weight_map, &expected);
let tensor_meta = collect_header_meta(&shard_dir, weight_map, &expected);
let model_meta = spec
.build_meta(&tensor_meta)
.unwrap_or_else(|e| panic!("failed to build Qwen35ModelMeta: {e}"));
let meta_linear = model_meta
.layers
.iter()
.filter(|layer| {
matches!(
layer.attention,
xserv_model::qwen35_moe::Qwen35AttentionMeta::Linear(_)
)
})
.count();
let meta_full = model_meta.layers.len() - meta_linear;
println!("model_meta_layers={}", model_meta.layers.len());
println!("model_meta_linear_layers={meta_linear}");
println!("model_meta_full_attention_layers={meta_full}");
}
fn collect_header_meta(
shard_dir: &Path,
weight_map: &serde_json::Map<String, Value>,
expected: &[Qwen35TensorSpec],
) -> BTreeMap<String, Qwen35TensorMeta> {
let mut by_shard = BTreeMap::<String, Vec<&Qwen35TensorSpec>>::new();
for spec in expected {
if let Some(shard) = weight_map.get(&spec.name).and_then(|v| v.as_str()) {
by_shard.entry(shard.to_string()).or_default().push(spec);
}
}
let mut out = BTreeMap::new();
for (shard, specs) in by_shard {
let path = shard_dir.join(&shard);
if !path.exists() {
continue;
}
let header = read_safetensors_header(&path);
for spec in specs {
if let Some(meta) = header.get(&spec.name) {
out.insert(
spec.name.clone(),
Qwen35TensorMeta {
name: spec.name.clone(),
dtype: meta
.get("dtype")
.and_then(|v| v.as_str())
.unwrap_or("?")
.to_string(),
shape: meta
.get("shape")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_u64().map(|n| n as usize))
.collect::<Vec<_>>()
})
.unwrap_or_default(),
},
);
}
}
}
out
}
fn validate_headers(
shard_dir: &Path,
weight_map: &serde_json::Map<String, Value>,
expected: &[Qwen35TensorSpec],
) {
let mut by_shard = BTreeMap::<String, Vec<&Qwen35TensorSpec>>::new();
for spec in expected {
if let Some(shard) = weight_map.get(&spec.name).and_then(|v| v.as_str()) {
by_shard.entry(shard.to_string()).or_default().push(spec);
}
}
let mut checked = 0usize;
let mut skipped_shards = 0usize;
let mut errors = Vec::new();
for (shard, specs) in by_shard {
let path = shard_dir.join(&shard);
if !path.exists() {
skipped_shards += 1;
continue;
}
let header = read_safetensors_header(&path);
for spec in specs {
checked += 1;
match header.get(&spec.name) {
Some(meta) => {
let dtype = meta.get("dtype").and_then(|v| v.as_str()).unwrap_or("?");
let shape = meta
.get("shape")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_u64().map(|n| n as usize))
.collect::<Vec<_>>()
})
.unwrap_or_default();
if dtype != spec.dtype || shape != spec.shape {
errors.push(format!(
"{} expected {} {:?}, got {} {:?}",
spec.name, spec.dtype, spec.shape, dtype, shape
));
}
}
None => errors.push(format!("{} missing from shard header {shard}", spec.name)),
}
}
}
println!("header_checked_tensors={checked}");
println!("header_skipped_missing_shards={skipped_shards}");
if errors.is_empty() {
println!("header_check=ok");
} else {
println!("header_check=errors {}", errors.len());
for err in errors.iter().take(50) {
println!("header_error={err}");
}
std::process::exit(4);
}
}
fn read_safetensors_header(path: &Path) -> serde_json::Map<String, Value> {
let mut file =
File::open(path).unwrap_or_else(|e| panic!("failed to open {}: {e}", path.display()));
let mut len_bytes = [0u8; 8];
file.read_exact(&mut len_bytes)
.unwrap_or_else(|e| panic!("failed to read header size from {}: {e}", path.display()));
let header_len = u64::from_le_bytes(len_bytes);
file.seek(SeekFrom::Start(8)).unwrap();
let mut header = vec![0u8; header_len as usize];
file.read_exact(&mut header)
.unwrap_or_else(|e| panic!("failed to read header from {}: {e}", path.display()));
serde_json::from_slice::<Value>(&header)
.unwrap_or_else(|e| panic!("failed to parse safetensors header {}: {e}", path.display()))
.as_object()
.cloned()
.expect("safetensors header must be a JSON object")
}

195
crates/xserv-model/src/bin/smoke-qwen35-full-attn.rs Normal file
View File

@@ -0,0 +1,195 @@
use std::path::PathBuf;
use half::bf16;
use xserv_model::{ModelConfig, Qwen35Moe, Qwen35MoeSpec, loader};
use xserv_tensor::{Device, Tensor};
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: smoke-qwen35-full-attn <model-dir> [layer] [seq_len]");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(3);
let seq_len: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(2);
let device = 0;
xserv_cuda::device::set_device(device).unwrap();
xserv_model::init_kernels();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let spec = Qwen35MoeSpec::from_config(&config);
assert_eq!(
spec.layer_kinds[layer],
xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention,
"layer {layer} is not a full-attention layer"
);
let weight_map = read_weight_map(&model_dir);
let keys = [
format!("model.language_model.layers.{layer}.self_attn.q_proj.weight"),
format!("model.language_model.layers.{layer}.self_attn.k_proj.weight"),
format!("model.language_model.layers.{layer}.self_attn.v_proj.weight"),
format!("model.language_model.layers.{layer}.self_attn.o_proj.weight"),
format!("model.language_model.layers.{layer}.self_attn.q_norm.weight"),
format!("model.language_model.layers.{layer}.self_attn.k_norm.weight"),
];
let mut shards: Vec<String> = keys
.iter()
.filter_map(|k| weight_map.get(k).cloned())
.collect();
shards.sort();
shards.dedup();
let mut tensors = std::collections::HashMap::new();
for shard in &shards {
tensors.extend(loader::load_safetensors(
&model_dir.join(shard),
Device::Cpu,
));
}
let weights = Qwen35Moe::take_full_attention_weights(&mut tensors, layer, Device::Cuda(device));
let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
let n_rot =
(spec.head_dim as f64 * config.text().partial_rotary_factor.unwrap_or(0.25)) as usize;
let rope_theta = config.rope_theta_value().unwrap_or(10_000_000.0) as f32;
let positions: Vec<u32> = (0..seq_len as u32).collect();
let out = Qwen35Moe::forward_full_attention_layer(
&spec,
&input,
&weights,
&positions,
n_rot,
rope_theta,
config.ln_eps(),
);
let q_rope_cpu = partial_rope_cpu(
&out.q_normed,
seq_len,
spec.num_heads,
spec.head_dim,
n_rot,
rope_theta,
)
.to_device(Device::Cuda(device));
let k_rope_cpu = partial_rope_cpu(
&out.k_normed,
seq_len,
spec.num_kv_heads,
spec.head_dim,
n_rot,
rope_theta,
)
.to_device(Device::Cuda(device));
println!("layer={layer}");
println!("input_shape={:?}", input.shape());
println!("q_full_shape={:?}", out.q_full.shape());
println!("query_shape={:?}", out.query.shape());
println!("gate_shape={:?}", out.gate.shape());
println!("k_shape={:?}", out.k.shape());
println!("v_shape={:?}", out.v.shape());
println!("q_normed_shape={:?}", out.q_normed.shape());
println!("k_normed_shape={:?}", out.k_normed.shape());
println!("n_rot={n_rot} rope_theta={rope_theta}");
println!("q_rope_cpu_shape={:?}", q_rope_cpu.shape());
println!("q_rope_gpu_shape={:?}", out.q_rope.shape());
println!("k_rope_cpu_shape={:?}", k_rope_cpu.shape());
println!("k_rope_gpu_shape={:?}", out.k_rope.shape());
println!("sample_q={:?}", sample_bf16(&out.q_normed, 8));
println!("sample_q_rope_cpu={:?}", sample_bf16(&q_rope_cpu, 8));
println!("sample_q_rope_gpu={:?}", sample_bf16(&out.q_rope, 8));
let token1_offset = spec.num_heads * spec.head_dim;
println!(
"sample_q_token1={:?}",
sample_bf16_offset(&out.q_normed, token1_offset, 8)
);
println!(
"sample_q_rope_cpu_token1={:?}",
sample_bf16_offset(&q_rope_cpu, token1_offset, 8)
);
println!(
"sample_q_rope_gpu_token1={:?}",
sample_bf16_offset(&out.q_rope, token1_offset, 8)
);
println!("sample_gate={:?}", sample_bf16(&out.gate, 8));
println!("attn_out_shape={:?}", out.attn_out.shape());
println!("attn_merged_shape={:?}", out.attn_merged.shape());
println!("gate_sigmoid_shape={:?}", out.gate_sigmoid.shape());
println!("gated_attn_shape={:?}", out.gated_attn.shape());
println!("projected_shape={:?}", out.projected.shape());
println!("sample_attn={:?}", sample_bf16(&out.attn_merged, 8));
println!("sample_projected={:?}", sample_bf16(&out.projected, 8));
}
fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
let index_path = model_dir.join("model.safetensors.index.json");
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()));
index
.get("weight_map")
.and_then(|v| v.as_object())
.expect("index must contain weight_map")
.iter()
.map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
.collect()
}
fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
let mut data = Vec::with_capacity(seq_len * hidden);
for i in 0..seq_len * hidden {
let x = ((i % 97) as f32 - 48.0) / 128.0;
data.push(bf16::from_f32(x));
}
Tensor::from_slice(&data, &[seq_len, hidden])
}
fn partial_rope_cpu(
x: &Tensor,
seq_len: usize,
num_heads: usize,
head_dim: usize,
n_rot: usize,
theta: f32,
) -> Tensor {
assert_eq!(x.shape(), &[seq_len * num_heads, head_dim]);
assert!(n_rot <= head_dim && n_rot % 2 == 0);
let cpu = x.to_device(Device::Cpu);
let src = cpu.as_slice::<bf16>();
let mut out = src.to_vec();
let half = n_rot / 2;
for s in 0..seq_len {
for h in 0..num_heads {
let base = (s * num_heads + h) * head_dim;
for i in 0..half {
let freq = 1.0f32 / theta.powf((2 * i) as f32 / n_rot as f32);
let angle = s as f32 * freq;
let (sin, cos) = angle.sin_cos();
let x0 = src[base + i].to_f32();
let x1 = src[base + i + half].to_f32();
out[base + i] = bf16::from_f32(x0 * cos - x1 * sin);
out[base + i + half] = bf16::from_f32(x1 * cos + x0 * sin);
}
}
}
Tensor::from_slice(&out, &[seq_len * num_heads, head_dim])
}
fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
sample_bf16_offset(t, 0, n)
}
fn sample_bf16_offset(t: &Tensor, offset: usize, n: usize) -> Vec<f32> {
let cpu = t.to_device(Device::Cpu);
cpu.as_slice::<bf16>()
.iter()
.skip(offset)
.take(n)
.map(|v| v.to_f32())
.collect()
}

332
crates/xserv-model/src/bin/smoke-qwen35-layer.rs Normal file
View File

@@ -0,0 +1,332 @@
use std::path::PathBuf;
use half::bf16;
use xserv_kernels::{
GemmBackend, add, matmul,
moe::{moe_sparse_gemv_bf16, moe_topk_softmax, moe_weighted_sum_sparse},
mul, row_scale_bf16, sigmoid, silu,
};
use xserv_model::{ModelConfig, Qwen35FullAttentionWeights, Qwen35Moe, Qwen35MoeSpec, loader};
use xserv_tensor::{Device, Tensor};
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: smoke-qwen35-layer <model-dir> [layer]");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
let device = 0;
let seq_len = 1;
xserv_cuda::device::set_device(device).unwrap();
xserv_model::init_kernels();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let spec = Qwen35MoeSpec::from_config(&config);
let weight_map = read_weight_map(&model_dir);
let mut keys = layer_common_keys(layer);
match spec.layer_kinds[layer] {
xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
keys.extend(linear_attention_keys(layer));
}
xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
keys.extend(full_attention_keys(layer));
}
}
keys.extend(moe_keys(layer));
let mut shards: Vec<String> = keys
.iter()
.filter_map(|k| weight_map.get(k).cloned())
.collect();
shards.sort();
shards.dedup();
let mut tensors = std::collections::HashMap::new();
for shard in &shards {
tensors.extend(loader::load_safetensors(
&model_dir.join(shard),
Device::Cpu,
));
}
let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
let p = format!("model.language_model.layers.{layer}");
let input_norm = take_tensor(&mut tensors, &format!("{p}.input_layernorm.weight"), device);
let post_norm = take_tensor(
&mut tensors,
&format!("{p}.post_attention_layernorm.weight"),
device,
);
let normed = xserv_kernels::rmsnorm(&input, &input_norm, config.ln_eps());
let attn_projected = match spec.layer_kinds[layer] {
xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
let weights = Qwen35Moe::take_linear_attention_projection_weights(
&mut tensors,
layer,
Device::Cuda(device),
);
let projected = Qwen35Moe::project_linear_attention(&spec, &normed, &weights);
let mut state = Tensor::zeros(
&[
spec.linear_value_heads,
spec.linear_value_dim,
spec.linear_value_dim,
],
xserv_tensor::DType::BF16,
Device::Cuda(device),
);
let recurrent =
Qwen35Moe::linear_attention_recurrent_step(&spec, &projected, &weights, &mut state);
let (_, out) = Qwen35Moe::finish_linear_attention(
&spec,
&recurrent,
&projected.z,
&weights,
config.ln_eps(),
);
out
}
xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
let weights = Qwen35FullAttentionWeights {
q_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.q_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
k_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.k_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
v_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.v_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
o_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.o_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
q_norm: take_tensor(
&mut tensors,
&format!("{p}.self_attn.q_norm.weight"),
device,
),
k_norm: take_tensor(
&mut tensors,
&format!("{p}.self_attn.k_norm.weight"),
device,
),
};
let n_rot = (spec.head_dim as f64 * config.text().partial_rotary_factor.unwrap_or(0.25))
as usize;
let positions = [0u32];
Qwen35Moe::forward_full_attention_layer(
&spec,
&normed,
&weights,
&positions,
n_rot,
config.rope_theta_value().unwrap_or(10_000_000.0) as f32,
config.ln_eps(),
)
.projected
}
};
let attn_residual = add(&input, &attn_projected);
let moe_in = xserv_kernels::rmsnorm(&attn_residual, &post_norm, config.ln_eps());
let moe_out = run_moe(&mut tensors, layer, &moe_in, &spec, device);
let layer_out = add(&attn_residual, &moe_out);
println!("layer={layer}");
println!("kind={:?}", spec.layer_kinds[layer]);
println!("input_shape={:?}", input.shape());
println!("normed_shape={:?}", normed.shape());
println!("attn_projected_shape={:?}", attn_projected.shape());
println!("attn_residual_shape={:?}", attn_residual.shape());
println!("moe_in_shape={:?}", moe_in.shape());
println!("moe_out_shape={:?}", moe_out.shape());
println!("layer_out_shape={:?}", layer_out.shape());
println!("sample_attn={:?}", sample_bf16(&attn_projected, 8));
println!("sample_moe={:?}", sample_bf16(&moe_out, 8));
println!("sample_layer_out={:?}", sample_bf16(&layer_out, 8));
}
fn run_moe(
tensors: &mut std::collections::HashMap<String, Tensor>,
layer: usize,
input: &Tensor,
spec: &Qwen35MoeSpec,
device: u32,
) -> Tensor {
let p = format!("model.language_model.layers.{layer}.mlp");
let router_w = take_tensor(tensors, &format!("{p}.gate.weight"), device)
.transpose(0, 1)
.contiguous();
let shared_gate_w = take_tensor(
tensors,
&format!("{p}.shared_expert.gate_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous();
let shared_up_w = take_tensor(
tensors,
&format!("{p}.shared_expert.up_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous();
let shared_down_w = take_tensor(
tensors,
&format!("{p}.shared_expert.down_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous();
let shared_router_w = take_tensor(tensors, &format!("{p}.shared_expert_gate.weight"), device)
.transpose(0, 1)
.contiguous();
let expert_gate_up =
take_tensor(tensors, &format!("{p}.experts.gate_up_proj"), device).contiguous();
let expert_down = take_tensor(tensors, &format!("{p}.experts.down_proj"), device).contiguous();
let router_logits = matmul(input, &router_w, GemmBackend::CuBlas);
let (topk_ids, topk_weights) =
moe_topk_softmax(&router_logits, spec.num_experts, spec.experts_per_token);
let shared_gate = matmul(input, &shared_gate_w, GemmBackend::CuBlas);
let shared_up = matmul(input, &shared_up_w, GemmBackend::CuBlas);
let shared_act = mul(&silu(&shared_gate), &shared_up);
let shared_down = matmul(&shared_act, &shared_down_w, GemmBackend::CuBlas);
let shared_router = sigmoid(&matmul(input, &shared_router_w, GemmBackend::CuBlas));
let shared_out = row_scale_bf16(&shared_down, &shared_router);
let gate_up = moe_sparse_gemv_bf16(
input,
&expert_gate_up,
&topk_ids,
spec.experts_per_token,
0,
spec.num_experts,
false,
);
let gate = gate_up.narrow(2, 0, spec.moe_intermediate).contiguous();
let up = gate_up
.narrow(2, spec.moe_intermediate, spec.moe_intermediate)
.contiguous();
let routed_act = mul(&silu(&gate), &up);
let down = moe_sparse_gemv_bf16(
&routed_act.reshape(&[spec.experts_per_token, spec.moe_intermediate]),
&expert_down,
&topk_ids,
spec.experts_per_token,
0,
spec.num_experts,
true,
);
let routed_out = moe_weighted_sum_sparse(&down, &topk_ids, &topk_weights, 0, spec.num_experts);
add(&routed_out, &shared_out)
}
fn take_tensor(
tensors: &mut std::collections::HashMap<String, Tensor>,
name: &str,
device: u32,
) -> Tensor {
tensors
.remove(name)
.unwrap_or_else(|| panic!("missing tensor {name}"))
.to_device(Device::Cuda(device))
}
fn layer_common_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}");
vec![
format!("{p}.input_layernorm.weight"),
format!("{p}.post_attention_layernorm.weight"),
]
}
fn linear_attention_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}.linear_attn");
vec![
format!("{p}.in_proj_qkv.weight"),
format!("{p}.in_proj_z.weight"),
format!("{p}.in_proj_a.weight"),
format!("{p}.in_proj_b.weight"),
format!("{p}.conv1d.weight"),
format!("{p}.A_log"),
format!("{p}.dt_bias"),
format!("{p}.norm.weight"),
format!("{p}.out_proj.weight"),
]
}
fn full_attention_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}.self_attn");
vec![
format!("{p}.q_proj.weight"),
format!("{p}.k_proj.weight"),
format!("{p}.v_proj.weight"),
format!("{p}.o_proj.weight"),
format!("{p}.q_norm.weight"),
format!("{p}.k_norm.weight"),
]
}
fn moe_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}.mlp");
vec![
format!("{p}.gate.weight"),
format!("{p}.shared_expert.gate_proj.weight"),
format!("{p}.shared_expert.up_proj.weight"),
format!("{p}.shared_expert.down_proj.weight"),
format!("{p}.shared_expert_gate.weight"),
format!("{p}.experts.gate_up_proj"),
format!("{p}.experts.down_proj"),
]
}
fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
let index_path = model_dir.join("model.safetensors.index.json");
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()));
index
.get("weight_map")
.and_then(|v| v.as_object())
.expect("index must contain weight_map")
.iter()
.map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
.collect()
}
fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
let mut data = Vec::with_capacity(seq_len * hidden);
for i in 0..seq_len * hidden {
let x = ((i % 97) as f32 - 48.0) / 128.0;
data.push(bf16::from_f32(x));
}
Tensor::from_slice(&data, &[seq_len, hidden])
}
fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
let cpu = t.to_device(Device::Cpu);
cpu.as_slice::<bf16>()
.iter()
.take(n)
.map(|v| v.to_f32())
.collect()
}

164
crates/xserv-model/src/bin/smoke-qwen35-linear-attn.rs Normal file
View File

@@ -0,0 +1,164 @@
use std::path::PathBuf;
use half::bf16;
use xserv_model::{ModelConfig, Qwen35Moe, Qwen35MoeSpec, loader};
use xserv_tensor::{Device, Tensor};
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: smoke-qwen35-linear-attn <model-dir> [layer] [seq_len]");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
let seq_len: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(2);
let device = 0;
xserv_cuda::device::set_device(device).unwrap();
xserv_model::init_kernels();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let spec = Qwen35MoeSpec::from_config(&config);
assert_eq!(
spec.layer_kinds[layer],
xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention,
"layer {layer} is not a linear-attention layer"
);
let weight_map = read_weight_map(&model_dir);
let keys = [
format!("model.language_model.layers.{layer}.linear_attn.in_proj_qkv.weight"),
format!("model.language_model.layers.{layer}.linear_attn.in_proj_z.weight"),
format!("model.language_model.layers.{layer}.linear_attn.in_proj_a.weight"),
format!("model.language_model.layers.{layer}.linear_attn.in_proj_b.weight"),
format!("model.language_model.layers.{layer}.linear_attn.conv1d.weight"),
format!("model.language_model.layers.{layer}.linear_attn.A_log"),
format!("model.language_model.layers.{layer}.linear_attn.dt_bias"),
format!("model.language_model.layers.{layer}.linear_attn.norm.weight"),
format!("model.language_model.layers.{layer}.linear_attn.out_proj.weight"),
];
let mut shards: Vec<String> = keys
.iter()
.filter_map(|k| weight_map.get(k).cloned())
.collect();
shards.sort();
shards.dedup();
let mut tensors = std::collections::HashMap::new();
for shard in &shards {
tensors.extend(loader::load_safetensors(
&model_dir.join(shard),
Device::Cpu,
));
}
let weights = Qwen35Moe::take_linear_attention_projection_weights(
&mut tensors,
layer,
Device::Cuda(device),
);
let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
let out = Qwen35Moe::project_linear_attention(&spec, &input, &weights);
let mut recurrent_state = Tensor::zeros(
&[
spec.linear_value_heads,
spec.linear_value_dim,
spec.linear_value_dim,
],
xserv_tensor::DType::BF16,
Device::Cuda(device),
);
let recurrent_out = if seq_len == 1 {
Some(Qwen35Moe::linear_attention_recurrent_step(
&spec,
&out,
&weights,
&mut recurrent_state,
))
} else {
None
};
let finished = recurrent_out
.as_ref()
.map(|r| Qwen35Moe::finish_linear_attention(&spec, r, &out.z, &weights, config.ln_eps()));
println!("layer={layer}");
println!("input_shape={:?}", input.shape());
println!("qkv_shape={:?}", out.qkv.shape());
println!("z_shape={:?}", out.z.shape());
println!("a_shape={:?}", out.a.shape());
println!("b_shape={:?}", out.b.shape());
println!("beta_shape={:?}", out.beta.shape());
println!("alpha_softplus_shape={:?}", out.alpha_softplus.shape());
println!("q_shape={:?}", out.q.shape());
println!("k_shape={:?}", out.k.shape());
println!("v_shape={:?}", out.v.shape());
println!("conv_shape={:?}", out.conv.shape());
println!("q_conv_shape={:?}", out.q_conv.shape());
println!("k_conv_shape={:?}", out.k_conv.shape());
println!("v_conv_shape={:?}", out.v_conv.shape());
println!("sample_q={:?}", sample_bf16(&out.q, 8));
println!("sample_k={:?}", sample_bf16(&out.k, 8));
println!("sample_v={:?}", sample_bf16(&out.v, 8));
println!("sample_conv={:?}", sample_bf16(&out.conv, 8));
println!("sample_q_conv={:?}", sample_bf16(&out.q_conv, 8));
println!("sample_k_conv={:?}", sample_bf16(&out.k_conv, 8));
println!("sample_v_conv={:?}", sample_bf16(&out.v_conv, 8));
println!("sample_z={:?}", sample_bf16(&out.z, 8));
println!("sample_a={:?}", sample_bf16(&out.a, 8));
println!("sample_b={:?}", sample_bf16(&out.b, 8));
println!("sample_beta={:?}", sample_bf16(&out.beta, 8));
println!(
"sample_alpha_softplus={:?}",
sample_bf16(&out.alpha_softplus, 8)
);
if let Some(recurrent_out) = recurrent_out {
println!("recurrent_out_shape={:?}", recurrent_out.shape());
println!("recurrent_state_shape={:?}", recurrent_state.shape());
println!("sample_recurrent_out={:?}", sample_bf16(&recurrent_out, 8));
println!(
"sample_recurrent_state={:?}",
sample_bf16(&recurrent_state, 8)
);
}
if let Some((norm_gated, projected)) = finished {
println!("norm_gated_shape={:?}", norm_gated.shape());
println!("projected_shape={:?}", projected.shape());
println!("sample_norm_gated={:?}", sample_bf16(&norm_gated, 8));
println!("sample_projected={:?}", sample_bf16(&projected, 8));
}
}
fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
let index_path = model_dir.join("model.safetensors.index.json");
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()));
index
.get("weight_map")
.and_then(|v| v.as_object())
.expect("index must contain weight_map")
.iter()
.map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
.collect()
}
fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
let mut data = Vec::with_capacity(seq_len * hidden);
for i in 0..seq_len * hidden {
let x = ((i % 97) as f32 - 48.0) / 128.0;
data.push(bf16::from_f32(x));
}
Tensor::from_slice(&data, &[seq_len, hidden])
}
fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
let cpu = t.to_device(Device::Cpu);
cpu.as_slice::<bf16>()
.iter()
.take(n)
.map(|v| v.to_f32())
.collect()
}

193
crates/xserv-model/src/bin/smoke-qwen35-moe.rs Normal file
View File

@@ -0,0 +1,193 @@
use std::path::PathBuf;
use half::bf16;
use xserv_kernels::{
GemmBackend, add, matmul,
moe::{moe_sparse_gemv_bf16, moe_topk_softmax, moe_weighted_sum_sparse},
mul, row_scale_bf16, sigmoid, silu,
};
use xserv_model::{ModelConfig, Qwen35MoeSpec, loader};
use xserv_tensor::{Device, Tensor};
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: smoke-qwen35-moe <model-dir> [layer] [seq_len]");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
let seq_len: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(1);
let device = 0;
xserv_cuda::device::set_device(device).unwrap();
xserv_model::init_kernels();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let spec = Qwen35MoeSpec::from_config(&config);
let weight_map = read_weight_map(&model_dir);
let keys = [
format!("model.language_model.layers.{layer}.mlp.gate.weight"),
format!("model.language_model.layers.{layer}.mlp.shared_expert.gate_proj.weight"),
format!("model.language_model.layers.{layer}.mlp.shared_expert.up_proj.weight"),
format!("model.language_model.layers.{layer}.mlp.shared_expert.down_proj.weight"),
format!("model.language_model.layers.{layer}.mlp.shared_expert_gate.weight"),
format!("model.language_model.layers.{layer}.mlp.experts.gate_up_proj"),
format!("model.language_model.layers.{layer}.mlp.experts.down_proj"),
];
let mut shards: Vec<String> = keys
.iter()
.filter_map(|k| weight_map.get(k).cloned())
.collect();
shards.sort();
shards.dedup();
let mut tensors = std::collections::HashMap::new();
for shard in &shards {
tensors.extend(loader::load_safetensors(
&model_dir.join(shard),
Device::Cpu,
));
}
let mut take = |name: &str| -> Tensor {
tensors
.remove(name)
.unwrap_or_else(|| panic!("missing tensor {name}"))
.to_device(Device::Cuda(device))
};
let p = format!("model.language_model.layers.{layer}.mlp");
let router_w = take(&format!("{p}.gate.weight"))
.transpose(0, 1)
.contiguous();
let shared_gate_w = take(&format!("{p}.shared_expert.gate_proj.weight"))
.transpose(0, 1)
.contiguous();
let shared_up_w = take(&format!("{p}.shared_expert.up_proj.weight"))
.transpose(0, 1)
.contiguous();
let shared_down_w = take(&format!("{p}.shared_expert.down_proj.weight"))
.transpose(0, 1)
.contiguous();
let shared_router_w = take(&format!("{p}.shared_expert_gate.weight"))
.transpose(0, 1)
.contiguous();
let expert_gate_up = take(&format!("{p}.experts.gate_up_proj")).contiguous();
let expert_down = take(&format!("{p}.experts.down_proj")).contiguous();
let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
let router_logits = matmul(&input, &router_w, GemmBackend::CuBlas);
let (topk_ids, topk_weights) =
moe_topk_softmax(&router_logits, spec.num_experts, spec.experts_per_token);
let shared_gate = matmul(&input, &shared_gate_w, GemmBackend::CuBlas);
let shared_up = matmul(&input, &shared_up_w, GemmBackend::CuBlas);
let shared_act = mul(&silu(&shared_gate), &shared_up);
let shared_down = matmul(&shared_act, &shared_down_w, GemmBackend::CuBlas);
let shared_router = sigmoid(&matmul(&input, &shared_router_w, GemmBackend::CuBlas));
let shared_out = row_scale_bf16(&shared_down, &shared_router);
let gate_up = moe_sparse_gemv_bf16(
&input,
&expert_gate_up,
&topk_ids,
spec.experts_per_token,
0,
spec.num_experts,
false,
);
let gate = gate_up.narrow(2, 0, spec.moe_intermediate).contiguous();
let up = gate_up
.narrow(2, spec.moe_intermediate, spec.moe_intermediate)
.contiguous();
let routed_act = mul(&silu(&gate), &up);
let down = moe_sparse_gemv_bf16(
&routed_act.reshape(&[seq_len * spec.experts_per_token, spec.moe_intermediate]),
&expert_down,
&topk_ids,
spec.experts_per_token,
0,
spec.num_experts,
true,
);
let routed_out = moe_weighted_sum_sparse(&down, &topk_ids, &topk_weights, 0, spec.num_experts);
let moe_out = add(&routed_out, &shared_out);
println!("layer={layer}");
println!("input_shape={:?}", input.shape());
println!("router_logits_shape={:?}", router_logits.shape());
println!("topk_ids_shape={:?}", topk_ids.shape());
println!("topk_weights_shape={:?}", topk_weights.shape());
println!("shared_gate_shape={:?}", shared_gate.shape());
println!("shared_up_shape={:?}", shared_up.shape());
println!("shared_act_shape={:?}", shared_act.shape());
println!("shared_down_shape={:?}", shared_down.shape());
println!("shared_router_shape={:?}", shared_router.shape());
println!("shared_out_shape={:?}", shared_out.shape());
println!("gate_up_shape={:?}", gate_up.shape());
println!("routed_act_shape={:?}", routed_act.shape());
println!("down_shape={:?}", down.shape());
println!("routed_out_shape={:?}", routed_out.shape());
println!("moe_out_shape={:?}", moe_out.shape());
println!("sample_router={:?}", sample_bf16(&router_logits, 8));
println!(
"sample_topk_ids={:?}",
sample_i32_raw(&topk_ids, spec.experts_per_token)
);
println!(
"sample_topk_weights={:?}",
sample_f32(&topk_weights, spec.experts_per_token)
);
println!("sample_shared_router={:?}", sample_bf16(&shared_router, 4));
println!("sample_shared_out={:?}", sample_bf16(&shared_out, 8));
println!("sample_routed_out={:?}", sample_bf16(&routed_out, 8));
println!("sample_moe_out={:?}", sample_bf16(&moe_out, 8));
}
fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
let index_path = model_dir.join("model.safetensors.index.json");
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()));
index
.get("weight_map")
.and_then(|v| v.as_object())
.expect("index must contain weight_map")
.iter()
.map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
.collect()
}
fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
let mut data = Vec::with_capacity(seq_len * hidden);
for i in 0..seq_len * hidden {
let x = ((i % 97) as f32 - 48.0) / 128.0;
data.push(bf16::from_f32(x));
}
Tensor::from_slice(&data, &[seq_len, hidden])
}
fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
let cpu = t.to_device(Device::Cpu);
cpu.as_slice::<bf16>()
.iter()
.take(n)
.map(|v| v.to_f32())
.collect()
}
fn sample_f32(t: &Tensor, n: usize) -> Vec<f32> {
let cpu = t.to_device(Device::Cpu);
cpu.as_slice::<f32>().iter().take(n).copied().collect()
}
fn sample_i32_raw(t: &Tensor, n: usize) -> Vec<i32> {
let cpu = t.to_device(Device::Cpu);
let bytes = cpu.as_raw_bytes();
(0..n)
.map(|i| {
let start = i * 4;
i32::from_ne_bytes(bytes[start..start + 4].try_into().unwrap())
})
.collect()
}

411
crates/xserv-model/src/bin/smoke-qwen35-prefix.rs Normal file
View File

@@ -0,0 +1,411 @@
use std::path::PathBuf;
use half::bf16;
use xserv_kernels::{
GemmBackend, add, embedding, matmul,
moe::{moe_sparse_gemv_bf16, moe_topk_softmax, moe_weighted_sum_sparse},
mul, row_scale_bf16, sigmoid, silu,
};
use xserv_model::{ModelConfig, Qwen35FullAttentionWeights, Qwen35Moe, Qwen35MoeSpec, loader};
use xserv_tensor::{Device, Tensor};
use xserv_tokenizer::Tokenizer;
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: smoke-qwen35-prefix <model-dir> [num_layers] [token_id]");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let num_layers: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(4);
let token_id: Option<u32> = args.get(3).and_then(|s| s.parse().ok());
let device = 0;
xserv_cuda::device::set_device(device).unwrap();
xserv_model::init_kernels();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
let spec = Qwen35MoeSpec::from_config(&config);
let weight_map = read_weight_map(&model_dir);
let mut final_tensors = load_keys(
&model_dir,
&weight_map,
&[
"model.language_model.embed_tokens.weight".to_string(),
"model.language_model.norm.weight".to_string(),
"lm_head.weight".to_string(),
],
);
let mut hidden = if let Some(token_id) = token_id {
let embed = take_tensor(
&mut final_tensors,
"model.language_model.embed_tokens.weight",
device,
);
println!("input_token_id={token_id}");
println!("input_token_text={:?}", tokenizer.decode(&[token_id]));
embedding(&embed, &[token_id])
} else {
make_deterministic_input(1, config.hidden()).to_device(Device::Cuda(device))
};
println!("prefix_layers={num_layers}");
println!("input_shape={:?}", hidden.shape());
for layer in 0..num_layers {
hidden = run_layer(
&model_dir,
&weight_map,
&config,
&spec,
layer,
&hidden,
device,
);
println!(
"layer={layer} kind={:?} out_shape={:?}",
spec.layer_kinds[layer],
hidden.shape()
);
println!("sample_layer_{layer}={:?}", sample_bf16(&hidden, 8));
}
let final_norm = take_tensor(
&mut final_tensors,
"model.language_model.norm.weight",
device,
);
let lm_head_t = take_tensor(&mut final_tensors, "lm_head.weight", device)
.transpose(0, 1)
.contiguous();
let normed = xserv_kernels::rmsnorm(&hidden, &final_norm, config.ln_eps());
let logits = matmul(&normed, &lm_head_t, GemmBackend::CuBlas);
let top = xserv_kernels::argmax_bf16_single(&logits);
println!("final_normed_shape={:?}", normed.shape());
println!("logits_shape={:?}", logits.shape());
println!("top_token={top}");
println!("top_token_text={:?}", tokenizer.decode(&[top]));
println!("sample_logits={:?}", sample_bf16(&logits, 8));
}
fn load_keys(
model_dir: &std::path::Path,
weight_map: &std::collections::HashMap<String, String>,
keys: &[String],
) -> std::collections::HashMap<String, Tensor> {
let mut shards: Vec<String> = keys
.iter()
.filter_map(|k| weight_map.get(k).cloned())
.collect();
shards.sort();
shards.dedup();
let mut tensors = std::collections::HashMap::new();
for shard in &shards {
tensors.extend(loader::load_safetensors(
&model_dir.join(shard),
Device::Cpu,
));
}
tensors
}
fn run_layer(
model_dir: &std::path::Path,
weight_map: &std::collections::HashMap<String, String>,
config: &ModelConfig,
spec: &Qwen35MoeSpec,
layer: usize,
input: &Tensor,
device: u32,
) -> Tensor {
let mut keys = layer_common_keys(layer);
match spec.layer_kinds[layer] {
xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
keys.extend(linear_attention_keys(layer));
}
xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
keys.extend(full_attention_keys(layer));
}
}
keys.extend(moe_keys(layer));
let mut shards: Vec<String> = keys
.iter()
.filter_map(|k| weight_map.get(k).cloned())
.collect();
shards.sort();
shards.dedup();
let mut tensors = std::collections::HashMap::new();
for shard in &shards {
tensors.extend(loader::load_safetensors(
&model_dir.join(shard),
Device::Cpu,
));
}
let p = format!("model.language_model.layers.{layer}");
let input_norm = take_tensor(&mut tensors, &format!("{p}.input_layernorm.weight"), device);
let post_norm = take_tensor(
&mut tensors,
&format!("{p}.post_attention_layernorm.weight"),
device,
);
let normed = xserv_kernels::rmsnorm(input, &input_norm, config.ln_eps());
let attn_projected = match spec.layer_kinds[layer] {
xserv_model::qwen35_moe::Qwen35LayerKind::LinearAttention => {
let weights = Qwen35Moe::take_linear_attention_projection_weights(
&mut tensors,
layer,
Device::Cuda(device),
);
let projected = Qwen35Moe::project_linear_attention(spec, &normed, &weights);
let mut state = Tensor::zeros(
&[
spec.linear_value_heads,
spec.linear_value_dim,
spec.linear_value_dim,
],
xserv_tensor::DType::BF16,
Device::Cuda(device),
);
let recurrent =
Qwen35Moe::linear_attention_recurrent_step(spec, &projected, &weights, &mut state);
let (_, out) = Qwen35Moe::finish_linear_attention(
spec,
&recurrent,
&projected.z,
&weights,
config.ln_eps(),
);
out
}
xserv_model::qwen35_moe::Qwen35LayerKind::FullAttention => {
let weights = Qwen35FullAttentionWeights {
q_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.q_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
k_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.k_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
v_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.v_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
o_w_t: take_tensor(
&mut tensors,
&format!("{p}.self_attn.o_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous(),
q_norm: take_tensor(
&mut tensors,
&format!("{p}.self_attn.q_norm.weight"),
device,
),
k_norm: take_tensor(
&mut tensors,
&format!("{p}.self_attn.k_norm.weight"),
device,
),
};
let n_rot = (spec.head_dim as f64 * config.text().partial_rotary_factor.unwrap_or(0.25))
as usize;
let positions = [0u32];
Qwen35Moe::forward_full_attention_layer(
spec,
&normed,
&weights,
&positions,
n_rot,
config.rope_theta_value().unwrap_or(10_000_000.0) as f32,
config.ln_eps(),
)
.projected
}
};
let attn_residual = add(input, &attn_projected);
let moe_in = xserv_kernels::rmsnorm(&attn_residual, &post_norm, config.ln_eps());
let moe_out = run_moe(&mut tensors, layer, &moe_in, spec, device);
add(&attn_residual, &moe_out)
}
fn run_moe(
tensors: &mut std::collections::HashMap<String, Tensor>,
layer: usize,
input: &Tensor,
spec: &Qwen35MoeSpec,
device: u32,
) -> Tensor {
let p = format!("model.language_model.layers.{layer}.mlp");
let router_w = take_tensor(tensors, &format!("{p}.gate.weight"), device)
.transpose(0, 1)
.contiguous();
let shared_gate_w = take_tensor(
tensors,
&format!("{p}.shared_expert.gate_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous();
let shared_up_w = take_tensor(
tensors,
&format!("{p}.shared_expert.up_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous();
let shared_down_w = take_tensor(
tensors,
&format!("{p}.shared_expert.down_proj.weight"),
device,
)
.transpose(0, 1)
.contiguous();
let shared_router_w = take_tensor(tensors, &format!("{p}.shared_expert_gate.weight"), device)
.transpose(0, 1)
.contiguous();
let expert_gate_up =
take_tensor(tensors, &format!("{p}.experts.gate_up_proj"), device).contiguous();
let expert_down = take_tensor(tensors, &format!("{p}.experts.down_proj"), device).contiguous();
let router_logits = matmul(input, &router_w, GemmBackend::CuBlas);
let (topk_ids, topk_weights) =
moe_topk_softmax(&router_logits, spec.num_experts, spec.experts_per_token);
let shared_gate = matmul(input, &shared_gate_w, GemmBackend::CuBlas);
let shared_up = matmul(input, &shared_up_w, GemmBackend::CuBlas);
let shared_act = mul(&silu(&shared_gate), &shared_up);
let shared_down = matmul(&shared_act, &shared_down_w, GemmBackend::CuBlas);
let shared_router = sigmoid(&matmul(input, &shared_router_w, GemmBackend::CuBlas));
let shared_out = row_scale_bf16(&shared_down, &shared_router);
let gate_up = moe_sparse_gemv_bf16(
input,
&expert_gate_up,
&topk_ids,
spec.experts_per_token,
0,
spec.num_experts,
false,
);
let gate = gate_up.narrow(2, 0, spec.moe_intermediate).contiguous();
let up = gate_up
.narrow(2, spec.moe_intermediate, spec.moe_intermediate)
.contiguous();
let routed_act = mul(&silu(&gate), &up);
let down = moe_sparse_gemv_bf16(
&routed_act.reshape(&[spec.experts_per_token, spec.moe_intermediate]),
&expert_down,
&topk_ids,
spec.experts_per_token,
0,
spec.num_experts,
true,
);
let routed_out = moe_weighted_sum_sparse(&down, &topk_ids, &topk_weights, 0, spec.num_experts);
add(&routed_out, &shared_out)
}
fn take_tensor(
tensors: &mut std::collections::HashMap<String, Tensor>,
name: &str,
device: u32,
) -> Tensor {
tensors
.remove(name)
.unwrap_or_else(|| panic!("missing tensor {name}"))
.to_device(Device::Cuda(device))
}
fn layer_common_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}");
vec![
format!("{p}.input_layernorm.weight"),
format!("{p}.post_attention_layernorm.weight"),
]
}
fn linear_attention_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}.linear_attn");
vec![
format!("{p}.in_proj_qkv.weight"),
format!("{p}.in_proj_z.weight"),
format!("{p}.in_proj_a.weight"),
format!("{p}.in_proj_b.weight"),
format!("{p}.conv1d.weight"),
format!("{p}.A_log"),
format!("{p}.dt_bias"),
format!("{p}.norm.weight"),
format!("{p}.out_proj.weight"),
]
}
fn full_attention_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}.self_attn");
vec![
format!("{p}.q_proj.weight"),
format!("{p}.k_proj.weight"),
format!("{p}.v_proj.weight"),
format!("{p}.o_proj.weight"),
format!("{p}.q_norm.weight"),
format!("{p}.k_norm.weight"),
]
}
fn moe_keys(layer: usize) -> Vec<String> {
let p = format!("model.language_model.layers.{layer}.mlp");
vec![
format!("{p}.gate.weight"),
format!("{p}.shared_expert.gate_proj.weight"),
format!("{p}.shared_expert.up_proj.weight"),
format!("{p}.shared_expert.down_proj.weight"),
format!("{p}.shared_expert_gate.weight"),
format!("{p}.experts.gate_up_proj"),
format!("{p}.experts.down_proj"),
]
}
fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
let index_path = model_dir.join("model.safetensors.index.json");
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()));
index
.get("weight_map")
.and_then(|v| v.as_object())
.expect("index must contain weight_map")
.iter()
.map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
.collect()
}
fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
let mut data = Vec::with_capacity(seq_len * hidden);
for i in 0..seq_len * hidden {
let x = ((i % 97) as f32 - 48.0) / 128.0;
data.push(bf16::from_f32(x));
}
Tensor::from_slice(&data, &[seq_len, hidden])
}
fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
let cpu = t.to_device(Device::Cpu);
cpu.as_slice::<bf16>()
.iter()
.take(n)
.map(|v| v.to_f32())
.collect()
}

8
tools/bench/client.py
View File

@@ -110,12 +110,16 @@ async def chat_stream(
choice = choices[0]
delta = choice.get("delta") or {}
content = delta.get("content")
if content:
reasoning_content = delta.get("reasoning_content")
token_text = content if content is not None else reasoning_content
# Exclude the role announcement, but count every generated-token
# chunk, including an empty UTF-8 fragment, in TTFT/TPOT.
if token_text is not None and "role" not in delta:
now = time.perf_counter()
if res.ttft_s < 0:
res.ttft_s = now - t_start
res.chunk_times.append(now)
res.text += content
res.text += token_text
if choice.get("finish_reason"):
res.finish_reason = choice["finish_reason"]
except Exception as e: # noqa: BLE001 — surface any failure to the report

10
tools/bench/config.py
View File

@@ -25,10 +25,7 @@ class SystemEndpoint:
model_id: str # what to put in the request body's "model" field
api_key: str | None = None # llama-server doesn't need one; xserv ignores it
# Extra fields merged into every request body for this system. Used to keep
# the two engines in the SAME generation mode — xserv hardcodes Qwen3
# thinking OFF (empty <think></think> in its prompt builder), so we disable
# thinking on llama-server via chat_template_kwargs to match. Both engines
# ignore unknown fields, so this is safe.
# both engines in the same chat-template generation mode.
extra_body: dict | None = None
# Process supervision is optional — if base_url is already serving, we skip launch.
launch_cmd: list[str] | None = None
@@ -49,7 +46,12 @@ class BenchConfig:
quality_max_tokens_aime: int = 16384
quality_max_tokens_gsm8k: int = 2048
quality_limit: int | None = None # subsample for smoke tests; None = all
quality_seed: int | None = None
quality_temperature: float = 0.0
quality_top_k: int = 0
quality_top_p: float = 1.0
quality_presence_penalty: float = 0.0
quality_repetition_penalty: float = 1.0
request_timeout_s: float = 1800.0

25
tools/bench/quality.py
View File

@@ -14,6 +14,7 @@ extra moving parts aren't worth it for the first iteration.
from __future__ import annotations
import asyncio
import random
import statistics
import time
from dataclasses import asdict, dataclass
@@ -38,6 +39,7 @@ class QualityRow:
n_total: int
n_correct: int
n_errors: int
n_length: int
accuracy: float
mean_completion_tokens: float
mean_ttft_ms: float
@@ -58,7 +60,9 @@ class QualityCase:
tpot_ms: float
e2e_s: float
error: str | None
finish_reason: str | None
response_preview: str
response_text: str
async def _run_one_task(
@@ -66,7 +70,10 @@ async def _run_one_task(
) -> tuple[QualityRow, list[QualityCase]]:
problems = task_mod.load()
if cfg.quality_limit is not None:
problems = problems[: cfg.quality_limit]
if cfg.quality_seed is not None and cfg.quality_limit < len(problems):
problems = random.Random(cfg.quality_seed).sample(problems, cfg.quality_limit)
else:
problems = problems[: cfg.quality_limit]
print(f"[quality] {ep.name} / {task_name}: {len(problems)} problems "
f"(max_tokens={max_tokens})")
@@ -75,13 +82,20 @@ async def _run_one_task(
async with httpx.AsyncClient(timeout=cfg.request_timeout_s) as client:
for prob in problems:
messages = task_mod.make_messages(prob["problem"])
extra_body = dict(ep.extra_body or {})
extra_body.update({
"top_k": cfg.quality_top_k,
"top_p": cfg.quality_top_p,
"presence_penalty": cfg.quality_presence_penalty,
"repetition_penalty": cfg.quality_repetition_penalty,
})
r = await chat_stream(
client, ep.base_url, ep.model_id, messages,
max_tokens=max_tokens,
temperature=cfg.quality_temperature,
api_key=ep.api_key,
timeout=cfg.request_timeout_s,
extra_body=ep.extra_body,
extra_body=extra_body,
)
pred = task_mod.extract_answer(r.text) if r.error is None else None
correct = task_mod.score(pred, prob["answer"]) if r.error is None else False
@@ -91,8 +105,9 @@ async def _run_one_task(
correct=correct, completion_tokens=r.completion_tokens,
ttft_ms=r.ttft_s * 1000 if r.ttft_s > 0 else -1.0,
tpot_ms=r.tpot_s * 1000 if r.tpot_s > 0 else -1.0,
e2e_s=r.e2e_s, error=r.error,
e2e_s=r.e2e_s, error=r.error, finish_reason=r.finish_reason,
response_preview=(r.text or "")[:240].replace("\n", " "),
response_text=r.text or "",
))
mark = "✓" if correct else ("E" if r.error else "✗")
print(f" [{mark}] {prob['id']:>4s} gold={prob['answer']:>6s} "
@@ -109,7 +124,9 @@ async def _run_one_task(
n_total=len(cases),
n_correct=correct,
n_errors=errors,
accuracy=correct / max(len(cases) - errors, 1),
n_length=sum(1 for c in cases if c.finish_reason == "length"),
# Transport/runtime failures are incorrect attempts, not exclusions.
accuracy=correct / max(len(cases), 1),
mean_completion_tokens=statistics.mean(c.completion_tokens for c in ok) if ok else 0.0,
mean_ttft_ms=statistics.mean(c.ttft_ms for c in ok if c.ttft_ms > 0) if ok else -1.0,
mean_tpot_ms=statistics.mean(c.tpot_ms for c in ok if c.tpot_ms > 0) if ok else -1.0,

12
tools/bench/report.py
View File

@@ -32,7 +32,11 @@ def _speed_table(rows: list[dict[str, Any]]) -> str:
by = {(r["system"], r["scenario"]): r for r in rows}
out = []
out.append("| scenario | metric | " + " | ".join(systems) + " | speedup (xserv ÷ llama.cpp) |")
out.append(
"| scenario | metric | "
+ " | ".join(systems)
+ " | xserv relative performance (higher is better) |"
)
out.append("|---|---|" + "|".join(["---"] * (len(systems) + 1)) + "|")
metrics = [
@@ -68,13 +72,13 @@ def _quality_table(rows: list[dict[str, Any]]) -> str:
for r in rows:
by_task.setdefault(r["task"], []).append(r)
out: list[str] = []
out.append("| task | system | n | correct | accuracy | mean tokens | TTFT (ms) | TPOT (ms/tok) | wall (s) |")
out.append("|---|---|---|---|---|---|---|---|---|")
out.append("| task | system | n | correct | accuracy | length | mean tokens | TTFT (ms) | TPOT (ms/tok) | wall (s) |")
out.append("|---|---|---|---|---|---|---|---|---|---|")
for task, task_rows in by_task.items():
for r in task_rows:
out.append(
f"| {task} | {r['system']} | {r['n_total']} | {r['n_correct']} | "
f"{r['accuracy'] * 100:.1f}% | {r['mean_completion_tokens']:.0f} | "
f"{r['accuracy'] * 100:.1f}% | {r['n_length']} | {r['mean_completion_tokens']:.0f} | "
f"{_fmt(r['mean_ttft_ms'])} | {_fmt(r['mean_tpot_ms'], 2)} | {r['wall_s']:.1f} |"
)
return "\n".join(out) + "\n"

33
tools/bench/runner.py
View File

@@ -88,6 +88,15 @@ def parse_args() -> argparse.Namespace:
p.add_argument("--quality-tasks", default="aime2025,gsm8k")
p.add_argument("--quality-limit", type=int, default=None,
help="Cap problems per task (smoke test). None = all problems.")
p.add_argument("--quality-seed", type=int, default=None,
help="Fixed seed for a random quality subset; default uses dataset prefix.")
p.add_argument("--quality-max-tokens-aime", type=int, default=16384)
p.add_argument("--quality-max-tokens-gsm8k", type=int, default=2048)
p.add_argument("--quality-temperature", type=float, default=0.0)
p.add_argument("--quality-top-k", type=int, default=0)
p.add_argument("--quality-top-p", type=float, default=1.0)
p.add_argument("--quality-presence-penalty", type=float, default=0.0)
p.add_argument("--quality-repetition-penalty", type=float, default=1.0)
p.add_argument("--speed-prompts", type=int, default=8)
p.add_argument("--speed-max-tokens", type=int, default=128)
p.add_argument("--speed-concurrency", default="1,2,4,8")
@@ -99,12 +108,16 @@ def parse_args() -> argparse.Namespace:
def build_endpoints(args) -> list[SystemEndpoint]:
wanted = set(s.strip() for s in args.systems.split(",") if s.strip())
eps: list[SystemEndpoint] = []
thinking_extra_body = None if args.enable_thinking else {
"chat_template_kwargs": {"enable_thinking": False}
}
if SYSTEM_XSERV in wanted:
if args.xserv_base_url:
eps.append(SystemEndpoint(
name=SYSTEM_XSERV, base_url=args.xserv_base_url,
model_id=args.xserv_model_id, launch_cmd=None,
extra_body=thinking_extra_body,
))
else:
model_dir = args.xserv_model or os.environ.get("XSERV_MODEL_DIR")
@@ -120,19 +133,15 @@ def build_endpoints(args) -> list[SystemEndpoint]:
),
health_path="/health",
ready_timeout_s=1200.0,
extra_body=thinking_extra_body,
))
# Match xserv's hardcoded thinking-OFF mode unless explicitly overridden.
llama_extra_body = None if args.enable_thinking else {
"chat_template_kwargs": {"enable_thinking": False}
}
if SYSTEM_LLAMA_CPP in wanted:
if args.llama_base_url:
eps.append(SystemEndpoint(
name=SYSTEM_LLAMA_CPP, base_url=args.llama_base_url,
model_id=args.llama_model_id, launch_cmd=None,
extra_body=llama_extra_body,
extra_body=thinking_extra_body,
))
else:
gguf = args.llama_gguf or os.environ.get("LLAMA_GGUF")
@@ -161,7 +170,7 @@ def build_endpoints(args) -> list[SystemEndpoint]:
# llama-server's health endpoint also returns 200 only when model is loaded.
health_path="/health",
ready_timeout_s=1200.0,
extra_body=llama_extra_body,
extra_body=thinking_extra_body,
))
return eps
@@ -194,7 +203,15 @@ def main() -> None:
speed_prompts=args.speed_prompts,
speed_max_tokens=args.speed_max_tokens,
speed_concurrency=tuple(int(c) for c in args.speed_concurrency.split(",") if c.strip()),
quality_max_tokens_aime=args.quality_max_tokens_aime,
quality_max_tokens_gsm8k=args.quality_max_tokens_gsm8k,
quality_limit=args.quality_limit,
quality_seed=args.quality_seed,
quality_temperature=args.quality_temperature,
quality_top_k=args.quality_top_k,
quality_top_p=args.quality_top_p,
quality_presence_penalty=args.quality_presence_penalty,
quality_repetition_penalty=args.quality_repetition_penalty,
)
os.makedirs(args.out_dir, exist_ok=True)
@@ -229,7 +246,7 @@ def main() -> None:
speed_raw=speed_raw,
quality_rows=q_rows_to_dicts(quality_rows) if quality_rows else [],
quality_cases=cases_to_dicts(quality_cases) if quality_cases else [],
env=collect_env(),
env={**collect_env(), "benchmark_args": vars(args)},
)

16
tools/bench/speed.py
View File

@@ -35,6 +35,7 @@ class SpeedRow:
scenario: str # e.g. "single/short", "concurrent-4"
requests: int
completion_tokens_total: int
prompt_tokens_mean: float
wall_s: float
ttft_ms_p50: float
ttft_ms_p95: float
@@ -64,6 +65,7 @@ def _summarize(system: str, scenario: str, results: list[StreamResult], wall_s:
scenario=scenario,
requests=len(results),
completion_tokens_total=total_tokens,
prompt_tokens_mean=(statistics.mean(r.prompt_tokens for r in ok) if ok else 0.0),
wall_s=wall_s,
ttft_ms_p50=_percentile(ttft_ms, 50),
ttft_ms_p95=_percentile(ttft_ms, 95),
@@ -82,6 +84,18 @@ async def run_single_stream(
rows: list[SpeedRow] = []
raw: list[dict[str, Any]] = []
for bucket, prompt in SPEED_PROMPTS.items():
# Prefill kernels/graphs can be shape-specific. Warm each prompt shape
# twice so p95 does not accidentally report one-time graph setup.
for _ in range(2):
await chat_concurrent(
ep.base_url, ep.model_id, [[{"role": "user", "content": prompt}]],
max_tokens=cfg.speed_max_tokens,
temperature=0.0,
api_key=ep.api_key,
timeout=cfg.request_timeout_s,
concurrency=1,
extra_body=ep.extra_body,
)
messages = [[{"role": "user", "content": prompt}]] * cfg.speed_prompts
results, wall = await chat_concurrent(
ep.base_url, ep.model_id, messages,
@@ -98,6 +112,7 @@ async def run_single_stream(
"system": ep.name, "scenario": f"single/{bucket}", "i": i,
"ttft_s": r.ttft_s, "tpot_s": r.tpot_s,
"completion_tokens": r.completion_tokens,
"prompt_tokens": r.prompt_tokens,
"e2e_s": r.e2e_s, "error": r.error,
"finish_reason": r.finish_reason,
})
@@ -131,6 +146,7 @@ async def run_concurrent(
"system": ep.name, "scenario": f"concurrent-{c}", "i": i,
"ttft_s": r.ttft_s, "tpot_s": r.tpot_s,
"completion_tokens": r.completion_tokens,
"prompt_tokens": r.prompt_tokens,
"e2e_s": r.e2e_s, "error": r.error,
"finish_reason": r.finish_reason,
})