124 lines
7.2 KiB
Markdown
124 lines
7.2 KiB
Markdown
# 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.
|