500 lines
16 KiB
Rust
500 lines
16 KiB
Rust
//! MTPLX runtime vector SDPA, BF16/D256 on the evaluation M5 Max (applegpu_g17s).
|
|
//! Original 1/2-pass kernels. Full attention and surrounding model routing are open.
|
|
use super::super::gpu::{Buffer, MetalConstant};
|
|
use super::array::{Array, Dtype, Layout};
|
|
use super::stream::{Device, Stream};
|
|
use super::{bytes, dispatch_specialized, encoder, ops, tensor, views};
|
|
|
|
pub(super) fn vector(
|
|
q: &Array,
|
|
k: &Array,
|
|
v: &Array,
|
|
mask: Option<&Array>,
|
|
scale: f32,
|
|
causal: bool,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
let qs = q.layout().shape().to_vec();
|
|
let ks = k.layout().shape().to_vec();
|
|
if stream.device() != Device::Gpu
|
|
|| qs.len() != 4
|
|
|| ks.len() != 4
|
|
|| k.layout().shape() != v.layout().shape()
|
|
|| qs[0] != ks[0]
|
|
|| qs[3] != 256
|
|
|| ks[3] != 256
|
|
|| ks[1] <= 0
|
|
|| qs[1] % ks[1] != 0
|
|
|| !(1..=8).contains(&qs[2])
|
|
|| qs[2] > ks[2]
|
|
|| qs[2] * (qs[1] / ks[1]) > 32
|
|
|| [q, k, v]
|
|
.iter()
|
|
.any(|a| a.layout().dtype() != Dtype::BF16 || a.layout().size() == 0)
|
|
|| !scale.is_finite()
|
|
|| (causal && mask.is_some())
|
|
{
|
|
return Err(
|
|
"vector SDPA requires compatible BF16/D256 GPU inputs, S<=8 and S*GQA<=32".into(),
|
|
);
|
|
}
|
|
let mut inputs = vec![q.clone(), k.clone(), v.clone()];
|
|
if let Some(mask) = mask {
|
|
let dtype = mask.layout().dtype();
|
|
if dtype.promote(Dtype::BF16) != Dtype::BF16 {
|
|
return Err("SDPA mask must promote to BF16".into());
|
|
}
|
|
let mask = ops::astype(
|
|
mask,
|
|
if dtype == Dtype::Bool {
|
|
dtype
|
|
} else {
|
|
Dtype::BF16
|
|
},
|
|
false,
|
|
stream,
|
|
)?;
|
|
inputs.push(ops::broadcast_to(
|
|
&mask,
|
|
&[qs[0], qs[1], qs[2], ks[2]],
|
|
stream,
|
|
)?);
|
|
}
|
|
Array::make_operation_with_inputs(
|
|
stream,
|
|
&inputs,
|
|
Layout::new(&qs, Dtype::BF16)?,
|
|
ops::Operation::SdpaVector { scale, causal },
|
|
)
|
|
}
|
|
|
|
fn copy(input: &Array) -> Result<Array, String> {
|
|
let layout = input.layout();
|
|
let out = Array::new(
|
|
layout.shape(),
|
|
layout.dtype(),
|
|
Buffer::mtplx_bytes(layout.nbytes() as u64)?,
|
|
)?;
|
|
drop(layout);
|
|
views::general_copy_inplace(input, &out)?;
|
|
Ok(out)
|
|
}
|
|
|
|
fn blocks(n: i32, simds: i32, override_value: i32, arch: Option<char>) -> Result<i32, String> {
|
|
let default = if arch == Some('d') {
|
|
if simds <= 2 && n > 8192 {
|
|
256
|
|
} else if simds >= 6 && n >= 65536 {
|
|
1024
|
|
} else if simds >= 6 && n >= 16384 {
|
|
512
|
|
} else {
|
|
128
|
|
}
|
|
} else if arch != Some('s') {
|
|
if simds >= 4 { 64 } else { 32 }
|
|
} else if n <= 1024 || simds <= 4 {
|
|
64
|
|
} else if n <= 8192 {
|
|
128
|
|
} else if n <= 32768 {
|
|
256
|
|
} else if n <= 65536 {
|
|
512
|
|
} else {
|
|
1024
|
|
};
|
|
if override_value > 0 {
|
|
override_value
|
|
.checked_add(31)
|
|
.map(|v| v / 32 * 32)
|
|
.ok_or("SDPA block count overflow".into())
|
|
} else {
|
|
Ok(default)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple M5 Max; actual pinned fused SDPA kernels"]
|
|
fn mtplx_sdpa_vector_matches_reference() {
|
|
use super::super::{
|
|
gpu::Context,
|
|
qwen_mtplx_tests::{capture_dispatches, pattern, verify},
|
|
};
|
|
use super::{allocator, configure_sources, stream::Streams};
|
|
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();
|
|
assert!(
|
|
std::env::var("MLX_SDPA_BLOCKS").is_err(),
|
|
"reference receipt uses default blocks"
|
|
);
|
|
for (n, s, over, want) in [
|
|
(1024, 24, 0, 64),
|
|
(16383, 24, 0, 256),
|
|
(16384, 24, 0, 256),
|
|
(65535, 24, 0, 512),
|
|
(65536, 24, 0, 512),
|
|
(65537, 24, 0, 1024),
|
|
(8193, 2, 0, 64),
|
|
(1, 1, 33, 64),
|
|
(1, 1, 31, 32),
|
|
] {
|
|
assert_eq!(blocks(n, s, over, Some('s')).unwrap(), want);
|
|
}
|
|
let cases = include_str!("../../../../tests/fixtures/mtplx-sdpa-vector.jsonl")
|
|
.lines()
|
|
.map(|s| serde_json::from_str::<serde_json::Value>(s).unwrap())
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(cases.len(), 48);
|
|
let bf16 = |shape: &[i32], salt| {
|
|
Array::new(
|
|
shape,
|
|
Dtype::BF16,
|
|
pattern(shape.iter().map(|&v| v as u32).product(), salt),
|
|
)
|
|
.unwrap()
|
|
};
|
|
for prepared in [false, true] {
|
|
for case in &cases {
|
|
super::eval::execution_tests::observed(|| {
|
|
let [batch, heads, rows, total] =
|
|
["batch", "heads", "rows", "total"].map(|k| case[k].as_i64().unwrap() as i32);
|
|
let step = if case["copied"].as_bool().unwrap() {
|
|
2
|
|
} else {
|
|
1
|
|
};
|
|
let mode = case["mode"].as_str().unwrap();
|
|
let q = ops::transpose(
|
|
&ops::slice(
|
|
&bf16(&[batch, rows, heads, 256 * step], 103),
|
|
&[0, 0, 0, 0],
|
|
&[batch, rows, heads, 256 * step],
|
|
&[1, 1, 1, step],
|
|
gpu,
|
|
)
|
|
.unwrap(),
|
|
&[0, 2, 1, 3],
|
|
gpu,
|
|
)
|
|
.unwrap();
|
|
let kv = |salt| {
|
|
ops::slice(
|
|
&bf16(&[batch, 2, total + 7, 256 * step], salt),
|
|
&[0, 0, 2, 0],
|
|
&[batch, 2, total + 2, 256 * step],
|
|
&[1, 1, 1, step],
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
};
|
|
let (k, v) = (kv(104), kv(105));
|
|
let mask = if matches!(mode, "bool" | "add") {
|
|
let data = (0..rows * total)
|
|
.flat_map(|i| {
|
|
if mode == "bool" {
|
|
vec![u8::from(i % 3 != 1)]
|
|
} else {
|
|
(if i % 3 != 1 { 0u16 } else { 0xc120 })
|
|
.to_le_bytes()
|
|
.to_vec()
|
|
}
|
|
})
|
|
.collect::<Vec<_>>();
|
|
Some(
|
|
Array::new(
|
|
&[1, 1, rows, total],
|
|
if mode == "bool" {
|
|
Dtype::Bool
|
|
} else {
|
|
Dtype::BF16
|
|
},
|
|
super::scalar_buffer(&data).unwrap(),
|
|
)
|
|
.unwrap(),
|
|
)
|
|
} else {
|
|
None
|
|
};
|
|
if prepared {
|
|
ops::evaluate(&streams, &[q.clone(), k.clone(), v.clone()], gpu, false)
|
|
.unwrap();
|
|
}
|
|
let out = if mode == "causal" {
|
|
vector(&q, &k, &v, None, 0.0625, true, gpu)
|
|
} else {
|
|
super::attention::scaled_dot_product(&q, &k, &v, mask.as_ref(), 0.0625, gpu)
|
|
}
|
|
.unwrap();
|
|
drop((q, k, v, mask));
|
|
let (_, calls) = capture_dispatches(|| {
|
|
ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap()
|
|
});
|
|
let names = calls
|
|
.iter()
|
|
.filter(|r| r.0.starts_with("sdpa_vector"))
|
|
.map(|r| r.0.as_str())
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(
|
|
names,
|
|
if total >= 1024 {
|
|
vec![
|
|
"sdpa_vector_2pass_1_bfloat16_t_256_256",
|
|
"sdpa_vector_2pass_2_bfloat16_t_256",
|
|
]
|
|
} else {
|
|
vec!["sdpa_vector_bfloat16_t_256_256"]
|
|
}
|
|
);
|
|
let mut data = vec![0; out.layout().nbytes()];
|
|
out.buffer().read(out.offset(), &mut data).unwrap();
|
|
verify(
|
|
case["kernel"].as_str().unwrap(),
|
|
&[(
|
|
&super::scalar_buffer(&data).unwrap(),
|
|
out.layout().size() as u32,
|
|
)],
|
|
);
|
|
});
|
|
}
|
|
}
|
|
streams.clear_streams().unwrap();
|
|
}
|
|
|
|
pub(super) fn evaluate(
|
|
inputs: &[Array],
|
|
out: &Array,
|
|
scale: f32,
|
|
causal: bool,
|
|
) -> Result<(), String> {
|
|
let qp = &inputs[0];
|
|
let ql = qp.layout();
|
|
let shape = ql.shape().to_vec();
|
|
let st = ql.strides();
|
|
let bidx = if shape[0] == 1 { 1 } else { 0 };
|
|
let q_ok = ql.flags().row_contiguous
|
|
|| ((shape[0] == 1 || shape[1] == 1)
|
|
&& st[3] == 1
|
|
&& st[2] == i64::from(shape[3]) * i64::from(shape[bidx])
|
|
&& st[bidx] == i64::from(shape[3]));
|
|
drop(ql);
|
|
let q = if q_ok { qp.clone() } else { copy(qp)? };
|
|
let mut copies = Vec::new();
|
|
let copy_unless =
|
|
|a: &Array, predicate: bool, copies: &mut Vec<Array>| -> Result<Array, String> {
|
|
if predicate {
|
|
Ok(a.clone())
|
|
} else {
|
|
let c = copy(a)?;
|
|
copies.push(c.clone());
|
|
Ok(c)
|
|
}
|
|
};
|
|
let kv_ok = |a: &Array| {
|
|
let l = a.layout();
|
|
let s = l.shape();
|
|
let t = l.strides();
|
|
t[3] == 1 && (s[0] == 1 || s[1] == 1 || t[0] == t[1] * i64::from(s[1]))
|
|
};
|
|
let k = copy_unless(&inputs[1], kv_ok(&inputs[1]), &mut copies)?;
|
|
let v = copy_unless(&inputs[2], kv_ok(&inputs[2]), &mut copies)?;
|
|
// A non-copied q has the extra descriptor held by the primitive, just as
|
|
// the original's local `array q` does. Only a unique copied query donates.
|
|
let ql = q.layout();
|
|
if q.is_donatable() && ql.flags().row_contiguous && ql.size() == out.layout().size() {
|
|
out.copy_shared_buffer(&q, ql.strides(), ql.flags(), q.data_size(), 0)?;
|
|
} else {
|
|
if !q_ok {
|
|
copies.push(q.clone());
|
|
}
|
|
let nbytes = out.layout().nbytes();
|
|
out.set_data(Buffer::mtplx_bytes(nbytes as u64)?)?;
|
|
}
|
|
drop(ql);
|
|
let mask = if let Some(m) = inputs.get(3) {
|
|
let l = m.layout();
|
|
let s = l.shape();
|
|
let t = l.strides();
|
|
let ok = l.flags().row_contiguous
|
|
|| shape[0] == 1
|
|
|| shape[1] == 1
|
|
|| t[0] == t[1] * i64::from(s[1]);
|
|
drop(l);
|
|
Some(copy_unless(m, ok, &mut copies)?)
|
|
} else {
|
|
None
|
|
};
|
|
let kl = k.layout();
|
|
let vl = v.layout();
|
|
let n = kl.shape()[2];
|
|
let gqa = shape[1] / kl.shape()[1];
|
|
let kh = kl.strides()[if kl.shape()[1] == 1 { 0 } else { 1 }] as u64;
|
|
let ks = kl.strides()[2] as u64;
|
|
let vh = vl.strides()[if vl.shape()[1] == 1 { 0 } else { 1 }] as u64;
|
|
let vs = vl.strides()[2] as u64;
|
|
let arch = encoder::arch_suffix();
|
|
let two = (matches!(arch, Some('s' | 'd')) && n >= 1024) || (gqa > 1 && n >= 4096);
|
|
let bool_mask = mask
|
|
.as_ref()
|
|
.is_some_and(|a| a.layout().dtype() == Dtype::Bool);
|
|
let mut constants = [
|
|
(20, mask.is_some()),
|
|
(21, !q.layout().flags().row_contiguous),
|
|
(22, causal && shape[2] > 1),
|
|
(23, bool_mask),
|
|
(24, mask.is_some() && !bool_mask),
|
|
(25, false),
|
|
]
|
|
.map(|(index, value)| MetalConstant {
|
|
index,
|
|
value: u32::from(value),
|
|
kind: 0,
|
|
})
|
|
.into_iter()
|
|
.collect::<Vec<_>>();
|
|
let mask_layout = mask.as_ref().map(Array::layout);
|
|
let mask_strides = mask_layout
|
|
.as_ref()
|
|
.map(|l| -> Result<[i32; 3], String> {
|
|
let s = l.shape();
|
|
let t = l.strides();
|
|
let raw = [
|
|
if s[3] > 1 { t[3] } else { 0 },
|
|
if s[2] > 1 { t[2] } else { 0 },
|
|
if s[1] > 1 {
|
|
t[1]
|
|
} else if s[0] > 1 {
|
|
t[0]
|
|
} else {
|
|
0
|
|
},
|
|
];
|
|
let [a, b, c] = raw.map(|v| i32::try_from(v).map_err(|_| "SDPA mask stride overflow"));
|
|
Ok([a?, b?, c?])
|
|
})
|
|
.transpose()?;
|
|
let qb = q.buffer();
|
|
let kb = k.buffer();
|
|
let vb = v.buffer();
|
|
let ob = out.buffer();
|
|
let mb = mask.as_ref().map(Array::buffer);
|
|
let base = vec![
|
|
tensor(0, &qb)
|
|
.at_byte_offset(q.offset())
|
|
.input(q.data_size()),
|
|
tensor(1, &kb)
|
|
.at_byte_offset(k.offset())
|
|
.input(k.data_size()),
|
|
tensor(2, &vb)
|
|
.at_byte_offset(v.offset())
|
|
.input(v.data_size()),
|
|
];
|
|
let mask_bindings = |offset: u32| {
|
|
let mut bindings = Vec::new();
|
|
if let (Some(m), Some(ms)) = (&mb, &mask_strides) {
|
|
bindings.push(
|
|
tensor(offset + u32::from(!bool_mask), m)
|
|
.at_byte_offset(mask.as_ref().unwrap().offset())
|
|
.input(mask.as_ref().unwrap().data_size()),
|
|
);
|
|
bindings.extend([
|
|
bytes(offset + 2, &ms[0]),
|
|
bytes(offset + 3, &ms[1]),
|
|
bytes(offset + 4, &ms[2]),
|
|
]);
|
|
}
|
|
bindings
|
|
};
|
|
if two {
|
|
let count = blocks(
|
|
n,
|
|
gqa * shape[2],
|
|
{
|
|
let value = unsafe { libc::getenv(c"MLX_SDPA_BLOCKS".as_ptr()) };
|
|
if value.is_null() {
|
|
0
|
|
} else {
|
|
unsafe { libc::atoi(value) }
|
|
}
|
|
},
|
|
arch,
|
|
)?;
|
|
constants.push(MetalConstant {
|
|
index: 26,
|
|
value: count as u32,
|
|
kind: 1,
|
|
});
|
|
let rows = shape[..3].iter().map(|&v| v as u64).product::<u64>();
|
|
let temporaries = [
|
|
Buffer::mtplx_bytes(rows * count as u64 * 256 * 2)?,
|
|
Buffer::mtplx_bytes(rows * count as u64 * 4)?,
|
|
Buffer::mtplx_bytes(rows * count as u64 * 4)?,
|
|
];
|
|
encoder::add_temporaries(&temporaries);
|
|
let [partial, sums, maxs] = &temporaries;
|
|
let mut bindings = base;
|
|
bindings.extend([
|
|
tensor(3, partial).output((rows * count as u64 * 256) as usize),
|
|
tensor(4, sums).output((rows * count as u64) as usize),
|
|
tensor(5, maxs).output((rows * count as u64) as usize),
|
|
bytes(7, &n),
|
|
bytes(8, &kh),
|
|
bytes(9, &ks),
|
|
bytes(10, &vh),
|
|
bytes(11, &vs),
|
|
bytes(12, &scale),
|
|
]);
|
|
bindings.extend(mask_bindings(13));
|
|
dispatch_specialized(
|
|
"sdpa_vector_2pass_1_bfloat16_t_256_256",
|
|
&bindings,
|
|
&constants,
|
|
[kl.shape()[1] as u32, shape[0] as u32, count as u32],
|
|
[32, gqa as u32, shape[2] as u32],
|
|
)?;
|
|
dispatch_specialized(
|
|
"sdpa_vector_2pass_2_bfloat16_t_256",
|
|
&[
|
|
tensor(0, partial).input((rows * count as u64 * 256) as usize),
|
|
tensor(1, sums).input((rows * count as u64) as usize),
|
|
tensor(2, maxs).input((rows * count as u64) as usize),
|
|
tensor(3, &ob)
|
|
.at_byte_offset(out.offset())
|
|
.output(out.data_size()),
|
|
bytes(4, &count),
|
|
],
|
|
&[],
|
|
[(shape[0] * shape[1]) as u32, shape[2] as u32, 1],
|
|
[1024, 1, 1],
|
|
)?;
|
|
} else {
|
|
let mut bindings = base;
|
|
bindings.extend([
|
|
tensor(3, &ob)
|
|
.at_byte_offset(out.offset())
|
|
.output(out.data_size()),
|
|
bytes(4, &gqa),
|
|
bytes(5, &n),
|
|
bytes(6, &kh),
|
|
bytes(7, &ks),
|
|
bytes(8, &vh),
|
|
bytes(9, &vs),
|
|
bytes(10, &scale),
|
|
]);
|
|
bindings.extend(mask_bindings(11));
|
|
dispatch_specialized(
|
|
"sdpa_vector_bfloat16_t_256_256",
|
|
&bindings,
|
|
&constants,
|
|
[(shape[0] * shape[1]) as u32, shape[2] as u32, 1],
|
|
[1024, 1, 1],
|
|
)?;
|
|
}
|
|
for copy in &copies {
|
|
encoder::add_temporaries(std::slice::from_ref(&*copy.buffer()));
|
|
}
|
|
Ok(())
|
|
}
|