refactor: move KV attention dispatch into the cache (#2171 Phase 4a)

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).

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 로그, 변경 없음).

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