feat(core): matrix-absorbed MLA decode over a compressed-latent KV cache by inureyes · Pull Request #995 · lablup/mlxcel

inureyes · 2026-08-02T03:43:00Z

Summary

Adds mlxcel_core::mla, a shared matrix-absorbed MLA implementation over a compressed-latent KV cache, and wires deepseek_v2 onto it behind MLXCEL_MLA_ABSORBED (default off). Stage 1 (cache format plus absorbed decode from composed MLX ops) and Stage 2 (split-KV reusing issue #898's merge kernel unchanged) both land here; Stage 2's partial producer is composed from MLX ops, not a fused kernel, and is not a speed path.

A correction to the issue's premise, up front

The issue says the MLA families cache the decompressed K/V. That is true for deepseek_v2, minicpm3, and dots1. It is not true for deepseek_v3 and deepseek_v32: both have carried matrix absorption since before this work, decomposing kv_b_proj into embed_q / unembed_out in sanitize_weights and caching (kv_latent, k_pe) with the rope score passed to SDPA as an additive mask. So the actual gap was V2, MiniCPM3, and dots1, plus the fact that the transform existed twice with no shared home. This PR builds that shared home and moves V2 onto it. V3/V3.2 are untouched (converging them is a follow-up with real regression risk that deserves its own PR), and MiniCPM3/dots1 are untouched and still decompressed.

What changed

src/lib/mlxcel-core/src/mla/absorb.rs: load-time fold of kv_b_proj into W_UK [1, H, qk_nope, r] for the query path and W_UV^T [1, H, r, v_head] for the output path. Dequantizes when quantized (including the NVFP4 global_scale sidecar, folded as a scalar since (x @ W^T) * s2 == x @ (W * s2)^T), refuses a biased layer because the identity has no term for a bias, and refuses a geometry that does not partition the matrix. Nothing on disk is touched.
src/lib/mlxcel-core/src/mla/cache.rs: MlaLatentCache, a checked view over KVCache that puts ckv in the key slot and kpe in the value slot. supports() is split out of wrap() so a caller can fall back (the borrow checker keeps the &mut alive through the Err arm of a Result carrying that borrow's lifetime, which would otherwise make the fallback unexpressible). Also the latent_bytes_per_token / decompressed_bytes_per_token accounting the KV-memory table is computed from.
src/lib/mlxcel-core/src/mla/decode.rs: Stage 1. SDPA(q_absorbed, ckv, ckv, scale, mask = scale * q_pe @ kpe^T + causal). The rope half of the score is additive in exactly SDPA's mask position, so the identity is exact and MLX's fused attention kernel still does the softmax and the value accumulation. Same shape deepseek_v3.rs already uses.
src/lib/mlxcel-core/src/mla/split_kv.rs: Stage 2. MlaSplitPlan cuts the latent range; each chunk produces a normalized partial and its log2 LSE; paged_attention_merge_states merges them.
src/lib/mlxcel-core/src/mla/stats.rs: per-path call counters and a path=count summary.
src/models/deepseek_v2.rs: MLAAttention gains an optional fold and a forward_absorbed branch; load_mla_attention builds the fold only under the flag; report_absorption prints the load-time decision.
examples/mla_absorbed_decode_bench.rs: three-arm decode benchmark plus the KV-memory table.
docs/mla-absorbed-decode.md, indexed from docs/README.md, plus both variables in docs/environment-variables.md.

Why the cache is a view over `KVCache` and not a new type

LanguageModel::make_caches returns Vec<KVCache> and forward takes &mut [KVCache], so a genuinely new cache type forces the family into the ModelOwnedSequenceState escape hatch the SSM hybrids use, which costs supports_batching, the paged decode backend, and the generic prompt-cache donation path. The FP16 update path already stores keys and values as independent [B, H, L, D] buffers reading each side's head dim from its own shape, so a 512-wide "key" beside a 64-wide "value" is already expressible. offset, live_start, trim, trim_front, nbytes, and can_trim_prompt_cache therefore keep their existing meanings on latent rows, and the slot assignment matches what deepseek_v3.rs has always done so both describe the same buffer.

How #898's merge kernel is reused, and whether its contract held

Reused verbatim. No change to paged_attention_v2_merge.cpp, to paged_attention_v2.h, to the ffi::paged_attention_merge_states signature, or to the C++ launcher. MLA fits the documented contract as H = num_heads, D = kv_lora_rank, M = batch, N = batch * chunks, and o_indptr request-major. The contract held unchanged; the only thing the new caller had to honour rather than discover is the log2 LSE clause, which is the one that fails silently: a natural-log LSE still merges and returns a plausible wrong weighted average. It is a named constant (split_kv::LOG2_E) with a comment, and merge_rejects_natural_log_lse_units is a negative control that asserts the natural-log variant does not match a host f64 reference while the log2 variant does. split_kv_matches_the_decompressed_reference_at_every_chunk_length restates, at the MLA caller, the regrouping property paged_v2::launch_tests::merge_is_associative_across_regroupings pins for the kernel.

Verified vs unverified

Verified, on synthetic tensors with MLA geometry against a host-side f64 reference that performs the pre-absorption computation: the absorption identity in f32 and f16; that the additive causal mask survives being merged into the rope score (with a positive control asserting the mask actually changes the answer, so the masked test cannot pass on a no-op); the up-projection direction of the fold; the (ckv, kpe) packing round-tripping through a real KVCache including the append boundary; that a latent cache stays trimmable so speculative decode's reject path removes one latent row per token; the split-and-merge decomposition at four chunk lengths and against the unsplit Stage 1 answer; the log2 LSE negative control; and, through deepseek_v2's own MLAAttention::forward, that a 6-token prefill plus 4 decode steps produce the same block output (through o_proj) with and without the fold, step for step, and that the resulting cache is more than 3x smaller.

Not verified: anything against a real MLA checkpoint. No DeepSeek-family checkpoint is available on this host (the locally available models are Llama-3.1-8B, Llama-3.2-1B, Qwen2.5-0.5B, Qwen3-0.6B, and gemma-3-1b, none of them MLA, and a DeepSeek-V3-class checkpoint does not fit in 128 GB). So the issue's token-parity gate on pinned DeepSeek oracles is not satisfied by this PR, and the flag stays default-off for exactly that reason. Nothing here should be read as real-model validation.

No benchmarks were run. The harness is written and smoke-tested for correctness on a toy geometry only; all throughput numbers are for the reviewer to produce.

CUDA

No CUDA was written in this PR, so there is no new unvalidated CUDA to declare. Stage 2 reuses #898's merge kernel, which already ships a structurally parallel CUDA port that #898 declared unvalidated; that status is unchanged. The issue asked for a fused MLA partial kernel with CUDA first, and that is deliberately not in this PR: a fused partial kernel is the thing that would make Stage 2 faster, and it should be written after Stage 1's measurements justify it (as the issue itself scopes it) and against a host that can validate at least one of the two backends. The interactions the issue lists are unaffected: kv_a_layernorm is untouched, so the CUDA RMSNorm overlay (#829) sits where it sat; nothing here assumes CUDA graph replay.

How to run the benchmark, and how to prove which path it took

export DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer
cargo run --release --features metal,accelerate --example mla_absorbed_decode_bench -- \
    --contexts 4096,16384,32768 --batches 1,4 --steps 64 --warmup 16

Three arms per (context, batch): decompressed, absorbed, split_kv. Every row ends in a paths= field read from mlxcel_core::mla::stats, taken and reset around the measured region, and the harness prints a WARNING on any row whose counters disagree with its arm's label. A correct run looks like decompressed=64 absorbed_prefill=0 absorbed_composed=0 absorbed_split_kv=0 on the first arm and ... absorbed_composed=64 ... on the second; an "absorbed" row reading absorbed_composed=0 is measuring the fallback and must not be reported. split_kv is expected to be slower than absorbed here, because its partials are composed from MLX ops rather than fused; that arm exists to show the decomposition running, not to be quoted.

For a real checkpoint, MLXCEL_MLA_ABSORBED=1 mlxcel generate -m models/<deepseek> ... prints one stdout line at load stating how many layers folded and the bytes/token before and after. A run whose line reads 0/N layers is running the fallback. Both signals go to stdout rather than through tracing because the mlxcel CLI installs no tracing subscriber, so tracing::info! on this path emits nothing at any RUST_LOG.

Test plan

cargo test --release -p mlxcel-core --lib --features metal,accelerate mla:: — 33 passed, 0 failed
cargo test --release -p mlxcel-core --lib --features metal,accelerate paged_v2:: — 55 passed, 0 failed (the reused merge kernel and its regrouping property)
cargo test --release --lib --features metal,accelerate models::deepseek — 32 passed, 0 failed (V2 wiring parity plus the untouched V3 / V3.2 suites)
cargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings — clean
cargo clippy --lib --tests --example mla_absorbed_decode_bench --features metal,accelerate -- -D warnings — clean
cargo fmt --all -- --check — clean
python3 scripts/ci/check_cross_repo_refs.py — exit 0
Benchmark harness smoke-run on a toy geometry: all three arms execute and each reports its own path

The full cargo test -p mlxcel-core sweep was not run: it exceeds this environment's watchdog. Its known pre-existing failure on this base commit is fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape (tracked as #964), untouched here.

Left out, deliberately

Fused Stage 2 partial kernel (Metal and CUDA). MiniCPM3 and dots1 wiring. Converging deepseek_v3 / deepseek_v32 onto the shared module. Absorbed prefill (prefill up-projects through the still-loaded kv_b_proj, which keeps prefill arithmetic identical to the baseline but makes a chunked prefill re-up-project rows it already handled). Paged-backed and non-FP16 latent caches, both declined with a message. Any default-on change: the flag stays off until real-checkpoint token parity exists.

Closes #907

DeepSeek-family MLA compresses K and V into one latent of width kv_lora_rank (512) plus a shared rope stream of width qk_rope_head_dim (64). `src/models/deepseek_v2.rs`, `minicpm3.rs`, and `dots1.rs` up-project that latent through `kv_b_proj` and cache the decompressed per-head K/V, which costs `num_heads * (qk_head_dim + v_head_dim)` per token instead of 576. `deepseek_v3.rs` and `deepseek_v32.rs` already absorb, each with their own copy of the transform. This adds `mlxcel_core::mla`, the shared implementation the remaining families can be folded onto and the existing two can converge on. `mla::absorb` folds `kv_b_proj` at load time into `W_UK [1, H, qk_nope, r]` for the query path and `W_UV^T [1, H, r, v_head]` for the output path, dequantizing when the layer is quantized and refusing a biased layer (the identity has no term for a bias). Nothing on disk is touched. `mla::cache` documents and checks the packing that lets the latent cache reuse `KVCache` rather than forcing a new cache type: the FP16 update path already stores keys and values as independent `[B, H, L, D]` buffers reading each side's head dim from its own shape, so `ckv` goes in the key slot and `kpe` in the value slot, matching what deepseek_v3.rs has always done. That keeps `LanguageModel::make_caches -> Vec<KVCache>`, trim, trim_front, nbytes, batching, and prompt-cache donation working unchanged. Non-FP16 modes and paged backing are declined with a message. `mla::decode` is Stage 1: `SDPA(q_absorbed, ckv, ckv, scale, mask = scale * q_pe @ kpe^T + causal)`. The rope half of the score is additive in exactly SDPA's mask position, so the identity is exact and the fused MLX attention kernel still does the softmax and the value accumulation. `mla::split_kv` is Stage 2: it cuts the latent range into chunks that each produce a normalized partial and its log2 LSE, then merges them with issue #898's `paged_attention_merge_states`, reused unchanged (no edit to the kernel, the FFI signature, or the C++ launcher). The partial producer here is composed from MLX ops: correct, and the reference a fused partial kernel is validated against, but not a speed win on its own. `mla::stats` records which path every MLA attention call took. Issue #899 shipped a fused decode path that never activated and whose benchmark compared the fallback against itself; these counters make an arm's executed path provable from its output. `tracing` is deliberately unused because the `mlxcel` CLI installs no subscriber. Verified on synthetic tensors with MLA geometry against a host-side f64 reference that performs the pre-absorption computation. No MLA checkpoint fits the development host, so real-checkpoint token parity is not covered here. `merge_rejects_natural_log_lse_units` is the negative control on the one merge-contract clause that fails silently: a natural-log LSE still merges and produces a plausible wrong answer. Off unless `MLXCEL_MLA_ABSORBED` is set; with it unset nothing here is constructed and no weight is dequantized. Refs #907

…t cache Wires `src/models/deepseek_v2.rs` onto `mlxcel_core::mla` behind `MLXCEL_MLA_ABSORBED`, and adds the decode benchmark harness. `load_mla_attention` folds `kv_b_proj` only when the flag is on, so an unset environment dequantizes nothing and allocates nothing. The original `kv_b_proj` stays loaded either way: the absorbed prefill up-projects the live latent window through it and runs the same dense SDPA the decompressed path runs, which keeps prefill arithmetic identical to the baseline and leaves the fold responsible only for decode. The cost of that choice is that a chunked prefill re-up-projects rows it already up-projected on the previous chunk, which is where absorption's decode-only scope is paid for. `MLAAttention::forward` keeps the decompressed path byte-for-byte and branches to `forward_absorbed` only when a fold exists and `MlaLatentCache::supports` accepts the cache. The check is per cache rather than per model because the KV cache mode is a runtime flag; its answer is stable for the life of a cache, so a run that declines on step one declines on every step and the cache never mixes latent and decompressed rows. The prefill mask policy is deliberately identical to the decompressed path so this change cannot alter which steps are causal. `report_absorption` prints one stdout line per load stating how many layers folded, the bytes/token/layer before and after, and whether split-KV is on, plus the decline reason when the fold did not take. `tracing` would not work here: the `mlxcel` CLI installs no subscriber, so a `tracing::info!` on this path emits nothing at any `RUST_LOG`. A run whose flag is on but whose line says `0/27 layers` is visibly running the fallback. `MlaLatentCache::supports` is split out of `wrap` so a caller can fall back: the borrow checker keeps the `&mut KVCache` alive through the `Err` arm of a `Result` carrying that borrow's lifetime, which would make the fallback unexpressible. `examples/mla_absorbed_decode_bench.rs` measures one MLA attention block per decode step across three arms (decompressed, absorbed, split_kv) over a context and batch sweep, and prints the KV-memory table the issue asks for. Every arm prints its `mla::stats` counters, taken and reset around the measured region, and warns on the row itself when the counters disagree with the arm's label. Synthetic tensors, not a real checkpoint: MLA is a DeepSeek-family feature and no MLA checkpoint fits the development host, so these are the throughput of the layer that changed rather than end-to-end model throughput. Tests: `absorbed_mla_attention_matches_the_decompressed_block_step_for_step` drives the family's own forward over a 6-token prefill plus 4 decode steps with the same synthetic weights, one attention folded and one not, and compares every step's block output through `o_proj`. `absorbed_attention_leaves_a_much_smaller_cache_behind` asserts the cache actually shrank, so a wiring regression that silently kept the decompressed layout fails. Not covered: token parity against a real DeepSeek checkpoint. Refs #907

…erify Adds `docs/mla-absorbed-decode.md` and indexes it from `docs/README.md`, plus the two new variables in `docs/environment-variables.md`. The page covers the absorption identity, the KV-memory table it buys and the fixed weight memory it costs, the `(ckv, kpe)` cache packing and why it is a view over `KVCache` rather than a new cache type, which KV cache modes are declined and why, why prefill stays decompressed, and how Stage 2 reuses issue #898's merge kernel including the log2 LSE clause that fails silently. It also records which families this reaches: `deepseek_v2` is wired here, `deepseek_v3` and `deepseek_v32` have carried their own absorption since before this work and are unchanged, and `minicpm3` and `dots1` are still on the decompressed path. Two accuracy points are stated rather than implied: the benchmark measures one MLA attention block on synthetic tensors, not end-to-end model throughput, and no row should be reported unless its `paths=` field names its own arm. Also fixes a clippy `collapsible_if` in the `deepseek_v2` dispatch and a `print_literal` in the harness header, and adds `a_latent_cache_stays_trimmable_for_speculative_decode`, which pins that a rejected draft token removes one latent row so speculative decode keeps working against the new layout. Refs #907

Adds the performance-validation report issue #907 requires, covering the decompressed, absorbed and split-KV arms at context 4096/16384/32768 and batch 1/4, with every row carrying the path counters it was measured on. The memory result is structural rather than a timing: 8.89x fewer KV bytes per token per layer, which moves the maximum context under a 16 GiB KV budget from 62137 to 552336 tokens. Decode latency crosses over where the mechanism predicts. Absorption trades a larger per-step matmul against the folded projection for much lower KV traffic, so it loses 14 to 15 percent at 4096 tokens, wins 1.18x to 1.33x at 16384, and reaches 1.80x at batch 1 with 32768 tokens. Anyone enabling the flag on short contexts should expect it to cost them. Split-KV is slower than plain absorbed in every configuration, which was stated in advance and is expected: its partials are composed MLX ops rather than a fused kernel, so it pays the split without the kernel that would repay it. It ships as a correct implementation that exercises reuse of the #898 merge kernel, not as a speed path. Records what is not verified. No MLA checkpoint is available on this host and a V3-class model does not fit in 128 GB, so the issue's pinned-oracle token-parity gate is unsatisfied and the flag defaults to off. Everything else is verified against a host f64 pre-absorption reference on synthetic MLA geometry. Also records the premise correction: deepseek_v3 and deepseek_v32 already ship matrix absorption, so the actual gap was the V2 family plus the transform existing twice with no shared home.

`deepseek_v2` emitted token id 64279 (`}'_`) for every generated token on `models/deepseek-v2-lite-4bit` while mlx-lm decoded the same weights correctly. A per-stage numerical diff against the mlx-lm 0.31.3 reference on identical prompt ids localized it in one pass: in layer 0 every tensor feeding attention matched bit for bit (`compressed_kv`, the `kv_a_layernorm` latent, `kv_b_proj` output, `k_nope`, `values`, the RoPE'd `q_pe` and `k_pe`, and the concatenated `queries` and `keys`, all at 0.0 relative RMS), and the first stage to leave the noise floor was the SDPA output itself at 6.0e-1. Broken down per query position, the last prompt position matched to 1.9e-6 while position 0 was off by 8.5e-1, which is the signature of a missing causal mask rather than a wrong projection or a bad weight. Root cause: `MLAAttention::forward` had the mask policy inverted. The `l == 1` decode branch called `causal_attention`, where causality is vacuous, while the `l > 1` prefill branch handed the caller's mask straight to the unmasked `attention_from_ptr`. Every generation path (CLI text prefill, the VLM embeddings prefill, chunked prefill) calls the model with `mask == None` and expects the model to build its own causal mask, as documented on `pad_tokens_for_prefill` in `src/lib/mlxcel-core/src/generate.rs` (line 162) and as `llama3.rs`, `deepseek.rs` and `dots1.rs` implement it. So the one branch that needed a causal mask never got one: prefill ran fully bidirectional, every layer wrote future-contaminated K/V into the cache, and the first sampled token was already wrong. `deepseek_v3.rs` carries a comment describing the identical defect being fixed there under issue #525; `deepseek_v2.rs` never got the same treatment. Route a maskless multi-token prefill through `causal_attention` in both the decompressed path and the matrix-absorbed prefill added by #995. An explicit mask (the padded-prefill case) still goes to the masked call, and decode is unchanged because `causal_attention` takes its single-query maskless fast path at `l == 1`. After the fix every dumped stage sits inside f16 noise (1.3e-6 at the layer-0 SDPA, rising to 1.7e-3 at the logits after 27 layers of accumulation) and the 20-token greedy continuation matches the reference id for id. `tests/deepseek_v2_greedy_parity.rs` closes the coverage gap that let this ship: nothing in the tree ever generated from a real `deepseek_v2` text checkpoint. It has two gates, the reference ids, and a reference-free assertion that adding five later prompt tokens does not move the first position's logits. The old code failed the second by 1.35 relative RMS against a 2e-2 threshold, so the gate has a 67x margin and stays meaningful for any checkpoint of this family. `docs/supported-models.md` now names the validation checkpoint and states the maskless-prefill contract for all four DeepSeek text families, including that V3 and V3.2 have no local checkpoint and carry code inspection rather than a real-checkpoint gate. Closes #991

…che (lablup#995) ## Summary Adds `mlxcel_core::mla`, a shared matrix-absorbed MLA implementation over a compressed-latent KV cache, and wires `deepseek_v2` onto it behind `MLXCEL_MLA_ABSORBED` (default off). Stage 1 (cache format plus absorbed decode from composed MLX ops) and Stage 2 (split-KV reusing issue lablup#898's merge kernel unchanged) both land here; Stage 2's partial producer is composed from MLX ops, not a fused kernel, and is not a speed path. ## A correction to the issue's premise, up front The issue says the MLA families cache the decompressed K/V. That is true for `deepseek_v2`, `minicpm3`, and `dots1`. It is **not** true for `deepseek_v3` and `deepseek_v32`: both have carried matrix absorption since before this work, decomposing `kv_b_proj` into `embed_q` / `unembed_out` in `sanitize_weights` and caching `(kv_latent, k_pe)` with the rope score passed to SDPA as an additive mask. So the actual gap was V2, MiniCPM3, and dots1, plus the fact that the transform existed twice with no shared home. This PR builds that shared home and moves V2 onto it. V3/V3.2 are untouched (converging them is a follow-up with real regression risk that deserves its own PR), and MiniCPM3/dots1 are untouched and still decompressed. ## What changed - `src/lib/mlxcel-core/src/mla/absorb.rs`: load-time fold of `kv_b_proj` into `W_UK [1, H, qk_nope, r]` for the query path and `W_UV^T [1, H, r, v_head]` for the output path. Dequantizes when quantized (including the NVFP4 `global_scale` sidecar, folded as a scalar since `(x @ W^T) * s2 == x @ (W * s2)^T`), refuses a biased layer because the identity has no term for a bias, and refuses a geometry that does not partition the matrix. Nothing on disk is touched. - `src/lib/mlxcel-core/src/mla/cache.rs`: `MlaLatentCache`, a checked view over `KVCache` that puts `ckv` in the key slot and `kpe` in the value slot. `supports()` is split out of `wrap()` so a caller can fall back (the borrow checker keeps the `&mut` alive through the `Err` arm of a `Result` carrying that borrow's lifetime, which would otherwise make the fallback unexpressible). Also the `latent_bytes_per_token` / `decompressed_bytes_per_token` accounting the KV-memory table is computed from. - `src/lib/mlxcel-core/src/mla/decode.rs`: Stage 1. `SDPA(q_absorbed, ckv, ckv, scale, mask = scale * q_pe @ kpe^T + causal)`. The rope half of the score is additive in exactly SDPA's mask position, so the identity is exact and MLX's fused attention kernel still does the softmax and the value accumulation. Same shape `deepseek_v3.rs` already uses. - `src/lib/mlxcel-core/src/mla/split_kv.rs`: Stage 2. `MlaSplitPlan` cuts the latent range; each chunk produces a normalized partial and its log2 LSE; `paged_attention_merge_states` merges them. - `src/lib/mlxcel-core/src/mla/stats.rs`: per-path call counters and a `path=count` summary. - `src/models/deepseek_v2.rs`: `MLAAttention` gains an optional fold and a `forward_absorbed` branch; `load_mla_attention` builds the fold only under the flag; `report_absorption` prints the load-time decision. - `examples/mla_absorbed_decode_bench.rs`: three-arm decode benchmark plus the KV-memory table. - `docs/mla-absorbed-decode.md`, indexed from `docs/README.md`, plus both variables in `docs/environment-variables.md`. ## Why the cache is a view over `KVCache` and not a new type `LanguageModel::make_caches` returns `Vec<KVCache>` and `forward` takes `&mut [KVCache]`, so a genuinely new cache type forces the family into the `ModelOwnedSequenceState` escape hatch the SSM hybrids use, which costs `supports_batching`, the paged decode backend, and the generic prompt-cache donation path. The FP16 update path already stores keys and values as independent `[B, H, L, D]` buffers reading each side's head dim from its own shape, so a 512-wide "key" beside a 64-wide "value" is already expressible. `offset`, `live_start`, `trim`, `trim_front`, `nbytes`, and `can_trim_prompt_cache` therefore keep their existing meanings on latent rows, and the slot assignment matches what `deepseek_v3.rs` has always done so both describe the same buffer. ## How lablup#898's merge kernel is reused, and whether its contract held Reused verbatim. No change to `paged_attention_v2_merge.cpp`, to `paged_attention_v2.h`, to the `ffi::paged_attention_merge_states` signature, or to the C++ launcher. MLA fits the documented contract as `H = num_heads`, `D = kv_lora_rank`, `M = batch`, `N = batch * chunks`, and `o_indptr` request-major. The contract held unchanged; the only thing the new caller had to honour rather than discover is the **log2 LSE** clause, which is the one that fails silently: a natural-log LSE still merges and returns a plausible wrong weighted average. It is a named constant (`split_kv::LOG2_E`) with a comment, and `merge_rejects_natural_log_lse_units` is a negative control that asserts the natural-log variant does *not* match a host f64 reference while the log2 variant does. `split_kv_matches_the_decompressed_reference_at_every_chunk_length` restates, at the MLA caller, the regrouping property `paged_v2::launch_tests::merge_is_associative_across_regroupings` pins for the kernel. ## Verified vs unverified **Verified, on synthetic tensors with MLA geometry against a host-side f64 reference that performs the pre-absorption computation:** the absorption identity in f32 and f16; that the additive causal mask survives being merged into the rope score (with a positive control asserting the mask actually changes the answer, so the masked test cannot pass on a no-op); the up-projection direction of the fold; the `(ckv, kpe)` packing round-tripping through a real `KVCache` including the append boundary; that a latent cache stays trimmable so speculative decode's reject path removes one latent row per token; the split-and-merge decomposition at four chunk lengths and against the unsplit Stage 1 answer; the log2 LSE negative control; and, through `deepseek_v2`'s own `MLAAttention::forward`, that a 6-token prefill plus 4 decode steps produce the same block output (through `o_proj`) with and without the fold, step for step, and that the resulting cache is more than 3x smaller. **Not verified:** anything against a real MLA checkpoint. No DeepSeek-family checkpoint is available on this host (the locally available models are Llama-3.1-8B, Llama-3.2-1B, Qwen2.5-0.5B, Qwen3-0.6B, and gemma-3-1b, none of them MLA, and a DeepSeek-V3-class checkpoint does not fit in 128 GB). So the issue's token-parity gate on pinned DeepSeek oracles is **not** satisfied by this PR, and the flag stays default-off for exactly that reason. Nothing here should be read as real-model validation. **No benchmarks were run.** The harness is written and smoke-tested for correctness on a toy geometry only; all throughput numbers are for the reviewer to produce. ## CUDA No CUDA was written in this PR, so there is no new unvalidated CUDA to declare. Stage 2 reuses lablup#898's merge kernel, which already ships a structurally parallel CUDA port that lablup#898 declared unvalidated; that status is unchanged. The issue asked for a fused MLA partial kernel with CUDA first, and that is deliberately not in this PR: a fused partial kernel is the thing that would make Stage 2 faster, and it should be written after Stage 1's measurements justify it (as the issue itself scopes it) and against a host that can validate at least one of the two backends. The interactions the issue lists are unaffected: `kv_a_layernorm` is untouched, so the CUDA RMSNorm overlay (lablup#829) sits where it sat; nothing here assumes CUDA graph replay. ## How to run the benchmark, and how to prove which path it took ``` export DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer cargo run --release --features metal,accelerate --example mla_absorbed_decode_bench -- \ --contexts 4096,16384,32768 --batches 1,4 --steps 64 --warmup 16 ``` Three arms per (context, batch): `decompressed`, `absorbed`, `split_kv`. Every row ends in a `paths=` field read from `mlxcel_core::mla::stats`, taken and reset around the measured region, and the harness prints a `WARNING` on any row whose counters disagree with its arm's label. A correct run looks like `decompressed=64 absorbed_prefill=0 absorbed_composed=0 absorbed_split_kv=0` on the first arm and `... absorbed_composed=64 ...` on the second; an "absorbed" row reading `absorbed_composed=0` is measuring the fallback and must not be reported. `split_kv` is expected to be slower than `absorbed` here, because its partials are composed from MLX ops rather than fused; that arm exists to show the decomposition running, not to be quoted. For a real checkpoint, `MLXCEL_MLA_ABSORBED=1 mlxcel generate -m models/<deepseek> ...` prints one stdout line at load stating how many layers folded and the bytes/token before and after. A run whose line reads `0/N layers` is running the fallback. Both signals go to stdout rather than through `tracing` because the `mlxcel` CLI installs no tracing subscriber, so `tracing::info!` on this path emits nothing at any `RUST_LOG`. ## Test plan - [x] `cargo test --release -p mlxcel-core --lib --features metal,accelerate mla::` — 33 passed, 0 failed - [x] `cargo test --release -p mlxcel-core --lib --features metal,accelerate paged_v2::` — 55 passed, 0 failed (the reused merge kernel and its regrouping property) - [x] `cargo test --release --lib --features metal,accelerate models::deepseek` — 32 passed, 0 failed (V2 wiring parity plus the untouched V3 / V3.2 suites) - [x] `cargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings` — clean - [x] `cargo clippy --lib --tests --example mla_absorbed_decode_bench --features metal,accelerate -- -D warnings` — clean - [x] `cargo fmt --all -- --check` — clean - [x] `python3 scripts/ci/check_cross_repo_refs.py` — exit 0 - [x] Benchmark harness smoke-run on a toy geometry: all three arms execute and each reports its own path The full `cargo test -p mlxcel-core` sweep was not run: it exceeds this environment's watchdog. Its known pre-existing failure on this base commit is `fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape` (tracked as lablup#964), untouched here. ## Left out, deliberately Fused Stage 2 partial kernel (Metal and CUDA). MiniCPM3 and dots1 wiring. Converging `deepseek_v3` / `deepseek_v32` onto the shared module. Absorbed prefill (prefill up-projects through the still-loaded `kv_b_proj`, which keeps prefill arithmetic identical to the baseline but makes a chunked prefill re-up-project rows it already handled). Paged-backed and non-FP16 latent caches, both declined with a message. Any default-on change: the flag stays off until real-checkpoint token parity exists. Closes lablup#907

`deepseek_v2` emitted token id 64279 (`}'_`) for every generated token on `models/deepseek-v2-lite-4bit` while mlx-lm decoded the same weights correctly. A per-stage numerical diff against the mlx-lm 0.31.3 reference on identical prompt ids localized it in one pass: in layer 0 every tensor feeding attention matched bit for bit (`compressed_kv`, the `kv_a_layernorm` latent, `kv_b_proj` output, `k_nope`, `values`, the RoPE'd `q_pe` and `k_pe`, and the concatenated `queries` and `keys`, all at 0.0 relative RMS), and the first stage to leave the noise floor was the SDPA output itself at 6.0e-1. Broken down per query position, the last prompt position matched to 1.9e-6 while position 0 was off by 8.5e-1, which is the signature of a missing causal mask rather than a wrong projection or a bad weight. Root cause: `MLAAttention::forward` had the mask policy inverted. The `l == 1` decode branch called `causal_attention`, where causality is vacuous, while the `l > 1` prefill branch handed the caller's mask straight to the unmasked `attention_from_ptr`. Every generation path (CLI text prefill, the VLM embeddings prefill, chunked prefill) calls the model with `mask == None` and expects the model to build its own causal mask, as documented on `pad_tokens_for_prefill` in `src/lib/mlxcel-core/src/generate.rs` (line 162) and as `llama3.rs`, `deepseek.rs` and `dots1.rs` implement it. So the one branch that needed a causal mask never got one: prefill ran fully bidirectional, every layer wrote future-contaminated K/V into the cache, and the first sampled token was already wrong. `deepseek_v3.rs` carries a comment describing the identical defect being fixed there under issue lablup#525; `deepseek_v2.rs` never got the same treatment. Route a maskless multi-token prefill through `causal_attention` in both the decompressed path and the matrix-absorbed prefill added by lablup#995. An explicit mask (the padded-prefill case) still goes to the masked call, and decode is unchanged because `causal_attention` takes its single-query maskless fast path at `l == 1`. After the fix every dumped stage sits inside f16 noise (1.3e-6 at the layer-0 SDPA, rising to 1.7e-3 at the logits after 27 layers of accumulation) and the 20-token greedy continuation matches the reference id for id. `tests/deepseek_v2_greedy_parity.rs` closes the coverage gap that let this ship: nothing in the tree ever generated from a real `deepseek_v2` text checkpoint. It has two gates, the reference ids, and a reference-free assertion that adding five later prompt tokens does not move the first position's logits. The old code failed the second by 1.35 relative RMS against a 2e-2 threshold, so the gate has a 67x margin and stays meaningful for any checkpoint of this family. `docs/supported-models.md` now names the validation checkpoint and states the maskless-prefill contract for all four DeepSeek text families, including that V3 and V3.2 have no local checkpoint and carry code inspection rather than a real-checkpoint gate. Closes lablup#991

inureyes added 4 commits August 2, 2026 12:46

inureyes force-pushed the feature/issue-907-mla-absorbed-decode branch from 53cacc7 to 73a4e90 Compare August 2, 2026 03:49

inureyes merged commit a93d365 into main Aug 2, 2026
5 checks passed

inureyes mentioned this pull request Aug 2, 2026

fix(models): make the deepseek_v2 prefill causal #998

Merged

5 tasks done

inureyes mentioned this pull request Aug 3, 2026

Epic: Fused paged decode, sorting-free sampling, and serving performance techniques #909

Closed

inureyes deleted the feature/issue-907-mla-absorbed-decode branch August 4, 2026 12:08

inureyes self-assigned this Aug 31, 2026

Navigation Menu

Repository navigation

feat(core): matrix-absorbed MLA decode over a compressed-latent KV cache - #995

feat(core): matrix-absorbed MLA decode over a compressed-latent KV cache#995
inureyes merged 4 commits into
mainfrom
feature/issue-907-mla-absorbed-decode

inureyes commented Aug 2, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

inureyes commented Aug 2, 2026

Summary

A correction to the issue's premise, up front

What changed

Why the cache is a view over KVCache and not a new type

How #898's merge kernel is reused, and whether its contract held

Verified vs unverified

CUDA

How to run the benchmark, and how to prove which path it took

Test plan

Left out, deliberately

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Why the cache is a view over `KVCache` and not a new type