From b74981ee7ae751213a6aa0318c826b82289d822f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 18:19:09 +0900 Subject: [PATCH 1/7] refactor: move KV attention dispatch into the cache (#2171 Phase 4a) Add `KVCache::attend` and `cache::attend_batched` in `src/lib/mlxcel-core/src/cache/attend.rs`: the cache appends the step's K/V and picks the kernel from the storage behind it. A pool-backed cache runs the #899 pooled paged entry at a single unmasked token (fused v2 or the ADR 0001 gather fallback, chosen inside) and the pool intercept plus SDPA otherwise; a dense cache runs the Turbo dequant-first variants its mode and env gates already select, or fused SDPA (causal helper for unmasked prefill, masked SDPA otherwise). A masked step never takes the paged single-token kernel, matching the guard the models carried. Qwen3 and Llama 3 (and Qwen2, Qwen2.5, Helium, the VLM text backbones that reuse their attention) stop branching on `is_paged_backed()` and stop reading `DecodeBatchContext`; their `forward_split_attention`, block and model batched forwards drop the context parameter and the trait default handles `forward_batched_with_context`. The dense-pointer paged route over dense `keys`/`values` is removed from those two forwards: it was unreachable there, since an FP16 paged layout gets pool-backed caches and Int8 or Turbo paged layouts fail its FP16 check. The kernels stay for the model-owned families (Gemma 3, Llama 4, `model_owned`), which still read the context; ADR 0008 records the rule, the dense single-sequence storage default, and what is left for Phase 4b. `MlaLatentCache::attend` gives the MLA latent cache the same entry over its own absorbed attention. In a batch, rows that the pooled launch declines now go through `attend` per row, so batched decode over Turbo caches takes the same dequant-first variants as single-sequence decode (same exact attention, different op order from the old full-dequant loop); FP16 dense and pool-backed routes are op for op what the models ran, pinned by `cache::attend::attend_tests` and `mla::decode_tests`. Refs #2171, #2166, ADR 0007, ADR 0008. --- ...0008-kv-attention-dispatch-in-the-cache.md | 51 +++ docs/adr/README.md | 1 + docs/turbo-kv-cache.md | 4 +- src/lib/mlxcel-core/src/cache.rs | 14 +- src/lib/mlxcel-core/src/cache/attend.rs | 238 +++++++++++++ src/lib/mlxcel-core/src/cache/attend_tests.rs | 326 ++++++++++++++++++ src/lib/mlxcel-core/src/generate.rs | 10 +- src/lib/mlxcel-core/src/mla/cache.rs | 37 ++ src/lib/mlxcel-core/src/mla/decode_tests.rs | 1 + src/models/helium.rs | 16 +- src/models/llama3.rs | 278 ++------------- src/models/qwen3.rs | 273 ++------------- src/server/batch/scheduler/decode_tick.rs | 5 + 13 files changed, 735 insertions(+), 519 deletions(-) create mode 100644 docs/adr/0008-kv-attention-dispatch-in-the-cache.md create mode 100644 src/lib/mlxcel-core/src/cache/attend.rs create mode 100644 src/lib/mlxcel-core/src/cache/attend_tests.rs diff --git a/docs/adr/0008-kv-attention-dispatch-in-the-cache.md b/docs/adr/0008-kv-attention-dispatch-in-the-cache.md new file mode 100644 index 000000000..e87d5c855 --- /dev/null +++ b/docs/adr/0008-kv-attention-dispatch-in-the-cache.md @@ -0,0 +1,51 @@ +# ADR 0008: Attention dispatch is a property of the KV cache + +**Status:** Accepted (2026-10-07). Epic #2166 Phase 4a (issue #2171). Amends [ADR 0007](0007-unified-batch-native-engine.md): it fills in the "KV storage" row of the component table (one KV interface over dense and paged storage) and records the kernel-selection rule and the single-sequence storage default that row left to this phase. ADR 0007's decision table stands unchanged. + +## Context + +The CLI and the server ran different attention kernels for the same model and the same request, and the choice was made in each model forward by the call site. `src/models/qwen3.rs` and `src/models/llama3.rs` each carried `l == 1 && mask.is_none() && cache.is_paged_backed()` in the single-sequence forward and a batched twin that read a scheduler-supplied `DecodeBatchContext` (`is_paged_decode()`, `paged_block_size`, `use_native_paged_kernel`). Two families could disagree on when the paged kernel applies, a storage-policy change meant editing model files, and a numerics fix in one kernel (#2128, #2162 in cuDNN SDPA) could reach one front end and not the other. ADR 0007's "same input, same output" invariant cannot hold across B=1 dense and B=1 paged until the storage behind a cache, not the model, picks the kernel. + +Measured inputs (ADR 0007, GB10, 2026-10-07): a lone sequence on the paged backend decodes 6.5 to 23 percent slower than on dense storage in every cell (Qwen3-1.7B and Llama-3.2-1B, 256 and 8192-token prompts), while paged TTFT at 8192 tokens was 8 percent lower on Llama and unchanged on Qwen3. + +## Decision + +### One entry, storage decides + +`KVCache::attend(&mut self, q, k, v, scale, mask)` appends the step's K/V and runs attention through the kernel that suits the storage behind the cache; `cache::attend_batched(q, k, v, &mut [&mut KVCache], scale, mask)` is the batched counterpart over one layer's caches. Both live in `src/lib/mlxcel-core/src/cache/attend.rs`. The models hand over projected, rotated Q/K/V and read nothing back to choose a kernel. + +The rule: + +| storage behind the cache | single-token, unmasked step | any other step | +|---|---|---| +| pool-backed paged (`KVCache::new_paged`, what `CachePool` wires for a sequence on the paged backend) | `paged_batch_decode_attention`: the fused v2 kernel above the token floors, else the ADR 0001 gather-then-SDPA fallback, both chosen inside that entry | pool append, gathered visible window, fused SDPA | +| dense Turbo4Asym / Turbo4 / Turbo4Delegated | the dequant-first SDPA variants the cache mode and its env gates already select (ADR 0002) | dense update, full dequant, SDPA | +| dense FP16 / Int8 | dense update, fused SDPA | dense update, causal SDPA (no mask) or masked SDPA | + +A masked step (prefill, speculative or MTP verify) never takes the paged single-token kernel. A family with `supports_paged_decode_backend() == false` never gets a pool-backed cache, so it always runs dense. In a batch, a single-token unmasked step over pool-backed caches is one whole-batch launch; any other shape, and any batch the pooled entry declines before writing, runs `attend` per row, so a row decodes through the same route whether it was scheduled alone or in a batch. + +Dispatch is a match on the cache's own fields. A `trait KvStorage` with a virtual call per layer per token was rejected, as ADR 0004 rejected per-op dynamic dispatch on the hot loop; the cache already resolves its quantization modes the same way. + +### The storage policy is set where the cache is built + +`CachePool::allocate_with_layout` decides storage once per sequence: pool-backed caches for a dense-natural family on the paged backend with an FP16 layout, dense caches otherwise. That allocation is the per-sequence storage policy. For the families on `attend`, `DecodeBatchContext` no longer selects a kernel; the scheduler still builds it for the model-owned families (Gemma 3, Llama 4, and the `model_owned` dispatch helpers), whose per-sequence state is not a `KVCache` and which keep their dense-pointer paged kernels until they move onto the entry. + +The dense-pointer paged route (`paged_decode_attention_dense_compat` over dense `keys`/`values` with a block table derived from the batch metadata) is removed from the Qwen3 and Llama 3 batched forwards. It was unreachable there: with an FP16 paged layout the pool wires pool-backed caches, and Int8 or Turbo paged layouts fail its FP16 check. The kernels stay for the model-owned families. + +### Single-sequence storage default + +Dense is the single-sequence KV storage, as ADR 0007 measured. It applies wherever the engine has one sequence: the CLI, `--max-batch-size 1`, and the Phase 4b/5 engine when a request runs alone. The server's storage policy is unchanged in this phase: `--decode-storage-backend` with `auto` resolving to paged for batched serving (`max_batch_size > 1`, a batching family that supports the paged backend), so a lone request on a batched server still decodes on the paged pool. Migrating a lone sequence between storages as the batch grows and shrinks is not in scope; it would need a detach-and-adopt across storages per transition, and ADR 0007's numbers do not yet say whether the win at B=1 pays for that copy. + +## Consequences + +- `rg 'is_paged_backed\(\)|is_paged_decode\(\)' src/models` returns nothing for Qwen3, Llama 3 and the families that reuse their attention (Qwen2, Qwen2.5, Helium, the VLM text backbones). The paged-vs-dense decision has one home, so the ADR 0007 parity harness compares one dispatch under two storages. +- Batched decode over Turbo caches now takes the same dequant-first variants as single-sequence decode instead of a full dequant per step; the result is the same exact attention through a different op order, so it is not bit-identical to the previous batched route. Pool-backed and FP16 dense routes are op-for-op what the models ran before, pinned by `cache::attend::attend_tests`. +- MLA latent caches keep their own attention and expose the same entry: `MlaLatentCache::attend` appends the step's latent and rope rows and runs the absorbed step (`mla::decode_tests` pins it). +- `RotatingKVCache` (Gemma 3's sliding window) is untouched; its decode undo log (#2182) and the model-owned families are the follow-up this phase leaves for 4b. +- Prompt-cache detach and adopt, the disaggregated handoff (`is_paged_backed()` in `src/distributed/disaggregated/handoff_impl.rs`) and `CachePool::sync_paged_state_with_dense` are unchanged: they read storage, which is still the cache's property. + +## References + +- Epic #2166, issue #2171, and the Phase 0 baseline in `docs/benchmark_results/unified-engine-baseline-gb10-2026-10-07.md`. +- [ADR 0001](0001-paged-attention-gather-vs-fused-kernel.md) (gather fallback), [ADR 0002](0002-turbo-kv-split-dequant-vs-fused.md) (Turbo dequant-first), [ADR 0004](0004-compute-backend-session-seam-and-stablehlo-family.md) (no per-op dynamic dispatch), [ADR 0007](0007-unified-batch-native-engine.md) (the engine this phase serves). +- `src/lib/mlxcel-core/src/cache/attend.rs`, `src/lib/mlxcel-core/src/cache/paged_batch_decode.rs` (#899), `src/lib/mlxcel-core/src/cache.rs` (`CachePool::allocate_with_layout`). diff --git a/docs/adr/README.md b/docs/adr/README.md index c34528a69..38a6b1332 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -11,3 +11,4 @@ This directory holds Architecture Decision Records (ADRs) for mlxcel. Each ADR c - [ADR 0005: Mixed prefill/decode step execution](0005-mixed-prefill-decode-step-execution.md). Accepted. Rejects ragged model-level mixing and kernel-level mixing (#908, epic #909) after finding that the tick alternation they were meant to replace does not exist: a chunked prefill parked next to a live decode batch yields every tick and starves until the batch drains, so decode streams are never blocked by prefill and the recoverable work is bounded to one decode step's linear layers per chunk. Extracts the policy into `src/server/batch/tick_policy.rs` so the tests stop asserting against a copy, ships the scheduling half behind `MLXCEL_MIXED_STEP` as a measurement instrument, and splits the starvation fix into a follow-up. That follow-up shipped as issue #1011: `--prefill-grant-interval N` (default 16) grants the parked prefill one tick after every N contended decode ticks, bounding a `C`-chunk prompt at `C * (N + 1)` ticks, and fills the ADR's open ITL cell. - [ADR 0006: SafeTensors shard discovery falls back to a directory glob past a stale index](0006-safetensors-shard-discovery-globs-past-a-stale-index.md). Accepted. Records why `collect_shard_paths` globs the directory when `model.safetensors.index.json` names shards that are not on disk, rather than failing the load: an audit of the affected mlx-community repositories (`scripts/audit_hf_index.py`) finds the published index is the pre-quantization one to the byte, that none of the shards it names exist, and that twelve such checkpoints are in the benchmark set. Names the pipeline-parallel partial loader as the path the tolerance does not cover. - [ADR 0007: One batch-native decode engine for the CLI and the server](0007-unified-batch-native-engine.md). Accepted. Epic #2166 Phase 0: one engine in `mlxcel-core` (open, prefill, batched step, close; a single sequence is a batch of one) with one KV interface over dense and paged storage, one sampling step and one finish step, `BatchScheduler` on top, `mlxcel run` and chat as in-process clients of the server engine, and a thin raw `generate`, after llama.cpp's `llama_decode(llama_batch)` / `llama-cli`-on-`llama-server` design. Supersedes ADR 0004's deferral of the KV, batching and scheduler abstraction. Holds the CLI-vs-server defaults decision table (semantic rows decided with reasons, speed rows decided by measurement), the single-stream regression threshold, and the baseline tools `make engine-parity` and `make bench-engine`. +- [ADR 0008: Attention dispatch is a property of the KV cache](0008-kv-attention-dispatch-in-the-cache.md). Accepted. Epic #2166 Phase 4a (#2171), amending ADR 0007's KV storage row: `KVCache::attend` and `cache::attend_batched` append the step's K/V and pick the kernel from the storage behind the cache (pool-backed paged runs the #899 pooled entry at an unmasked single token, dense runs the Turbo dequant-first variants or fused SDPA), so Qwen3, Llama 3 and the families that reuse their attention no longer branch on `is_paged_backed()` or read `DecodeBatchContext`. Records the kernel-selection table, the removal of the unreachable dense-pointer paged route from those forwards, and the dense single-sequence storage default with the storage-migration question left open. diff --git a/docs/turbo-kv-cache.md b/docs/turbo-kv-cache.md index a97d7c4f3..02da68f3b 100644 --- a/docs/turbo-kv-cache.md +++ b/docs/turbo-kv-cache.md @@ -454,8 +454,8 @@ llama3 Fp16 dense-natural-cache allocations at B=1/B=2, plus a CUDA qwen3 arm that requires greedy-token equality and bounded logit RMS with TF32 disabled. That evidence does not claim byte identity for VLM front ends, model-owned caches, Turbo/quantized-KV modes, or every CUDA reduction geometry. The live -batched path uses the native block-table decode kernel -(`DecodeBatchContext::use_native_paged_kernel`, set by the scheduler); a +batched path is the pooled paged decode of #899, dispatched from the cache +itself (`KVCache::attend`, ADR 0008) rather than by the model forward; a gather-then-SDPA path is kept as a correctness reference. At batch 4 the native kernel runs at 276 tok/s for a 512-token prompt and 84 tok/s for a 4096-token prompt, versus 146 and 7.7 tok/s for the gather reference (1.9x and 10.9x). The diff --git a/src/lib/mlxcel-core/src/cache.rs b/src/lib/mlxcel-core/src/cache.rs index ee7d8c9ff..1af414dba 100644 --- a/src/lib/mlxcel-core/src/cache.rs +++ b/src/lib/mlxcel-core/src/cache.rs @@ -75,6 +75,9 @@ //! path and is intended as an opt-in performance comparison, not the default //! compressed-only memory target. +/// One attention entry per cache: storage decides the kernel (issue #2171, +/// ADR 0008). +mod attend; pub mod batch_quant; mod decode_undo; #[cfg(test)] @@ -117,6 +120,7 @@ pub mod turbo; #[path = "cache/turbo_tests.rs"] mod turbo_tests; +pub use attend::attend_batched; pub use batch_quant::{ BatchKvQuantConfig, BatchQuantizedKVCache, BatchTurboQuantKVCache, DEFAULT_KV_GROUP_SIZE, KvQuantScheme, @@ -690,11 +694,11 @@ impl KVCache { /// Whether this cache writes/reads through a shared [`PagedBlockPool`] /// (built via [`Self::new_paged`]) instead of its own dense buffers. /// - /// The batched-decode dispatch in each transformer model uses this to skip - /// the native dense-pointer paged kernel (which reads `keys`/`values`, both - /// `None` here) and fall through to the per-sequence `update_and_fetch` - /// loop, whose pool intercept transparently writes to the pool and gathers - /// the visible window. See the model `forward_split_attention` dispatch. + /// Attention dispatch reads this inside [`Self::attend`] and + /// [`attend_batched`] (#2171, ADR 0008); model forwards no longer do. The + /// remaining readers are storage bookkeeping: `CachePool` detach, adopt + /// and length sync, the MLA latent-cache guard, and the disaggregated + /// handoff. #[inline] pub fn is_paged_backed(&self) -> bool { self.paged_backing.is_some() diff --git a/src/lib/mlxcel-core/src/cache/attend.rs b/src/lib/mlxcel-core/src/cache/attend.rs new file mode 100644 index 000000000..15df85e87 --- /dev/null +++ b/src/lib/mlxcel-core/src/cache/attend.rs @@ -0,0 +1,238 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! One attention entry per KV cache (issue #2171, epic #2166 Phase 4a). +//! +//! Which kernel answers a step is a property of the storage behind the cache, +//! not of the model forward that calls it. Before this module every transformer +//! family that could run on the server's paged pool carried the same branch in +//! its own attention block (`l == 1 && mask.is_none() && cache.is_paged_backed()`, +//! plus a batched twin that read a scheduler-supplied `DecodeBatchContext`), so +//! two families could disagree on when the paged kernel applies and a storage +//! policy change meant editing model files. +//! +//! [`KVCache::attend`] and [`attend_batched`] own that decision now. The rule, +//! recorded in ADR 0008: +//! +//! | storage behind the cache | single-token, unmasked step | any other step | +//! |---|---|---| +//! | pool-backed paged (`new_paged`) | [`paged_batch_decode_attention`]: fused v2 or the ADR 0001 gather fallback, chosen inside | pool append, gathered window, fused SDPA | +//! | dense Turbo4Asym / Turbo4 / Turbo4Delegated | the dequant-first SDPA variants the cache mode and its env gates already select | dense update, full dequant, SDPA | +//! | dense FP16 / Int8 | dense update, fused SDPA | dense update, causal or masked SDPA | +//! +//! A masked step (prefill, speculative or MTP verify) never takes the paged +//! single-token kernel; that is the `mask.is_none()` guard the models used to +//! carry. Dispatch is a match on the cache's own fields rather than a trait +//! object: ADR 0004 rejected per-op dynamic dispatch on the decode hot loop, +//! and the cache already resolves its quantization modes the same way. +//! +//! Storage is chosen once, where the cache is built: `CachePool` wires +//! pool-backed caches for a sequence on the paged backend and dense caches +//! otherwise, so the per-sequence storage policy the scheduler used to carry +//! in `DecodeBatchContext` is the cache itself for these families. + +use cxx::UniquePtr; + +use super::{KVCache, paged_batch_decode_attention, turbo}; +use crate::ffi::{self, MlxArray}; + +impl KVCache { + /// Append this step's K/V and run attention through the kernel that suits + /// the storage behind this cache. + /// + /// `q` is `[B, Hq, L, D]`; `new_keys` and `new_values` are `[B, Hkv, L, D]` + /// with RoPE and any Q/K norm already applied; `mask` is the optional + /// additive mask the model built (`None` for a plain decode step and for + /// causal prefill, which runs through the causal helper). Returns + /// `[B, Hq, L, D]`. + /// + /// Cache mutation is exactly what the pre-existing paths performed: the + /// pooled launch appends through `write_paged`, every dense route goes + /// through [`Self::update`] or [`Self::update_and_fetch`], and nothing here + /// touches `offset` on its own. No host sync and no allocation beyond the + /// shape read. + /// + /// Used by: `models::qwen3::Attention::forward`, + /// `models::llama3::Attention::forward` (and the families that reuse them: + /// Qwen2, Qwen2.5, Helium, the VLM text backbones), [`attend_batched`]. + pub fn attend( + &mut self, + q: &MlxArray, + new_keys: UniquePtr, + new_values: UniquePtr, + scale: f32, + mask: Option<&MlxArray>, + ) -> UniquePtr { + let q_len = query_len(q); + // Pool-backed single-token decode (#899). The batch-1 launch is the + // `[1, H, 1, D]` shape the whole-batch entry point serves, and the + // server reaches it through the single-sequence forward whenever one + // request is decoding. The call declines (`None`) before touching the + // pool for anything it cannot serve, so the dense route below still + // runs its pool intercept (append plus gather) unchanged. + if q_len == 1 + && mask.is_none() + && self.paged_backing.is_some() + && let Some(out) = paged_batch_decode_attention( + q, + &new_keys, + &new_values, + &mut [&mut *self], + scale, + 0.0, + ) + { + return out; + } + self.attend_dense(q, q_len, new_keys, new_values, scale, mask) + } + + /// Dense-storage attention: the Turbo decode variants at a single token, + /// otherwise update, fetch the live window and run SDPA. + fn attend_dense( + &mut self, + q: &MlxArray, + q_len: i32, + new_keys: UniquePtr, + new_values: UniquePtr, + scale: f32, + mask: Option<&MlxArray>, + ) -> UniquePtr { + if q_len == 1 { + // Turbo4Asym (FP16 K, 4-bit V) decodes dequant-first by default: + // exact and 3-6x faster than the lossy sparse-V weighted sum, which + // stays reachable behind `MLXCEL_TURBO4_ASYM_DEQUANT_SDPA=0`. + // Symmetric Turbo4 and Turbo4Delegated mirror mlx-swift-lm's + // dequant-first policy. Each gate is one cached `OnceLock` read. + if turbo::sparse_v::turbo4_asym_dequant_sdpa_enabled() + && self.turbo4_asym_dequant_sdpa_available() + { + return self.update_and_turbo4_asym_dequant_sdpa_attention( + q, new_keys, new_values, scale, mask, + ); + } + if self.sparse_v_available() { + return self + .update_and_sparse_v_attention(q, new_keys, new_values, scale, mask) + .expect("update_and_sparse_v_attention returned None despite sparse_v_available"); + } + if turbo::sparse_v::turbo4_dequant_sdpa_enabled() && self.turbo4_dequant_sdpa_available() + { + return self.update_and_turbo4_dequant_sdpa_attention( + q, new_keys, new_values, scale, mask, + ); + } + if turbo::sparse_v::turbo4_delegated_compressed_attention_enabled() + && self.turbo4_delegated_available() + { + return self.update_and_turbo4_delegated_attention( + q, new_keys, new_values, scale, mask, + ); + } + } + let (cache_k, cache_v) = self.update_and_fetch(new_keys, new_values); + match mask { + // Causal prefill without an explicit mask uses the shared causal + // helper; decode and masked or padded prefill use the unified + // dispatcher, which also accepts a null mask. + None if q_len > 1 => crate::causal_attention(q, &cache_k, &cache_v, scale, 0.0, 0), + _ => { + let mask_ptr = mask.map_or(std::ptr::null(), |m| m as *const MlxArray); + // SAFETY: `mask_ptr` is null or points at `mask`, a borrow that + // outlives this call; `attention_from_ptr` only reads it. + unsafe { + crate::layers::attention_from_ptr( + q, &cache_k, &cache_v, scale, mask_ptr, 0.0, 0, + ) + } + } + } + } +} + +/// Batched counterpart of [`KVCache::attend`]: one layer's caches for every +/// sequence of a decode or batched-prefill step. +/// +/// `q_batched` is `[B, Hq, T, D]`, `k_batched` and `v_batched` are +/// `[B, Hkv, T, D]`, `caches[b]` is sequence `b`'s cache for this layer, and +/// `mask` is the optional `[B, T, S]` additive mask whose row `b` goes to +/// sequence `b`. Returns `[B, Hq, T, D]`; the batch-axis concat commutes with +/// the transpose and reshape the caller applies afterward. +/// +/// A single-token unmasked step over pool-backed caches is one whole-batch +/// launch ([`paged_batch_decode_attention`], which declines before writing +/// when it cannot serve the batch). Every other shape, and a declined batch, +/// goes through [`KVCache::attend`] per sequence, so a row decodes the same way +/// whether it was scheduled alone or in a batch. +/// +/// Used by: `models::qwen3::Attention::forward_split_attention`, +/// `models::llama3::Attention::forward_split_attention`. +pub fn attend_batched( + q_batched: &MlxArray, + k_batched: &MlxArray, + v_batched: &MlxArray, + caches: &mut [&mut KVCache], + scale: f32, + mask: Option<&MlxArray>, +) -> UniquePtr { + let batch = caches.len(); + assert!(batch > 0, "attend_batched: empty batch"); + let seq_len = query_len(q_batched); + if seq_len == 1 + && mask.is_none() + && caches[0].is_paged_backed() + && let Some(out) = + paged_batch_decode_attention(q_batched, k_batched, v_batched, caches, scale, 0.0) + { + return out; + } + + let mut outputs: Vec> = Vec::with_capacity(batch); + for (b, cache) in caches.iter_mut().enumerate() { + let q_b = slice_row(q_batched, b); + let k_b = slice_row(k_batched, b); + let v_b = slice_row(v_batched, b); + let mask_b = mask.map(|m| { + let row = ffi::slice(m, &[b as i32, 0, 0], &[b as i32 + 1, seq_len, i32::MAX]); + ffi::squeeze_axis(&row, 0) + }); + outputs.push(cache.attend(&q_b, k_b, v_b, scale, mask_b.as_deref())); + } + let mut outputs = outputs.into_iter(); + let mut result = outputs.next().expect("a non-empty batch has a first row"); + for out in outputs { + result = crate::concatenate(&result, &out, 0); + } + result +} + +/// Query length of a `[B, H, L, D]` tensor (a shape read, no sync). +#[inline] +fn query_len(q: &MlxArray) -> i32 { + ffi::array_shape(q)[2] +} + +/// Slice row `b` out of a `[B, H, T, D]` tensor, keeping the batch axis. +#[inline] +fn slice_row(arr: &MlxArray, b: usize) -> UniquePtr { + ffi::slice( + arr, + &[b as i32, 0, 0, 0], + &[b as i32 + 1, i32::MAX, i32::MAX, i32::MAX], + ) +} + +#[cfg(test)] +#[path = "attend_tests.rs"] +mod attend_tests; diff --git a/src/lib/mlxcel-core/src/cache/attend_tests.rs b/src/lib/mlxcel-core/src/cache/attend_tests.rs new file mode 100644 index 000000000..71ad8d9ef --- /dev/null +++ b/src/lib/mlxcel-core/src/cache/attend_tests.rs @@ -0,0 +1,326 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Tests for the cache-owned attention entry (issue #2171). +//! +//! The dense half pins [`KVCache::attend`] and [`attend_batched`] against the +//! `update_and_fetch` plus SDPA pair the models used to spell out, op for op, +//! so the outputs are bit-identical. The paged half drives a real +//! [`PagedBlockPool`] through [`KVCache::new_paged`] and checks that an +//! unmasked single-token step is served by the pooled entry (the production +//! counters move) while a masked one is not, and that both agree with the +//! gather-then-SDPA baseline computed from the same post-write pool state. + +use std::cell::RefCell; +use std::rc::Rc; + +use super::*; +use crate::cache::{PagedBlockPool, PagedKvLayout, PagedSequenceState, paged_batch_decode_stats}; +use crate::dtype; + +const PAGE: usize = 32; +const HEADS: i32 = 4; +const KV_HEADS: i32 = 2; +const DIM: i32 = 64; +const SCALE: f32 = 0.125; + +/// xorshift64* in [-1, 1), deterministic so a failure reproduces exactly. +struct Rng(u64); + +impl Rng { + fn new(seed: u64) -> Self { + Self(seed | 1) + } + fn next_f32(&mut self) -> f32 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + ((self.0 >> 40) as f32 / (1u32 << 24) as f32) * 2.0 - 1.0 + } +} + +fn random_array(rng: &mut Rng, shape: &[i32]) -> UniquePtr { + let n: usize = shape.iter().map(|d| *d as usize).product(); + let data: Vec = (0..n).map(|_| rng.next_f32()).collect(); + let f32_arr = ffi::from_slice_f32(&data, shape); + ffi::astype(&f32_arr, dtype::FLOAT16) +} + +fn to_vec_f32(arr: &MlxArray) -> Vec { + let as_f32 = ffi::astype(arr, dtype::FLOAT32); + ffi::eval(&as_f32); + ffi::array_to_raw_bytes(&as_f32) + .chunks_exact(4) + .map(|c| f32::from_ne_bytes(c.try_into().expect("4-byte chunk"))) + .collect() +} + +fn relative_rms(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len(), "outputs must have the same length"); + let mut num = 0.0f64; + let mut den = 0.0f64; + for (x, y) in a.iter().zip(b.iter()) { + num += ((x - y) as f64).powi(2); + den += (*y as f64).powi(2); + } + ((num / den.max(1e-12)) as f32).sqrt() +} + +/// `(q, k, v)` for one step of `len` tokens over `batch` rows. +fn step(rng: &mut Rng, batch: i32, len: i32) -> (UniquePtr, UniquePtr, UniquePtr) { + ( + random_array(rng, &[batch, HEADS, len, DIM]), + random_array(rng, &[batch, KV_HEADS, len, DIM]), + random_array(rng, &[batch, KV_HEADS, len, DIM]), + ) +} + +/// The pair every model used to spell out: dense update, live window, SDPA. +fn reference( + cache: &mut KVCache, + q: &MlxArray, + k: UniquePtr, + v: UniquePtr, + mask: Option<&MlxArray>, +) -> UniquePtr { + let (ck, cv) = cache.update_and_fetch(k, v); + let q_len = ffi::array_shape(q)[2]; + if q_len > 1 && mask.is_none() { + crate::causal_attention(q, &ck, &cv, SCALE, 0.0, 0) + } else { + crate::layers::attention(q, &ck, &cv, SCALE, mask, 0.0, 0) + } +} + +fn fresh_pool( + num_layers: usize, +) -> ( + Rc>, + Vec>>, +) { + let bytes_per_block = PAGE * KV_HEADS as usize * DIM as usize * 2; + let layout = PagedKvLayout::uniform(num_layers, PAGE, bytes_per_block).expect("valid layout"); + let mut pool = PagedBlockPool::new(layout.clone()); + pool.set_slab_blocks(512) + .expect("fresh pool has no storage"); + let states = (0..4) + .map(|_| Rc::new(RefCell::new(PagedSequenceState::new(&layout)))) + .collect(); + (Rc::new(RefCell::new(pool)), states) +} + +#[test] +fn dense_prefill_then_decode_matches_update_and_fetch_plus_sdpa() { + let mut rng = Rng::new(0x2171); + let mut got = KVCache::new(); + let mut want = KVCache::new(); + + // Causal prefill of 8 tokens, no explicit mask. + let (q, k, v) = step(&mut rng, 1, 8); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + let a = got.attend(&q, k, v, SCALE, None); + let b = reference(&mut want, &q, k2, v2, None); + assert_eq!(to_vec_f32(&a), to_vec_f32(&b), "prefill output differs"); + assert_eq!(got.offset, want.offset); + assert_eq!(got.offset, 8); + + // Three decode steps. + for _ in 0..3 { + let (q, k, v) = step(&mut rng, 1, 1); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + let a = got.attend(&q, k, v, SCALE, None); + let b = reference(&mut want, &q, k2, v2, None); + assert_eq!(to_vec_f32(&a), to_vec_f32(&b), "decode output differs"); + assert_eq!(got.offset, want.offset); + } + assert_eq!(got.offset, 11); +} + +#[test] +fn dense_masked_step_matches_masked_sdpa() { + let mut rng = Rng::new(0x2172); + let mut got = KVCache::new(); + let mut want = KVCache::new(); + let (q, k, v) = step(&mut rng, 1, 4); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + got.attend(&q, k, v, SCALE, None); + reference(&mut want, &q, k2, v2, None); + + // A two-token verify step with an explicit causal mask over 6 keys. + let (q, k, v) = step(&mut rng, 1, 2); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + let mask = crate::utils::create_causal_mask(2, 4); + let a = got.attend(&q, k, v, SCALE, Some(&mask)); + let b = reference(&mut want, &q, k2, v2, Some(&mask)); + assert_eq!(to_vec_f32(&a), to_vec_f32(&b), "masked output differs"); + assert_eq!(got.offset, 6); +} + +#[test] +fn batched_dense_matches_per_row_attend() { + let mut rng = Rng::new(0x2173); + let mut c0 = KVCache::new(); + let mut c1 = KVCache::new(); + let mut r0 = KVCache::new(); + let mut r1 = KVCache::new(); + for (c, r) in [(&mut c0, &mut r0), (&mut c1, &mut r1)] { + let (q, k, v) = step(&mut rng, 1, 5); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + c.attend(&q, k, v, SCALE, None); + reference(r, &q, k2, v2, None); + } + + let (q, k, v) = step(&mut rng, 2, 1); + let mut rows = Vec::new(); + for (b, r) in [&mut r0, &mut r1].into_iter().enumerate() { + let q_b = slice_row(&q, b); + let k_b = slice_row(&k, b); + let v_b = slice_row(&v, b); + rows.push(reference(r, &q_b, k_b, v_b, None)); + } + let want = crate::concatenate(&rows[0], &rows[1], 0); + + let mut caches: Vec<&mut KVCache> = vec![&mut c0, &mut c1]; + let got = attend_batched(&q, &k, &v, &mut caches, SCALE, None); + assert_eq!(ffi::array_shape(&got), vec![2, HEADS, 1, DIM]); + assert_eq!(to_vec_f32(&got), to_vec_f32(&want)); + assert_eq!(c0.offset, 6); + assert_eq!(c1.offset, 6); +} + +#[test] +fn batched_masked_rows_get_their_own_mask_row() { + let mut rng = Rng::new(0x2174); + let mut c0 = KVCache::new(); + let mut c1 = KVCache::new(); + let mut r0 = KVCache::new(); + let mut r1 = KVCache::new(); + + // Padded batched prefill: 3 tokens, mask `[B, 3, 3]`, row 1 masks its + // first key so the two rows must not share a mask. + let (q, k, v) = step(&mut rng, 2, 3); + let mut mask_data = vec![0.0f32; 2 * 3 * 3]; + for t in 0..3 { + for s in 0..3 { + if s > t { + mask_data[t * 3 + s] = -1e9; + mask_data[9 + t * 3 + s] = -1e9; + } + } + } + for t in 0..3 { + mask_data[9 + t * 3] = -1e9; + } + let mask = ffi::astype(&ffi::from_slice_f32(&mask_data, &[2, 3, 3]), dtype::FLOAT16); + + let mut rows = Vec::new(); + for (b, r) in [&mut r0, &mut r1].into_iter().enumerate() { + let q_b = slice_row(&q, b); + let k_b = slice_row(&k, b); + let v_b = slice_row(&v, b); + let m_b = ffi::squeeze_axis( + &ffi::slice(&mask, &[b as i32, 0, 0], &[b as i32 + 1, 3, i32::MAX]), + 0, + ); + rows.push(reference(r, &q_b, k_b, v_b, Some(&m_b))); + } + let want = crate::concatenate(&rows[0], &rows[1], 0); + + let mut caches: Vec<&mut KVCache> = vec![&mut c0, &mut c1]; + let got = attend_batched(&q, &k, &v, &mut caches, SCALE, Some(&mask)); + assert_eq!(to_vec_f32(&got), to_vec_f32(&want)); +} + +#[test] +fn paged_cache_single_token_goes_through_the_pooled_entry() { + let (pool, states) = fresh_pool(1); + let mut rng = Rng::new(0x2175); + let mut paged = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let mut baseline = KVCache::new_paged(pool, states[1].clone(), 0); + + // Prefill both sequences with the same 40 tokens (crosses a page). + let (q, k, v) = step(&mut rng, 1, 40); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + let before = paged_batch_decode_stats(); + let a = paged.attend(&q, k, v, SCALE, None); + let b = reference(&mut baseline, &q, k2, v2, None); + let after = paged_batch_decode_stats(); + assert_eq!(to_vec_f32(&a), to_vec_f32(&b), "paged prefill differs"); + assert_eq!( + after.v2_launches + after.gather_fallbacks, + before.v2_launches + before.gather_fallbacks, + "a multi-token step must not take the pooled single-token entry" + ); + assert_eq!(paged.offset, 40); + + // Unmasked decode: the pooled entry serves it (one launch or one gather + // fallback, either way a served step) and agrees with the gather path. + let (q, k, v) = step(&mut rng, 1, 1); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + let before = paged_batch_decode_stats(); + let a = paged.attend(&q, k, v, SCALE, None); + let after = paged_batch_decode_stats(); + assert_eq!( + after.v2_launches + after.gather_fallbacks, + before.v2_launches + before.gather_fallbacks + 1, + "an unmasked single-token step on a pool-backed cache is served by the pooled entry" + ); + let b = reference(&mut baseline, &q, k2, v2, None); + let rms = relative_rms(&to_vec_f32(&a), &to_vec_f32(&b)); + assert!(rms < 2e-2, "paged decode diverged from the gather baseline: rms {rms}"); + assert_eq!(paged.offset, 41); + assert_eq!(baseline.offset, 41); + + // Masked single token (a verify step): never the pooled kernel, and the + // pool still receives the row through the dense route's intercept. + let (q, k, v) = step(&mut rng, 1, 1); + let (k2, v2) = (ffi::copy(&k), ffi::copy(&v)); + let mask = crate::utils::create_causal_mask(1, 41); + let before = paged_batch_decode_stats(); + let a = paged.attend(&q, k, v, SCALE, Some(&mask)); + let after = paged_batch_decode_stats(); + assert_eq!( + after.v2_launches + after.gather_fallbacks, + before.v2_launches + before.gather_fallbacks, + "a masked step must not take the pooled single-token entry" + ); + let b = reference(&mut baseline, &q, k2, v2, Some(&mask)); + assert_eq!(to_vec_f32(&a), to_vec_f32(&b), "masked paged step differs"); + assert_eq!(paged.offset, 42); +} + +#[test] +fn batched_paged_decode_is_one_pooled_launch() { + let (pool, states) = fresh_pool(1); + let mut rng = Rng::new(0x2176); + let mut c0 = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let mut c1 = KVCache::new_paged(pool, states[1].clone(), 0); + for c in [&mut c0, &mut c1] { + let (q, k, v) = step(&mut rng, 1, 12); + c.attend(&q, k, v, SCALE, None); + } + let (q, k, v) = step(&mut rng, 2, 1); + let before = paged_batch_decode_stats(); + let mut caches: Vec<&mut KVCache> = vec![&mut c0, &mut c1]; + let out = attend_batched(&q, &k, &v, &mut caches, SCALE, None); + let after = paged_batch_decode_stats(); + assert_eq!(ffi::array_shape(&out), vec![2, HEADS, 1, DIM]); + assert_eq!( + after.v2_launches + after.gather_fallbacks, + before.v2_launches + before.gather_fallbacks + 1, + "the batch is served by exactly one pooled call" + ); + assert_eq!(c0.offset, 13); + assert_eq!(c1.offset, 13); +} diff --git a/src/lib/mlxcel-core/src/generate.rs b/src/lib/mlxcel-core/src/generate.rs index 2363a089d..48d84abc6 100644 --- a/src/lib/mlxcel-core/src/generate.rs +++ b/src/lib/mlxcel-core/src/generate.rs @@ -1159,7 +1159,15 @@ pub enum DecodeStorageBackend { Paged, } -/// Optional scheduler/runtime context for batched decode dispatch. +/// Per-step storage policy the scheduler hands to batched decode. +/// +/// Since #2171 (ADR 0008) the families whose per-sequence state is a +/// [`crate::cache::KVCache`] (Qwen3, Llama 3 and the families that reuse their +/// attention) read nothing from this: storage is chosen when `CachePool` +/// builds the cache, and [`crate::cache::KVCache::attend`] picks the kernel +/// from the storage behind it. The model-owned families (Gemma 3, Llama 4 and +/// the `model_owned` dispatch helpers) still consume it for their +/// dense-pointer paged kernels. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct DecodeBatchContext { pub storage_backend: DecodeStorageBackend, diff --git a/src/lib/mlxcel-core/src/mla/cache.rs b/src/lib/mlxcel-core/src/mla/cache.rs index 4ddb7554c..debee3000 100644 --- a/src/lib/mlxcel-core/src/mla/cache.rs +++ b/src/lib/mlxcel-core/src/mla/cache.rs @@ -59,6 +59,8 @@ use cxx::UniquePtr; use crate::cache::{KVCache, KVCacheMode}; use crate::ffi::MlxArray; +use crate::mla::absorb::MlaAbsorbedProjections; +use crate::mla::decode::absorbed_decode; use crate::mla::MlaGeometry; /// Bytes one token of the compressed-latent cache costs, per layer. @@ -254,6 +256,41 @@ impl<'a> MlaLatentCache<'a> { self.inner.update_and_fetch(ckv, kpe) } + /// Append one step's latent and rope rows and run absorbed attention over + /// the live window: the MLA analogue of [`KVCache::attend`] (#2171). + /// + /// A latent cache keeps its own attention (`absorbed_decode`: the score is + /// split into the absorbed latent term and the rope term, and the kernel + /// sees the latent as both K and V), so it cannot share the dense + /// dispatch. It shares the entry: the caller hands over this step's rows + /// and the query halves, and never reads the cache back to choose a + /// kernel. + /// + /// `q_nope` is `[B, H, L, qk_nope_head_dim]` (not yet absorbed), `q_pe` is + /// `[B, H, L, qk_rope_head_dim]` (already rotated), `ckv` and `kpe` are + /// this step's `[B, 1, L, kv_lora_rank]` and `[B, 1, L, qk_rope_head_dim]` + /// rows, `mask` is the optional additive causal or padding mask. Returns + /// `[B, H, L, v_head_dim]`. + /// + /// Used by: the MLA families once they move onto the latent cache (#907 + /// Stage 2); `mla::decode_tests` pins it against the decompressed + /// reference. + #[must_use] + #[allow(clippy::too_many_arguments)] + pub fn attend( + &mut self, + q_nope: &MlxArray, + q_pe: &MlxArray, + ckv: UniquePtr, + kpe: UniquePtr, + proj: &MlaAbsorbedProjections, + scale: f32, + mask: Option<&MlxArray>, + ) -> UniquePtr { + let (ckv_all, kpe_all) = self.update_and_fetch(ckv, kpe); + absorbed_decode(q_nope, q_pe, &ckv_all, &kpe_all, proj, scale, mask) + } + /// Bytes per token this cache costs, at the element size it stores. #[must_use] pub const fn bytes_per_token(&self, bytes_per_element: usize) -> usize { diff --git a/src/lib/mlxcel-core/src/mla/decode_tests.rs b/src/lib/mlxcel-core/src/mla/decode_tests.rs index bf864a0bc..0a5fd9bc7 100644 --- a/src/lib/mlxcel-core/src/mla/decode_tests.rs +++ b/src/lib/mlxcel-core/src/mla/decode_tests.rs @@ -26,6 +26,7 @@ //! real DeepSeek model. use super::*; +use crate::cache::KVCache; use crate::dtype; use crate::mla::absorb::MlaAbsorbedProjections; use crate::mla::testkit::{MlaFixture, TINY, max_rel_error, serial, to_vec_f32}; diff --git a/src/models/helium.rs b/src/models/helium.rs index ee4987391..24866f9e0 100644 --- a/src/models/helium.rs +++ b/src/models/helium.rs @@ -50,7 +50,7 @@ //! crossing the cxx bridge is an uncatchable `std::terminate`, not a Rust error, //! so a check that happens at the first forward pass is not a check. -use mlxcel_core::generate::{DecodeBatchContext, LanguageModel}; +use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::KVCache; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; @@ -950,19 +950,7 @@ impl LanguageModel for HeliumModel { batch_caches: &mut [&mut [KVCache]], mask: Option<&MlxArray>, ) -> UniquePtr { - self.inner - .forward_batched_impl(input_ids, batch_caches, mask, None) - } - - fn forward_batched_with_context( - &self, - input_ids: &MlxArray, - batch_caches: &mut [&mut [KVCache]], - mask: Option<&MlxArray>, - context: Option<&DecodeBatchContext>, - ) -> UniquePtr { - self.inner - .forward_batched_impl(input_ids, batch_caches, mask, context) + self.inner.forward_batched_impl(input_ids, batch_caches, mask) } fn supports_batched_prefill(&self) -> bool { diff --git a/src/models/llama3.rs b/src/models/llama3.rs index 84db19fd6..36d09435f 100644 --- a/src/models/llama3.rs +++ b/src/models/llama3.rs @@ -17,10 +17,10 @@ //! This implements the standard Llama architecture for dense models //! like Llama 3.1 8B Instruct. -use mlxcel_core::cache::{BatchedAttentionMetadata, PagedDecodeMetadata, SequenceId}; -use mlxcel_core::generate::{DecodeBatchContext, LanguageModel}; +use mlxcel_core::cache::{BatchedAttentionMetadata, SequenceId}; +use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::{ - FusedQKVLinear, KVCache, KVCacheMode, RMSNorm, UnifiedEmbedding, UnifiedLinear, + FusedQKVLinear, KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear, }; use mlxcel_core::utils::pipeline_hint; use mlxcel_core::weights::WeightMap; @@ -716,82 +716,12 @@ impl Attention { (q, k, v) }; - // Decode-case (l == 1) attention dispatch for the Turbo quantized cache - // modes. Prefill (l > 1) builds the cache from scratch and falls through - // to the standard masked/causal paths below. This module's `Attention` - // is also re-exported as Qwen2 / Qwen2.5's attention, so the routing here - // covers those families too. - // - // Turbo4Asym (FP16 K + 4-bit V) decodes via dequant-first native SDPA by - // default: V is dequantized to FP16 transiently and fed with the FP16 K - // to native SDPA. This is exact and ~3-6x faster than the lossy sparse-V - // weighted-sum path, which stays reachable behind - // `MLXCEL_TURBO4_ASYM_DEQUANT_SDPA=0` for A/B and fallback. Symmetric - // Turbo4 and Turbo4Delegated mirror mlx-swift-lm's dequant-first policy. - // Each gate is parsed once and cached in a `OnceLock`. - let use_turbo4_asym_dequant_sdpa = - mlxcel_core::cache::turbo::sparse_v::turbo4_asym_dequant_sdpa_enabled(); - let use_delegated_compressed = - mlxcel_core::cache::turbo::sparse_v::turbo4_delegated_compressed_attention_enabled(); - let use_turbo4_dequant_sdpa = - mlxcel_core::cache::turbo::sparse_v::turbo4_dequant_sdpa_enabled(); - // Single-sequence pooled paged decode (#899). The server takes this - // path whenever the active batch is one request: `decode_single_step` - // calls the single-sequence `forward`, never `forward_split_attention`. - // Without this branch the two single-sequence scenarios in the issue's - // benchmark matrix (16K and 32K) could never reach the fused kernel, and - // neither could any moment in a batched run where only one request is - // still decoding. The batch-1 launch is exactly the `[1, H, 1, D]` shape - // the whole-batch entry point already serves. - let pooled_single = if l == 1 && mask.is_none() && cache.is_paged_backed() { - mlxcel_core::cache::paged_batch_decode_attention( - &q, - &k, - &v, - &mut [&mut *cache], - self.scale, - 0.0, - ) - } else { - None - }; - - let attn_out = if let Some(out) = pooled_single { - out - } else if l == 1 - && use_turbo4_asym_dequant_sdpa - && cache.turbo4_asym_dequant_sdpa_available() - { - // Default Turbo4Asym decode: dequantize the 4-bit V to FP16 and run - // native SDPA with the FP16 K. Exact full-dequant attention. - cache.update_and_turbo4_asym_dequant_sdpa_attention(&q, k, v, self.scale, mask) - } else if l == 1 && cache.sparse_v_available() { - // Sparse-V fallback: only taken when the dequant-SDPA gate above is - // disabled (`MLXCEL_TURBO4_ASYM_DEQUANT_SDPA=0`). - cache - .update_and_sparse_v_attention(&q, k, v, self.scale, mask) - .expect("update_and_sparse_v_attention returned None despite sparse_v_available") - } else if l == 1 && use_turbo4_dequant_sdpa && cache.turbo4_dequant_sdpa_available() { - cache.update_and_turbo4_dequant_sdpa_attention(&q, k, v, self.scale, mask) - } else if l == 1 && use_delegated_compressed && cache.turbo4_delegated_available() { - // The helper always produces an attention output: it routes - // through the fused Metal kernel when available and falls - // through to the graph-only reference path otherwise. - cache.update_and_turbo4_delegated_attention(&q, k, v, self.scale, mask) - } else if l > 1 && mask.is_none() { - // Prefill: use causal masking - let (cache_k, cache_v) = cache.update_and_fetch(k, v); - mlxcel_core::causal_attention(&q, &cache_k, &cache_v, self.scale, 0.0, 0) - } else { - // Single token or explicit mask - let (cache_k, cache_v) = cache.update_and_fetch(k, v); - let mask_ptr = mask.map(|m| m as *const _).unwrap_or(std::ptr::null()); - unsafe { - mlxcel_core::layers::attention_from_ptr( - &q, &cache_k, &cache_v, self.scale, mask_ptr, 0.0, 0, - ) - } - }; + // Attention dispatch is a property of the cache (#2171, ADR 0008): the + // cache appends this step's K/V and picks the kernel from the storage + // behind it (the pooled paged entry at a single unmasked token, the + // Turbo dequant-first variants, or fused SDPA). Prefill (l > 1) and + // masked steps never take the paged single-token kernel. + let attn_out = cache.attend(&q, k, v, self.scale, mask); // Transpose back and reshape let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); @@ -819,12 +749,10 @@ impl Attention { caches: &mut [&mut KVCache], metadata: &BatchedAttentionMetadata, mask: Option<&MlxArray>, - decode_context: Option<&DecodeBatchContext>, ) -> UniquePtr { let b = caches.len(); let seq_len = mlxcel_core::array_shape(q_batched)[1]; debug_assert_eq!(metadata.len(), b); - let mut attn_outputs: Vec> = Vec::with_capacity(b); let q_batched = mlxcel_core::reshape( q_batched, @@ -850,162 +778,23 @@ impl Attention { let (q_batched, k_batched) = self.apply_rope_batched(&q_batched, &k_batched, &metadata.rope_offsets); - // Pool-backed paged decode (#899): the whole batch in one fused launch. - // This is the production path for `--decode-storage paged`; it appends - // every sequence's new K/V to the shared `PagedBlockPool` and runs - // attention once over the batch, replacing the per-sequence - // `update_and_fetch` + `gather_visible` + SDPA loop below. It declines - // (leaving the pool untouched) for anything it cannot serve, including - // dense-backed batches, so the loop below is still reached unchanged. - // Multi-token steps (batched prefill, speculative / MTP verify) never - // enter it: `seq_len != 1`. - if seq_len == 1 - && mask.is_none() - && decode_context.is_some_and(|context| context.is_paged_decode()) - && let Some(attn_out) = mlxcel_core::cache::paged_batch_decode_attention( - &q_batched, &k_batched, &v_batched, caches, self.scale, 0.0, - ) - { - let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); - return mlxcel_core::reshape( - &attn_out, - &[b as i32, seq_len, self.num_heads * self.head_dim], - ); - } - - let paged_decode = decode_context.and_then(|context| { - if seq_len != 1 || mask.is_some() || !context.is_paged_decode() { - return None; - } - if caches.iter().any(|cache| cache.mode != KVCacheMode::Fp16) { - return None; - } - // Pool-backed caches keep no dense `keys`/`values` buffers for the - // dense-compat paged kernel to read. Since #899 they are normally - // served by `paged_batch_decode_attention` above; reaching here - // means that path declined, so fall through to the per-sequence - // `update_and_fetch` loop below, whose pool intercept writes into - // the shared `PagedBlockPool` and gathers the visible window back. - if caches.iter().any(|cache| cache.is_paged_backed()) { - return None; - } - let metadata = - PagedDecodeMetadata::from_attention_metadata(metadata, context.paged_block_size) - .ok()?; - Some((context.use_native_paged_kernel, metadata)) - }); - - if let Some((use_native_kernel, paged_metadata)) = paged_decode { - tracing::debug!( - batch_size = b, - block_size = paged_metadata.block_size, - native_kernel = use_native_kernel, - "Llama3 paged decode attention dispatch" - ); - let mut cache_keys: Vec<*const MlxArray> = Vec::with_capacity(b); - let mut cache_values: Vec<*const MlxArray> = Vec::with_capacity(b); - - for (i, cache) in caches.iter_mut().enumerate() { - let k_i = mlxcel_core::slice( - &k_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - let v_i = mlxcel_core::slice( - &v_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - cache.update(k_i, v_i); - cache_keys.push(cache.keys.as_ref().unwrap().as_ref().unwrap() as *const MlxArray); - cache_values - .push(cache.values.as_ref().unwrap().as_ref().unwrap() as *const MlxArray); - } - - let attn_out = if use_native_kernel { - mlxcel_core::layers::paged_decode_attention_dense_compat( - &q_batched, - &cache_keys, - &cache_values, - &paged_metadata, - self.scale, - ) - } else { - mlxcel_core::layers::paged_decode_attention_dense_fallback( - &q_batched, - &cache_keys, - &cache_values, - &paged_metadata, - self.scale, - ) - } - .expect("valid llama3 paged decode attention inputs"); - - let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); - return mlxcel_core::reshape( - &attn_out, - &[b as i32, seq_len, self.num_heads * self.head_dim], - ); - } - - for (i, cache) in caches.iter_mut().enumerate() { - // Slice [B, heads, T, dim] -> [1, heads, T, dim] for sequence i. - let q_i = mlxcel_core::slice( - &q_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - let k_i = mlxcel_core::slice( - &k_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - let v_i = mlxcel_core::slice( - &v_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - - // Update KV cache - let (cache_k, cache_v) = cache.update_and_fetch(k_i, v_i); - - let mask_i = mask.map(|m| { - let sliced = - mlxcel_core::slice(m, &[i as i32, 0, 0], &[i as i32 + 1, seq_len, i32::MAX]); - mlxcel_core::squeeze_axis(&sliced, 0) - }); - - // Causal prefill without explicit masks can use the shared causal - // helper; masked/padded prefill and decode keep using the unified - // attention dispatcher. - let attn_out = if seq_len > 1 && mask_i.is_none() { - mlxcel_core::causal_attention(&q_i, &cache_k, &cache_v, self.scale, 0.0, 0) - } else { - let mask_ptr = mask_i - .as_ref() - .map(|m| m.as_ref().unwrap() as *const _) - .unwrap_or(std::ptr::null()); - unsafe { - mlxcel_core::layers::attention_from_ptr( - &q_i, &cache_k, &cache_v, self.scale, mask_ptr, 0.0, 0, - ) - } - }; - - // Transpose back: [1, n_heads, T, head_dim] -> [1, T, n_heads * head_dim] - let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); - let attn_out = - mlxcel_core::reshape(&attn_out, &[1, seq_len, self.num_heads * self.head_dim]); - - attn_outputs.push(attn_out); - } + // One batched attention entry (#2171, ADR 0008): a single-token + // unmasked step over pool-backed caches is one whole-batch launch + // (#899); every other shape, and a batch that launch declines, runs + // each row through `KVCache::attend`, so a row decodes the same way + // whether it was scheduled alone or in a batch. Multi-token steps + // (batched prefill, speculative / MTP verify) never reach the paged + // single-token kernel. + let attn_out = mlxcel_core::cache::attend_batched( + &q_batched, &k_batched, &v_batched, caches, self.scale, mask, + ); - // Concatenate along batch dim: B * [1, T, hidden] -> [B, T, hidden] - let mut result = attn_outputs.remove(0); - for attn_out in attn_outputs { - result = mlxcel_core::concatenate(&result, &attn_out, 0); - } - result + // Transpose back: [B, n_heads, T, head_dim] -> [B, T, n_heads * head_dim] + let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); + mlxcel_core::reshape( + &attn_out, + &[b as i32, seq_len, self.num_heads * self.head_dim], + ) } /// Build the dense attention block from a checkpoint. @@ -1226,7 +1015,6 @@ impl TransformerBlock { x: &MlxArray, caches: &mut [&mut KVCache], mask: Option<&MlxArray>, - decode_context: Option<&DecodeBatchContext>, ) -> UniquePtr { // Batched pre-attention norm let normed = self.input_layernorm.forward(x); @@ -1246,7 +1034,6 @@ impl TransformerBlock { caches, &metadata, mask, - decode_context, ); // Batched output projection @@ -1408,7 +1195,6 @@ impl Llama3Model { input_ids: &MlxArray, batch_caches: &mut [&mut [KVCache]], mask: Option<&MlxArray>, - decode_context: Option<&DecodeBatchContext>, ) -> UniquePtr { let b = batch_caches.len(); @@ -1423,7 +1209,7 @@ impl Llama3Model { .map(|caches| &mut caches[layer_idx]) .collect(); - h = self.layers[layer_idx].forward_batched(&h, &mut layer_caches, mask, decode_context); + h = self.layers[layer_idx].forward_batched(&h, &mut layer_caches, mask); } // Batched final norm: [B, 1, hidden_dim] @@ -1596,17 +1382,7 @@ impl LanguageModel for Llama3Model { batch_caches: &mut [&mut [KVCache]], mask: Option<&MlxArray>, ) -> UniquePtr { - self.forward_batched_impl(input_ids, batch_caches, mask, None) - } - - fn forward_batched_with_context( - &self, - input_ids: &MlxArray, - batch_caches: &mut [&mut [KVCache]], - mask: Option<&MlxArray>, - context: Option<&DecodeBatchContext>, - ) -> UniquePtr { - self.forward_batched_impl(input_ids, batch_caches, mask, context) + self.forward_batched_impl(input_ids, batch_caches, mask) } fn supports_batched_prefill(&self) -> bool { diff --git a/src/models/qwen3.rs b/src/models/qwen3.rs index 92c01b301..04f392461 100644 --- a/src/models/qwen3.rs +++ b/src/models/qwen3.rs @@ -18,10 +18,10 @@ //! - Q/K normalization (RMSNorm after projection, before RoPE) //! - Explicit head_dim in config -use mlxcel_core::cache::{BatchedAttentionMetadata, PagedDecodeMetadata}; -use mlxcel_core::generate::{DecodeBatchContext, LanguageModel}; +use mlxcel_core::cache::BatchedAttentionMetadata; +use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::{ - FusedQKVLinear, KVCache, KVCacheMode, RMSNorm, UnifiedEmbedding, UnifiedLinear, + FusedQKVLinear, KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear, }; use mlxcel_core::utils::pipeline_hint; use mlxcel_core::weights::WeightMap; @@ -326,80 +326,12 @@ impl Attention { (q, k, v) }; - // Decode-case (l == 1) attention dispatch for the Turbo quantized cache - // modes. Prefill (l > 1) builds the cache from scratch and falls through - // to the standard masked/causal paths below. - // - // Turbo4Asym (FP16 K + 4-bit V) decodes via dequant-first native SDPA by - // default: V is dequantized to FP16 transiently and fed with the FP16 K - // to native SDPA. This is exact and ~3-6x faster than the lossy sparse-V - // weighted-sum path, which stays reachable behind - // `MLXCEL_TURBO4_ASYM_DEQUANT_SDPA=0` for A/B and fallback. Symmetric - // Turbo4 and Turbo4Delegated mirror mlx-swift-lm's dequant-first policy. - // Each gate is parsed once and cached in a `OnceLock`. - let use_turbo4_asym_dequant_sdpa = - mlxcel_core::cache::turbo::sparse_v::turbo4_asym_dequant_sdpa_enabled(); - let use_delegated_compressed = - mlxcel_core::cache::turbo::sparse_v::turbo4_delegated_compressed_attention_enabled(); - let use_turbo4_dequant_sdpa = - mlxcel_core::cache::turbo::sparse_v::turbo4_dequant_sdpa_enabled(); - // Single-sequence pooled paged decode (#899). The server takes this - // path whenever the active batch is one request: `decode_single_step` - // calls the single-sequence `forward`, never `forward_split_attention`. - // Without this branch the two single-sequence scenarios in the issue's - // benchmark matrix (16K and 32K) could never reach the fused kernel, and - // neither could any moment in a batched run where only one request is - // still decoding. The batch-1 launch is exactly the `[1, H, 1, D]` shape - // the whole-batch entry point already serves. - let pooled_single = if l == 1 && mask.is_none() && cache.is_paged_backed() { - mlxcel_core::cache::paged_batch_decode_attention( - &q, - &k, - &v, - &mut [&mut *cache], - self.scale, - 0.0, - ) - } else { - None - }; - - let attn_out = if let Some(out) = pooled_single { - out - } else if l == 1 - && use_turbo4_asym_dequant_sdpa - && cache.turbo4_asym_dequant_sdpa_available() - { - // Default Turbo4Asym decode: dequantize the 4-bit V to FP16 and run - // native SDPA with the FP16 K. Exact full-dequant attention. - cache.update_and_turbo4_asym_dequant_sdpa_attention(&q, k, v, self.scale, mask) - } else if l == 1 && cache.sparse_v_available() { - // Sparse-V fallback: only taken when the dequant-SDPA gate above is - // disabled (`MLXCEL_TURBO4_ASYM_DEQUANT_SDPA=0`). The helper consumes - // k/v, fills the packed cache, and runs the fused kernel (or graph - // fallback). When `sparse_v_available` is true it always returns Some. - cache - .update_and_sparse_v_attention(&q, k, v, self.scale, mask) - .expect("update_and_sparse_v_attention returned None despite sparse_v_available") - } else if l == 1 && use_turbo4_dequant_sdpa && cache.turbo4_dequant_sdpa_available() { - cache.update_and_turbo4_dequant_sdpa_attention(&q, k, v, self.scale, mask) - } else if l == 1 && use_delegated_compressed && cache.turbo4_delegated_available() { - // The helper always produces an attention output: it routes - // through the fused Metal kernel when available and falls - // through to the graph-only reference path otherwise. - cache.update_and_turbo4_delegated_attention(&q, k, v, self.scale, mask) - } else if l > 1 && mask.is_none() { - let (cache_k, cache_v) = cache.update_and_fetch(k, v); - mlxcel_core::causal_attention(&q, &cache_k, &cache_v, self.scale, 0.0, 0) - } else { - let (cache_k, cache_v) = cache.update_and_fetch(k, v); - let mask_ptr = mask.map(|m| m as *const _).unwrap_or(std::ptr::null()); - unsafe { - mlxcel_core::layers::attention_from_ptr( - &q, &cache_k, &cache_v, self.scale, mask_ptr, 0.0, 0, - ) - } - }; + // Attention dispatch is a property of the cache (#2171, ADR 0008): the + // cache appends this step's K/V and picks the kernel from the storage + // behind it (the pooled paged entry at a single unmasked token, the + // Turbo dequant-first variants, or fused SDPA). Prefill (l > 1) and + // masked steps never take the paged single-token kernel. + let attn_out = cache.attend(&q, k, v, self.scale, mask); // Transpose back and reshape let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); @@ -429,12 +361,10 @@ impl Attention { caches: &mut [&mut KVCache], metadata: &BatchedAttentionMetadata, mask: Option<&MlxArray>, - decode_context: Option<&DecodeBatchContext>, ) -> UniquePtr { let b = caches.len(); let seq_len = mlxcel_core::array_shape(q_batched)[1]; debug_assert_eq!(metadata.len(), b); - let mut attn_outputs: Vec> = Vec::with_capacity(b); let q_batched = mlxcel_core::reshape( q_batched, @@ -459,159 +389,23 @@ impl Attention { let (q_batched, k_batched) = self.apply_rope_batched(&q_batched, &k_batched, &metadata.rope_offsets); - // Pool-backed paged decode (#899): the whole batch in one fused launch. - // This is the production path for `--decode-storage paged`; it appends - // every sequence's new K/V to the shared `PagedBlockPool` and runs - // attention once over the batch, replacing the per-sequence - // `update_and_fetch` + `gather_visible` + SDPA loop below. It declines - // (leaving the pool untouched) for anything it cannot serve, including - // dense-backed batches, so the loop below is still reached unchanged. - // Multi-token steps (batched prefill, speculative / MTP verify) never - // enter it: `seq_len != 1`. - if seq_len == 1 - && mask.is_none() - && decode_context.is_some_and(|context| context.is_paged_decode()) - && let Some(attn_out) = mlxcel_core::cache::paged_batch_decode_attention( - &q_batched, &k_batched, &v_batched, caches, self.scale, 0.0, - ) - { - let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); - return mlxcel_core::reshape( - &attn_out, - &[b as i32, seq_len, self.num_heads * self.head_dim], - ); - } - - let paged_decode = decode_context.and_then(|context| { - if seq_len != 1 || mask.is_some() || !context.is_paged_decode() { - return None; - } - if caches.iter().any(|cache| cache.mode != KVCacheMode::Fp16) { - return None; - } - // Pool-backed caches keep no dense `keys`/`values` buffers for the - // dense-compat paged kernel to read. Since #899 they are normally - // served by `paged_batch_decode_attention` above; reaching here - // means that path declined, so fall through to the per-sequence - // `update_and_fetch` loop below, whose pool intercept writes into - // the shared `PagedBlockPool` and gathers the visible window back. - if caches.iter().any(|cache| cache.is_paged_backed()) { - return None; - } - let metadata = - PagedDecodeMetadata::from_attention_metadata(metadata, context.paged_block_size) - .ok()?; - Some((context.use_native_paged_kernel, metadata)) - }); - - if let Some((use_native_kernel, paged_metadata)) = paged_decode { - tracing::debug!( - batch_size = b, - block_size = paged_metadata.block_size, - native_kernel = use_native_kernel, - "Qwen3 paged decode attention dispatch" - ); - let mut cache_keys: Vec<*const MlxArray> = Vec::with_capacity(b); - let mut cache_values: Vec<*const MlxArray> = Vec::with_capacity(b); - - for (i, cache) in caches.iter_mut().enumerate() { - let k_i = mlxcel_core::slice( - &k_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - let v_i = mlxcel_core::slice( - &v_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - cache.update(k_i, v_i); - cache_keys.push(cache.keys.as_ref().unwrap().as_ref().unwrap() as *const MlxArray); - cache_values - .push(cache.values.as_ref().unwrap().as_ref().unwrap() as *const MlxArray); - } - - let attn_out = if use_native_kernel { - mlxcel_core::layers::paged_decode_attention_dense_compat( - &q_batched, - &cache_keys, - &cache_values, - &paged_metadata, - self.scale, - ) - } else { - mlxcel_core::layers::paged_decode_attention_dense_fallback( - &q_batched, - &cache_keys, - &cache_values, - &paged_metadata, - self.scale, - ) - } - .expect("valid qwen3 paged decode attention inputs"); - - let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); - return mlxcel_core::reshape( - &attn_out, - &[b as i32, seq_len, self.num_heads * self.head_dim], - ); - } - - for (i, cache) in caches.iter_mut().enumerate() { - // Slice [B, heads, T, dim] -> [1, heads, T, dim] for sequence i. - let q_i = mlxcel_core::slice( - &q_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - let k_i = mlxcel_core::slice( - &k_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - let v_i = mlxcel_core::slice( - &v_batched, - &[i as i32, 0, 0, 0], - &[i as i32 + 1, i32::MAX, i32::MAX, i32::MAX], - ); - - // Update KV cache - let (cache_k, cache_v) = cache.update_and_fetch(k_i, v_i); - - let mask_i = mask.map(|m| { - let sliced = - mlxcel_core::slice(m, &[i as i32, 0, 0], &[i as i32 + 1, seq_len, i32::MAX]); - mlxcel_core::squeeze_axis(&sliced, 0) - }); - - let attn_out = if seq_len > 1 && mask_i.is_none() { - mlxcel_core::causal_attention(&q_i, &cache_k, &cache_v, self.scale, 0.0, 0) - } else { - let mask_ptr = mask_i - .as_ref() - .map(|m| m.as_ref().unwrap() as *const _) - .unwrap_or(std::ptr::null()); - unsafe { - mlxcel_core::layers::attention_from_ptr( - &q_i, &cache_k, &cache_v, self.scale, mask_ptr, 0.0, 0, - ) - } - }; - - // Transpose back: [1, n_heads, T, head_dim] -> [1, T, n_heads * head_dim] - let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); - let attn_out = - mlxcel_core::reshape(&attn_out, &[1, seq_len, self.num_heads * self.head_dim]); - - attn_outputs.push(attn_out); - } + // One batched attention entry (#2171, ADR 0008): a single-token + // unmasked step over pool-backed caches is one whole-batch launch + // (#899); every other shape, and a batch that launch declines, runs + // each row through `KVCache::attend`, so a row decodes the same way + // whether it was scheduled alone or in a batch. Multi-token steps + // (batched prefill, speculative / MTP verify) never reach the paged + // single-token kernel. + let attn_out = mlxcel_core::cache::attend_batched( + &q_batched, &k_batched, &v_batched, caches, self.scale, mask, + ); - // Concatenate along batch dim: B * [1, T, hidden] -> [B, T, hidden] - let mut result = attn_outputs.remove(0); - for attn_out in attn_outputs { - result = mlxcel_core::concatenate(&result, &attn_out, 0); - } - result + // Transpose back: [B, n_heads, T, head_dim] -> [B, T, n_heads * head_dim] + let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); + mlxcel_core::reshape( + &attn_out, + &[b as i32, seq_len, self.num_heads * self.head_dim], + ) } pub fn from_weights( @@ -783,7 +577,6 @@ impl TransformerBlock { x: &MlxArray, caches: &mut [&mut KVCache], mask: Option<&MlxArray>, - decode_context: Option<&DecodeBatchContext>, ) -> UniquePtr { // Batched pre-attention norm let normed = self.input_layernorm.forward(x); @@ -803,7 +596,6 @@ impl TransformerBlock { caches, &metadata, mask, - decode_context, ); // Batched output projection @@ -968,7 +760,6 @@ impl Qwen3Model { input_ids: &MlxArray, batch_caches: &mut [&mut [KVCache]], mask: Option<&MlxArray>, - decode_context: Option<&DecodeBatchContext>, ) -> UniquePtr { let b = batch_caches.len(); @@ -983,7 +774,7 @@ impl Qwen3Model { .map(|caches| &mut caches[layer_idx]) .collect(); - h = self.layers[layer_idx].forward_batched(&h, &mut layer_caches, mask, decode_context); + h = self.layers[layer_idx].forward_batched(&h, &mut layer_caches, mask); } // Batched final norm: [B, 1, hidden_dim] @@ -1155,17 +946,7 @@ impl LanguageModel for Qwen3Model { batch_caches: &mut [&mut [KVCache]], mask: Option<&MlxArray>, ) -> UniquePtr { - self.forward_batched_impl(input_ids, batch_caches, mask, None) - } - - fn forward_batched_with_context( - &self, - input_ids: &MlxArray, - batch_caches: &mut [&mut [KVCache]], - mask: Option<&MlxArray>, - context: Option<&DecodeBatchContext>, - ) -> UniquePtr { - self.forward_batched_impl(input_ids, batch_caches, mask, context) + self.forward_batched_impl(input_ids, batch_caches, mask) } fn supports_batched_prefill(&self) -> bool { diff --git a/src/server/batch/scheduler/decode_tick.rs b/src/server/batch/scheduler/decode_tick.rs index 0909b1951..f4d042ae7 100644 --- a/src/server/batch/scheduler/decode_tick.rs +++ b/src/server/batch/scheduler/decode_tick.rs @@ -314,6 +314,11 @@ impl BatchScheduler { /// Batched decode-storage context for the active backend. Shared by the /// synchronous batched decode and the lookahead prime so both drive the /// identical dense / native-paged execution path. + /// + /// Only the model-owned families (Gemma 3, Llama 4) still select a kernel + /// from this. `KVCache`-backed families decide inside + /// `KVCache::attend` from the storage the pool wired for the sequence + /// (#2171, ADR 0008), so for them the value is informational. pub(super) fn decode_batch_context(&self) -> DecodeBatchContext { match self.decode_storage_backend { DecodeStorageBackend::Auto | DecodeStorageBackend::Dense => { From 6798da230516adfaf7593456545a1fd26b24d103 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 18:19:57 +0900 Subject: [PATCH 2/7] chore: cargo fmt the attend module and its tests --- src/lib/mlxcel-core/src/cache/attend.rs | 12 +++++++----- src/lib/mlxcel-core/src/cache/attend_tests.rs | 15 +++++++++++++-- src/lib/mlxcel-core/src/mla/cache.rs | 2 +- src/models/helium.rs | 3 ++- src/models/llama3.rs | 15 ++++----------- src/models/qwen3.rs | 15 ++++----------- 6 files changed, 31 insertions(+), 31 deletions(-) diff --git a/src/lib/mlxcel-core/src/cache/attend.rs b/src/lib/mlxcel-core/src/cache/attend.rs index 15df85e87..5ece67d08 100644 --- a/src/lib/mlxcel-core/src/cache/attend.rs +++ b/src/lib/mlxcel-core/src/cache/attend.rs @@ -125,9 +125,12 @@ impl KVCache { if self.sparse_v_available() { return self .update_and_sparse_v_attention(q, new_keys, new_values, scale, mask) - .expect("update_and_sparse_v_attention returned None despite sparse_v_available"); + .expect( + "update_and_sparse_v_attention returned None despite sparse_v_available", + ); } - if turbo::sparse_v::turbo4_dequant_sdpa_enabled() && self.turbo4_dequant_sdpa_available() + if turbo::sparse_v::turbo4_dequant_sdpa_enabled() + && self.turbo4_dequant_sdpa_available() { return self.update_and_turbo4_dequant_sdpa_attention( q, new_keys, new_values, scale, mask, @@ -136,9 +139,8 @@ impl KVCache { if turbo::sparse_v::turbo4_delegated_compressed_attention_enabled() && self.turbo4_delegated_available() { - return self.update_and_turbo4_delegated_attention( - q, new_keys, new_values, scale, mask, - ); + return self + .update_and_turbo4_delegated_attention(q, new_keys, new_values, scale, mask); } } let (cache_k, cache_v) = self.update_and_fetch(new_keys, new_values); diff --git a/src/lib/mlxcel-core/src/cache/attend_tests.rs b/src/lib/mlxcel-core/src/cache/attend_tests.rs index 71ad8d9ef..8b171b6d2 100644 --- a/src/lib/mlxcel-core/src/cache/attend_tests.rs +++ b/src/lib/mlxcel-core/src/cache/attend_tests.rs @@ -78,7 +78,15 @@ fn relative_rms(a: &[f32], b: &[f32]) -> f32 { } /// `(q, k, v)` for one step of `len` tokens over `batch` rows. -fn step(rng: &mut Rng, batch: i32, len: i32) -> (UniquePtr, UniquePtr, UniquePtr) { +fn step( + rng: &mut Rng, + batch: i32, + len: i32, +) -> ( + UniquePtr, + UniquePtr, + UniquePtr, +) { ( random_array(rng, &[batch, HEADS, len, DIM]), random_array(rng, &[batch, KV_HEADS, len, DIM]), @@ -278,7 +286,10 @@ fn paged_cache_single_token_goes_through_the_pooled_entry() { ); let b = reference(&mut baseline, &q, k2, v2, None); let rms = relative_rms(&to_vec_f32(&a), &to_vec_f32(&b)); - assert!(rms < 2e-2, "paged decode diverged from the gather baseline: rms {rms}"); + assert!( + rms < 2e-2, + "paged decode diverged from the gather baseline: rms {rms}" + ); assert_eq!(paged.offset, 41); assert_eq!(baseline.offset, 41); diff --git a/src/lib/mlxcel-core/src/mla/cache.rs b/src/lib/mlxcel-core/src/mla/cache.rs index debee3000..fb20fee69 100644 --- a/src/lib/mlxcel-core/src/mla/cache.rs +++ b/src/lib/mlxcel-core/src/mla/cache.rs @@ -59,9 +59,9 @@ use cxx::UniquePtr; use crate::cache::{KVCache, KVCacheMode}; use crate::ffi::MlxArray; +use crate::mla::MlaGeometry; use crate::mla::absorb::MlaAbsorbedProjections; use crate::mla::decode::absorbed_decode; -use crate::mla::MlaGeometry; /// Bytes one token of the compressed-latent cache costs, per layer. /// diff --git a/src/models/helium.rs b/src/models/helium.rs index 24866f9e0..3f66bde0b 100644 --- a/src/models/helium.rs +++ b/src/models/helium.rs @@ -950,7 +950,8 @@ impl LanguageModel for HeliumModel { batch_caches: &mut [&mut [KVCache]], mask: Option<&MlxArray>, ) -> UniquePtr { - self.inner.forward_batched_impl(input_ids, batch_caches, mask) + self.inner + .forward_batched_impl(input_ids, batch_caches, mask) } fn supports_batched_prefill(&self) -> bool { diff --git a/src/models/llama3.rs b/src/models/llama3.rs index 36d09435f..c333c9f6a 100644 --- a/src/models/llama3.rs +++ b/src/models/llama3.rs @@ -19,9 +19,7 @@ use mlxcel_core::cache::{BatchedAttentionMetadata, SequenceId}; use mlxcel_core::generate::LanguageModel; -use mlxcel_core::layers::{ - FusedQKVLinear, KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear, -}; +use mlxcel_core::layers::{FusedQKVLinear, KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear}; use mlxcel_core::utils::pipeline_hint; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; @@ -1027,14 +1025,9 @@ impl TransformerBlock { // Per-sequence attention still owns cache mutation, but positional // metadata and RoPE now stay on a batched path. - let attn_concat = self.self_attn.forward_split_attention( - &q, - &k, - &v, - caches, - &metadata, - mask, - ); + let attn_concat = self + .self_attn + .forward_split_attention(&q, &k, &v, caches, &metadata, mask); // Batched output projection let attn_out = self.self_attn.o_proj.forward(&attn_concat); diff --git a/src/models/qwen3.rs b/src/models/qwen3.rs index 04f392461..33fe23f74 100644 --- a/src/models/qwen3.rs +++ b/src/models/qwen3.rs @@ -20,9 +20,7 @@ use mlxcel_core::cache::BatchedAttentionMetadata; use mlxcel_core::generate::LanguageModel; -use mlxcel_core::layers::{ - FusedQKVLinear, KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear, -}; +use mlxcel_core::layers::{FusedQKVLinear, KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear}; use mlxcel_core::utils::pipeline_hint; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; @@ -589,14 +587,9 @@ impl TransformerBlock { // Per-sequence attention still owns cache mutation, but positional // metadata and RoPE now stay on a batched path. - let attn_concat = self.self_attn.forward_split_attention( - &q, - &k, - &v, - caches, - &metadata, - mask, - ); + let attn_concat = self + .self_attn + .forward_split_attention(&q, &k, &v, caches, &metadata, mask); // Batched output projection let attn_out = self.self_attn.o_proj.forward(&attn_concat); From a400b30a9921b48bf4cd9504b3f01132adfcd57a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 18:22:48 +0900 Subject: [PATCH 3/7] test: pin MlaLatentCache::attend against the decompressed reference --- src/lib/mlxcel-core/src/mla/decode_tests.rs | 26 +++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/lib/mlxcel-core/src/mla/decode_tests.rs b/src/lib/mlxcel-core/src/mla/decode_tests.rs index 0a5fd9bc7..bc9f906ab 100644 --- a/src/lib/mlxcel-core/src/mla/decode_tests.rs +++ b/src/lib/mlxcel-core/src/mla/decode_tests.rs @@ -102,6 +102,32 @@ fn absorbed_attention_respects_an_additive_causal_mask() { ); } +#[test] +fn latent_cache_attend_matches_the_decompressed_reference() { + let _guard = serial(); + // The cache-owned entry (#2171): hand the latent cache this step's rows and + // the query halves, and it appends then runs absorbed attention over the + // live window. Same identity as `absorbed_decode`, pinned against the same + // host f64 reference so a wrong window or a dropped append cannot hide. + let fx = MlaFixture::new(TINY, 2, 1, 37, 0xA77E); + let dt = dtype::FLOAT32; + let proj = MlaAbsorbedProjections::from_dense(&fx.kv_b_array(dt), fx.geometry).unwrap(); + let mut inner = KVCache::new(); + let mut cache = MlaLatentCache::wrap(&mut inner, fx.geometry).unwrap(); + let out = cache.attend( + &fx.q_nope_array(dt), + &fx.q_pe_array(dt), + fx.ckv_array(dt), + fx.kpe_array(dt), + &proj, + fx.scale, + None, + ); + assert_eq!(cache.offset(), 37, "attend appends the step's rows"); + let err = max_rel_error(&to_vec_f32(&out), &fx.decompressed_reference(false)); + assert!(err < F32_TOL, "latent cache attend drifted by {err}"); +} + #[test] fn expand_latent_reproduces_the_up_projection() { // The prefill direction of the same fold. `k_nope` and `v` must equal what From c22e08b455df495ce0f157122c23ebd1da74fbc3 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 18:41:47 +0900 Subject: [PATCH 4/7] refactor: route DeepSeek V2 absorbed decode through MlaLatentCache::attend MlaLatentCache::attend had no production caller: deepseek_v2 still spelled out update_and_fetch plus the split-KV or Stage 1 choice itself, so the new entry was a parallel unused path, against #2171's "no parallel unused path" criterion. The entry now carries the decode policy the model ran (Stage 2 split-KV first when enabled and the plan accepts, Stage 1 otherwise) and deepseek_v2's single-token absorbed step calls it; prefill keeps up-projecting through the model's kv_b_proj. Also pins the Turbo batched numerics change: a new attend test asserts each batched Turbo4Asym row is bit-identical to the single-sequence dequant-first call and within float noise of the old full-dequant route. ADR 0008 now states that the removed dense-compat route did serve the FP16 boundary layers of Turbo and --kv-bits paged batches (Boundary-V and --kv-skip-last-layer), which now run per-row SDPA. Validated: clippy -D warnings (test-fast, cuda, lib/tests/bins), fmt, and mlxcel-core cache::attend, mla::, decode_undo, paged_batch_decode, paged_detach, detach tests. Refs #2171 --- ...0008-kv-attention-dispatch-in-the-cache.md | 6 +- src/lib/mlxcel-core/src/cache/attend_tests.rs | 64 ++++++++++++++++++- src/lib/mlxcel-core/src/mla/cache.rs | 30 ++++++++- src/models/deepseek_v2.rs | 35 +++------- 4 files changed, 102 insertions(+), 33 deletions(-) diff --git a/docs/adr/0008-kv-attention-dispatch-in-the-cache.md b/docs/adr/0008-kv-attention-dispatch-in-the-cache.md index e87d5c855..2cd2ad1c1 100644 --- a/docs/adr/0008-kv-attention-dispatch-in-the-cache.md +++ b/docs/adr/0008-kv-attention-dispatch-in-the-cache.md @@ -30,7 +30,7 @@ Dispatch is a match on the cache's own fields. A `trait KvStorage` with a virtua `CachePool::allocate_with_layout` decides storage once per sequence: pool-backed caches for a dense-natural family on the paged backend with an FP16 layout, dense caches otherwise. That allocation is the per-sequence storage policy. For the families on `attend`, `DecodeBatchContext` no longer selects a kernel; the scheduler still builds it for the model-owned families (Gemma 3, Llama 4, and the `model_owned` dispatch helpers), whose per-sequence state is not a `KVCache` and which keep their dense-pointer paged kernels until they move onto the entry. -The dense-pointer paged route (`paged_decode_attention_dense_compat` over dense `keys`/`values` with a block table derived from the batch metadata) is removed from the Qwen3 and Llama 3 batched forwards. It was unreachable there: with an FP16 paged layout the pool wires pool-backed caches, and Int8 or Turbo paged layouts fail its FP16 check. The kernels stay for the model-owned families. +The dense-pointer paged route (`paged_decode_attention_dense_compat` over dense `keys`/`values` with a block table derived from the batch metadata) is removed from the Qwen3 and Llama 3 batched forwards. With an FP16 paged layout the pool wires pool-backed caches, which that route already skipped, so it only ever served the FP16 layers of a non-FP16 paged layout, whose caches are dense: the Boundary-V layers a Turbo mode keeps at FP16 (`MLXCEL_KV_BOUNDARY_V_LAYERS`, two at each end by default) and the last layer `--kv-skip-last-layer` (default on) keeps at FP16 under `--kv-bits`. Those layers now run `attend` per row (dense update, then SDPA through `attention_from_ptr`), the route the quantized layers of the same batch already took. The removed C++ entry was itself a per-row loop of block slices, a concat and a direct `fast::scaled_dot_product_attention` call, so this is the same attention without the concat copy, but not bit-identical to it. The kernels stay for the model-owned families. ### Single-sequence storage default @@ -39,8 +39,8 @@ Dense is the single-sequence KV storage, as ADR 0007 measured. It applies wherev ## Consequences - `rg 'is_paged_backed\(\)|is_paged_decode\(\)' src/models` returns nothing for Qwen3, Llama 3 and the families that reuse their attention (Qwen2, Qwen2.5, Helium, the VLM text backbones). The paged-vs-dense decision has one home, so the ADR 0007 parity harness compares one dispatch under two storages. -- Batched decode over Turbo caches now takes the same dequant-first variants as single-sequence decode instead of a full dequant per step; the result is the same exact attention through a different op order, so it is not bit-identical to the previous batched route. Pool-backed and FP16 dense routes are op-for-op what the models ran before, pinned by `cache::attend::attend_tests`. -- MLA latent caches keep their own attention and expose the same entry: `MlaLatentCache::attend` appends the step's latent and rope rows and runs the absorbed step (`mla::decode_tests` pins it). +- Batched decode over Turbo caches now takes the same dequant-first variants as single-sequence decode instead of a full dequant per step; the result is the same exact attention through a different op order, so it is not bit-identical to the previous batched route. Pool-backed routes and FP16 dense routes on the dense backend are op-for-op what the models ran before, pinned by `cache::attend::attend_tests` (which also pins a batched Turbo4Asym row to the single-sequence dequant-first call); the FP16 boundary layers of a quantized paged batch changed kernel as described above. +- MLA latent caches keep their own attention and expose the same entry: `MlaLatentCache::attend` appends the step's latent and rope rows and runs the absorbed step, Stage 2 split-KV first when enabled; DeepSeek V2's absorbed decode calls it (`mla::decode_tests` pins it), while its prefill keeps up-projecting through the model's `kv_b_proj`. - `RotatingKVCache` (Gemma 3's sliding window) is untouched; its decode undo log (#2182) and the model-owned families are the follow-up this phase leaves for 4b. - Prompt-cache detach and adopt, the disaggregated handoff (`is_paged_backed()` in `src/distributed/disaggregated/handoff_impl.rs`) and `CachePool::sync_paged_state_with_dense` are unchanged: they read storage, which is still the cache's property. diff --git a/src/lib/mlxcel-core/src/cache/attend_tests.rs b/src/lib/mlxcel-core/src/cache/attend_tests.rs index 8b171b6d2..b2163a4e5 100644 --- a/src/lib/mlxcel-core/src/cache/attend_tests.rs +++ b/src/lib/mlxcel-core/src/cache/attend_tests.rs @@ -26,7 +26,9 @@ use std::cell::RefCell; use std::rc::Rc; use super::*; -use crate::cache::{PagedBlockPool, PagedKvLayout, PagedSequenceState, paged_batch_decode_stats}; +use crate::cache::{ + KVCacheMode, PagedBlockPool, PagedKvLayout, PagedSequenceState, paged_batch_decode_stats, +}; use crate::dtype; const PAGE: usize = 32; @@ -335,3 +337,63 @@ fn batched_paged_decode_is_one_pooled_launch() { assert_eq!(c0.offset, 13); assert_eq!(c1.offset, 13); } + +/// Batched decode over Turbo4Asym caches takes the same dequant-first route a +/// lone sequence takes (#2171): each row of [`attend_batched`] is bit-identical +/// to that row's single-sequence dequant-first call, and stays within float +/// noise of the full-dequant `update_and_fetch` plus SDPA pair the batched +/// loop ran before. Pins the behavior change ADR 0008 records, so a future +/// batched Turbo route cannot silently diverge from the single-sequence one. +#[test] +fn batched_turbo4_asym_rows_match_single_sequence_dequant_first() { + let mut rng = Rng::new(0x2174); + let mode = KVCacheMode::Turbo4Asym; + let mut c0 = KVCache::new_with_mode(mode); + let mut c1 = KVCache::new_with_mode(mode); + let mut s0 = KVCache::new_with_mode(mode); + let mut s1 = KVCache::new_with_mode(mode); + let mut f0 = KVCache::new_with_mode(mode); + let mut f1 = KVCache::new_with_mode(mode); + for (c, s, f) in [(&mut c0, &mut s0, &mut f0), (&mut c1, &mut s1, &mut f1)] { + let (q, k, v) = step(&mut rng, 1, 6); + c.attend(&q, ffi::copy(&k), ffi::copy(&v), SCALE, None); + s.attend(&q, ffi::copy(&k), ffi::copy(&v), SCALE, None); + reference(f, &q, k, v, None); + } + + let (q, k, v) = step(&mut rng, 2, 1); + let mut single_rows = Vec::new(); + let mut full_rows = Vec::new(); + for (b, (s, f)) in [(&mut s0, &mut f0), (&mut s1, &mut f1)] + .into_iter() + .enumerate() + { + let q_b = slice_row(&q, b); + single_rows.push(s.update_and_turbo4_asym_dequant_sdpa_attention( + &q_b, + slice_row(&k, b), + slice_row(&v, b), + SCALE, + None, + )); + full_rows.push(reference(f, &q_b, slice_row(&k, b), slice_row(&v, b), None)); + } + let single = crate::concatenate(&single_rows[0], &single_rows[1], 0); + let full = crate::concatenate(&full_rows[0], &full_rows[1], 0); + + let mut caches: Vec<&mut KVCache> = vec![&mut c0, &mut c1]; + let got = attend_batched(&q, &k, &v, &mut caches, SCALE, None); + let got = to_vec_f32(&got); + assert_eq!( + got, + to_vec_f32(&single), + "batched row left the dequant-first route" + ); + let drift = relative_rms(&got, &to_vec_f32(&full)); + assert!( + drift < 1e-2, + "dequant-first drifted {drift} from the full-dequant reference" + ); + assert_eq!(c0.offset, 7); + assert_eq!(c1.offset, 7); +} diff --git a/src/lib/mlxcel-core/src/mla/cache.rs b/src/lib/mlxcel-core/src/mla/cache.rs index fb20fee69..c68124771 100644 --- a/src/lib/mlxcel-core/src/mla/cache.rs +++ b/src/lib/mlxcel-core/src/mla/cache.rs @@ -62,6 +62,7 @@ use crate::ffi::MlxArray; use crate::mla::MlaGeometry; use crate::mla::absorb::MlaAbsorbedProjections; use crate::mla::decode::absorbed_decode; +use crate::mla::{MlaSplitPlan, absorbed_decode_split_kv}; /// Bytes one token of the compressed-latent cache costs, per layer. /// @@ -272,9 +273,15 @@ impl<'a> MlaLatentCache<'a> { /// rows, `mask` is the optional additive causal or padding mask. Returns /// `[B, H, L, v_head_dim]`. /// - /// Used by: the MLA families once they move onto the latent cache (#907 - /// Stage 2); `mla::decode_tests` pins it against the decompressed - /// reference. + /// A single-token unmasked step takes the Stage 2 split-KV decode when + /// [`crate::mla::split_kv_enabled`] is on and its plan accepts the shape, + /// and Stage 1 [`absorbed_decode`] otherwise, the same order the model + /// forward ran before this entry existed. + /// + /// Used by: `models::deepseek_v2::Attention::forward_absorbed` (decode + /// steps; prefill up-projects the window through the model's `kv_b_proj` + /// and stays in the model); `mla::decode_tests` pins it against the + /// decompressed reference. #[must_use] #[allow(clippy::too_many_arguments)] pub fn attend( @@ -288,6 +295,23 @@ impl<'a> MlaLatentCache<'a> { mask: Option<&MlxArray>, ) -> UniquePtr { let (ckv_all, kpe_all) = self.update_and_fetch(ckv, kpe); + let q_shape = crate::ffi::array_shape(q_nope); + if mask.is_none() && q_shape.len() == 4 && q_shape[2] == 1 && crate::mla::split_kv_enabled() + { + let plan = MlaSplitPlan::heuristic( + q_shape[0].max(0) as usize, + q_shape[1].max(0) as usize, + self.seq_len(), + crate::paged_v2::device_target_ctas(), + ); + // A plan that declines falls through to Stage 1 rather than + // failing the step. + if let Ok(out) = + absorbed_decode_split_kv(q_nope, q_pe, &ckv_all, &kpe_all, proj, scale, &plan) + { + return out; + } + } absorbed_decode(q_nope, q_pe, &ckv_all, &kpe_all, proj, scale, mask) } diff --git a/src/models/deepseek_v2.rs b/src/models/deepseek_v2.rs index 4ec397651..775e45b69 100644 --- a/src/models/deepseek_v2.rs +++ b/src/models/deepseek_v2.rs @@ -25,9 +25,7 @@ use crate::models::switch_layers::validate_expert_quantization_params; use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::{KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear}; -use mlxcel_core::mla::{ - self, MlaAbsorbedProjections, MlaDecodePath, MlaGeometry, MlaLatentCache, MlaSplitPlan, -}; +use mlxcel_core::mla::{self, MlaAbsorbedProjections, MlaDecodePath, MlaGeometry, MlaLatentCache}; use mlxcel_core::utils::{repeat_kv, silu, slice_axis}; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr, concatenate}; @@ -536,32 +534,17 @@ impl MLAAttention { let mut view = MlaLatentCache::wrap(cache, geometry) .expect("supports() was checked immediately before this wrap"); - let (ckv_all, kpe_all) = view.update_and_fetch(ckv, k_pe); - let kv_len = view.seq_len(); let output = if l == 1 { - // Stage 2 first when asked for, Stage 1 otherwise. A split plan that - // declines falls through to Stage 1 rather than failing the step. - let split = if mla::split_kv_enabled() { - let plan = MlaSplitPlan::heuristic( - b.max(0) as usize, - self.num_heads, - kv_len, - mlxcel_core::paged_v2::device_target_ctas(), - ); - mla::absorbed_decode_split_kv( - q_nope, &q_pe, &ckv_all, &kpe_all, proj, self.scale, &plan, - ) - .ok() - } else { - None - }; - split.unwrap_or_else(|| { - // Decode sees the whole cache, so no mask, matching the - // `causal_attention` call the decompressed path makes here. - mla::absorbed_decode(q_nope, &q_pe, &ckv_all, &kpe_all, proj, self.scale, None) - }) + // The latent cache's own attention entry (#2171, ADR 0008): it + // appends this step's rows and runs Stage 2 split-KV when asked + // for, Stage 1 otherwise. Decode sees the whole cache, so no mask, + // matching the `causal_attention` call the decompressed path makes + // here. + view.attend(q_nope, &q_pe, ckv, k_pe, proj, self.scale, None) } else { + let (ckv_all, kpe_all) = view.update_and_fetch(ckv, k_pe); + let kv_len = view.seq_len(); mla::stats::record(MlaDecodePath::AbsorbedPrefill); // Up-project the live latent window back into per-head K and V. // `kv_b_proj` wants `[B, S, kv_lora_rank]`, the shape before the From e0e127b7443acb7617cf72ff21325215d85f94f8 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 18:50:55 +0900 Subject: [PATCH 5/7] test: pin attend_batched on empty, one-row and mixed-storage batches Add four cases to cache::attend::attend_tests. An empty batch panics with "empty batch". A one-row dense batch matches the per-row attend bit for bit, and a one-row pool-backed batch is served by exactly one pooled call, which is the shape the lookahead prime and forward_batched hand over without a DecodeBatchContext. A batch that mixes a pool-backed row with a dense row, in either order, makes the pooled entry decline before it writes, and each row then matches an independent attend on its own storage, with only the pool-backed row counted as a pooled call. Validation: - gpu-lock run --tag pr2210-final -- cargo test --profile test-fast --features cuda -p mlxcel-core --lib cache::attend -- --test-threads=1: 11 passed Refs #2171 --- src/lib/mlxcel-core/src/cache/attend_tests.rs | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/src/lib/mlxcel-core/src/cache/attend_tests.rs b/src/lib/mlxcel-core/src/cache/attend_tests.rs index b2163a4e5..053c41c30 100644 --- a/src/lib/mlxcel-core/src/cache/attend_tests.rs +++ b/src/lib/mlxcel-core/src/cache/attend_tests.rs @@ -397,3 +397,126 @@ fn batched_turbo4_asym_rows_match_single_sequence_dequant_first() { assert_eq!(c0.offset, 7); assert_eq!(c1.offset, 7); } + +#[test] +#[should_panic(expected = "empty batch")] +fn batched_attend_rejects_an_empty_batch() { + let mut rng = Rng::new(0x2177); + let (q, k, v) = step(&mut rng, 1, 1); + let mut caches: Vec<&mut KVCache> = Vec::new(); + attend_batched(&q, &k, &v, &mut caches, SCALE, None); +} + +/// A one-row batch is what the lookahead prime and `forward_batched` hand over +/// when no `DecodeBatchContext` exists. Dense storage matches the per-row +/// `attend` bit for bit. +#[test] +fn batched_dense_one_row_matches_attend() { + let mut rng = Rng::new(0x2178); + let mut batched = KVCache::new(); + let mut single = KVCache::new(); + let (q, k, v) = step(&mut rng, 1, 5); + batched.attend(&q, ffi::copy(&k), ffi::copy(&v), SCALE, None); + single.attend(&q, k, v, SCALE, None); + + let (q, k, v) = step(&mut rng, 1, 1); + let want = single.attend(&q, ffi::copy(&k), ffi::copy(&v), SCALE, None); + let mut caches: Vec<&mut KVCache> = vec![&mut batched]; + let got = attend_batched(&q, &k, &v, &mut caches, SCALE, None); + assert_eq!(ffi::array_shape(&got), vec![1, HEADS, 1, DIM]); + assert_eq!(to_vec_f32(&got), to_vec_f32(&want)); + assert_eq!(batched.offset, 6); +} + +/// The pooled launch keys on `caches[0]`, so a one-row pool-backed batch +/// reaches it at a single-token step with no scheduler context. +#[test] +fn batched_paged_one_row_is_one_pooled_launch() { + let (pool, states) = fresh_pool(1); + let mut rng = Rng::new(0x2179); + let mut paged = KVCache::new_paged(pool, states[0].clone(), 0); + let (q, k, v) = step(&mut rng, 1, 12); + paged.attend(&q, k, v, SCALE, None); + + let (q, k, v) = step(&mut rng, 1, 1); + let before = paged_batch_decode_stats(); + let mut caches: Vec<&mut KVCache> = vec![&mut paged]; + let out = attend_batched(&q, &k, &v, &mut caches, SCALE, None); + let after = paged_batch_decode_stats(); + assert_eq!(ffi::array_shape(&out), vec![1, HEADS, 1, DIM]); + assert_eq!( + after.v2_launches + after.gather_fallbacks, + before.v2_launches + before.gather_fallbacks + 1, + "a one-row pool-backed batch is served by the pooled entry" + ); + assert_eq!(paged.offset, 13); +} + +/// A batch that mixes a pool-backed row with a dense row cannot share one +/// launch. The pooled entry declines before it writes to the pool, and every +/// row then runs `attend` on its own storage, in either row order. +#[test] +fn batched_mixed_storage_runs_each_row_on_its_own_route() { + for paged_first in [true, false] { + let (pool, states) = fresh_pool(1); + let mut rng = Rng::new(0x217a); + let mut paged = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let mut paged_ref = KVCache::new_paged(pool, states[1].clone(), 0); + let mut dense = KVCache::new(); + let mut dense_ref = KVCache::new(); + + let (q, k, v) = step(&mut rng, 1, 12); + paged.attend(&q, ffi::copy(&k), ffi::copy(&v), SCALE, None); + paged_ref.attend(&q, ffi::copy(&k), ffi::copy(&v), SCALE, None); + let (q, k, v) = step(&mut rng, 1, 12); + dense.attend(&q, ffi::copy(&k), ffi::copy(&v), SCALE, None); + dense_ref.attend(&q, k, v, SCALE, None); + + let (q, k, v) = step(&mut rng, 2, 1); + let (paged_row, dense_row) = if paged_first { (0, 1) } else { (1, 0) }; + let want_paged = paged_ref.attend( + &slice_row(&q, paged_row), + slice_row(&k, paged_row), + slice_row(&v, paged_row), + SCALE, + None, + ); + let want_dense = dense_ref.attend( + &slice_row(&q, dense_row), + slice_row(&k, dense_row), + slice_row(&v, dense_row), + SCALE, + None, + ); + + let before = paged_batch_decode_stats(); + let mut caches: Vec<&mut KVCache> = if paged_first { + vec![&mut paged, &mut dense] + } else { + vec![&mut dense, &mut paged] + }; + let got = attend_batched(&q, &k, &v, &mut caches, SCALE, None); + let after = paged_batch_decode_stats(); + assert_eq!(ffi::array_shape(&got), vec![2, HEADS, 1, DIM]); + assert_eq!( + after.v2_launches + after.gather_fallbacks, + before.v2_launches + before.gather_fallbacks + 1, + "only the pool-backed row reaches the pooled entry (paged_first={paged_first})" + ); + + let got_paged = to_vec_f32(&slice_row(&got, paged_row)); + let got_dense = to_vec_f32(&slice_row(&got, dense_row)); + assert_eq!( + got_dense, + to_vec_f32(&want_dense), + "dense row differs (paged_first={paged_first})" + ); + assert_eq!( + got_paged, + to_vec_f32(&want_paged), + "pool-backed row differs (paged_first={paged_first})" + ); + assert_eq!(paged.offset, 13); + assert_eq!(dense.offset, 13); + } +} From 52734e8f38e04a201a85fae65965b56a95d2576a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 18:54:15 +0900 Subject: [PATCH 6/7] docs: describe cache-owned attention dispatch and its numeric changes Record where the dense-versus-paged attention choice lives after #2171. CONTINUOUS_BATCHING.md gains a section covering the cache entry (KVCache::attend, attend_batched, MlaLatentCache::attend), the pooled launch keying on caches[0].is_paged_backed() so callers without a DecodeBatchContext also reach the fused launch, the dequant-first variants now used for batched Turbo rows, and the per-row dense update plus SDPA that replaces the removed dense-compat route for FP16 layers under a non-FP16 paged layout (same math, no concat copy, not bit-identical). turbo-kv-cache.md repeats the Turbo and Boundary-V consequences, architecture.md names cache/attend.rs, and environment-variables.md no longer says use_native_paged_kernel governs the dense-compat decode for Qwen3 and Llama 3, since only the model-owned families still use it. The stale is_paged_backed() remark in the paged_scheduler_parity doc comment now points at attend_batched. There is no Korean docs tree in this checkout, so no translation was updated. ADR 0008 and its index entry were already part of the PR. Validation: - python3 scripts/insert_apache_header.py --check: ok - cargo test --profile test-fast --features cuda --test vlm_wrapper_capability_delegation --test dead_doc_pointers: 4 passed Refs #2171 --- docs/CONTINUOUS_BATCHING.md | 27 +++++++++++++++++++++++++++ docs/architecture.md | 5 ++++- docs/environment-variables.md | 2 +- docs/turbo-kv-cache.md | 8 ++++++++ tests/paged_scheduler_parity.rs | 5 +++-- 5 files changed, 43 insertions(+), 4 deletions(-) diff --git a/docs/CONTINUOUS_BATCHING.md b/docs/CONTINUOUS_BATCHING.md index c2497decd..611a39985 100644 --- a/docs/CONTINUOUS_BATCHING.md +++ b/docs/CONTINUOUS_BATCHING.md @@ -383,6 +383,33 @@ of total work: at 1024 tokens of context per request the kernel runs 0.91x at batch 1 but 1.41x at batch 4 and 1.47x at batch 8, since a batched launch spreads the same chunk count over more requests and amortizes the merge pass. +#### Where the dense-versus-paged choice lives + +The kernel is picked by the cache, not by the model forward +([ADR 0008](adr/0008-kv-attention-dispatch-in-the-cache.md), issue #2171). +`KVCache::attend` appends the step's K/V and runs attention for the storage +behind that cache, and `cache::attend_batched` does the same for one layer of a +batch. Qwen3, Llama 3 and the families that reuse their attention (Qwen2, +Qwen2.5, Helium, the VLM text backbones) call these and no longer branch on +`is_paged_backed()` or read the scheduler's `DecodeBatchContext`; DeepSeek V2's +absorbed decode goes through `MlaLatentCache::attend`. Gemma 3, Llama 4 and the +other model-owned families keep their own dense-pointer paged kernels for now. + +- The pooled launch keys on `caches[0].is_paged_backed()` at an unmasked + single-token step, so a caller with no `DecodeBatchContext` (the lookahead + prime, `forward_batched`) reaches the fused launch too. A batch that mixes + pool-backed and dense rows declines before it writes and runs each row through + `attend` on its own storage. +- Batched decode over Turbo caches takes the dequant-first variants that + single-sequence decode already used, instead of a full dequant per step. The + attention is the same exact math in a different op order, so it is not + bit-identical to the old batched route. +- The FP16 layers of a non-FP16 paged layout (the Boundary-V layers a Turbo mode + keeps at FP16, and the last layer under `--kv-skip-last-layer`) hold dense + caches. They now run a per-row dense update plus SDPA instead of the removed + `paged_decode_attention_dense_compat` route: the same attention without the + concat copy, but not bit-identical to it. + #### Seeing which path ran The server announces the first occurrence of each distinct dispatch outcome at diff --git a/docs/architecture.md b/docs/architecture.md index 40ac7ec63..2092b2c23 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -36,7 +36,10 @@ src/ - `src/lib/mlxcel-core/src/lib.rs` — `cxx::bridge` definitions and crate exports. - `src/lib/mlxcel-core/src/cache.rs` and `src/lib/mlxcel-core/src/cache/` — FP16/INT8/TurboQuant KV cache variants, - paged cache layout, detach/adopt helpers, and cache tests. + paged cache layout, detach/adopt helpers, and cache tests. `cache/attend.rs` + holds `KVCache::attend` and `cache::attend_batched`, the one attention entry + that picks the dense or paged kernel from the storage behind the cache + ([ADR 0008](adr/0008-kv-attention-dispatch-in-the-cache.md)). - `src/lib/mlxcel-core/src/ops.rs`, `src/lib/mlxcel-core/src/dtype.rs`, `src/lib/mlxcel-core/src/streams.rs` — wrappers around common MLX operations and runtime concepts. - `src/lib/mlxcel-core/src/sampling.rs` — penalties and token sampling shared by CLI/server paths. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index f10d58f39..09719ba7c 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -565,7 +565,7 @@ recommended as normal deployment settings. | `MLXCEL_DISABLE_SOFTCAP_GQA_DECODE_GROUPED` | `1` disables, `0` enables | unset | Legacy rollback knob for grouped softcap-GQA decode, checked before the enable variable above. Set to `1` to restore the repeat-based path. | | `MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS` | truthy disables | maskless path on | Disables the single-query maskless attention path. | | `MLXCEL_EXPERIMENTAL_BOOL_CAUSAL_MASK` | truthy enables | off | Enables an experimental boolean causal-mask path. | -| `MLXCEL_PAGED_ATTENTION_NATIVE` | `1`/`true`/`on`/`yes` force the native kernel; `0`/`false`/`off`/`no` force gather (case-insensitive); unset or any other value defers to the adaptive selector | selector-governed (server: the issue #899 token floors; library entry point on Metal: native only for batch>=4 + ctx<=4096 + single-slab; on CUDA: native for any single-slab layer since #634; on ROCm: the CUDA rule since #2068) | Overrides the fused paged-attention decode kernel (Metal since epic #116 Phase 6/#123, since #634 also CUDA via `mx.fast.cuda_kernel`, and since #2068 also ROCm via `fast::hip_kernel`). **Two consumers**, per `resolve_dispatch_decision` and `resolve_paged_v2_dispatch` in `src/lib/mlxcel-core/src/layers.rs`: (1) the library-only `paged_decode_attention_pooled` entry point, where since #331 an unset value no longer means "always gather" (`select_pooled_paged_dispatch` picks the kernel only inside the regime ADR 0001 measured it winning) and this variable force-pins either arm for A/B testing; and (2) **the server's pool-backed batched paged decode**, where issue #899 made the fused v2 kernel the production path and named this variable's force-off values its kill switch: setting one restores the pre-#899 gather-then-SDPA behaviour end to end, and a force-on value pins v2 for every servable shape, bypassing the measured token floors (`MLXCEL_PAGED_V2_MIN_KV_TOKENS`, `MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST`). The override is checked before the selector in both consumers, and a forced dispatch still goes through the kernel's structural declines (single-slab layer, servable geometry, non-empty batch). History: #710 retired the *library* entry point to a library-only API, which is where the "not a server knob" reading came from; #899 gave the variable its second, server-side consumer. The separate `DecodeBatchContext::use_native_paged_kernel` request still governs the dense-compat block-table decode, a different path. See [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md) and [Continuous batching](CONTINUOUS_BATCHING.md#seeing-which-path-ran). | +| `MLXCEL_PAGED_ATTENTION_NATIVE` | `1`/`true`/`on`/`yes` force the native kernel; `0`/`false`/`off`/`no` force gather (case-insensitive); unset or any other value defers to the adaptive selector | selector-governed (server: the issue #899 token floors; library entry point on Metal: native only for batch>=4 + ctx<=4096 + single-slab; on CUDA: native for any single-slab layer since #634; on ROCm: the CUDA rule since #2068) | Overrides the fused paged-attention decode kernel (Metal since epic #116 Phase 6/#123, since #634 also CUDA via `mx.fast.cuda_kernel`, and since #2068 also ROCm via `fast::hip_kernel`). **Two consumers**, per `resolve_dispatch_decision` and `resolve_paged_v2_dispatch` in `src/lib/mlxcel-core/src/layers.rs`: (1) the library-only `paged_decode_attention_pooled` entry point, where since #331 an unset value no longer means "always gather" (`select_pooled_paged_dispatch` picks the kernel only inside the regime ADR 0001 measured it winning) and this variable force-pins either arm for A/B testing; and (2) **the server's pool-backed batched paged decode**, where issue #899 made the fused v2 kernel the production path and named this variable's force-off values its kill switch: setting one restores the pre-#899 gather-then-SDPA behaviour end to end, and a force-on value pins v2 for every servable shape, bypassing the measured token floors (`MLXCEL_PAGED_V2_MIN_KV_TOKENS`, `MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST`). The override is checked before the selector in both consumers, and a forced dispatch still goes through the kernel's structural declines (single-slab layer, servable geometry, non-empty batch). History: #710 retired the *library* entry point to a library-only API, which is where the "not a server knob" reading came from; #899 gave the variable its second, server-side consumer. Since #2171 the pool-backed paged decode is dispatched from the KV cache (`KVCache::attend`, [ADR 0008](adr/0008-kv-attention-dispatch-in-the-cache.md)) for Qwen3, Llama 3 and the families that reuse their attention, so for them the scheduler's `DecodeBatchContext::use_native_paged_kernel` request no longer selects a kernel; it still gates the dense-pointer paged decode of the model-owned families (Gemma 3, Llama 4). See [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md) and [Continuous batching](CONTINUOUS_BATCHING.md#seeing-which-path-ran). | | `MLXCEL_PAGED_ATTENTION_V2` | `1`/`true`/`on`/`yes` enable (case-insensitive, trimmed); unset or anything else disables | off | **Diagnostic.** Issue #898's comparison gate: makes the **library-only** v1 entry point (`PagedBlockPool::paged_decode_fused`, reached from `paged_decode_attention_pooled`) try the fused v2 kernel first. It does **not** gate the server's production v2 decode, which is issue #899, is default-on, and is governed by the token floors and `MLXCEL_PAGED_ATTENTION_NATIVE` instead. Read once per process, and a typo degrades to v1 rather than enabling a path nobody asked for. See [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md). | | `MLXCEL_PAGED_DECODE_V2_CHUNK` | positive integer pages | unset (planned or autotuned chunk size) | **Diagnostic.** Pins `pages_per_chunk` for every paged decode v2 launch, overriding both the plan's occupancy heuristic and any autotuned value (issue #906), matching the `MLXCEL_PAGED_DECODE_SPLITS` / `MLXCEL_QMM_TILE_*` escape-hatch convention. The plan clamps the value into the feasible range, so an over-large value is safe; a non-positive or unparseable value warns and is ignored in favour of the planned size. Read once per process. | | `MLXCEL_PAGED_DECODE_V2_TARGET_CTAS` | positive integer | derived (Apple: `gpu_core_count * 8`, floored at 64; every other host including CUDA: `512`) | **Diagnostic.** Moves the occupancy target the v2 plan binary-searches its chunk size against, as opposed to `MLXCEL_PAGED_DECODE_V2_CHUNK`, which pins the chunk size itself. This is the knob to turn when the derived device parallelism is wrong. The Apple figure is the `gpu_core_count` device-scale proxy, not a calibrated core count, and the fixed non-Apple target is **unvalidated on CUDA** (a CUDA-specific target should come from the SM count), so both are starting points. A non-positive or unparseable value warns and uses the derived target. Read once per process. | diff --git a/docs/turbo-kv-cache.md b/docs/turbo-kv-cache.md index 02da68f3b..358061560 100644 --- a/docs/turbo-kv-cache.md +++ b/docs/turbo-kv-cache.md @@ -162,6 +162,14 @@ mlxcel generate -m models/ \ `0` disables the policy. Values larger than half the layer count are clamped by the runtime. The flag is inert for `fp16` and `int8` cache modes. +Under continuous batching the attention route follows the storage behind each +layer's cache ([ADR 0008](adr/0008-kv-attention-dispatch-in-the-cache.md)). +Turbo rows of a batched decode step take the same dequant-first variants as +single-sequence decode, and the layers Boundary-V (or `--kv-skip-last-layer`) +keeps at FP16 under a non-FP16 paged layout hold dense caches that run a +per-row dense update plus SDPA. Both are the same exact attention in a different +op order, so they are not bit-identical to the pre-#2171 batched routes. + ## Symmetric Turbo4 allowlist K-side quantization can strongly affect softmax quality. The code therefore has diff --git a/tests/paged_scheduler_parity.rs b/tests/paged_scheduler_parity.rs index a7662ed33..0d38f2418 100644 --- a/tests/paged_scheduler_parity.rs +++ b/tests/paged_scheduler_parity.rs @@ -55,8 +55,9 @@ //! exactly as `execute_batched_decode` dispatches it. For pool-backed caches //! this reaches the #899 whole-batch pooled decode helper, including its //! gather fallback below the fused dispatch floor; if that helper declines, -//! the model's `is_paged_backed()` guard still keeps the dense-pointer native -//! kernel away from placeholder dense buffers. +//! `cache::attend_batched` runs each row through `KVCache::attend` on its own +//! storage (ADR 0008), so no dense-pointer kernel ever sees a placeholder +//! dense buffer. //! //! # Running //! From d0374496faa9cd862f7d0aa2f5f1e99fcb3f9096 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 18:57:10 +0900 Subject: [PATCH 7/7] docs: add technical report for PR #2210 --- ...attention-dispatch-in-cache-20261007.en.md | 60 +++++++++++++++++++ ...attention-dispatch-in-cache-20261007.ko.md | 60 +++++++++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.en.md create mode 100644 TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.ko.md diff --git a/TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.en.md b/TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.en.md new file mode 100644 index 000000000..d4c88882b --- /dev/null +++ b/TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.en.md @@ -0,0 +1,60 @@ +# Technical Report: PR #2210 - KV attention dispatch inside the cache + +**Date**: 2026-10-07 + +**Status**: Implemented and verified on GB10; pending merge. + +**Languages**: Rust, Markdown (ADR, docs) + +**Risk Level**: Medium. Every Qwen3, Llama 3, Helium and DeepSeek V2 attention step now goes through one cache entry. FP16 dense and pool-backed routes run the same ops as before, and real-checkpoint parity is unchanged. Two batched routes (FP16 boundary layers under a non-FP16 paged layout, and Turbo rows) now compute the same attention in a different op order. + +## Executive Summary + +Before this PR, each model forward decided which attention kernel to run by asking the cache what storage it had (`is_paged_backed()`) and reading the scheduler's `DecodeBatchContext`. Epic #2166 needs one engine where storage is a property of the sequence, not a branch in every model. This PR, Phase 4a (#2171), adds `KVCache::attend` and `cache::attend_batched`, which append the step's K/V and choose the kernel from the storage behind the cache, plus `MlaLatentCache::attend` for the MLA latent cache. Qwen3, Llama 3 (and the families that reuse their attention), Helium and DeepSeek V2's absorbed decode call these entries. ADR 0008 records the design and defers Gemma 3, Llama 4 and the `model_owned` helpers to Phase 4b. + +## 1. Problem Statement + +- The kernel choice was repeated in each model's `forward_split_attention` and batched forward: pool-backed caches took the pooled paged kernel, dense caches under a paged layout took a dense-pointer paged route (`paged_decode_attention_dense_compat`), Turbo modes took dequant-first variants, everything else took SDPA. +- The CLI and the server ran different kernels for the same model and request because each call site chose its own route. +- Adding or changing a storage meant editing every family that branched on it. + +## 2. Change Summary + +- **`src/lib/mlxcel-core/src/cache/attend.rs`:** + - `KVCache::attend(q, k, v, scale, mask)`. Pool-backed cache at a single unmasked token: the #899 pooled entry, which declines before writing. Pool-backed otherwise: pool intercept plus SDPA. Dense Turbo: the dequant-first variants (ADR 0002). Dense FP16/Int8: dense update plus fused SDPA. + - `attend_batched(q, k, v, caches, scale, mask)`: one pooled launch for the batch when `caches[0]` is pool-backed and the step is a single unmasked token; rows it declines go through `attend` one by one. +- **`MlaLatentCache::attend`** holds the absorbed MLA decode policy: split-KV when enabled and the plan accepts, else `absorbed_decode`. +- **Models:** qwen3, llama3 and helium call `attend` / `attend_batched` and drop the `DecodeBatchContext` parameter from their batched forwards. deepseek_v2's single-token absorbed decode calls `MlaLatentCache::attend`. +- **ADR 0008** and docs (`CONTINUOUS_BATCHING.md`, `turbo-kv-cache.md`, `environment-variables.md`, `architecture.md`). +- **Tests:** `attend_tests.rs` (dispatch per storage, Turbo batched rows vs single-sequence, empty, one-row and mixed-storage batches), MLA `attend` against the decompressed reference. + +## 3. Technical Decisions + +- **A match on the cache's fields, not a `KvStorage` trait object.** The cache already resolves its quantization mode by matching; a virtual call per layer per token is what ADR 0004 rejected for the hot loop. +- **Storage is chosen where the cache is built.** `CachePool::allocate_with_layout` gives a dense-natural family on the paged backend pool-backed caches only for an FP16 layout. That allocation is the per-sequence storage policy, so `DecodeBatchContext` no longer selects kernels for the families on `attend`. The ADR 0007 measured default (dense for a lone sequence) holds wherever the worker runs at `max_batch_size=1`; the server's `--decode-storage-backend` policy is unchanged, and moving a sequence between storages as the batch grows is out of scope. +- **Remove the dense-pointer route from Qwen3 and Llama 3.** It only served the FP16 layers a non-FP16 paged layout keeps dense (Turbo Boundary-V layers, the last layer under `--kv-skip-last-layer`). The removed C++ entry was a per-row loop of block slices, a concat and SDPA; per-row dense update plus SDPA is the same attention without the concat copy. The kernels stay for the model-owned families until Phase 4b. +- **Accept the Turbo batched numerics change.** Rows the pooled launch declines now take the dequant-first variants, the same route single-sequence decode already used. This moves batched output toward ADR 0007's "same input, same output" invariant instead of away from it; a test pins batched rows bit-identical to the single-sequence call. +- **Keep the per-row retry after a refused batch.** A variant that skipped the pooled attempt per row would fall back to the pool intercept plus gather, which ADR 0001 measured at 2 to 3 times SDPA cost past 4K tokens. + +## 4. Validation + +- Unit tests: mlxcel-core `cache::attend` (11), `mla::` (42), `cache::paged_batch_decode`, `cache::decode_undo`, `cache::paged_detach`; root modules `models::qwen3::`, `models::llama3::`, `models::helium`, `models::deepseek_v2`, scheduler, prompt-cache, model-owned lookahead, `distributed::disaggregated` (189), each run alone. +- clippy (`-D warnings`, root and mlxcel-core tests), `cargo fmt --check`, the license-header check, `vlm_wrapper_capability_delegation`, `dead_doc_pointers`. +- Real checkpoints, release `mlxcel-engine-parity` under `MLXCEL_SDPA_DETERMINISTIC=1`, before and after, and again after merging `main` (#2204): Qwen3-1.7B 4-bit and Llama-3.2-1B 4-bit, CLI vs server dense vs server paged identical; the only divergences are the documented prompt-cache-hit pairs, unchanged. + +## 5. Residual Risks + +- **Turbo plus paged batched serving was not run on a real checkpoint.** The boundary-layer and Turbo-row routes are covered by unit tests only. The epic's end-of-run measurement should include one Turbo paged batch for throughput. +- **DeepSeek absorbed decode was not run on a real checkpoint.** It is opt-in (`MLXCEL_MLA_ABSORBED`), deepseek-v2-lite is known broken on GB10, and the change moves existing calls. The block-level unit test covers it. +- **The ignored real-checkpoint paged parity tests** (`paged_scheduler_parity`, `serving_handoff_parity_tests`) compile but were not run. +- **Model-owned families still read `DecodeBatchContext`** until Phase 4b. + +## 6. Learning Points + +- **"Unreachable" needs a check against every layout, not the default one.** The first draft called the dense-compat route unreachable because FP16 paged layouts get pool-backed caches. Mixed layouts, where most layers are quantized and a few are kept at FP16, still reached it. +- **A new entry point is only finished when production code calls it.** `MlaLatentCache::attend` passed its own tests while DeepSeek still did its own dispatch; the review caught it by searching for callers. + +## 7. Related + +- Epic #2166, issue #2171, ADR 0007, ADR 0008. +- #2172 (Phase 4b, the deferred families), ADR 0001 (paged gather cost), ADR 0002 (Turbo dequant-first), ADR 0004 (no per-op dynamic dispatch), #899 (pooled paged entry), #2182 (rotating decode-undo log, untouched). diff --git a/TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.ko.md b/TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.ko.md new file mode 100644 index 000000000..e58b5a215 --- /dev/null +++ b/TECHNICAL_REPORTS/2210-kv-attention-dispatch-in-cache-20261007.ko.md @@ -0,0 +1,60 @@ +# 기술 보고서: PR #2210 - KV attention 분기를 캐시 안으로 + +**작성일**: 2026-10-07 + +**상태**: GB10에서 구현하고 검증했으며 머지 대기 중이다. + +**언어**: Rust, Markdown(ADR, 문서) + +**위험도**: 중간. Qwen3, Llama 3, Helium, DeepSeek V2의 attention 단계가 모두 캐시의 진입점 하나를 거친다. FP16 dense 경로와 pool 기반 경로는 이전과 같은 연산을 수행하고, 실제 체크포인트 패리티도 그대로다. 배치 경로 두 곳(FP16이 아닌 paged 레이아웃에서 FP16로 남는 경계 층, Turbo 행)은 같은 attention을 다른 연산 순서로 계산한다. + +## 요약 + +이 PR 전에는 모델 forward가 직접 캐시의 저장소 종류(`is_paged_backed()`)를 묻고 스케줄러의 `DecodeBatchContext`를 읽어 attention 커널을 골랐다. epic #2166은 저장소가 모델마다 갈라지는 분기가 아니라 시퀀스의 속성인 엔진 하나를 목표로 한다. Phase 4a(#2171)인 이 PR은 `KVCache::attend`와 `cache::attend_batched`를 추가한다. 두 함수는 그 단계의 K/V를 붙이고 캐시 뒤의 저장소를 보고 커널을 고른다. MLA latent 캐시에는 `MlaLatentCache::attend`를 추가했다. Qwen3, Llama 3(그리고 그 attention을 재사용하는 계열), Helium, DeepSeek V2의 absorbed decode가 이 진입점을 호출한다. 설계는 ADR 0008에 남겼고, Gemma 3, Llama 4, `model_owned` 헬퍼는 Phase 4b로 미뤘다. + +## 1. 문제 정의 + +- 커널 선택이 모델마다 `forward_split_attention`과 배치 forward에 반복되어 있었다. pool 기반 캐시는 pooled paged 커널로, paged 레이아웃 아래의 dense 캐시는 dense 포인터 paged 경로(`paged_decode_attention_dense_compat`)로, Turbo 모드는 dequant-first 변형으로, 나머지는 SDPA로 갔다. +- 호출 지점마다 경로를 따로 고르니 같은 모델, 같은 요청에서도 CLI와 서버가 다른 커널을 돌렸다. +- 저장소를 추가하거나 바꾸려면 그 저장소로 분기하는 모델 계열을 모두 고쳐야 했다. + +## 2. 변경 요약 + +- **`src/lib/mlxcel-core/src/cache/attend.rs`** + - `KVCache::attend(q, k, v, scale, mask)`. pool 기반 캐시에서 mask 없는 한 토큰이면 #899 pooled 진입점을 쓰고, 이 진입점은 쓰기 전에 거절 여부를 정한다. pool 기반의 나머지 경우는 pool intercept와 SDPA. dense Turbo는 dequant-first 변형(ADR 0002). dense FP16/Int8은 dense 갱신과 fused SDPA. + - `attend_batched(q, k, v, caches, scale, mask)`: `caches[0]`이 pool 기반이고 mask 없는 한 토큰 단계면 배치 전체를 pooled 호출 한 번으로 처리하고, 거절된 행은 하나씩 `attend`로 보낸다. +- **`MlaLatentCache::attend`**는 absorbed MLA decode 정책을 맡는다. split-KV가 켜져 있고 계획이 받아들이면 split-KV, 아니면 `absorbed_decode`다. +- **모델**: qwen3, llama3, helium은 `attend` / `attend_batched`를 호출하고 배치 forward에서 `DecodeBatchContext` 인자를 뺐다. deepseek_v2의 한 토큰 absorbed decode는 `MlaLatentCache::attend`를 호출한다. +- **ADR 0008**과 문서(`CONTINUOUS_BATCHING.md`, `turbo-kv-cache.md`, `environment-variables.md`, `architecture.md`). +- **테스트**: `attend_tests.rs`(저장소별 분기, Turbo 배치 행과 단일 시퀀스 비교, 빈 배치, 한 행 배치, 저장소가 섞인 배치), MLA `attend`와 압축 해제 기준 구현 비교. + +## 3. 기술적 선택과 그 이유 + +- **`KvStorage` trait 객체 대신 캐시 필드에 대한 match.** 캐시는 이미 양자화 모드를 match로 정한다. 층마다, 토큰마다 가상 호출을 하는 방식은 ADR 0004가 hot loop에서 배제했다. +- **저장소는 캐시를 만들 때 정한다.** `CachePool::allocate_with_layout`은 paged 백엔드의 dense 계열에 FP16 레이아웃일 때만 pool 기반 캐시를 준다. 이 할당이 시퀀스별 저장소 정책이므로, `attend`로 옮긴 계열에서는 `DecodeBatchContext`가 더는 커널을 고르지 않는다. ADR 0007에서 측정으로 정한 기본값(시퀀스가 하나면 dense)은 워커가 `max_batch_size=1`로 도는 곳에서 성립한다. 서버의 `--decode-storage-backend` 정책은 바꾸지 않았고, 배치가 커질 때 시퀀스를 다른 저장소로 옮기는 일은 범위 밖이다. +- **Qwen3와 Llama 3에서 dense 포인터 경로를 제거한다.** 이 경로는 FP16이 아닌 paged 레이아웃이 dense로 남겨 두는 FP16 층(Turbo Boundary-V 층, `--kv-skip-last-layer`일 때 마지막 층)만 처리했다. 제거한 C++ 진입점은 행마다 블록을 잘라 이어 붙인 뒤 SDPA를 부르는 루프였다. 행마다 dense 갱신 후 SDPA를 부르면 같은 attention을 이어 붙이기 복사 없이 계산한다. 이 커널은 Phase 4b 전까지 모델 소유 계열을 위해 남긴다. +- **Turbo 배치의 수치 변화는 받아들인다.** pooled 호출이 거절한 행은 이제 dequant-first 변형을 탄다. 단일 시퀀스 decode가 이미 쓰던 경로다. 배치 출력이 ADR 0007의 '같은 입력, 같은 출력' 불변식에서 멀어지는 것이 아니라 가까워지며, 배치 행이 단일 시퀀스 호출과 비트 단위로 같다는 것을 테스트로 고정했다. +- **배치가 거절된 뒤 행별로 다시 시도하는 동작은 유지한다.** 행마다 pooled 시도를 건너뛰는 변형을 만들면 pool intercept와 gather 경로로 떨어지는데, ADR 0001은 4K 토큰을 넘으면 이 경로가 SDPA보다 2~3배 느리다고 측정했다. + +## 4. 검증 + +- 단위 테스트: mlxcel-core `cache::attend`(11), `mla::`(42), `cache::paged_batch_decode`, `cache::decode_undo`, `cache::paged_detach`. 루트 크레이트는 `models::qwen3::`, `models::llama3::`, `models::helium`, `models::deepseek_v2`, 스케줄러, prompt-cache, 모델 소유 lookahead, `distributed::disaggregated`(189)를 모듈마다 따로 돌렸다. +- clippy(`-D warnings`, 루트와 mlxcel-core 테스트), `cargo fmt --check`, 라이선스 헤더 검사, `vlm_wrapper_capability_delegation`, `dead_doc_pointers`. +- 실제 체크포인트: release `mlxcel-engine-parity`를 `MLXCEL_SDPA_DETERMINISTIC=1`로 변경 전후와 `main`(#2204) 병합 후에 돌렸다. Qwen3-1.7B 4-bit와 Llama-3.2-1B 4-bit 모두 CLI, 서버 dense, 서버 paged 출력이 같았고, 문서화된 prompt-cache 적중 쌍만 이전과 똑같이 갈라졌다. + +## 5. 남은 위험 + +- **Turbo와 paged를 함께 쓰는 배치 서빙은 실제 체크포인트로 돌리지 않았다.** 경계 층 경로와 Turbo 행 경로는 단위 테스트로만 확인했다. epic 마지막 측정에 Turbo paged 배치 처리량을 한 번 넣어야 한다. +- **DeepSeek absorbed decode도 실제 체크포인트로 돌리지 않았다.** 이 경로는 `MLXCEL_MLA_ABSORBED`로 켜야 하고, deepseek-v2-lite는 GB10에서 이미 깨져 있으며, 이번 변경은 기존 호출을 옮긴 것이다. 블록 단위 단위 테스트가 이를 확인한다. +- **기본으로 ignore된 실제 체크포인트 paged 패리티 테스트**(`paged_scheduler_parity`, `serving_handoff_parity_tests`)는 컴파일만 확인했다. +- **모델 소유 계열은 Phase 4b 전까지 `DecodeBatchContext`를 읽는다.** + +## 6. 학습 포인트 + +- **'도달 불가'는 기본 레이아웃 하나가 아니라 모든 레이아웃에서 확인해야 한다.** 첫 초안은 FP16 paged 레이아웃이 pool 기반 캐시를 받으니 dense-compat 경로에 도달할 수 없다고 적었다. 대부분의 층은 양자화하고 일부 층만 FP16로 남기는 혼합 레이아웃은 여전히 그 경로를 탔다. +- **새 진입점은 실제 코드가 호출해야 끝난다.** `MlaLatentCache::attend`는 자체 테스트를 통과했지만 DeepSeek는 여전히 자기 분기를 쓰고 있었다. 리뷰에서 호출 지점을 검색해 찾아냈다. + +## 7. 관련 항목 + +- epic #2166, 이슈 #2171, ADR 0007, ADR 0008. +- #2172(Phase 4b, 미룬 계열), ADR 0001(paged gather 비용), ADR 0002(Turbo dequant-first), ADR 0004(op 단위 동적 분기 배제), #899(pooled paged 진입점), #2182(회전 캐시 decode-undo 로그, 변경 없음).