Qwen row-map expert gather (#112)

The sorted Qwen gate/up path now retains the uint32 token row map, builds GPU per-expert row/tile prefix sums, and dispatches the pinned 64×64×64 NAX kernel. The kernel reads original token rows and emits SwiGLU directly. Down projection continues through the #110 guarded stock gather. The unsorted decode/verify path is unchanged. Packed and separate gate/up weights preserve their original activation rounding and fallback ordering.

Oracle: MTPLX 9882703f3105363ddc37eca9f97aa09a1d387112, MLX 0.32.2. Installed Bare-Speed artifact: 74559cdf34fbfc0b593de72d17e93f37fd4f9ea7; artifact hashes and machine details are in ../issue104-ar/environment.json. Native pre-port baseline: 8965e77ea173c932e00f70759d0a83ac1c5072f0.

Only GPU shader source is shipped from the oracle: the MTPLX kernel body and its MLX 0.32.2 GPU template/functor closure, with MIT licenses in metal/. tools/issue112-extract-metal.py reproduces those files from the pinned reference environment. All scheduling, policy, graph construction, checking, model loading and sessions remain Rust. There is no production Python/MLX/C runtime dependency and the pinned runtime metallib is unchanged.

Route and failure contract

All original tensor/layout predicates, affine BF16/F16, b4/b8, groups 32/64/128, N/K multiples of 64, and the 4,096-routed-row threshold apply.
Below-threshold, non-NAX and ineligible inputs retain the stock path. MTPLX_MOE_SORTED_GATHER_KERNEL=0 and the GPU-family rehearsal switch are resolved by Qwen policy without mutating the process environment.
Whole-first-call canaries are keyed by compiled instantiation and the under/from-64-rows-per-expert regime. They compare uint16 bits, including signed zero, then reuse the checked result. Compile/dispatch failures or mismatches disable that key for the process with a reason and counters.
A failed packed SwiGLU canary tries the separately checked plain row-map gather before guarded stock fallback, as upstream does. Separate SwitchGLU returns to its original up/gate/compiled-activation chain.
The only generalized helper is Qwen's GatherQMM: its existing BF16 behavior is retained, with the F16/group-128 specializations needed for the oracle's supported canary/fallback contract. All callers are in the Qwen module.

Isolated verification

The native synthetic matrix covers 72 cases: BF16/F16 × b4/b8 × g32/64/128 × packed/separate × 4,096/8,193/32,769 routed rows. It includes empty expert runs, partial tiles, both canary regimes, a strided uint32 row map and int64 expert IDs. Every result equals guarded stock and the pinned oracle bit for bit; 48 first-use canaries pass. Actual installed layer-0 projections additionally cover 410/2,048/3,277/4,095/4,096/8,192 tokens (12 packed/separate cases).

Failure injection uses actual invalid Metal source, a changed output sign bit, and explicitly differing positive/negative zero. Each disables its key and subsequent calls use exact stock output without another canary. Metadata and policy tests cover all exclusion predicates and switch spellings. The full model off-switch run checks the unchanged output/state and absence of kernel activation.

The full-model large-prefix control reuses #110's 7,374-token cold prompt and 3,278-token warm suffix at chunks 2,048 and 4,096. It checks every logit and cache array against the pre-port, pinned-oracle receipts, plus cancellation and retry. Explicit hidden-state tests use an MTP-capable runtime with cold 4,096-token and warm 2,048-token-suffix prefills; an AR-only runtime correctly omits warm hidden arrays in the oracle. The AR timing tests compare complete state after one generated token. State reads occur after measured intervals.

Paired release measurements

M5 Max, 128 GiB; serial before/after processes, three repetitions per chunk and phase, identical installed artifacts and token IDs. Both default boundary splitting (tail interval 256) and full chunks (tail interval 0) are measured. The latter exercises the wide row regime and guarded partial tails. The release baseline uses unchanged committed production source with only the same measurement test module added. No other GPU workload or build ran during these measurements.

Cells show median ± sample standard deviation across all three repetitions. Repetition zero includes first-use initialization/canary work; this is not an isolated measurement of canary overhead. All individual times and means are in release-performance.json. TTFT starts at the request call and ends at its first token callback, after model loading.

Tail policy / chunk / phase	Before TTFT s	After TTFT s	Before prefill tok/s	After prefill tok/s
Default / 2048 / cold	2.959 ± 0.473	2.953 ± 0.480	1385.831 ± 183.340	1388.955 ± 186.819
Default / 2048 / warm	1.839 ± 0.075	1.900 ± 0.088	1115.426 ± 45.227	1079.363 ± 52.619
Default / 4096 / cold	3.690 ± 0.063	3.684 ± 0.029	1111.739 ± 19.009	1113.593 ± 8.843
Default / 4096 / warm	2.012 ± 0.013	1.989 ± 0.019	1019.529 ± 6.863	1031.159 ± 10.182
Full / 2048 / cold	2.411 ± 0.475	2.504 ± 0.428	1699.510 ± 259.818	1636.600 ± 239.080
Full / 2048 / warm	1.262 ± 0.063	1.334 ± 0.052	1623.785 ± 78.593	1536.347 ± 64.230
Full / 4096 / cold	3.100 ± 0.116	2.999 ± 0.136	1321.520 ± 49.755	1366.360 ± 58.429
Full / 4096 / warm	1.767 ± 0.057	1.634 ± 0.086	1159.717 ± 36.444	1254.103 ± 61.907

Peak physical footprint: default 84.724 → 84.669 GB; full chunks 88.579 → 87.952 GB. The results are mixed and do not establish an overall speedup. In particular, the default 256-token tail grid often keeps the small stock route; setting chunk size 4096 does not imply every body call has that width. The port fulfills the current-route obligation without assuming the upstream isolated-kernel speedup applies to this server benchmark.

Reproduction

Set DS4_QWEN38_ARTIFACTS to the installed artifact directory. Build the release test executable with cargo test --release --all-features --bin ds4-server issue112 --no-run, then copy the executable reported by Cargo outside target before running it. Run each ignored test in its own serial process under target/debug/test-supervisor MEMORY_MIB 60 120 --command TEST_EXECUTABLE FILTER --ignored --nocapture --test-threads=1:

issue112_all_kernel_variants and issue112_canary_failures: 4096 MiB.
issue112_installed_gate_up: 16384 MiB.
issue112_full_model_state: 114688 MiB, DS4_ISSUE110_DIR=OUTPUT_DIRECTORY, MTPLX_GDN_BOUNDARY_TAIL_INTERVAL=0.
issue112_hidden_prefill: 114688 MiB, DS4_ISSUE112_DIR=OUTPUT_DIRECTORY, MTPLX_GDN_BOUNDARY_TAIL_INTERVAL=0.
issue112_paired_prefill: 114688 MiB, DS4_ISSUE112_DIR=OUTPUT_DIRECTORY, optionally tail interval 0; DS4_ISSUE112_REPEATS defaults to 3. Also run a one-repetition full-chunk control with MTPLX_MOE_SORTED_GATHER_KERNEL=0.

Keep child stdout/stderr separate from the supervisor's resource output to avoid interleaved JSON writes. Under the isolated MLX 0.32.2 reference Python, run tools/issue112-reference.py PINNED_CHECKOUT NATIVE_VARIANT_OUTPUT, tools/issue112-installed-reference.py PINNED_CHECKOUT MODEL_DIR NATIVE_OUTPUT, and tools/issue112-model-reference.py PINNED_CHECKOUT MODEL_DIR NATIVE_DIR ORACLE_DIR (add --hidden for the hidden-state run; use matching tail policy). The full-model oracle uses the server's lazy-snapshot policy for cache fingerprints, matching the native session entry; a deep clone adds zero and would change negative-zero key bits during measurement itself.

For the before arm, build commit 8965e77 with only issue112_bench.rs and its #[cfg(test)] mod issue112_bench; registration added. Do not use an off-switch run as the only baseline: the preserved pre-port executable is the baseline. Both arms use the same test harness. Run timing arms without other builds or GPU work; correctness checks may run alongside CPU-only compilation.

Completion review

All 72 synthetic and 12 installed projection receipts match guarded stock and pinned MTPLX; compile/sign-bit/signed-zero failures pass sticky-fallback checks. All four large-prefix snapshots match #110's pre-port/oracle fixtures. The 24 release before/after state comparisons and four explicit hidden-state comparisons are exact. The direct pinned oracle matches all four representative release states and all four hidden-prefill states, including every cache bit. The off-switch matches all four baseline states with zero kernel activation.

DeepSeek's resident_multi_session_switching_preserves_each_kv_frontier and GLM's glm53_two_row_verifier_restores_recurrent_state both pass on the preserved before binary and the final after binary with installed models. Source review confirms no DS4/GLM model, sampling, cache, shader, or orchestration path changed. Their DS4 oracle behavior remains assigned to those families; this is the Qwen MTPLX route port, not a new DS4Server-only feature. Non-NAX predicates are tested; no physical M1–M4 run is claimed.

All required commit gates pass: formatting, warning-free Clippy, signed app bundle, and all-features tests (312 application + 8 tool tests; 234 explicitly ignored hardware/reference tests). Relevant ignored tests above were run separately under the resource supervisor. Machine-readable evidence, exact state fingerprints, gate logs, resource summaries and binary/source provenance are stored beside this file. No acceptance item is deferred to a stub.

9.4 KiB Raw Blame History Unescape Escape