1615 lines
62 KiB
Rust
1615 lines
62 KiB
Rust
//! QSAIndexer shared projection/preparation/cache prefix from qwen4_exp.py.
|
|
//! Original compiled/fused/eager routing shares one persistent cache and core.
|
|
use super::array::{Array, Dtype};
|
|
use super::binary::{Binary, binary};
|
|
use super::cache::QsaCache;
|
|
use super::mlp::Linear;
|
|
use super::qsa_core::{Config, Core, Source};
|
|
use super::stream::Stream;
|
|
use super::{indexing, ops, qsa_prepare, reduce, rope};
|
|
|
|
pub(super) struct Indexer<'a> {
|
|
pub(super) projection: Option<Linear<'a>>,
|
|
pub(super) heads: i32,
|
|
pub(super) kv_heads: i32,
|
|
pub(super) head_dim: i32,
|
|
pub(super) block_topk: i32,
|
|
pub(super) ratio: i32,
|
|
pub(super) q_norm: &'a Array,
|
|
pub(super) k_norm: &'a Array,
|
|
pub(super) inv_freq: &'a Array,
|
|
pub(super) mrope_axes: Option<&'a Array>,
|
|
pub(super) eps: f64,
|
|
pub(super) scaling: f64,
|
|
pub(super) scratch: u64,
|
|
pub(super) compiled: Option<(Vec<usize>, Core<'a>)>,
|
|
}
|
|
|
|
pub(super) struct PreparedRows {
|
|
pub(super) query: Array,
|
|
pub(super) pooled: Array,
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) total: i32,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "issue107_indexer.rs"]
|
|
mod issue107;
|
|
|
|
impl<'a> Indexer<'a> {
|
|
#[expect(
|
|
clippy::too_many_arguments,
|
|
reason = "original static compiled-route predicate"
|
|
)]
|
|
fn compiled_supported(
|
|
&self,
|
|
source: &Array,
|
|
cache: &QsaCache,
|
|
pos: i32,
|
|
supplied: bool,
|
|
decode: bool,
|
|
mode: super::QsaCompiledMode,
|
|
options: &super::QsaIndexerOptions,
|
|
stream: Stream,
|
|
) -> bool {
|
|
use super::{QsaCompiledMode, stream::Device};
|
|
let s = source.layout();
|
|
if !options.fused
|
|
|| !options.compiled
|
|
|| stream.device() != Device::Gpu
|
|
|| s.shape().len() != 3
|
|
|| s.shape()[0] != 1
|
|
|| s.shape()[1] <= 0
|
|
|| self.kv_heads != 1
|
|
{
|
|
return false;
|
|
}
|
|
let rows = s.shape()[1];
|
|
if !decode
|
|
&& ((mode == QsaCompiledMode::UpdateOnly && options.phase_prefill)
|
|
|| (rows as u32 >= options.prefill_min_rows
|
|
&& (!matches!(
|
|
mode,
|
|
QsaCompiledMode::PrefillBlocks | QsaCompiledMode::UpdateOnly
|
|
) || rows as u32 != options.prefill_compile_rows)))
|
|
{
|
|
return false;
|
|
}
|
|
if cache.ratio != self.ratio
|
|
|| pos != cache.kv.offset
|
|
|| !matches!(s.dtype(), Dtype::BF16 | Dtype::F16 | Dtype::F32)
|
|
|| !(1..=128).contains(&self.head_dim)
|
|
|| !(1..=512).contains(&self.block_topk)
|
|
|| self.ratio <= 0
|
|
{
|
|
return false;
|
|
}
|
|
let freq = self.inv_freq.layout();
|
|
if freq.shape().len() != 1
|
|
|| freq.dtype() != Dtype::F32
|
|
|| freq.shape()[0] <= 0
|
|
|| i64::from(freq.shape()[0]) * 2 > i64::from(self.head_dim)
|
|
|| [self.q_norm, self.k_norm]
|
|
.iter()
|
|
.any(|a| a.layout().shape() != [self.head_dim] || a.layout().dtype() != s.dtype())
|
|
{
|
|
return false;
|
|
}
|
|
let width = (i64::from(self.heads) + i64::from(self.kv_heads)) * i64::from(self.head_dim);
|
|
if supplied {
|
|
if i64::from(s.shape()[2]) != width {
|
|
return false;
|
|
}
|
|
} else {
|
|
let Some(projection) = self.projection else {
|
|
return false;
|
|
};
|
|
let weight = projection.weight().layout();
|
|
if weight.shape().len() != 2 || i64::from(weight.shape()[0]) != width {
|
|
return false;
|
|
}
|
|
let same_dtype = match projection {
|
|
Linear::Dense(_) => weight.dtype() == s.dtype(),
|
|
Linear::Quantized(a) => {
|
|
a.scales.layout().dtype() == s.dtype() && a.biases.layout().dtype() == s.dtype()
|
|
}
|
|
};
|
|
if !same_dtype {
|
|
return false;
|
|
}
|
|
}
|
|
for backing in [&cache.raw_keys, &cache.pooled].into_iter().flatten() {
|
|
let a = backing.borrow();
|
|
let l = a.layout();
|
|
if l.shape().len() != 3
|
|
|| l.shape()[0] != 1
|
|
|| l.shape()[2] != self.head_dim
|
|
|| l.dtype() != s.dtype()
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
if pos > 0
|
|
&& cache
|
|
.raw_keys
|
|
.as_ref()
|
|
.is_none_or(|a| a.borrow().layout().shape()[1] < pos)
|
|
{
|
|
return false;
|
|
}
|
|
if let Some(pool) = &cache.pooled {
|
|
if !(0..=pool.borrow().layout().shape()[1]).contains(&cache.pooled_len) {
|
|
return false;
|
|
}
|
|
} else if cache.pooled_len != 0 {
|
|
return false;
|
|
}
|
|
let logical = (i64::from(pos) + i64::from(rows)) / i64::from(self.ratio);
|
|
if logical - i64::from(cache.pooled_len).min(logical)
|
|
> (i64::from(rows) + i64::from(self.ratio) - 1) / i64::from(self.ratio)
|
|
{
|
|
return false;
|
|
}
|
|
if decode
|
|
&& !matches!(mode, QsaCompiledMode::UpdateOnly | QsaCompiledMode::Blocks)
|
|
&& options.gather_decode
|
|
{
|
|
return false;
|
|
}
|
|
matches!(
|
|
mode,
|
|
QsaCompiledMode::UpdateOnly | QsaCompiledMode::PrefillBlocks
|
|
) || s.dtype() != Dtype::F32
|
|
|| super::sparse_attention::nax_available()
|
|
}
|
|
|
|
fn selector_supported(&self, q: &Array, p: &Array, prefill: bool, stream: Stream) -> bool {
|
|
let q = q.layout();
|
|
let p = p.layout();
|
|
stream.device() == super::stream::Device::Gpu
|
|
&& [q.dtype(), p.dtype()]
|
|
.iter()
|
|
.all(|d| matches!(d, Dtype::BF16 | Dtype::F16 | Dtype::F32))
|
|
&& q.shape().len() == 4
|
|
&& p.shape().len() == 3
|
|
&& q.shape()[0] == 1
|
|
&& p.shape()[0] == 1
|
|
&& q.shape()[1] > i32::from(prefill)
|
|
&& q.shape()[2] > 0
|
|
&& q.shape()[3] > 0
|
|
&& p.shape()[1] > 0
|
|
&& p.shape()[2] == q.shape()[3]
|
|
&& (1..=512).contains(&self.block_topk)
|
|
&& self.ratio > 0
|
|
&& (prefill
|
|
|| (q.dtype() != Dtype::F32 && p.dtype() != Dtype::F32)
|
|
|| super::sparse_attention::nax_available())
|
|
}
|
|
|
|
/// QSAIndexer.__call__/_call_rows. Eligibility alone chooses a lane;
|
|
/// errors after selection propagate, never rerun through an eager fallback.
|
|
#[expect(
|
|
clippy::too_many_arguments,
|
|
reason = "original indexer inputs and evaluator"
|
|
)]
|
|
pub(super) fn call(
|
|
&mut self,
|
|
hidden: &Array,
|
|
pos: i32,
|
|
cache: &mut QsaCache,
|
|
qk_rows: Option<&Array>,
|
|
options: &super::QsaIndexerOptions,
|
|
streams: &super::stream::Streams,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
stream: Stream,
|
|
) -> Result<(Option<super::attention::EagerSelection>, &'static str), String> {
|
|
use super::{
|
|
QsaSelectMode,
|
|
attention::{self, EagerSelection},
|
|
qsa_prefill,
|
|
};
|
|
let shape = hidden.layout().shape().to_vec();
|
|
if shape.len() != 3
|
|
|| shape[0] != 1
|
|
|| shape[1] <= 0
|
|
|| pos < 0
|
|
|| self.ratio <= 0
|
|
|| self.block_topk <= 0
|
|
{
|
|
return Err("QSA serves a positive single sequence with valid geometry".into());
|
|
}
|
|
let rows = shape[1];
|
|
let decode = rows == 1;
|
|
let total = pos.checked_add(rows).ok_or("QSA position overflow")?;
|
|
let mode = options.compiled_mode(
|
|
rows as u32,
|
|
total as u32,
|
|
self.ratio as u32,
|
|
self.block_topk as u32,
|
|
);
|
|
let vision = vision.filter(|_| {
|
|
self.mrope_axes.is_some() && crate::engine::qwen_vision::cache_policy() != 0
|
|
});
|
|
if vision.is_none()
|
|
&& self.compiled_supported(
|
|
qk_rows.unwrap_or(hidden),
|
|
cache,
|
|
pos,
|
|
qk_rows.is_some(),
|
|
decode,
|
|
mode,
|
|
options,
|
|
stream,
|
|
)
|
|
{
|
|
return self
|
|
.call_compiled(
|
|
hidden,
|
|
pos,
|
|
cache,
|
|
qk_rows,
|
|
mode,
|
|
options.prefill_workspace,
|
|
stream,
|
|
)
|
|
.map(|a| (a, "compiled"));
|
|
}
|
|
let Some(prepared) =
|
|
self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, vision, stream)?
|
|
else {
|
|
return Ok((None, "update_only"));
|
|
};
|
|
let q = prepared.query;
|
|
let pooled = prepared.pooled;
|
|
let logical = pooled.layout().shape()[1];
|
|
let large_prefill = !decode && options.large_prefill(rows as u32, pos as u32);
|
|
if large_prefill && self.selector_supported(&q, &pooled, true, stream) {
|
|
let mut values = qsa_prefill::select(
|
|
&q,
|
|
&pooled,
|
|
[pos, total, logical].map(qsa_prepare::Position::Host),
|
|
self.block_topk,
|
|
self.ratio,
|
|
None,
|
|
QsaSelectMode::Blocks,
|
|
options.prefill_workspace,
|
|
stream,
|
|
)?
|
|
.into_iter();
|
|
attention::prefill_count("metal_selector")?;
|
|
return Ok((
|
|
Some(EagerSelection::PrefillBlocks(
|
|
values.next().ok_or("QSA missing prefill ids")?,
|
|
values.next().ok_or("QSA missing prefill validity")?,
|
|
)),
|
|
"metal_prefill",
|
|
));
|
|
}
|
|
let backing = cache.pooled.as_ref().map(|a| a.borrow().clone());
|
|
let fused = options.fused
|
|
&& (decode || (rows as u32) < options.prefill_min_rows)
|
|
&& !(decode && options.gather_decode)
|
|
&& backing
|
|
.as_ref()
|
|
.is_some_and(|p| self.selector_supported(&q, p, false, stream));
|
|
if !fused {
|
|
return attention::select_eager(
|
|
&q,
|
|
pos,
|
|
cache,
|
|
total,
|
|
self.block_topk,
|
|
i32::try_from(options.tile_rows).map_err(|_| "QSA tile overflow")?,
|
|
&super::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,
|
|
},
|
|
streams,
|
|
stream,
|
|
)
|
|
.map(|a| (Some(a), "eager"));
|
|
}
|
|
let mode = if decode && options.flash {
|
|
QsaSelectMode::Blocks
|
|
} else if !decode
|
|
&& options.gather
|
|
&& rows as u32 <= options.gather_max_rows
|
|
&& total as u32 >= options.gather_min_context
|
|
{
|
|
QsaSelectMode::RowTokens
|
|
} else {
|
|
QsaSelectMode::DenseMask
|
|
};
|
|
let mut values = self
|
|
.select_fused(
|
|
&q,
|
|
pos,
|
|
backing.as_ref().unwrap(),
|
|
total,
|
|
logical,
|
|
mode,
|
|
stream,
|
|
)?
|
|
.into_iter();
|
|
let first = values.next().ok_or("QSA missing fused selection")?;
|
|
let out = match mode {
|
|
QsaSelectMode::Blocks => {
|
|
let width = first.layout().shape()[1];
|
|
EagerSelection::FlashBlocks(
|
|
ops::squeeze(
|
|
&ops::slice(&first, &[0, 0], &[1, width], &[1, 1], stream)?,
|
|
Some(&[0]),
|
|
stream,
|
|
)?,
|
|
((pos + 1) / self.ratio) * self.ratio,
|
|
)
|
|
}
|
|
QsaSelectMode::RowTokens => {
|
|
EagerSelection::RowTokens(first, values.next().ok_or("QSA missing fused validity")?)
|
|
}
|
|
QsaSelectMode::DenseMask => EagerSelection::Dense(first),
|
|
};
|
|
Ok((Some(out), "legacy_fused"))
|
|
}
|
|
|
|
/// Identity seal is of captured parameter descriptors, not their bytes,
|
|
/// call shapes or environment. Replacing a leaf resets the whole bank.
|
|
fn compiled_core(&mut self, prefill_workspace: u64) -> Result<&mut Core<'a>, String> {
|
|
let mut signature = vec![self.q_norm.id(), self.k_norm.id(), self.inv_freq.id()];
|
|
if let Some(projection) = self.projection {
|
|
signature.push(projection.weight().id());
|
|
}
|
|
if let Some(Linear::Quantized(layer)) = self.projection {
|
|
signature.extend([layer.scales.id(), layer.biases.id()]);
|
|
}
|
|
if self
|
|
.compiled
|
|
.as_ref()
|
|
.is_none_or(|(old, _)| *old != signature)
|
|
{
|
|
let core = Core::new(
|
|
Config {
|
|
heads: self.heads,
|
|
kv_heads: self.kv_heads,
|
|
dim: self.head_dim,
|
|
topk: self.block_topk,
|
|
ratio: self.ratio,
|
|
eps: self.eps,
|
|
scale: self.scaling,
|
|
scratch: self.scratch,
|
|
prefill_workspace,
|
|
..Config::default()
|
|
},
|
|
[
|
|
self.q_norm.clone(),
|
|
self.k_norm.clone(),
|
|
self.inv_freq.clone(),
|
|
],
|
|
self.projection,
|
|
)?;
|
|
self.compiled = Some((signature, core));
|
|
}
|
|
Ok(&mut self.compiled.as_mut().unwrap().1)
|
|
}
|
|
|
|
/// Original _call_rows_compiled: reserve first, commit explicit outputs
|
|
/// without evaluating frontiers. Attention alone advances kv.offset.
|
|
#[expect(
|
|
clippy::too_many_arguments,
|
|
reason = "original compiled indexer inputs and resolved workspace"
|
|
)]
|
|
pub(super) fn call_compiled(
|
|
&mut self,
|
|
hidden: &Array,
|
|
pos: i32,
|
|
cache: &mut QsaCache,
|
|
qk_rows: Option<&Array>,
|
|
mode: super::QsaCompiledMode,
|
|
prefill_workspace: u64,
|
|
stream: Stream,
|
|
) -> Result<Option<super::attention::EagerSelection>, String> {
|
|
use super::{
|
|
QsaCompiledMode,
|
|
attention::{self, EagerSelection},
|
|
cache::object,
|
|
};
|
|
use qsa_prepare::Position;
|
|
let rows = hidden
|
|
.layout()
|
|
.shape()
|
|
.get(1)
|
|
.copied()
|
|
.ok_or("QSA hidden rank")?;
|
|
let total = pos.checked_add(rows).ok_or("QSA total overflow")?;
|
|
if self.ratio <= 0 {
|
|
return Err("QSA ratio must be positive".into());
|
|
}
|
|
let logical = total / self.ratio;
|
|
let value = qk_rows.unwrap_or(hidden);
|
|
let (raw, pool) = cache.ensure_compiled_backings(
|
|
value.layout().dtype(),
|
|
pos,
|
|
rows,
|
|
self.head_dim,
|
|
stream,
|
|
)?;
|
|
let out = self.compiled_core(prefill_workspace)?.select(
|
|
if qk_rows.is_some() {
|
|
Source::QkRows
|
|
} else {
|
|
Source::Hidden
|
|
},
|
|
value,
|
|
&raw.borrow(),
|
|
&pool.borrow(),
|
|
[pos, total, logical, cache.pooled_len.min(logical)].map(Position::Host),
|
|
mode,
|
|
None,
|
|
stream,
|
|
)?;
|
|
cache.raw_keys = Some(object(out.raw_keys));
|
|
cache.pooled = Some(object(out.pooled));
|
|
cache.pooled_len = logical;
|
|
cache.pooled_f32_t = None;
|
|
let mut selected = out.selection.into_iter();
|
|
Ok(Some(match mode {
|
|
QsaCompiledMode::UpdateOnly => return Ok(None),
|
|
QsaCompiledMode::Blocks => {
|
|
let ids = selected.next().ok_or("QSA missing block ids")?;
|
|
let mut end = ids.layout().shape().to_vec();
|
|
end[0] = 1;
|
|
let first =
|
|
ops::slice(&ids, &vec![0; end.len()], &end, &vec![1; end.len()], stream)?;
|
|
EagerSelection::FlashBlocks(
|
|
ops::squeeze(&first, Some(&[0]), stream)?,
|
|
((pos + 1) / self.ratio) * self.ratio,
|
|
)
|
|
}
|
|
QsaCompiledMode::PrefillBlocks => {
|
|
attention::prefill_count("compiled_selector")?;
|
|
EagerSelection::PrefillBlocks(
|
|
selected.next().ok_or("QSA missing prefill ids")?,
|
|
selected.next().ok_or("QSA missing prefill validity")?,
|
|
)
|
|
}
|
|
QsaCompiledMode::RowTokens => EagerSelection::RowTokens(
|
|
selected.next().ok_or("QSA missing row ids")?,
|
|
selected.next().ok_or("QSA missing row validity")?,
|
|
),
|
|
QsaCompiledMode::DenseMask => {
|
|
let mask = selected.next().ok_or("QSA missing dense mask")?;
|
|
EagerSelection::Dense(ops::slice(
|
|
&mask,
|
|
&[0; 4],
|
|
&[1, 1, rows, total],
|
|
&[1; 4],
|
|
stream,
|
|
)?)
|
|
}
|
|
}))
|
|
}
|
|
|
|
/// Original _select_fused: chunk query rows (32MiB score scratch), never
|
|
/// logical history; keep dense specialization at backing capacity+tail.
|
|
#[expect(
|
|
clippy::too_many_arguments,
|
|
reason = "original selector mode and frontier arguments"
|
|
)]
|
|
pub(super) fn select_fused(
|
|
&self,
|
|
q: &Array,
|
|
pos: i32,
|
|
pooled: &Array,
|
|
total: i32,
|
|
logical: i32,
|
|
mode: super::QsaSelectMode,
|
|
stream: Stream,
|
|
) -> Result<Vec<Array>, String> {
|
|
use super::{QsaSelectMode, qsa_select};
|
|
let rows = q.layout().shape().get(1).copied().ok_or("QSA query rank")?;
|
|
let blocks = pooled
|
|
.layout()
|
|
.shape()
|
|
.get(1)
|
|
.copied()
|
|
.ok_or("QSA pooled rank")?;
|
|
if rows <= 0 {
|
|
return Err("QSA fused rows must be positive".into());
|
|
}
|
|
let chunk_rows = u64::try_from(rows)
|
|
.unwrap()
|
|
.min((self.scratch / (blocks.max(1) as u64 * 4)).max(1))
|
|
as i32;
|
|
let capacity = if mode == QsaSelectMode::DenseMask {
|
|
Some(
|
|
blocks
|
|
.checked_add(1)
|
|
.and_then(|n| n.checked_mul(self.ratio))
|
|
.ok_or("QSA dense capacity overflow")?,
|
|
)
|
|
} else {
|
|
None
|
|
};
|
|
let mut chunks = Vec::new();
|
|
for start in (0..rows).step_by(chunk_rows as usize) {
|
|
let mut end = q.layout().shape().to_vec();
|
|
end[1] = start.saturating_add(chunk_rows).min(rows);
|
|
let q = ops::slice(q, &[0, start, 0, 0], &end, &[1; 4], stream)?;
|
|
chunks.push(qsa_select::select(
|
|
&q,
|
|
pooled,
|
|
[
|
|
qsa_prepare::Position::Host(
|
|
pos.checked_add(start).ok_or("QSA position overflow")?,
|
|
),
|
|
qsa_prepare::Position::Host(total),
|
|
qsa_prepare::Position::Host(logical),
|
|
],
|
|
self.block_topk,
|
|
self.ratio,
|
|
capacity,
|
|
mode,
|
|
stream,
|
|
)?);
|
|
}
|
|
let mut outputs = if chunks.len() == 1 {
|
|
chunks.pop().unwrap()
|
|
} else {
|
|
(0..chunks[0].len())
|
|
.map(|leaf| {
|
|
ops::concatenate(
|
|
&chunks.iter().map(|c| c[leaf].clone()).collect::<Vec<_>>(),
|
|
if mode == QsaSelectMode::DenseMask {
|
|
2
|
|
} else {
|
|
0
|
|
},
|
|
stream,
|
|
)
|
|
})
|
|
.collect::<Result<Vec<_>, _>>()?
|
|
};
|
|
if mode == QsaSelectMode::DenseMask {
|
|
outputs[0] = ops::slice(&outputs[0], &[0; 4], &[1, 1, rows, total], &[1; 4], stream)?;
|
|
}
|
|
Ok(outputs)
|
|
}
|
|
|
|
fn prepare_queries(
|
|
&self,
|
|
q: &Array,
|
|
pos: i32,
|
|
fused: bool,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
if vision.is_none()
|
|
&& fused
|
|
&& qsa_prepare::supported(q, self.q_norm, self.inv_freq, 4, stream)
|
|
{
|
|
return qsa_prepare::queries(
|
|
q,
|
|
self.q_norm,
|
|
self.inv_freq,
|
|
qsa_prepare::Position::Host(pos),
|
|
self.eps,
|
|
self.scaling,
|
|
stream,
|
|
);
|
|
}
|
|
let q = ops::rms_norm(q, Some(self.q_norm), self.eps as f32, stream)?;
|
|
let positions = indexing::arange_range_i32(pos, pos + q.layout().shape()[1], stream)?;
|
|
let (cos, sin) = match (vision, self.mrope_axes) {
|
|
(Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin(
|
|
&positions,
|
|
table,
|
|
delta,
|
|
self.inv_freq,
|
|
axes,
|
|
self.scaling,
|
|
stream,
|
|
)?,
|
|
_ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?,
|
|
};
|
|
rope::apply_partial(&q, &cos, &sin, stream)
|
|
}
|
|
|
|
fn pool_keys(
|
|
&self,
|
|
fresh: &Array,
|
|
old: i32,
|
|
total: i32,
|
|
fused: bool,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
if vision.is_none()
|
|
&& fused
|
|
&& qsa_prepare::supported(fresh, self.k_norm, self.inv_freq, 3, stream)
|
|
{
|
|
return qsa_prepare::pool(
|
|
fresh,
|
|
self.k_norm,
|
|
self.inv_freq,
|
|
qsa_prepare::Position::Host(old),
|
|
self.ratio,
|
|
self.eps,
|
|
self.scaling,
|
|
stream,
|
|
);
|
|
}
|
|
let fresh = ops::reshape(fresh, &[1, total - old, self.ratio, self.head_dim], stream)?;
|
|
let pooled = reduce::mean_f32(
|
|
&ops::astype(&fresh, Dtype::F32, false, stream)?,
|
|
&[2],
|
|
false,
|
|
stream,
|
|
)?;
|
|
let pooled = ops::astype(&pooled, fresh.layout().dtype(), false, stream)?;
|
|
let pooled = ops::rms_norm(&pooled, Some(self.k_norm), self.eps as f32, stream)?;
|
|
let ratio = Array::new(
|
|
&[],
|
|
Dtype::I32,
|
|
super::scalar_buffer(&self.ratio.to_le_bytes())?,
|
|
)?;
|
|
let positions = binary(
|
|
&indexing::arange_range_i32(old, total, stream)?,
|
|
&ratio,
|
|
Binary::Multiply,
|
|
stream,
|
|
)?;
|
|
let (cos, sin) = match (vision, self.mrope_axes) {
|
|
(Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin(
|
|
&positions,
|
|
table,
|
|
delta,
|
|
self.inv_freq,
|
|
axes,
|
|
self.scaling,
|
|
stream,
|
|
)?,
|
|
_ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?,
|
|
};
|
|
ops::squeeze(
|
|
&rope::apply_partial(
|
|
&ops::expand_dims(&pooled, &[2], stream)?,
|
|
&cos,
|
|
&sin,
|
|
stream,
|
|
)?,
|
|
Some(&[2]),
|
|
stream,
|
|
)
|
|
}
|
|
|
|
pub(super) fn extend_pooled(
|
|
&self,
|
|
cache: &mut QsaCache,
|
|
total: i32,
|
|
fused: bool,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
stream: Stream,
|
|
) -> Result<Option<Array>, String> {
|
|
let nb = total / self.ratio;
|
|
let old = cache.pooled_len.min(nb);
|
|
if nb > old {
|
|
let raw = cache
|
|
.raw_keys
|
|
.as_ref()
|
|
.ok_or("QSA pooling requires raw cache")?
|
|
.borrow();
|
|
let fresh = ops::slice(
|
|
&raw,
|
|
&[0, old * self.ratio, 0],
|
|
&[1, nb * self.ratio, self.head_dim],
|
|
&[1; 3],
|
|
stream,
|
|
)?;
|
|
drop(raw);
|
|
cache.write_pooled(
|
|
&self.pool_keys(&fresh, old, nb, fused, vision, stream)?,
|
|
old,
|
|
nb,
|
|
stream,
|
|
)?;
|
|
}
|
|
if nb == 0 {
|
|
return Ok(None);
|
|
}
|
|
let pooled = cache
|
|
.pooled
|
|
.as_ref()
|
|
.ok_or("QSA pooled cache missing")?
|
|
.borrow();
|
|
ops::slice(&pooled, &[0; 3], &[1, nb, self.head_dim], &[1; 3], stream).map(Some)
|
|
}
|
|
|
|
/// Common noncompiled _call_rows prefix. None means the complete-block
|
|
/// budget fits: cache updates still happen, but lazy query work is unused.
|
|
/// Cache.kv.offset is advanced by Attention, never by this indexer.
|
|
#[expect(
|
|
clippy::too_many_arguments,
|
|
reason = "original preparation inputs and request positions"
|
|
)]
|
|
pub(super) fn prepare_rows(
|
|
&self,
|
|
hidden: &Array,
|
|
pos: i32,
|
|
cache: &mut QsaCache,
|
|
qk_rows: Option<&Array>,
|
|
fused: bool,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
stream: Stream,
|
|
) -> Result<Option<PreparedRows>, String> {
|
|
let shape = hidden.layout().shape().to_vec();
|
|
if shape.len() != 3
|
|
|| shape[0] != 1
|
|
|| shape[1] <= 0
|
|
|| pos < 0
|
|
|| cache.kv.offset != pos
|
|
|| self.heads <= 0
|
|
|| self.head_dim <= 0
|
|
|| self.ratio <= 0
|
|
|| self.block_topk <= 0
|
|
|| cache.ratio != self.ratio
|
|
{
|
|
return Err("invalid QSA indexer dimensions/frontier".into());
|
|
}
|
|
let total = pos.checked_add(shape[1]).ok_or("QSA position overflow")?;
|
|
let width = self
|
|
.heads
|
|
.checked_mul(self.head_dim)
|
|
.ok_or("QSA query width overflow")?;
|
|
let full_width = width
|
|
.checked_add(self.head_dim)
|
|
.ok_or("QSA projection width overflow")?;
|
|
let qk = match qk_rows {
|
|
Some(rows) => rows.clone(),
|
|
None => self
|
|
.projection
|
|
.ok_or("QSA hidden source requires a projection")?
|
|
.apply(hidden, stream)?,
|
|
};
|
|
if qk.layout().shape() != [1, shape[1], full_width] {
|
|
return Err("QSA projection must have (heads+1)*head_dim columns".into());
|
|
}
|
|
let parts = ops::split(&qk, &[width], -1, stream)?;
|
|
let query = ops::reshape(&parts[0], &[1, shape[1], self.heads, self.head_dim], stream)?;
|
|
let keys = ops::reshape(&parts[1], &[1, shape[1], self.head_dim], stream)?;
|
|
let query = self.prepare_queries(&query, pos, fused, vision, stream)?;
|
|
cache.write_raw(&keys, stream)?;
|
|
let pooled = self.extend_pooled(cache, total, fused, vision, stream)?;
|
|
if total / self.ratio <= self.block_topk {
|
|
return Ok(None);
|
|
}
|
|
Ok(Some(PreparedRows {
|
|
query,
|
|
pooled: pooled.ok_or("QSA selection requires pooled keys")?,
|
|
total,
|
|
}))
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Metal and original ongoing indexer routing receipts; no weights"]
|
|
fn mtplx_qsa_canonical_indexer_routes_match_reference() {
|
|
use super::super::{
|
|
gpu::Context,
|
|
qwen_mtplx_tests::{capture_dispatches, pattern, qsa_indexer_options},
|
|
};
|
|
use super::{
|
|
QsaCompiledMode, allocator,
|
|
attention::EagerSelection,
|
|
cache::object,
|
|
mlp,
|
|
stream::{Device, Streams},
|
|
};
|
|
use sha2::{Digest, Sha256};
|
|
super::configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
|
|
let gpu = streams.default_stream(Device::Gpu).unwrap();
|
|
let cases = include_str!("../../../../tests/fixtures/mtplx-custom-kernels.jsonl")
|
|
.lines()
|
|
.chain(include_str!("../../../../tests/fixtures/mtplx-qsa-indexer-connected.jsonl").lines())
|
|
.map(|l| serde_json::from_str::<serde_json::Value>(l).unwrap())
|
|
.filter(|c| c["kernel"].as_str().unwrap().starts_with("qsa_indexer_"))
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(cases.len(), 200);
|
|
let weight = mlp::fixture_linear(640, 2560, 4, 64, 75);
|
|
let q_norm = Array::new(&[128], Dtype::BF16, pattern(128, 72)).unwrap();
|
|
let k_norm = Array::new(&[128], Dtype::BF16, pattern(128, 73)).unwrap();
|
|
let freqs = cases[0]["frequencies"]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.flat_map(|v| (v.as_f64().unwrap() as f32).to_le_bytes())
|
|
.collect::<Vec<_>>();
|
|
let inv_freq = Array::new(&[32], Dtype::F32, super::scalar_buffer(&freqs).unwrap()).unwrap();
|
|
let make = || Indexer {
|
|
projection: Some(Linear::Quantized(&weight)),
|
|
heads: 4,
|
|
kv_heads: 1,
|
|
head_dim: 128,
|
|
block_topk: 512,
|
|
ratio: 4,
|
|
q_norm: &q_norm,
|
|
k_norm: &k_norm,
|
|
inv_freq: &inv_freq,
|
|
mrope_axes: None,
|
|
eps: 1e-6,
|
|
scaling: 1.,
|
|
scratch: 32 * 1024 * 1024,
|
|
compiled: None,
|
|
};
|
|
let mut indexer = make();
|
|
let mut cache = QsaCache::new(4);
|
|
let source = |shape: &[i32], salt| {
|
|
Array::new(
|
|
shape,
|
|
Dtype::BF16,
|
|
pattern(shape.iter().product::<i32>() as u32, salt),
|
|
)
|
|
.unwrap()
|
|
};
|
|
let digest = |a: &Array| {
|
|
let a =
|
|
ops::contiguous(&ops::astype(a, Dtype::F32, false, gpu).unwrap(), false, gpu).unwrap();
|
|
ops::evaluate(&streams, std::slice::from_ref(&a), gpu, false).unwrap();
|
|
let mut bytes = vec![0; a.layout().nbytes()];
|
|
a.buffer().read(a.offset(), &mut bytes).unwrap();
|
|
Sha256::digest(bytes)
|
|
.iter()
|
|
.map(|b| format!("{b:02x}"))
|
|
.collect::<String>()
|
|
};
|
|
for c in &cases {
|
|
super::eval::execution_tests::observed(|| {
|
|
let n = |s: &str| c[s].as_i64().unwrap() as i32;
|
|
let step = n("step");
|
|
let rows = n("rows");
|
|
let pos = n("pos");
|
|
let options = qsa_indexer_options(&c["options"]);
|
|
if step == 0 {
|
|
indexer = make();
|
|
cache = QsaCache::new(4);
|
|
let prefix = n("prefix");
|
|
if prefix > 0 {
|
|
cache
|
|
.write_raw(&source(&[1, prefix, 128], 88), gpu)
|
|
.unwrap();
|
|
indexer
|
|
.extend_pooled(&mut cache, prefix, options.fused, None, gpu)
|
|
.unwrap();
|
|
let kv = ops::transpose(&source(&[1, prefix, 2, 256], 89), &[0, 2, 1, 3], gpu)
|
|
.unwrap();
|
|
cache.kv.update_and_fetch(&kv, &kv, gpu).unwrap();
|
|
let leaves = [
|
|
&cache.raw_keys,
|
|
&cache.pooled,
|
|
&cache.pooled_f32_t,
|
|
&cache.kv.keys,
|
|
&cache.kv.values,
|
|
]
|
|
.into_iter()
|
|
.flatten()
|
|
.map(|a| a.borrow().clone())
|
|
.collect::<Vec<_>>();
|
|
ops::evaluate(&streams, &leaves, gpu, false).unwrap();
|
|
}
|
|
}
|
|
assert_eq!(cache.kv.offset, pos);
|
|
let hidden = source(&[1, rows, 2560], 77 + step as u32);
|
|
let qk = c["supplied"].as_bool().unwrap().then(|| {
|
|
ops::slice(
|
|
&source(&[1, rows, n("stride")], 121 + step as u32),
|
|
&[0; 3],
|
|
&[1, rows, 640],
|
|
&[1; 3],
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
});
|
|
let ((selected, lane), dispatches) = capture_dispatches(|| {
|
|
indexer
|
|
.call(
|
|
&hidden,
|
|
pos,
|
|
&mut cache,
|
|
qk.as_ref(),
|
|
&options,
|
|
&streams,
|
|
None,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
});
|
|
assert_eq!(lane, c["lane"].as_str().unwrap(), "{}", c["kernel"]);
|
|
if matches!(lane, "compiled" | "metal_prefill" | "legacy_fused") {
|
|
assert!(dispatches.is_empty(), "{} must stay lazy", c["kernel"]);
|
|
}
|
|
if let Some(counts) = c.get("graph_counts") {
|
|
let actual =
|
|
indexer
|
|
.compiled
|
|
.as_ref()
|
|
.map_or(serde_json::json!([0, 0, 0]), |(_, core)| {
|
|
let r = core.report();
|
|
serde_json::json!([r["compiled_calls"], r["traces"], r["entry_count"]])
|
|
});
|
|
assert_eq!(actual, *counts, "{}", c["kernel"]);
|
|
}
|
|
let (kind, outputs, tail) = match selected {
|
|
None => ("none", vec![], None),
|
|
Some(EagerSelection::Dense(a)) => ("dense", vec![a], None),
|
|
Some(EagerSelection::PrefillBlocks(a, b)) => ("flash_prefill", vec![a, b], None),
|
|
Some(EagerSelection::FlashBlocks(a, t)) => ("flash", vec![a], Some(t)),
|
|
Some(EagerSelection::RowTokens(a, b)) => ("gather_rows", vec![a, b], None),
|
|
Some(EagerSelection::DecodeTokens(a)) => ("decode_tokens", vec![a], None),
|
|
};
|
|
assert_eq!(kind, c["kind"].as_str().unwrap(), "{}", c["kernel"]);
|
|
assert_eq!(tail, c["tail"].as_i64().map(|v| v as i32));
|
|
let mut leaves = outputs.clone();
|
|
leaves.extend(
|
|
[&cache.raw_keys, &cache.pooled, &cache.pooled_f32_t]
|
|
.into_iter()
|
|
.flatten()
|
|
.map(|a| a.borrow().clone()),
|
|
);
|
|
ops::evaluate(&streams, &leaves, gpu, false).unwrap();
|
|
assert_eq!(
|
|
serde_json::json!(outputs.iter().map(&digest).collect::<Vec<_>>()),
|
|
c["sha256"],
|
|
"{} output",
|
|
c["kernel"]
|
|
);
|
|
assert_eq!(cache.kv.offset, pos);
|
|
assert_eq!(cache.pooled_len, n("pooled_len"));
|
|
for (key, leaf, axis) in [
|
|
("raw_cap", &cache.raw_keys, 1),
|
|
("pooled_cap", &cache.pooled, 1),
|
|
("mirror_cap", &cache.pooled_f32_t, 3),
|
|
] {
|
|
assert_eq!(
|
|
leaf.as_ref()
|
|
.map_or(0, |a| a.borrow().layout().shape()[axis]),
|
|
n(key),
|
|
"{} {key}",
|
|
c["kernel"]
|
|
);
|
|
}
|
|
for (key, a) in [
|
|
("raw", &cache.raw_keys),
|
|
("pooled", &cache.pooled),
|
|
("mirror", &cache.pooled_f32_t),
|
|
] {
|
|
let hash = a.as_ref().map(|a| digest(&a.borrow()));
|
|
assert_eq!(
|
|
serde_json::json!(hash),
|
|
c["state_sha256"][key],
|
|
"{} {key}",
|
|
c["kernel"]
|
|
);
|
|
}
|
|
let kv = ops::transpose(
|
|
&source(&[1, rows, 2, 256], 140 + step as u32),
|
|
&[0, 2, 1, 3],
|
|
gpu,
|
|
)
|
|
.unwrap();
|
|
let (k, v) = cache.kv.update_and_fetch(&kv, &kv, gpu).unwrap();
|
|
ops::evaluate(
|
|
&streams,
|
|
&[k.borrow().clone(), v.borrow().clone()],
|
|
gpu,
|
|
false,
|
|
)
|
|
.unwrap();
|
|
});
|
|
}
|
|
// Reuse the original 720 pre-reservation eligibility receipts, including
|
|
// inconsistent frontiers and present-but-nonbucketed or empty backings.
|
|
let routing: serde_json::Value = serde_json::from_str(include_str!(
|
|
"../../../../tests/fixtures/mtplx-qsa-compiled-routes.json"
|
|
))
|
|
.unwrap();
|
|
let indexer = make();
|
|
let mut checked = 0;
|
|
for c in routing["cases"].as_array().unwrap() {
|
|
let options = qsa_indexer_options(&routing["options"][c[0].as_u64().unwrap() as usize]);
|
|
let state = &routing["states"][c[1].as_u64().unwrap() as usize];
|
|
let n = |i: usize| state[i].as_i64().unwrap_or(0) as i32;
|
|
let mut cache = QsaCache::new(4);
|
|
cache.kv.offset = n(2);
|
|
cache.pooled_len = n(5);
|
|
cache.raw_keys = state[3]
|
|
.as_i64()
|
|
.map(|cap| object(Array::unallocated(&[1, cap as i32, 128], Dtype::BF16).unwrap()));
|
|
cache.pooled = state[4]
|
|
.as_i64()
|
|
.map(|cap| object(Array::unallocated(&[1, cap as i32, 128], Dtype::BF16).unwrap()));
|
|
let mode = match c[2].as_str().unwrap() {
|
|
"update_only" => QsaCompiledMode::UpdateOnly,
|
|
"blocks" => QsaCompiledMode::Blocks,
|
|
"prefill_blocks" => QsaCompiledMode::PrefillBlocks,
|
|
"row_tokens" => QsaCompiledMode::RowTokens,
|
|
"dense_mask" => QsaCompiledMode::DenseMask,
|
|
_ => panic!("unknown mode"),
|
|
};
|
|
assert_eq!(
|
|
options.compiled_mode(n(1) as u32, (n(0) + n(1)) as u32, 4, 512),
|
|
mode
|
|
);
|
|
for supplied in [false, true] {
|
|
let source =
|
|
Array::unallocated(&[1, n(1), if supplied { 640 } else { 2560 }], Dtype::BF16)
|
|
.unwrap();
|
|
let (accepted, dispatches) = capture_dispatches(|| {
|
|
indexer.compiled_supported(
|
|
&source,
|
|
&cache,
|
|
n(0),
|
|
supplied,
|
|
n(1) == 1,
|
|
mode,
|
|
&options,
|
|
gpu,
|
|
)
|
|
});
|
|
assert!(dispatches.is_empty());
|
|
assert_eq!(
|
|
accepted,
|
|
c[if supplied { 4 } else { 3 }].as_bool().unwrap(),
|
|
"{c} {state} supplied{supplied}"
|
|
);
|
|
checked += 1;
|
|
}
|
|
}
|
|
assert_eq!(checked, 720);
|
|
assert!(indexer.compiled.is_none());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Metal and original indexer lifecycle receipts; no model weights"]
|
|
fn mtplx_qsa_indexer_compiled_commit_matches_reference() {
|
|
use super::super::{
|
|
gpu::Context,
|
|
qwen_mtplx_tests::{capture_dispatches, pattern},
|
|
};
|
|
use super::{
|
|
QsaCompiledMode, allocator,
|
|
attention::{self, EagerSelection},
|
|
mlp,
|
|
stream::{Device, Streams},
|
|
};
|
|
use sha2::{Digest, Sha256};
|
|
super::configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
|
|
let gpu = streams.default_stream(Device::Gpu).unwrap();
|
|
let cases = include_str!("../../../../tests/fixtures/mtplx-qsa-indexer-commit.jsonl")
|
|
.lines()
|
|
.map(|l| serde_json::from_str::<serde_json::Value>(l).unwrap())
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(cases.len(), 26);
|
|
let renewed = |a: &Array| {
|
|
Array::new(
|
|
a.layout().shape(),
|
|
a.layout().dtype(),
|
|
a.buffer()
|
|
.view(a.offset(), a.layout().nbytes() as u64)
|
|
.unwrap(),
|
|
)
|
|
.unwrap()
|
|
};
|
|
let norm = |salt| Array::new(&[128], Dtype::BF16, pattern(128, salt)).unwrap();
|
|
let replacement = norm(72);
|
|
let q_norm = [norm(72), replacement.clone(), replacement];
|
|
let k_norm = [norm(73), norm(73)];
|
|
let invalid = Array::unallocated(&[127], Dtype::BF16).unwrap();
|
|
let bytes = cases[0]["frequencies"]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.flat_map(|v| (v.as_f64().unwrap() as f32).to_le_bytes())
|
|
.collect::<Vec<_>>();
|
|
let freq = Array::new(&[32], Dtype::F32, super::scalar_buffer(&bytes).unwrap()).unwrap();
|
|
let freq = [renewed(&freq), freq];
|
|
let weight = Array::new(&[640, 128], Dtype::BF16, pattern(640 * 128, 231)).unwrap();
|
|
let weights = [renewed(&weight), weight];
|
|
let mut layers = vec![mlp::fixture_linear(640, 128, 4, 64, 75)];
|
|
for leaf in 0..3 {
|
|
let a = layers.last().unwrap();
|
|
layers.push(mlp::QuantizedLinear {
|
|
weight: if leaf == 0 {
|
|
renewed(&a.weight)
|
|
} else {
|
|
a.weight.clone()
|
|
},
|
|
scales: if leaf == 1 {
|
|
renewed(&a.scales)
|
|
} else {
|
|
a.scales.clone()
|
|
},
|
|
biases: if leaf == 2 {
|
|
renewed(&a.biases)
|
|
} else {
|
|
a.biases.clone()
|
|
},
|
|
bits: a.bits,
|
|
group: a.group,
|
|
});
|
|
}
|
|
for retain in [true, false] {
|
|
for quantized in [false, true] {
|
|
let mut indexer = Indexer {
|
|
projection: Some(if quantized {
|
|
Linear::Quantized(&layers[0])
|
|
} else {
|
|
Linear::Dense(&weights[0])
|
|
}),
|
|
heads: 4,
|
|
kv_heads: 1,
|
|
head_dim: 128,
|
|
block_topk: 8,
|
|
ratio: 4,
|
|
q_norm: &q_norm[0],
|
|
k_norm: &k_norm[0],
|
|
inv_freq: &freq[0],
|
|
mrope_axes: None,
|
|
eps: 1e-6,
|
|
scaling: 1.,
|
|
scratch: 32 * 1024 * 1024,
|
|
compiled: None,
|
|
};
|
|
let mut cache = QsaCache::new(4);
|
|
for c in cases.iter().filter(|c| c["quantized"] == quantized) {
|
|
super::eval::execution_tests::observed(|| {
|
|
let n = |key: &str| c[key].as_i64().unwrap() as i32;
|
|
let step = n("step");
|
|
let rows = n("rows");
|
|
if step >= 2 {
|
|
indexer.scratch = 1024 * 1024;
|
|
}
|
|
if step == 3 {
|
|
indexer.q_norm = &q_norm[1];
|
|
}
|
|
if step == 4 {
|
|
indexer.k_norm = &k_norm[1];
|
|
}
|
|
if step == 5 {
|
|
indexer.inv_freq = &freq[1];
|
|
}
|
|
if step == 6 {
|
|
indexer.projection = Some(if quantized {
|
|
Linear::Quantized(&layers[1])
|
|
} else {
|
|
Linear::Dense(&weights[1])
|
|
});
|
|
}
|
|
if quantized && step == 7 {
|
|
indexer.projection = Some(Linear::Quantized(&layers[2]));
|
|
}
|
|
if quantized && step == 8 {
|
|
indexer.projection = Some(Linear::Quantized(&layers[3]));
|
|
}
|
|
if step == 9 {
|
|
cache.reserve_indexer_capacity(2048, 512, gpu).unwrap();
|
|
}
|
|
if step == 10 {
|
|
assert_eq!(cache.trim(3), 3);
|
|
}
|
|
let pos = cache.kv.offset;
|
|
assert_eq!(pos, n("pos"));
|
|
let workspace = if step >= 2 {
|
|
2 * 1024 * 1024
|
|
} else {
|
|
128 * 1024 * 1024
|
|
};
|
|
if step == 11 {
|
|
let report = indexer.compiled.as_ref().unwrap().1.report();
|
|
let signature = indexer.compiled.as_ref().unwrap().0.clone();
|
|
indexer.q_norm = &invalid;
|
|
assert!(indexer.compiled_core(workspace).is_err());
|
|
assert_eq!(indexer.compiled.as_ref().unwrap().0, signature);
|
|
assert_eq!(indexer.compiled.as_ref().unwrap().1.report(), report);
|
|
indexer.q_norm = &q_norm[1];
|
|
}
|
|
if step == 12 {
|
|
// A new Rust handle of the same descriptor is not a new leaf.
|
|
indexer.q_norm = &q_norm[2];
|
|
}
|
|
let hidden = Array::new(
|
|
&[1, rows, 128],
|
|
Dtype::BF16,
|
|
pattern(rows as u32 * 128, 300 + step as u32),
|
|
)
|
|
.unwrap();
|
|
let supplied = c["supplied"].as_bool().unwrap();
|
|
let qk = supplied.then(|| {
|
|
let a = Array::new(
|
|
&[1, rows, 704],
|
|
Dtype::BF16,
|
|
pattern(rows as u32 * 704, 320 + step as u32),
|
|
)
|
|
.unwrap();
|
|
ops::slice(&a, &[0; 3], &[1, rows, 640], &[1; 3], gpu).unwrap()
|
|
});
|
|
if cache.pooled.is_some() {
|
|
cache.pooled_f32_view(cache.pooled_len, gpu).unwrap();
|
|
}
|
|
let old = [cache.raw_keys.clone(), cache.pooled.clone()];
|
|
let before = old
|
|
.iter()
|
|
.flatten()
|
|
.map(|a| a.borrow().clone())
|
|
.collect::<Vec<_>>();
|
|
ops::evaluate(&streams, &before, gpu, false).unwrap();
|
|
drop(before);
|
|
let old = if retain {
|
|
Some(old)
|
|
} else {
|
|
drop(old);
|
|
None
|
|
};
|
|
let mode = match c["mode"].as_str().unwrap() {
|
|
"update_only" => QsaCompiledMode::UpdateOnly,
|
|
"blocks" => QsaCompiledMode::Blocks,
|
|
"prefill_blocks" => QsaCompiledMode::PrefillBlocks,
|
|
"row_tokens" => QsaCompiledMode::RowTokens,
|
|
"dense_mask" => QsaCompiledMode::DenseMask,
|
|
_ => panic!("mode"),
|
|
};
|
|
let engagement = attention::qsa_prefill_engagement()
|
|
.unwrap()
|
|
.get("compiled_selector")
|
|
.copied()
|
|
.unwrap_or(0);
|
|
let (result, dispatches) = capture_dispatches(|| {
|
|
indexer
|
|
.call_compiled(
|
|
&hidden,
|
|
pos,
|
|
&mut cache,
|
|
qk.as_ref(),
|
|
mode,
|
|
workspace,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
});
|
|
assert!(dispatches.is_empty(), "cache commit must stay lazy");
|
|
assert_eq!(cache.kv.offset, pos);
|
|
assert_eq!(cache.pooled_len, n("pooled_len"));
|
|
assert!(cache.pooled_f32_t.is_none());
|
|
assert_eq!(
|
|
indexer.compiled.as_ref().unwrap().1.report(),
|
|
c["report"],
|
|
"{step}/{quantized}"
|
|
);
|
|
assert_eq!(
|
|
attention::qsa_prefill_engagement()
|
|
.unwrap()
|
|
.get("compiled_selector")
|
|
.copied()
|
|
.unwrap_or(0),
|
|
engagement + u64::from(mode == QsaCompiledMode::PrefillBlocks)
|
|
);
|
|
let (kind, mut outputs, tail) = match result {
|
|
None => ("none", vec![], None),
|
|
Some(EagerSelection::FlashBlocks(a, tail)) => {
|
|
("flash", vec![a], Some(tail))
|
|
}
|
|
Some(EagerSelection::PrefillBlocks(a, b)) => {
|
|
("flash_prefill", vec![a, b], None)
|
|
}
|
|
Some(EagerSelection::RowTokens(a, b)) => ("gather_rows", vec![a, b], None),
|
|
Some(EagerSelection::Dense(a)) => ("dense", vec![a], None),
|
|
_ => panic!("unexpected return contract"),
|
|
};
|
|
assert_eq!(kind, c["kind"].as_str().unwrap());
|
|
assert_eq!(tail, c["tail"].as_i64().map(|v| v as i32));
|
|
outputs.extend([
|
|
cache.raw_keys.as_ref().unwrap().borrow().clone(),
|
|
cache.pooled.as_ref().unwrap().borrow().clone(),
|
|
]);
|
|
if let Some(old) = old {
|
|
for (old, new) in old.iter().zip([&cache.raw_keys, &cache.pooled]) {
|
|
if let Some(old) = old {
|
|
assert!(!std::rc::Rc::ptr_eq(old, new.as_ref().unwrap()));
|
|
outputs.push(old.borrow().clone());
|
|
}
|
|
}
|
|
}
|
|
for (a, shape) in outputs.iter().zip(c["shapes"].as_array().unwrap()) {
|
|
assert_eq!(serde_json::json!(a.layout().shape()), *shape);
|
|
}
|
|
let outputs = outputs
|
|
.iter()
|
|
.map(|a| {
|
|
ops::contiguous(
|
|
&ops::astype(a, Dtype::F32, false, gpu).unwrap(),
|
|
false,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
})
|
|
.collect::<Vec<_>>();
|
|
ops::evaluate(&streams, &outputs, gpu, false).unwrap();
|
|
for (leaf, a) in outputs.iter().enumerate() {
|
|
let mut bytes = vec![0; a.layout().nbytes()];
|
|
a.buffer().read(a.offset(), &mut bytes).unwrap();
|
|
let hash = Sha256::digest(bytes)
|
|
.iter()
|
|
.map(|b| format!("{b:02x}"))
|
|
.collect::<String>();
|
|
assert_eq!(
|
|
hash,
|
|
c["sha256"][leaf].as_str().unwrap(),
|
|
"step{step} quant{quantized} retain{retain} leaf{leaf}"
|
|
);
|
|
}
|
|
// A real attention invocation advances this after indexer commit.
|
|
cache.kv.offset = pos + rows;
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple Metal, original stateful QSAIndexer receipts, no model weights"]
|
|
fn mtplx_qsa_indexer_prefix_matches_reference() {
|
|
use super::super::{
|
|
gpu::Context,
|
|
qwen_mtplx_tests::{capture_dispatches, pattern, verify_typed},
|
|
};
|
|
use super::{
|
|
QsaEagerOptions, allocator,
|
|
attention::{self, EagerSelection},
|
|
configure_sources, mlp,
|
|
stream::{Device, Streams},
|
|
};
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
|
|
let gpu = streams.default_stream(Device::Gpu).unwrap();
|
|
let tf32 = std::env::var("MLX_ENABLE_TF32").ok().as_deref() != Some("0");
|
|
let cases = include_str!("../../../../tests/fixtures/mtplx-qsa-indexer-prefix.jsonl")
|
|
.lines()
|
|
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
|
|
.filter(|case| case["tf32"].as_bool() == Some(tf32))
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(cases.len(), 153);
|
|
for prepared in [false, true] {
|
|
for tag in ["bfloat16", "float16", "float32"] {
|
|
let dtype = match tag {
|
|
"bfloat16" => Dtype::BF16,
|
|
"float16" => Dtype::F16,
|
|
_ => Dtype::F32,
|
|
};
|
|
for profile in 0..if dtype == Dtype::BF16 { 7 } else { 5 } {
|
|
let sequence = cases
|
|
.iter()
|
|
.filter(|c| c["dtype"] == tag && c["profile"] == profile)
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(sequence.len(), 9);
|
|
let first = sequence[0];
|
|
let dim = first["dim"].as_i64().unwrap() as i32;
|
|
let rot = first["rot"].as_i64().unwrap() as i32;
|
|
let fused = first["fused"].as_bool().unwrap();
|
|
let kind = first["kind"].as_str().unwrap();
|
|
let source = |shape: &[i32], salt: u32, dtype: Dtype| {
|
|
let count = shape.iter().product::<i32>();
|
|
ops::astype(
|
|
&Array::new(shape, Dtype::BF16, pattern(count as u32, salt)).unwrap(),
|
|
dtype,
|
|
false,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
};
|
|
let weight = source(&[5 * dim, 64], 187, dtype);
|
|
let quant = match kind {
|
|
"q4" => Some(mlp::fixture_linear(5 * dim as u32, 64, 4, 64, 187)),
|
|
"q8" => Some(mlp::fixture_linear(5 * dim as u32, 64, 8, 64, 187)),
|
|
_ => None,
|
|
};
|
|
let norm_dtype = if kind == "promoted" {
|
|
Dtype::F32
|
|
} else {
|
|
dtype
|
|
};
|
|
let q_norm = source(&[dim], 185, norm_dtype);
|
|
let k_norm = source(&[dim], 186, norm_dtype);
|
|
let frequencies = (0..rot / 2)
|
|
.flat_map(|i| ((1. / f64::from(i + 1)) as f32).to_le_bytes())
|
|
.collect::<Vec<_>>();
|
|
let inv_freq = Array::new(
|
|
&[rot / 2],
|
|
Dtype::F32,
|
|
super::scalar_buffer(&frequencies).unwrap(),
|
|
)
|
|
.unwrap();
|
|
let indexer = Indexer {
|
|
projection: Some(
|
|
quant
|
|
.as_ref()
|
|
.map_or(Linear::Dense(&weight), Linear::Quantized),
|
|
),
|
|
heads: 4,
|
|
kv_heads: 1,
|
|
head_dim: dim,
|
|
block_topk: 512,
|
|
ratio: 4,
|
|
q_norm: &q_norm,
|
|
k_norm: &k_norm,
|
|
inv_freq: &inv_freq,
|
|
mrope_axes: None,
|
|
eps: 1e-6,
|
|
scaling: 1.125,
|
|
scratch: 32 * 1024 * 1024,
|
|
compiled: None,
|
|
};
|
|
let options = QsaEagerOptions {
|
|
large_prefill: false,
|
|
flash: false,
|
|
gather_decode: true,
|
|
gather: true,
|
|
gather_max_rows: 8,
|
|
gather_min_context: 0,
|
|
};
|
|
let mut cache = QsaCache::new(4);
|
|
cache.reserve_indexer_capacity(256, 256, gpu).unwrap();
|
|
for case in sequence {
|
|
super::eval::execution_tests::observed(|| {
|
|
let num = |s: &str| case[s].as_i64().unwrap() as i32;
|
|
let rows = num("rows");
|
|
let step = num("step");
|
|
if num("trim") > 0 {
|
|
assert_eq!(cache.trim(num("trim")), num("trim"));
|
|
}
|
|
let pos = cache.kv.offset;
|
|
assert_eq!(pos, num("pos"));
|
|
let old = cache.pooled_len;
|
|
let hidden = source(&[1, rows, 64], 190 + step as u32 * 7, dtype);
|
|
let qk = if kind == "supplied" {
|
|
Some(
|
|
ops::slice(
|
|
&source(&[1, rows, 5 * dim * 2], 191 + step as u32 * 7, dtype),
|
|
&[0; 3],
|
|
&[1, rows, 5 * dim * 2],
|
|
&[1, 1, 2],
|
|
gpu,
|
|
)
|
|
.unwrap(),
|
|
)
|
|
} else {
|
|
None
|
|
};
|
|
if prepared {
|
|
let mut inputs = vec![
|
|
hidden.clone(),
|
|
q_norm.clone(),
|
|
k_norm.clone(),
|
|
weight.clone(),
|
|
];
|
|
if let Some(qk) = &qk {
|
|
inputs.push(qk.clone());
|
|
}
|
|
ops::evaluate(&streams, &inputs, gpu, false).unwrap();
|
|
}
|
|
let (prefix, calls) = capture_dispatches(|| {
|
|
indexer
|
|
.prepare_rows(
|
|
&hidden,
|
|
pos,
|
|
&mut cache,
|
|
qk.as_ref(),
|
|
fused,
|
|
None,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
});
|
|
assert!(calls.is_empty(), "QSA prefix must remain lazy");
|
|
assert_eq!(cache.kv.offset, num("offset"));
|
|
assert_eq!(cache.pooled_len, num("pooled_len"));
|
|
let mut outputs = vec![cache.raw_keys.as_ref().unwrap().borrow().clone()];
|
|
if let Some(pooled) = &cache.pooled {
|
|
outputs.extend([
|
|
pooled.borrow().clone(),
|
|
cache.pooled_f32_t.as_ref().unwrap().borrow().clone(),
|
|
]);
|
|
}
|
|
let lane = if let Some(prefix) = prefix {
|
|
assert_eq!(prefix.total, pos + rows);
|
|
assert_eq!(prefix.pooled.layout().shape(), [1, cache.pooled_len, dim]);
|
|
let selection = attention::select_eager(
|
|
&prefix.query,
|
|
pos,
|
|
&mut cache,
|
|
prefix.total,
|
|
512,
|
|
0,
|
|
&options,
|
|
&streams,
|
|
gpu,
|
|
)
|
|
.unwrap();
|
|
match selection {
|
|
EagerSelection::RowTokens(ids, valid) => {
|
|
outputs.extend([ids, valid]);
|
|
"gather_rows"
|
|
}
|
|
EagerSelection::DecodeTokens(ids) => {
|
|
outputs.push(ids);
|
|
"decode_tokens"
|
|
}
|
|
EagerSelection::Dense(mask) => {
|
|
outputs.push(mask);
|
|
"dense"
|
|
}
|
|
_ => panic!("unexpected eager route"),
|
|
}
|
|
} else {
|
|
"none"
|
|
};
|
|
assert_eq!(lane, case["lane"].as_str().unwrap());
|
|
for (i, a) in outputs.iter().enumerate() {
|
|
let shape = case["shapes"][i]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.map(|d| d.as_i64().unwrap() as i32)
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(a.layout().shape(), shape, "{} output{i}", case["kernel"]);
|
|
}
|
|
drop((hidden, qk));
|
|
let outputs = outputs
|
|
.iter()
|
|
// NumPy's reference tobytes is logical row order;
|
|
// the cache's F32 mirror is a transposed view.
|
|
.map(|a| {
|
|
ops::contiguous(
|
|
&ops::astype(a, Dtype::F32, false, gpu).unwrap(),
|
|
false,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
})
|
|
.collect::<Vec<_>>();
|
|
let (_, calls) = capture_dispatches(|| {
|
|
ops::evaluate(&streams, &outputs, gpu, false).unwrap()
|
|
});
|
|
let custom = fused && dim <= 128 && norm_dtype == dtype;
|
|
assert_eq!(
|
|
calls
|
|
.iter()
|
|
.filter(|c| c.0.contains("qsa_prepare_q"))
|
|
.count(),
|
|
usize::from(custom && lane != "none")
|
|
);
|
|
assert_eq!(
|
|
calls.iter().filter(|c| c.0.contains("qsa_pool_k")).count(),
|
|
usize::from(custom && cache.pooled_len > old)
|
|
);
|
|
let buffers = outputs.iter().map(Array::buffer).collect::<Vec<_>>();
|
|
let checks = buffers
|
|
.iter()
|
|
.zip(&outputs)
|
|
.map(|(b, a)| (&**b, a.layout().size() as u32, 4))
|
|
.collect::<Vec<_>>();
|
|
verify_typed(case["kernel"].as_str().unwrap(), &checks);
|
|
cache.kv.offset = pos + rows; // Attention's frontier advance.
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|