Issue 110: guarded stock expert gathers
Qwen's stock MoE gathers now mirror pinned MTPLX's
moe_sorted_gather._call_pad / gather_qmm wrapper. On hardware that actually
uses NAX, an expert-sorted affine transposed [rows, 1, K] call above 32,767
rows is padded to the next multiple of 64 unless already aligned. Activations
are zero padded, the final expert ID is repeated, explicit left indices append
an arange in their own dtype, and the output is sliced back to its real rows.
Both separate/fused Gate/Up and Down use this entry point. The ordinary runtime
primitive and pinned 0.32.2 shader remain unchanged.
The hardware test deliberately ignores the MTPLX route-fallback switch, as the stock MLX shader does. Unit controls exercise non-NAX and ineligible shapes; no physical M1–M4 hardware result is claimed. This implementation is confined to Qwen; DS4/GLM have no callers of the new wrapper. It is an oracle correctness fix, not a new DS4Server feature.
Old checkpoint manifests lack both producer hardware and prefill chunk sizes. Files with at least 3,277 tokens may contain affected rows and are rejected before state restoration unless they carry the new guard policy. Short old files remain compatible; new files retain the existing container format. The admission test verifies rejection leaves healthy live state unchanged.
Verification
- Original initialized-NaN low-level sentinels still demonstrate the pinned shader defect. The same five token counts (3,276 / 3,277 / 3,404 / 4,095 / 4,096) now additionally execute the real guarded graph and require every returned row to equal the exact expected BF16 value, with original shape.
- Installed layer-0 weights exercise 512 experts, fused Gate/Up K=2560/N=1280 and Down K=640/N=2560, group64/4-bit. All output bytes match the pinned guarded oracle at 20,480 / 32,760 / 32,770 / 34,040 / 40,950 / 40,960 rows. Explicit reverse left indices at 34,040 rows also match. Sentinel weights control group32/4-bit; the unchanged runtime-gather suites cover other dispatches.
- The current optimized oracle's row-map SwiGLU equals its guarded stock chain byte-for-byte at 2,048 / 3,277 / 4,095 / 4,096 tokens, with successful first-use canaries in both narrow/wide regimes and actual kernel engagement recorded. This is an oracle control. The native optimized route belongs to #112.
- Allocation refusal releases the failed padded graph; a fresh valid retry computes all rows exactly and returns allocator activity to zero.
- Full-model cold and warm checks use actual server execution defaults, greedy
generation, 7,374 cold tokens, then a 3,278-token suffix. Both 2,048 and 4,096
chunk sizes are exercised.
MTPLX_GDN_BOUNDARY_TAIL_INTERVAL=0is explicit: default boundary splitting would divide the final partial span into 256-row pieces and conceal this reproduction. Cold and warm affected body spans are 3,277 tokens; cancellation occurs after a settled affected chunk and preserves the committed state for retry. The default 2,048 size remains a control.
All four full-model cases match the oracle byte-for-byte in final logits, token IDs and every saved trunk/cache tensor. The required formatting, warnings-denied Clippy, signed bundle and ordinary tests pass (309 application + 8 tool tests). The official focused upstream suite passes 16 tests, including both block equivalence variants with the canary/kernel active. One ordinary-suite run hit the unrelated numeric Git-revision YAML fixture flake tracked as #115; the unchanged suite passed on rerun.
The oracle is MTPLX 9882703f3105363ddc37eca9f97aa09a1d387112 with MLX exactly
0.32.2 and the unchanged installed Bare-Speed revision from
../issue104-ar/environment.json. No runtime/assets/default chunk changes or
performance claims accompany this correctness fix. See checks.json and
adjacent receipts for actual completed results.
Reproduction
local-eval-results/mtplx-reference-env-0.32.2/bin/python \
tools/issue110-reference.py /path/to/pinned/mtplx /path/to/model \
> /tmp/issue110-oracle.jsonl
cargo test --all-features --bin ds4-server --no-run
cargo build --bin test-supervisor
export DS4_QWEN38_ARTIFACTS=/path/to/model
export DS4_ISSUE110_ORACLE=/tmp/issue110-oracle.jsonl
# Use the executable printed by cargo; run GPU processes serially.
target/debug/test-supervisor 4096 60 120 --command TEST_EXECUTABLE \
issue103_sorted_gather_receipts --ignored --nocapture --test-threads=1
target/debug/test-supervisor 16384 60 120 --command TEST_EXECUTABLE \
issue110_actual_projection_rows_match_oracle --ignored --nocapture --test-threads=1
export DS4_ISSUE110_DIR=/tmp/issue110-product
export MTPLX_GDN_BOUNDARY_TAIL_INTERVAL=0
target/debug/test-supervisor 114688 60 120 --command TEST_EXECUTABLE \
issue110_large_cold_and_warm_prefill --ignored --nocapture --test-threads=1
local-eval-results/mtplx-reference-env-0.32.2/bin/python \
tools/issue110-model-reference.py /path/to/pinned/mtplx \
"$DS4_QWEN38_ARTIFACTS" "$DS4_ISSUE110_DIR" /tmp/issue110-model-oracle