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