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

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()
}