DS4Server/docs/issue113-head-chunk

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.