194 lines
7.5 KiB
Rust
194 lines
7.5 KiB
Rust
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()
|
|
}
|