Parent: #103. Related: #110 (owns the MLX 0.32.2 row-guard correctness fix; this issue owns the current-route port).
Why this exists
Oracle assignment (user): MTPLX pinned 9882703f3105363ddc37eca9f97aa09a1d387112 is the parity reference for Qwen 3.8 Flash Next. DS4 remains the reference for DeepSeek and GLM. For Qwen, the native Rust/Metal implementation must follow the current MTPLX execution route, not only its numerical results. The Flash-Next prefill MoE route changed upstream after our pinned reference, and DS4Server still runs the old route. Upstream states the new route is bit-identical to the guarded stock chain, so this is a route/performance parity obligation, not a confirmed numerical bug. It is a required port, not an optional optimization.
moe_sorted_gather.sort_rows(x, indices): argsort of the flattened expert ids, inv_order = argsort(order), row_map = order // top_k (uint32). It returns tokens [n_tokens, 1, K]without copying x[order // top_k].
Gate/up: moe_sorted_gather.swiglu_rows(...). When mtplx/kernels/moe_sorted_gather_nax.pyapplies(...), one NAX kernel reads token rows in place through row_map and writes silu(gate) * up directly. Otherwise it falls back to gather_qmm_rows (row copy + guarded stock gather_qmm) + split + nn.silu(gate) * up.
Down: moe_sorted_gather.switch_linear(down_proj, h, idx, sorted_indices=True), i.e. MLX's stock sorted kernel behind the #110 row guard. Upstream measured the row-map kernel slower for the already-sorted down projection (0.88x–0.94x), so down stays on the stock kernel.
_scatter_unsort(y, inv_order, ...).
Kernel contract (moe_sorted_gather_nax.py docstring and applies):
Threadgroups own 64 sorted rows of one expert (BM=BN=BK=64, 128 threads); per-expert tile/row prefix tables (_schedule); expert found by binary search. 32-bit row positions, so no 32,767-row bound.
SwiGLU epilogue: 32 gate rows + 32 matching up rows per threadgroup. Accumulators are rounded to the activation dtype, then MLX's Sigmoid/Multiply with the same per-op rounding as nn.silu(gate) * up. Supports the fused [gate | up] weight and separate gate/up weights.
Engages only when all hold: NAX GPU (nax_detect.nax_available()), affine mode, bf16/fp16 activations with matching scales/biases dtype, biases present, 4 or 8 bits, group 32/64/128, tokens shape [n,1,K], uint32 row_map, K and N multiples of 64, and routed rows >= min_rows() = 4,096 on MLX < 0.32.3 (16,384 on >= 0.32.3; we stay on 0.32.2). The default 2,048-token chunk (20,480 routed rows) qualifies.
First-use canary per compiled instantiation and per row regime (under / from 64 rows per expert): run the whole first call through kernel and stock chain, compare bit for bit (signed zero counts). On mismatch or compile failure, disable the kernel for that instantiation/regime for the process, with a printed reason and counter, and use the stock chain.
Off switch MTPLX_MOE_SORTED_GATHER_KERNEL=0.
Upstream history: aa68928d1e2e82d95bf4492fa0a7bd38e0bfdd5c (initial), 07261055, ae2746a5 (threshold under 0.32.3), 97cea671 (whole-first-call canary per regime), 205a5bbd, 724fcc8c (keep stock path when headers are unreadable).
Current DS4Server route
moe.rsswitched (src/engine/metal/qwen_mtplx/moe.rs:54-110): sort when indices.size >= 64, copy rows with indexing::take, stock gather for packed gate/up (or separate up and gate), split, mlp::silu/swiglu, stock gather for down, unsort with take. This matches the pre-aa68928d upstream chain.
Native implementation scope
Rust: row-map sort equivalent to sort_rows (same ordering, uint32 map, inverse order); dispatch with the exact applies predicates and the 4,096-row threshold; per-instantiation/regime canary with process-lifetime disable and reason/counter; MTPLX_MOE_SORTED_GATHER_KERNEL=0 off switch through the Qwen policy (no process-env mutation); guarded stock fallback via #110's guard.
Metal: a new .metal kernel for the row-map SwiGLU gate/up gather (packed and separate gate/up), bit-identical to the guarded stock chain. Upstream compiles its kernel at runtime from the installed MLX headers; DS4Server has no MLX install and must ship its own source. Record MLX MIT provenance (see metal/MLX-LM-LICENSE.txt, metal/MTPLX-*) or implement independently. Do not patch or replace mtplx-runtime-0.32.2.metallib.
Down projection remains the stock sorted kernel behind the #110 guard.
Qwen-only. M1–M4 (no NAX), decode/verify widths (< 64 routed rows) and ineligible shapes keep the existing route. No DS4-family MoE/kernel changes; runtime 0.32.2 and default chunk unchanged.
Dependencies
#110's guard must exist first or together: it protects down and every stock fallback of this route.
Acceptance
Kernel vs guarded stock chain bit-exact (including signed zero) for gate/up at default 2,048 tokens, 4,096 and larger chunks, partial chunks, both row regimes, packed and separate gate/up, the shipped Bare-Speed bits/group sizes, plus synthetic coverage of the other supported bits/groups and bf16/fp16.
Predicates: below-threshold, non-NAX and ineligible shapes take the old route. Canary failure injection disables the kernel and falls back. The off switch works.
Full-model cold/warm prefill logits/hidden/state equal to the pre-port native route and to pinned MTPLX on identical artifacts, at chunk 2,048 and 4,096.
Paired serial Bare-Speed M5 Max measurements (TTFT, prefill tok/s, memory peak, variance) before/after. Upstream reports 1.47x/1.31x/1.23x gate/up kernel speedup at 2,048/4,096/8,192-token chunks; record our numbers, do not assume them.
Reference control already run during review: upstream tests/test_moe_sorted_gather_nax.py::test_flash_next_moe_block_is_bit_identical_with_and_without_the_kernel (combine=1/0) passed with kernel/canary active on M5 Max / MLX 0.32.2.
DSv4/GLM outputs and state unchanged. Repository gates before every commit: cargo fmt --all -- --check, cargo clippy --all-targets --all-features -- -D warnings, make bundle, cargo test --all-features.
Implementation awaits the user's instruction.
Parent: #103. Related: #110 (owns the MLX 0.32.2 row-guard correctness fix; this issue owns the current-route port).
## Why this exists
Oracle assignment (user): **MTPLX pinned `9882703f3105363ddc37eca9f97aa09a1d387112` is the parity reference for Qwen 3.8 Flash Next. DS4 remains the reference for DeepSeek and GLM.** For Qwen, the native Rust/Metal implementation must follow the *current* MTPLX execution route, not only its numerical results. The Flash-Next prefill MoE route changed upstream after our pinned reference, and DS4Server still runs the old route. Upstream states the new route is bit-identical to the guarded stock chain, so this is a route/performance parity obligation, not a confirmed numerical bug. It is a required port, not an optional optimization.
## Current upstream route (MTPLX `9882703`)
`qwen4_exp` switch-GLU `__call__` → `_sorted` (mtplx/models/qwen4_exp.py, ~1545-1578), for >= 64 routed rows:
- `moe_sorted_gather.sort_rows(x, indices)`: argsort of the flattened expert ids, `inv_order = argsort(order)`, `row_map = order // top_k` (uint32). It returns `tokens [n_tokens, 1, K]` **without** copying `x[order // top_k]`.
- Gate/up: `moe_sorted_gather.swiglu_rows(...)`. When `mtplx/kernels/moe_sorted_gather_nax.py` `applies(...)`, one NAX kernel reads token rows in place through `row_map` and writes `silu(gate) * up` directly. Otherwise it falls back to `gather_qmm_rows` (row copy + guarded stock `gather_qmm`) + split + `nn.silu(gate) * up`.
- Down: `moe_sorted_gather.switch_linear(down_proj, h, idx, sorted_indices=True)`, i.e. MLX's stock sorted kernel **behind the #110 row guard**. Upstream measured the row-map kernel slower for the already-sorted down projection (0.88x–0.94x), so down stays on the stock kernel.
- `_scatter_unsort(y, inv_order, ...)`.
Kernel contract (`moe_sorted_gather_nax.py` docstring and `applies`):
- Threadgroups own 64 sorted rows of one expert (BM=BN=BK=64, 128 threads); per-expert tile/row prefix tables (`_schedule`); expert found by binary search. 32-bit row positions, so no 32,767-row bound.
- SwiGLU epilogue: 32 gate rows + 32 matching up rows per threadgroup. Accumulators are rounded to the activation dtype, then MLX's `Sigmoid`/`Multiply` with the same per-op rounding as `nn.silu(gate) * up`. Supports the fused `[gate | up]` weight and separate gate/up weights.
- Engages only when all hold: NAX GPU (`nax_detect.nax_available()`), affine mode, bf16/fp16 activations with matching scales/biases dtype, biases present, 4 or 8 bits, group 32/64/128, `tokens` shape `[n,1,K]`, uint32 `row_map`, K and N multiples of 64, and **routed rows >= `min_rows()` = 4,096 on MLX < 0.32.3** (16,384 on >= 0.32.3; we stay on 0.32.2). The default 2,048-token chunk (20,480 routed rows) qualifies.
- First-use canary per compiled instantiation and per row regime (under / from 64 rows per expert): run the whole first call through kernel and stock chain, compare bit for bit (signed zero counts). On mismatch or compile failure, disable the kernel for that instantiation/regime for the process, with a printed reason and counter, and use the stock chain.
- Off switch `MTPLX_MOE_SORTED_GATHER_KERNEL=0`.
- Upstream history: `aa68928d1e2e82d95bf4492fa0a7bd38e0bfdd5c` (initial), `07261055`, `ae2746a5` (threshold under 0.32.3), `97cea671` (whole-first-call canary per regime), `205a5bbd`, `724fcc8c` (keep stock path when headers are unreadable).
## Current DS4Server route
`moe.rs` `switched` (src/engine/metal/qwen_mtplx/moe.rs:54-110): sort when `indices.size >= 64`, copy rows with `indexing::take`, stock `gather` for packed gate/up (or separate up and gate), split, `mlp::silu`/`swiglu`, stock `gather` for down, unsort with `take`. This matches the pre-`aa68928d` upstream chain.
## Native implementation scope
- Rust: row-map sort equivalent to `sort_rows` (same ordering, uint32 map, inverse order); dispatch with the exact `applies` predicates and the 4,096-row threshold; per-instantiation/regime canary with process-lifetime disable and reason/counter; `MTPLX_MOE_SORTED_GATHER_KERNEL=0` off switch through the Qwen policy (no process-env mutation); guarded stock fallback via #110's guard.
- Metal: a new `.metal` kernel for the row-map SwiGLU gate/up gather (packed and separate gate/up), bit-identical to the guarded stock chain. Upstream compiles its kernel at runtime from the installed MLX headers; DS4Server has no MLX install and must ship its own source. Record MLX MIT provenance (see `metal/MLX-LM-LICENSE.txt`, `metal/MTPLX-*`) or implement independently. Do not patch or replace `mtplx-runtime-0.32.2.metallib`.
- Down projection remains the stock sorted kernel behind the #110 guard.
- Qwen-only. M1–M4 (no NAX), decode/verify widths (< 64 routed rows) and ineligible shapes keep the existing route. No DS4-family MoE/kernel changes; runtime 0.32.2 and default chunk unchanged.
## Dependencies
#110's guard must exist first or together: it protects down and every stock fallback of this route.
## Acceptance
- Kernel vs guarded stock chain bit-exact (including signed zero) for gate/up at default 2,048 tokens, 4,096 and larger chunks, partial chunks, both row regimes, packed and separate gate/up, the shipped Bare-Speed bits/group sizes, plus synthetic coverage of the other supported bits/groups and bf16/fp16.
- Predicates: below-threshold, non-NAX and ineligible shapes take the old route. Canary failure injection disables the kernel and falls back. The off switch works.
- Full-model cold/warm prefill logits/hidden/state equal to the pre-port native route and to pinned MTPLX on identical artifacts, at chunk 2,048 and 4,096.
- Paired serial Bare-Speed M5 Max measurements (TTFT, prefill tok/s, memory peak, variance) before/after. Upstream reports 1.47x/1.31x/1.23x gate/up kernel speedup at 2,048/4,096/8,192-token chunks; record our numbers, do not assume them.
- Reference control already run during review: upstream `tests/test_moe_sorted_gather_nax.py::test_flash_next_moe_block_is_bit_identical_with_and_without_the_kernel` (combine=1/0) passed with kernel/canary active on M5 Max / MLX 0.32.2.
- DSv4/GLM outputs and state unchanged. Repository gates before every commit: `cargo fmt --all -- --check`, `cargo clippy --all-targets --all-features -- -D warnings`, `make bundle`, `cargo test --all-features`.
Implementation awaits the user's instruction.
Qwen now follows pinned MTPLX's row-map gate/up route: original argsort/inverse ordering and uint32 row map, GPU expert row/tile prefix scheduling, 64×64×64 NAX kernel, and packed/separate SwiGLU with exact activation rounding. Down remains guarded stock. The 4096-row threshold, all BF16/F16/bits/group/layout predicates, off switch and GPU-family fallback are wired through Qwen policy. Whole-first-call bitwise canaries cover both row regimes, retain process-lifetime failure state and diagnostics, and reproduce the plain-gather/stock fallback sequence. GPU shader sources carry MLX/MTPLX MIT provenance; orchestration remains Rust and the runtime metallib is unchanged.
Isolated verification on the installed Bare-Speed artifact / M5 Max / MLX 0.32.2:
72 synthetic cases and 12 actual-weight cases are bit-exact against guarded stock and pinned MTPLX, including BF16/F16, b4/b8, g32/64/128, packed/separate, both regimes, strided maps, int64 IDs, empty experts, partial tiles and up to 8192-token projections.
Actual compile failure, changed sign bit, and +0/-0 mismatch injections prove sticky disable and exact fallback; off-switch full-model states match baseline with no kernel activation.
Large cold/warm sessions at chunks 2048/4096 match the pre-port/oracle complete state, including cancellation/retry. All 24 release before/after snapshots and four explicit hidden-state comparisons match exactly. Direct pinned MTPLX comparisons match all logits, hidden arrays and cache bits for both chunk sizes and cold/warm phases. Snapshot fingerprinting uses the matching server lazy-snapshot policy.
Installed DeepSeek session-switching and GLM verifier/recurrent-state controls pass on both before and after executables. Review confirms the production changes are Qwen-only; this is MTPLX route parity, not a new DS4Server-only feature.
All four commit gates pass: fmt, Clippy with -D warnings, make bundle, and all-features tests (312 application + 8 tool tests). Relevant ignored Metal/reference checks ran separately and serially.
Paired release measurements include first-use work, three repetitions, TTFT, prefill throughput, peak footprint and variance under default and unsplit-tail policies. Results are mixed; no overall speedup is claimed. Full-chunk peak footprint was 88.579 → 87.952 GB. Raw measurements, exact state/kernel receipts, provenance, gates, limitations and reproduction steps are in docs/issue112-row-map.
The complete issue scope has been implemented and reviewed; no stub or deferred acceptance item remains.
Implemented and pushed in [7b894fbfeed81d2944b9cf479e60638ef04987c7](https://git.rfc1437.de/hugo/DS4Server/commit/7b894fbfeed81d2944b9cf479e60638ef04987c7).
Qwen now follows pinned MTPLX's row-map gate/up route: original argsort/inverse ordering and uint32 row map, GPU expert row/tile prefix scheduling, 64×64×64 NAX kernel, and packed/separate SwiGLU with exact activation rounding. Down remains guarded stock. The 4096-row threshold, all BF16/F16/bits/group/layout predicates, off switch and GPU-family fallback are wired through Qwen policy. Whole-first-call bitwise canaries cover both row regimes, retain process-lifetime failure state and diagnostics, and reproduce the plain-gather/stock fallback sequence. GPU shader sources carry MLX/MTPLX MIT provenance; orchestration remains Rust and the runtime metallib is unchanged.
Isolated verification on the installed Bare-Speed artifact / M5 Max / MLX 0.32.2:
- 72 synthetic cases and 12 actual-weight cases are bit-exact against guarded stock and pinned MTPLX, including BF16/F16, b4/b8, g32/64/128, packed/separate, both regimes, strided maps, int64 IDs, empty experts, partial tiles and up to 8192-token projections.
- Actual compile failure, changed sign bit, and +0/-0 mismatch injections prove sticky disable and exact fallback; off-switch full-model states match baseline with no kernel activation.
- Large cold/warm sessions at chunks 2048/4096 match the pre-port/oracle complete state, including cancellation/retry. All 24 release before/after snapshots and four explicit hidden-state comparisons match exactly. Direct pinned MTPLX comparisons match all logits, hidden arrays and cache bits for both chunk sizes and cold/warm phases. Snapshot fingerprinting uses the matching server lazy-snapshot policy.
- Installed DeepSeek session-switching and GLM verifier/recurrent-state controls pass on both before and after executables. Review confirms the production changes are Qwen-only; this is MTPLX route parity, not a new DS4Server-only feature.
- All four commit gates pass: fmt, Clippy with -D warnings, make bundle, and all-features tests (312 application + 8 tool tests). Relevant ignored Metal/reference checks ran separately and serially.
Paired release measurements include first-use work, three repetitions, TTFT, prefill throughput, peak footprint and variance under default and unsplit-tail policies. Results are mixed; no overall speedup is claimed. Full-chunk peak footprint was 88.579 → 87.952 GB. Raw measurements, exact state/kernel receipts, provenance, gates, limitations and reproduction steps are in [docs/issue112-row-map](https://git.rfc1437.de/hugo/DS4Server/src/commit/7b894fbfeed81d2944b9cf479e60638ef04987c7/docs/issue112-row-map).
The complete issue scope has been implemented and reviewed; no stub or deferred acceptance item remains.
Blocking a user prevents them from interacting with repositories, such as opening or commenting on pull requests or issues. Learn more about blocking a user.
Parent: #103. Related: #110 (owns the MLX 0.32.2 row-guard correctness fix; this issue owns the current-route port).
Why this exists
Oracle assignment (user): MTPLX pinned
9882703f3105363ddc37eca9f97aa09a1d387112is the parity reference for Qwen 3.8 Flash Next. DS4 remains the reference for DeepSeek and GLM. For Qwen, the native Rust/Metal implementation must follow the current MTPLX execution route, not only its numerical results. The Flash-Next prefill MoE route changed upstream after our pinned reference, and DS4Server still runs the old route. Upstream states the new route is bit-identical to the guarded stock chain, so this is a route/performance parity obligation, not a confirmed numerical bug. It is a required port, not an optional optimization.Current upstream route (MTPLX
9882703)qwen4_expswitch-GLU__call__→_sorted(mtplx/models/qwen4_exp.py, ~1545-1578), for >= 64 routed rows:moe_sorted_gather.sort_rows(x, indices): argsort of the flattened expert ids,inv_order = argsort(order),row_map = order // top_k(uint32). It returnstokens [n_tokens, 1, K]without copyingx[order // top_k].moe_sorted_gather.swiglu_rows(...). Whenmtplx/kernels/moe_sorted_gather_nax.pyapplies(...), one NAX kernel reads token rows in place throughrow_mapand writessilu(gate) * updirectly. Otherwise it falls back togather_qmm_rows(row copy + guarded stockgather_qmm) + split +nn.silu(gate) * up.moe_sorted_gather.switch_linear(down_proj, h, idx, sorted_indices=True), i.e. MLX's stock sorted kernel behind the #110 row guard. Upstream measured the row-map kernel slower for the already-sorted down projection (0.88x–0.94x), so down stays on the stock kernel._scatter_unsort(y, inv_order, ...).Kernel contract (
moe_sorted_gather_nax.pydocstring andapplies):_schedule); expert found by binary search. 32-bit row positions, so no 32,767-row bound.Sigmoid/Multiplywith the same per-op rounding asnn.silu(gate) * up. Supports the fused[gate | up]weight and separate gate/up weights.nax_detect.nax_available()), affine mode, bf16/fp16 activations with matching scales/biases dtype, biases present, 4 or 8 bits, group 32/64/128,tokensshape[n,1,K], uint32row_map, K and N multiples of 64, and routed rows >=min_rows()= 4,096 on MLX < 0.32.3 (16,384 on >= 0.32.3; we stay on 0.32.2). The default 2,048-token chunk (20,480 routed rows) qualifies.MTPLX_MOE_SORTED_GATHER_KERNEL=0.aa68928d1e2e82d95bf4492fa0a7bd38e0bfdd5c(initial),07261055,ae2746a5(threshold under 0.32.3),97cea671(whole-first-call canary per regime),205a5bbd,724fcc8c(keep stock path when headers are unreadable).Current DS4Server route
moe.rsswitched(src/engine/metal/qwen_mtplx/moe.rs:54-110): sort whenindices.size >= 64, copy rows withindexing::take, stockgatherfor packed gate/up (or separate up and gate), split,mlp::silu/swiglu, stockgatherfor down, unsort withtake. This matches the pre-aa68928dupstream chain.Native implementation scope
sort_rows(same ordering, uint32 map, inverse order); dispatch with the exactappliespredicates and the 4,096-row threshold; per-instantiation/regime canary with process-lifetime disable and reason/counter;MTPLX_MOE_SORTED_GATHER_KERNEL=0off switch through the Qwen policy (no process-env mutation); guarded stock fallback via #110's guard..metalkernel for the row-map SwiGLU gate/up gather (packed and separate gate/up), bit-identical to the guarded stock chain. Upstream compiles its kernel at runtime from the installed MLX headers; DS4Server has no MLX install and must ship its own source. Record MLX MIT provenance (seemetal/MLX-LM-LICENSE.txt,metal/MTPLX-*) or implement independently. Do not patch or replacemtplx-runtime-0.32.2.metallib.Dependencies
#110's guard must exist first or together: it protects down and every stock fallback of this route.
Acceptance
tests/test_moe_sorted_gather_nax.py::test_flash_next_moe_block_is_bit_identical_with_and_without_the_kernel(combine=1/0) passed with kernel/canary active on M5 Max / MLX 0.32.2.cargo fmt --all -- --check,cargo clippy --all-targets --all-features -- -D warnings,make bundle,cargo test --all-features.Implementation awaits the user's instruction.
Implemented and pushed in 7b894fbfeed81d2944b9cf479e60638ef04987c7.
Qwen now follows pinned MTPLX's row-map gate/up route: original argsort/inverse ordering and uint32 row map, GPU expert row/tile prefix scheduling, 64×64×64 NAX kernel, and packed/separate SwiGLU with exact activation rounding. Down remains guarded stock. The 4096-row threshold, all BF16/F16/bits/group/layout predicates, off switch and GPU-family fallback are wired through Qwen policy. Whole-first-call bitwise canaries cover both row regimes, retain process-lifetime failure state and diagnostics, and reproduce the plain-gather/stock fallback sequence. GPU shader sources carry MLX/MTPLX MIT provenance; orchestration remains Rust and the runtime metallib is unchanged.
Isolated verification on the installed Bare-Speed artifact / M5 Max / MLX 0.32.2:
Paired release measurements include first-use work, three repetitions, TTFT, prefill throughput, peak footprint and variance under default and unsplit-tail policies. Results are mixed; no overall speedup is claimed. Full-chunk peak footprint was 88.579 → 87.952 GB. Raw measurements, exact state/kernel receipts, provenance, gates, limitations and reproduction steps are in docs/issue112-row-map.
The complete issue scope has been implemented and reviewed; no stub or deferred acceptance item remains.