hugo/DS4Server
1
0 Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Preserve image-aware sparse QSA for Qwen vision

Browse Source
This commit is contained in:
Georg Bauer committed 2026-10-04 14:03:12 +02:00
1 parent a8cad88100
commit dccf81971f
26 files changed
+1215 -44
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/engine/metal/qwen_mtplx/qsa_indexer.rs
+79 -20
View File
@@ -18,6 +18,7 @@ pub(super) struct Indexer<'a> {
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,
@@ -31,6 +32,10 @@ pub(super) struct PreparedRows {
pub(super) total: i32,
}
#[cfg(test)]
#[path = "issue107_indexer.rs"]
mod issue107;
impl<'a> Indexer<'a> {
#[expect(
clippy::too_many_arguments,
@@ -195,6 +200,7 @@ impl<'a> Indexer<'a> {
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::{
@@ -221,16 +227,21 @@ impl<'a> Indexer<'a> {
self.ratio as u32,
self.block_topk as u32,
);
if self.compiled_supported(
qk_rows.unwrap_or(hidden),
cache,
pos,
qk_rows.is_some(),
decode,
mode,
options,
stream,
) {
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,
@@ -244,7 +255,7 @@ impl<'a> Indexer<'a> {
.map(|a| (a, "compiled"));
}
let Some(prepared) =
self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, stream)?
self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, vision, stream)?
else {
return Ok((None, "update_only"));
};
@@ -573,9 +584,13 @@ impl<'a> Indexer<'a> {
q: &Array,
pos: i32,
fused: bool,
vision: Option<(Option<&Array>, i32)>,
stream: Stream,
) -> Result<Array, String> {
if fused && qsa_prepare::supported(q, self.q_norm, self.inv_freq, 4, stream) {
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,
@@ -588,7 +603,18 @@ impl<'a> Indexer<'a> {
}
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) = rope::cos_sin(&positions, self.inv_freq, self.scaling, 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)
}
@@ -598,9 +624,13 @@ impl<'a> Indexer<'a> {
old: i32,
total: i32,
fused: bool,
vision: Option<(Option<&Array>, i32)>,
stream: Stream,
) -> Result<Array, String> {
if fused && qsa_prepare::supported(fresh, self.k_norm, self.inv_freq, 3, stream) {
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,
@@ -632,7 +662,18 @@ impl<'a> Indexer<'a> {
Binary::Multiply,
stream,
)?;
let (cos, sin) = rope::cos_sin(&positions, self.inv_freq, self.scaling, 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)?,
@@ -650,6 +691,7 @@ impl<'a> Indexer<'a> {
cache: &mut QsaCache,
total: i32,
fused: bool,
vision: Option<(Option<&Array>, i32)>,
stream: Stream,
) -> Result<Option<Array>, String> {
let nb = total / self.ratio;
@@ -669,7 +711,7 @@ impl<'a> Indexer<'a> {
)?;
drop(raw);
cache.write_pooled(
&self.pool_keys(&fresh, old, nb, fused, stream)?,
&self.pool_keys(&fresh, old, nb, fused, vision, stream)?,
old,
nb,
stream,
@@ -689,6 +731,10 @@ impl<'a> Indexer<'a> {
/// 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,
@@ -696,6 +742,7 @@ impl<'a> Indexer<'a> {
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();
@@ -733,9 +780,9 @@ impl<'a> Indexer<'a> {
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, stream)?;
let query = self.prepare_queries(&query, pos, fused, vision, stream)?;
cache.write_raw(&keys, stream)?;
let pooled = self.extend_pooled(cache, total, fused, stream)?;
let pooled = self.extend_pooled(cache, total, fused, vision, stream)?;
if total / self.ratio <= self.block_topk {
return Ok(None);
}
@@ -793,6 +840,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() {
q_norm: &q_norm,
k_norm: &k_norm,
inv_freq: &inv_freq,
mrope_axes: None,
eps: 1e-6,
scaling: 1.,
scratch: 32 * 1024 * 1024,
@@ -835,7 +883,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() {
.write_raw(&source(&[1, prefix, 128], 88), gpu)
.unwrap();
indexer
.extend_pooled(&mut cache, prefix, options.fused, gpu)
.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();
@@ -875,6 +923,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() {
qk.as_ref(),
&options,
&streams,
None,
gpu,
)
.unwrap()
@@ -1111,6 +1160,7 @@ fn mtplx_qsa_indexer_compiled_commit_matches_reference() {
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,
@@ -1399,6 +1449,7 @@ fn mtplx_qsa_indexer_prefix_matches_reference() {
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,
@@ -1454,7 +1505,15 @@ fn mtplx_qsa_indexer_prefix_matches_reference() {
}
let (prefix, calls) = capture_dispatches(|| {
indexer
.prepare_rows(&hidden, pos, &mut cache, qk.as_ref(), fused, gpu)
.prepare_rows(
&hidden,
pos,
&mut cache,
qk.as_ref(),
fused,
None,
gpu,
)
.unwrap()
});
assert!(calls.is_empty(), "QSA prefix must remain lazy");