515 lines
18 KiB
Rust
515 lines
18 KiB
Rust
//! Qwen SparseMoeBlock as ordinary Array nodes, with the original two
|
|
//! CustomKernel nodes only where MTPLX selects fused decode/verify.
|
|
use super::array::{Array, Dtype, Layout};
|
|
use super::binary::{Binary, binary};
|
|
use super::mlp::{self, QuantizedLinear};
|
|
use super::stream::{Device, Stream, Streams};
|
|
use super::{indexing, linear, normalization, ops, reduce, sort, unary};
|
|
|
|
pub(super) enum GateUp {
|
|
Fused(QuantizedLinear),
|
|
Separate {
|
|
gate: QuantizedLinear,
|
|
up: QuantizedLinear,
|
|
},
|
|
}
|
|
|
|
pub(super) struct FeedForward {
|
|
pub(super) input: GateUp,
|
|
pub(super) down: QuantizedLinear,
|
|
}
|
|
|
|
fn gather(
|
|
input: &Array,
|
|
layer: &QuantizedLinear,
|
|
ids: &Array,
|
|
sorted: bool,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
linear::guarded_gather_qmm(
|
|
input,
|
|
[&layer.weight, &layer.scales, &layer.biases],
|
|
[None, Some(ids)],
|
|
[layer.bits, layer.group],
|
|
sorted,
|
|
streams,
|
|
stream,
|
|
)
|
|
}
|
|
|
|
impl FeedForward {
|
|
fn shared(&self, input: &Array, stream: Stream) -> Result<Array, String> {
|
|
match &self.input {
|
|
GateUp::Fused(gu) => mlp::fused_gate_up_mlp(input, gu, &self.down, stream),
|
|
GateUp::Separate { gate, up } => {
|
|
// Qwen3NextMLP evaluates gate first, unlike SwitchGLU below.
|
|
let gate = gate.apply(input, stream)?;
|
|
let up = up.apply(input, stream)?;
|
|
self.down.apply(&mlp::swiglu(&gate, &up, stream)?, stream)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn switched(
|
|
&self,
|
|
input: &Array,
|
|
indices: &Array,
|
|
sorted_gather: bool,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
let do_sort = indices.layout().size() >= 64;
|
|
let mut x = if do_sort {
|
|
input.clone()
|
|
} else {
|
|
ops::expand_dims(input, &[-2, -3], stream)?
|
|
};
|
|
let mut ids = indices.clone();
|
|
let mut inverse = None;
|
|
let mut row_map = None;
|
|
if do_sort {
|
|
let size =
|
|
i32::try_from(indices.layout().size()).map_err(|_| "MoE index size overflow")?;
|
|
ids = ops::reshape(&ids, &[size], stream)?;
|
|
let order = sort::sort(&ids, Some(-1), true, None, stream)?;
|
|
inverse = Some(sort::sort(&order, Some(-1), true, None, stream)?);
|
|
let topk = *indices.layout().shape().last().ok_or("MoE indices rank")? as u32;
|
|
let buffer = super::super::gpu::Buffer::mtplx_bytes(4)?;
|
|
buffer.write(0, &topk.to_le_bytes())?;
|
|
let divisor = Array::new(&[], Dtype::U32, buffer)?;
|
|
let rows = super::binary::floor_divide_integer(&order, &divisor, stream)?;
|
|
let k = *input.layout().shape().last().ok_or("MoE input rank")?;
|
|
let flat_rows = i32::try_from(input.layout().size() / k as usize)
|
|
.map_err(|_| "MoE row count overflow")?;
|
|
x = ops::reshape(&x, &[flat_rows, 1, k], stream)?;
|
|
row_map = Some(rows);
|
|
ids = indexing::take(&ids, &order, 0, stream)?;
|
|
}
|
|
let hidden = if let Some(map) = row_map {
|
|
let (gate, up) = match &self.input {
|
|
GateUp::Fused(gu) => (gu, None),
|
|
GateUp::Separate { gate, up } => (gate, Some(up)),
|
|
};
|
|
super::moe_sorted::swiglu_rows(
|
|
&x,
|
|
&map,
|
|
gate,
|
|
up,
|
|
&ids,
|
|
sorted_gather,
|
|
streams,
|
|
stream,
|
|
)?
|
|
} else {
|
|
match &self.input {
|
|
GateUp::Fused(gu) => {
|
|
let gu = gather(&x, gu, &ids, do_sort, streams, stream)?;
|
|
let n = *gu.layout().shape().last().ok_or("MoE gate/up rank")?;
|
|
if n % 2 != 0 {
|
|
return Err("MoE fused gate/up width must be even".into());
|
|
}
|
|
let pair = ops::split(&gu, &[n / 2], -1, stream)?;
|
|
binary(
|
|
&mlp::silu(&pair[0], stream)?,
|
|
&pair[1],
|
|
Binary::Multiply,
|
|
stream,
|
|
)?
|
|
}
|
|
GateUp::Separate { gate, up } => {
|
|
let up = gather(&x, up, &ids, do_sort, streams, stream)?;
|
|
let gate = gather(&x, gate, &ids, do_sort, streams, stream)?;
|
|
mlp::swiglu(&gate, &up, stream)?
|
|
}
|
|
}
|
|
};
|
|
let mut output = gather(&hidden, &self.down, &ids, do_sort, streams, stream)?;
|
|
if let Some(inverse) = inverse {
|
|
output = indexing::take(&output, &inverse, 0, stream)?;
|
|
let mut shape = indices.layout().shape().to_vec();
|
|
shape.extend_from_slice(&output.layout().shape()[1..]);
|
|
output = ops::reshape(&output, &shape, stream)?;
|
|
}
|
|
ops::squeeze(&output, Some(&[-2]), stream)
|
|
}
|
|
}
|
|
|
|
pub(super) struct SparseMoe {
|
|
pub(super) experts: FeedForward,
|
|
pub(super) shared: FeedForward,
|
|
pub(super) router: QuantizedLinear,
|
|
pub(super) shared_gate: QuantizedLinear,
|
|
pub(super) topk: i32,
|
|
pub(super) normalize: bool,
|
|
pub(super) sorted_gather: bool,
|
|
pub(super) fused_decode: bool,
|
|
pub(super) fused_verify: bool,
|
|
}
|
|
|
|
impl SparseMoe {
|
|
fn fused_rows(&self, input: &Array) -> Option<u32> {
|
|
let GateUp::Fused(gu) = &self.experts.input else {
|
|
return None;
|
|
};
|
|
if gu.bits != 4
|
|
|| !matches!(gu.group, 32 | 64)
|
|
|| self.experts.down.bits != 4
|
|
|| !matches!(self.experts.down.group, 32 | 64)
|
|
{
|
|
return None;
|
|
}
|
|
let l = input.layout();
|
|
let shape = l.shape();
|
|
let rows = shape[shape.len() - 2];
|
|
let k = shape[shape.len() - 1];
|
|
if l.size() != rows as usize * k as usize {
|
|
return None;
|
|
}
|
|
((rows == 1 && self.fused_decode) || ((2..=4).contains(&rows) && self.fused_verify))
|
|
.then_some(rows as u32)
|
|
}
|
|
|
|
pub(super) fn apply(
|
|
&self,
|
|
input: &Array,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
if stream.device() != Device::Gpu
|
|
|| input.layout().dtype() != Dtype::BF16
|
|
|| input.layout().shape().len() < 2
|
|
|| input.layout().size() == 0
|
|
{
|
|
return Err("SparseMoe requires nonempty BF16 GPU rows".into());
|
|
}
|
|
let fused = self.fused_rows(input);
|
|
let gates = normalization::softmax_last(&self.router.apply(input, stream)?, true, stream)?;
|
|
let shape = gates.layout().shape().to_vec();
|
|
let rank = shape.len();
|
|
let experts = shape[rank - 1];
|
|
if self.topk < 1 || self.topk > experts {
|
|
return Err("invalid MoE top-k".into());
|
|
}
|
|
let partition = sort::sort(&gates, Some(-1), true, Some(-self.topk), stream)?;
|
|
let mut starts = vec![0; rank];
|
|
starts[rank - 1] = experts - self.topk;
|
|
let indices = ops::slice(&partition, &starts, &shape, &vec![1; rank], stream)?;
|
|
let mut scores = indexing::take_along_axis(&gates, &indices, -1, stream)?;
|
|
if self.normalize {
|
|
scores = binary(
|
|
&scores,
|
|
&reduce::sum(&scores, &[-1], true, stream)?,
|
|
Binary::Divide,
|
|
stream,
|
|
)?;
|
|
}
|
|
let y = if let Some(rows) = fused {
|
|
let GateUp::Fused(gu) = &self.experts.input else {
|
|
unreachable!()
|
|
};
|
|
fused_experts(
|
|
input,
|
|
gu,
|
|
&self.experts.down,
|
|
&indices,
|
|
&scores,
|
|
rows,
|
|
self.topk,
|
|
stream,
|
|
)?
|
|
} else {
|
|
let switched =
|
|
self.experts
|
|
.switched(input, &indices, self.sorted_gather, streams, stream)?;
|
|
let scores = ops::expand_dims(&scores, &[-1], stream)?;
|
|
reduce::sum(
|
|
&binary(&switched, &scores, Binary::Multiply, stream)?,
|
|
&[-2],
|
|
false,
|
|
stream,
|
|
)?
|
|
};
|
|
let gate = || {
|
|
unary::unary(
|
|
&self.shared_gate.apply(input, stream)?,
|
|
unary::Unary::Sigmoid,
|
|
stream,
|
|
)
|
|
};
|
|
let (gate, shared) = if fused.is_some() {
|
|
(gate()?, self.shared.shared(input, stream)?)
|
|
} else {
|
|
let shared = self.shared.shared(input, stream)?;
|
|
(gate()?, shared)
|
|
};
|
|
let y = binary(
|
|
&y,
|
|
&binary(&gate, &shared, Binary::Multiply, stream)?,
|
|
Binary::Add,
|
|
stream,
|
|
)?;
|
|
if fused.is_some() {
|
|
ops::astype(&y, input.layout().dtype(), false, stream)
|
|
} else {
|
|
Ok(y)
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)] // Original two CustomKernel call inputs and template dimensions.
|
|
fn fused_experts(
|
|
input: &Array,
|
|
gu: &QuantizedLinear,
|
|
down: &QuantizedLinear,
|
|
indices: &Array,
|
|
scores: &Array,
|
|
rows: u32,
|
|
topk: i32,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
if topk != 10
|
|
|| gu.weight.layout().shape().get(1) != Some(&1280)
|
|
|| down.weight.layout().shape().get(1) != Some(&2560)
|
|
|| input.layout().shape().last() != Some(&2560)
|
|
{
|
|
return Err("fused MoE exports require the installed top10/2560/640 shape".into());
|
|
}
|
|
let flatten = |a: &Array| ops::reshape(a, &[a.layout().size() as i32], stream);
|
|
let x = if rows == 1 {
|
|
flatten(input)?
|
|
} else {
|
|
ops::reshape(input, &[rows as i32, 2560], stream)?
|
|
};
|
|
let ids = ops::astype(&flatten(indices)?, Dtype::U32, false, stream)?;
|
|
let scores = ops::astype(&flatten(scores)?, Dtype::F32, false, stream)?;
|
|
let hshape = if rows == 1 {
|
|
vec![10, 640]
|
|
} else {
|
|
vec![rows as i32, 10, 640]
|
|
};
|
|
let h = Array::make_operation_with_inputs(
|
|
stream,
|
|
&[
|
|
x,
|
|
gu.weight.clone(),
|
|
gu.scales.clone(),
|
|
gu.biases.clone(),
|
|
ids.clone(),
|
|
],
|
|
Layout::new(&hshape, Dtype::BF16)?,
|
|
ops::Operation::MoeGluH {
|
|
rows,
|
|
group: gu.group,
|
|
},
|
|
)?;
|
|
let yshape = if rows == 1 {
|
|
vec![2560]
|
|
} else {
|
|
vec![rows as i32, 2560]
|
|
};
|
|
let y = Array::make_operation_with_inputs(
|
|
stream,
|
|
&[
|
|
h,
|
|
down.weight.clone(),
|
|
down.scales.clone(),
|
|
down.biases.clone(),
|
|
ids,
|
|
scores,
|
|
],
|
|
Layout::new(&yshape, Dtype::BF16)?,
|
|
ops::Operation::MoeDownY {
|
|
rows,
|
|
group: down.group,
|
|
},
|
|
)?;
|
|
ops::reshape(&y, input.layout().shape(), stream)
|
|
}
|
|
|
|
pub(super) fn fused_evaluate(
|
|
inputs: &[Array],
|
|
outputs: &[Array],
|
|
rows: u32,
|
|
group: u32,
|
|
down: bool,
|
|
) -> Result<(), String> {
|
|
if !(1..=4).contains(&rows) || !matches!(group, 32 | 64) {
|
|
return Err("invalid fused MoE specialization".into());
|
|
}
|
|
let suffix = if rows == 1 {
|
|
format!("g{group}_bf16")
|
|
} else {
|
|
format!("m_m{rows}_g{group}_bf16")
|
|
};
|
|
let name = format!(
|
|
"kernel_qwen_mtplx_moe_{}_{suffix}",
|
|
if down { "down_y" } else { "glu_h" }
|
|
);
|
|
let count = rows * if down { 2560 } else { 10 * 640 };
|
|
super::gdn::custom_evaluate(inputs, outputs, |bindings| {
|
|
super::dispatch_geometry(
|
|
&name,
|
|
bindings,
|
|
&[],
|
|
[count.div_ceil(32) * 1024, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)
|
|
})
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple M5 Max; full MoE Array graph, no model weights; use test-supervisor"]
|
|
fn mtplx_graph_sparse_moe_matches_reference() {
|
|
use super::super::{
|
|
gpu::Context,
|
|
qwen_mtplx_tests::{capture_dispatches, pattern, verify},
|
|
};
|
|
use super::{allocator, configure_sources};
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
|
|
let gpu = streams.default_stream(Device::Gpu).unwrap();
|
|
let fixtures = include_str!("../../../../tests/fixtures/mtplx-custom-kernels.jsonl")
|
|
.lines()
|
|
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
|
|
.filter(|row| row["kernel"].as_str().unwrap().starts_with("moe_"))
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(fixtures.len(), 220);
|
|
for prepared in [false, true] {
|
|
for row in &fixtures {
|
|
let [tokens, experts, k, intermediate, gu_group, down_group] = [
|
|
"tokens",
|
|
"experts",
|
|
"k",
|
|
"intermediate",
|
|
"gu_group",
|
|
"down_group",
|
|
]
|
|
.map(|key| row[key].as_u64().unwrap() as u32);
|
|
eprintln!("MoE Array graph prepared={prepared} {}", row["kernel"]);
|
|
let pack = |e: u32, n: u32, k: u32, group: u32, salt: u32, expert: bool| {
|
|
let mut q = mlp::fixture_linear(e * n, k, 4, group, salt);
|
|
if expert {
|
|
for a in [&mut q.weight, &mut q.scales, &mut q.biases] {
|
|
let width = *a.layout().shape().last().unwrap();
|
|
*a = ops::reshape(a, &[e as i32, n as i32, width], gpu).unwrap();
|
|
}
|
|
}
|
|
q
|
|
};
|
|
let ff = |expert: bool, separate: bool| {
|
|
let e = if expert { experts } else { 1 };
|
|
let salt = if expert { 39 } else { 46 };
|
|
FeedForward {
|
|
input: if separate {
|
|
GateUp::Separate {
|
|
gate: pack(e, intermediate, k, gu_group, salt, expert),
|
|
up: pack(
|
|
e,
|
|
intermediate,
|
|
k,
|
|
down_group,
|
|
if expert { 45 } else { 54 },
|
|
expert,
|
|
),
|
|
}
|
|
} else {
|
|
GateUp::Fused(pack(e, 2 * intermediate, k, gu_group, salt, expert))
|
|
},
|
|
down: pack(
|
|
e,
|
|
k,
|
|
intermediate,
|
|
down_group,
|
|
if expert { 41 } else { 48 },
|
|
expert,
|
|
),
|
|
}
|
|
};
|
|
let model = SparseMoe {
|
|
experts: ff(true, row["separate_experts"].as_bool().unwrap_or(false)),
|
|
shared: ff(false, row["separate_shared"].as_bool().unwrap_or(false)),
|
|
router: mlp::fixture_linear(experts, k, 8, 64, 50),
|
|
shared_gate: mlp::fixture_linear(1, k, 4, 64, 52),
|
|
topk: 10,
|
|
normalize: row["normalize"].as_bool().unwrap(),
|
|
sorted_gather: true,
|
|
fused_decode: row["fused"].as_bool().unwrap_or(false),
|
|
fused_verify: row["fused"].as_bool().unwrap_or(false),
|
|
};
|
|
if prepared {
|
|
let mut parameters = Vec::new();
|
|
let mut add = |q: &QuantizedLinear| {
|
|
parameters.extend([q.weight.clone(), q.scales.clone(), q.biases.clone()]);
|
|
};
|
|
for ff in [&model.experts, &model.shared] {
|
|
match &ff.input {
|
|
GateUp::Fused(gu) => add(gu),
|
|
GateUp::Separate { gate, up } => {
|
|
add(gate);
|
|
add(up);
|
|
}
|
|
}
|
|
add(&ff.down);
|
|
}
|
|
add(&model.router);
|
|
add(&model.shared_gate);
|
|
ops::evaluate(&streams, ¶meters, gpu, false).unwrap();
|
|
}
|
|
let input = Array::new(
|
|
&[1, tokens as i32, k as i32],
|
|
Dtype::BF16,
|
|
pattern(tokens * k, 44),
|
|
)
|
|
.unwrap();
|
|
let output = model.apply(&input, &streams, gpu).unwrap();
|
|
let fused = row["fused"].as_bool().unwrap_or(false)
|
|
&& !row["separate_experts"].as_bool().unwrap_or(false)
|
|
&& tokens <= 4;
|
|
assert_eq!(model.fused_rows(&input).is_some(), fused);
|
|
// Graph edges, not an external model/parameter table, keep inputs alive.
|
|
drop((input, model));
|
|
let (_, dispatches) = capture_dispatches(|| {
|
|
ops::evaluate(&streams, std::slice::from_ref(&output), gpu, false).unwrap();
|
|
});
|
|
let actual = dispatches
|
|
.into_iter()
|
|
.filter(|call| call.0.starts_with("kernel_qwen_mtplx_moe_"))
|
|
.collect::<Vec<_>>();
|
|
let expected = if fused {
|
|
[
|
|
("glu_h", gu_group, tokens * 10 * 640),
|
|
("down_y", down_group, tokens * 2560),
|
|
]
|
|
.map(|(name, group, rows)| {
|
|
let suffix = if tokens == 1 {
|
|
format!("g{group}_bf16")
|
|
} else {
|
|
format!("m_m{tokens}_g{group}_bf16")
|
|
};
|
|
(
|
|
format!("kernel_qwen_mtplx_moe_{name}_{suffix}"),
|
|
[rows.div_ceil(32) * 1024, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)
|
|
})
|
|
.to_vec()
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
assert_eq!(
|
|
actual, expected,
|
|
"original CustomKernel dispatch geometry/order"
|
|
);
|
|
verify(
|
|
row["kernel"].as_str().unwrap(),
|
|
&[(&output.buffer(), tokens * k)],
|
|
);
|
|
}
|
|
}
|
|
streams.clear_streams().unwrap();
|
|
}
|