Match Qwen head-chunked attention execution

docs/issue113-head-chunk/README.md

@@ -0,0 +1,123 @@
+# 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.