7656 lines
243 KiB
Rust
7656 lines
243 KiB
Rust
//! MTPLX model operations using its pinned custom and runtime 0.32.2 shaders.
|
|
//! Rust owns selection and dispatch. No runtime C/C++ library is linked.
|
|
use super::*;
|
|
use std::cell::RefCell;
|
|
#[cfg(test)]
|
|
use std::cell::{Cell, Ref};
|
|
use std::ffi::CString;
|
|
use std::rc::Rc;
|
|
|
|
pub(super) mod allocator;
|
|
mod ar;
|
|
mod array;
|
|
mod attention;
|
|
mod binary;
|
|
mod binding;
|
|
mod cache;
|
|
mod compile_graph;
|
|
mod compiled;
|
|
mod context_copy;
|
|
mod decoder;
|
|
mod dense;
|
|
mod encoder;
|
|
mod eval;
|
|
mod execution;
|
|
mod fused;
|
|
mod gdn;
|
|
mod graph_sync;
|
|
mod hyper;
|
|
mod indexing;
|
|
#[cfg(test)]
|
|
mod issue103_analysis;
|
|
#[cfg(test)]
|
|
mod issue104;
|
|
mod linear;
|
|
mod load;
|
|
mod mlp;
|
|
mod model;
|
|
mod moe;
|
|
mod mtp;
|
|
mod ngram;
|
|
mod ngram_sidecar;
|
|
mod ngram_stage;
|
|
mod normalization;
|
|
mod ops;
|
|
mod owner;
|
|
mod ple;
|
|
mod policy;
|
|
mod prefill;
|
|
mod qsa_attention;
|
|
mod qsa_compile;
|
|
mod qsa_core;
|
|
#[cfg(test)]
|
|
mod qsa_graph;
|
|
mod qsa_indexer;
|
|
mod qsa_prefill;
|
|
mod qsa_prepare;
|
|
mod qsa_select;
|
|
mod random;
|
|
mod reduce;
|
|
mod request;
|
|
mod residency;
|
|
mod rope;
|
|
mod sampling;
|
|
mod scheduler;
|
|
mod sdpa;
|
|
mod session_cache;
|
|
mod session_checkpoint;
|
|
mod session_codec;
|
|
mod sort;
|
|
mod sparse_attention;
|
|
mod speculative;
|
|
mod stream;
|
|
#[cfg(test)]
|
|
mod submission;
|
|
mod unary;
|
|
mod views;
|
|
mod vision;
|
|
mod vision_kernels;
|
|
mod weights;
|
|
#[cfg(test)]
|
|
pub(super) use encoder::Scope;
|
|
pub(super) use encoder::active_allocator;
|
|
pub(in crate::engine) use owner::QwenExecutor;
|
|
#[cfg(test)]
|
|
pub(super) use qsa_graph::{QsaCacheInput, QsaGraphBank};
|
|
pub(in crate::engine) use request::LiveSession as QwenResidentState;
|
|
#[cfg(test)]
|
|
pub(super) use stream::{Device as StreamDevice, Streams};
|
|
|
|
struct Binding<'a> {
|
|
raw: MetalBinding,
|
|
allocation: Option<BufferAllocationHold>,
|
|
access: Option<encoder::Access>,
|
|
_borrow: std::marker::PhantomData<&'a ()>,
|
|
}
|
|
|
|
fn tensor(index: u32, buffer: &Buffer) -> Binding<'_> {
|
|
Binding {
|
|
raw: MetalBinding {
|
|
tensor: buffer.raw(),
|
|
bytes: std::ptr::null(),
|
|
size: 0,
|
|
index,
|
|
offset: 0,
|
|
},
|
|
allocation: Some(buffer.hold_allocation()),
|
|
access: None,
|
|
_borrow: std::marker::PhantomData,
|
|
}
|
|
}
|
|
|
|
fn bytes<T: ?Sized>(index: u32, value: &T) -> Binding<'_> {
|
|
Binding {
|
|
raw: MetalBinding {
|
|
tensor: std::ptr::null(),
|
|
bytes: std::ptr::from_ref(value).cast(),
|
|
size: u32::try_from(std::mem::size_of_val(value)).expect("Metal inline bytes size"),
|
|
index,
|
|
offset: 0,
|
|
},
|
|
allocation: None,
|
|
access: None,
|
|
_borrow: std::marker::PhantomData,
|
|
}
|
|
}
|
|
|
|
fn scalar_buffer(value: &[u8]) -> Result<Buffer, String> {
|
|
let buffer = Buffer::mtplx_bytes(value.len() as u64)?;
|
|
buffer.write(0, value)?;
|
|
Ok(buffer)
|
|
}
|
|
|
|
fn compile_custom(
|
|
name: &str,
|
|
source: impl FnOnce() -> Result<String, String>,
|
|
) -> Result<(), String> {
|
|
let name = CString::new(name).map_err(|e| e.to_string())?;
|
|
if unsafe { ds4_gpu_mtplx_compile(name.as_ptr(), std::ptr::null()) } == 0 {
|
|
let source = CString::new(source()?).map_err(|e| e.to_string())?;
|
|
check(
|
|
unsafe { ds4_gpu_mtplx_compile(name.as_ptr(), source.as_ptr()) },
|
|
"compiling original MTPLX custom kernel",
|
|
)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn qsa_output_args(section: &serde_json::Value) -> Result<Vec<String>, String> {
|
|
section["outputs"]
|
|
.as_array()
|
|
.ok_or("invalid QSA output names")?
|
|
.iter()
|
|
.map(|name| {
|
|
let name = name.as_str().ok_or("invalid QSA output name")?;
|
|
let ty = if name.ends_with("ids") {
|
|
"int"
|
|
} else if name == "adjusted_scores" {
|
|
"float"
|
|
} else {
|
|
"bool"
|
|
};
|
|
Ok(format!("device {ty}* {name}"))
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn qsa_shader_source(header: &str, name: &str, args: &[String], body: &str) -> String {
|
|
let args = args
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, a)| format!("{a} [[buffer({i})]]"))
|
|
.collect::<Vec<_>>()
|
|
.join(",\n");
|
|
format!(
|
|
"{header}kernel void {name}(\n{args},\nuint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],\nuint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]) {{{body}}}\n"
|
|
)
|
|
}
|
|
|
|
fn dispatch(
|
|
name: &str,
|
|
bindings: &[Binding<'_>],
|
|
groups: [u32; 3],
|
|
threads: [u32; 3],
|
|
) -> Result<(), String> {
|
|
dispatch_specialized(name, bindings, &[], groups, threads)
|
|
}
|
|
|
|
fn dispatch_specialized(
|
|
name: &str,
|
|
bindings: &[Binding<'_>],
|
|
constants: &[MetalConstant],
|
|
groups: [u32; 3],
|
|
threads: [u32; 3],
|
|
) -> Result<(), String> {
|
|
dispatch_geometry(name, bindings, constants, groups, threads, false)
|
|
}
|
|
|
|
fn dispatch_geometry(
|
|
name: &str,
|
|
bindings: &[Binding<'_>],
|
|
constants: &[MetalConstant],
|
|
grid: [u32; 3],
|
|
threads: [u32; 3],
|
|
individual_threads: bool,
|
|
) -> Result<(), String> {
|
|
if let Some(result) = encoder::with_active(|encoder| {
|
|
encoder.dispatch(name, bindings, constants, grid, threads, individual_threads)
|
|
}) {
|
|
return result;
|
|
}
|
|
dispatch_geometry_on(
|
|
name,
|
|
bindings,
|
|
constants,
|
|
grid,
|
|
threads,
|
|
individual_threads,
|
|
None,
|
|
)
|
|
}
|
|
|
|
fn dispatch_geometry_on(
|
|
name: &str,
|
|
bindings: &[Binding<'_>],
|
|
constants: &[MetalConstant],
|
|
grid: [u32; 3],
|
|
threads: [u32; 3],
|
|
individual_threads: bool,
|
|
target: Option<(*mut c_void, *mut c_void, bool)>,
|
|
) -> Result<(), String> {
|
|
let name = CString::new(name).map_err(|e| e.to_string())?;
|
|
let raw = bindings.iter().map(|b| b.raw).collect::<Vec<_>>();
|
|
let mut seen = std::collections::HashSet::new();
|
|
let allocations = bindings
|
|
.iter()
|
|
.filter_map(|b| b.allocation.as_ref())
|
|
.filter(|a| seen.insert(a.identity()))
|
|
.cloned()
|
|
.collect::<Vec<_>>();
|
|
unsafe extern "C" fn release(context: *mut c_void) {
|
|
// The bridge takes this Box exactly once at callback registration,
|
|
// including when a submitted standalone command subsequently fails.
|
|
drop(unsafe { Box::from_raw(context.cast::<Vec<BufferAllocationHold>>()) });
|
|
}
|
|
let mut owner = Box::into_raw(Box::new(allocations)).cast::<c_void>();
|
|
let (command, encoder, barrier) =
|
|
target.unwrap_or((std::ptr::null_mut(), std::ptr::null_mut(), false));
|
|
let result = unsafe {
|
|
ds4_gpu_mtplx_encode(
|
|
name.as_ptr(),
|
|
raw.as_ptr(),
|
|
raw.len() as u32,
|
|
constants.as_ptr(),
|
|
constants.len() as u32,
|
|
grid.as_ptr(),
|
|
threads.as_ptr(),
|
|
u32::from(individual_threads),
|
|
&mut owner,
|
|
Some(release),
|
|
command,
|
|
encoder,
|
|
u32::from(barrier),
|
|
)
|
|
};
|
|
if !owner.is_null() {
|
|
// Validation/pipeline/encoder failure: no ownership was transferred.
|
|
unsafe { release(owner) };
|
|
}
|
|
let result = check(result, &format!("dispatching MTPLX {name:?}"));
|
|
#[cfg(test)]
|
|
if result.is_ok() {
|
|
super::qwen_mtplx_tests::record_dispatch(
|
|
name.to_str().unwrap(),
|
|
grid,
|
|
threads,
|
|
individual_threads,
|
|
);
|
|
}
|
|
result
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
#[cfg(test)]
|
|
pub(super) enum SortDtype {
|
|
U32,
|
|
Bf16,
|
|
F32,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct RouterOutput {
|
|
pub(super) probabilities: Buffer,
|
|
/// Last ten columns of argsort, not a contiguous [rows,10] allocation.
|
|
pub(super) indices: Buffer,
|
|
pub(super) index_row_stride: u32,
|
|
pub(super) scores: Buffer,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// SparseMoeBlock's precise softmax -> argpartition -> take_along_axis ->
|
|
/// optional normalization. Consumes donatable logits; retains GPU scratch.
|
|
#[cfg(test)]
|
|
pub(super) fn router(
|
|
logits: Buffer,
|
|
[rows, experts]: [u32; 2],
|
|
normalize: bool,
|
|
) -> Result<RouterOutput, String> {
|
|
let count = u64::from(rows) * u64::from(experts);
|
|
if rows == 0
|
|
|| experts < 10
|
|
|| count > i32::MAX as u64
|
|
|| unsafe { ds4_gpu_tensor_bytes(logits.raw()) } < count * 2
|
|
{
|
|
return Err("invalid MTPLX top-10 router shape or buffer".into());
|
|
}
|
|
let group = if experts <= 4096 {
|
|
experts.div_ceil(4).div_ceil(32) * 32
|
|
} else {
|
|
1024 // Pinned looped-softmax pipeline maximum on M5 Max.
|
|
};
|
|
let threads = rows
|
|
.checked_mul(group)
|
|
.ok_or("MTPLX softmax grid overflow")?;
|
|
dispatch_geometry(
|
|
if experts <= 4096 {
|
|
"block_softmax_precise_bfloat16"
|
|
} else {
|
|
"looped_softmax_precise_bfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &logits).input(count as usize),
|
|
tensor(1, &logits).output(count as usize),
|
|
bytes(2, &experts),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[group, 1, 1],
|
|
true,
|
|
)?;
|
|
let order = Buffer::mtplx_bytes(count * 4)?;
|
|
let mut temporary = argsort(&logits, &order, [rows, experts], SortDtype::Bf16)?;
|
|
let offset = u64::from(experts - 10) * 4;
|
|
let indices = order.view(offset, count * 4 - offset)?;
|
|
let scores = Buffer::mtplx_bytes(u64::from(rows) * 10 * 2)?;
|
|
let grid = [1, 10, rows];
|
|
dispatch_geometry(
|
|
if rows == 1 || experts == 10 {
|
|
"kernel_qwen_mtplx_gather_axis_bf16_int_11"
|
|
} else {
|
|
"kernel_qwen_mtplx_gather_axis_bf16_int_10"
|
|
},
|
|
&[
|
|
tensor(0, &logits).input(count as usize),
|
|
tensor(1, &indices).input(((u64::from(rows) - 1) * u64::from(experts) + 10) as usize),
|
|
tensor(2, &scores).output((u64::from(rows) * 10) as usize),
|
|
bytes(3, &rows),
|
|
bytes(4, &u64::from(experts)),
|
|
bytes(5, &u64::from(experts)),
|
|
bytes(6, &1_u64),
|
|
bytes(7, &1_u32),
|
|
bytes(8, &experts),
|
|
bytes(9, &1_u64),
|
|
bytes(10, &1_u64),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
if normalize {
|
|
let sums = Buffer::mtplx_bytes(u64::from(rows) * 2)?;
|
|
dispatch_geometry(
|
|
"row_reduce_small_1_reduce_sumbfloat16",
|
|
&[
|
|
tensor(0, &scores).input((u64::from(rows) * 10) as usize),
|
|
tensor(1, &sums).output(rows as usize),
|
|
bytes(2, &10_u64),
|
|
bytes(3, &1_u64),
|
|
bytes(4, &rows),
|
|
bytes(5, &10_u64),
|
|
bytes(6, &1_u32),
|
|
bytes(7, &0_u32),
|
|
bytes(8, &0_u64),
|
|
bytes(9, &0_u32),
|
|
],
|
|
&[],
|
|
[rows, 1, 1],
|
|
[rows.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
if rows == 1 {
|
|
dispatch_geometry(
|
|
"vs_Dividebfloat16",
|
|
&[
|
|
tensor(0, &scores).input((u64::from(rows) * 10) as usize),
|
|
tensor(1, &sums).input(rows as usize),
|
|
tensor(2, &scores).output((u64::from(rows) * 10) as usize),
|
|
bytes(3, &10_u32),
|
|
],
|
|
&[],
|
|
[10, 1, 1],
|
|
[10, 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let grid = [10, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_Dividebfloat16",
|
|
&[
|
|
tensor(0, &scores).input((u64::from(rows) * 10) as usize),
|
|
tensor(1, &sums).input(rows as usize),
|
|
tensor(2, &scores).output((u64::from(rows) * 10) as usize),
|
|
bytes(3, &[10_u64, 1]),
|
|
bytes(4, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
temporary.push(sums);
|
|
}
|
|
temporary.push(order);
|
|
Ok(RouterOutput {
|
|
probabilities: logits,
|
|
indices,
|
|
index_row_stride: experts,
|
|
scores,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
/// Contiguous last-axis argsort used by routing and _gather_sort. The reference
|
|
/// also implements argpartition by this complete sort, not a local top-k kernel.
|
|
#[cfg(test)]
|
|
pub(super) fn argsort(
|
|
input: &Buffer,
|
|
output: &Buffer,
|
|
[rows, width]: [u32; 2],
|
|
dtype: SortDtype,
|
|
) -> Result<Vec<Buffer>, String> {
|
|
use array::{Array, Dtype};
|
|
let dtype = match dtype {
|
|
SortDtype::U32 => Dtype::U32,
|
|
SortDtype::Bf16 => Dtype::BF16,
|
|
SortDtype::F32 => Dtype::F32,
|
|
};
|
|
let count = u64::from(rows) * u64::from(width);
|
|
let item_bytes = dtype.itemsize() as u64;
|
|
if rows == 0
|
|
|| width == 0
|
|
|| count > i32::MAX as u64
|
|
|| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < count * item_bytes
|
|
|| unsafe { ds4_gpu_tensor_bytes(output.raw()) } < count * 4
|
|
{
|
|
return Err("invalid MTPLX contiguous argsort shape or buffer".into());
|
|
}
|
|
let shape = if rows == 1 {
|
|
vec![width as i32]
|
|
} else {
|
|
vec![rows as i32, width as i32]
|
|
};
|
|
let input = Array::new(&shape, dtype, input.view(0, count * item_bytes)?)?;
|
|
let output = Array::new(&shape, Dtype::U32, output.view(0, count * 4)?)?;
|
|
sort::sort_into(&input, &output, shape.len() - 1, true)
|
|
}
|
|
|
|
fn block_dims(grid: [u32; 3]) -> [u32; 3] {
|
|
block_dims_pow2(grid, 10)
|
|
}
|
|
|
|
fn block_dims_pow2(grid: [u32; 3], pow2: u32) -> [u32; 3] {
|
|
let mut group = [1, 1, 1];
|
|
let mut used = 0;
|
|
loop {
|
|
let before = used;
|
|
for axis in 0..3 {
|
|
if grid[axis] >= group[axis] * 2 {
|
|
group[axis] *= 2;
|
|
used += 1;
|
|
}
|
|
if used == pow2 {
|
|
break;
|
|
}
|
|
}
|
|
if used == pow2 || used == before {
|
|
break;
|
|
}
|
|
}
|
|
group
|
|
}
|
|
|
|
/// Actual compiled nn.silu(gate) followed by the reference's separate * up.
|
|
/// Input is the fused [rows,2*intermediate] BF16 pack, with strided half views.
|
|
/// Keep returned SiLU storage alive until the enclosing GPU batch completes.
|
|
#[cfg(test)]
|
|
pub(super) fn gate_up_bf16(
|
|
input: &Buffer,
|
|
output: &Buffer,
|
|
[rows, intermediate]: [u32; 2],
|
|
) -> Result<Buffer, String> {
|
|
let count = u64::from(rows) * u64::from(intermediate);
|
|
if rows == 0
|
|
|| intermediate == 0
|
|
|| !intermediate.is_multiple_of(2)
|
|
|| count > i32::MAX as u64
|
|
|| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < count * 4
|
|
|| unsafe { ds4_gpu_tensor_bytes(output.raw()) } < count * 2
|
|
{
|
|
return Err("invalid MTPLX fused gate/up buffer or shape".into());
|
|
}
|
|
let activated = Buffer::mtplx_bytes(count * 2)?;
|
|
let up = input.view(
|
|
u64::from(intermediate) * 2,
|
|
count * 4 - u64::from(intermediate) * 2,
|
|
)?;
|
|
let count = count as u32;
|
|
if rows == 1 {
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"kernel_qwen_mtplx_silu_bf16_contiguous"
|
|
} else {
|
|
"kernel_qwen_mtplx_silu_bf16_contiguous_n"
|
|
},
|
|
&[
|
|
tensor(0, input).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
|
|
tensor(1, &activated).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"vv_Multiplybfloat16"
|
|
} else {
|
|
"vvn_Multiplybfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &activated).input(count as usize),
|
|
tensor(1, &up).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
|
|
tensor(2, output).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let shape = [rows, intermediate];
|
|
let in_strides = [u64::from(intermediate) * 2, 1];
|
|
let out_strides = [u64::from(intermediate), 1];
|
|
let grid = [intermediate, rows, 1];
|
|
let group = block_dims(grid);
|
|
dispatch_geometry(
|
|
"kernel_qwen_mtplx_silu_bf16_strided_2",
|
|
&[
|
|
tensor(0, input).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
|
|
bytes(1, &in_strides),
|
|
tensor(2, &activated).output(count as usize),
|
|
bytes(3, &shape),
|
|
],
|
|
&[],
|
|
grid,
|
|
group,
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
"g2_Multiplybfloat16",
|
|
&[
|
|
tensor(0, &activated).input(count as usize),
|
|
tensor(1, &up).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
|
|
tensor(2, output).output(count as usize),
|
|
bytes(3, &out_strides),
|
|
bytes(4, &in_strides),
|
|
],
|
|
&[],
|
|
grid,
|
|
group,
|
|
true,
|
|
)?;
|
|
}
|
|
Ok(activated)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn rms_norm_bf16(
|
|
input: &Buffer,
|
|
weight: Option<&Buffer>,
|
|
output: &Buffer,
|
|
[rows, width]: [u32; 2],
|
|
eps: f32,
|
|
) -> Result<(), String> {
|
|
let count = u64::from(rows) * u64::from(width);
|
|
if rows == 0
|
|
|| width == 0
|
|
|| count > i32::MAX as u64
|
|
|| !eps.is_finite()
|
|
|| eps < 0.0
|
|
|| [input, output]
|
|
.iter()
|
|
.any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < count * 2)
|
|
|| weight.is_some_and(|w| unsafe { ds4_gpu_tensor_bytes(w.raw()) } < u64::from(width) * 2)
|
|
{
|
|
return Err("invalid MTPLX BF16 RMS norm".into());
|
|
}
|
|
let default_weight = weight
|
|
.is_none()
|
|
.then(|| scalar_buffer(&0x3f80_u16.to_le_bytes()))
|
|
.transpose()?;
|
|
let bound_weight = weight.or(default_weight.as_ref()).unwrap();
|
|
normalization::dispatch(
|
|
[
|
|
(input, 0, count as usize),
|
|
(
|
|
bound_weight,
|
|
0,
|
|
if weight.is_some() { width as usize } else { 1 },
|
|
),
|
|
(output, 0, count as usize),
|
|
],
|
|
array::Dtype::BF16,
|
|
width,
|
|
u32::from(weight.is_some()),
|
|
eps,
|
|
)
|
|
}
|
|
|
|
/// The shared HC/PLE GroupedRMSNorm: unweighted per-group RMS followed by
|
|
/// separate full-feature weighting. Folding the weight into RMS changes rounding.
|
|
#[cfg(test)]
|
|
pub(super) fn grouped_rms_norm(
|
|
input: &Buffer,
|
|
weight: &Buffer,
|
|
[rows, width, group_width]: [u32; 3],
|
|
eps: f32,
|
|
) -> Result<Buffer, String> {
|
|
if group_width == 0
|
|
|| !width.is_multiple_of(group_width)
|
|
|| unsafe { ds4_gpu_tensor_bytes(weight.raw()) } < u64::from(width) * 2
|
|
{
|
|
return Err("invalid MTPLX grouped RMS shape or weight".into());
|
|
}
|
|
let count = u64::from(rows) * u64::from(width);
|
|
if count == 0 || count > i32::MAX as u64 {
|
|
return Err("invalid MTPLX grouped RMS size".into());
|
|
}
|
|
let output = Buffer::mtplx_bytes(count * 2)?;
|
|
rms_norm_bf16(
|
|
input,
|
|
None,
|
|
&output,
|
|
[count as u32 / group_width, group_width],
|
|
eps,
|
|
)?;
|
|
if rows == 1 {
|
|
let work = if width < 65536 { 1 } else { 4 };
|
|
let threads = width.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"vv_Multiplybfloat16"
|
|
} else {
|
|
"vvn_Multiplybfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &output).input(count as usize),
|
|
tensor(1, weight).input(width as usize),
|
|
tensor(2, &output).output(count as usize),
|
|
bytes(3, &width),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let grid = [width, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_Multiplybfloat16",
|
|
&[
|
|
tensor(0, &output).input(count as usize),
|
|
tensor(1, weight).input(width as usize),
|
|
tensor(2, &output).output(count as usize),
|
|
bytes(3, &[u64::from(width), 1]),
|
|
bytes(4, &[0_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
Ok(output)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn cast_bf16_f32(input: &Buffer, output: &Buffer, count: u32, to_f32: bool) -> Result<(), String> {
|
|
let (in_bytes, out_bytes) = if to_f32 { (2, 4) } else { (4, 2) };
|
|
if count == 0
|
|
|| count > i32::MAX as u32
|
|
|| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < u64::from(count) * in_bytes
|
|
|| unsafe { ds4_gpu_tensor_bytes(output.raw()) } < u64::from(count) * out_bytes
|
|
{
|
|
return Err("invalid MTPLX BF16/FP32 cast".into());
|
|
}
|
|
let work = if count < 65536 {
|
|
1
|
|
} else if to_f32 {
|
|
2
|
|
} else {
|
|
4
|
|
};
|
|
let threads = count.div_ceil(work);
|
|
let dtype = if to_f32 {
|
|
"bfloat16float32"
|
|
} else {
|
|
"float32bfloat16"
|
|
};
|
|
dispatch_geometry(
|
|
&format!("v{}_copy{dtype}", if work == 1 { "" } else { "n" }),
|
|
&[
|
|
tensor(0, input).input(count as usize),
|
|
tensor(1, output).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct GatedNormOutput {
|
|
pub(super) hidden: Buffer,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// SigmoidRMSNormGated, including the projected Z view's token stride.
|
|
/// RMS output rounds to BF16 before its FP32 sigmoid/multiply epilogue.
|
|
#[cfg(test)]
|
|
pub(super) fn sigmoid_rms_norm(
|
|
input: &Buffer,
|
|
weight: &Buffer,
|
|
gate: Option<&Buffer>,
|
|
[tokens, heads, width, gate_stride]: [u32; 4],
|
|
eps: f32,
|
|
) -> Result<GatedNormOutput, String> {
|
|
let row_width = heads
|
|
.checked_mul(width)
|
|
.ok_or("MTPLX gated RMS width overflow")?;
|
|
let count = tokens
|
|
.checked_mul(row_width)
|
|
.filter(|n| *n > 0 && *n <= i32::MAX as u32)
|
|
.ok_or("invalid MTPLX gated RMS count")?;
|
|
let rows = tokens
|
|
.checked_mul(heads)
|
|
.ok_or("MTPLX gated RMS row overflow")?;
|
|
let hidden = Buffer::mtplx_bytes(u64::from(count) * 2)?;
|
|
rms_norm_bf16(input, Some(weight), &hidden, [rows, width], eps)?;
|
|
let Some(gate) = gate else {
|
|
return Ok(GatedNormOutput {
|
|
hidden,
|
|
_temporaries: Vec::new(),
|
|
});
|
|
};
|
|
if gate_stride < row_width
|
|
|| unsafe { ds4_gpu_tensor_bytes(gate.raw()) }
|
|
< (u64::from(tokens - 1) * u64::from(gate_stride) + u64::from(row_width)) * 2
|
|
{
|
|
return Err("invalid MTPLX gated RMS Z view".into());
|
|
}
|
|
let g = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
if tokens == 1 || gate_stride == row_width {
|
|
cast_bf16_f32(gate, &g, count, true)?;
|
|
} else {
|
|
let grid = [row_width, tokens, 1];
|
|
dispatch_geometry(
|
|
"g2_copybfloat16float32",
|
|
&[
|
|
tensor(0, gate).input(
|
|
(u64::from(tokens - 1) * u64::from(gate_stride) + u64::from(row_width))
|
|
as usize,
|
|
),
|
|
tensor(1, &g).output(count as usize),
|
|
bytes(3, &[u64::from(gate_stride), 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let work = if count < 65536 { 1 } else { 2 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"v_Sigmoidfloat32float32"
|
|
} else {
|
|
"vn_Sigmoidfloat32float32"
|
|
},
|
|
&[
|
|
tensor(0, &g).input(count as usize),
|
|
tensor(1, &g).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
let x = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
cast_bf16_f32(&hidden, &x, count, true)?;
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"vv_Multiplyfloat32"
|
|
} else {
|
|
"vvn_Multiplyfloat32"
|
|
},
|
|
&[
|
|
tensor(0, &g).input(count as usize),
|
|
tensor(1, &x).input(count as usize),
|
|
tensor(2, &x).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
cast_bf16_f32(&x, &hidden, count, false)?;
|
|
Ok(GatedNormOutput {
|
|
hidden,
|
|
_temporaries: vec![g, x],
|
|
})
|
|
}
|
|
|
|
/// Row gather used by MTPLX's _gather_sort and _scatter_unsort. Indices must
|
|
/// originate from the checked routing/sort path, not unchecked external data.
|
|
#[cfg(test)]
|
|
pub(super) fn gather_front(
|
|
input: &Buffer,
|
|
indices: &Buffer,
|
|
output: &Buffer,
|
|
[rows, width, routes]: [u32; 3],
|
|
dtype: SortDtype,
|
|
) -> Result<(), String> {
|
|
let (tag, item_bytes) = match dtype {
|
|
SortDtype::U32 => ("u32", 4),
|
|
SortDtype::Bf16 => ("bf16", 2),
|
|
SortDtype::F32 => ("f32", 4),
|
|
};
|
|
let src_count = u64::from(rows) * u64::from(width);
|
|
let out_count = u64::from(routes) * u64::from(width);
|
|
if [rows, width, routes]
|
|
.iter()
|
|
.any(|v| *v == 0 || *v > i32::MAX as u32)
|
|
|| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < src_count * item_bytes
|
|
|| unsafe { ds4_gpu_tensor_bytes(indices.raw()) } < u64::from(routes) * 4
|
|
|| unsafe { ds4_gpu_tensor_bytes(output.raw()) } < out_count * item_bytes
|
|
{
|
|
return Err("invalid MTPLX gather-front shape or buffer".into());
|
|
}
|
|
let work = if width > 8 && item_bytes < 4 { 2 } else { 1 };
|
|
let grid = [width.div_ceil(work), routes, 1];
|
|
let group = block_dims(grid);
|
|
let loc = if src_count > i32::MAX as u64 || out_count > i32::MAX as u64 {
|
|
"int64_t"
|
|
} else {
|
|
"int"
|
|
};
|
|
dispatch_geometry(
|
|
&format!("kernel_qwen_mtplx_gather_front_{tag}_{loc}_{work}"),
|
|
&[
|
|
tensor(0, input).input(src_count as usize),
|
|
tensor(1, indices).input(routes as usize),
|
|
tensor(2, output).output(out_count as usize),
|
|
bytes(3, &u64::from(width)),
|
|
bytes(4, &rows),
|
|
],
|
|
&[],
|
|
grid,
|
|
group,
|
|
true,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct GatherSorted {
|
|
pub(super) x: Buffer,
|
|
pub(super) indices: Buffer,
|
|
pub(super) inverse: Buffer,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// The connected reference _gather_sort path. Retain the result through GPU
|
|
/// completion; the caller applies the reference's >=64-route sort condition.
|
|
#[cfg(test)]
|
|
pub(super) fn gather_sort(
|
|
input: &Buffer,
|
|
indices: &Buffer,
|
|
[tokens, top_k, width]: [u32; 3],
|
|
) -> Result<GatherSorted, String> {
|
|
let routes = tokens
|
|
.checked_mul(top_k)
|
|
.filter(|v| *v > 0 && *v <= i32::MAX as u32)
|
|
.ok_or("invalid MTPLX route count")?;
|
|
if width == 0
|
|
|| width > i32::MAX as u32
|
|
|| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < u64::from(tokens) * u64::from(width) * 2
|
|
{
|
|
return Err("invalid MTPLX gather-sort input".into());
|
|
}
|
|
let order = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
|
|
let inverse = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
|
|
let token_rows = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
|
|
let sorted_indices = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
|
|
let x = Buffer::mtplx_bytes(u64::from(routes) * u64::from(width) * 2)?;
|
|
let mut temporary = argsort(indices, &order, [1, routes], SortDtype::U32)?;
|
|
temporary.extend(argsort(&order, &inverse, [1, routes], SortDtype::U32)?);
|
|
let divisor = scalar_buffer(&top_k.to_le_bytes())?;
|
|
let work = if routes < 65536 { 1 } else { 2 };
|
|
let threads = routes.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"vs_Divideuint32"
|
|
} else {
|
|
"vsn_Divideuint32"
|
|
},
|
|
&[
|
|
tensor(0, &order).input(routes as usize),
|
|
tensor(1, &divisor).input(1),
|
|
tensor(2, &token_rows).output(routes as usize),
|
|
bytes(3, &routes),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
gather_front(
|
|
input,
|
|
&token_rows,
|
|
&x,
|
|
[tokens, width, routes],
|
|
SortDtype::Bf16,
|
|
)?;
|
|
gather_front(
|
|
indices,
|
|
&order,
|
|
&sorted_indices,
|
|
[routes, 1, routes],
|
|
SortDtype::U32,
|
|
)?;
|
|
temporary.extend([order, token_rows]);
|
|
Ok(GatherSorted {
|
|
x,
|
|
indices: sorted_indices,
|
|
inverse,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
fn qmv_limit(k: u32, n: u32) -> u32 {
|
|
if k <= 2048 && n <= 2048 {
|
|
33
|
|
} else if k <= 4096 && n <= 4096 {
|
|
25
|
|
} else {
|
|
13
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, PartialEq)]
|
|
pub(super) struct AffinePlan {
|
|
pub(super) name: String,
|
|
pub(super) groups: [u32; 3],
|
|
pub(super) threads: [u32; 3],
|
|
pub(super) splits: u32,
|
|
pub(super) parameters: Vec<u32>,
|
|
}
|
|
|
|
/// Contiguous, non-batched affine x @ W.T on the requested M5 Max.
|
|
/// Selection follows quantized.cpp at 1f8e74e, not a local speed heuristic.
|
|
pub(super) fn affine_plan(
|
|
shape: [u32; 3],
|
|
bits: u32,
|
|
group: u32,
|
|
nax: bool,
|
|
) -> Result<AffinePlan, String> {
|
|
affine_plan_batched(shape, bits, group, nax, 1)
|
|
}
|
|
|
|
/// QuantizedMatmul uses M=x.shape(-2) and the original batch strides when
|
|
/// x is not row-contiguous. Only the B=1 matrix route may use Split-K.
|
|
fn affine_plan_batched(
|
|
[m, n, k]: [u32; 3],
|
|
bits: u32,
|
|
group: u32,
|
|
nax: bool,
|
|
batches: u32,
|
|
) -> Result<AffinePlan, String> {
|
|
if [m, n, k].iter().any(|d| *d == 0 || *d > i32::MAX as u32)
|
|
|| !matches!(bits, 4 | 8)
|
|
|| !matches!(group, 32 | 64)
|
|
|| !k.is_multiple_of(group)
|
|
|| batches == 0
|
|
|| u64::from(m)
|
|
.checked_mul(u64::from(n))
|
|
.and_then(|v| v.checked_mul(u64::from(batches)))
|
|
.is_none_or(|v| v > i32::MAX as u64)
|
|
{
|
|
return Err("invalid MTPLX BF16 affine shape or installed-pack quantization".into());
|
|
}
|
|
let suffix = format!("bfloat16_t_gs_{group}_b_{bits}");
|
|
let batch = u8::from(batches > 1);
|
|
let limit = qmv_limit(k, n);
|
|
let mut plan = AffinePlan {
|
|
name: String::new(),
|
|
groups: [1, 1, 1],
|
|
threads: [32, 2, 1],
|
|
splits: 1,
|
|
parameters: vec![k, n],
|
|
};
|
|
if m < limit {
|
|
if matches!(k, 64 | 128) {
|
|
plan.name = format!("affine_qmv_quad_{suffix}_d_{k}_batch_{batch}");
|
|
plan.groups = [m, n.div_ceil(64), batches];
|
|
plan.threads = [32, 1, 1];
|
|
} else if m >= 2 {
|
|
let vectors = m.div_ceil(m.div_ceil(5));
|
|
plan.name = format!("affine_qmv_wide_{suffix}_nv_{vectors}_kl_8_batch_{batch}");
|
|
plan.groups = [m.div_ceil(vectors), n.div_ceil(8), batches];
|
|
plan.parameters.push(m);
|
|
} else {
|
|
let fast = n.is_multiple_of(8) && k.is_multiple_of((32 / bits) * 2 * 32);
|
|
plan.name = format!(
|
|
"affine_qmv{}_{suffix}_batch_{batch}",
|
|
if fast { "_fast" } else { "" }
|
|
);
|
|
plan.groups = [m, n.div_ceil(8), batches];
|
|
}
|
|
return Ok(plan);
|
|
}
|
|
let tiles = u64::from(n.div_ceil(32)) * u64::from(m.div_ceil(32));
|
|
let alignment = group.max(32);
|
|
let mut splits = if batches == 1 {
|
|
((512 / tiles) as u32).max(1).min(k / alignment)
|
|
} else {
|
|
1
|
|
};
|
|
while splits > 1 && !k.is_multiple_of(splits * alignment) {
|
|
splits -= 1;
|
|
}
|
|
plan.splits = splits;
|
|
plan.parameters.push(m);
|
|
plan.threads = [32, 2, 2];
|
|
if splits > 1 {
|
|
plan.name = format!("affine_qmm_t_splitk_{suffix}_alN_{}", n.is_multiple_of(32));
|
|
plan.groups = [n.div_ceil(32), m.div_ceil(32), splits];
|
|
plan.parameters.extend([k / splits, m * n]);
|
|
} else if nax && k.is_multiple_of(64) {
|
|
let bm = if m <= 32 { 32 } else { 64 };
|
|
plan.name = format!(
|
|
"affine_qmm_t_nax_{suffix}_bm{bm}_bn64_bk64_wm2_wn2_alN_{}_batch_{batch}",
|
|
n.is_multiple_of(64)
|
|
);
|
|
plan.groups = [n.div_ceil(64), m.div_ceil(bm), batches];
|
|
} else {
|
|
plan.name = format!(
|
|
"affine_qmm_t_{suffix}_alN_{}_batch_{batch}",
|
|
n.is_multiple_of(32)
|
|
);
|
|
plan.groups = [n.div_ceil(32), m.div_ceil(32), batches];
|
|
}
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Normalized x[X,M,K], affine weights[E,N,K], index vectors[R], y[R,M,N].
|
|
/// For sorted=true x is already gathered in rhs order and lhs is identity,
|
|
/// matching _gather_sort and the reference's implicit-lhs sorted dispatch.
|
|
/// Indices originate from the checked router/sort path, not unchecked user data.
|
|
#[cfg(test)]
|
|
pub(super) fn gather_bf16(
|
|
buffers: [&Buffer; 7],
|
|
shape: [u32; 6],
|
|
quant: [u32; 2],
|
|
sorted: bool,
|
|
nax: bool,
|
|
) -> Result<(), String> {
|
|
gather_bf16_broadcast(buffers, shape, quant, sorted, nax, 1, [shape[5]; 2])
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn gather_bf16_broadcast(
|
|
buffers: [&Buffer; 7],
|
|
[x_batches, m, n, k, experts, routes]: [u32; 6],
|
|
[bits, group]: [u32; 2],
|
|
sorted: bool,
|
|
nax: bool,
|
|
lhs_repeat: u32,
|
|
[rhs_width, rhs_stride]: [u32; 2],
|
|
) -> Result<(), String> {
|
|
affine_plan([m, n, k], bits, group, nax)?;
|
|
if [x_batches, experts, routes]
|
|
.iter()
|
|
.any(|d| *d == 0 || *d > i32::MAX as u32)
|
|
|| u64::from(m) * u64::from(routes) * u64::from(n) > i32::MAX as u64
|
|
|| (sorted && x_batches != routes)
|
|
|| lhs_repeat == 0
|
|
|| !routes.is_multiple_of(lhs_repeat)
|
|
|| (sorted && lhs_repeat != 1)
|
|
|| rhs_width == 0
|
|
|| rhs_stride < rhs_width
|
|
|| !routes.is_multiple_of(rhs_width)
|
|
|| (rhs_stride != rhs_width && lhs_repeat != 1 && lhs_repeat != rhs_width)
|
|
|| (sorted && rhs_stride != rhs_width)
|
|
{
|
|
return Err("invalid normalized MTPLX gather shape".into());
|
|
}
|
|
let [w, s, b, x, lhs, rhs, y] = buffers;
|
|
let packed_cols = k / (32 / bits);
|
|
let param_cols = k / group;
|
|
let w_bytes = (u64::from(experts) * u64::from(n)).checked_mul(u64::from(packed_cols) * 4);
|
|
let p_bytes = (u64::from(experts) * u64::from(n)).checked_mul(u64::from(param_cols) * 2);
|
|
let expected = [
|
|
w_bytes,
|
|
p_bytes,
|
|
p_bytes,
|
|
(u64::from(x_batches) * u64::from(m)).checked_mul(u64::from(k) * 2),
|
|
Some(u64::from(routes / lhs_repeat) * 4),
|
|
Some(
|
|
(u64::from(routes / rhs_width - 1) * u64::from(rhs_stride) + u64::from(rhs_width)) * 4,
|
|
),
|
|
Some(u64::from(routes) * u64::from(m) * u64::from(n) * 2),
|
|
];
|
|
for (buffer, size) in buffers.iter().zip(expected) {
|
|
if size.is_none_or(|size| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < size) {
|
|
return Err("MTPLX gather buffer is smaller than its shape".into());
|
|
}
|
|
}
|
|
let elements: [usize; 7] =
|
|
std::array::from_fn(|i| (expected[i].unwrap() / [4, 2, 2, 2, 4, 4, 2][i]) as usize);
|
|
let strided_rhs = rhs_stride != rhs_width && rhs_width < routes;
|
|
let batch_dim = if strided_rhs || (lhs_repeat > 1 && lhs_repeat < routes) {
|
|
2_u32
|
|
} else {
|
|
1
|
|
};
|
|
let batch_shape = if batch_dim == 2 {
|
|
let width = if strided_rhs { rhs_width } else { lhs_repeat };
|
|
[routes / width, width]
|
|
} else {
|
|
[routes, 1]
|
|
};
|
|
let lhs_strides = if strided_rhs && lhs_repeat == 1 {
|
|
[u64::from(rhs_width), 1]
|
|
} else if lhs_repeat > 1 && batch_dim == 1 {
|
|
[0, 0]
|
|
} else {
|
|
[1_u64, 0]
|
|
};
|
|
let rhs_strides = if strided_rhs {
|
|
[u64::from(rhs_stride), 1]
|
|
} else if batch_dim == 1 {
|
|
[1_u64, 0]
|
|
} else {
|
|
[u64::from(lhs_repeat), 1]
|
|
};
|
|
let x_shape = [x_batches, m, k];
|
|
let x_strides = [u64::from(m) * u64::from(k), u64::from(k), 1];
|
|
let w_shape = [experts, n, packed_cols];
|
|
let w_strides = [
|
|
u64::from(n) * u64::from(packed_cols),
|
|
u64::from(packed_cols),
|
|
1,
|
|
];
|
|
let p_strides = [
|
|
u64::from(n) * u64::from(param_cols),
|
|
u64::from(param_cols),
|
|
1,
|
|
];
|
|
let batch_shape = batch_shape.map(|d| d as i32);
|
|
let lhs_strides = lhs_strides.map(|d| d as i64);
|
|
let rhs_strides = rhs_strides.map(|d| d as i64);
|
|
|
|
gather_bf16_layout(
|
|
[
|
|
tensor(0, w).input(elements[0]),
|
|
tensor(1, s).input(elements[1]),
|
|
tensor(2, b).input(elements[2]),
|
|
tensor(3, x).input(elements[3]),
|
|
tensor(4, lhs).input(elements[4]),
|
|
tensor(5, rhs).input(elements[5]),
|
|
tensor(6, y).output(elements[6]),
|
|
],
|
|
[m, n, k, experts, routes, routes],
|
|
[bits, group],
|
|
sorted,
|
|
nax,
|
|
[
|
|
&x_shape.map(|d| d as i32),
|
|
&w_shape.map(|d| d as i32),
|
|
&batch_shape[..batch_dim as usize],
|
|
],
|
|
[
|
|
&x_strides.map(|d| d as i64),
|
|
&w_strides.map(|d| d as i64),
|
|
&p_strides.map(|d| d as i64),
|
|
&p_strides.map(|d| d as i64),
|
|
&lhs_strides[..batch_dim as usize],
|
|
&rhs_strides[..batch_dim as usize],
|
|
],
|
|
)
|
|
}
|
|
|
|
/// Shared original GatherQMM dispatch; array batch strides are passed verbatim,
|
|
/// while the two broadcast index layouts have jointly collapsed dimensions.
|
|
fn gather_bf16_layout(
|
|
bindings: [Binding<'_>; 7],
|
|
[m, n, k, experts, routes, sorted_rows]: [u32; 6],
|
|
[bits, group]: [u32; 2],
|
|
sorted: bool,
|
|
nax: bool,
|
|
[x_shape, w_shape, batch_shape]: [&[i32]; 3],
|
|
[
|
|
x_strides,
|
|
w_strides,
|
|
s_strides,
|
|
b_strides,
|
|
lhs_strides,
|
|
rhs_strides,
|
|
]: [&[i64]; 6],
|
|
) -> Result<(), String> {
|
|
let suffix = format!("bfloat16_t_gs_{group}_b_{bits}");
|
|
if m == 1 && routes >= 16 && sorted && routes / experts >= 4 {
|
|
let (bm, bn, bk, wm, wn) = if nax {
|
|
(
|
|
if sorted_rows / experts < 64 { 32 } else { 64 },
|
|
64,
|
|
64,
|
|
2,
|
|
2,
|
|
)
|
|
} else {
|
|
(16, 32, 32, 1, 2)
|
|
};
|
|
let name = format!(
|
|
"affine_gather_qmm_rhs{}_nt_{suffix}_bm_{bm}_bn_{bn}_bk_{bk}_wm_{wm}_wn_{wn}",
|
|
if nax { "_nax" } else { "" }
|
|
);
|
|
let constants = [
|
|
(200, sorted_rows.is_multiple_of(bm)),
|
|
(201, n.is_multiple_of(bn)),
|
|
(202, k.is_multiple_of(bk)),
|
|
]
|
|
.map(|(index, value)| MetalConstant {
|
|
index,
|
|
value: u32::from(value),
|
|
kind: 0,
|
|
});
|
|
|
|
let [w, s, b, x, _lhs, rhs, y] = bindings;
|
|
let mut sorted_bindings = [x, w, s, b, rhs, y];
|
|
for (index, binding) in sorted_bindings.iter_mut().enumerate() {
|
|
binding.raw.index = index as u32;
|
|
}
|
|
let mut bindings = Vec::from(sorted_bindings);
|
|
bindings.extend([bytes(6, &sorted_rows), bytes(7, &n), bytes(8, &k)]);
|
|
return dispatch_specialized(
|
|
&name,
|
|
&bindings,
|
|
&constants,
|
|
[n.div_ceil(bn), sorted_rows.div_ceil(bm), 1],
|
|
[32, wn, wm],
|
|
);
|
|
}
|
|
let matrix = m >= qmv_limit(k, n);
|
|
let (name, groups, threads) = if matrix {
|
|
if nax && k.is_multiple_of(64) {
|
|
(
|
|
format!(
|
|
"affine_gather_qmm_t_nax_{suffix}_bm64_bn64_bk64_wm2_wn2_alN_{}",
|
|
n.is_multiple_of(64)
|
|
),
|
|
[n.div_ceil(64), m.div_ceil(64), routes],
|
|
[32, 2, 2],
|
|
)
|
|
} else {
|
|
(
|
|
format!("affine_gather_qmm_t_{suffix}_alN_{}", n.is_multiple_of(32)),
|
|
[n.div_ceil(32), m.div_ceil(32), routes],
|
|
[32, 2, 2],
|
|
)
|
|
}
|
|
} else {
|
|
let fast = n.is_multiple_of(8) && k.is_multiple_of((32 / bits) * 2 * 32);
|
|
(
|
|
format!(
|
|
"affine_gather_qmv{}_{suffix}",
|
|
if fast { "_fast" } else { "" }
|
|
),
|
|
[m, n.div_ceil(8), routes],
|
|
[32, 2, 1],
|
|
)
|
|
};
|
|
|
|
let x_dim = (x_shape.len() - 2) as u32;
|
|
let w_dim = (w_shape.len() - 2) as u32;
|
|
let batch_dim = batch_shape.len() as u32;
|
|
let mut bindings = Vec::from(bindings);
|
|
bindings.extend([bytes(7, &k), bytes(8, &n)]);
|
|
if matrix {
|
|
bindings.push(bytes(9, &m));
|
|
}
|
|
let base = bindings.len() as u32;
|
|
bindings.extend([
|
|
bytes(base, &x_dim),
|
|
bytes(base + 1, x_shape),
|
|
bytes(base + 2, x_strides),
|
|
bytes(base + 3, &w_dim),
|
|
bytes(base + 4, w_shape),
|
|
bytes(base + 5, w_strides),
|
|
bytes(base + 6, s_strides),
|
|
bytes(base + 7, b_strides),
|
|
bytes(base + 8, &batch_dim),
|
|
bytes(base + 9, batch_shape),
|
|
bytes(base + 10, lhs_strides),
|
|
bytes(base + 11, rhs_strides),
|
|
]);
|
|
dispatch(&name, &bindings, groups, threads)
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
#[cfg(test)]
|
|
pub(super) enum GateUpWeights<'a> {
|
|
Fused {
|
|
pack: [&'a Buffer; 3],
|
|
quant: [u32; 2],
|
|
},
|
|
Separate {
|
|
gate: [&'a Buffer; 3],
|
|
up: [&'a Buffer; 3],
|
|
gate_quant: [u32; 2],
|
|
up_quant: [u32; 2],
|
|
},
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct SwitchWeights<'a> {
|
|
pub(super) input: GateUpWeights<'a>,
|
|
pub(super) down: [&'a Buffer; 3],
|
|
pub(super) down_quant: [u32; 2],
|
|
}
|
|
|
|
/// Actual shapeless compiled swiglu for two contiguous projection outputs.
|
|
#[cfg(test)]
|
|
fn swiglu_bf16(gate: &Buffer, up: &Buffer, output: &Buffer, count: u32) -> Result<(), String> {
|
|
if count == 0
|
|
|| count > i32::MAX as u32
|
|
|| [gate, up, output]
|
|
.iter()
|
|
.any(|buffer| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < u64::from(count) * 2)
|
|
{
|
|
return Err("invalid MTPLX contiguous SwiGLU input".into());
|
|
}
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"kernel_qwen_mtplx_swiglu_bf16_contiguous"
|
|
} else {
|
|
"kernel_qwen_mtplx_swiglu_bf16_contiguous_n"
|
|
},
|
|
&[
|
|
tensor(0, gate).input(count as usize),
|
|
tensor(1, up).input(count as usize),
|
|
tensor(2, output).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct SwitchOutput {
|
|
pub(super) hidden: Buffer,
|
|
_temporaries: Vec<Buffer>,
|
|
_sorted: Option<GatherSorted>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn flatten_indices(input: &Buffer, [rows, width, stride]: [u32; 3]) -> Result<Buffer, String> {
|
|
let count = u64::from(rows) * u64::from(width);
|
|
if rows == 0
|
|
|| width == 0
|
|
|| stride < width
|
|
|| count > i32::MAX as u64
|
|
|| unsafe { ds4_gpu_tensor_bytes(input.raw()) }
|
|
< (u64::from(rows - 1) * u64::from(stride) + u64::from(width)) * 4
|
|
{
|
|
return Err("invalid MTPLX index flatten view".into());
|
|
}
|
|
if rows == 1 || stride == width {
|
|
return input.view(0, count * 4);
|
|
}
|
|
let flat = Buffer::mtplx_bytes(count * 4)?;
|
|
let grid = [width, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_copyuint32uint32",
|
|
&[
|
|
tensor(0, input)
|
|
.input((u64::from(rows - 1) * u64::from(stride) + u64::from(width)) as usize),
|
|
tensor(1, &flat).output(count as usize),
|
|
bytes(3, &[u64::from(stride), 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
Ok(flat)
|
|
}
|
|
|
|
/// Complete SwitchGLU or _FusedGateUpSwitchGLU with already-routed expert indices.
|
|
/// Own the result until GPU completion; router/shared-expert combination belongs
|
|
/// to SparseMoeBlock, not this reference submodule.
|
|
#[cfg(test)]
|
|
pub(super) fn switch_glu(
|
|
weights: SwitchWeights<'_>,
|
|
input: &Buffer,
|
|
indices: &Buffer,
|
|
[tokens, top_k, k, intermediate, experts]: [u32; 5],
|
|
) -> Result<SwitchOutput, String> {
|
|
switch_glu_strided(
|
|
weights,
|
|
input,
|
|
indices,
|
|
[tokens, top_k, k, intermediate, experts],
|
|
top_k,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn switch_glu_strided(
|
|
weights: SwitchWeights<'_>,
|
|
input: &Buffer,
|
|
indices: &Buffer,
|
|
[tokens, top_k, k, intermediate, experts]: [u32; 5],
|
|
index_stride: u32,
|
|
) -> Result<SwitchOutput, String> {
|
|
let routes = tokens
|
|
.checked_mul(top_k)
|
|
.filter(|v| *v > 0 && *v <= i32::MAX as u32)
|
|
.ok_or("invalid MTPLX switch route count")?;
|
|
if [k, intermediate, experts]
|
|
.iter()
|
|
.any(|v| *v == 0 || *v > i32::MAX as u32)
|
|
|| u64::from(routes) * u64::from(intermediate) * 2 > i32::MAX as u64
|
|
|| u64::from(routes) * u64::from(k) > i32::MAX as u64
|
|
|| index_stride < top_k
|
|
|| unsafe { ds4_gpu_tensor_bytes(indices.raw()) }
|
|
< (u64::from(tokens - 1) * u64::from(index_stride) + u64::from(top_k)) * 4
|
|
{
|
|
return Err("invalid MTPLX switch dimensions".into());
|
|
}
|
|
let arange = |size: u32| -> Result<Buffer, String> {
|
|
let result = Buffer::mtplx_bytes(u64::from(size) * 4)?;
|
|
dispatch_geometry(
|
|
"arangeuint32",
|
|
&[
|
|
bytes(0, &0_u32),
|
|
bytes(1, &1_u32),
|
|
tensor(2, &result).output(size as usize),
|
|
],
|
|
&[],
|
|
[size, 1, 1],
|
|
[size.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
Ok(result)
|
|
};
|
|
// _gather_sort flattens the strided top-k view. Only that flatten needs
|
|
// materialization; the unsorted GatherQMM path accepts original strides.
|
|
let flat_indices = if routes >= 64 {
|
|
Some(flatten_indices(indices, [tokens, top_k, index_stride])?)
|
|
} else {
|
|
None
|
|
};
|
|
let sorted = if routes >= 64 {
|
|
Some(gather_sort(
|
|
input,
|
|
flat_indices.as_ref().unwrap_or(indices),
|
|
[tokens, top_k, k],
|
|
)?)
|
|
} else {
|
|
None
|
|
};
|
|
let route_range = arange(routes)?;
|
|
let token_range = if sorted.is_none() {
|
|
Some(arange(tokens)?)
|
|
} else {
|
|
None
|
|
};
|
|
let (x, rhs, lhs, batches, repeat) = if let Some(sorted) = &sorted {
|
|
(&sorted.x, &sorted.indices, &route_range, routes, 1)
|
|
} else {
|
|
(input, indices, token_range.as_ref().unwrap(), tokens, top_k)
|
|
};
|
|
let rhs_layout = if sorted.is_some() {
|
|
[routes; 2]
|
|
} else {
|
|
[top_k, index_stride]
|
|
};
|
|
let h = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 2)?;
|
|
let y = Buffer::mtplx_bytes(u64::from(routes) * u64::from(k) * 2)?;
|
|
let activation_temporary = match weights.input {
|
|
GateUpWeights::Fused {
|
|
pack: [w, s, b],
|
|
quant,
|
|
} => {
|
|
let gu = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 4)?;
|
|
gather_bf16_broadcast(
|
|
[w, s, b, x, lhs, rhs, &gu],
|
|
[batches, 1, intermediate * 2, k, experts, routes],
|
|
quant,
|
|
sorted.is_some(),
|
|
true,
|
|
repeat,
|
|
rhs_layout,
|
|
)?;
|
|
let activated = gate_up_bf16(&gu, &h, [routes, intermediate])?;
|
|
vec![gu, activated]
|
|
}
|
|
GateUpWeights::Separate {
|
|
gate,
|
|
up,
|
|
gate_quant,
|
|
up_quant,
|
|
} => {
|
|
let gate_out = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 2)?;
|
|
let up_out = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 2)?;
|
|
// SwitchGLU constructs Up before Gate; Qwen3NextMLP uses the reverse.
|
|
for ([w, s, b], quant, output) in
|
|
[(up, up_quant, &up_out), (gate, gate_quant, &gate_out)]
|
|
{
|
|
gather_bf16_broadcast(
|
|
[w, s, b, x, lhs, rhs, output],
|
|
[batches, 1, intermediate, k, experts, routes],
|
|
quant,
|
|
sorted.is_some(),
|
|
true,
|
|
repeat,
|
|
rhs_layout,
|
|
)?;
|
|
}
|
|
swiglu_bf16(&gate_out, &up_out, &h, routes * intermediate)?;
|
|
vec![gate_out, up_out]
|
|
}
|
|
};
|
|
let [w, s, b] = weights.down;
|
|
gather_bf16_broadcast(
|
|
[w, s, b, &h, &route_range, rhs, &y],
|
|
[routes, 1, k, intermediate, experts, routes],
|
|
weights.down_quant,
|
|
sorted.is_some(),
|
|
true,
|
|
1,
|
|
rhs_layout,
|
|
)?;
|
|
let hidden = if let Some(sorted) = &sorted {
|
|
let unsorted = Buffer::mtplx_bytes(u64::from(routes) * u64::from(k) * 2)?;
|
|
gather_front(
|
|
&y,
|
|
&sorted.inverse,
|
|
&unsorted,
|
|
[routes, k, routes],
|
|
SortDtype::Bf16,
|
|
)?;
|
|
unsorted
|
|
} else {
|
|
y.view(0, u64::from(routes) * u64::from(k) * 2)?
|
|
};
|
|
let mut temporary = vec![route_range, h, y];
|
|
temporary.extend(activation_temporary);
|
|
temporary.extend(token_range);
|
|
temporary.extend(flat_indices);
|
|
Ok(SwitchOutput {
|
|
hidden,
|
|
_temporaries: temporary,
|
|
_sorted: sorted,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct MoeWeights<'a> {
|
|
pub(super) experts: SwitchWeights<'a>,
|
|
pub(super) shared: SwitchWeights<'a>,
|
|
pub(super) router: [&'a Buffer; 3],
|
|
pub(super) router_quant: [u32; 2],
|
|
pub(super) shared_gate: [&'a Buffer; 3],
|
|
pub(super) shared_gate_quant: [u32; 2],
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct MoeOutput {
|
|
pub(super) hidden: Buffer,
|
|
pub(super) used_fused: bool,
|
|
_router: RouterOutput,
|
|
_switch: Option<SwitchOutput>,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// Connected SparseMoeBlock with separate or sanitize-fused expert/shared MLPs. The two flags
|
|
/// are MTPLX_FUSED_MOE_DECODE / VERIFY; both default off in the reference.
|
|
#[cfg(test)]
|
|
pub(super) fn sparse_moe(
|
|
weights: MoeWeights<'_>,
|
|
input: &Buffer,
|
|
[tokens, experts, k, intermediate, shared_intermediate]: [u32; 5],
|
|
normalize: bool,
|
|
[fused_decode, fused_verify]: [bool; 2],
|
|
) -> Result<MoeOutput, String> {
|
|
let count = u64::from(tokens) * u64::from(k);
|
|
if tokens == 0
|
|
|| experts < 10
|
|
|| k == 0
|
|
|| count * 10 > i32::MAX as u64
|
|
|| u64::from(tokens) * u64::from(experts) > i32::MAX as u64
|
|
|| [intermediate, shared_intermediate]
|
|
.iter()
|
|
.any(|d| *d == 0 || u64::from(tokens) * u64::from(*d) * 20 > i32::MAX as u64)
|
|
{
|
|
return Err("invalid MTPLX sparse MoE shape".into());
|
|
}
|
|
let logits = Buffer::mtplx_bytes(u64::from(tokens) * u64::from(experts) * 2)?;
|
|
let [w, s, b] = weights.router;
|
|
let [bits, group] = weights.router_quant;
|
|
let mut temporary = affine_bf16(
|
|
[w, s, b, input, &logits],
|
|
[tokens, experts, k],
|
|
bits,
|
|
group,
|
|
true,
|
|
)?;
|
|
let routing = router(logits, [tokens, experts], normalize)?;
|
|
let used_fused = ((tokens == 1 && fused_decode) || ((2..=4).contains(&tokens) && fused_verify))
|
|
&& matches!(
|
|
weights.experts.input,
|
|
GateUpWeights::Fused {
|
|
quant: [4, 32 | 64],
|
|
..
|
|
}
|
|
)
|
|
&& weights.experts.down_quant[0] == 4
|
|
&& matches!(weights.experts.down_quant[1], 32 | 64);
|
|
let hidden = Buffer::mtplx_bytes(count * 2)?;
|
|
let switched = if used_fused {
|
|
if k != 2560 || intermediate != 640 {
|
|
return Err("MTPLX fused MoE requires the installed 2560/640 model shape".into());
|
|
}
|
|
temporary.extend(fused_experts(
|
|
&weights.experts,
|
|
input,
|
|
&routing,
|
|
&hidden,
|
|
[tokens, experts],
|
|
)?);
|
|
None
|
|
} else {
|
|
let switched = switch_glu_strided(
|
|
weights.experts,
|
|
input,
|
|
&routing.indices,
|
|
[tokens, 10, k, intermediate, experts],
|
|
routing.index_row_stride,
|
|
)?;
|
|
let grid = [k, tokens * 10, 1];
|
|
dispatch_geometry(
|
|
"g2_Multiplybfloat16",
|
|
&[
|
|
tensor(0, &switched.hidden).input((count * 10) as usize),
|
|
tensor(1, &routing.scores).input((u64::from(tokens) * 10) as usize),
|
|
tensor(2, &switched.hidden).output((count * 10) as usize),
|
|
bytes(3, &[u64::from(k), 1]),
|
|
bytes(4, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let gx = k.div_ceil(4).min(32);
|
|
dispatch(
|
|
"col_reduce_small_1_reduce_sumbfloat16",
|
|
&[
|
|
tensor(0, &switched.hidden).input((count * 10) as usize),
|
|
tensor(1, &hidden).output(count as usize),
|
|
bytes(2, &10_u64),
|
|
bytes(3, &u64::from(k)),
|
|
bytes(4, &tokens),
|
|
bytes(5, &(u64::from(k) * 10)),
|
|
bytes(6, &1_u32),
|
|
bytes(7, &10_u32),
|
|
bytes(8, &u64::from(k)),
|
|
bytes(9, &1_u32),
|
|
bytes(10, &1_u64),
|
|
],
|
|
[k.div_ceil(4).div_ceil(gx), tokens, 1],
|
|
[gx, 8, 1],
|
|
)?;
|
|
Some(switched)
|
|
};
|
|
// The fused branch evaluates sigmoid(shared_gate(x)) before shared_expert(x);
|
|
// the parent stock branch evaluates the MLP first. Keep both source orders.
|
|
let gate = || {
|
|
shared_gate_sigmoid(
|
|
weights.shared_gate,
|
|
weights.shared_gate_quant,
|
|
input,
|
|
[tokens, k],
|
|
)
|
|
};
|
|
let early_gate = if used_fused { Some(gate()?) } else { None };
|
|
let shared_h = Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 2)?;
|
|
let shared = Buffer::mtplx_bytes(count * 2)?;
|
|
match weights.shared.input {
|
|
GateUpWeights::Fused {
|
|
pack: [w, s, b],
|
|
quant: [bits, group],
|
|
} => {
|
|
let gu = Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 4)?;
|
|
temporary.extend(affine_bf16(
|
|
[w, s, b, input, &gu],
|
|
[tokens, shared_intermediate * 2, k],
|
|
bits,
|
|
group,
|
|
true,
|
|
)?);
|
|
temporary.push(gate_up_bf16(&gu, &shared_h, [tokens, shared_intermediate])?);
|
|
temporary.push(gu);
|
|
}
|
|
GateUpWeights::Separate {
|
|
gate,
|
|
up,
|
|
gate_quant,
|
|
up_quant,
|
|
} => {
|
|
let gate_out =
|
|
Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 2)?;
|
|
let up_out =
|
|
Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 2)?;
|
|
for ([w, s, b], [bits, group], output) in
|
|
[(gate, gate_quant, &gate_out), (up, up_quant, &up_out)]
|
|
{
|
|
temporary.extend(affine_bf16(
|
|
[w, s, b, input, output],
|
|
[tokens, shared_intermediate, k],
|
|
bits,
|
|
group,
|
|
true,
|
|
)?);
|
|
}
|
|
swiglu_bf16(&gate_out, &up_out, &shared_h, tokens * shared_intermediate)?;
|
|
temporary.extend([gate_out, up_out]);
|
|
}
|
|
}
|
|
let [w, s, b] = weights.shared.down;
|
|
let [bits, group] = weights.shared.down_quant;
|
|
temporary.extend(affine_bf16(
|
|
[w, s, b, &shared_h, &shared],
|
|
[tokens, k, shared_intermediate],
|
|
bits,
|
|
group,
|
|
true,
|
|
)?);
|
|
let (shared_gate, gate_temporary) = match early_gate {
|
|
Some(result) => result,
|
|
None => gate()?,
|
|
};
|
|
temporary.extend(gate_temporary);
|
|
if tokens == 1 {
|
|
let work = if k < 65536 { 1 } else { 4 };
|
|
let threads = k.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"sv_Multiplybfloat16"
|
|
} else {
|
|
"svn_Multiplybfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &shared_gate).input(tokens as usize),
|
|
tensor(1, &shared).input(count as usize),
|
|
tensor(2, &shared).output(count as usize),
|
|
bytes(3, &k),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let grid = [k, tokens, 1];
|
|
dispatch_geometry(
|
|
"g2_Multiplybfloat16",
|
|
&[
|
|
tensor(0, &shared_gate).input(tokens as usize),
|
|
tensor(1, &shared).input(count as usize),
|
|
tensor(2, &shared).output(count as usize),
|
|
bytes(3, &[1_u64, 0]),
|
|
bytes(4, &[u64::from(k), 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let count = count as u32;
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"vv_Addbfloat16"
|
|
} else {
|
|
"vvn_Addbfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &hidden).input(count as usize),
|
|
tensor(1, &shared).input(count as usize),
|
|
tensor(2, &hidden).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
temporary.extend([shared_h, shared, shared_gate]);
|
|
Ok(MoeOutput {
|
|
hidden,
|
|
used_fused,
|
|
_router: routing,
|
|
_switch: switched,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn shared_gate_sigmoid(
|
|
[w, s, b]: [&Buffer; 3],
|
|
[bits, group]: [u32; 2],
|
|
input: &Buffer,
|
|
[tokens, k]: [u32; 2],
|
|
) -> Result<(Buffer, Vec<Buffer>), String> {
|
|
let gate = Buffer::mtplx_bytes(u64::from(tokens) * 2)?;
|
|
let temporary = affine_bf16([w, s, b, input, &gate], [tokens, 1, k], bits, group, true)?;
|
|
let work = if tokens < 65536 { 1 } else { 4 };
|
|
let threads = tokens.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"v_Sigmoidbfloat16bfloat16"
|
|
} else {
|
|
"vn_Sigmoidbfloat16bfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &gate).input(tokens as usize),
|
|
tensor(1, &gate).output(tokens as usize),
|
|
bytes(2, &tokens),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
Ok((gate, temporary))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn fused_experts(
|
|
weights: &SwitchWeights<'_>,
|
|
input: &Buffer,
|
|
routing: &RouterOutput,
|
|
output: &Buffer,
|
|
[tokens, experts]: [u32; 2],
|
|
) -> Result<Vec<Buffer>, String> {
|
|
let GateUpWeights::Fused {
|
|
pack: gu_pack,
|
|
quant: [4, gu_group],
|
|
} = weights.input
|
|
else {
|
|
return Err("MTPLX fused MoE requires a fused Q4 gate/up pack".into());
|
|
};
|
|
// Entry conditions are checked by SparseMoeBlock, including the installed
|
|
// dimensions for these template instantiations. Validate all six packs.
|
|
for (pack, n, k, group) in [
|
|
(gu_pack, 1280_u64, 2560_u64, gu_group),
|
|
(weights.down, 2560, 640, weights.down_quant[1]),
|
|
] {
|
|
let count = u64::from(experts) * n * k;
|
|
for (buffer, size) in pack.into_iter().zip([
|
|
count / 2,
|
|
count / u64::from(group) * 2,
|
|
count / u64::from(group) * 2,
|
|
]) {
|
|
if unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < size {
|
|
return Err("MTPLX fused MoE pack smaller than shape".into());
|
|
}
|
|
}
|
|
}
|
|
let indices = flatten_indices(&routing.indices, [tokens, 10, routing.index_row_stride])?;
|
|
let routes = tokens * 10;
|
|
let scores = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
|
|
dispatch_geometry(
|
|
"v_copybfloat16float32",
|
|
&[
|
|
tensor(0, &routing.scores).input(routes as usize),
|
|
tensor(1, &scores).output(routes as usize),
|
|
bytes(2, &routes),
|
|
],
|
|
&[],
|
|
[routes, 1, 1],
|
|
[routes, 1, 1],
|
|
true,
|
|
)?;
|
|
let h = Buffer::mtplx_bytes(u64::from(routes) * 640 * 2)?;
|
|
let suffix = |group| {
|
|
if tokens == 1 {
|
|
format!("g{group}_bf16")
|
|
} else {
|
|
format!("m_m{tokens}_g{group}_bf16")
|
|
}
|
|
};
|
|
let [w, s, b] = gu_pack;
|
|
dispatch(
|
|
&format!("kernel_qwen_mtplx_moe_glu_h_{}", suffix(gu_group)),
|
|
&[
|
|
tensor(0, input).input(tokens as usize * 2560),
|
|
tensor(1, w).input(experts as usize * 1280 * 2560 / 8),
|
|
tensor(2, s).input(experts as usize * 1280 * 2560 / gu_group as usize),
|
|
tensor(3, b).input(experts as usize * 1280 * 2560 / gu_group as usize),
|
|
tensor(4, &indices).input(routes as usize),
|
|
tensor(5, &h).output(routes as usize * 640),
|
|
],
|
|
[(routes * 640).div_ceil(32), 1, 1],
|
|
[1024, 1, 1],
|
|
)?;
|
|
let [w, s, b] = weights.down;
|
|
dispatch(
|
|
&format!(
|
|
"kernel_qwen_mtplx_moe_down_y_{}",
|
|
suffix(weights.down_quant[1])
|
|
),
|
|
&[
|
|
tensor(0, &h).input(routes as usize * 640),
|
|
tensor(1, w).input(experts as usize * 2560 * 640 / 8),
|
|
tensor(2, s).input(experts as usize * 2560 * 640 / weights.down_quant[1] as usize),
|
|
tensor(3, b).input(experts as usize * 2560 * 640 / weights.down_quant[1] as usize),
|
|
tensor(4, &indices).input(routes as usize),
|
|
tensor(5, &scores).input(routes as usize),
|
|
tensor(6, output).output(tokens as usize * 2560),
|
|
],
|
|
[(tokens * 2560).div_ceil(32), 1, 1],
|
|
[1024, 1, 1],
|
|
)?;
|
|
Ok(vec![indices, scores, h])
|
|
}
|
|
|
|
/// Keep returned intermediate buffers alive until the enclosing command batch completes.
|
|
#[cfg(test)]
|
|
pub(super) fn affine_bf16(
|
|
buffers: [&Buffer; 5],
|
|
shape: [u32; 3],
|
|
bits: u32,
|
|
group: u32,
|
|
nax: bool,
|
|
) -> Result<Vec<Buffer>, String> {
|
|
affine_bf16_layout(buffers, shape, bits, group, nax, None)
|
|
}
|
|
|
|
fn affine_bf16_layout(
|
|
buffers: [&Buffer; 5],
|
|
shape: [u32; 3],
|
|
bits: u32,
|
|
group: u32,
|
|
nax: bool,
|
|
input_layout: Option<[(u64, usize); 4]>,
|
|
) -> Result<Vec<Buffer>, String> {
|
|
let [m, n, k] = shape;
|
|
let plan = affine_plan(shape, bits, group, nax)?;
|
|
let expected = [
|
|
u64::from(n) * u64::from(k) / (8 / u64::from(bits)),
|
|
u64::from(n) * u64::from(k / group) * 2,
|
|
u64::from(n) * u64::from(k / group) * 2,
|
|
u64::from(m) * u64::from(k) * 2,
|
|
u64::from(m) * u64::from(n) * 2,
|
|
];
|
|
let elements: [usize; 5] = std::array::from_fn(|i| (expected[i] / [4, 2, 2, 2, 2][i]) as usize);
|
|
let input_layout = input_layout.unwrap_or_else(|| std::array::from_fn(|i| (0, elements[i])));
|
|
for (i, (buffer, size)) in buffers.iter().zip(expected).enumerate() {
|
|
let offset = input_layout.get(i).map_or(0, |layout| layout.0);
|
|
if offset
|
|
.checked_add(size)
|
|
.is_none_or(|end| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < end)
|
|
{
|
|
return Err("MTPLX affine binding is smaller than its shape".into());
|
|
}
|
|
}
|
|
let mut temporary = Vec::new();
|
|
if plan.splits > 1 {
|
|
temporary.push(Buffer::mtplx_bytes(expected[4] * u64::from(plan.splits))?);
|
|
encoder::add_temporaries(&temporary);
|
|
}
|
|
let output = temporary.first().unwrap_or(buffers[4]);
|
|
let mut bindings = buffers[..4]
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, b)| {
|
|
tensor(i as u32, b)
|
|
.at_byte_offset(input_layout[i].0)
|
|
.input(input_layout[i].1)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
bindings.push(tensor(4, output).output(elements[4] * plan.splits as usize));
|
|
bindings.extend(
|
|
plan.parameters
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, p)| bytes(i as u32 + 5, p)),
|
|
);
|
|
dispatch(&plan.name, &bindings, plan.groups, plan.threads)?;
|
|
if plan.splits > 1
|
|
&& let Some(extra) = reduce_split_bf16(output, buffers[4], plan.splits, m * n)?
|
|
{
|
|
temporary.push(extra);
|
|
}
|
|
Ok(temporary)
|
|
}
|
|
|
|
fn reduce_split_bf16(
|
|
input: &Buffer,
|
|
output: &Buffer,
|
|
splits: u32,
|
|
width: u32,
|
|
) -> Result<Option<Buffer>, String> {
|
|
// ColReduceArgs for contiguous [split, M, N], reducing axis 0.
|
|
let size = u64::from(splits);
|
|
let stride = u64::from(width);
|
|
let zero = 0_u32;
|
|
let zero_stride = 0_u64;
|
|
let ndim = 1_u32;
|
|
let non_col = 1_u64;
|
|
let two_pass = splits > 256 && width / 32 < 1024;
|
|
let intermediate = if two_pass {
|
|
Some(Buffer::mtplx_bytes(stride * 32 * 2)?)
|
|
} else {
|
|
None
|
|
};
|
|
if let Some(buffer) = &intermediate {
|
|
encoder::add_temporaries(std::slice::from_ref(buffer));
|
|
}
|
|
let mut bindings = vec![
|
|
tensor(0, input).input((size * stride) as usize),
|
|
tensor(1, intermediate.as_ref().unwrap_or(output))
|
|
.output((stride * if two_pass { 32 } else { 1 }) as usize),
|
|
bytes(2, &size),
|
|
bytes(3, &stride),
|
|
bytes(4, &zero),
|
|
bytes(5, &zero_stride),
|
|
bytes(6, &zero),
|
|
bytes(7, &splits),
|
|
bytes(8, &stride),
|
|
bytes(9, &ndim),
|
|
bytes(10, &non_col),
|
|
];
|
|
let large = if size * stride > i32::MAX as u64 {
|
|
"_large"
|
|
} else {
|
|
""
|
|
};
|
|
if splits < 32 {
|
|
let x = width.div_ceil(4).min(32);
|
|
dispatch(
|
|
&format!("col_reduce_small{large}_1_reduce_sumbfloat16"),
|
|
&bindings,
|
|
[width.div_ceil(4).div_ceil(x), 1, 1],
|
|
[x, splits.min(8), 1],
|
|
)?;
|
|
} else if two_pass {
|
|
bindings.push(bytes(11, &non_col));
|
|
dispatch(
|
|
&format!("col_reduce_2pass{large}_1_32_32_reduce_sumbfloat16"),
|
|
&bindings,
|
|
[width.div_ceil(32), 32, 1],
|
|
[256, 1, 1],
|
|
)?;
|
|
reduce_split_bf16(intermediate.as_ref().unwrap(), output, 32, width)?;
|
|
} else {
|
|
dispatch(
|
|
&format!("col_reduce_looped{large}_1_32_32_reduce_sumbfloat16"),
|
|
&bindings,
|
|
[width.div_ceil(32), 1, 1],
|
|
[256, 1, 1],
|
|
)?;
|
|
}
|
|
Ok(intermediate)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct ConvNormOutput {
|
|
pub(super) q: Buffer,
|
|
pub(super) k: Buffer,
|
|
pub(super) v: Buffer,
|
|
pub(super) conv: Buffer,
|
|
/// Pre-convolution rows retained by the reference's verify capture.
|
|
pub(super) qkv: Buffer,
|
|
pub(super) used_fused: bool,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// Stock GatedDeltaNet._l2norm, including its BF16 boundary before Q scaling.
|
|
#[cfg(test)]
|
|
fn gdn_l2_norm(input: &Buffer, output: &Buffer, rows: u32, q: bool) -> Result<Vec<Buffer>, String> {
|
|
let count = rows * 2048;
|
|
let heads = rows * 16;
|
|
let xf = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
let sq = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
let sum = Buffer::mtplx_bytes(u64::from(heads) * 4)?;
|
|
if rows == 1 {
|
|
cast_bf16_f32(input, &xf, count, true)?;
|
|
} else {
|
|
let grid = [2048, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_copybfloat16float32",
|
|
&[
|
|
tensor(0, input).input(((rows - 1) * 10240 + 2048) as usize),
|
|
tensor(1, &xf).output(count as usize),
|
|
bytes(3, &[10240_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let work = if count < 65536 { 1 } else { 2 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"vv_Multiplyfloat32"
|
|
} else {
|
|
"vvn_Multiplyfloat32"
|
|
},
|
|
&[
|
|
tensor(0, &xf).input(count as usize),
|
|
tensor(1, &xf).input(count as usize),
|
|
tensor(2, &sq).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
if heads >= 32 {
|
|
dispatch_geometry(
|
|
"row_reduce_simple_sumfloat32",
|
|
&[
|
|
tensor(0, &sq).input(count as usize),
|
|
tensor(1, &sum).output(heads as usize),
|
|
bytes(2, &128_u64),
|
|
bytes(3, &u64::from(heads)),
|
|
],
|
|
&[],
|
|
[32, heads.div_ceil(4), 1],
|
|
[32, 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
dispatch_geometry(
|
|
"row_reduce_looped_1_reduce_sumfloat32",
|
|
&[
|
|
tensor(0, &sq).input(count as usize),
|
|
tensor(1, &sum).output(heads as usize),
|
|
bytes(2, &128_u64),
|
|
bytes(3, &1_u64),
|
|
bytes(4, &heads),
|
|
bytes(5, &128_u64),
|
|
bytes(6, &1_u32),
|
|
bytes(7, &0_u32),
|
|
bytes(8, &0_u64),
|
|
bytes(9, &0_u32),
|
|
],
|
|
&[],
|
|
[32, heads, 1],
|
|
[32, 1, 1],
|
|
true,
|
|
)?;
|
|
}
|
|
let epsilon = scalar_buffer(&1e-6_f32.to_le_bytes())?;
|
|
let work = if heads < 65536 { 1 } else { 2 };
|
|
let threads = heads.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"vs_Addfloat32"
|
|
} else {
|
|
"vsn_Addfloat32"
|
|
},
|
|
&[
|
|
tensor(0, &sum).input(heads as usize),
|
|
tensor(1, &epsilon).input(1),
|
|
tensor(2, &sum).output(heads as usize),
|
|
bytes(3, &heads),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"v_Rsqrtfloat32float32"
|
|
} else {
|
|
"vn_Rsqrtfloat32float32"
|
|
},
|
|
&[
|
|
tensor(0, &sum).input(heads as usize),
|
|
tensor(1, &sum).output(heads as usize),
|
|
bytes(2, &heads),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
let grid = [128, heads, 1];
|
|
dispatch_geometry(
|
|
"g2_Multiplyfloat32",
|
|
&[
|
|
tensor(0, &xf).input(count as usize),
|
|
tensor(1, &sum).input(heads as usize),
|
|
tensor(2, &xf).output(count as usize),
|
|
bytes(3, &[128_u64, 1]),
|
|
bytes(4, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
cast_bf16_f32(&xf, output, count, false)?;
|
|
if q {
|
|
// Weak Python scalar converts to BF16 BEFORE Multiplybfloat16.
|
|
let bits = 128_f32.powf(-0.5).to_bits();
|
|
let scale = ((bits + 0x7fff + ((bits >> 16) & 1)) >> 16) as u16;
|
|
let scale = scalar_buffer(&scale.to_le_bytes())?;
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"sv_Multiplybfloat16"
|
|
} else {
|
|
"svn_Multiplybfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &scale).input(1),
|
|
tensor(1, output).input(count as usize),
|
|
tensor(2, output).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
}
|
|
Ok(vec![xf, sq, sum])
|
|
}
|
|
|
|
/// GDN conv/state and Q/K/V preparation for B=1. `fused` contains the
|
|
/// model's two conv-norm flags; cache absence/lengths/mask suppress fusion.
|
|
/// `length` is the cache's optional one-element INT32 valid-length tensor.
|
|
#[cfg(test)]
|
|
pub(super) fn gdn_conv_norm(
|
|
input: &Buffer,
|
|
weight: &Buffer,
|
|
[rows, stride]: [u32; 2],
|
|
[state, length]: [Option<&Buffer>; 2],
|
|
mask: Option<&Buffer>,
|
|
fused: [bool; 2],
|
|
has_cache: bool,
|
|
) -> Result<ConvNormOutput, String> {
|
|
let count = u64::from(rows) * 10240;
|
|
if rows == 0
|
|
|| stride < 10240
|
|
|| (u64::from(rows) + 3) * 10240 > i32::MAX as u64
|
|
|| unsafe { ds4_gpu_tensor_bytes(input.raw()) }
|
|
< (u64::from(rows - 1) * u64::from(stride) + 10240) * 2
|
|
|| unsafe { ds4_gpu_tensor_bytes(weight.raw()) } < 10240 * 4 * 2
|
|
|| state.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < 3 * 10240 * 2)
|
|
|| mask.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(rows))
|
|
|| length.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < 4)
|
|
{
|
|
return Err("invalid MTPLX GDN conv inputs".into());
|
|
}
|
|
let mut temporary = Vec::new();
|
|
let input_span = ((rows - 1) as usize * stride as usize) + 10240;
|
|
let qkv = if let Some(mask) = mask {
|
|
let qkv = Buffer::mtplx_bytes(count * 2)?;
|
|
let zero = scalar_buffer(&0_u16.to_le_bytes())?;
|
|
let grid = [10240, rows, 1];
|
|
dispatch_geometry(
|
|
if rows == 1 {
|
|
"g1_Selectbfloat16"
|
|
} else {
|
|
"g2_Selectbfloat16"
|
|
},
|
|
&[
|
|
tensor(0, mask).input(rows as usize),
|
|
tensor(1, input).input(input_span),
|
|
tensor(2, &zero).input(1),
|
|
tensor(3, &qkv).output(count as usize),
|
|
bytes(4, &if rows == 1 { [0_u64, 0] } else { [1_u64, 0] }),
|
|
bytes(
|
|
5,
|
|
&if rows == 1 {
|
|
[1_u64, 0]
|
|
} else {
|
|
[u64::from(stride), 1]
|
|
},
|
|
),
|
|
bytes(6, &[0_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
qkv
|
|
} else {
|
|
input.view(0, (u64::from(rows - 1) * u64::from(stride) + 10240) * 2)?
|
|
};
|
|
let stride = if mask.is_some() { 10240 } else { stride };
|
|
let qkv_span = (rows - 1) as usize * stride as usize + 10240;
|
|
let q = Buffer::mtplx_bytes(u64::from(rows) * 2048 * 2)?;
|
|
let k = Buffer::mtplx_bytes(u64::from(rows) * 2048 * 2)?;
|
|
let v;
|
|
let conv;
|
|
let used_fused = has_cache
|
|
&& mask.is_none()
|
|
&& length.is_none()
|
|
&& ((rows == 1 && fused[0]) || ((2..=6).contains(&rows) && fused[1]));
|
|
if used_fused {
|
|
v = Buffer::mtplx_bytes(u64::from(rows) * 6144 * 2)?;
|
|
conv = Buffer::mtplx_bytes(3 * 10240 * 2)?;
|
|
let packed;
|
|
let source = if rows == 1 || stride == 10240 {
|
|
&qkv
|
|
} else {
|
|
packed = Buffer::mtplx_bytes(count * 2)?;
|
|
let grid = [10240, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, &qkv).input(qkv_span),
|
|
tensor(1, &packed).output(count as usize),
|
|
bytes(3, &[u64::from(stride), 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
temporary.push(packed);
|
|
temporary.last().unwrap()
|
|
};
|
|
let source = source.view(0, count * 2)?;
|
|
let zero;
|
|
let state = if let Some(state) = state {
|
|
state
|
|
} else {
|
|
zero = Buffer::mtplx_bytes(3 * 10240 * 2)?;
|
|
let scalar = scalar_buffer(&0_u16.to_le_bytes())?;
|
|
// CustomKernel materializes the broadcast scalar with General copy.
|
|
dispatch_geometry(
|
|
"g1_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, &scalar).input(1),
|
|
tensor(1, &zero).output(3 * 10240),
|
|
bytes(3, &0_u64),
|
|
],
|
|
&[],
|
|
[3 * 10240, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)?;
|
|
temporary.push(zero);
|
|
temporary.last().unwrap()
|
|
};
|
|
let state = state.view(0, 3 * 10240 * 2)?;
|
|
dispatch(
|
|
&if rows == 1 {
|
|
"kernel_qwen_mtplx_gdn_conv_norm_bf16".into()
|
|
} else {
|
|
format!("kernel_qwen_mtplx_gdn_conv_norm_rows_s{rows}_bf16")
|
|
},
|
|
&[
|
|
tensor(0, &source).input(count as usize),
|
|
tensor(1, &state).input(3 * 10240),
|
|
tensor(2, weight).input(4 * 10240),
|
|
tensor(3, &q).output(rows as usize * 2048),
|
|
tensor(4, &k).output(rows as usize * 2048),
|
|
tensor(5, &v).output(rows as usize * 6144),
|
|
tensor(6, &conv).output(3 * 10240),
|
|
],
|
|
[10, 1, 1],
|
|
[1024, 1, 1],
|
|
)?;
|
|
encoder::add_temporaries(&temporary);
|
|
} else {
|
|
conv = Buffer::mtplx_bytes(3 * 10240 * 2)?;
|
|
let window = Buffer::mtplx_bytes((count + 3 * 10240) * 2)?;
|
|
let tail = window.view(3 * 10240 * 2, count * 2)?;
|
|
// Original concatenate_gpu uses concurrent disjoint out_slice writers.
|
|
let concurrent = encoder::concurrent_scope();
|
|
let zero = state
|
|
.is_none()
|
|
.then(|| scalar_buffer(&0_u16.to_le_bytes()))
|
|
.transpose()?;
|
|
// Concatenate uses GeneralGeneral copies, even for contiguous slices.
|
|
dispatch_geometry(
|
|
"gg1_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, state.or(zero.as_ref()).unwrap()).input(if state.is_some() {
|
|
3 * 10240
|
|
} else {
|
|
1
|
|
}),
|
|
tensor(1, &window).output(3 * 10240),
|
|
bytes(3, &u64::from(state.is_some())),
|
|
bytes(4, &1_u64),
|
|
],
|
|
&[],
|
|
[3 * 10240, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)?;
|
|
let contiguous = rows == 1 || stride == 10240;
|
|
let grid = if contiguous {
|
|
[count as u32, 1, 1]
|
|
} else {
|
|
[10240, rows, 1]
|
|
};
|
|
dispatch_geometry(
|
|
if contiguous {
|
|
"gg1_copybfloat16bfloat16"
|
|
} else {
|
|
"gg2_copybfloat16bfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &qkv).input(qkv_span),
|
|
tensor(1, &tail).output(count as usize),
|
|
bytes(
|
|
3,
|
|
&if contiguous {
|
|
[1_u64, 0]
|
|
} else {
|
|
[u64::from(stride), 1]
|
|
},
|
|
),
|
|
bytes(
|
|
4,
|
|
&if contiguous {
|
|
[1_u64, 0]
|
|
} else {
|
|
[10240_u64, 1]
|
|
},
|
|
),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
drop(concurrent);
|
|
if let Some(length) = length {
|
|
let zero = scalar_buffer(&0_i32.to_le_bytes())?;
|
|
let upper = scalar_buffer(&(rows as i32).to_le_bytes())?;
|
|
let end = Buffer::mtplx_bytes(4)?;
|
|
let positions = Buffer::mtplx_bytes(12)?;
|
|
let range = Buffer::mtplx_bytes(12)?;
|
|
dispatch_geometry(
|
|
"vs_Maximumint32",
|
|
&[
|
|
tensor(0, length).input(1),
|
|
tensor(1, &zero).input(1),
|
|
tensor(2, &end).output(1),
|
|
bytes(3, &1_u32),
|
|
],
|
|
&[],
|
|
[1, 1, 1],
|
|
[1, 1, 1],
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
"vs_Minimumint32",
|
|
&[
|
|
tensor(0, &end).input(1),
|
|
tensor(1, &upper).input(1),
|
|
tensor(2, &end).output(1),
|
|
bytes(3, &1_u32),
|
|
],
|
|
&[],
|
|
[1, 1, 1],
|
|
[1, 1, 1],
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &0_i32),
|
|
bytes(1, &1_i32),
|
|
tensor(2, &range).output(3),
|
|
],
|
|
&[],
|
|
[3, 1, 1],
|
|
[3, 1, 1],
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
"sv_Addint32",
|
|
&[
|
|
tensor(0, &end).input(1),
|
|
tensor(1, &range).input(3),
|
|
tensor(2, &positions).output(3),
|
|
bytes(3, &3_u32),
|
|
],
|
|
&[],
|
|
[3, 1, 1],
|
|
[3, 1, 1],
|
|
true,
|
|
)?;
|
|
let grid = [10240, 3, 1];
|
|
dispatch_geometry(
|
|
"kernel_qwen_mtplx_gather_axis_bf16_idxi32_int_10",
|
|
&[
|
|
tensor(0, &window).input(count as usize + 3 * 10240),
|
|
tensor(1, &positions).input(3),
|
|
tensor(2, &conv).output(3 * 10240),
|
|
bytes(3, &10240_u32),
|
|
bytes(4, &1_u64),
|
|
bytes(5, &0_u64),
|
|
bytes(6, &1_u64),
|
|
bytes(7, &0_u32),
|
|
bytes(8, &(rows + 3)),
|
|
bytes(9, &10240_u64),
|
|
bytes(10, &1_u64),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
temporary.extend([end, positions, range]);
|
|
} else {
|
|
let view = window.view(u64::from(rows) * 10240 * 2, 3 * 10240 * 2)?;
|
|
dispatch_geometry(
|
|
"g1_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, &view).input(3 * 10240),
|
|
tensor(1, &conv).output(3 * 10240),
|
|
bytes(3, &1_u64),
|
|
],
|
|
&[],
|
|
[3 * 10240, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)?;
|
|
}
|
|
let convolved = Buffer::mtplx_bytes(count * 2)?;
|
|
let grid = [10240, rows, 1];
|
|
dispatch_geometry(
|
|
"depthwise_conv_1d_bfloat16",
|
|
&[
|
|
tensor(0, &window).input(count as usize + 3 * 10240),
|
|
tensor(1, weight).input(4 * 10240),
|
|
tensor(2, &convolved).output(count as usize),
|
|
bytes(3, &[(rows + 3) * 10240, 10240, 1]),
|
|
bytes(4, &4_u32),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let count = count as u32;
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"kernel_qwen_mtplx_silu_bf16_contiguous"
|
|
} else {
|
|
"kernel_qwen_mtplx_silu_bf16_contiguous_n"
|
|
},
|
|
&[
|
|
tensor(0, &convolved).input(count as usize),
|
|
tensor(1, &convolved).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
temporary.extend(gdn_l2_norm(&convolved, &q, rows, true)?);
|
|
let key = convolved.view(2048 * 2, u64::from(count - 2048) * 2)?;
|
|
temporary.extend(gdn_l2_norm(&key, &k, rows, false)?);
|
|
let value = convolved.view(4096 * 2, u64::from(count - 4096) * 2)?;
|
|
if rows == 1 {
|
|
v = value;
|
|
} else {
|
|
v = Buffer::mtplx_bytes(u64::from(rows) * 6144 * 2)?;
|
|
let grid = [6144, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, &value).input((rows - 1) as usize * 10240 + 6144),
|
|
tensor(1, &v).output(rows as usize * 6144),
|
|
bytes(3, &[10240_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
temporary.extend([window, convolved]);
|
|
}
|
|
Ok(ConvNormOutput {
|
|
q,
|
|
k,
|
|
v,
|
|
conv,
|
|
qkv,
|
|
used_fused,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct DeltaUpdateOutput {
|
|
pub(super) hidden: Buffer,
|
|
pub(super) delta: Buffer,
|
|
pub(super) g: Buffer,
|
|
pub(super) beta: Buffer,
|
|
_initial_state: Option<Buffer>,
|
|
}
|
|
|
|
/// Actual gated_delta_update inference branch, including compute_g, sigmoid,
|
|
/// optional zero initial state and mask. Q/K/V are contiguous as required by
|
|
/// the reference custom-kernel dispatcher; A/B retain projection row strides.
|
|
#[cfg(test)]
|
|
pub(super) fn gated_delta_update(
|
|
[q, k, v]: [&Buffer; 3],
|
|
[a, b, a_log, dt_bias]: [&Buffer; 4],
|
|
[rows, gate_stride]: [u32; 2],
|
|
state: Option<&Buffer>,
|
|
mask: Option<&Buffer>,
|
|
) -> Result<DeltaUpdateOutput, String> {
|
|
let count = u64::from(rows) * 48;
|
|
let state_count = 48 * 128 * 128_u32;
|
|
if rows == 0 || gate_stride < 48 || u64::from(rows) * 6144 > i32::MAX as u64
|
|
|| [q, k].iter().any(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(rows) * 2048 * 2)
|
|
|| unsafe { ds4_gpu_tensor_bytes(v.raw()) } < u64::from(rows) * 6144 * 2
|
|
|| [a, b].iter().any(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) }
|
|
< (u64::from(rows - 1) * u64::from(gate_stride) + 48) * 2)
|
|
|| [a_log, dt_bias].iter().any(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < 48 * 2)
|
|
|| state.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(state_count) * 4)
|
|
|| mask.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(rows))
|
|
{ return Err("invalid MTPLX gated delta inputs".into()); }
|
|
let gate_span = (rows - 1) as usize * gate_stride as usize + 48;
|
|
let beta = Buffer::mtplx_bytes(count * 2)?;
|
|
let g = Buffer::mtplx_bytes(count * 4)?;
|
|
if rows == 1 || gate_stride == 48 {
|
|
let count = count as u32;
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"v_Sigmoidbfloat16bfloat16"
|
|
} else {
|
|
"vn_Sigmoidbfloat16bfloat16"
|
|
},
|
|
&[
|
|
tensor(0, b).input(gate_span),
|
|
tensor(1, &beta).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let grid = [48, rows, 1];
|
|
dispatch_geometry(
|
|
"gn1_Sigmoidbfloat16bfloat16",
|
|
&[
|
|
tensor(0, b).input(gate_span),
|
|
tensor(1, &beta).output(count as usize),
|
|
bytes(2, &[rows, 48]),
|
|
bytes(3, &[u64::from(gate_stride), 1]),
|
|
bytes(4, &2_u32),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
if rows == 1 {
|
|
dispatch_geometry(
|
|
"kernel_qwen_mtplx_compute_g_bf16_contiguous",
|
|
&[
|
|
tensor(0, a_log).input(48),
|
|
tensor(1, a).input(gate_span),
|
|
tensor(2, dt_bias).input(48),
|
|
tensor(3, &g).output(count as usize),
|
|
bytes(4, &48_u32),
|
|
],
|
|
&[],
|
|
[48, 1, 1],
|
|
[48, 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let grid = [48, rows, 1];
|
|
dispatch_geometry(
|
|
"kernel_qwen_mtplx_compute_g_bf16_strided_2",
|
|
&[
|
|
tensor(0, a_log).input(48),
|
|
tensor(1, a).input(gate_span),
|
|
tensor(2, dt_bias).input(48),
|
|
bytes(3, &[0_u64, 1, u64::from(gate_stride), 1, 0, 1]),
|
|
tensor(4, &g).output(count as usize),
|
|
bytes(5, &[rows, 48]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let hidden = Buffer::mtplx_bytes(u64::from(rows) * 6144 * 2)?;
|
|
let delta = Buffer::mtplx_bytes(u64::from(state_count) * 4)?;
|
|
let initial_state = if state.is_none() {
|
|
let zero = Buffer::mtplx_bytes(u64::from(state_count) * 4)?;
|
|
let scalar = scalar_buffer(&0_f32.to_le_bytes())?;
|
|
dispatch_geometry(
|
|
"g1_copyfloat32float32",
|
|
&[
|
|
tensor(0, &scalar).input(1),
|
|
tensor(1, &zero).output(state_count as usize),
|
|
bytes(3, &0_u64),
|
|
],
|
|
&[],
|
|
[state_count, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)?;
|
|
Some(zero)
|
|
} else {
|
|
None
|
|
};
|
|
let state = state.or(initial_state.as_ref()).unwrap();
|
|
let length = scalar_buffer(&(rows as i32).to_le_bytes())?;
|
|
let mut bindings = vec![
|
|
tensor(0, q).input(rows as usize * 2048),
|
|
tensor(1, k).input(rows as usize * 2048),
|
|
tensor(2, v).input(rows as usize * 6144),
|
|
tensor(3, &g).input(count as usize),
|
|
tensor(4, &beta).input(count as usize),
|
|
tensor(5, state).input(state_count as usize),
|
|
tensor(6, &length).input(1),
|
|
];
|
|
if let Some(mask) = mask {
|
|
bindings.push(tensor(7, mask).input(rows as usize));
|
|
}
|
|
let index = bindings.len() as u32;
|
|
bindings.extend([
|
|
tensor(index, &hidden).output(rows as usize * 6144),
|
|
tensor(index + 1, &delta).output(state_count as usize),
|
|
]);
|
|
dispatch(
|
|
if mask.is_some() {
|
|
"kernel_qwen_mtplx_gated_delta_masked_bf16"
|
|
} else {
|
|
"kernel_qwen_mtplx_gated_delta_bf16"
|
|
},
|
|
&bindings,
|
|
[1, 32, 48],
|
|
[32, 4, 1],
|
|
)?;
|
|
if let Some(initial_state) = &initial_state {
|
|
encoder::add_temporaries(std::slice::from_ref(initial_state));
|
|
}
|
|
Ok(DeltaUpdateOutput {
|
|
hidden,
|
|
delta,
|
|
g,
|
|
beta,
|
|
_initial_state: initial_state,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_copy_span(shape: [u32; 3], strides: [u64; 3]) -> Result<usize, String> {
|
|
if shape.contains(&0) {
|
|
return Ok(0);
|
|
}
|
|
shape
|
|
.into_iter()
|
|
.zip(strides)
|
|
.try_fold(1_usize, |span, (dim, stride)| {
|
|
usize::try_from(stride)
|
|
.ok()
|
|
.and_then(|stride| (dim as usize - 1).checked_mul(stride))
|
|
.and_then(|n| span.checked_add(n))
|
|
.ok_or_else(|| "QSA copy element span overflow".into())
|
|
})
|
|
}
|
|
|
|
/// GeneralGeneral copy used by the reference's slice updates/concatenations.
|
|
/// Collapse only dimensions contiguous in BOTH source and destination.
|
|
#[cfg(test)]
|
|
fn qsa_copy(
|
|
source: Option<&Buffer>,
|
|
target: &Buffer,
|
|
shape: [u32; 3],
|
|
source_strides: [u64; 3],
|
|
target_strides: [u64; 3],
|
|
width: u32,
|
|
output_elements: usize,
|
|
) -> Result<(), String> {
|
|
qsa_copy_typed(
|
|
source,
|
|
target,
|
|
shape,
|
|
source_strides,
|
|
target_strides,
|
|
if width == 2 { "bfloat16" } else { "float32" },
|
|
output_elements,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_copy_typed(
|
|
source: Option<&Buffer>,
|
|
target: &Buffer,
|
|
shape: [u32; 3],
|
|
source_strides: [u64; 3],
|
|
target_strides: [u64; 3],
|
|
dtype: &str,
|
|
output_elements: usize,
|
|
) -> Result<(), String> {
|
|
if shape.contains(&0) {
|
|
return Ok(());
|
|
}
|
|
let itemsize = match dtype {
|
|
"bool_" => 1,
|
|
"bfloat16" => 2,
|
|
"float32" | "int32" | "uint32" => 4,
|
|
_ => return Err("unsupported QSA copy dtype".into()),
|
|
};
|
|
let zero = source
|
|
.is_none()
|
|
.then(|| scalar_buffer(&[0; 4][..itemsize]))
|
|
.transpose()?;
|
|
qsa_copy_offsets(
|
|
[
|
|
tensor(0, source.or(zero.as_ref()).unwrap()).input(if source.is_some() {
|
|
qsa_copy_span(shape, source_strides)?
|
|
} else {
|
|
1
|
|
}),
|
|
tensor(1, target).output(output_elements),
|
|
],
|
|
shape,
|
|
source_strides,
|
|
target_strides,
|
|
dtype,
|
|
[None; 2],
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_copy_offsets(
|
|
buffers: [Binding<'_>; 2],
|
|
shape: [u32; 3],
|
|
source_strides: [u64; 3],
|
|
target_strides: [u64; 3],
|
|
dtype: &str,
|
|
offsets: [Option<&Buffer>; 2],
|
|
) -> Result<(), String> {
|
|
if shape.contains(&0) {
|
|
return Ok(());
|
|
}
|
|
let mut dims: Vec<(u32, u64, u64)> = Vec::new();
|
|
for ((n, s), d) in shape.into_iter().zip(source_strides).zip(target_strides) {
|
|
if n == 1 {
|
|
continue;
|
|
}
|
|
if let Some(last) = dims.last_mut()
|
|
&& last.1 == u64::from(n) * s
|
|
&& last.2 == u64::from(n) * d
|
|
{
|
|
last.0 = last.0.checked_mul(n).ok_or("QSA copy shape overflow")?;
|
|
last.1 = s;
|
|
last.2 = d;
|
|
} else {
|
|
dims.push((n, s, d));
|
|
}
|
|
}
|
|
if dims.is_empty() {
|
|
dims.push((1, 0, 0));
|
|
}
|
|
let mut ss = [0_u64; 3];
|
|
let mut ds = [0_u64; 3];
|
|
let mut grid = [1_u32; 3];
|
|
for (i, (n, s, d)) in dims.iter().enumerate() {
|
|
ss[i] = *s;
|
|
ds[i] = *d;
|
|
grid[dims.len() - 1 - i] = *n;
|
|
}
|
|
let dynamic = offsets.iter().any(Option::is_some);
|
|
let zero = 0_i64;
|
|
let [source, target] = buffers;
|
|
let bindings = [
|
|
source,
|
|
target,
|
|
bytes(3, &ss),
|
|
bytes(4, &ds),
|
|
offsets[0].map_or_else(|| bytes(6, &zero), |b| tensor(6, b).input(1)),
|
|
offsets[1].map_or_else(|| bytes(7, &zero), |b| tensor(7, b).input(1)),
|
|
];
|
|
dispatch_geometry(
|
|
&format!(
|
|
"gg{}{}_copy{dtype}{dtype}",
|
|
dims.len(),
|
|
if dynamic { "_dynamic" } else { "" }
|
|
),
|
|
&bindings[..if dynamic { 6 } else { 4 }],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)
|
|
}
|
|
|
|
/// Original DynamicSlice/Update row-copy stage for B1 BF16 indexer leaves.
|
|
/// Start leaves are int32[1] graph values with validated in-bounds semantics;
|
|
/// they stay on-device, just as in the reference's private GPU primitives.
|
|
/// Returns offset temporaries so callers retain the entire asynchronous chain.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_dynamic_copy_rows(
|
|
[source, target]: [&Buffer; 2],
|
|
[source_stride, target_stride]: [u64; 2],
|
|
[rows, width]: [u32; 2],
|
|
starts: [Option<&Buffer>; 2],
|
|
data_sizes: [usize; 2],
|
|
) -> Result<Vec<Buffer>, String> {
|
|
if rows == 0
|
|
|| width == 0
|
|
|| rows.checked_mul(width).is_none_or(|n| n > i32::MAX as u32)
|
|
|| source_stride < u64::from(width)
|
|
|| target_stride < u64::from(width)
|
|
|| starts
|
|
.iter()
|
|
.flatten()
|
|
.any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
|
|
|| [source, target]
|
|
.into_iter()
|
|
.zip([source_stride, target_stride])
|
|
.any(|(b, stride)| {
|
|
u64::from(rows - 1)
|
|
.checked_mul(stride)
|
|
.and_then(|n| n.checked_add(u64::from(width)))
|
|
.and_then(|n| n.checked_mul(2))
|
|
.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n)
|
|
})
|
|
{
|
|
return Err("invalid QSA dynamic copy buffers or geometry".into());
|
|
}
|
|
for ((buffer, size), stride) in [source, target]
|
|
.into_iter()
|
|
.zip(data_sizes)
|
|
.zip([source_stride, target_stride])
|
|
{
|
|
if size < qsa_copy_span([1, rows, width], [0, stride, 1])?
|
|
|| size
|
|
.checked_mul(2)
|
|
.is_none_or(|n| n as u64 > unsafe { ds4_gpu_tensor_bytes(buffer.raw()) })
|
|
{
|
|
return Err("invalid QSA dynamic copy array data_size".into());
|
|
}
|
|
}
|
|
let mut offsets = [None, None];
|
|
for (i, (start, stride)) in starts
|
|
.into_iter()
|
|
.zip([source_stride, target_stride])
|
|
.enumerate()
|
|
{
|
|
if let Some(start) = start {
|
|
let indices = array::Array::new(&[1], array::Dtype::I32, start.view(0, 4)?)?;
|
|
let stride = i64::try_from(stride).map_err(|_| "QSA dynamic copy stride overflow")?;
|
|
let offset = views::dynamic_offset(&indices, &[0, stride, 1], &[1])?;
|
|
offsets[i] = Some(offset.buffer().view(0, 8)?);
|
|
}
|
|
}
|
|
qsa_copy_offsets(
|
|
[
|
|
tensor(0, source).input(data_sizes[0]),
|
|
tensor(1, target).output(data_sizes[1]),
|
|
],
|
|
[1, rows, width],
|
|
[0, source_stride, 1],
|
|
[0, target_stride, 1],
|
|
"bfloat16",
|
|
[offsets[0].as_ref(), offsets[1].as_ref()],
|
|
)?;
|
|
Ok(offsets.into_iter().flatten().collect())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_zero(count: u32, width: u32) -> Result<Buffer, String> {
|
|
let output = Buffer::mtplx_bytes(u64::from(count) * u64::from(width))?;
|
|
let zero = scalar_buffer(&[0; 4][..width as usize])?;
|
|
let work = if count < 65536 { 1 } else { 8 / width };
|
|
let threads = count.div_ceil(work);
|
|
let dtype = if width == 2 { "bfloat16" } else { "float32" };
|
|
dispatch_geometry(
|
|
&format!("s{}_copy{dtype}{dtype}", if work == 1 { "" } else { "n" }),
|
|
&[
|
|
tensor(0, &zero).input(1),
|
|
tensor(1, &output).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
Ok(output)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_grown_capacity(end: u32, current: u32) -> Result<u32, String> {
|
|
end.checked_add(255)
|
|
.map(|n| n / 256 * 256)
|
|
.zip(current.checked_mul(2))
|
|
.map(|(a, b)| a.max(b))
|
|
.filter(|cap| *cap <= i32::MAX as u32 / 512)
|
|
.ok_or_else(|| "QSA capacity overflow".into())
|
|
}
|
|
|
|
/// Four logical state leaves, sharing immutable backing until a later write.
|
|
/// Python array identity and its separately shared storage are different things:
|
|
/// __setitem__ overwrites the array descriptor, whereas slicing creates an array.
|
|
#[cfg(test)]
|
|
pub(super) struct QsaArray {
|
|
storage: array::Array,
|
|
pub(super) head_stride: Cell<u64>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl QsaArray {
|
|
pub(super) fn new(buffer: Buffer, head_stride: u64) -> Result<Rc<Self>, String> {
|
|
let elements = i32::try_from(unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } / 2)
|
|
.map_err(|_| "QSA backing dimension overflow")?;
|
|
Ok(Rc::new(Self {
|
|
storage: array::Array::new(&[elements], array::Dtype::BF16, buffer)?,
|
|
head_stride: Cell::new(head_stride),
|
|
}))
|
|
}
|
|
|
|
pub(super) fn buffer(&self) -> Ref<'_, Buffer> {
|
|
self.storage.buffer()
|
|
}
|
|
|
|
fn slice(&self) -> Rc<Self> {
|
|
Rc::new(Self {
|
|
storage: self.storage.share_data(),
|
|
head_stride: Cell::new(self.head_stride.get()),
|
|
})
|
|
}
|
|
|
|
fn shared_storage(&self) -> bool {
|
|
self.storage.shared_storage()
|
|
}
|
|
|
|
// Installed BF16 vector-copy contract: descriptor and Data must both be
|
|
// exclusive, and at most 16 KiB of the root allocation may be unused.
|
|
fn can_donate_bf16(array: &Rc<Self>, input: &Buffer, output_bytes: u64) -> bool {
|
|
let storage = array.buffer();
|
|
Rc::strong_count(array) == 1
|
|
&& array.storage.is_donatable()
|
|
&& storage.identity() == input.identity()
|
|
&& input.unique_data()
|
|
&& unsafe { ds4_gpu_tensor_bytes(input.raw()) } >= output_bytes
|
|
&& input.allocation_bytes().saturating_sub(output_bytes) <= 16384
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
#[cfg(test)]
|
|
pub(super) struct QsaSnapshot {
|
|
pub(super) leaves: [Option<Rc<QsaArray>>; 4], // keys, values, raw, pooled
|
|
pub(super) offset: u32,
|
|
pub(super) pooled_len: u32,
|
|
}
|
|
|
|
/// Installed single-sequence QSACache + KVCache geometry (KV2x256, index128).
|
|
/// Array aliases follow MTPLX's state setter; distinct views use storage COW.
|
|
#[derive(Default)]
|
|
#[cfg(test)]
|
|
pub(super) struct QsaCache {
|
|
pub(super) keys: Option<Rc<QsaArray>>,
|
|
pub(super) values: Option<Rc<QsaArray>>,
|
|
pub(super) raw: Option<Rc<QsaArray>>,
|
|
pub(super) pooled: Option<Rc<QsaArray>>,
|
|
pub(super) mirror: Option<Rc<Buffer>>,
|
|
pub(super) offset: u32,
|
|
pub(super) kv_cap: u32,
|
|
pub(super) kv_head_stride: u64,
|
|
pub(super) raw_cap: u32,
|
|
pub(super) pooled_cap: u32,
|
|
pub(super) pooled_len: u32,
|
|
pub(super) mirror_cap: u32,
|
|
pub(super) mirror_strides: [u64; 2], // D, blocks; initial transpose is a VIEW
|
|
reserved_raw: u32,
|
|
reserved_pooled: u32,
|
|
}
|
|
|
|
/// Original precompute_qsa_replay_capacity for the installed ratio4/step256.
|
|
/// The fixed ceil(window/4) staging window can exceed the logical frontier.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_replay_capacity(
|
|
start: u32,
|
|
window: u32,
|
|
existing: [u32; 2],
|
|
) -> Result<[u32; 2], String> {
|
|
let end = start
|
|
.checked_add(window)
|
|
.ok_or("QSA replay frontier overflow")?;
|
|
let staging = window.div_ceil(4);
|
|
let staging_tokens = staging.checked_mul(4).ok_or("QSA replay window overflow")?;
|
|
let [raw, pooled] = [
|
|
existing[0].max(end).max(staging_tokens),
|
|
existing[1].max(end / 4).max(staging),
|
|
]
|
|
.map(|required| {
|
|
if required == 0 {
|
|
return Ok(0);
|
|
}
|
|
required
|
|
.max(256)
|
|
.checked_next_power_of_two()
|
|
.filter(|n| *n <= i32::MAX as u32 / 512)
|
|
.ok_or_else(|| "QSA replay capacity overflow".to_owned())
|
|
});
|
|
Ok([raw?, pooled?])
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl QsaCache {
|
|
fn grow_index(
|
|
buffer: &mut Option<Rc<QsaArray>>,
|
|
cap: &mut u32,
|
|
target: u32,
|
|
) -> Result<(), String> {
|
|
if target == 0 || target > i32::MAX as u32 / 512 {
|
|
return Err("invalid QSA capacity".into());
|
|
}
|
|
let grown = qsa_zero(target * 128, 2)?;
|
|
if let Some(old) = buffer.as_ref() {
|
|
qsa_copy(
|
|
Some(&old.buffer()),
|
|
&grown,
|
|
[1, *cap, 128],
|
|
[0, 128, 1],
|
|
[0, 128, 1],
|
|
2,
|
|
target as usize * 128,
|
|
)?;
|
|
}
|
|
*buffer = Some(QsaArray::new(grown, 0)?);
|
|
*cap = target;
|
|
Ok(())
|
|
}
|
|
|
|
fn index_cow(buffer: &Option<Rc<QsaArray>>, cap: u32) -> Result<(), String> {
|
|
if let Some(old) = buffer.as_ref()
|
|
&& old.shared_storage()
|
|
{
|
|
let copy = Buffer::mtplx_bytes(u64::from(cap) * 128 * 2)?;
|
|
let count = cap * 128;
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"v_copybfloat16bfloat16"
|
|
} else {
|
|
"vn_copybfloat16bfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &old.buffer()).input(count as usize),
|
|
tensor(1, ©).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
old.storage.replace(copy);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn reserve(&mut self, raw: u32, pooled: u32) -> Result<(), String> {
|
|
let raw = raw.max(self.raw_cap).max(self.reserved_raw);
|
|
let pooled = pooled.max(self.pooled_cap).max(self.reserved_pooled);
|
|
if raw < self.offset
|
|
|| pooled < self.pooled_len
|
|
|| raw > i32::MAX as u32 / 512
|
|
|| pooled > i32::MAX as u32 / 512
|
|
{
|
|
return Err("invalid QSA reservation".into());
|
|
}
|
|
self.reserved_raw = raw;
|
|
self.reserved_pooled = pooled;
|
|
if self.raw.is_some() && raw > self.raw_cap {
|
|
Self::grow_index(&mut self.raw, &mut self.raw_cap, raw)?;
|
|
}
|
|
if self.pooled.is_some() && pooled > self.pooled_cap {
|
|
Self::grow_index(&mut self.pooled, &mut self.pooled_cap, pooled)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// _ensure_compiled_backings includes pending reservations and materializes
|
|
/// pristine leaves. It neither advances KV.offset nor the pooled frontier.
|
|
pub(super) fn ensure_compiled_backings(&mut self, rows: u32) -> Result<(), String> {
|
|
if rows == 0 {
|
|
return Err("compiled QSA requires positive rows".into());
|
|
}
|
|
let [raw, pooled] = qsa_replay_capacity(
|
|
self.offset,
|
|
rows,
|
|
[
|
|
self.raw_cap.max(self.reserved_raw),
|
|
self.pooled_cap.max(self.reserved_pooled),
|
|
],
|
|
)?;
|
|
self.reserve(raw, pooled.max(256))?;
|
|
if self.raw.is_none() {
|
|
Self::grow_index(&mut self.raw, &mut self.raw_cap, self.reserved_raw)?;
|
|
}
|
|
if self.pooled.is_none() {
|
|
Self::grow_index(&mut self.pooled, &mut self.pooled_cap, self.reserved_pooled)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn write_raw(
|
|
&mut self,
|
|
keys: &Buffer,
|
|
rows: u32,
|
|
stride: u64,
|
|
) -> Result<(), String> {
|
|
let end = self
|
|
.offset
|
|
.checked_add(rows)
|
|
.filter(|n| *n <= i32::MAX as u32 / 512)
|
|
.ok_or("QSA raw range overflow")?;
|
|
let bytes_needed = u64::from(rows.saturating_sub(1))
|
|
.checked_mul(stride)
|
|
.and_then(|n| n.checked_add(128))
|
|
.and_then(|n| n.checked_mul(2));
|
|
if rows == 0
|
|
|| stride < 128
|
|
|| bytes_needed.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(keys.raw()) } < n)
|
|
{
|
|
return Err("invalid QSA raw input".into());
|
|
}
|
|
if self.raw.is_none() || end > self.raw_cap {
|
|
let target = qsa_grown_capacity(end, self.raw_cap)?.max(self.reserved_raw);
|
|
Self::grow_index(&mut self.raw, &mut self.raw_cap, target)?;
|
|
} else {
|
|
Self::index_cow(&self.raw, self.raw_cap)?;
|
|
}
|
|
let target = self
|
|
.raw
|
|
.as_ref()
|
|
.unwrap()
|
|
.buffer()
|
|
.view(u64::from(self.offset) * 128 * 2, u64::from(rows) * 128 * 2)?;
|
|
qsa_copy(
|
|
Some(keys),
|
|
&target,
|
|
[1, rows, 128],
|
|
[0, stride, 1],
|
|
[0, 128, 1],
|
|
2,
|
|
self.raw_cap as usize * 128,
|
|
)
|
|
}
|
|
|
|
pub(super) fn write_pooled(
|
|
&mut self,
|
|
blocks: &Buffer,
|
|
start: u32,
|
|
total: u32,
|
|
) -> Result<(), String> {
|
|
if total <= start
|
|
|| total > i32::MAX as u32 / 512
|
|
|| unsafe { ds4_gpu_tensor_bytes(blocks.raw()) } < u64::from(total - start) * 128 * 2
|
|
{
|
|
return Err("invalid QSA pooled input".into());
|
|
}
|
|
if self.pooled.is_none() || total > self.pooled_cap {
|
|
let target = qsa_grown_capacity(total, self.pooled_cap)?.max(self.reserved_pooled);
|
|
Self::grow_index(&mut self.pooled, &mut self.pooled_cap, target)?;
|
|
} else {
|
|
Self::index_cow(&self.pooled, self.pooled_cap)?;
|
|
}
|
|
let target = self.pooled.as_ref().unwrap().buffer().view(
|
|
u64::from(start) * 128 * 2,
|
|
u64::from(total - start) * 128 * 2,
|
|
)?;
|
|
qsa_copy(
|
|
Some(blocks),
|
|
&target,
|
|
[1, total - start, 128],
|
|
[0, 128, 1],
|
|
[0, 128, 1],
|
|
2,
|
|
self.pooled_cap as usize * 128,
|
|
)?;
|
|
if self.mirror.is_none() {
|
|
self.rebuild_mirror()?;
|
|
} else {
|
|
if self.mirror_cap < self.pooled_cap {
|
|
let grown = qsa_zero(self.pooled_cap * 128, 4)?;
|
|
qsa_copy(
|
|
self.mirror.as_deref(),
|
|
&grown,
|
|
[1, 128, self.mirror_cap],
|
|
[0, self.mirror_strides[0], self.mirror_strides[1]],
|
|
[0, u64::from(self.pooled_cap), 1],
|
|
4,
|
|
self.pooled_cap as usize * 128,
|
|
)?;
|
|
self.mirror = Some(Rc::new(grown));
|
|
self.mirror_cap = self.pooled_cap;
|
|
self.mirror_strides = [u64::from(self.pooled_cap), 1];
|
|
}
|
|
let f32blocks = Buffer::mtplx_bytes(u64::from(total - start) * 128 * 4)?;
|
|
cast_bf16_f32(blocks, &f32blocks, (total - start) * 128, true)?;
|
|
let at = u64::from(start) * self.mirror_strides[1] * 4;
|
|
let span = (127 * self.mirror_strides[0]
|
|
+ u64::from(total - start - 1) * self.mirror_strides[1]
|
|
+ 1)
|
|
* 4;
|
|
let target = self.mirror.as_ref().unwrap().view(at, span)?;
|
|
qsa_copy(
|
|
Some(&f32blocks),
|
|
&target,
|
|
[1, 128, total - start],
|
|
[0, 1, 128],
|
|
[0, self.mirror_strides[0], self.mirror_strides[1]],
|
|
4,
|
|
self.mirror_cap as usize * 128,
|
|
)?;
|
|
}
|
|
self.pooled_len = total;
|
|
Ok(())
|
|
}
|
|
|
|
fn rebuild_mirror(&mut self) -> Result<(), String> {
|
|
let pooled = self
|
|
.pooled
|
|
.as_ref()
|
|
.ok_or("QSA mirror needs pooled state")?;
|
|
let mirror = Buffer::mtplx_bytes(u64::from(self.pooled_cap) * 128 * 4)?;
|
|
cast_bf16_f32(&pooled.buffer(), &mirror, self.pooled_cap * 128, true)?;
|
|
self.mirror = Some(Rc::new(mirror));
|
|
self.mirror_cap = self.pooled_cap;
|
|
self.mirror_strides = [1, 128];
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn pooled_f32_view(&mut self, nb: u32) -> Result<(&Buffer, [u64; 2]), String> {
|
|
if nb == 0 || nb > self.pooled_cap {
|
|
return Err("invalid QSA mirror frontier".into());
|
|
}
|
|
if self.mirror.is_none() || self.mirror_cap < nb {
|
|
self.rebuild_mirror()?;
|
|
}
|
|
Ok((self.mirror.as_deref().unwrap(), self.mirror_strides))
|
|
}
|
|
|
|
pub(super) fn extend_pooled(
|
|
&mut self,
|
|
total: u32,
|
|
norm: &Buffer,
|
|
freq: &Buffer,
|
|
eager: bool,
|
|
) -> Result<Vec<Buffer>, String> {
|
|
let nb = total / 4;
|
|
let old = self.pooled_len.min(nb);
|
|
if total > self.raw_cap {
|
|
return Err("QSA pool exceeds raw backing".into());
|
|
}
|
|
if nb > old {
|
|
let fresh = self
|
|
.raw
|
|
.as_ref()
|
|
.ok_or("QSA pool needs raw keys")?
|
|
.buffer()
|
|
.view(
|
|
u64::from(old) * 4 * 128 * 2,
|
|
u64::from(nb - old) * 4 * 128 * 2,
|
|
)?;
|
|
let (blocks, mut temporary) = if eager {
|
|
qsa_pool_keys_eager([&fresh, norm, freq], nb - old, [128, 1], old as i32)?
|
|
} else {
|
|
(
|
|
qsa_pool_keys([&fresh, norm, freq], nb - old, [128, 1], old as i32)?,
|
|
Vec::new(),
|
|
)
|
|
};
|
|
self.write_pooled(&blocks, old, nb)?;
|
|
temporary.push(blocks);
|
|
return Ok(temporary);
|
|
}
|
|
Ok(Vec::new())
|
|
}
|
|
|
|
pub(super) fn update_kv(
|
|
&mut self,
|
|
keys: &Buffer,
|
|
values: &Buffer,
|
|
rows: u32,
|
|
strides: [u64; 2],
|
|
) -> Result<(), String> {
|
|
let end = self
|
|
.offset
|
|
.checked_add(rows)
|
|
.filter(|n| *n <= i32::MAX as u32 / 512)
|
|
.ok_or("QSA KV range overflow")?;
|
|
let bytes_needed = u64::from(rows.saturating_sub(1))
|
|
.checked_mul(strides[1])
|
|
.and_then(|n| n.checked_add(strides[0]))
|
|
.and_then(|n| n.checked_add(256))
|
|
.and_then(|n| n.checked_mul(2));
|
|
if rows == 0
|
|
|| bytes_needed.is_none_or(|n| {
|
|
[keys, values]
|
|
.iter()
|
|
.any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n)
|
|
})
|
|
{
|
|
return Err("invalid QSA KV input".into());
|
|
}
|
|
if self.keys.is_none() || end > self.kv_cap {
|
|
let keep = if !self.offset.is_multiple_of(256) {
|
|
self.offset
|
|
} else {
|
|
self.kv_cap
|
|
};
|
|
let extra = rows.div_ceil(256) * 256;
|
|
let cap = keep
|
|
.checked_add(extra)
|
|
.filter(|n| *n <= i32::MAX as u32 / 512)
|
|
.ok_or("QSA KV allocation overflow")?;
|
|
for buffer in [&mut self.keys, &mut self.values] {
|
|
let grown = if buffer.is_none() {
|
|
qsa_zero(cap * 512, 2)?
|
|
} else {
|
|
Buffer::mtplx_bytes(u64::from(cap) * 512 * 2)?
|
|
};
|
|
let concurrent = encoder::concurrent_scope();
|
|
if let Some(old) = buffer.as_ref() {
|
|
qsa_copy(
|
|
Some(&old.buffer()),
|
|
&grown,
|
|
[2, keep, 256],
|
|
[self.kv_head_stride, 256, 1],
|
|
[u64::from(cap) * 256, 256, 1],
|
|
2,
|
|
keep as usize * 512,
|
|
)?;
|
|
}
|
|
if buffer.is_some() {
|
|
let tail = grown.view(
|
|
u64::from(keep) * 256 * 2,
|
|
(u64::from(cap) * 256 + u64::from(extra) * 256) * 2,
|
|
)?;
|
|
qsa_copy(
|
|
None,
|
|
&tail,
|
|
[2, extra, 256],
|
|
[0; 3],
|
|
[u64::from(cap) * 256, 256, 1],
|
|
2,
|
|
extra as usize * 512,
|
|
)?;
|
|
}
|
|
drop(concurrent);
|
|
*buffer = Some(QsaArray::new(grown, u64::from(cap) * 256)?);
|
|
}
|
|
self.kv_cap = cap;
|
|
self.kv_head_stride = u64::from(cap) * 256;
|
|
} else {
|
|
// Distinct slices can prevent donation; aliases of the same array
|
|
// object observe its descriptor replacement instead.
|
|
let need_copy = [&self.keys, &self.values]
|
|
.iter()
|
|
.any(|b| b.as_ref().unwrap().shared_storage())
|
|
|| self.kv_head_stride != u64::from(self.kv_cap) * 256;
|
|
if need_copy {
|
|
for buffer in [&mut self.keys, &mut self.values] {
|
|
let copy = Buffer::mtplx_bytes(u64::from(self.kv_cap) * 512 * 2)?;
|
|
if self.kv_head_stride == u64::from(self.kv_cap) * 256 {
|
|
let count = self.kv_cap * 512;
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"v_copybfloat16bfloat16"
|
|
} else {
|
|
"vn_copybfloat16bfloat16"
|
|
},
|
|
&[
|
|
tensor(0, &buffer.as_ref().unwrap().buffer()).input(
|
|
(self.kv_head_stride + u64::from(self.kv_cap) * 256) as usize,
|
|
),
|
|
tensor(1, ©).output(self.kv_cap as usize * 512),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let grid = [self.kv_cap * 256, 2, 1];
|
|
dispatch_geometry(
|
|
"g2_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, &buffer.as_ref().unwrap().buffer()).input(
|
|
(self.kv_head_stride + u64::from(self.kv_cap) * 256) as usize,
|
|
),
|
|
tensor(1, ©).output(self.kv_cap as usize * 512),
|
|
bytes(3, &[self.kv_head_stride, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let array = buffer.as_ref().unwrap();
|
|
array.storage.replace(copy);
|
|
array.head_stride.set(u64::from(self.kv_cap) * 256);
|
|
}
|
|
self.kv_head_stride = u64::from(self.kv_cap) * 256;
|
|
}
|
|
}
|
|
for (source, target) in [keys, values].into_iter().zip([&self.keys, &self.values]) {
|
|
let span = (self.kv_head_stride + u64::from(rows) * 256) * 2;
|
|
let target = target
|
|
.as_ref()
|
|
.unwrap()
|
|
.buffer()
|
|
.view(u64::from(self.offset) * 256 * 2, span)?;
|
|
qsa_copy(
|
|
Some(source),
|
|
&target,
|
|
[2, rows, 256],
|
|
[strides[0], strides[1], 1],
|
|
[self.kv_head_stride, 256, 1],
|
|
2,
|
|
self.kv_cap as usize * 512,
|
|
)?;
|
|
}
|
|
self.offset = end;
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn trim(&mut self, n: u32) -> u32 {
|
|
let trimmed = n.min(self.offset);
|
|
self.offset -= trimmed;
|
|
self.pooled_len = self.pooled_len.min(self.offset / 4);
|
|
trimmed
|
|
}
|
|
|
|
pub(super) fn nbytes(&self) -> u64 {
|
|
// QSACache reports logical backing capacities, excluding its FP32 mirror.
|
|
u64::from(if self.keys.is_some() { self.kv_cap } else { 0 }) * 2048
|
|
+ u64::from(self.raw_cap) * 256
|
|
+ u64::from(self.pooled_cap) * 256
|
|
}
|
|
|
|
pub(super) fn snapshot(&self) -> QsaSnapshot {
|
|
let nb = self.pooled_len.min(self.offset / 4);
|
|
let kv = |array: &Option<Rc<QsaArray>>| {
|
|
array.as_ref().map(|a| {
|
|
if self.offset == self.kv_cap {
|
|
a.clone()
|
|
} else {
|
|
a.slice()
|
|
}
|
|
})
|
|
};
|
|
QsaSnapshot {
|
|
leaves: [
|
|
kv(&self.keys),
|
|
kv(&self.values),
|
|
self.raw.as_ref().map(|a| a.slice()),
|
|
if nb == 0 {
|
|
None
|
|
} else {
|
|
self.pooled.as_ref().map(|a| a.slice())
|
|
},
|
|
],
|
|
offset: self.offset,
|
|
pooled_len: nb,
|
|
}
|
|
}
|
|
|
|
pub(super) fn restore(&mut self, state: &QsaSnapshot) -> Result<(), String> {
|
|
let head_stride = state.leaves[0].as_ref().map_or(0, |a| a.head_stride.get());
|
|
let kv_span = head_stride
|
|
.checked_add(u64::from(state.offset) * 256)
|
|
.and_then(|n| n.checked_mul(2));
|
|
if state.offset>i32::MAX as u32/512 || state.pooled_len>state.offset/4
|
|
|| state.leaves[0].is_some()!=state.leaves[1].is_some()
|
|
|| (state.offset>0 && state.leaves[0].is_none())
|
|
|| (state.pooled_len>0 && state.leaves[3].is_none())
|
|
|| state.leaves[..2].iter().flatten().any(|b|kv_span.is_none_or(|n|unsafe{ds4_gpu_tensor_bytes(b.buffer().raw())}<n))
|
|
|| state.leaves[2].as_ref().is_some_and(|b|unsafe{ds4_gpu_tensor_bytes(b.buffer().raw())}<u64::from(state.offset)*128*2)
|
|
|| state.leaves[3].as_ref().is_some_and(|b|unsafe{ds4_gpu_tensor_bytes(b.buffer().raw())}<u64::from(state.pooled_len)*128*2) {
|
|
return Err("invalid QSA snapshot".into());
|
|
}
|
|
[self.keys, self.values, self.raw, self.pooled] = state.leaves.clone();
|
|
self.offset = state.offset;
|
|
self.kv_cap = state.offset;
|
|
self.kv_head_stride = head_stride;
|
|
self.raw_cap = if self.raw.is_some() { state.offset } else { 0 };
|
|
self.pooled_cap = state.pooled_len;
|
|
self.pooled_len = state.pooled_len;
|
|
self.reserved_raw = self.raw_cap;
|
|
self.reserved_pooled = self.pooled_cap;
|
|
self.mirror = None;
|
|
self.mirror_cap = 0;
|
|
self.mirror_strides = [0; 2];
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub(super) enum QsaSelectMode {
|
|
Blocks,
|
|
DenseMask,
|
|
RowTokens,
|
|
}
|
|
|
|
impl QsaSelectMode {
|
|
fn name(self) -> &'static str {
|
|
match self {
|
|
Self::Blocks => "blocks",
|
|
Self::DenseMask => "dense_mask",
|
|
Self::RowTokens => "row_tokens",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Original selector specialization for the installed H4/D128/ratio4 model.
|
|
/// Logical frontiers remain device scalars, not shader specializations.
|
|
#[derive(Clone, Copy)]
|
|
#[cfg(test)]
|
|
pub(super) struct QsaSelectPlan {
|
|
pub(super) rows: u32,
|
|
pub(super) blocks: u32,
|
|
pub(super) topk: u32,
|
|
pub(super) pos: u32,
|
|
pub(super) total: u32,
|
|
pub(super) logical: u32,
|
|
pub(super) dense_width: u32,
|
|
pub(super) mode: QsaSelectMode,
|
|
pub(super) tf32: bool,
|
|
pub(super) input_f32: [bool; 2],
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl QsaSelectPlan {
|
|
fn widths(self) -> Vec<(u32, u32)> {
|
|
match self.mode {
|
|
QsaSelectMode::Blocks => vec![(self.topk, 4), (self.topk, 1), (self.topk, 4)],
|
|
QsaSelectMode::DenseMask => vec![(self.dense_width, 1)],
|
|
QsaSelectMode::RowTokens => vec![(self.topk * 4 + 4, 4), (self.topk * 4 + 4, 1)],
|
|
}
|
|
}
|
|
|
|
fn canonical(self) -> qsa_select::Kernel {
|
|
use array::Dtype;
|
|
qsa_select::Kernel {
|
|
heads: 4,
|
|
dim: 128,
|
|
blocks: self.blocks as i32,
|
|
topk: self.topk as i32,
|
|
ratio: 4,
|
|
dense_width: if self.mode == QsaSelectMode::DenseMask {
|
|
self.dense_width as i32
|
|
} else {
|
|
0
|
|
},
|
|
mode: self.mode,
|
|
tf32: self.tf32,
|
|
dtypes: self
|
|
.input_f32
|
|
.map(|f32| if f32 { Dtype::F32 } else { Dtype::BF16 }),
|
|
}
|
|
}
|
|
|
|
pub(super) fn source(self) -> Result<String, String> {
|
|
self.canonical().source()
|
|
}
|
|
}
|
|
|
|
/// One original fused selector dispatch, including its otherwise private scratch.
|
|
#[cfg(test)]
|
|
pub(super) struct QsaSelection {
|
|
pub(super) outputs: Vec<Buffer>,
|
|
pub(super) row_widths: Vec<(u32, u32)>, // elements, storage bytes
|
|
pub(super) row_strides: Vec<u32>,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// One original fused selector dispatch, including its otherwise private scratch.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_select(
|
|
[q, pooled]: [&Buffer; 2],
|
|
q_strides: [u64; 3],
|
|
pooled_strides: [u64; 2],
|
|
plan: QsaSelectPlan,
|
|
) -> Result<QsaSelection, String> {
|
|
if plan.rows == 0
|
|
|| plan.rows > i32::MAX as u32 / 512
|
|
|| !(1..=512).contains(&plan.topk)
|
|
|| plan.total > i32::MAX as u32
|
|
|| plan
|
|
.pos
|
|
.checked_add(plan.rows)
|
|
.is_none_or(|n| n > plan.total)
|
|
|| plan.logical > plan.blocks
|
|
|| plan.logical != plan.total / 4
|
|
|| (plan.mode == QsaSelectMode::DenseMask && plan.dense_width < plan.total)
|
|
{
|
|
return Err("invalid QSA selector geometry or frontier".into());
|
|
}
|
|
let positions =
|
|
[plan.pos, plan.total, plan.logical].map(|value| scalar_buffer(&value.to_le_bytes()));
|
|
let [pos, total, logical] = positions;
|
|
let positions = [pos?, total?, logical?];
|
|
let mut out = qsa_select_device(
|
|
[q, pooled],
|
|
q_strides,
|
|
pooled_strides,
|
|
plan,
|
|
[&positions[0], &positions[1], &positions[2]],
|
|
)?;
|
|
out._temporaries.extend(positions);
|
|
Ok(out)
|
|
}
|
|
|
|
/// Fixed selector geometry with explicit int32[1] graph frontier leaves.
|
|
/// plan.pos/total/logical are unused here; only the device leaves are consumed.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_select_device(
|
|
[q, pooled]: [&Buffer; 2],
|
|
q_strides: [u64; 3],
|
|
pooled_strides: [u64; 2],
|
|
plan: QsaSelectPlan,
|
|
[pos, total, logical]: [&Buffer; 3],
|
|
) -> Result<QsaSelection, String> {
|
|
if plan.rows == 0
|
|
|| plan.rows > i32::MAX as u32 / 512
|
|
|| plan.blocks > i32::MAX as u32
|
|
|| !(1..=512).contains(&plan.topk)
|
|
|| [pos, total, logical]
|
|
.into_iter()
|
|
.any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
|
|
{
|
|
return Err("invalid QSA device selector geometry or leaves".into());
|
|
}
|
|
for (buffer, shape, strides, f32) in [
|
|
(q, [plan.rows, 4, 128], q_strides, plan.input_f32[0]),
|
|
(
|
|
pooled,
|
|
[1, plan.blocks, 128],
|
|
[0, pooled_strides[0], pooled_strides[1]],
|
|
plan.input_f32[1],
|
|
),
|
|
] {
|
|
let span = shape
|
|
.into_iter()
|
|
.zip(strides)
|
|
.try_fold(0_u64, |n, (dim, stride)| {
|
|
u64::from(dim.saturating_sub(1))
|
|
.checked_mul(stride)
|
|
.and_then(|v| n.checked_add(v))
|
|
})
|
|
.and_then(|n| n.checked_add(1))
|
|
.and_then(|n| n.checked_mul(if f32 { 4 } else { 2 }));
|
|
if !shape.contains(&0)
|
|
&& span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < n)
|
|
{
|
|
return Err("QSA selector input exceeds backing buffer".into());
|
|
}
|
|
}
|
|
use array::{Array, Dtype};
|
|
let kernel = plan.canonical();
|
|
let q = buffer_array(
|
|
q,
|
|
&[1, plan.rows as i32, 4, 128],
|
|
&[0, q_strides[0], q_strides[1], q_strides[2]],
|
|
kernel.dtypes[0],
|
|
)?;
|
|
let pooled = buffer_array(
|
|
pooled,
|
|
&[1, plan.blocks as i32, 128],
|
|
&[0, pooled_strides[0], pooled_strides[1]],
|
|
kernel.dtypes[1],
|
|
)?;
|
|
let inputs = vec![
|
|
q,
|
|
pooled,
|
|
Array::new(&[1], Dtype::I32, pos.view(0, 4)?)?,
|
|
Array::new(&[1], Dtype::I32, total.view(0, 4)?)?,
|
|
Array::new(&[1], Dtype::I32, logical.view(0, 4)?)?,
|
|
];
|
|
let arrays = kernel
|
|
.layouts(plan.rows as i32)?
|
|
.into_iter()
|
|
.map(|layout| {
|
|
Array::new(
|
|
layout.shape(),
|
|
layout.dtype(),
|
|
Buffer::mtplx_bytes((layout.nbytes() as u64).max(4))?,
|
|
)
|
|
})
|
|
.collect::<Result<Vec<_>, String>>()?;
|
|
qsa_select::encode(&inputs, &arrays, kernel)?;
|
|
let mut buffers = arrays
|
|
.iter()
|
|
.map(|a| a.buffer().view(0, (a.layout().nbytes() as u64).max(4)))
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
let scratch = buffers.pop().ok_or("QSA scratch sibling missing")?;
|
|
let widths = plan.widths();
|
|
Ok(QsaSelection {
|
|
outputs: buffers,
|
|
row_strides: widths.iter().map(|(w, _)| *w).collect(),
|
|
row_widths: widths,
|
|
_temporaries: vec![scratch],
|
|
})
|
|
}
|
|
|
|
/// Preserve an existing Buffer view when delegating to the canonical backend.
|
|
#[cfg(test)]
|
|
fn buffer_array(
|
|
buffer: &Buffer,
|
|
shape: &[i32],
|
|
strides: &[u64],
|
|
dtype: array::Dtype,
|
|
) -> Result<array::Array, String> {
|
|
use array::{Array, Flags};
|
|
let strides = strides
|
|
.iter()
|
|
.map(|&s| i64::try_from(s).map_err(|_| "QSA stride overflow"))
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
let size = if shape.contains(&0) {
|
|
0
|
|
} else {
|
|
shape.iter().zip(&strides).try_fold(1_i64, |n, (&d, &s)| {
|
|
n.checked_add(
|
|
i64::from(d.saturating_sub(1))
|
|
.checked_mul(s)
|
|
.ok_or("QSA span overflow")?,
|
|
)
|
|
.ok_or("QSA span overflow")
|
|
})? as usize
|
|
};
|
|
let a = Array::unallocated(shape, dtype)?;
|
|
let (_, row, col) = a.layout().contiguity(&strides)?;
|
|
a.set_strided_data(
|
|
buffer.view(0, (size * dtype.itemsize()) as u64)?,
|
|
size,
|
|
&strides,
|
|
Flags {
|
|
contiguous: row || col,
|
|
row_contiguous: row,
|
|
col_contiguous: col,
|
|
},
|
|
0,
|
|
)?;
|
|
Ok(a)
|
|
}
|
|
|
|
/// QSAIndexer._select_fused: bound score scratch per dispatch, concatenate
|
|
/// original typed outputs, then return a zero-copy logical dense-mask slice.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_select_chunked(
|
|
input: [&Buffer; 2],
|
|
q_strides: [u64; 3],
|
|
pooled_strides: [u64; 2],
|
|
mut plan: QsaSelectPlan,
|
|
) -> Result<QsaSelection, String> {
|
|
if plan.rows == 0
|
|
|| plan.rows > i32::MAX as u32 / 512
|
|
|| !(1..=512).contains(&plan.topk)
|
|
|| plan.total > i32::MAX as u32
|
|
|| plan
|
|
.pos
|
|
.checked_add(plan.rows)
|
|
.is_none_or(|n| n > plan.total)
|
|
|| plan.logical > plan.blocks
|
|
|| plan.logical != plan.total / 4
|
|
{
|
|
return Err("invalid chunked QSA geometry".into());
|
|
}
|
|
if plan.mode == QsaSelectMode::DenseMask {
|
|
plan.dense_width = plan
|
|
.blocks
|
|
.checked_add(1)
|
|
.and_then(|n| n.checked_mul(4))
|
|
.ok_or("QSA dense capacity overflow")?;
|
|
}
|
|
let chunk_rows = u64::from(plan.rows)
|
|
.min(((32 * 1024 * 1024) / (u64::from(plan.blocks.max(1)) * 4)).max(1))
|
|
as u32;
|
|
let mut chunks = Vec::new();
|
|
for start in (0..plan.rows).step_by(chunk_rows as usize) {
|
|
let rows = chunk_rows.min(plan.rows - start);
|
|
let at = u64::from(start)
|
|
.checked_mul(q_strides[0])
|
|
.and_then(|n| n.checked_mul(if plan.input_f32[0] { 4 } else { 2 }))
|
|
.ok_or("QSA query slice overflow")?;
|
|
let bytes = unsafe { ds4_gpu_tensor_bytes(input[0].raw()) }
|
|
.checked_sub(at)
|
|
.ok_or("QSA query slice outside backing")?;
|
|
let query = input[0].view(at, bytes)?;
|
|
let pos = plan
|
|
.pos
|
|
.checked_add(start)
|
|
.ok_or("QSA query position overflow")?;
|
|
chunks.push(qsa_select(
|
|
[&query, input[1]],
|
|
q_strides,
|
|
pooled_strides,
|
|
QsaSelectPlan { rows, pos, ..plan },
|
|
)?);
|
|
}
|
|
let mut result = qsa_concat_selection(chunks, plan, chunk_rows)?;
|
|
if plan.mode == QsaSelectMode::DenseMask {
|
|
result.row_widths[0].0 = plan.total;
|
|
}
|
|
Ok(result)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_concat_selection(
|
|
mut chunks: Vec<QsaSelection>,
|
|
plan: QsaSelectPlan,
|
|
chunk_rows: u32,
|
|
) -> Result<QsaSelection, String> {
|
|
Ok(if chunks.len() == 1 {
|
|
chunks.pop().unwrap()
|
|
} else {
|
|
let widths = plan.widths();
|
|
let outputs = widths
|
|
.iter()
|
|
.map(|(width, bytes)| {
|
|
Buffer::mtplx_bytes(u64::from(plan.rows) * u64::from(*width) * u64::from(*bytes))
|
|
})
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
for (i, (target, (width, bytes))) in outputs.iter().zip(&widths).enumerate() {
|
|
let dtype = if *bytes == 1 {
|
|
"bool_"
|
|
} else if i == 0 {
|
|
"int32"
|
|
} else {
|
|
"float32"
|
|
};
|
|
let concurrent = encoder::concurrent_scope();
|
|
for (chunk_index, chunk) in chunks.iter().enumerate() {
|
|
let start = chunk_index as u32 * chunk_rows;
|
|
let rows = chunk_rows.min(plan.rows - start);
|
|
let view = target.view(
|
|
u64::from(start) * u64::from(*width) * u64::from(*bytes),
|
|
u64::from(rows) * u64::from(*width) * u64::from(*bytes),
|
|
)?;
|
|
qsa_copy_typed(
|
|
Some(&chunk.outputs[i]),
|
|
&view,
|
|
[1, rows, *width],
|
|
[0, u64::from(*width), 1],
|
|
[0, u64::from(*width), 1],
|
|
dtype,
|
|
(rows) as usize * (*width) as usize,
|
|
)?;
|
|
}
|
|
drop(concurrent);
|
|
}
|
|
let mut temporaries = Vec::new();
|
|
for chunk in chunks {
|
|
temporaries.extend(chunk.outputs);
|
|
temporaries.extend(chunk._temporaries);
|
|
}
|
|
QsaSelection {
|
|
outputs,
|
|
row_strides: widths.iter().map(|(w, _)| *w).collect(),
|
|
row_widths: widths,
|
|
_temporaries: temporaries,
|
|
}
|
|
})
|
|
}
|
|
|
|
/// The producer-aware 128MiB score-plane budget, including the original
|
|
/// alignment-down rule. The general FP32 producer needs H4+1 planes.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prefill_chunk_rows(
|
|
rows: u32,
|
|
blocks: u32,
|
|
budget: u64,
|
|
alignment: u32,
|
|
mpp: bool,
|
|
) -> Result<u32, String> {
|
|
if rows == 0 || blocks == 0 || budget == 0 || alignment == 0 {
|
|
return Err("invalid QSA prefill budget".into());
|
|
}
|
|
let per_row = u64::from(blocks) * 4 * if mpp { 1 } else { 5 };
|
|
let mut chunk = u64::from(rows).min((budget / per_row).max(1)) as u32;
|
|
if chunk >= alignment && chunk < rows {
|
|
chunk = alignment.max(chunk / alignment * alignment);
|
|
}
|
|
Ok(chunk)
|
|
}
|
|
|
|
/// Original General copy (including cast), after runtime dimension collapse.
|
|
#[cfg(test)]
|
|
fn qsa_copy_f32(
|
|
input: &Buffer,
|
|
shape: [u32; 3],
|
|
strides: [u64; 3],
|
|
bf16: bool,
|
|
) -> Result<Buffer, String> {
|
|
qsa_copy_general(
|
|
input,
|
|
shape,
|
|
strides,
|
|
if bf16 { "bfloat16" } else { "float32" },
|
|
"float32",
|
|
4,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_copy_general(
|
|
input: &Buffer,
|
|
shape: [u32; 3],
|
|
strides: [u64; 3],
|
|
source_type: &str,
|
|
output_type: &str,
|
|
output_bytes: u32,
|
|
) -> Result<Buffer, String> {
|
|
let count = shape
|
|
.into_iter()
|
|
.try_fold(1_u32, |n, d| n.checked_mul(d))
|
|
.filter(|n| *n > 0 && *n <= i32::MAX as u32)
|
|
.ok_or("QSA General copy exceeds its index range")?;
|
|
let output = Buffer::mtplx_bytes(u64::from(count) * u64::from(output_bytes))?;
|
|
let mut dims: Vec<(u32, u64)> = Vec::new();
|
|
for (d, s) in shape.into_iter().zip(strides).filter(|(d, _)| *d > 1) {
|
|
if let Some(last) = dims.last_mut()
|
|
&& last.1 == u64::from(d) * s
|
|
{
|
|
last.0 *= d;
|
|
last.1 = s;
|
|
} else {
|
|
dims.push((d, s));
|
|
}
|
|
}
|
|
let mut grid = [1; 3];
|
|
let mut ss = [0_u64; 3];
|
|
for (i, &(d, s)) in dims.iter().enumerate() {
|
|
grid[dims.len() - 1 - i] = d;
|
|
ss[i] = s;
|
|
}
|
|
dispatch_geometry(
|
|
&format!("g{}_copy{source_type}{output_type}", dims.len()),
|
|
&[
|
|
tensor(0, input).input(qsa_copy_span(shape, strides)?),
|
|
tensor(1, &output).output(count as usize),
|
|
bytes(3, &ss),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
Ok(output)
|
|
}
|
|
|
|
// AsType Vector preserves physically contiguous layouts, including transposes
|
|
// and broadcasts; General materializes all logical entries.
|
|
#[cfg(test)]
|
|
fn qsa_cast_f32(input: &Buffer, shape: [u32; 3], strides: &mut [u64; 3]) -> Result<Buffer, String> {
|
|
let count = shape.into_iter().try_fold(1_u32, |n, d| n.checked_mul(d));
|
|
let span = shape
|
|
.into_iter()
|
|
.zip(*strides)
|
|
.try_fold(1_u64, |n, (d, s)| {
|
|
u64::from(d.saturating_sub(1))
|
|
.checked_mul(s)
|
|
.and_then(|v| n.checked_add(v))
|
|
})
|
|
.and_then(|n| n.checked_mul(2));
|
|
if shape.contains(&0)
|
|
|| count.is_none_or(|n| n > i32::MAX as u32)
|
|
|| span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
|
|
{
|
|
return Err("invalid QSA cast input".into());
|
|
}
|
|
let mut dims = shape
|
|
.into_iter()
|
|
.zip(*strides)
|
|
.filter(|(d, s)| *d > 1 && *s != 0)
|
|
.collect::<Vec<_>>();
|
|
dims.sort_unstable_by_key(|(_, s)| *s);
|
|
let mut size = 1_u64;
|
|
let contiguous = dims.into_iter().all(|(d, s)| {
|
|
let fits = s == size;
|
|
size *= u64::from(d);
|
|
fits
|
|
});
|
|
if contiguous {
|
|
let cast = Buffer::mtplx_bytes(size * 4)?;
|
|
cast_bf16_f32(input, &cast, size as u32, true)?;
|
|
Ok(cast)
|
|
} else {
|
|
let cast = qsa_copy_f32(input, shape, *strides, true)?;
|
|
*strides = [
|
|
u64::from(shape[1]) * u64::from(shape[2]),
|
|
u64::from(shape[2]),
|
|
1,
|
|
];
|
|
Ok(cast)
|
|
}
|
|
}
|
|
|
|
#[repr(C)]
|
|
struct QsaGemmParams {
|
|
matrix: [i32; 8],
|
|
batch_strides: [i64; 3],
|
|
swizzle: i32,
|
|
k_iterations: i32,
|
|
batch_ndim: i32,
|
|
padding: i32,
|
|
}
|
|
|
|
/// qsa_indexer_prefill_scores / _select_eager, installed H4/D128 geometry.
|
|
/// Preserve the pinned runtime's casts, layout checks, batch collapse and
|
|
/// M5 Max GEMV/Steel/NAX/Split-K selection, followed by per-head ReLU and sum.
|
|
/// The caller retains the returned temporaries through GPU completion.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prefill_f32_scores(
|
|
[q, pooled_t]: [&Buffer; 2],
|
|
[rows, blocks]: [u32; 2],
|
|
mut q_strides: [u64; 3],
|
|
pooled_strides: [u64; 2],
|
|
q_bf16: bool,
|
|
tf32: bool,
|
|
) -> Result<(Buffer, Vec<Buffer>), String> {
|
|
if rows == 0
|
|
|| blocks == 0
|
|
|| rows.checked_mul(512).is_none_or(|n| n > i32::MAX as u32)
|
|
|| rows
|
|
.checked_mul(4)
|
|
.and_then(|n| n.checked_mul(blocks))
|
|
.is_none_or(|n| n > i32::MAX as u32)
|
|
|| blocks.checked_mul(128).is_none_or(|n| n > i32::MAX as u32)
|
|
{
|
|
return Err("invalid QSA FP32 score geometry".into());
|
|
}
|
|
for (b, shape, strides, width) in [
|
|
(q, [rows, 4, 128], q_strides, if q_bf16 { 2 } else { 4 }),
|
|
(
|
|
pooled_t,
|
|
[1, 128, blocks],
|
|
[0, pooled_strides[0], pooled_strides[1]],
|
|
4,
|
|
),
|
|
] {
|
|
let span = shape
|
|
.into_iter()
|
|
.zip(strides)
|
|
.try_fold(1_u64, |n, (d, s)| {
|
|
u64::from(d - 1)
|
|
.checked_mul(s)
|
|
.and_then(|v| n.checked_add(v))
|
|
})
|
|
.and_then(|n| n.checked_mul(width));
|
|
if strides.iter().any(|s| *s > i32::MAX as u64)
|
|
|| span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n)
|
|
{
|
|
return Err("QSA FP32 input outside backing or index range".into());
|
|
}
|
|
}
|
|
let mut temporary = Vec::new();
|
|
let cast = if q_bf16 {
|
|
Some(qsa_cast_f32(q, [rows, 4, 128], &mut q_strides)?)
|
|
} else {
|
|
None
|
|
};
|
|
let a = cast.as_ref().unwrap_or(q);
|
|
let per_head_size = rows as usize * 4 * blocks as usize;
|
|
// Matmul allocates its output before check_transpose's input copies.
|
|
let per_head = Buffer::mtplx_bytes(per_head_size as u64 * 4)?;
|
|
let copied_a = if q_strides[2] != 1 && q_strides[1] != 1 {
|
|
let copied = qsa_copy_f32(a, [rows, 4, 128], q_strides, false)?;
|
|
q_strides = [512, 128, 1];
|
|
Some(copied)
|
|
} else {
|
|
None
|
|
};
|
|
let a = copied_a.as_ref().unwrap_or(a);
|
|
let ta = q_strides[2] != 1;
|
|
let lda = q_strides[if ta { 2 } else { 1 }] as i32;
|
|
let [bd, bn_stride] = pooled_strides;
|
|
let b_normal = bn_stride == 1 && (blocks != 1 || bd == u64::from(blocks));
|
|
let b_transposed = bd == 1 && (blocks != 1 || bn_stride == 128);
|
|
// check_transpose sees the already broadcast B. A General copy therefore
|
|
// duplicates B across rows and prevents collapsing the rows into M.
|
|
let copied_b = if b_normal || b_transposed {
|
|
None
|
|
} else {
|
|
Some(qsa_copy_f32(
|
|
pooled_t,
|
|
[rows, 128, blocks],
|
|
[0, bd, bn_stride],
|
|
false,
|
|
)?)
|
|
};
|
|
let b = copied_b.as_ref().unwrap_or(pooled_t);
|
|
let tb = copied_b.is_none() && !b_normal;
|
|
let ldb = if copied_b.is_some() {
|
|
blocks
|
|
} else if tb {
|
|
bn_stride as u32
|
|
} else {
|
|
bd as u32
|
|
} as i32;
|
|
let mut batch = rows;
|
|
let mut m = 4_u32;
|
|
let mut a_batch = q_strides[0] as i64;
|
|
let mut b_batch = if copied_b.is_some() {
|
|
i64::from(blocks) * 128
|
|
} else {
|
|
0
|
|
};
|
|
if rows > 1 && !ta && q_strides[1] == 128 && a_batch == 512 && b_batch == 0 {
|
|
m *= rows;
|
|
batch = 1;
|
|
a_batch = 0;
|
|
b_batch = 0;
|
|
}
|
|
let a_size = qsa_copy_span([rows, 4, 128], q_strides)?;
|
|
let b_size = if copied_b.is_some() {
|
|
rows as usize * 128 * blocks as usize
|
|
} else {
|
|
qsa_copy_span([1, 128, blocks], [0, bd, bn_stride])?
|
|
};
|
|
let trans = format!(
|
|
"{}{}",
|
|
if ta { 't' } else { 'n' },
|
|
if tb { 't' } else { 'n' }
|
|
);
|
|
if blocks == 1 {
|
|
let (bm, bn, sm, sn) = if ta {
|
|
(
|
|
1,
|
|
if m >= 2048 {
|
|
16
|
|
} else if m >= 512 {
|
|
4
|
|
} else {
|
|
2
|
|
},
|
|
8,
|
|
4,
|
|
)
|
|
} else if 128 >= 16 * m {
|
|
(1, 8, 1, 32)
|
|
} else {
|
|
(if m >= 4096 { 8 } else { 4 }, 1, 1, 32)
|
|
};
|
|
let out_per_group = if ta { bn * sn * 4 } else { bm * sm * 4 };
|
|
dispatch(
|
|
&format!(
|
|
"gemv{}float32_bm{bm}_bn{bn}_sm{sm}_sn{sn}_tm4_tn4_nc0_axpby0",
|
|
if ta { "_t_" } else { "_" }
|
|
),
|
|
&[
|
|
tensor(0, a).input(a_size),
|
|
tensor(1, b).input(b_size),
|
|
tensor(3, &per_head).output(per_head_size),
|
|
bytes(4, &128_i32),
|
|
bytes(5, &m),
|
|
bytes(6, &lda),
|
|
bytes(9, &1_i32),
|
|
bytes(10, &batch),
|
|
bytes(11, &b_batch),
|
|
bytes(12, &a_batch),
|
|
],
|
|
[m.div_ceil(out_per_group), 1, batch],
|
|
[32, bn, bm],
|
|
)?;
|
|
} else if batch == 1 && ((!tf32 && m.max(blocks) <= 128) || (tf32 && m.max(blocks) < 64)) {
|
|
let (bm, bn, bk, partitions, partition_size) = if tf32 {
|
|
(64, 64, 256, 2, 64)
|
|
} else {
|
|
let parts = (8 / (m.div_ceil(32) * blocks.div_ceil(32)))
|
|
.next_power_of_two()
|
|
.clamp(2, 32);
|
|
(
|
|
if m < 40 { 16 } else { 32 },
|
|
if blocks < 40 { 16 } else { 32 },
|
|
16,
|
|
parts,
|
|
128 / parts,
|
|
)
|
|
};
|
|
let scratch_size = partitions as usize * m as usize * blocks as usize;
|
|
let scratch = Buffer::mtplx_bytes(scratch_size as u64 * 4)?;
|
|
let tn = blocks.div_ceil(bn);
|
|
let tm = m.div_ceil(bm);
|
|
let params = [
|
|
m as i32,
|
|
blocks as i32,
|
|
128,
|
|
lda,
|
|
ldb,
|
|
blocks as i32,
|
|
tn as i32,
|
|
tm as i32,
|
|
partitions as i32,
|
|
(m * blocks) as i32,
|
|
partition_size as i32,
|
|
0,
|
|
(partition_size / bk) as i32,
|
|
];
|
|
let constants = [
|
|
(200, m.is_multiple_of(bm)),
|
|
(201, blocks.is_multiple_of(bn)),
|
|
]
|
|
.map(|(index, value)| MetalConstant {
|
|
index,
|
|
value: u32::from(value),
|
|
kind: 0,
|
|
});
|
|
let name = if tf32 {
|
|
format!("steel_gemm_splitk_nax_{trans}_float32_float32_bm64_bn64_bk256_wm2_wn2")
|
|
} else {
|
|
format!(
|
|
"steel_gemm_splitk_{trans}_float32_float32_bm{bm}_bn{bn}_bk16_wm2_wn2_MN_{}aligned_K_taligned",
|
|
if m.is_multiple_of(bm) && blocks.is_multiple_of(bn) {
|
|
't'
|
|
} else {
|
|
'n'
|
|
}
|
|
)
|
|
};
|
|
dispatch_specialized(
|
|
&name,
|
|
&[
|
|
tensor(0, a).input(a_size),
|
|
tensor(1, b).input(b_size),
|
|
tensor(2, &scratch).output(scratch_size),
|
|
bytes(3, ¶ms),
|
|
],
|
|
if tf32 { &constants } else { &[] },
|
|
if tf32 {
|
|
[tn * tm * partitions, 1, 1]
|
|
} else {
|
|
[tn, tm, partitions]
|
|
},
|
|
[32, 2, 2],
|
|
)?;
|
|
let grid = [blocks, m, 1];
|
|
dispatch_geometry(
|
|
"steel_gemm_splitk_accum_float32_float32",
|
|
&[
|
|
tensor(0, &scratch).input(scratch_size),
|
|
tensor(1, &per_head).output(per_head_size),
|
|
bytes(2, &partitions),
|
|
bytes(3, &(m * blocks)),
|
|
bytes(4, &blocks),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
encoder::add_temporaries(std::slice::from_ref(&scratch));
|
|
temporary.push(scratch);
|
|
} else {
|
|
let (bn, bk, wn, swizzle) = if tf32 {
|
|
(128, 256, 4, 2)
|
|
} else {
|
|
(64, 16, 2, 0)
|
|
};
|
|
let tn = blocks.div_ceil(bn);
|
|
let tm = m.div_ceil(64);
|
|
let params = QsaGemmParams {
|
|
matrix: [
|
|
m as i32,
|
|
blocks as i32,
|
|
128,
|
|
lda,
|
|
ldb,
|
|
blocks as i32,
|
|
tn as i32,
|
|
tm as i32,
|
|
],
|
|
batch_strides: [a_batch, b_batch, i64::from(m) * i64::from(blocks)],
|
|
swizzle,
|
|
k_iterations: 128 / bk,
|
|
batch_ndim: 1,
|
|
padding: 0,
|
|
};
|
|
let constants = [
|
|
(10, false),
|
|
(100, false),
|
|
(110, false),
|
|
(200, m.is_multiple_of(64)),
|
|
(201, blocks.is_multiple_of(bn)),
|
|
(202, !tf32),
|
|
]
|
|
.map(|(index, value)| MetalConstant {
|
|
index,
|
|
value: u32::from(value),
|
|
kind: 0,
|
|
});
|
|
dispatch_specialized(
|
|
&format!(
|
|
"steel_gemm_fused_{}{trans}_float32_float32_bm64_bn{bn}_bk{bk}_wm2_wn{wn}",
|
|
if tf32 { "nax_" } else { "" }
|
|
),
|
|
&[
|
|
tensor(0, a).input(a_size),
|
|
tensor(1, b).input(b_size),
|
|
tensor(3, &per_head).output(per_head_size),
|
|
bytes(4, ¶ms),
|
|
],
|
|
&constants,
|
|
[tn * (1 << swizzle), tm.div_ceil(1 << swizzle), batch],
|
|
[32, wn, 2],
|
|
)?;
|
|
}
|
|
encoder::add_temporaries(copied_a.as_slice());
|
|
encoder::add_temporaries(copied_b.as_slice());
|
|
temporary.extend(cast);
|
|
temporary.extend(copied_a);
|
|
temporary.extend(copied_b);
|
|
let zero = scalar_buffer(&0_f32.to_le_bytes())?;
|
|
let count = rows * 4 * blocks;
|
|
let work = if count < 65536 { 1 } else { 2 };
|
|
dispatch_geometry(
|
|
&format!("vs{}_Maximumfloat32", if work == 1 { "" } else { "n" }),
|
|
&[
|
|
tensor(0, &per_head).input(per_head_size),
|
|
tensor(1, &zero).input(1),
|
|
tensor(2, &per_head).output(per_head_size),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[count.div_ceil(work), 1, 1],
|
|
[count.div_ceil(work).min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
let scores_size = rows as usize * blocks as usize;
|
|
let scores = Buffer::mtplx_bytes(scores_size as u64 * 4)?;
|
|
let ndim = u32::from(rows > 1);
|
|
let outer_shape = if rows > 1 { rows } else { 0 };
|
|
let outer_stride = if rows > 1 { u64::from(blocks) * 4 } else { 0 };
|
|
if blocks == 1 {
|
|
dispatch_geometry(
|
|
"row_reduce_small_1_reduce_sumfloat32",
|
|
&[
|
|
tensor(0, &per_head).input(per_head_size),
|
|
tensor(1, &scores).output(scores_size),
|
|
bytes(2, &4_u64),
|
|
bytes(3, &1_u64),
|
|
bytes(4, &outer_shape),
|
|
bytes(5, &outer_stride),
|
|
bytes(6, &ndim),
|
|
bytes(7, &0_i32),
|
|
bytes(8, &0_u64),
|
|
bytes(9, &0_i32),
|
|
],
|
|
&[],
|
|
[rows, 1, 1],
|
|
[rows.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
} else {
|
|
let gx = blocks.div_ceil(4).min(32);
|
|
dispatch(
|
|
"col_reduce_small_1_reduce_sumfloat32",
|
|
&[
|
|
tensor(0, &per_head).input(per_head_size),
|
|
tensor(1, &scores).output(scores_size),
|
|
bytes(2, &4_u64),
|
|
bytes(3, &u64::from(blocks)),
|
|
bytes(4, &outer_shape),
|
|
bytes(5, &outer_stride),
|
|
bytes(6, &ndim),
|
|
bytes(7, &4_u32),
|
|
bytes(8, &u64::from(blocks)),
|
|
bytes(9, &1_i32),
|
|
bytes(10, &1_u64),
|
|
],
|
|
[blocks.div_ceil(4).div_ceil(gx), rows, 1],
|
|
[gx, 4, 1],
|
|
)?;
|
|
}
|
|
let divisor = scalar_buffer(&(128_f64.sqrt() as f32).to_le_bytes())?;
|
|
let count = rows * blocks;
|
|
let work = if count < 65536 { 1 } else { 2 };
|
|
dispatch_geometry(
|
|
&format!("vs{}_Dividefloat32", if work == 1 { "" } else { "n" }),
|
|
&[
|
|
tensor(0, &scores).input(scores_size),
|
|
tensor(1, &divisor).input(1),
|
|
tensor(2, &scores).output(scores_size),
|
|
bytes(3, &count),
|
|
],
|
|
&[],
|
|
[count.div_ceil(work), 1, 1],
|
|
[count.div_ceil(work).min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
temporary.push(per_head);
|
|
Ok((scores, temporary))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_vector(
|
|
prefix: &str,
|
|
op: &str,
|
|
bindings: &[Binding<'_>],
|
|
count: u32,
|
|
width: u32,
|
|
) -> Result<(), String> {
|
|
let work = if count < 65536 { 1 } else { 8 / width };
|
|
let suffix = if work == 1 || prefix == "ss" { "" } else { "n" };
|
|
dispatch_geometry(
|
|
&format!("{prefix}{suffix}_{op}"),
|
|
bindings,
|
|
&[],
|
|
[count.div_ceil(work), 1, 1],
|
|
[count.div_ceil(work).min(1024), 1, 1],
|
|
true,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct QsaEagerTopk {
|
|
rows: u32,
|
|
pos_start: u32,
|
|
tiled: bool,
|
|
/// Untiled: last K columns of ArgPartition, stride N. Tiled: concat, stride K.
|
|
pub(super) indices: Buffer,
|
|
pub(super) valid: Buffer,
|
|
pub(super) masked_scores: Option<Buffer>,
|
|
pub(super) positions: Buffer,
|
|
pub(super) complete_blocks: Buffer,
|
|
pub(super) width: u32,
|
|
pub(super) row_stride: u32,
|
|
pub(super) index_row_stride: u32,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// _select_eager, including _tiled_topk's synchronous per-tile eval and concat.
|
|
/// ArgPartition in the pinned runtime is a complete argsort, not a replacement
|
|
/// custom selector. The pooled FP32 mirror and tie vector are shared by tiles.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_eager_topk(
|
|
input: [&Buffer; 2],
|
|
q_strides: [u64; 3],
|
|
pooled_t_strides: [u64; 2],
|
|
plan: QsaSelectPlan,
|
|
tile_rows: u32,
|
|
mut commands: Option<&mut Commands>,
|
|
) -> Result<QsaEagerTopk, String> {
|
|
if plan
|
|
.pos
|
|
.checked_add(plan.rows)
|
|
.is_none_or(|n| n > plan.total)
|
|
|| plan.total > i32::MAX as u32
|
|
|| plan.blocks != plan.total / 4
|
|
|| !plan.input_f32[1]
|
|
|| !(1..=512).contains(&plan.topk)
|
|
|| plan.rows == 0
|
|
|| plan.blocks == 0
|
|
|| plan
|
|
.rows
|
|
.checked_mul(plan.blocks)
|
|
.is_none_or(|n| n > i32::MAX as u32)
|
|
{
|
|
return Err("invalid QSA eager frontier".into());
|
|
}
|
|
let rows = plan.rows;
|
|
let blocks = plan.blocks;
|
|
let positions = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
|
|
let complete = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
|
|
let blk = Buffer::mtplx_bytes(u64::from(blocks) * 4)?;
|
|
let tie = Buffer::mtplx_bytes(u64::from(blocks) * 4)?;
|
|
let one = scalar_buffer(&1_i32.to_le_bytes())?;
|
|
let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
|
|
let tie_scale = scalar_buffer(&1e-12_f32.to_le_bytes())?;
|
|
let neg = scalar_buffer(&f32::NEG_INFINITY.to_le_bytes())?;
|
|
// Original vector dispatch threshold for int32/float32; SS entries have
|
|
// no size argument, which the shared Metal binder may harmlessly leave set.
|
|
let vector = |prefix: &str, op: &str, bindings: &[Binding<'_>], count: u32| {
|
|
qsa_vector(prefix, op, bindings, count, 4)
|
|
};
|
|
for (out, count, start) in [(&positions, rows, plan.pos), (&blk, blocks, 0)] {
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &start),
|
|
bytes(1, &1_i32),
|
|
tensor(2, out).output(count as usize),
|
|
],
|
|
&[],
|
|
[count, 1, 1],
|
|
[count.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
}
|
|
let prefix = if rows == 1 { "ss" } else { "vs" };
|
|
vector(
|
|
prefix,
|
|
"Addint32",
|
|
&[
|
|
tensor(0, &positions).input(rows as usize),
|
|
tensor(1, &one).input(1),
|
|
tensor(2, &complete).output(rows as usize),
|
|
bytes(3, &rows),
|
|
],
|
|
rows,
|
|
)?;
|
|
vector(
|
|
prefix,
|
|
"Divideint32",
|
|
&[
|
|
tensor(0, &complete).input(rows as usize),
|
|
tensor(1, &ratio).input(1),
|
|
tensor(2, &complete).output(rows as usize),
|
|
bytes(3, &rows),
|
|
],
|
|
rows,
|
|
)?;
|
|
let work = if blocks < 65536 { 1 } else { 2 };
|
|
dispatch_geometry(
|
|
&format!("v{}_copyint32float32", if work == 1 { "" } else { "n" }),
|
|
&[
|
|
tensor(0, &blk).input(blocks as usize),
|
|
tensor(1, &tie).output(blocks as usize),
|
|
bytes(2, &blocks),
|
|
],
|
|
&[],
|
|
[blocks.div_ceil(work), 1, 1],
|
|
[blocks.div_ceil(work).min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
vector(
|
|
if blocks == 1 { "ss" } else { "vs" },
|
|
"Multiplyfloat32",
|
|
&[
|
|
tensor(0, &tie).input(blocks as usize),
|
|
tensor(1, &tie_scale).input(1),
|
|
tensor(2, &tie).output(blocks as usize),
|
|
bytes(3, &blocks),
|
|
],
|
|
blocks,
|
|
)?;
|
|
let width = plan.topk.min(blocks);
|
|
let tiled = rows > 1 && tile_rows > 0 && tile_rows < rows;
|
|
let (indices, valid, masked_scores, mut temporary) = if tiled {
|
|
let mut parts = Vec::new();
|
|
for start in (0..rows).step_by(tile_rows as usize) {
|
|
let len = tile_rows.min(rows - start);
|
|
let complete_view = complete.view(u64::from(start) * 4, u64::from(len) * 4)?;
|
|
let valid = qsa_eager_valid(&blk, &complete_view, len, blocks)?;
|
|
let item_bytes = if plan.input_f32[0] { 4 } else { 2 };
|
|
let q_len = u64::from(len - 1)
|
|
.checked_mul(q_strides[0])
|
|
.and_then(|n| n.checked_add(3_u64.checked_mul(q_strides[1])?))
|
|
.and_then(|n| n.checked_add(127_u64.checked_mul(q_strides[2])?))
|
|
.and_then(|n| n.checked_add(1))
|
|
.and_then(|n| n.checked_mul(item_bytes))
|
|
.ok_or("QSA eager tile query span overflow")?;
|
|
let offset = u64::from(start)
|
|
.checked_mul(q_strides[0])
|
|
.and_then(|n| n.checked_mul(item_bytes))
|
|
.ok_or("QSA eager tile query offset overflow")?;
|
|
let q = input[0].view(offset, q_len)?;
|
|
let (part, _scores, _scratch) = qsa_eager_rank(
|
|
[&q, input[1]],
|
|
q_strides,
|
|
pooled_t_strides,
|
|
QsaSelectPlan {
|
|
rows: len,
|
|
pos: plan.pos + start,
|
|
..plan
|
|
},
|
|
&valid,
|
|
&tie,
|
|
&neg,
|
|
)?;
|
|
// mx.eval(top_t) completes before constructing the next tile.
|
|
// Keep only the evaluated ArgPartition view/backing in parts;
|
|
// score intermediates are not retained for the entire query chunk.
|
|
if let Some(result) = encoder::with_active(|encoder| encoder.synchronize()) {
|
|
result?;
|
|
} else {
|
|
commands
|
|
.as_deref_mut()
|
|
.ok_or("QSA eager eval requires an active command owner")?
|
|
.eval_boundary()?;
|
|
}
|
|
parts.push((part, len));
|
|
}
|
|
let indices = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 4)?;
|
|
let mut start = 0;
|
|
let mut temporary = Vec::new();
|
|
let concurrent = encoder::concurrent_scope();
|
|
for (part, len) in parts {
|
|
let dst = indices.view(
|
|
u64::from(start) * u64::from(width) * 4,
|
|
u64::from(len) * u64::from(width) * 4,
|
|
)?;
|
|
qsa_copy_typed(
|
|
Some(&part),
|
|
&dst,
|
|
[1, len, width],
|
|
[0, u64::from(blocks), 1],
|
|
[0, u64::from(width), 1],
|
|
"uint32",
|
|
(len) as usize * (width) as usize,
|
|
)?;
|
|
temporary.push(part);
|
|
start += len;
|
|
}
|
|
drop(concurrent);
|
|
// Global valid is a different lazy expression from valid_t; evaluate
|
|
// it for the epilogue after the per-tile evals, as in the reference.
|
|
let valid = qsa_eager_valid(&blk, &complete, rows, blocks)?;
|
|
(indices, valid, None, temporary)
|
|
} else {
|
|
let valid = qsa_eager_valid(&blk, &complete, rows, blocks)?;
|
|
let (indices, scores, temporary) =
|
|
qsa_eager_rank(input, q_strides, pooled_t_strides, plan, &valid, &tie, &neg)?;
|
|
(indices, valid, Some(scores), temporary)
|
|
};
|
|
temporary.extend([blk, tie, one, ratio, tie_scale, neg]);
|
|
Ok(QsaEagerTopk {
|
|
rows,
|
|
pos_start: plan.pos,
|
|
tiled,
|
|
indices,
|
|
valid,
|
|
masked_scores,
|
|
positions,
|
|
complete_blocks: complete,
|
|
width,
|
|
row_stride: blocks,
|
|
index_row_stride: if tiled { width } else { blocks },
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_eager_valid(
|
|
blk: &Buffer,
|
|
complete: &Buffer,
|
|
rows: u32,
|
|
blocks: u32,
|
|
) -> Result<Buffer, String> {
|
|
let count = rows * blocks;
|
|
let valid = Buffer::mtplx_bytes(u64::from(count))?;
|
|
if rows == 1 || blocks == 1 {
|
|
qsa_vector(
|
|
if count == 1 {
|
|
"ss"
|
|
} else if rows == 1 {
|
|
"vs"
|
|
} else {
|
|
"sv"
|
|
},
|
|
"Lessint32",
|
|
&[
|
|
tensor(0, blk).input(blocks as usize),
|
|
tensor(1, complete).input(rows as usize),
|
|
tensor(2, &valid).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
} else {
|
|
let grid = [blocks, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_Lessint32",
|
|
&[
|
|
tensor(0, blk).input(blocks as usize),
|
|
tensor(1, complete).input(rows as usize),
|
|
tensor(2, &valid).output(count as usize),
|
|
bytes(3, &[0_u64, 1]),
|
|
bytes(4, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
Ok(valid)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_eager_rank(
|
|
input: [&Buffer; 2],
|
|
q_strides: [u64; 3],
|
|
pooled_t_strides: [u64; 2],
|
|
plan: QsaSelectPlan,
|
|
valid: &Buffer,
|
|
tie: &Buffer,
|
|
neg: &Buffer,
|
|
) -> Result<(Buffer, Buffer, Vec<Buffer>), String> {
|
|
let rows = plan.rows;
|
|
let blocks = plan.blocks;
|
|
let count = rows * blocks;
|
|
let (scores, mut temporary) = qsa_prefill_f32_scores(
|
|
input,
|
|
[rows, blocks],
|
|
q_strides,
|
|
pooled_t_strides,
|
|
!plan.input_f32[0],
|
|
plan.tf32,
|
|
)?;
|
|
qsa_vector(
|
|
if count == 1 { "v" } else { "vs" },
|
|
"Selectfloat32",
|
|
&[
|
|
tensor(0, valid).input(count as usize),
|
|
tensor(1, &scores).input(count as usize),
|
|
tensor(2, neg).input(1),
|
|
tensor(3, &scores).output(count as usize),
|
|
bytes(4, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
if rows == 1 || blocks == 1 {
|
|
qsa_vector(
|
|
if count == 1 {
|
|
"ss"
|
|
} else if blocks == 1 {
|
|
"vs"
|
|
} else {
|
|
"vv"
|
|
},
|
|
"Subtractfloat32",
|
|
&[
|
|
tensor(0, &scores).input(count as usize),
|
|
tensor(1, tie).input(blocks as usize),
|
|
tensor(2, &scores).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
} else {
|
|
let grid = [blocks, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_Subtractfloat32",
|
|
&[
|
|
tensor(0, &scores).input(count as usize),
|
|
tensor(1, tie).input(blocks as usize),
|
|
tensor(2, &scores).output(count as usize),
|
|
bytes(3, &[u64::from(blocks), 1]),
|
|
bytes(4, &[0_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let sorted = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
temporary.extend(argsort(&scores, &sorted, [rows, blocks], SortDtype::F32)?);
|
|
let width = plan.topk.min(blocks);
|
|
let offset = u64::from(blocks - width) * 4;
|
|
let indices = sorted.view(offset, u64::from(count) * 4 - offset)?;
|
|
temporary.push(sorted);
|
|
Ok((indices, scores, temporary))
|
|
}
|
|
|
|
/// The eager selector's output switches. `large_prefill` is the common
|
|
/// phase/rows/earliest-position eligibility decision, also used by the fused
|
|
/// and compiled selectors; it must not be inferred from rows > 1 alone.
|
|
pub(super) struct QsaEagerOptions {
|
|
pub(super) large_prefill: bool,
|
|
pub(super) flash: bool,
|
|
pub(super) gather_decode: bool,
|
|
pub(super) gather: bool,
|
|
pub(super) gather_max_rows: u32,
|
|
pub(super) gather_min_context: u32,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) enum QsaIndexerSelection {
|
|
Dense(QsaSelection),
|
|
PrefillBlocks(QsaSelection),
|
|
FlashBlocks(QsaSelection, u32),
|
|
DecodeTokens(QsaSelection),
|
|
RowTokens(QsaSelection),
|
|
}
|
|
|
|
/// Complete _select_eager return-branch order. Inactive lazy expressions
|
|
/// (especially the dense scatter mask in flash/gather modes) are not encoded.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_eager_output(
|
|
top: QsaEagerTopk,
|
|
total: u32,
|
|
options: QsaEagerOptions,
|
|
) -> Result<QsaIndexerSelection, String> {
|
|
if top.pos_start.checked_add(top.rows) != Some(total)
|
|
|| top.row_stride != total / 4
|
|
|| top.row_stride <= 512
|
|
{
|
|
return Err("invalid active eager selector output frontier".into());
|
|
}
|
|
if top.rows > 1 && options.large_prefill {
|
|
return qsa_eager_blocks(top).map(QsaIndexerSelection::PrefillBlocks);
|
|
}
|
|
if top.rows == 1 && options.flash {
|
|
let (blocks, tail) = qsa_eager_decode(top, total, false)?;
|
|
return Ok(QsaIndexerSelection::FlashBlocks(blocks, tail));
|
|
}
|
|
if top.rows == 1 && options.gather_decode {
|
|
return qsa_eager_decode(top, total, true)
|
|
.map(|(tokens, _)| QsaIndexerSelection::DecodeTokens(tokens));
|
|
}
|
|
if top.rows > 1
|
|
&& !top.tiled
|
|
&& options.gather
|
|
&& top.rows <= options.gather_max_rows
|
|
&& total >= options.gather_min_context
|
|
{
|
|
return qsa_eager_gather_rows(top).map(QsaIndexerSelection::RowTokens);
|
|
}
|
|
qsa_eager_dense(top, total).map(QsaIndexerSelection::Dense)
|
|
}
|
|
|
|
/// Eager flash_prefill epilogue: chronological int32 blocks, gathered validity,
|
|
/// then invalid slots replaced by zero. No custom selector substitution.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_eager_blocks(top: QsaEagerTopk) -> Result<QsaSelection, String> {
|
|
let rows = top.rows;
|
|
let width = top.width;
|
|
let count = rows
|
|
.checked_mul(width)
|
|
.filter(|n| *n > 0 && *n <= i32::MAX as u32)
|
|
.ok_or("invalid eager QSA block output shape")?;
|
|
let (sorted, ids) = qsa_eager_sorted_ids(&top)?;
|
|
let indices64 = Buffer::mtplx_bytes(u64::from(count) * 8)?;
|
|
let valid = Buffer::mtplx_bytes(u64::from(count))?;
|
|
let work = if count < 65536 { 1 } else { 2 };
|
|
// int64 output has one element per thread even above the vector threshold.
|
|
dispatch_geometry(
|
|
"v_copyint32int64",
|
|
&[
|
|
tensor(0, &sorted).input(count as usize),
|
|
tensor(1, &indices64).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[count, 1, 1],
|
|
[count.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
let grid = [1, width, rows];
|
|
dispatch_geometry(
|
|
"kernel_qwen_mtplx_gather_axis_bool_idxi64_int_11",
|
|
&[
|
|
tensor(0, &top.valid).input(rows as usize * top.row_stride as usize),
|
|
tensor(1, &indices64).input(count as usize),
|
|
tensor(2, &valid).output(count as usize),
|
|
bytes(3, &rows),
|
|
bytes(4, &u64::from(top.row_stride)),
|
|
bytes(5, &u64::from(width)),
|
|
bytes(6, &1_u64),
|
|
bytes(7, &1_u32),
|
|
bytes(8, &top.row_stride),
|
|
bytes(9, &1_u64),
|
|
bytes(10, &1_u64),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let zero = scalar_buffer(&0_i32.to_le_bytes())?;
|
|
dispatch_geometry(
|
|
&format!(
|
|
"{}{}_Selectint32",
|
|
if count == 1 { "v" } else { "vs" },
|
|
if work == 1 { "" } else { "n" }
|
|
),
|
|
&[
|
|
tensor(0, &valid).input(count as usize),
|
|
tensor(1, &sorted).input(count as usize),
|
|
tensor(2, &zero).input(1),
|
|
tensor(3, &sorted).output(count as usize),
|
|
bytes(4, &count),
|
|
],
|
|
&[],
|
|
[count.div_ceil(work), 1, 1],
|
|
[count.div_ceil(work).min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
let mut temporary = top._temporaries;
|
|
temporary.extend(top.masked_scores);
|
|
temporary.extend([
|
|
ids,
|
|
indices64,
|
|
top.indices,
|
|
top.valid,
|
|
top.positions,
|
|
top.complete_blocks,
|
|
]);
|
|
Ok(QsaSelection {
|
|
outputs: vec![sorted, valid],
|
|
row_widths: vec![(width, 4), (width, 1)],
|
|
row_strides: vec![width; 2],
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_eager_sorted_ids(top: &QsaEagerTopk) -> Result<(Buffer, Buffer), String> {
|
|
let width = top.width;
|
|
let rows = top.rows;
|
|
let ids = qsa_eager_cast_indices(top, "int32", 4)?;
|
|
let sorted = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 4)?;
|
|
let bn = width.div_ceil(4).next_power_of_two().clamp(32, 512);
|
|
let segment = if rows == 1 { i32::MAX as u32 } else { width };
|
|
dispatch(
|
|
&format!("c_block_sort_int32_int32_bn{bn}_tn4"),
|
|
&[
|
|
tensor(0, &ids).input(rows as usize * width as usize),
|
|
tensor(1, &sorted).output(rows as usize * width as usize),
|
|
bytes(2, &width),
|
|
bytes(3, &1_u32),
|
|
bytes(4, &1_u32),
|
|
bytes(5, &segment),
|
|
bytes(6, &segment),
|
|
],
|
|
[1, rows, 1],
|
|
[bn, 1, 1],
|
|
)?;
|
|
Ok((sorted, ids))
|
|
}
|
|
|
|
/// Untiled eager decode epilogues. The dead selected-mask expression in the
|
|
/// Python source is not evaluated by these lazy output branches.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_eager_decode(
|
|
top: QsaEagerTopk,
|
|
total: u32,
|
|
gather: bool,
|
|
) -> Result<(QsaSelection, u32), String> {
|
|
if top.rows != 1
|
|
|| top.pos_start.checked_add(1) != Some(total)
|
|
|| total / 4 != top.row_stride
|
|
|| top.row_stride <= 512
|
|
{
|
|
return Err("invalid active QSA eager decode frontier".into());
|
|
}
|
|
let tail_start = total / 4 * 4;
|
|
let k = top.width;
|
|
let (sorted, cast) = qsa_eager_sorted_ids(&top)?;
|
|
let mut temporary = top._temporaries;
|
|
temporary.extend(top.masked_scores);
|
|
temporary.extend([
|
|
cast,
|
|
top.indices,
|
|
top.valid,
|
|
top.positions,
|
|
top.complete_blocks,
|
|
]);
|
|
let (out, width) = if !gather {
|
|
(sorted, k)
|
|
} else {
|
|
let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
|
|
qsa_vector(
|
|
"vs",
|
|
"Multiplyint32",
|
|
&[
|
|
tensor(0, &sorted).input(k as usize),
|
|
tensor(1, &ratio).input(1),
|
|
tensor(2, &sorted).output(k as usize),
|
|
bytes(3, &k),
|
|
],
|
|
k,
|
|
4,
|
|
)?;
|
|
let ratio_range = Buffer::mtplx_bytes(16)?;
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &0_u32),
|
|
bytes(1, &1_u32),
|
|
tensor(2, &ratio_range).output(4),
|
|
],
|
|
&[],
|
|
[4, 1, 1],
|
|
[4, 1, 1],
|
|
true,
|
|
)?;
|
|
let block_tokens = Buffer::mtplx_bytes(u64::from(k) * 16)?;
|
|
let grid = [4, k, 1];
|
|
dispatch_geometry(
|
|
"g2_Addint32",
|
|
&[
|
|
tensor(0, &sorted).input(k as usize),
|
|
tensor(1, &ratio_range).input(4),
|
|
tensor(2, &block_tokens).output(k as usize * 4),
|
|
bytes(3, &[1_u64, 0]),
|
|
bytes(4, &[0_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let tail_len = total - tail_start;
|
|
let width = k * 4 + tail_len;
|
|
let tail = if tail_len > 0 {
|
|
let tail = Buffer::mtplx_bytes(u64::from(tail_len) * 4)?;
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &tail_start),
|
|
bytes(1, &1_u32),
|
|
tensor(2, &tail).output(tail_len as usize),
|
|
],
|
|
&[],
|
|
[tail_len, 1, 1],
|
|
[tail_len, 1, 1],
|
|
true,
|
|
)?;
|
|
Some(tail)
|
|
} else {
|
|
None
|
|
};
|
|
let out = Buffer::mtplx_bytes(u64::from(width) * 4)?;
|
|
let concurrent = encoder::concurrent_scope();
|
|
qsa_copy_typed(
|
|
Some(&block_tokens),
|
|
&out,
|
|
[1, 1, k * 4],
|
|
[0, 0, 1],
|
|
[0, 0, 1],
|
|
"int32",
|
|
(k * 4) as usize,
|
|
)?;
|
|
if let Some(tail) = tail {
|
|
let dst = out.view(u64::from(k) * 16, u64::from(tail_len) * 4)?;
|
|
qsa_copy_typed(
|
|
Some(&tail),
|
|
&dst,
|
|
[1, 1, tail_len],
|
|
[0, 0, 1],
|
|
[0, 0, 1],
|
|
"int32",
|
|
(tail_len) as usize,
|
|
)?;
|
|
temporary.push(tail);
|
|
}
|
|
drop(concurrent);
|
|
temporary.extend([sorted, ratio, ratio_range, block_tokens]);
|
|
(out, width)
|
|
};
|
|
Ok((
|
|
QsaSelection {
|
|
outputs: vec![out],
|
|
row_widths: vec![(width, 4)],
|
|
row_strides: vec![width],
|
|
_temporaries: temporary,
|
|
},
|
|
tail_start,
|
|
))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_eager_cast_indices(
|
|
top: &QsaEagerTopk,
|
|
dtype: &str,
|
|
bytes_per_item: u32,
|
|
) -> Result<Buffer, String> {
|
|
let count = top.rows * top.width;
|
|
let input_size = (top.rows as usize - 1) * top.index_row_stride as usize + top.width as usize;
|
|
let output = Buffer::mtplx_bytes(u64::from(count) * u64::from(bytes_per_item))?;
|
|
if top.rows == 1 || top.width == top.index_row_stride {
|
|
qsa_vector(
|
|
"v",
|
|
&format!("copyuint32{dtype}"),
|
|
&[
|
|
tensor(0, &top.indices).input(input_size),
|
|
tensor(1, &output).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
count,
|
|
bytes_per_item,
|
|
)?;
|
|
} else {
|
|
let grid = [top.width, top.rows, 1];
|
|
dispatch_geometry(
|
|
&format!("g2_copyuint32{dtype}"),
|
|
&[
|
|
tensor(0, &top.indices).input(input_size),
|
|
tensor(1, &output).output(count as usize),
|
|
bytes(3, &[u64::from(top.index_row_stride), 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
Ok(output)
|
|
}
|
|
|
|
/// _select_eager's rows-gather epilogue. Keeps argsort order (unlike the
|
|
/// chronological flash block contract), then appends each row's visible tail.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_eager_gather_rows(top: QsaEagerTopk) -> Result<QsaSelection, String> {
|
|
let rows = top.rows;
|
|
let k = top.width;
|
|
let count = rows * k;
|
|
let width = k * 4 + 4;
|
|
if rows <= 1
|
|
|| top.row_stride <= 512
|
|
|| rows.checked_mul(width).is_none_or(|n| n > i32::MAX as u32)
|
|
{
|
|
return Err("invalid active QSA gather-rows geometry".into());
|
|
}
|
|
let idx64 = qsa_eager_cast_indices(&top, "int64", 8)?;
|
|
let blk_ok = Buffer::mtplx_bytes(u64::from(count))?;
|
|
let grid = [1, k, rows];
|
|
dispatch_geometry(
|
|
"kernel_qwen_mtplx_gather_axis_bool_idxi64_int_11",
|
|
&[
|
|
tensor(0, &top.valid).input(rows as usize * top.row_stride as usize),
|
|
tensor(1, &idx64).input(count as usize),
|
|
tensor(2, &blk_ok).output(count as usize),
|
|
bytes(3, &rows),
|
|
bytes(4, &u64::from(top.row_stride)),
|
|
bytes(5, &u64::from(k)),
|
|
bytes(6, &1_u64),
|
|
bytes(7, &1_u32),
|
|
bytes(8, &top.row_stride),
|
|
bytes(9, &1_u64),
|
|
bytes(10, &1_u64),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let idx32 = qsa_eager_cast_indices(&top, "int32", 4)?;
|
|
let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
|
|
qsa_vector(
|
|
"vs",
|
|
"Multiplyint32",
|
|
&[
|
|
tensor(0, &idx32).input(count as usize),
|
|
tensor(1, &ratio).input(1),
|
|
tensor(2, &idx32).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
let ratio_range = Buffer::mtplx_bytes(16)?;
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &0_u32),
|
|
bytes(1, &1_u32),
|
|
tensor(2, &ratio_range).output(4),
|
|
],
|
|
&[],
|
|
[4, 1, 1],
|
|
[4, 1, 1],
|
|
true,
|
|
)?;
|
|
let tokens = Buffer::mtplx_bytes(u64::from(count) * 16)?;
|
|
let blocks_ok = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
let grid = [4, count, 1];
|
|
dispatch_geometry(
|
|
"g2_Addint32",
|
|
&[
|
|
tensor(0, &idx32).input(count as usize),
|
|
tensor(1, &ratio_range).input(4),
|
|
tensor(2, &tokens).output(count as usize * 4),
|
|
bytes(3, &[1_u64, 0]),
|
|
bytes(4, &[0_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
"g2_copybool_bool_",
|
|
&[
|
|
tensor(0, &blk_ok).input(count as usize),
|
|
tensor(1, &blocks_ok).output(count as usize * 4),
|
|
bytes(3, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let tail_start = top.complete_blocks;
|
|
let tail_ratio = scalar_buffer(&4_i32.to_le_bytes())?;
|
|
qsa_vector(
|
|
"vs",
|
|
"Multiplyint32",
|
|
&[
|
|
tensor(0, &tail_start).input(rows as usize),
|
|
tensor(1, &tail_ratio).input(1),
|
|
tensor(2, &tail_start).output(rows as usize),
|
|
bytes(3, &rows),
|
|
],
|
|
rows,
|
|
4,
|
|
)?;
|
|
let tail_range = Buffer::mtplx_bytes(16)?;
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &0_u32),
|
|
bytes(1, &1_u32),
|
|
tensor(2, &tail_range).output(4),
|
|
],
|
|
&[],
|
|
[4, 1, 1],
|
|
[4, 1, 1],
|
|
true,
|
|
)?;
|
|
let tail_tokens = Buffer::mtplx_bytes(u64::from(rows) * 16)?;
|
|
let tail_ok = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
|
|
let grid = [4, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_Addint32",
|
|
&[
|
|
tensor(0, &tail_start).input(rows as usize),
|
|
tensor(1, &tail_range).input(4),
|
|
tensor(2, &tail_tokens).output(rows as usize * 4),
|
|
bytes(3, &[1_u64, 0]),
|
|
bytes(4, &[0_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
dispatch_geometry(
|
|
"g2_LessEqualint32",
|
|
&[
|
|
tensor(0, &tail_tokens).input(rows as usize * 4),
|
|
tensor(1, &top.positions).input(rows as usize),
|
|
tensor(2, &tail_ok).output(rows as usize * 4),
|
|
bytes(3, &[4_u64, 1]),
|
|
bytes(4, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let ids = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 4)?;
|
|
let valid = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width))?;
|
|
for (block, tail, out, dtype, b) in [
|
|
(&tokens, &tail_tokens, &ids, "int32", 4_u64),
|
|
(&blocks_ok, &tail_ok, &valid, "bool_", 1),
|
|
] {
|
|
let _concurrent = encoder::concurrent_scope();
|
|
qsa_copy_typed(
|
|
Some(block),
|
|
out,
|
|
[1, rows, k * 4],
|
|
[0, u64::from(k) * 4, 1],
|
|
[0, u64::from(width), 1],
|
|
dtype,
|
|
(rows) as usize * (k * 4) as usize,
|
|
)?;
|
|
let offset = u64::from(k) * 4 * b;
|
|
let dst = out.view(offset, u64::from(rows) * u64::from(width) * b - offset)?;
|
|
qsa_copy_typed(
|
|
Some(tail),
|
|
&dst,
|
|
[1, rows, 4],
|
|
[0, 4, 1],
|
|
[0, u64::from(width), 1],
|
|
dtype,
|
|
(rows) as usize * 4,
|
|
)?;
|
|
}
|
|
let count = rows * width;
|
|
let zero = scalar_buffer(&0_i32.to_le_bytes())?;
|
|
qsa_vector(
|
|
"vs",
|
|
"Selectint32",
|
|
&[
|
|
tensor(0, &valid).input(count as usize),
|
|
tensor(1, &ids).input(count as usize),
|
|
tensor(2, &zero).input(1),
|
|
tensor(3, &ids).output(count as usize),
|
|
bytes(4, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
let mut temporary = top._temporaries;
|
|
temporary.extend(top.masked_scores);
|
|
temporary.extend([
|
|
idx64,
|
|
idx32,
|
|
blk_ok,
|
|
blocks_ok,
|
|
ratio,
|
|
ratio_range,
|
|
tail_ratio,
|
|
tail_range,
|
|
tokens,
|
|
tail_start,
|
|
tail_tokens,
|
|
tail_ok,
|
|
top.indices,
|
|
top.valid,
|
|
top.positions,
|
|
]);
|
|
Ok(QsaSelection {
|
|
outputs: vec![ids, valid],
|
|
row_widths: vec![(width, 4), (width, 1)],
|
|
row_strides: vec![width; 2],
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
/// Remaining dense-mask epilogue of _select_eager, after the active indexer
|
|
/// threshold (N>512). Uses original ScatterAxis/Repeat/Concatenate operations.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_eager_dense(top: QsaEagerTopk, total: u32) -> Result<QsaSelection, String> {
|
|
let rows = top.rows;
|
|
let blocks = top.row_stride;
|
|
if blocks <= 512
|
|
|| total / 4 != blocks
|
|
|| top.pos_start.checked_add(rows) != Some(total)
|
|
|| total > i32::MAX as u32
|
|
|| rows.checked_mul(total).is_none_or(|n| n > i32::MAX as u32)
|
|
{
|
|
return Err("invalid active QSA dense-mask geometry".into());
|
|
}
|
|
let k = top.width;
|
|
let index_count = rows * k;
|
|
let count = rows * blocks;
|
|
let indices64 = qsa_eager_cast_indices(&top, "int64", 8)?;
|
|
let selected = Buffer::mtplx_bytes(u64::from(count))?;
|
|
let zero = scalar_buffer(&[0_u8])?;
|
|
let one = scalar_buffer(&[1_u8])?;
|
|
qsa_vector(
|
|
"s",
|
|
"copybool_bool_",
|
|
&[
|
|
tensor(0, &zero).input(1),
|
|
tensor(1, &selected).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
count,
|
|
1,
|
|
)?;
|
|
let grid = [1, k, rows];
|
|
dispatch_geometry(
|
|
if index_count == 1 {
|
|
"kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_11"
|
|
} else {
|
|
"kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_01"
|
|
},
|
|
&[
|
|
tensor(0, &one).input(1),
|
|
tensor(1, &indices64).input(index_count as usize),
|
|
tensor(2, &selected).output(count as usize),
|
|
bytes(3, &rows),
|
|
bytes(4, &0_u64),
|
|
bytes(5, &u64::from(k)),
|
|
bytes(6, &1_u64),
|
|
bytes(7, &1_u32),
|
|
bytes(8, &blocks),
|
|
bytes(9, &0_u64),
|
|
bytes(10, &1_u64),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
qsa_vector(
|
|
"vv",
|
|
"BitwiseAndbool_",
|
|
&[
|
|
tensor(0, &selected).input(count as usize),
|
|
tensor(1, &top.valid).input(count as usize),
|
|
tensor(2, &selected).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
1,
|
|
)?;
|
|
// repeat inserts a broadcast axis and its final reshape is a General copy.
|
|
let repeated = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
let grid = [4, count, 1];
|
|
dispatch_geometry(
|
|
"g2_copybool_bool_",
|
|
&[
|
|
tensor(0, &selected).input(count as usize),
|
|
tensor(1, &repeated).output(count as usize * 4),
|
|
bytes(3, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
let mut temporary = top._temporaries;
|
|
temporary.extend(top.masked_scores);
|
|
let mask = if total == blocks * 4 {
|
|
repeated
|
|
} else {
|
|
let pad_count = rows * (total - blocks * 4);
|
|
let pad = Buffer::mtplx_bytes(u64::from(pad_count))?;
|
|
let pad_zero = scalar_buffer(&[0_u8])?;
|
|
qsa_vector(
|
|
"s",
|
|
"copybool_bool_",
|
|
&[
|
|
tensor(0, &pad_zero).input(1),
|
|
tensor(1, &pad).output(pad_count as usize),
|
|
bytes(2, &pad_count),
|
|
],
|
|
pad_count,
|
|
1,
|
|
)?;
|
|
let padded = Buffer::mtplx_bytes(u64::from(rows) * u64::from(total))?;
|
|
let concurrent = encoder::concurrent_scope();
|
|
qsa_copy_typed(
|
|
Some(&repeated),
|
|
&padded,
|
|
[1, rows, blocks * 4],
|
|
[0, u64::from(blocks) * 4, 1],
|
|
[0, u64::from(total), 1],
|
|
"bool_",
|
|
(rows) as usize * (blocks * 4) as usize,
|
|
)?;
|
|
let tail = padded.view(
|
|
u64::from(blocks) * 4,
|
|
u64::from(rows) * u64::from(total) - u64::from(blocks) * 4,
|
|
)?;
|
|
qsa_copy_typed(
|
|
Some(&pad),
|
|
&tail,
|
|
[1, rows, total - blocks * 4],
|
|
[0, u64::from(total - blocks * 4), 1],
|
|
[0, u64::from(total), 1],
|
|
"bool_",
|
|
(rows) as usize * (total - blocks * 4) as usize,
|
|
)?;
|
|
drop(concurrent);
|
|
temporary.extend([repeated, pad, pad_zero]);
|
|
padded
|
|
};
|
|
let tpos = Buffer::mtplx_bytes(u64::from(total) * 4)?;
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &0_u32),
|
|
bytes(1, &1_i32),
|
|
tensor(2, &tpos).output(total as usize),
|
|
],
|
|
&[],
|
|
[total, 1, 1],
|
|
[total.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
let tail_start = top.complete_blocks;
|
|
let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
|
|
qsa_vector(
|
|
if rows == 1 { "ss" } else { "vs" },
|
|
"Multiplyint32",
|
|
&[
|
|
tensor(0, &tail_start).input(rows as usize),
|
|
tensor(1, &ratio).input(1),
|
|
tensor(2, &tail_start).output(rows as usize),
|
|
bytes(3, &rows),
|
|
],
|
|
rows,
|
|
4,
|
|
)?;
|
|
let count = rows * total;
|
|
let tail = Buffer::mtplx_bytes(u64::from(count))?;
|
|
let causal = Buffer::mtplx_bytes(u64::from(count))?;
|
|
for (op, rhs, out) in [
|
|
("GreaterEqualint32", &tail_start, &tail),
|
|
("LessEqualint32", &top.positions, &causal),
|
|
] {
|
|
if rows == 1 {
|
|
qsa_vector(
|
|
"vs",
|
|
op,
|
|
&[
|
|
tensor(0, &tpos).input(total as usize),
|
|
tensor(1, rhs).input(rows as usize),
|
|
tensor(2, out).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
} else {
|
|
let grid = [total, rows, 1];
|
|
dispatch_geometry(
|
|
&format!("g2_{op}"),
|
|
&[
|
|
tensor(0, &tpos).input(total as usize),
|
|
tensor(1, rhs).input(rows as usize),
|
|
tensor(2, out).output(count as usize),
|
|
bytes(3, &[0_u64, 1]),
|
|
bytes(4, &[1_u64, 0]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
}
|
|
qsa_vector(
|
|
"vv",
|
|
"BitwiseOrbool_",
|
|
&[
|
|
tensor(0, &mask).input(count as usize),
|
|
tensor(1, &tail).input(count as usize),
|
|
tensor(2, &mask).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
1,
|
|
)?;
|
|
qsa_vector(
|
|
"vv",
|
|
"BitwiseAndbool_",
|
|
&[
|
|
tensor(0, &mask).input(count as usize),
|
|
tensor(1, &causal).input(count as usize),
|
|
tensor(2, &mask).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
1,
|
|
)?;
|
|
temporary.extend([
|
|
selected,
|
|
zero,
|
|
one,
|
|
ratio,
|
|
indices64,
|
|
tpos,
|
|
tail_start,
|
|
tail,
|
|
causal,
|
|
top.indices,
|
|
top.valid,
|
|
top.positions,
|
|
]);
|
|
Ok(QsaSelection {
|
|
outputs: vec![mask],
|
|
row_widths: vec![(total, 1)],
|
|
row_strides: vec![total],
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
/// Original TensorOps score producer, not a replacement QMM implementation.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prefill_mpp_scores(
|
|
[q, pooled]: [&Buffer; 2],
|
|
[rows, blocks]: [u32; 2],
|
|
q_strides: [u64; 3],
|
|
pooled_strides: [u64; 2],
|
|
) -> Result<Buffer, String> {
|
|
let tiles = rows
|
|
.div_ceil(16)
|
|
.checked_mul(blocks.div_ceil(32))
|
|
.and_then(|n| n.checked_mul(128));
|
|
if rows == 0
|
|
|| blocks == 0
|
|
|| rows > i32::MAX as u32 / 512
|
|
|| blocks > i32::MAX as u32
|
|
|| tiles.is_none()
|
|
{
|
|
return Err("invalid QSA MPP score geometry".into());
|
|
}
|
|
for (b, shape, strides) in [
|
|
(q, [rows, 4, 128], q_strides),
|
|
(
|
|
pooled,
|
|
[1, blocks, 128],
|
|
[0, pooled_strides[0], pooled_strides[1]],
|
|
),
|
|
] {
|
|
let span = shape
|
|
.into_iter()
|
|
.zip(strides)
|
|
.try_fold(1_u64, |n, (d, s)| {
|
|
u64::from(d - 1)
|
|
.checked_mul(s)
|
|
.and_then(|v| n.checked_add(v))
|
|
})
|
|
.and_then(|n| n.checked_mul(2));
|
|
if span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n) {
|
|
return Err("QSA MPP input outside backing".into());
|
|
}
|
|
}
|
|
use array::{Array, Dtype};
|
|
let inputs = [
|
|
buffer_array(
|
|
q,
|
|
&[1, rows as i32, 4, 128],
|
|
&[0, q_strides[0], q_strides[1], q_strides[2]],
|
|
Dtype::BF16,
|
|
)?,
|
|
buffer_array(
|
|
pooled,
|
|
&[1, blocks as i32, 128],
|
|
&[0, pooled_strides[0], pooled_strides[1]],
|
|
Dtype::BF16,
|
|
)?,
|
|
];
|
|
let scores = Array::unallocated(&[rows as i32, blocks as i32], Dtype::F32)?;
|
|
qsa_prefill::mpp_evaluate(&inputs, &scores)?;
|
|
let buffer = scores.buffer().view(0, scores.layout().nbytes() as u64)?;
|
|
Ok(buffer)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl QsaSelectPlan {
|
|
fn prefill_plan(self) -> qsa_prefill::TopK {
|
|
qsa_prefill::TopK {
|
|
blocks: self.blocks,
|
|
topk: self.topk,
|
|
ratio: 4,
|
|
dense_width: if self.mode == QsaSelectMode::DenseMask {
|
|
self.dense_width
|
|
} else {
|
|
0
|
|
},
|
|
mode: self.mode,
|
|
}
|
|
}
|
|
pub(super) fn prefill_source(self) -> Result<String, String> {
|
|
self.prefill_plan().prefill_source()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prefill_topk(
|
|
scores: &Buffer,
|
|
strides: [u64; 2],
|
|
plan: QsaSelectPlan,
|
|
) -> Result<QsaSelection, String> {
|
|
if plan.rows == 0
|
|
|| plan.blocks == 0
|
|
|| plan.rows > i32::MAX as u32 / 512
|
|
|| !(1..=512).contains(&plan.topk)
|
|
|| plan.total > i32::MAX as u32
|
|
|| plan
|
|
.pos
|
|
.checked_add(plan.rows)
|
|
.is_none_or(|n| n > plan.total)
|
|
|| plan.logical > plan.blocks
|
|
|| plan.logical != plan.total / 4
|
|
|| plan
|
|
.rows
|
|
.checked_mul(plan.blocks)
|
|
.is_none_or(|n| n > i32::MAX as u32)
|
|
|| (plan.mode == QsaSelectMode::DenseMask && plan.dense_width < plan.total)
|
|
{
|
|
return Err("invalid QSA prefill selector geometry".into());
|
|
}
|
|
let positions = [
|
|
scalar_buffer(&plan.pos.to_le_bytes())?,
|
|
scalar_buffer(&plan.total.to_le_bytes())?,
|
|
scalar_buffer(&plan.logical.to_le_bytes())?,
|
|
];
|
|
let mut out = qsa_prefill_topk_device(
|
|
scores,
|
|
strides,
|
|
plan,
|
|
[&positions[0], &positions[1], &positions[2]],
|
|
)?;
|
|
out._temporaries.extend(positions);
|
|
Ok(out)
|
|
}
|
|
|
|
/// Original prefill top-k with moving frontiers kept on the GPU.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prefill_topk_device(
|
|
scores: &Buffer,
|
|
strides: [u64; 2],
|
|
plan: QsaSelectPlan,
|
|
[pos, total, logical]: [&Buffer; 3],
|
|
) -> Result<QsaSelection, String> {
|
|
if plan.rows == 0
|
|
|| plan.blocks == 0
|
|
|| plan.rows > i32::MAX as u32 / 512
|
|
|| !(1..=512).contains(&plan.topk)
|
|
|| plan
|
|
.rows
|
|
.checked_mul(plan.blocks)
|
|
.is_none_or(|n| n > i32::MAX as u32)
|
|
|| [pos, total, logical]
|
|
.into_iter()
|
|
.any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
|
|
{
|
|
return Err("invalid QSA device prefill geometry or leaves".into());
|
|
}
|
|
let span = u64::from(plan.rows - 1)
|
|
.checked_mul(strides[0])
|
|
.and_then(|n| {
|
|
u64::from(plan.blocks - 1)
|
|
.checked_mul(strides[1])
|
|
.and_then(|m| n.checked_add(m))
|
|
})
|
|
.and_then(|n| n.checked_add(1))
|
|
.and_then(|n| n.checked_mul(4));
|
|
if span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(scores.raw()) } < n) {
|
|
return Err("QSA score view outside backing".into());
|
|
}
|
|
use array::{Array, Dtype};
|
|
let inputs = [
|
|
buffer_array(
|
|
scores,
|
|
&[plan.rows as i32, plan.blocks as i32],
|
|
&strides,
|
|
Dtype::F32,
|
|
)?,
|
|
Array::new(&[1], Dtype::I32, pos.view(0, 4)?)?,
|
|
Array::new(&[1], Dtype::I32, total.view(0, 4)?)?,
|
|
Array::new(&[1], Dtype::I32, logical.view(0, 4)?)?,
|
|
];
|
|
let kernel = plan.prefill_plan();
|
|
let outputs = kernel
|
|
.layouts(plan.rows as i32)?
|
|
.into_iter()
|
|
.map(|l| Array::unallocated(l.shape(), l.dtype()))
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
qsa_prefill::topk_evaluate(&inputs, &outputs, kernel)?;
|
|
let outputs = outputs
|
|
.iter()
|
|
.map(|a| a.buffer().view(0, a.layout().nbytes() as u64))
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
let widths = plan.widths();
|
|
Ok(QsaSelection {
|
|
outputs,
|
|
row_strides: widths.iter().map(|(w, _)| *w).collect(),
|
|
row_widths: widths,
|
|
_temporaries: Vec::new(),
|
|
})
|
|
}
|
|
|
|
/// qsa_indexer_prefill_metal on M5: BF16 pairs select TensorOps, otherwise the
|
|
/// general FP32 expression. Dispatch failures never switch producers.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prefill(
|
|
input: [&Buffer; 2],
|
|
q_strides: [u64; 3],
|
|
pooled_strides: [u64; 2],
|
|
plan: QsaSelectPlan,
|
|
workspace: u64,
|
|
) -> Result<QsaSelection, String> {
|
|
if plan.rows <= 1
|
|
|| plan.rows > i32::MAX as u32 / 512
|
|
|| plan
|
|
.pos
|
|
.checked_add(plan.rows)
|
|
.is_none_or(|n| n > plan.total)
|
|
|| plan.logical > plan.blocks
|
|
|| plan.logical != plan.total / 4
|
|
|| !(1..=512).contains(&plan.topk)
|
|
{
|
|
return Err("invalid QSA prefill input".into());
|
|
}
|
|
let mpp = plan.input_f32 == [false; 2];
|
|
let chunk_rows = qsa_prefill_chunk_rows(plan.rows, plan.blocks, workspace, 32, mpp)?;
|
|
let mut pooled_s = [0, pooled_strides[0], pooled_strides[1]];
|
|
let pooled_cast = if !mpp && !plan.input_f32[1] {
|
|
Some(qsa_cast_f32(
|
|
input[1],
|
|
[1, plan.blocks, 128],
|
|
&mut pooled_s,
|
|
)?)
|
|
} else {
|
|
None
|
|
};
|
|
let pooled = pooled_cast.as_ref().unwrap_or(input[1]);
|
|
let mut chunks = Vec::new();
|
|
for start in (0..plan.rows).step_by(chunk_rows as usize) {
|
|
let rows = chunk_rows.min(plan.rows - start);
|
|
let offset = u64::from(start)
|
|
.checked_mul(q_strides[0])
|
|
.and_then(|n| n.checked_mul(if plan.input_f32[0] { 4 } else { 2 }))
|
|
.ok_or("QSA prefill offset overflow")?;
|
|
let bytes = unsafe { ds4_gpu_tensor_bytes(input[0].raw()) }
|
|
.checked_sub(offset)
|
|
.ok_or("QSA prefill slice outside backing")?;
|
|
let q = input[0].view(offset, bytes)?;
|
|
let (scores, scratch) = if mpp {
|
|
(
|
|
qsa_prefill_mpp_scores(
|
|
[&q, pooled],
|
|
[rows, plan.blocks],
|
|
q_strides,
|
|
pooled_strides,
|
|
)?,
|
|
Vec::new(),
|
|
)
|
|
} else {
|
|
qsa_prefill_f32_scores(
|
|
[&q, pooled],
|
|
[rows, plan.blocks],
|
|
q_strides,
|
|
[pooled_s[2], pooled_s[1]],
|
|
!plan.input_f32[0],
|
|
plan.tf32,
|
|
)?
|
|
};
|
|
let mut selected = qsa_prefill_topk(
|
|
&scores,
|
|
[u64::from(plan.blocks), 1],
|
|
QsaSelectPlan {
|
|
rows,
|
|
pos: plan.pos + start,
|
|
..plan
|
|
},
|
|
)?;
|
|
selected._temporaries.push(scores);
|
|
selected._temporaries.extend(scratch);
|
|
chunks.push(selected);
|
|
}
|
|
let mut selected = qsa_concat_selection(chunks, plan, chunk_rows)?;
|
|
selected._temporaries.extend(pooled_cast);
|
|
Ok(selected)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prefill_cached(
|
|
cache: &QsaCache,
|
|
prepared_q: &Buffer,
|
|
rows: u32,
|
|
mode: QsaSelectMode,
|
|
workspace: u64,
|
|
) -> Result<QsaSelection, String> {
|
|
let total = cache
|
|
.offset
|
|
.checked_add(rows)
|
|
.ok_or("QSA prefill cache position overflow")?;
|
|
let logical = total / 4;
|
|
if cache.pooled_len < logical {
|
|
return Err("QSA prefill needs complete pooled keys".into());
|
|
}
|
|
let pooled = cache
|
|
.pooled
|
|
.as_ref()
|
|
.ok_or("QSA prefill needs pooled backing")?;
|
|
// The large-prefill branch receives _extend_pooled's logical slice, not
|
|
// the reserved cache capacity used by the legacy decode selector.
|
|
qsa_prefill(
|
|
[prepared_q, &pooled.buffer()],
|
|
[512, 128, 1],
|
|
[128, 1],
|
|
QsaSelectPlan {
|
|
rows,
|
|
blocks: logical,
|
|
topk: 512,
|
|
pos: cache.offset,
|
|
total,
|
|
logical,
|
|
dense_width: total,
|
|
mode,
|
|
tf32: true,
|
|
input_f32: [false; 2],
|
|
},
|
|
workspace,
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_select_cached(
|
|
cache: &QsaCache,
|
|
prepared_q: &Buffer,
|
|
rows: u32,
|
|
mode: QsaSelectMode,
|
|
tf32: bool,
|
|
) -> Result<QsaSelection, String> {
|
|
let total = cache
|
|
.offset
|
|
.checked_add(rows)
|
|
.ok_or("QSA cache position overflow")?;
|
|
if cache.pooled_len < total / 4 {
|
|
return Err("QSA selector needs complete pooled keys".into());
|
|
}
|
|
let pooled = cache
|
|
.pooled
|
|
.as_ref()
|
|
.ok_or("QSA selector needs pooled backing")?;
|
|
qsa_select_chunked(
|
|
[prepared_q, &pooled.buffer()],
|
|
[512, 128, 1],
|
|
[128, 1],
|
|
QsaSelectPlan {
|
|
rows,
|
|
blocks: cache.pooled_cap,
|
|
topk: 512,
|
|
pos: cache.offset,
|
|
total,
|
|
logical: total / 4,
|
|
dense_width: 0,
|
|
mode,
|
|
tf32,
|
|
input_f32: [false; 2],
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Original eager _rope_cos_sin + _apply_partial_rope, for normalized BF16
|
|
/// QSA queries (H4) or pooled keys (H1). Every cast/concat stays explicit.
|
|
#[cfg(test)]
|
|
fn qsa_eager_rope(
|
|
normalized: Buffer,
|
|
freq: &Buffer,
|
|
[rows, heads]: [u32; 2],
|
|
[position, position_scale]: [i32; 2],
|
|
attention_scale: f32,
|
|
) -> Result<(Buffer, Vec<Buffer>), String> {
|
|
let count = rows * heads * 64;
|
|
let positions = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
|
|
dispatch_geometry(
|
|
"arangeint32",
|
|
&[
|
|
bytes(0, &position),
|
|
bytes(1, &1_i32),
|
|
tensor(2, &positions).output(rows as usize),
|
|
],
|
|
&[],
|
|
[rows, 1, 1],
|
|
[rows.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
if position_scale != 1 {
|
|
let scale = scalar_buffer(&position_scale.to_le_bytes())?;
|
|
qsa_vector(
|
|
if rows == 1 { "ss" } else { "vs" },
|
|
"Multiplyint32",
|
|
&[
|
|
tensor(0, &positions).input(rows as usize),
|
|
tensor(1, &scale).input(1),
|
|
tensor(2, &positions).output(rows as usize),
|
|
bytes(3, &rows),
|
|
],
|
|
rows,
|
|
4,
|
|
)?;
|
|
}
|
|
let fp_positions = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
|
|
qsa_vector(
|
|
"v",
|
|
"copyint32float32",
|
|
&[
|
|
tensor(0, &positions).input(rows as usize),
|
|
tensor(1, &fp_positions).output(rows as usize),
|
|
bytes(2, &rows),
|
|
],
|
|
rows,
|
|
4,
|
|
)?;
|
|
let angles = Buffer::mtplx_bytes(u64::from(rows) * 32 * 4)?;
|
|
if rows == 1 {
|
|
qsa_vector(
|
|
"sv",
|
|
"Multiplyfloat32",
|
|
&[
|
|
tensor(0, &fp_positions).input(rows as usize),
|
|
tensor(1, freq).input(32),
|
|
tensor(2, &angles).output(rows as usize * 32),
|
|
bytes(3, &32_u32),
|
|
],
|
|
32,
|
|
4,
|
|
)?;
|
|
} else {
|
|
let grid = [32, rows, 1];
|
|
dispatch_geometry(
|
|
"g2_Multiplyfloat32",
|
|
&[
|
|
tensor(0, &fp_positions).input(rows as usize),
|
|
tensor(1, freq).input(32),
|
|
tensor(2, &angles).output(rows as usize * 32),
|
|
bytes(3, &[1_u64, 0]),
|
|
bytes(4, &[0_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let emb = Buffer::mtplx_bytes(u64::from(rows) * 64 * 4)?;
|
|
let concurrent = encoder::concurrent_scope();
|
|
for offset in [0, 32 * 4] {
|
|
let dst = emb.view(offset, (u64::from(rows) * 64 - 32) * 4)?;
|
|
qsa_copy_typed(
|
|
Some(&angles),
|
|
&dst,
|
|
[1, rows, 32],
|
|
[0, 32, 1],
|
|
[0, 64, 1],
|
|
"float32",
|
|
(rows) as usize * 32,
|
|
)?;
|
|
}
|
|
drop(concurrent);
|
|
let cosine = Buffer::mtplx_bytes(u64::from(rows) * 64 * 4)?;
|
|
let sine = Buffer::mtplx_bytes(u64::from(rows) * 64 * 4)?;
|
|
for (op, out) in [("Cosfloat32float32", &cosine), ("Sinfloat32float32", &sine)] {
|
|
qsa_vector(
|
|
"v",
|
|
op,
|
|
&[
|
|
tensor(0, &emb).input(rows as usize * 64),
|
|
tensor(1, out).output(rows as usize * 64),
|
|
bytes(2, &(rows * 64)),
|
|
],
|
|
rows * 64,
|
|
4,
|
|
)?;
|
|
if attention_scale != 1.0 {
|
|
let scale = scalar_buffer(&attention_scale.to_le_bytes())?;
|
|
qsa_vector(
|
|
"vs",
|
|
"Multiplyfloat32",
|
|
&[
|
|
tensor(0, out).input(rows as usize * 64),
|
|
tensor(1, &scale).input(1),
|
|
tensor(2, out).output(rows as usize * 64),
|
|
bytes(3, &(rows * 64)),
|
|
],
|
|
rows * 64,
|
|
4,
|
|
)?;
|
|
}
|
|
}
|
|
let norm_bytes = u64::from(rows) * u64::from(heads) * 128 * 2;
|
|
let x2 = normalized.view(32 * 2, norm_bytes - 32 * 2)?;
|
|
let negative = Buffer::mtplx_bytes(u64::from(rows) * u64::from(heads) * 32 * 2)?;
|
|
if rows * heads == 1 {
|
|
qsa_vector(
|
|
"v",
|
|
"Negativebfloat16bfloat16",
|
|
&[
|
|
tensor(0, &x2).input(((rows * heads - 1) * 128 + 32) as usize),
|
|
tensor(1, &negative).output(rows as usize * heads as usize * 32),
|
|
bytes(2, &32_u32),
|
|
],
|
|
32,
|
|
2,
|
|
)?;
|
|
} else {
|
|
let grid = [32, rows * heads, 1];
|
|
dispatch_geometry(
|
|
"gn1_Negativebfloat16bfloat16",
|
|
&[
|
|
tensor(0, &x2).input(((rows * heads - 1) * 128 + 32) as usize),
|
|
tensor(1, &negative).output(rows as usize * heads as usize * 32),
|
|
bytes(2, &[rows * heads, 32]),
|
|
bytes(3, &[128_u64, 1]),
|
|
bytes(4, &2_u32),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
let rotated = Buffer::mtplx_bytes(u64::from(count) * 2)?;
|
|
let concurrent = encoder::concurrent_scope();
|
|
for (src, stride, offset) in [(&negative, 32_u64, 0), (&normalized, 128, 32 * 2)] {
|
|
let dst = rotated.view(offset, u64::from(count) * 2 - offset)?;
|
|
qsa_copy_typed(
|
|
Some(src),
|
|
&dst,
|
|
[1, rows * heads, 32],
|
|
[0, stride, 1],
|
|
[0, 64, 1],
|
|
"bfloat16",
|
|
(rows * heads) as usize * 32,
|
|
)?;
|
|
}
|
|
drop(concurrent);
|
|
let mut norm_strides = [u64::from(heads) * 128, 128, 1];
|
|
let x_float = qsa_cast_f32(&normalized, [rows, heads, 64], &mut norm_strides)?;
|
|
let rotated_float = Buffer::mtplx_bytes(u64::from(count) * 4)?;
|
|
cast_bf16_f32(&rotated, &rotated_float, count, true)?;
|
|
for (x, table) in [(&x_float, &cosine), (&rotated_float, &sine)] {
|
|
if heads == 1 {
|
|
qsa_vector(
|
|
"vv",
|
|
"Multiplyfloat32",
|
|
&[
|
|
tensor(0, x).input(count as usize),
|
|
tensor(1, table).input(rows as usize * 64),
|
|
tensor(2, x).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
} else if rows == 1 {
|
|
let grid = [64, heads, 1];
|
|
dispatch_geometry(
|
|
"g2_Multiplyfloat32",
|
|
&[
|
|
tensor(0, x).input(count as usize),
|
|
tensor(1, table).input(rows as usize * 64),
|
|
tensor(2, x).output(count as usize),
|
|
bytes(3, &[64_u64, 1]),
|
|
bytes(4, &[0_u64, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
} else {
|
|
let grid = [64, heads, rows];
|
|
dispatch_geometry(
|
|
"g3_Multiplyfloat32",
|
|
&[
|
|
tensor(0, x).input(count as usize),
|
|
tensor(1, table).input(rows as usize * 64),
|
|
tensor(2, x).output(count as usize),
|
|
bytes(3, &[u64::from(heads) * 64, 64, 1]),
|
|
bytes(4, &[64_u64, 0, 1]),
|
|
],
|
|
&[],
|
|
grid,
|
|
block_dims(grid),
|
|
true,
|
|
)?;
|
|
}
|
|
}
|
|
qsa_vector(
|
|
"vv",
|
|
"Addfloat32",
|
|
&[
|
|
tensor(0, &x_float).input(count as usize),
|
|
tensor(1, &rotated_float).input(count as usize),
|
|
tensor(2, &x_float).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
let rope = Buffer::mtplx_bytes(u64::from(count) * 2)?;
|
|
cast_bf16_f32(&x_float, &rope, count, false)?;
|
|
let output = Buffer::mtplx_bytes(norm_bytes)?;
|
|
let pass = normalized.view(64 * 2, norm_bytes - 64 * 2)?;
|
|
let concurrent = encoder::concurrent_scope();
|
|
for (src, stride, offset) in [(&rope, 64_u64, 0), (&pass, 128, 64 * 2)] {
|
|
let dst = output.view(offset, norm_bytes - offset)?;
|
|
qsa_copy_typed(
|
|
Some(src),
|
|
&dst,
|
|
[1, rows * heads, 64],
|
|
[0, stride, 1],
|
|
[0, 128, 1],
|
|
"bfloat16",
|
|
(rows * heads) as usize * 64,
|
|
)?;
|
|
}
|
|
drop(concurrent);
|
|
Ok((
|
|
output,
|
|
vec![
|
|
normalized,
|
|
positions,
|
|
fp_positions,
|
|
angles,
|
|
emb,
|
|
cosine,
|
|
sine,
|
|
negative,
|
|
rotated,
|
|
x_float,
|
|
rotated_float,
|
|
rope,
|
|
],
|
|
))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prepare_queries_eager(
|
|
[input, norm, freq]: [&Buffer; 3],
|
|
rows: u32,
|
|
strides: [u64; 3],
|
|
position: i32,
|
|
) -> Result<(Buffer, Vec<Buffer>), String> {
|
|
let span = u64::from(rows.saturating_sub(1))
|
|
.checked_mul(strides[0])
|
|
.and_then(|n| n.checked_add(3_u64.checked_mul(strides[1])?))
|
|
.and_then(|n| n.checked_add(127_u64.checked_mul(strides[2])?))
|
|
.and_then(|n| n.checked_add(1))
|
|
.and_then(|n| n.checked_mul(2));
|
|
if rows == 0
|
|
|| rows > i32::MAX as u32 / 512
|
|
|| position.checked_add(rows as i32 - 1).is_none()
|
|
|| span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
|
|
|| unsafe { ds4_gpu_tensor_bytes(freq.raw()) } < 128
|
|
{
|
|
return Err("invalid eager QSA query preparation".into());
|
|
}
|
|
// RMSNorm::eval_gpu materializes noncontiguous projection slices, then
|
|
// donates that copy to its output. A contiguous caller input is borrowed.
|
|
let copied = if strides[2] != 1 || strides[1] != 128 || (rows > 1 && strides[0] != 512) {
|
|
Some(qsa_copy_general(
|
|
input,
|
|
[rows, 4, 128],
|
|
strides,
|
|
"bfloat16",
|
|
"bfloat16",
|
|
2,
|
|
)?)
|
|
} else {
|
|
None
|
|
};
|
|
let normalized = if let Some(copy) = copied {
|
|
rms_norm_bf16(©, Some(norm), ©, [rows * 4, 128], 1e-6)?;
|
|
copy
|
|
} else {
|
|
let out = Buffer::mtplx_bytes(u64::from(rows) * 512 * 2)?;
|
|
rms_norm_bf16(input, Some(norm), &out, [rows * 4, 128], 1e-6)?;
|
|
out
|
|
};
|
|
qsa_eager_rope(normalized, freq, [rows, 4], [position, 1], 1.0)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_pool_keys_eager(
|
|
[input, norm, freq]: [&Buffer; 3],
|
|
blocks: u32,
|
|
[row_stride, dim_stride]: [u64; 2],
|
|
start: i32,
|
|
) -> Result<(Buffer, Vec<Buffer>), String> {
|
|
if blocks == 0
|
|
|| blocks > i32::MAX as u32 / 512
|
|
|| start
|
|
.checked_add(blocks as i32 - 1)
|
|
.and_then(|n| n.checked_mul(4))
|
|
.is_none()
|
|
|| unsafe { ds4_gpu_tensor_bytes(freq.raw()) } < 128
|
|
{
|
|
return Err("invalid eager QSA pooling".into());
|
|
}
|
|
let mut strides = [
|
|
row_stride
|
|
.checked_mul(4)
|
|
.ok_or("QSA pool stride overflow")?,
|
|
row_stride,
|
|
dim_stride,
|
|
];
|
|
let fresh = qsa_cast_f32(input, [blocks, 4, 128], &mut strides)?;
|
|
if strides != [512, 128, 1] {
|
|
return Err("unsupported eager pool physical layout".into());
|
|
}
|
|
let sum = Buffer::mtplx_bytes(u64::from(blocks) * 128 * 4)?;
|
|
let ndim = u32::from(blocks > 1);
|
|
dispatch(
|
|
"col_reduce_small_1_reduce_sumfloat32",
|
|
&[
|
|
tensor(0, &fresh).input(blocks as usize * 512),
|
|
tensor(1, &sum).output(blocks as usize * 128),
|
|
bytes(2, &4_u64),
|
|
bytes(3, &128_u64),
|
|
bytes(4, &if blocks > 1 { blocks } else { 0 }),
|
|
bytes(5, &if blocks > 1 { 512_u64 } else { 0 }),
|
|
bytes(6, &ndim),
|
|
bytes(7, &4_u32),
|
|
bytes(8, &128_u64),
|
|
bytes(9, &1_u32),
|
|
bytes(10, &1_u64),
|
|
],
|
|
[1, blocks, 1],
|
|
[32, 4, 1],
|
|
)?;
|
|
let count = blocks * 128;
|
|
let mean_scale = scalar_buffer(&0.25_f32.to_le_bytes())?;
|
|
qsa_vector(
|
|
"vs",
|
|
"Multiplyfloat32",
|
|
&[
|
|
tensor(0, &sum).input(count as usize),
|
|
tensor(1, &mean_scale).input(1),
|
|
tensor(2, &sum).output(count as usize),
|
|
bytes(3, &count),
|
|
],
|
|
count,
|
|
4,
|
|
)?;
|
|
let pooled = Buffer::mtplx_bytes(u64::from(count) * 2)?;
|
|
cast_bf16_f32(&sum, &pooled, count, false)?;
|
|
rms_norm_bf16(&pooled, Some(norm), &pooled, [blocks, 128], 1e-6)?;
|
|
let (out, mut temporary) = qsa_eager_rope(pooled, freq, [blocks, 1], [start, 4], 1.0)?;
|
|
temporary.extend([fresh, sum]);
|
|
Ok((out, temporary))
|
|
}
|
|
|
|
/// Original installed QSA indexer geometry: H4, D128, rotary64, ratio4,
|
|
/// epsilon1e-6 and attention scaling1. Strides are physical BF16 elements.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prepare_queries(
|
|
input: [&Buffer; 3],
|
|
rows: u32,
|
|
strides: [u64; 3],
|
|
pos_start: i32,
|
|
) -> Result<Buffer, String> {
|
|
if rows == 0 || rows > i32::MAX as u32 / 512 || pos_start.checked_add(rows as i32 - 1).is_none()
|
|
{
|
|
return Err("invalid QSA query position".into());
|
|
}
|
|
let pos = scalar_buffer(&pos_start.to_le_bytes())?;
|
|
qsa_prepare_queries_bound(input, rows, strides, tensor(6, &pos).input(1))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_prepare_queries_device(
|
|
input: [&Buffer; 3],
|
|
rows: u32,
|
|
strides: [u64; 3],
|
|
pos: &Buffer,
|
|
) -> Result<Buffer, String> {
|
|
if unsafe { ds4_gpu_tensor_bytes(pos.raw()) } < 4 {
|
|
return Err("invalid QSA device query position".into());
|
|
}
|
|
qsa_prepare_queries_bound(input, rows, strides, tensor(6, pos).input(1))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_prepare_queries_bound(
|
|
[input, norm, inv_freq]: [&Buffer; 3],
|
|
rows: u32,
|
|
[row_stride, head_stride, dim_stride]: [u64; 3],
|
|
pos_start: Binding<'_>,
|
|
) -> Result<Buffer, String> {
|
|
let end = u64::from(rows.saturating_sub(1))
|
|
.checked_mul(row_stride)
|
|
.and_then(|n| n.checked_add(3_u64.checked_mul(head_stride)?))
|
|
.and_then(|n| n.checked_add(127_u64.checked_mul(dim_stride)?))
|
|
.and_then(|n| n.checked_add(1))
|
|
.and_then(|n| n.checked_mul(2));
|
|
if rows == 0
|
|
|| rows > i32::MAX as u32 / 512
|
|
|| dim_stride == 0
|
|
|| end.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
|
|
|| unsafe { ds4_gpu_tensor_bytes(norm.raw()) } < 128 * 2
|
|
|| unsafe { ds4_gpu_tensor_bytes(inv_freq.raw()) } < 32 * 4
|
|
{
|
|
return Err("invalid MTPLX QSA query preparation buffers or shape".into());
|
|
}
|
|
let output = Buffer::mtplx_bytes(u64::from(rows) * 512 * 2)?;
|
|
dispatch(
|
|
"kernel_qwen_mtplx_qsa_prepare_q_bf16",
|
|
&[
|
|
tensor(0, input).input((end.unwrap() / 2) as usize),
|
|
bytes(1, &[0_u64, row_stride, head_stride, dim_stride]),
|
|
tensor(2, norm).input(128),
|
|
bytes(3, &1_u64),
|
|
tensor(4, inv_freq).input(32),
|
|
bytes(5, &1_u64),
|
|
pos_start,
|
|
tensor(7, &output).output(rows as usize * 512),
|
|
],
|
|
[rows * 4, 1, 1],
|
|
[32, 1, 1],
|
|
)?;
|
|
Ok(output)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_pool_keys(
|
|
inputs: [&Buffer; 3],
|
|
blocks: u32,
|
|
strides: [u64; 2],
|
|
block_start: i32,
|
|
) -> Result<Buffer, String> {
|
|
if blocks == 0
|
|
|| blocks > i32::MAX as u32 / 4
|
|
|| block_start
|
|
.checked_add(blocks as i32 - 1)
|
|
.and_then(|n| n.checked_mul(4))
|
|
.is_none()
|
|
{
|
|
return Err("invalid MTPLX QSA pooled-key position".into());
|
|
}
|
|
let start = scalar_buffer(&block_start.to_le_bytes())?;
|
|
qsa_pool_keys_bound(inputs, blocks, strides, tensor(6, &start).input(1))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_pool_keys_bound(
|
|
[input, norm, inv_freq]: [&Buffer; 3],
|
|
blocks: u32,
|
|
[row_stride, dim_stride]: [u64; 2],
|
|
block_start: Binding<'_>,
|
|
) -> Result<Buffer, String> {
|
|
let tokens = u64::from(blocks) * 4;
|
|
let end = tokens
|
|
.saturating_sub(1)
|
|
.checked_mul(row_stride)
|
|
.and_then(|n| n.checked_add(127_u64.checked_mul(dim_stride)?))
|
|
.and_then(|n| n.checked_add(1))
|
|
.and_then(|n| n.checked_mul(2));
|
|
if blocks == 0
|
|
|| tokens > i32::MAX as u64
|
|
|| dim_stride == 0
|
|
|| end.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
|
|
|| unsafe { ds4_gpu_tensor_bytes(norm.raw()) } < 128 * 2
|
|
|| unsafe { ds4_gpu_tensor_bytes(inv_freq.raw()) } < 32 * 4
|
|
{
|
|
return Err("invalid MTPLX QSA pooled-key preparation buffers or shape".into());
|
|
}
|
|
let output = Buffer::mtplx_bytes(u64::from(blocks) * 128 * 2)?;
|
|
dispatch(
|
|
"kernel_qwen_mtplx_qsa_pool_k_bf16",
|
|
&[
|
|
tensor(0, input).input((end.unwrap() / 2) as usize),
|
|
bytes(1, &[0_u64, row_stride, dim_stride]),
|
|
tensor(2, norm).input(128),
|
|
bytes(3, &1_u64),
|
|
tensor(4, inv_freq).input(32),
|
|
bytes(5, &1_u64),
|
|
block_start,
|
|
tensor(7, &output).output(blocks as usize * 128),
|
|
],
|
|
[blocks, 1, 1],
|
|
[32, 1, 1],
|
|
)?;
|
|
Ok(output)
|
|
}
|
|
|
|
fn qsa_compiled_clamp_name(max_new: u32, max_start: u32) -> String {
|
|
format!("kernel_qwen_mtplx_qsa_clamp_n{max_new}_p{max_start}")
|
|
}
|
|
|
|
/// Original compiler's three constant-CSE layouts, with only shape constants
|
|
/// and the exported symbol specialized. Dynamic frontier leaves stay inputs.
|
|
pub(super) fn qsa_compiled_clamp_source(max_new: u32, max_start: u32) -> Result<String, String> {
|
|
if max_new == 0 || max_new > i32::MAX as u32 || max_start > i32::MAX as u32 {
|
|
return Err("invalid compiled QSA clamp geometry".into());
|
|
}
|
|
let bank: serde_json::Value = serde_json::from_str(include_str!(
|
|
"../../../metal/mtplx-qsa-compiled-scalars.json"
|
|
))
|
|
.map_err(|e| e.to_string())?;
|
|
let kind = if max_start == 0 {
|
|
"zero"
|
|
} else if max_start == max_new {
|
|
"equal"
|
|
} else {
|
|
"distinct"
|
|
};
|
|
let variant = &bank["clamp"][kind];
|
|
let body = variant["kernels"]
|
|
.as_str()
|
|
.ok_or("missing original compiled QSA clamp")?;
|
|
let original_name = body
|
|
.split_once("host_name(\"")
|
|
.and_then(|(_, s)| s.split_once('"'))
|
|
.ok_or("invalid compiled QSA symbol")?
|
|
.0;
|
|
let mut body = body.replace(original_name, &qsa_compiled_clamp_name(max_new, max_start));
|
|
body = body.replace(
|
|
&format!("tmp_D = static_cast<int32_t>({});", variant["max_new"]),
|
|
&format!("tmp_D = static_cast<int32_t>({max_new});"),
|
|
);
|
|
if kind == "distinct" {
|
|
body = body.replace(
|
|
&format!("tmp_F = static_cast<int32_t>({});", variant["max_start"]),
|
|
&format!("tmp_F = static_cast<int32_t>({max_start});"),
|
|
);
|
|
}
|
|
Ok(body)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_compiled_frontiers(
|
|
[logical, pooled_len]: [&Buffer; 2],
|
|
[max_new, max_start]: [u32; 2],
|
|
) -> Result<[Buffer; 2], String> {
|
|
let block_start = qsa_compiled_block_start([logical, pooled_len], [max_new, max_start])?;
|
|
let raw_start = qsa_compiled_raw_start(&block_start)?;
|
|
Ok([block_start, raw_start])
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_compiled_block_start(
|
|
[logical, pooled_len]: [&Buffer; 2],
|
|
[max_new, max_start]: [u32; 2],
|
|
) -> Result<Buffer, String> {
|
|
qsa_compile::Scalar::Clamp { max_new, max_start }.buffer(&[logical, pooled_len])
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_compiled_raw_start(block_start: &Buffer) -> Result<Buffer, String> {
|
|
qsa_compile::Scalar::Multiply(4).buffer(&[block_start])
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_retained_bf16_copy(input: &Buffer, count: u32) -> Result<Buffer, String> {
|
|
let output = Buffer::mtplx_bytes(u64::from(count) * 2)?;
|
|
let work = if count < 65536 { 1 } else { 4 };
|
|
let threads = count.div_ceil(work);
|
|
dispatch_geometry(
|
|
if work == 1 {
|
|
"v_copybfloat16bfloat16"
|
|
} else {
|
|
"vn_copybfloat16bfloat16"
|
|
},
|
|
&[
|
|
tensor(0, input).input(count as usize),
|
|
tensor(1, &output).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[threads, 1, 1],
|
|
[threads.min(1024), 1, 1],
|
|
true,
|
|
)?;
|
|
Ok(output)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct QsaCompiledCacheWindow {
|
|
pub(super) raw: Buffer,
|
|
pub(super) pooled: Buffer,
|
|
pub(super) block_start: Buffer,
|
|
pub(super) raw_start: Buffer,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// Connected cache-maintenance subgraph of QSACompiledIndexerCore.from_qk_rows.
|
|
/// This entry retains the old leaves, so the reference's non-donation path is
|
|
/// required. Graph-bank ownership/donation and BFS scheduling are separate.
|
|
/// Position/logical/frontier are validated int32[1] device leaves; no readback.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_compiled_cache_window(
|
|
[qk, raw, pooled, norm, freq]: [&Buffer; 5],
|
|
[pos, logical, pooled_len]: [&Buffer; 3],
|
|
rows: u32,
|
|
qk_stride: u64,
|
|
[raw_cap, pooled_cap]: [u32; 2],
|
|
) -> Result<QsaCompiledCacheWindow, String> {
|
|
let max_new = rows.div_ceil(4);
|
|
let max_start = (raw_cap / 4).min(pooled_cap).checked_sub(max_new);
|
|
if rows == 0
|
|
|| rows > raw_cap
|
|
|| qk_stride < 640
|
|
|| max_start.is_none()
|
|
|| [raw_cap, pooled_cap]
|
|
.into_iter()
|
|
.any(|n| n < 256 || !n.is_power_of_two() || n > i32::MAX as u32 / 512)
|
|
|| [pos, logical, pooled_len]
|
|
.into_iter()
|
|
.any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
|
|
|| unsafe { ds4_gpu_tensor_bytes(raw.raw()) } < u64::from(raw_cap) * 256
|
|
|| unsafe { ds4_gpu_tensor_bytes(pooled.raw()) } < u64::from(pooled_cap) * 256
|
|
|| u64::from(rows - 1)
|
|
.checked_mul(qk_stride)
|
|
.and_then(|n| n.checked_add(640))
|
|
.and_then(|n| n.checked_mul(2))
|
|
.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(qk.raw()) } < n)
|
|
{
|
|
return Err("invalid compiled QSA cache-window inputs or geometry".into());
|
|
}
|
|
let raw_next = qsa_retained_bf16_copy(raw, raw_cap * 128)?;
|
|
let k = qk.view(512 * 2, (u64::from(rows - 1) * qk_stride + 128) * 2)?;
|
|
let mut temporary = qsa_dynamic_copy_rows(
|
|
[&k, &raw_next],
|
|
[qk_stride, 128],
|
|
[rows, 128],
|
|
[None, Some(pos)],
|
|
[
|
|
((rows - 1) as u64 * qk_stride + 128) as usize,
|
|
raw_cap as usize * 128,
|
|
],
|
|
)?;
|
|
let [block_start, raw_start] =
|
|
qsa_compiled_frontiers([logical, pooled_len], [max_new, max_start.unwrap()])?;
|
|
let window = Buffer::mtplx_bytes(u64::from(max_new) * 4 * 128 * 2)?;
|
|
temporary.extend(qsa_dynamic_copy_rows(
|
|
[&raw_next, &window],
|
|
[128, 128],
|
|
[max_new * 4, 128],
|
|
[Some(&raw_start), None],
|
|
[raw_cap as usize * 128, max_new as usize * 512],
|
|
)?);
|
|
let pool_window = qsa_pool_keys_bound(
|
|
[&window, norm, freq],
|
|
max_new,
|
|
[128, 1],
|
|
tensor(6, &block_start).input(1),
|
|
)?;
|
|
let pooled_next = qsa_retained_bf16_copy(pooled, pooled_cap * 128)?;
|
|
temporary.extend(qsa_dynamic_copy_rows(
|
|
[&pool_window, &pooled_next],
|
|
[128, 128],
|
|
[max_new, 128],
|
|
[None, Some(&block_start)],
|
|
[max_new as usize * 128, pooled_cap as usize * 128],
|
|
)?);
|
|
temporary.extend([k, window, pool_window]);
|
|
Ok(QsaCompiledCacheWindow {
|
|
raw: raw_next,
|
|
pooled: pooled_next,
|
|
block_start,
|
|
raw_start,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct QsaProjected {
|
|
pub(super) q: Buffer,
|
|
pub(super) raw_keys: Buffer,
|
|
_temporary: Vec<Buffer>,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
|
pub(super) enum QsaCompiledMode {
|
|
UpdateOnly,
|
|
Blocks,
|
|
PrefillBlocks,
|
|
RowTokens,
|
|
DenseMask,
|
|
}
|
|
|
|
pub(super) fn qsa_compiled_row_position_source(start: u32) -> Result<String, String> {
|
|
if start > i32::MAX as u32 {
|
|
return Err("invalid compiled QSA row offset".into());
|
|
}
|
|
let bank: serde_json::Value = serde_json::from_str(include_str!(
|
|
"../../../tests/fixtures/mtplx-qsa-select-jit.json"
|
|
))
|
|
.map_err(|e| e.to_string())?;
|
|
let body = bank[3]["kernels"]
|
|
.as_str()
|
|
.ok_or("missing original compiled row offset")?;
|
|
let symbol = body
|
|
.split_once("host_name(\"")
|
|
.and_then(|(_, s)| s.split_once('"'))
|
|
.ok_or("invalid compiled row offset symbol")?
|
|
.0;
|
|
Ok(body
|
|
.replace(symbol, &format!("kernel_qwen_mtplx_qsa_pos_add_{start}"))
|
|
.replace(
|
|
"tmp_C = static_cast<int32_t>(3);",
|
|
&format!("tmp_C = static_cast<int32_t>({start});"),
|
|
))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_compiled_row_position(pos: &Buffer, start: u32) -> Result<Buffer, String> {
|
|
qsa_compile::Scalar::Add(start).buffer(&[pos])
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct QsaIndexerWeights<'a> {
|
|
pub(super) projection: [&'a Buffer; 3],
|
|
pub(super) quant: [u32; 2],
|
|
pub(super) q_norm: &'a Buffer,
|
|
pub(super) k_norm: &'a Buffer,
|
|
pub(super) inv_freq: &'a Buffer,
|
|
}
|
|
|
|
/// Resolved MTPLX switches for both indexer routes. Prefill's
|
|
/// device-dependent AUTO value must be resolved before entering this path.
|
|
pub(super) struct QsaIndexerOptions {
|
|
pub(super) fused: bool,
|
|
pub(super) compiled: bool,
|
|
pub(super) prefill: bool,
|
|
pub(super) phase_prefill: bool,
|
|
pub(super) prefill_min_rows: u32,
|
|
pub(super) prefill_compile_rows: u32,
|
|
pub(super) prefill_min_context: u32,
|
|
pub(super) prefill_workspace: u64,
|
|
pub(super) tile_rows: u32,
|
|
pub(super) flash: bool,
|
|
pub(super) gather_decode: bool,
|
|
pub(super) gather: bool,
|
|
pub(super) gather_max_rows: u32,
|
|
pub(super) gather_min_context: u32,
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) tf32: bool,
|
|
}
|
|
|
|
impl QsaIndexerOptions {
|
|
fn large_prefill(&self, rows: u32, pos: u32) -> bool {
|
|
self.prefill
|
|
&& self.phase_prefill
|
|
&& rows >= self.prefill_min_rows
|
|
&& pos >= self.prefill_min_context
|
|
}
|
|
|
|
/// Original _compiled_mode; total is the post-update token frontier.
|
|
pub(super) fn compiled_mode(
|
|
&self,
|
|
rows: u32,
|
|
total: u32,
|
|
ratio: u32,
|
|
topk: u32,
|
|
) -> QsaCompiledMode {
|
|
if total / ratio <= topk {
|
|
QsaCompiledMode::UpdateOnly
|
|
} else if rows == 1 && self.flash {
|
|
QsaCompiledMode::Blocks
|
|
} else if rows != 1 && self.large_prefill(rows, total.saturating_sub(rows)) {
|
|
QsaCompiledMode::PrefillBlocks
|
|
} else if rows != 1
|
|
&& self.gather
|
|
&& rows <= self.gather_max_rows
|
|
&& total >= self.gather_min_context
|
|
{
|
|
QsaCompiledMode::RowTokens
|
|
} else {
|
|
QsaCompiledMode::DenseMask
|
|
}
|
|
}
|
|
|
|
/// Cache/phase portion of _compiled_route_supported, BEFORE capacity
|
|
/// bucketing. QsaCache fixes B1/BF16/D128/ratio4. This does not validate a
|
|
/// source tensor, projection/norm dtype or device; those remain entry-point
|
|
/// obligations. In particular it is not a generic tensor eligibility test.
|
|
#[cfg(test)]
|
|
pub(super) fn compiled_cache_mode(
|
|
&self,
|
|
cache: &QsaCache,
|
|
rows: u32,
|
|
pos: u32,
|
|
) -> Option<QsaCompiledMode> {
|
|
let total = pos.checked_add(rows)?;
|
|
if !self.fused || !self.compiled || rows == 0 || pos != cache.offset {
|
|
return None;
|
|
}
|
|
let mode = self.compiled_mode(rows, total, 4, 512);
|
|
if rows != 1 {
|
|
if mode == QsaCompiledMode::UpdateOnly && self.phase_prefill {
|
|
return None;
|
|
}
|
|
if rows >= self.prefill_min_rows
|
|
&& (!matches!(
|
|
mode,
|
|
QsaCompiledMode::PrefillBlocks | QsaCompiledMode::UpdateOnly
|
|
) || rows != self.prefill_compile_rows)
|
|
{
|
|
return None;
|
|
}
|
|
}
|
|
if pos > 0 && (cache.raw.is_none() || cache.raw_cap < pos) {
|
|
return None;
|
|
}
|
|
if cache.pooled.is_none() {
|
|
if cache.pooled_len != 0 {
|
|
return None;
|
|
}
|
|
} else if cache.pooled_len > cache.pooled_cap {
|
|
return None;
|
|
}
|
|
let logical = total / 4;
|
|
if logical - cache.pooled_len.min(logical) > rows.div_ceil(4) {
|
|
return None;
|
|
}
|
|
if rows == 1
|
|
&& !matches!(mode, QsaCompiledMode::UpdateOnly | QsaCompiledMode::Blocks)
|
|
&& self.gather_decode
|
|
{
|
|
return None;
|
|
}
|
|
Some(mode)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct QsaIndexed {
|
|
pub(super) selection: Option<QsaIndexerSelection>,
|
|
pub(super) lane: &'static str,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
/// Connected _call_rows host flow for the installed, typed B1/BF16 indexer.
|
|
/// Replays the parameter-bound indexer graph with its exclusive cache state.
|
|
/// Shared descriptors/Data retain the reference's copying path.
|
|
/// Eligibility rejection selects eager; dispatch errors are returned, not hidden.
|
|
#[allow(clippy::too_many_arguments)] // Explicit model inputs, cache and evaluator state.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_indexer_rows(
|
|
weights: QsaIndexerWeights<'_>,
|
|
hidden: &Buffer,
|
|
supplied_qk: Option<(&Buffer, u64)>,
|
|
rows: u32,
|
|
cache: &mut QsaCache,
|
|
options: &QsaIndexerOptions,
|
|
commands: Option<&mut Commands>,
|
|
graph: &mut QsaGraphBank,
|
|
) -> Result<QsaIndexed, String> {
|
|
let total = qsa_validate_indexer_options(rows, cache.offset, options)?;
|
|
let Some(mode) = options.compiled_cache_mode(cache, rows, cache.offset) else {
|
|
return qsa_indexer_rows_uncompiled(
|
|
weights,
|
|
hidden,
|
|
supplied_qk,
|
|
rows,
|
|
cache,
|
|
options,
|
|
commands,
|
|
);
|
|
};
|
|
cache.ensure_compiled_backings(rows)?;
|
|
let logical = total / 4;
|
|
let frontiers = [cache.offset, total, logical, cache.pooled_len.min(logical)]
|
|
.map(|value| scalar_buffer(&(value as i32).to_le_bytes()));
|
|
let [pos, end, blocks, frontier] = frontiers;
|
|
let frontiers = [pos?, end?, blocks?, frontier?];
|
|
let budget = if mode == QsaCompiledMode::PrefillBlocks {
|
|
options.prefill_workspace
|
|
} else {
|
|
32 * 1024 * 1024
|
|
};
|
|
let geometry = [rows, cache.raw_cap, cache.pooled_cap];
|
|
let (updated, selected) = graph.run(
|
|
weights,
|
|
supplied_qk.map_or(hidden, |(q, _)| q),
|
|
QsaCacheInput::State([cache.raw.as_mut().unwrap(), cache.pooled.as_mut().unwrap()]),
|
|
frontiers.each_ref(),
|
|
geometry,
|
|
supplied_qk.map(|(_, stride)| stride),
|
|
mode,
|
|
budget,
|
|
options.tf32,
|
|
)?;
|
|
// Commit leaves/frontier without waiting or reading GPU scalars. Attention
|
|
// alone advances KV.offset. A former eager mirror is now invalid.
|
|
cache.raw = Some(QsaArray::new(updated.raw, 0)?);
|
|
cache.pooled = Some(QsaArray::new(updated.pooled, 0)?);
|
|
cache.pooled_len = logical;
|
|
cache.mirror = None;
|
|
cache.mirror_cap = 0;
|
|
cache.mirror_strides = [0; 2];
|
|
let mut temporary = updated._temporaries;
|
|
temporary.extend(frontiers);
|
|
let selection = if mode == QsaCompiledMode::UpdateOnly {
|
|
None
|
|
} else {
|
|
let mut out = selected.ok_or("compiled QSA omitted its selection")?;
|
|
Some(match mode {
|
|
QsaCompiledMode::Blocks => {
|
|
out._temporaries.extend(out.outputs.drain(1..));
|
|
out.row_widths.truncate(1);
|
|
out.row_strides.truncate(1);
|
|
QsaIndexerSelection::FlashBlocks(out, total / 4 * 4)
|
|
}
|
|
QsaCompiledMode::PrefillBlocks => {
|
|
out._temporaries.extend(out.outputs.drain(2..));
|
|
out.row_widths.truncate(2);
|
|
out.row_strides.truncate(2);
|
|
QsaIndexerSelection::PrefillBlocks(out)
|
|
}
|
|
QsaCompiledMode::RowTokens => QsaIndexerSelection::RowTokens(out),
|
|
QsaCompiledMode::DenseMask => {
|
|
out.row_widths[0] = (total, 1);
|
|
QsaIndexerSelection::Dense(out)
|
|
}
|
|
QsaCompiledMode::UpdateOnly => unreachable!(),
|
|
})
|
|
};
|
|
Ok(QsaIndexed {
|
|
selection,
|
|
lane: "compiled",
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn qsa_validate_indexer_options(
|
|
rows: u32,
|
|
pos: u32,
|
|
options: &QsaIndexerOptions,
|
|
) -> Result<u32, String> {
|
|
let total = pos
|
|
.checked_add(rows)
|
|
.filter(|n| rows > 0 && *n <= i32::MAX as u32 / 512)
|
|
.ok_or("invalid QSA indexer frontier")?;
|
|
if options.prefill_min_rows < 2
|
|
|| options.prefill_compile_rows < 2
|
|
|| options.prefill_min_context < 2049
|
|
|| options.prefill_workspace < 1024 * 1024
|
|
|| options.gather_max_rows < 2
|
|
{
|
|
return Err("unresolved or invalid QSA indexer options".into());
|
|
}
|
|
Ok(total)
|
|
}
|
|
|
|
/// _call_rows after compiled-route rejection, for the installed B1 BF16 model.
|
|
/// Does not replace an eligible compiled route. Attention advances KV.offset
|
|
/// later; this call only writes the indexer's raw/pool state.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_indexer_rows_uncompiled(
|
|
weights: QsaIndexerWeights<'_>,
|
|
hidden: &Buffer,
|
|
supplied_qk: Option<(&Buffer, u64)>,
|
|
rows: u32,
|
|
cache: &mut QsaCache,
|
|
options: &QsaIndexerOptions,
|
|
commands: Option<&mut Commands>,
|
|
) -> Result<QsaIndexed, String> {
|
|
let total = qsa_validate_indexer_options(rows, cache.offset, options)?;
|
|
let mut temporary = Vec::new();
|
|
let mut projected = None;
|
|
let (qk, stride) = if let Some((qk, stride)) = supplied_qk {
|
|
let span = u64::from(rows - 1)
|
|
.checked_mul(stride)
|
|
.and_then(|n| n.checked_add(640))
|
|
.and_then(|n| n.checked_mul(2));
|
|
if stride < 640 || span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(qk.raw()) } < n) {
|
|
return Err("invalid supplied QSA projection view".into());
|
|
}
|
|
(qk, stride)
|
|
} else {
|
|
let output = Buffer::mtplx_bytes(u64::from(rows) * 640 * 2)?;
|
|
let [w, s, b] = weights.projection;
|
|
temporary.extend(affine_bf16(
|
|
[w, s, b, hidden, &output],
|
|
[rows, 640, 2560],
|
|
weights.quant[0],
|
|
weights.quant[1],
|
|
true,
|
|
)?);
|
|
projected = Some(output);
|
|
(projected.as_ref().unwrap(), 640)
|
|
};
|
|
// In the dense==sparse regime q preparation is a dead lazy expression.
|
|
// Projection/raw writes and completed-key pooling still belong to cache.
|
|
let query = if total / 4 > 512 {
|
|
if options.fused {
|
|
Some(qsa_prepare_queries(
|
|
[qk, weights.q_norm, weights.inv_freq],
|
|
rows,
|
|
[stride, 128, 1],
|
|
cache.offset as i32,
|
|
)?)
|
|
} else {
|
|
let (q, scratch) = qsa_prepare_queries_eager(
|
|
[qk, weights.q_norm, weights.inv_freq],
|
|
rows,
|
|
[stride, 128, 1],
|
|
cache.offset as i32,
|
|
)?;
|
|
temporary.extend(scratch);
|
|
Some(q)
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
let raw_len = (u64::from(rows - 1) * stride + 128) * 2;
|
|
let raw = qk.view(512 * 2, raw_len)?;
|
|
cache.write_raw(&raw, rows, stride)?;
|
|
temporary.extend(cache.extend_pooled(
|
|
total,
|
|
weights.k_norm,
|
|
weights.inv_freq,
|
|
!options.fused,
|
|
)?);
|
|
temporary.extend(projected);
|
|
let Some(query) = query else {
|
|
return Ok(QsaIndexed {
|
|
selection: None,
|
|
lane: "update_only",
|
|
_temporaries: temporary,
|
|
});
|
|
};
|
|
let decode = rows == 1;
|
|
let large_prefill = options.large_prefill(rows, cache.offset);
|
|
// All prepared B1/H4/D128 BF16 arrays meet the original M5 eligibility.
|
|
let (selection, lane) = if !decode && large_prefill {
|
|
let mut out = qsa_prefill_cached(
|
|
cache,
|
|
&query,
|
|
rows,
|
|
QsaSelectMode::Blocks,
|
|
options.prefill_workspace,
|
|
)?;
|
|
out._temporaries.extend(out.outputs.drain(2..));
|
|
out.row_widths.truncate(2);
|
|
out.row_strides.truncate(2);
|
|
(QsaIndexerSelection::PrefillBlocks(out), "metal_prefill")
|
|
} else if options.fused
|
|
&& (decode || rows < options.prefill_min_rows)
|
|
&& !(decode && options.gather_decode)
|
|
{
|
|
if decode && options.flash {
|
|
let mut out =
|
|
qsa_select_cached(cache, &query, rows, QsaSelectMode::Blocks, options.tf32)?;
|
|
out._temporaries.extend(out.outputs.drain(1..));
|
|
out.row_widths.truncate(1);
|
|
out.row_strides.truncate(1);
|
|
(
|
|
QsaIndexerSelection::FlashBlocks(out, total / 4 * 4),
|
|
"legacy_fused",
|
|
)
|
|
} else if !decode
|
|
&& options.gather
|
|
&& rows <= options.gather_max_rows
|
|
&& total >= options.gather_min_context
|
|
{
|
|
(
|
|
QsaIndexerSelection::RowTokens(qsa_select_cached(
|
|
cache,
|
|
&query,
|
|
rows,
|
|
QsaSelectMode::RowTokens,
|
|
options.tf32,
|
|
)?),
|
|
"legacy_fused",
|
|
)
|
|
} else {
|
|
(
|
|
QsaIndexerSelection::Dense(qsa_select_cached(
|
|
cache,
|
|
&query,
|
|
rows,
|
|
QsaSelectMode::DenseMask,
|
|
options.tf32,
|
|
)?),
|
|
"legacy_fused",
|
|
)
|
|
}
|
|
} else {
|
|
let pos = cache.offset;
|
|
let (pooled, strides) = cache.pooled_f32_view(total / 4)?;
|
|
let top = qsa_eager_topk(
|
|
[&query, pooled],
|
|
[512, 128, 1],
|
|
strides,
|
|
QsaSelectPlan {
|
|
rows,
|
|
blocks: total / 4,
|
|
topk: 512,
|
|
pos,
|
|
total,
|
|
logical: total / 4,
|
|
dense_width: total,
|
|
mode: QsaSelectMode::Blocks,
|
|
tf32: options.tf32,
|
|
input_f32: [false, true],
|
|
},
|
|
options.tile_rows,
|
|
commands,
|
|
)?;
|
|
(
|
|
qsa_eager_output(
|
|
top,
|
|
total,
|
|
QsaEagerOptions {
|
|
large_prefill,
|
|
flash: options.flash,
|
|
gather_decode: options.gather_decode,
|
|
gather: options.gather,
|
|
gather_max_rows: options.gather_max_rows,
|
|
gather_min_context: options.gather_min_context,
|
|
},
|
|
)?,
|
|
"eager",
|
|
)
|
|
};
|
|
temporary.push(query);
|
|
Ok(QsaIndexed {
|
|
selection: Some(selection),
|
|
lane,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
/// Shared-input indexer projection followed by the selected original fused or
|
|
/// eager preparation. Raw keys retain the projection stride for cache.write_raw.
|
|
#[cfg(test)]
|
|
pub(super) fn qsa_project_prepare(
|
|
pack: [&Buffer; 3],
|
|
[norm, inv_freq]: [&Buffer; 2],
|
|
input: &Buffer,
|
|
rows: u32,
|
|
quant: [u32; 2],
|
|
pos_start: i32,
|
|
eager: bool,
|
|
) -> Result<QsaProjected, String> {
|
|
if rows == 0 || rows > i32::MAX as u32 / 640 {
|
|
return Err("invalid MTPLX QSA projection rows".into());
|
|
}
|
|
let projected = Buffer::mtplx_bytes(u64::from(rows) * 640 * 2)?;
|
|
let [w, s, b] = pack;
|
|
let mut temporary = affine_bf16(
|
|
[w, s, b, input, &projected],
|
|
[rows, 640, 2560],
|
|
quant[0],
|
|
quant[1],
|
|
true,
|
|
)?;
|
|
let q = if eager {
|
|
let (q, scratch) = qsa_prepare_queries_eager(
|
|
[&projected, norm, inv_freq],
|
|
rows,
|
|
[640, 128, 1],
|
|
pos_start,
|
|
)?;
|
|
temporary.extend(scratch);
|
|
q
|
|
} else {
|
|
qsa_prepare_queries([&projected, norm, inv_freq], rows, [640, 128, 1], pos_start)?
|
|
};
|
|
let raw_keys = projected.view(512 * 2, (u64::from(rows) * 640 - 512) * 2)?;
|
|
temporary.push(projected);
|
|
Ok(QsaProjected {
|
|
q,
|
|
raw_keys,
|
|
_temporary: temporary,
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) enum GdnInputWeights<'a> {
|
|
Fused {
|
|
pack: [&'a Buffer; 3],
|
|
quant: [u32; 2],
|
|
},
|
|
/// Original qkv, z, b, a module order, each with its own quantization.
|
|
Separate {
|
|
packs: [[&'a Buffer; 3]; 4],
|
|
quant: [[u32; 2]; 4],
|
|
},
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct GdnWeights<'a> {
|
|
pub(super) input: GdnInputWeights<'a>,
|
|
pub(super) output: [&'a Buffer; 3],
|
|
pub(super) output_quant: [u32; 2], // bits, group
|
|
pub(super) conv: &'a Buffer,
|
|
pub(super) a_log: &'a Buffer,
|
|
pub(super) dt_bias: &'a Buffer,
|
|
pub(super) norm: &'a Buffer,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
struct GdnProjected {
|
|
qkv: Buffer,
|
|
z: Buffer,
|
|
b: Buffer,
|
|
a: Buffer,
|
|
strides: [u32; 3], // qkv, z, a/b
|
|
temporary: Vec<Buffer>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn gdn_project(
|
|
weights: &GdnInputWeights<'_>,
|
|
input: &Buffer,
|
|
rows: u32,
|
|
) -> Result<GdnProjected, String> {
|
|
let count = u64::from(rows) * 16480;
|
|
if rows == 0 || count > i32::MAX as u64 {
|
|
return Err("invalid MTPLX GDN projection rows".into());
|
|
}
|
|
let mut temporary = Vec::new();
|
|
match weights {
|
|
GdnInputWeights::Fused {
|
|
pack: [w, s, b],
|
|
quant: [bits, group],
|
|
} => {
|
|
let projected = Buffer::mtplx_bytes(count * 2)?;
|
|
temporary.extend(affine_bf16(
|
|
[w, s, b, input, &projected],
|
|
[rows, 16480, 2560],
|
|
*bits,
|
|
*group,
|
|
true,
|
|
)?);
|
|
let qkv = projected.view(0, (count - 6240) * 2)?;
|
|
let z = projected.view(10240 * 2, (count - 10336) * 2)?;
|
|
let b = projected.view(16384 * 2, (count - 16432) * 2)?;
|
|
let a = projected.view(16432 * 2, (count - 16432) * 2)?;
|
|
temporary.push(projected);
|
|
Ok(GdnProjected {
|
|
qkv,
|
|
z,
|
|
b,
|
|
a,
|
|
strides: [16480; 3],
|
|
temporary,
|
|
})
|
|
}
|
|
GdnInputWeights::Separate { packs, quant } => {
|
|
let mut project = |index: usize, width| {
|
|
let [w, s, b] = packs[index];
|
|
let [bits, group] = quant[index];
|
|
let output = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 2)?;
|
|
temporary.extend(affine_bf16(
|
|
[w, s, b, input, &output],
|
|
[rows, width, 2560],
|
|
bits,
|
|
group,
|
|
true,
|
|
)?);
|
|
Ok::<_, String>(output)
|
|
};
|
|
// Do not sort these by width or coalesce independently quantized modules.
|
|
let qkv = project(0, 10240)?;
|
|
let z = project(1, 6144)?;
|
|
let b = project(2, 48)?;
|
|
let a = project(3, 48)?;
|
|
Ok(GdnProjected {
|
|
qkv,
|
|
z,
|
|
b,
|
|
a,
|
|
strides: [10240, 6144, 48],
|
|
temporary,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct GdnDecodeOutput {
|
|
pub(super) hidden: Buffer,
|
|
pub(super) conv: Buffer,
|
|
pub(super) delta: Buffer,
|
|
// Like the reference encoder's temporaries, live through command completion.
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) struct GdnStagedOutput {
|
|
pub(super) hidden: Buffer,
|
|
pub(super) prepared: ConvNormOutput,
|
|
pub(super) recurrent: DeltaUpdateOutput,
|
|
_norm: Option<GatedNormOutput>,
|
|
_temporaries: Vec<Buffer>,
|
|
}
|
|
|
|
// Keep the fixed GPU-handle bundle inline, without an allocation per layer.
|
|
#[allow(clippy::large_enum_variant)]
|
|
#[cfg(test)]
|
|
pub(super) enum GdnForwardOutput {
|
|
Fused(GdnDecodeOutput),
|
|
Staged(GdnStagedOutput),
|
|
}
|
|
|
|
/// GatedDeltaNet's inference branch order for the installed B=1 geometry.
|
|
/// Flags: fused step, conv-single, conv-verify, fused output, verify capture.
|
|
/// State bindings: BF16 conv, FP32 delta, optional INT32 valid length.
|
|
#[cfg(test)]
|
|
pub(super) fn gdn_forward(
|
|
weights: &GdnWeights<'_>,
|
|
input: &Buffer,
|
|
rows: u32,
|
|
cache: [Option<&Buffer>; 3],
|
|
mask: Option<&Buffer>,
|
|
flags: [bool; 5],
|
|
has_cache: bool,
|
|
) -> Result<GdnForwardOutput, String> {
|
|
let [conv, delta, length] = if has_cache { cache } else { [None; 3] };
|
|
if rows == 1
|
|
&& has_cache
|
|
&& flags[0]
|
|
&& !flags[4]
|
|
&& mask.is_none()
|
|
&& length.is_none()
|
|
&& let Some(delta) = delta
|
|
{
|
|
gdn_decode(weights, input, conv, delta).map(GdnForwardOutput::Fused)
|
|
} else {
|
|
gdn_staged(
|
|
weights,
|
|
input,
|
|
rows,
|
|
[conv, delta, length],
|
|
mask,
|
|
[flags[1], flags[2], flags[3]],
|
|
has_cache,
|
|
)
|
|
.map(GdnForwardOutput::Staged)
|
|
}
|
|
}
|
|
|
|
/// Connected cold/prefill/verify forward with either input projection layout.
|
|
/// Fused-step selection remains the caller's earlier branch; these are the
|
|
/// conv-single, conv-verify and output-fusion flags from the actual model.
|
|
#[cfg(test)]
|
|
pub(super) fn gdn_staged(
|
|
weights: &GdnWeights<'_>,
|
|
input: &Buffer,
|
|
rows: u32,
|
|
[conv, delta, length]: [Option<&Buffer>; 3],
|
|
mask: Option<&Buffer>,
|
|
flags: [bool; 3],
|
|
has_cache: bool,
|
|
) -> Result<GdnStagedOutput, String> {
|
|
let GdnProjected {
|
|
qkv,
|
|
z,
|
|
b,
|
|
a,
|
|
strides,
|
|
mut temporary,
|
|
} = gdn_project(&weights.input, input, rows)?;
|
|
let prepared = gdn_conv_norm(
|
|
&qkv,
|
|
weights.conv,
|
|
[rows, strides[0]],
|
|
[conv, length],
|
|
mask,
|
|
[flags[0], flags[1]],
|
|
has_cache,
|
|
)?;
|
|
let recurrent = gated_delta_update(
|
|
[&prepared.q, &prepared.k, &prepared.v],
|
|
[&a, &b, weights.a_log, weights.dt_bias],
|
|
[rows, strides[2]],
|
|
delta,
|
|
mask,
|
|
)?;
|
|
let hidden = Buffer::mtplx_bytes(u64::from(rows) * 2560 * 2)?;
|
|
let norm = if rows == 1
|
|
&& flags[2]
|
|
&& weights.output_quant[0] == 4
|
|
&& [32, 64].contains(&weights.output_quant[1])
|
|
{
|
|
let [w, s, b] = weights.output;
|
|
dispatch(
|
|
&format!(
|
|
"kernel_qwen_mtplx_gdn_out_fused_gs{}_bf16",
|
|
weights.output_quant[1]
|
|
),
|
|
&[
|
|
tensor(0, &recurrent.hidden).input(6144),
|
|
tensor(1, &z).input(6144),
|
|
tensor(2, weights.norm).input(128),
|
|
tensor(3, w).input(2560 * 6144 / 8),
|
|
tensor(4, s).input(2560 * 6144 / weights.output_quant[1] as usize),
|
|
tensor(5, b).input(2560 * 6144 / weights.output_quant[1] as usize),
|
|
tensor(6, &hidden).output(2560),
|
|
],
|
|
[2560_u32.div_ceil(32), 1, 1],
|
|
[1024, 1, 1],
|
|
)?;
|
|
None
|
|
} else {
|
|
let norm = sigmoid_rms_norm(
|
|
&recurrent.hidden,
|
|
weights.norm,
|
|
Some(&z),
|
|
[rows, 48, 128, strides[1]],
|
|
1e-6,
|
|
)?;
|
|
let [w, s, b] = weights.output;
|
|
temporary.extend(affine_bf16(
|
|
[w, s, b, &norm.hidden, &hidden],
|
|
[rows, 2560, 6144],
|
|
weights.output_quant[0],
|
|
weights.output_quant[1],
|
|
true,
|
|
)?);
|
|
Some(norm)
|
|
};
|
|
temporary.extend([qkv, z, b, a]);
|
|
Ok(GdnStagedOutput {
|
|
hidden,
|
|
prepared,
|
|
recurrent,
|
|
_norm: norm,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|
|
|
|
/// The reference's live-state, unmasked, single-row fused-step branch.
|
|
/// Caller checks the model branch gate; cold/prefill/verify use different graphs.
|
|
/// Keep this result alive until the enclosing command batch completes.
|
|
#[cfg(test)]
|
|
pub(super) fn gdn_decode(
|
|
weights: &GdnWeights<'_>,
|
|
input: &Buffer,
|
|
conv: Option<&Buffer>,
|
|
delta: &Buffer,
|
|
) -> Result<GdnDecodeOutput, String> {
|
|
for (buffer, size) in [
|
|
(delta, 48 * 128 * 128 * 4),
|
|
(weights.conv, 10240 * 4 * 2),
|
|
(weights.a_log, 48 * 2),
|
|
(weights.dt_bias, 48 * 2),
|
|
(weights.norm, 128 * 2),
|
|
] {
|
|
if unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < size {
|
|
return Err("MTPLX GDN state/weight binding is too small".into());
|
|
}
|
|
}
|
|
if conv.is_some_and(|conv| unsafe { ds4_gpu_tensor_bytes(conv.raw()) } < 3 * 10240 * 2) {
|
|
return Err("MTPLX GDN conv state binding is too small".into());
|
|
}
|
|
let GdnProjected {
|
|
qkv,
|
|
z,
|
|
b,
|
|
a,
|
|
mut temporary,
|
|
..
|
|
} = gdn_project(&weights.input, input, 1)?;
|
|
let gated = Buffer::mtplx_bytes(6144 * 2)?;
|
|
let next_conv = Buffer::mtplx_bytes(3 * 10240 * 2)?;
|
|
let next_delta = Buffer::mtplx_floats(48 * 128 * 128)?;
|
|
// CustomKernel allocates outputs before materializing a broadcast input.
|
|
// Keep the cold conv copy after the input projection, as in the reference.
|
|
let initial_conv = if conv.is_none() {
|
|
let zero = Buffer::mtplx_bytes(3 * 10240 * 2)?;
|
|
let scalar = scalar_buffer(&0_u16.to_le_bytes())?;
|
|
dispatch_geometry(
|
|
"g1_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, &scalar).input(1),
|
|
tensor(1, &zero).output(3 * 10240),
|
|
bytes(3, &0_u64),
|
|
],
|
|
&[],
|
|
[3 * 10240, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)?;
|
|
Some(zero)
|
|
} else {
|
|
None
|
|
};
|
|
let conv = conv.or(initial_conv.as_ref()).unwrap();
|
|
gdn::dispatch_step(&[
|
|
tensor(0, &qkv).input(10240),
|
|
tensor(1, &z).input(6144),
|
|
tensor(2, &a).input(48),
|
|
tensor(3, &b).input(48),
|
|
tensor(4, conv).input(3 * 10240),
|
|
tensor(5, weights.conv).input(4 * 10240),
|
|
tensor(6, weights.a_log).input(48),
|
|
tensor(7, weights.dt_bias).input(48),
|
|
tensor(8, delta).input(48 * 128 * 128),
|
|
tensor(9, weights.norm).input(128),
|
|
tensor(10, &gated).output(6144),
|
|
tensor(11, &next_conv).output(3 * 10240),
|
|
tensor(12, &next_delta).output(48 * 128 * 128),
|
|
])?;
|
|
if let Some(initial_conv) = initial_conv {
|
|
encoder::add_temporaries(std::slice::from_ref(&initial_conv));
|
|
temporary.push(initial_conv);
|
|
}
|
|
let hidden = Buffer::mtplx_bytes(2560 * 2)?;
|
|
let [w, s, bias] = weights.output;
|
|
temporary.extend(affine_bf16(
|
|
[w, s, bias, &gated, &hidden],
|
|
[1, 2560, 6144],
|
|
weights.output_quant[0],
|
|
weights.output_quant[1],
|
|
true,
|
|
)?);
|
|
temporary.extend([qkv, z, b, a, gated]);
|
|
Ok(GdnDecodeOutput {
|
|
hidden,
|
|
conv: next_conv,
|
|
delta: next_delta,
|
|
_temporaries: temporary,
|
|
})
|
|
}
|