Preserve MTPLX head-chunked Qwen verify SDPA execution #113

Closed
opened 2026-10-04 12:33:12 +00:00 by hugo · 1 comment
Owner

Discovered while independently validating #109 against pinned MTPLX 9882703f3105363ddc37eca9f97aa09a1d387112 / MLX 0.32.2. This is a separate execution-route and numerical mismatch, not the attention-gate defect.

Reproduction and boundary

A real greedy Qwen request Count from one to ten. (low reasoning, seed 1, MTP depth 3, max 12, context-copy disabled) has a 46-token rendered prompt. Prefill logits/hidden and the first three decoder-layer outputs agree exactly. In its first target verification [760,1156,369,9859], the first QSA layer receives identical hyper-connection mixed input, injection weights and attention gate values, but its attention output already differs BEFORE gating. This propagates into verify hidden/logits and warm state.

Pinned qwen4_exp.py::_verify_sdpa defaults MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK ON. It divides query heads within their KV groups when q_len * (q_heads / kv_heads) > 32 and q_len <= 32, with heads_per_chunk = max(1, min(gqa, 32 // q_len)); masks are sliced/broadcast according to _sdpa_head_chunked, and output is concatenated in original head order. DS4Server qsa_attention::Attention.call instead always calls attention::scaled_dot_product, which sends the shipped 24-query/2-KV-head, dim256, four-row case to the ordinary fallback.

With the oracle's explicit MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK=0 control, #109's corrected native implementation matches every real cold/warm verify row, hidden/logit byte, generated token, final logits and hidden state. The default oracle still differs. This isolates the missing consumer route and does not authorize treating the off control as default parity.

Implementation

Port the pinned _verify_sdpa_head_chunk_plan and _sdpa_head_chunked execution route in Rust with existing native Metal primitives. Preserve the current default-on/explicit-off policy and exact eligibility, per-KV grouping, tail heads, broadcast/head-specific masks, primitive selection and output order. Use the upstream route's actual scope (do not infer eligibility solely from the function name). Preserve unchunked/text controls and all other model families. No Python/C execution in product and no global SDPA/math-flag change.

Audit numerical state compatibility because previously banked multi-row Qwen states contain the older attention result. Ensure obsolete affected persisted states cannot contaminate corrected warm turns; preserve unaffected formats where possible.

Acceptance

  • Independent oracle/native real-geometry fixtures around plan boundaries: single/multiple query rows, grouped and ungrouped heads, remainder heads, masks including head-specific masks, default on and explicit off, and unsupported/wide cases.
  • Verify actual dispatch/shape route, not merely approximately equal outputs; require exact reference bits where the primitive contract requires them.
  • Run the real cold/warm request above against the DEFAULT pinned oracle and compare all verify rows, final state, warm/disk reopen and generated tokens. Keep the off path as a control.
  • Preserve DS4/GLM and unrelated Qwen prefill, sparse selections, gates and cache offsets.
  • Run all repository commit gates; one complete commit, push, implementation/test comment, and closure.

Reproduction tools/receipts are being committed with #109 under tools/issue109-model-reference.py, the ignored issue109_verify_rows_and_warm_checkpoint native test and docs/issue109-verify-gate/. Their #109 comparison explicitly selects the existing unchunked consumer and does not claim this default route is implemented.

Discovered while independently validating #109 against pinned MTPLX `9882703f3105363ddc37eca9f97aa09a1d387112` / MLX 0.32.2. This is a separate execution-route and numerical mismatch, not the attention-gate defect. ## Reproduction and boundary A real greedy Qwen request `Count from one to ten.` (low reasoning, seed 1, MTP depth 3, max 12, context-copy disabled) has a 46-token rendered prompt. Prefill logits/hidden and the first three decoder-layer outputs agree exactly. In its first target verification `[760,1156,369,9859]`, the first QSA layer receives identical hyper-connection mixed input, injection weights and attention gate values, but its attention output already differs BEFORE gating. This propagates into verify hidden/logits and warm state. Pinned `qwen4_exp.py::_verify_sdpa` defaults `MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK` ON. It divides query heads within their KV groups when `q_len * (q_heads / kv_heads) > 32` and `q_len <= 32`, with `heads_per_chunk = max(1, min(gqa, 32 // q_len))`; masks are sliced/broadcast according to `_sdpa_head_chunked`, and output is concatenated in original head order. DS4Server `qsa_attention::Attention.call` instead always calls `attention::scaled_dot_product`, which sends the shipped 24-query/2-KV-head, dim256, four-row case to the ordinary fallback. With the oracle's explicit `MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK=0` control, #109's corrected native implementation matches every real cold/warm verify row, hidden/logit byte, generated token, final logits and hidden state. The default oracle still differs. This isolates the missing consumer route and does not authorize treating the off control as default parity. ## Implementation Port the pinned `_verify_sdpa_head_chunk_plan` and `_sdpa_head_chunked` execution route in Rust with existing native Metal primitives. Preserve the current default-on/explicit-off policy and exact eligibility, per-KV grouping, tail heads, broadcast/head-specific masks, primitive selection and output order. Use the upstream route's actual scope (do not infer eligibility solely from the function name). Preserve unchunked/text controls and all other model families. No Python/C execution in product and no global SDPA/math-flag change. Audit numerical state compatibility because previously banked multi-row Qwen states contain the older attention result. Ensure obsolete affected persisted states cannot contaminate corrected warm turns; preserve unaffected formats where possible. ## Acceptance - Independent oracle/native real-geometry fixtures around plan boundaries: single/multiple query rows, grouped and ungrouped heads, remainder heads, masks including head-specific masks, default on and explicit off, and unsupported/wide cases. - Verify actual dispatch/shape route, not merely approximately equal outputs; require exact reference bits where the primitive contract requires them. - Run the real cold/warm request above against the DEFAULT pinned oracle and compare all verify rows, final state, warm/disk reopen and generated tokens. Keep the off path as a control. - Preserve DS4/GLM and unrelated Qwen prefill, sparse selections, gates and cache offsets. - Run all repository commit gates; one complete commit, push, implementation/test comment, and closure. Reproduction tools/receipts are being committed with #109 under `tools/issue109-model-reference.py`, the ignored `issue109_verify_rows_and_warm_checkpoint` native test and `docs/issue109-verify-gate/`. Their #109 comparison explicitly selects the existing unchunked consumer and does not claim this default route is implemented.
hugo added the bug label 2026-10-04 12:33:12 +00:00
Author
Owner

Implemented and pushed in d9026aa7bfbc8f74ec82c120178c34c5f08510d3.

The Qwen dense consumer now reproduces pinned _verify_sdpa in native Rust: default-on policy, exact eligibility and per-KV head plan, remainder heads, original subcall order and concatenation. It uses existing native primitives; no global SDPA/math flag or DS4/GLM path changed.

Review against the actual pinned source resolved two details: the wrapper applies in every eligible phase, including short prefills; and it passes masks unchanged to each subcall, allowing compatible broadcasts and rejecting incompatible head-specific masks. For D256 rows 9–32, the wrapper still chunks but MLX 0.32.2 keeps its fallback primitive; we preserve that rather than forcing vector dispatch.

Checkpoint identity records the producing head-chunk policy. Missing metadata means the pre-port unchunked path. Policy differences reject MTP and potentially affected AR prefixes before tensor reads, while unaffected <=3-token AR prefixes remain compatible. Both directions and legacy/current metadata are tested. An actual #109 pre-port checkpoint is rejected under the corrected default and restores exactly with explicit off; corrected state survives disk reopen and reproduces the whole warm frontier.

Isolated verification:

  • 520 oracle plan cases; 216 real-D256 attention fixtures: all 159 successful results match raw BF16 bits and all 57 invalid cases match rejection. Subcall shapes, fused primitive counts, vector dispatches and fallback softmax counts match. Coverage includes grouped/ungrouped/remainder heads, batch two, strided/offset arrays, mask variants, default/off, plan boundaries and 4097-token KV.
  • The exact 46-token reported request, cold/warm/disk reopen, greedy seed 1 and depth 3: every verify input/output row, final logit/hidden byte, generated token, and trunk/MTP cache bit matches DEFAULT pinned MTPLX. Both turns have four verify forwards and 12 tokens; warm restores 58 tokens. No numeric or signed-zero exception is permitted.
  • Separate explicit-off native/oracle runs pass and preserve all retained #109 pre-port verify/final/cache evidence byte for byte.
  • QSA eager selection, dense masks, prefill gather, rows gather, unchunked SDPA, module TF32 off/on and gate/injection regressions pass. Installed DeepSeek session switching and GLM verifier/recurrent-state controls pass.
  • All four commit gates pass: fmt, Clippy -D warnings, make bundle, all-features tests (313 application + 8 tool tests). Relevant ignored Metal/reference checks ran separately under the supervisor.

Review, reproduction, exact receipts, snapshot-stage details and gate evidence. This is the pinned Qwen execution route, not a new DS4Server-only feature. The full scope is complete with no deferred acceptance item.

Implemented and pushed in [d9026aa7bfbc8f74ec82c120178c34c5f08510d3](https://git.rfc1437.de/hugo/DS4Server/commit/d9026aa7bfbc8f74ec82c120178c34c5f08510d3). The Qwen dense consumer now reproduces pinned `_verify_sdpa` in native Rust: default-on policy, exact eligibility and per-KV head plan, remainder heads, original subcall order and concatenation. It uses existing native primitives; no global SDPA/math flag or DS4/GLM path changed. Review against the actual pinned source resolved two details: the wrapper applies in every eligible phase, including short prefills; and it passes masks unchanged to each subcall, allowing compatible broadcasts and rejecting incompatible head-specific masks. For D256 rows 9–32, the wrapper still chunks but MLX 0.32.2 keeps its fallback primitive; we preserve that rather than forcing vector dispatch. Checkpoint identity records the producing head-chunk policy. Missing metadata means the pre-port unchunked path. Policy differences reject MTP and potentially affected AR prefixes before tensor reads, while unaffected <=3-token AR prefixes remain compatible. Both directions and legacy/current metadata are tested. An actual #109 pre-port checkpoint is rejected under the corrected default and restores exactly with explicit off; corrected state survives disk reopen and reproduces the whole warm frontier. Isolated verification: - 520 oracle plan cases; 216 real-D256 attention fixtures: all 159 successful results match raw BF16 bits and all 57 invalid cases match rejection. Subcall shapes, fused primitive counts, vector dispatches and fallback softmax counts match. Coverage includes grouped/ungrouped/remainder heads, batch two, strided/offset arrays, mask variants, default/off, plan boundaries and 4097-token KV. - The exact 46-token reported request, cold/warm/disk reopen, greedy seed 1 and depth 3: every verify input/output row, final logit/hidden byte, generated token, and trunk/MTP cache bit matches DEFAULT pinned MTPLX. Both turns have four verify forwards and 12 tokens; warm restores 58 tokens. No numeric or signed-zero exception is permitted. - Separate explicit-off native/oracle runs pass and preserve all retained #109 pre-port verify/final/cache evidence byte for byte. - QSA eager selection, dense masks, prefill gather, rows gather, unchunked SDPA, module TF32 off/on and gate/injection regressions pass. Installed DeepSeek session switching and GLM verifier/recurrent-state controls pass. - All four commit gates pass: fmt, Clippy -D warnings, make bundle, all-features tests (313 application + 8 tool tests). Relevant ignored Metal/reference checks ran separately under the supervisor. [Review, reproduction, exact receipts, snapshot-stage details and gate evidence](https://git.rfc1437.de/hugo/DS4Server/src/commit/d9026aa7bfbc8f74ec82c120178c34c5f08510d3/docs/issue113-head-chunk). This is the pinned Qwen execution route, not a new DS4Server-only feature. The full scope is complete with no deferred acceptance item.
hugo closed this issue 2026-10-04 14:45:51 +00:00
Sign in to join this conversation.