# Qwen head-chunked SDPA (#113)

The dense Qwen attention consumer now follows pinned MTPLX `_verify_sdpa`:
when enabled (default) and `rows * GQA > 32` with `rows <= 32`, split query
heads inside each KV group into `min(GQA, 32 / rows)` heads, preserve remainder
heads, call the existing SDPA primitive with that group's K/V, then concatenate
in original head order. The first engagement is logged. Explicit false/off/no/0
keeps the prior consumer. Policy is resolved through the existing Qwen lookup;
no process environment mutation or shared-family SDPA change was introduced.

Oracle: MTPLX `9882703f3105363ddc37eca9f97aa09a1d387112` / MLX 0.32.2.
Artifact: Bare-Speed `74559cdf34fbfc0b593de72d17e93f37fd4f9ea7`; environment
and artifact fingerprints in ../issue104-ar/environment.json. Pre-port native
commit: `7b894fbfeed81d2944b9cf479e60638ef04987c7`.

## Exact scope found during review

The wrapper's name does not restrict it to verification. It is the final dense
consumer after sparse/gather branches in `Attention.__call__`, in any phase.
Thus short prefills also use it. Existing sparse selections, gathered consumers,
gates, projection order, RoPE and cache offsets are unchanged.

The pinned `_sdpa_head_chunked` passes `mask` **unchanged** to each subcall. It
does not slice a full-query-head mask. Broadcast-compatible masks work; a mask
whose head dimension cannot broadcast to a chunk is rejected. The fixtures
preserve these errors as well as valid head-specific masks. This resolves the
issue body's tentative mask-slicing wording against the actual pinned source.

For installed head dimension 256, MLX 0.32.2 uses the vector primitive only
through eight query rows. The wrapper still splits heads for eligible rows
9–32, and those subcalls retain the existing fallback primitive. The port does
not force a vector route merely because the wrapper docstring describes it.
Wide rows above 32 pass through unchanged.

## Cache compatibility

Checkpoint identity now records the head-chunk policy. Missing metadata means
the actual old unchunked route. A policy difference rejects all MTP states and
AR states with four or more tokens, before reading or allocating tensor blobs.
MTP can retain a short frontier from a wider verify window; an AR prefix of at
most three tokens cannot have executed a three-row prefill body and remains
restorable. Existing format and unrelated identity checks remain intact. Policy
changes in both directions are covered, including current/legacy metadata.

The product test rejects an actual pre-port checkpoint produced during #109,
without changing the live session, then saves/reopens corrected state and
compares the complete warm frontier. The explicit-off control restores that
same actual old file exactly. Fake tiny checkpoint fixtures additionally cover
AR/MTP × 1/3/4-token frontiers × both producing/consuming policies and missing
metadata; affected files stop after header/manifest reads.

## Isolated oracle evidence

- 520 independent plan cases cover boundary/invalid head counts and rows.
- 216 installed-dimension attention cases include single/grouped/ungrouped
  heads, remainder heads, batch two, offset/strided Q/K/V, no/bool/additive/
  two-dimensional/full-head/chunk-local masks, default on and explicit off,
  short KV, and contexts through 4097 tokens. All 159 successful outputs match
  raw BF16 bits and all 57 invalid cases match oracle rejection. Native graph
  subcall shapes, fused primitive counts, vector dispatches and fallback
  softmax counts match the independently captured oracle route.
- The exact reported request (`Count from one to ten.`, low reasoning, seed 1,
  greedy, depth 3, max 12, context-copy disabled) is checked cold, warm and after
  disk reopen. Default and off controls have separate native/reference runs.
  Every verify input row, output logit/hidden byte, emitted token, final array
  and saved trunk/MTP cache bit is compared; no zero-sign or numeric tolerance
  exception is allowed in this issue's oracle tool.

Comparisons use the same snapshot stage as the server: lazy hybrid trunk and
cloned MTP head. Taking an extra deep clone while measuring a lazy key array
would add zero and change signed-zero bits. Warm suffix prefill's final hidden
row is compared after the same final-row selection used by the upstream owner
(`generation.py::_prefill_restored_prompt_suffix`), not against its unsliced
multi-row intermediate. Verify-row arrays themselves are compared in full.

## Reproduction

Generate the independent fixture using isolated MLX 0.32.2 Python:
`tools/issue113-reference.py PINNED_CHECKOUT tests/fixtures/mtplx-head-chunk.json`.
Build the release test executable with `cargo test --release --all-features
--bin ds4-server issue113 --no-run`, then copy the executable outside `target`
before running it. Run `issue113_plan_and_policy_match_pinned_oracle` normally.

Run the ignored tests serially under `target/debug/test-supervisor MEMORY_MIB
60 120 --command TEST_BINARY FILTER --ignored --nocapture --test-threads=1`:

- `issue113_head_chunk_outputs_and_dispatch`: 4096 MiB.
- `qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state`: 4096 MiB,
  with `DS4_QWEN38_CHECKPOINT_DIR` set to an empty output directory.
- `issue113_default_verify_and_checkpoint`: 114688 MiB, with
  `DS4_QWEN38_ARTIFACTS=MODEL_DIR`, `MTPLX_CONTEXT_COPY=0`,
  `DS4_ISSUE109_DIR=OUTPUT_DIRECTORY` (shared real-request harness), and
  `DS4_ISSUE113_OLD=ACTUAL_PRE_PORT_CHECKPOINT`. Repeat separately with
  `MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK=0` for the off control.

For each model run, use isolated reference Python to execute
`tools/issue113-model-reference.py PINNED_CHECKOUT MODEL_DIR NATIVE_OUTPUT
ORACLE_OUTPUT`, with the matching default/off switch. Keep child output
separate from supervisor resource logs to avoid interleaved JSON records.

## Completion review

The four-row reproduction `[760,1156,369,9859]` now takes four vector calls
(8+4 query heads for each of two KV groups). Its logits and hidden arrays differ
from the old ordinary fallback and match the default pinned oracle exactly.
Both cold and warm turns have four verify forwards and 12 emitted tokens; the
warm turn restores 58 tokens. All verify rows, prefill final-row state, final
logits/hidden, trunk arrays and MTP arrays match raw reference bits. Disk reopen
reproduces the entire warm state. The off route matches all retained #109
pre-port verify, final-logit/hidden and full cache-array evidence byte for byte.

Regression controls pass for the QSA eager selector, dense masks, tiled prefill
gather, row gather, unchunked SDPA, complete QSA modules with TF32 off/on, and
exhaustive gate/injection laws. Installed DeepSeek resident-session switching
and GLM verifier/recurrent-state controls pass. Review confirms no production
DS4/GLM code or shader changed and no shared-family helper was redirected.
This is the pinned Qwen execution route, not a DS4Server-only feature.

All commit gates pass: formatting, warning-free Clippy, signed application
bundle, and all-features tests (313 application + 8 tool tests, 236 explicitly
ignored hardware/reference tests). Relevant ignored tests ran separately under
the resource supervisor; their receipts and the complete gate outputs are
stored here. No implementation or acceptance item is stubbed or deferred.