Prevent uninitialized expert outputs for large unaligned M5 prefill chunks #110

Closed
opened 2026-10-03 09:20:33 +00:00 by hugo · 6 comments
Owner

Publication status: analysis commit dcc4c516f0e030566c67472f8f1be0baa21ab629 is durable locally; its push to origin/main awaits explicit approval after automatic review refused the default-branch push. Until published, linked repository evidence files are available locally under /Users/gb/Projects/DS4Server/docs/issue103-analysis/ and the committed test/tool paths, not yet at the remote file URLs.

Review status (updated 2026-10-03)

Confirmed defect or numerical parity difference; implementation still awaits verification. This description incorporates the independent review and the user's explicit model-specific oracle clarification. It supersedes the original issue scope and conflicting conclusions in the initial evidence report. The second review is addressed by explicitly separating guard correctness from the current optimized-route porting candidate. Original raw receipts/diagnostics remain unchanged; no production fix or new product/performance result is claimed; the additional reference-only test results are recorded below. The report and AGENTS.md have not been rewritten in this issue-only update.

Model-specific parity boundary

The user explicitly clarified the oracle assignment after independent review:

  • Qwen 3.8 Flash Next: behavioral/numerical parity with the pinned MTPLX 2.12.2 revision 9882703f3105363ddc37eca9f97aa09a1d387112, including token processing, RNG/sampling, context accounting, inference and cache/session state.
  • DSv4 Flash and GLM 5.3 Flash: preserve DS4 behavior. A Qwen fix must not change their token IDs, sampler/RNG sequences, numerical inference, context accounting or cache behavior.
  • Keep changes in the Qwen family/phase/dispatch where applicable. Before changing any shared helper, trace every caller; preserve DS4-backed callers exactly and provide DS4 regression controls. A helper shared by Qwen and GLM must not silently impose MTPLX semantics on GLM. Do not globally change the DS4 sampler, shared RNG, Unicode tables, sigmoid/fast-math rules, or non-Qwen checkpoint formats for this ticket.

This is the user's model-specific clarification of the blanket DS4-oracle wording in AGENTS.md, not permission to weaken DS4 parity. Production remains native Rust/Metal with the existing Objective-C bridge; MTPLX Python is an isolated oracle only.

Broken behavior and evidence

Confirmed correctness defect: current native sorted expert gathers use the MLX 0.32.2 NAX shader without the pad/slice guard. Caller: prefill chunk → MoE FeedForward::switched → linear::gather_qmm_bf16 → gather_evaluate → gather_bf16_layout → affine_gather_qmm_rhs_nax_* from the pinned metallib. Predicates: affected M5/NAX, m=1, sorted RHS/no left indices, routes>=16, routes/experts>=4 and applicable tile/function-constant alignment. Above 32767 routed rows with a non-tile-aligned count, leading outputs remain unwritten.

The sentinel receipt remains unchanged: BF16 ones, 4-bit group 32 unit weights, 64 experts, K=N=256,ten routes/token; expected output 256(BF16 0x4380), initialized NaN 0x7fc1. tokens/routes/bad rows: 3276/32760/0; 3277/32770/32; 3404/34040/1280; 4095/40950/8192; 4096/40960/0. Default 2048 means 20480 routed rows, below this correctness-defect threshold. Larger supported user chunks/partial chunks are affected. Full-model large-prompt output and both actual projection geometries are still unmeasured.

Current upstream route is broader than this guard: qwen4_exp._sorted uses sort_rows to produce a row map; gate/up uses swiglu_rows, which selects the fused row-map NAX SwiGLU kernel when eligible; down uses guarded stock switch_linear. The kernel threshold is 4096 routed rows on pinned MLX 0.32.2, so default 2048-token/20480-route prefill qualifies when dtype/quantization/device predicates and first-use canary pass. The guard protects Down and any stock Gate/Up fallback. DS4Server still uses copied rows and gather→split→SiLU/SwiGLU→gather.

That missing fused route is an explicit current-route/performance porting candidate, not by itself a demonstrated numerical error: MTPLX's own off switch, canary failure and ineligible predicates select the bit-equivalent guarded stock chain. No DS4Server same-machine performance-contract violation from this missing optimization has been measured. This issue owns the uninitialized-output bug and guarded stock fallback, not the whole optimized route.

Native implementation scope

Own the guard only: native host pad-to-64-row-tile/slice at applicable Qwen stock sorted-gather calls above 32,767 routed rows, zero padded activations, repeat last expert to preserve sorted order, pad lhs_indices with arange when present, then slice to original rows. Mirror b36b106/moe_sorted_gather.py and exact affected kernel/runtime predicates; keep allocations/index shapes/lifetimes correct. Include Gate/Up and Down on the current native stock chain. If a separately implemented optimized Gate/Up path is installed later, preserve this guard on Down and all stock fallbacks rather than dropping it as obsolete.

Scope boundary: do not fold a new row-map/fused SwiGLU kernel into this correctness fix. Track its separate Qwen porting/optimization candidate under #103, Stage J. Before any later implementation it needs its own focused scope: row-map sort; native .metal source/Rust dispatch; affine BF16/FP16, 4/8-bit,groups32/64/128 and K/N64 predicates; 4096-route threshold on 0.32.2; first-use bit-level canary per instantiation/regime; guarded stock fallback and MTPLX_MOE_SORTED_GATHER_KERNEL=0 off switch; unchanged guarded Down; verified MLX MIT provenance or independent source. No linked implementation issue is created or authorized by this description update. A missing numerical/performance acceptance result cannot be marked satisfied by fixing the guard alone.

The 0.32.2 shader stays pinned: no metallib patch/replacement, runtime upgrade, default chunk increase or artifact substitution. Protect M1–M4 and unaffected decode/verify dispatches. No DS4-family MoE/kernel behavior changes. No claim that the absence of an optional bit-identical optimization alone establishes a new bug; performance classification requires matched measurements.

Acceptance

Convert sentinel controls into exact every-row correctness, preserving 3276/3277/3404/4095/4096 token cases and their row bounds. Cover Gate/Up and Down actual shapes/group/bit layouts on the current native stock path, sorted padding/index/activation rules, output shape/order, affected dispatch gating, aligned/below-threshold controls, warm/partial chunks, allocation refusal/cancellation cleanup and valid retry. Report unavailable hardware explicitly rather than extrapolating.

Require native guarded Gate/Up and Down outputs to match pinned latest MTPLX on identical inputs/weights, including its supported guarded fallback. Also compare the latest eligible optimized Gate/Up oracle with its guarded stock chain bit for bit at default 2048 tokens/20480 routes and larger chunks; canary engagement must be recorded. That upstream optimized/fallback equivalence is an oracle control, not proof that DS4Server implements the optimized route or reaches the performance target. Preserve full-model cold/warm native/oracle parity at supported chunk 4096 and default 2048; do not introduce the unrelated optimized kernel as a hidden condition of this guard's completion.

Two official upstream test variants were executed during the subsequent review: test_flash_next_moe_block_is_bit_identical_with_and_without_the_kernel, combine=1/0, both passed with the NAX canary/kernel active. They use synthetic 64-expert BF16/group 32/4-bit weights, hidden 256/inter 64, 2100 tokens, top_k 10 at pinned MTPLX 9882703/MLX 0.32.2 on M5 Max. These are reference-only equivalence receipts, not full Bare-Speed/native/performance results. Reproduce from the immutable test below using the recorded reference environment; no need for this chat's scratch log.

The separately scoped optimized-route candidate must, if later authorized, pass native-vs-guarded-chain exactness, predicates/canary/fallback/off-switch, default 2048 and larger chunks, and matched TTFT/prefill/decode/memory measurements before claiming current-route/performance coverage. #110 cannot close those separate obligations. Keep runtime/assets/defaults and DSv4/GLM outputs/state unchanged.

Severity

High under supported larger user chunk settings: uninitialized leading expert rows are deterministically reproduced. Default 2048 is below the trigger; full-model large-prompt response remains unmeasured.

Current route and reference-only equivalence sources

# In the separate pinned MTPLX 9882703 checkout/reference environment (MLX 0.32.2):
python -m pytest -o addopts= -q -s tests/test_moe_sorted_gather_nax.py::test_flash_next_moe_block_is_bit_identical_with_and_without_the_kernel
# Verify both parameter variants pass AND the output reports kernel/canary on.

Shared evidence and reproduction

Analysis parent: #103. Affected DS4Server revision: cea8fc6f7c7aa7aeb97ebc303594ae4088ef6983. Compared MTPLX old e652d55e2652137a4abcf1312357abbf3eb9d692 (2.10.2) with refreshed 9882703f3105363ddc37eca9f97aa09a1d387112 (2.12.2), on MLX exactly 0.32.2. Same Bare-Speed artifact revision 74559cdf34fbfc0b593de72d17e93f37fd4f9ea7; unchanged hashes, compiler/OS/hardware and effective options in environment.json. Hardware: M5 Max 128 GiB/macOS 26.6.2. Synthetic model-free cases use explicit shapes/dtypes in the diagnostic, no template/image/prompt or context setting; model case uses raw IDs, seed0 and restored identical prompt snapshots.

Durable report and instructions, native test, oracle probe. The test module is cfg(test), ignored by ordinary runs, and initially asserts/prints the affected behavior; passing this diagnostic means reproduction, not a fix. Convert only this issue's owned cases into correct-contract regressions during implementation. Shared vision test has separate case ownership listed in the report.

cargo test --all-features --bin ds4-server issue103_sorted_gather_receipts --no-run
cargo build --bin test-supervisor
export DS4_QWEN38_ARTIFACTS='/path/to/unchanged/qwen3.8-flash-next'
export DS4_ISSUE103_ORACLE="$PWD/docs/issue103-analysis/latest-oracle.jsonl"
# Resolve the emitted test executable; do not pass cargo to the watchdog.
target/debug/test-supervisor 4096 60 120 --command /path/to/ds4_server-test-binary issue103_sorted_gather_receipts --ignored --nocapture --test-threads=1

Immutable upstream change and latest test. Padding/predicate reference: moe_sorted_gather.py. Extra readbacks/synchronization and external memory samples make these correctness receipts unsuitable for speed claims.

Validation and review boundary

  • Implement production behavior in Rust with existing Metal/Objective-C bridge only; Python remains an isolated oracle.
  • Run cargo fmt --all -- --check, cargo clippy --all-targets --all-features -- -D warnings, make bundle, cargo test --all-features before every future commit, plus the owned native/oracle and product regressions above.
  • Implementation awaits the user's independent-agent review and subsequent instruction. Creating this issue or closing #103 does not authorize a fix. Review the corrected model-specific oracle, active-route applicability and state ownership first; preserve DS4-backed families.
  • No arbitrary dependencies on sibling issues: coordinate shared source changes where needed; each issue owns its concrete failure and acceptance.
Publication status: analysis commit `dcc4c516f0e030566c67472f8f1be0baa21ab629` is durable locally; its push to origin/main awaits explicit approval after automatic review refused the default-branch push. Until published, linked repository evidence files are available locally under `/Users/gb/Projects/DS4Server/docs/issue103-analysis/` and the committed test/tool paths, not yet at the remote file URLs. ## Review status (updated 2026-10-03) Confirmed defect or numerical parity difference; implementation still awaits verification. This description incorporates the [independent review](https://git.rfc1437.de/hugo/DS4Server/issues/110#issuecomment-3073) and the user's explicit model-specific oracle clarification. It supersedes the original issue scope and conflicting conclusions in the initial evidence report. The [second review](https://git.rfc1437.de/hugo/DS4Server/issues/110#issuecomment-3099) is addressed by explicitly separating guard correctness from the current optimized-route porting candidate. Original raw receipts/diagnostics remain unchanged; no production fix or new product/performance result is claimed; the additional reference-only test results are recorded below. The report and AGENTS.md have not been rewritten in this issue-only update. ## Model-specific parity boundary The user explicitly clarified the oracle assignment after independent review: - Qwen 3.8 Flash Next: behavioral/numerical parity with the pinned MTPLX 2.12.2 revision `9882703f3105363ddc37eca9f97aa09a1d387112`, including token processing, RNG/sampling, context accounting, inference and cache/session state. - DSv4 Flash and GLM 5.3 Flash: preserve DS4 behavior. A Qwen fix must not change their token IDs, sampler/RNG sequences, numerical inference, context accounting or cache behavior. - Keep changes in the Qwen family/phase/dispatch where applicable. Before changing any shared helper, trace every caller; preserve DS4-backed callers exactly and provide DS4 regression controls. A helper shared by Qwen and GLM must not silently impose MTPLX semantics on GLM. Do not globally change the DS4 sampler, shared RNG, Unicode tables, sigmoid/fast-math rules, or non-Qwen checkpoint formats for this ticket. This is the user's model-specific clarification of the blanket DS4-oracle wording in AGENTS.md, not permission to weaken DS4 parity. Production remains native Rust/Metal with the existing Objective-C bridge; MTPLX Python is an isolated oracle only. ## Broken behavior and evidence Confirmed correctness defect: current native sorted expert gathers use the MLX 0.32.2 NAX shader without the pad/slice guard. Caller: prefill chunk → MoE FeedForward::switched → linear::gather_qmm_bf16 → gather_evaluate → gather_bf16_layout → affine_gather_qmm_rhs_nax_* from the pinned metallib. Predicates: affected M5/NAX, m=1, sorted RHS/no left indices, routes>=16, routes/experts>=4 and applicable tile/function-constant alignment. Above 32767 routed rows with a non-tile-aligned count, leading outputs remain unwritten. The sentinel receipt remains unchanged: BF16 ones, 4-bit group 32 unit weights, 64 experts, K=N=256,ten routes/token; expected output 256(BF16 0x4380), initialized NaN 0x7fc1. tokens/routes/bad rows: 3276/32760/0; 3277/32770/32; 3404/34040/1280; 4095/40950/8192; 4096/40960/0. Default 2048 means 20480 routed rows, below this correctness-defect threshold. Larger supported user chunks/partial chunks are affected. Full-model large-prompt output and both actual projection geometries are still unmeasured. Current upstream route is broader than this guard: qwen4_exp._sorted uses sort_rows to produce a row map; gate/up uses swiglu_rows, which selects the fused row-map NAX SwiGLU kernel when eligible; down uses guarded stock switch_linear. The kernel threshold is 4096 routed rows on pinned MLX 0.32.2, so default 2048-token/20480-route prefill qualifies when dtype/quantization/device predicates and first-use canary pass. The guard protects Down and any stock Gate/Up fallback. DS4Server still uses copied rows and gather→split→SiLU/SwiGLU→gather. That missing fused route is an explicit current-route/performance porting candidate, not by itself a demonstrated numerical error: MTPLX's own off switch, canary failure and ineligible predicates select the bit-equivalent guarded stock chain. No DS4Server same-machine performance-contract violation from this missing optimization has been measured. This issue owns the uninitialized-output bug and guarded stock fallback, not the whole optimized route. ## Native implementation scope Own the guard only: native host pad-to-64-row-tile/slice at applicable Qwen stock sorted-gather calls above 32,767 routed rows, zero padded activations, repeat last expert to preserve sorted order, pad lhs_indices with arange when present, then slice to original rows. Mirror b36b106/moe_sorted_gather.py and exact affected kernel/runtime predicates; keep allocations/index shapes/lifetimes correct. Include Gate/Up and Down on the current native stock chain. If a separately implemented optimized Gate/Up path is installed later, preserve this guard on Down and all stock fallbacks rather than dropping it as obsolete. Scope boundary: do not fold a new row-map/fused SwiGLU kernel into this correctness fix. Track its separate Qwen porting/optimization candidate under [#103, Stage J](https://git.rfc1437.de/hugo/DS4Server/issues/103#issuecomment-3100). Before any later implementation it needs its own focused scope: row-map sort; native .metal source/Rust dispatch; affine BF16/FP16, 4/8-bit,groups32/64/128 and K/N64 predicates; 4096-route threshold on 0.32.2; first-use bit-level canary per instantiation/regime; guarded stock fallback and MTPLX_MOE_SORTED_GATHER_KERNEL=0 off switch; unchanged guarded Down; verified MLX MIT provenance or independent source. No linked implementation issue is created or authorized by this description update. A missing numerical/performance acceptance result cannot be marked satisfied by fixing the guard alone. The 0.32.2 shader stays pinned: no metallib patch/replacement, runtime upgrade, default chunk increase or artifact substitution. Protect M1–M4 and unaffected decode/verify dispatches. No DS4-family MoE/kernel behavior changes. No claim that the absence of an optional bit-identical optimization alone establishes a new bug; performance classification requires matched measurements. ## Acceptance Convert sentinel controls into exact every-row correctness, preserving 3276/3277/3404/4095/4096 token cases and their row bounds. Cover Gate/Up and Down actual shapes/group/bit layouts on the current native stock path, sorted padding/index/activation rules, output shape/order, affected dispatch gating, aligned/below-threshold controls, warm/partial chunks, allocation refusal/cancellation cleanup and valid retry. Report unavailable hardware explicitly rather than extrapolating. Require native guarded Gate/Up and Down outputs to match pinned latest MTPLX on identical inputs/weights, including its supported guarded fallback. Also compare the latest eligible optimized Gate/Up oracle with its guarded stock chain bit for bit at default 2048 tokens/20480 routes and larger chunks; canary engagement must be recorded. That upstream optimized/fallback equivalence is an oracle control, not proof that DS4Server implements the optimized route or reaches the performance target. Preserve full-model cold/warm native/oracle parity at supported chunk 4096 and default 2048; do not introduce the unrelated optimized kernel as a hidden condition of this guard's completion. Two official upstream test variants were executed during the subsequent review: test_flash_next_moe_block_is_bit_identical_with_and_without_the_kernel, combine=1/0, both passed with the NAX canary/kernel active. They use synthetic 64-expert BF16/group 32/4-bit weights, hidden 256/inter 64, 2100 tokens, top_k 10 at pinned MTPLX 9882703/MLX 0.32.2 on M5 Max. These are reference-only equivalence receipts, not full Bare-Speed/native/performance results. Reproduce from the immutable test below using the recorded reference environment; no need for this chat's scratch log. The separately scoped optimized-route candidate must, if later authorized, pass native-vs-guarded-chain exactness, predicates/canary/fallback/off-switch, default 2048 and larger chunks, and matched TTFT/prefill/decode/memory measurements before claiming current-route/performance coverage. #110 cannot close those separate obligations. Keep runtime/assets/defaults and DSv4/GLM outputs/state unchanged. ## Severity High under supported larger user chunk settings: uninitialized leading expert rows are deterministically reproduced. Default 2048 is below the trigger; full-model large-prompt response remains unmeasured. ## Current route and reference-only equivalence sources - [Current Flash-Next model route](https://github.com/youssofal/mtplx/blob/9882703f3105363ddc37eca9f97aa09a1d387112/mtplx/models/qwen4_exp.py), [sort/guard/fallback entry points](https://github.com/youssofal/mtplx/blob/9882703f3105363ddc37eca9f97aa09a1d387112/mtplx/moe_sorted_gather.py). - [Fused kernel predicates, threshold and canary](https://github.com/youssofal/mtplx/blob/9882703f3105363ddc37eca9f97aa09a1d387112/mtplx/kernels/moe_sorted_gather_nax.py), [initial port aa68928d](https://github.com/youssofal/mtplx/commit/aa68928d1e2e82d95bf4492fa0a7bd38e0bfdd5c). - [Official block-equivalence test](https://github.com/youssofal/mtplx/blob/9882703f3105363ddc37eca9f97aa09a1d387112/tests/test_moe_sorted_gather_nax.py). ```sh # In the separate pinned MTPLX 9882703 checkout/reference environment (MLX 0.32.2): python -m pytest -o addopts= -q -s tests/test_moe_sorted_gather_nax.py::test_flash_next_moe_block_is_bit_identical_with_and_without_the_kernel # Verify both parameter variants pass AND the output reports kernel/canary on. ``` ## Shared evidence and reproduction Analysis parent: [#103](https://git.rfc1437.de/hugo/DS4Server/issues/103). Affected DS4Server revision: `cea8fc6f7c7aa7aeb97ebc303594ae4088ef6983`. Compared MTPLX old `e652d55e2652137a4abcf1312357abbf3eb9d692` (2.10.2) with refreshed `9882703f3105363ddc37eca9f97aa09a1d387112` (2.12.2), on MLX exactly 0.32.2. Same Bare-Speed artifact revision `74559cdf34fbfc0b593de72d17e93f37fd4f9ea7`; unchanged hashes, compiler/OS/hardware and effective options in [environment.json](https://git.rfc1437.de/hugo/DS4Server/src/commit/dcc4c516f0e030566c67472f8f1be0baa21ab629/docs/issue103-analysis/environment.json). Hardware: M5 Max 128 GiB/macOS 26.6.2. Synthetic model-free cases use explicit shapes/dtypes in the diagnostic, no template/image/prompt or context setting; model case uses raw IDs, seed0 and restored identical prompt snapshots. Durable [report and instructions](https://git.rfc1437.de/hugo/DS4Server/src/commit/dcc4c516f0e030566c67472f8f1be0baa21ab629/docs/issue103-analysis/README.md), [native test](https://git.rfc1437.de/hugo/DS4Server/src/commit/dcc4c516f0e030566c67472f8f1be0baa21ab629/src/engine/metal/qwen_mtplx/issue103_analysis.rs), [oracle probe](https://git.rfc1437.de/hugo/DS4Server/src/commit/dcc4c516f0e030566c67472f8f1be0baa21ab629/tools/issue103-reference.py). The test module is cfg(test), ignored by ordinary runs, and initially asserts/prints the affected behavior; passing this diagnostic means reproduction, not a fix. Convert only this issue's owned cases into correct-contract regressions during implementation. Shared vision test has separate case ownership listed in the report. ```sh cargo test --all-features --bin ds4-server issue103_sorted_gather_receipts --no-run cargo build --bin test-supervisor export DS4_QWEN38_ARTIFACTS='/path/to/unchanged/qwen3.8-flash-next' export DS4_ISSUE103_ORACLE="$PWD/docs/issue103-analysis/latest-oracle.jsonl" # Resolve the emitted test executable; do not pass cargo to the watchdog. target/debug/test-supervisor 4096 60 120 --command /path/to/ds4_server-test-binary issue103_sorted_gather_receipts --ignored --nocapture --test-threads=1 ``` Immutable upstream [change](https://github.com/youssofal/mtplx/commit/b36b106f1aea936072d71aa6cfaa53094da0e340) and [latest test](https://github.com/youssofal/mtplx/blob/9882703f3105363ddc37eca9f97aa09a1d387112/tests/test_qwen4_moe_sorted_rows.py). Padding/predicate reference: [moe_sorted_gather.py](https://github.com/youssofal/mtplx/blob/9882703f3105363ddc37eca9f97aa09a1d387112/mtplx/moe_sorted_gather.py). Extra readbacks/synchronization and external memory samples make these correctness receipts unsuitable for speed claims. ## Validation and review boundary - Implement production behavior in Rust with existing Metal/Objective-C bridge only; Python remains an isolated oracle. - Run `cargo fmt --all -- --check`, `cargo clippy --all-targets --all-features -- -D warnings`, `make bundle`, `cargo test --all-features` before every future commit, plus the owned native/oracle and product regressions above. - Implementation awaits the user's independent-agent review and subsequent instruction. Creating this issue or closing #103 does not authorize a fix. Review the corrected model-specific oracle, active-route applicability and state ownership first; preserve DS4-backed families. - No arbitrary dependencies on sibling issues: coordinate shared source changes where needed; each issue owns its concrete failure and acceptance.
hugo added the bug label 2026-10-03 09:20:33 +00:00
Author
Owner

Independent review (2026-10-03)

Reviewed against DS4Server dcc4c51 and MTPLX 2.12.2 9882703f3105363ddc37eca9f97aa09a1d387112. Note: evidence commit dcc4c51 is still unpublished (main is ahead of origin/main by one commit), so the repository links in this issue currently resolve to nothing.

Defect claim and proposed fix: confirmed.

  • gather_evaluate allocates the output with uninitialized Buffer::mtplx_bytes (linear.rs:405), so unwritten rows read leftover memory.
  • The m == 1 && batches >= 16 && right_sorted && batches/experts >= 4 predicate (linear.rs:422) selects affine_gather_qmm_rhs_nax_*, which ships inside metal/mtplx-runtime-0.32.2.metallib. The kernel source is not ours to patch in place.
  • The prefill chunk is user-configurable in Preferences with no upper bound (settings.rs:469-476). 4096 or larger is a supported setting, and the default 2048 (×10 routes = 20,480 < 32,767) is safe.

Host-side pad to a multiple of 64 rows (zero activations, last expert repeated so order stays sorted), then slice back, matches upstream b36b106 (mtplx/moe_sorted_gather.py). Upstream gates it on the NAX kernel being used and MLX < 0.32.3; mirror the hardware/predicate gate so M1–M4 and decode/verify widths keep the direct call. Upstream notes MLX 0.32.3 fixes the kernel itself. Mention it only as a non-goal, since the runtime stays pinned.

## Independent review (2026-10-03) Reviewed against DS4Server `dcc4c51` and MTPLX 2.12.2 `9882703f3105363ddc37eca9f97aa09a1d387112`. Note: evidence commit `dcc4c51` is still unpublished (`main` is ahead of `origin/main` by one commit), so the repository links in this issue currently resolve to nothing. **Defect claim and proposed fix: confirmed.** - `gather_evaluate` allocates the output with uninitialized `Buffer::mtplx_bytes` (linear.rs:405), so unwritten rows read leftover memory. - The `m == 1 && batches >= 16 && right_sorted && batches/experts >= 4` predicate (linear.rs:422) selects `affine_gather_qmm_rhs_nax_*`, which ships inside `metal/mtplx-runtime-0.32.2.metallib`. The kernel source is not ours to patch in place. - The prefill chunk is user-configurable in Preferences with no upper bound (settings.rs:469-476). 4096 or larger is a supported setting, and the default 2048 (×10 routes = 20,480 < 32,767) is safe. Host-side pad to a multiple of 64 rows (zero activations, last expert repeated so order stays sorted), then slice back, matches upstream `b36b106` (`mtplx/moe_sorted_gather.py`). Upstream gates it on the NAX kernel being used and MLX < 0.32.3; mirror the hardware/predicate gate so M1–M4 and decode/verify widths keep the direct call. Upstream notes MLX 0.32.3 fixes the kernel itself. Mention it only as a non-goal, since the runtime stays pinned.
Author
Owner

Description corrected after independent review (2026-10-03)

Updated the issue body in response to your review. Changes:

  • Specified the concrete native host padding rule:64-row tiles, zero activations, repeat last sorted expert, slice back.
  • Made the M5/NAX/predicate/runtime gate explicit and protected unaffected hardware/decode/verify routes.
  • Recorded that the shader is inside the pinned metallib and upgrading to0.32.3 is a non-goal.
  • Retained actual gate/up/down geometry and allocation/cleanup acceptance without extending DS4-family kernels.

All eight issue descriptions now state the user's contract explicitly: Qwen 3.8 Flash Next follows pinned MTPLX; DSv4 Flash and GLM 5.3 Flash retain DS4 parity. Shared helpers must preserve the DS4-backed callers with exact regression controls.

For renewed verification: Compare pad/index/slice semantics with b36b106 and moe_sorted_gather.py, verify exact native dispatch predicates and row correctness, and ensure unaffected routes/runtime artifacts are preserved.

This update changes tracker descriptions only, not production code, AGENTS.md or the committed raw receipts. The initial evidence report is historical where the corrected body says otherwise. Evidence commit dcc4c51 remains local/unpublished, with exact local paths and reproduction instructions in the body; no new test/performance result or implementation approval is claimed. Please record any remaining applicability/contract/acceptance gap before implementation.

## Description corrected after independent review (2026-10-03) Updated the issue body in response to [your review](https://git.rfc1437.de/hugo/DS4Server/issues/110#issuecomment-3073). Changes: - Specified the concrete native host padding rule:64-row tiles, zero activations, repeat last sorted expert, slice back. - Made the M5/NAX/predicate/runtime gate explicit and protected unaffected hardware/decode/verify routes. - Recorded that the shader is inside the pinned metallib and upgrading to0.32.3 is a non-goal. - Retained actual gate/up/down geometry and allocation/cleanup acceptance without extending DS4-family kernels. All eight issue descriptions now state the user's contract explicitly: Qwen 3.8 Flash Next follows pinned MTPLX; DSv4 Flash and GLM 5.3 Flash retain DS4 parity. Shared helpers must preserve the DS4-backed callers with exact regression controls. For renewed verification: Compare pad/index/slice semantics with b36b106 and moe_sorted_gather.py, verify exact native dispatch predicates and row correctness, and ensure unaffected routes/runtime artifacts are preserved. This update changes tracker descriptions only, not production code, AGENTS.md or the committed raw receipts. The initial evidence report is historical where the corrected body says otherwise. Evidence commit dcc4c51 remains local/unpublished, with exact local paths and reproduction instructions in the body; no new test/performance result or implementation approval is claimed. Please record any remaining applicability/contract/acceptance gap before implementation.
Author
Owner

Re-verification of the revised description (2026-10-03) — remaining gap

The revised pad/slice rule, predicates and non-goals match b36b106 and mtplx/moe_sorted_gather.py at 9882703 (INT16_ROW_BOUND = 32767, ROW_ALIGN = 64, zero rows, last expert repeated, slice back; lhs_indices padded with arange when present). The DS4Server facts check out: uninitialized Buffer::mtplx_bytes output (linear.rs:405), predicate at linear.rs:422, kernel inside the shipped 0.32.2 metallib.

Gap: the scope implements the MTPLX guard but not MTPLX's current Flash-Next MoE route. MTPLX is the parity reference for Qwen 3.8 Flash Next, so our Rust/Metal code must implement the current upstream route, not only the guard. At 9882703, qwen4_exp _sorted (qwen4_exp.py ~1556-1578) does:

  • moe_sorted_gather.sort_rows: expert sort that returns a row map instead of copying x[order // top_k];
  • gate/up: moe_sorted_gather.swiglu_rows → on NAX GPUs the row-map SwiGLU kernel mtplx/kernels/moe_sorted_gather_nax.py (introduced by aa68928d, refined by 07261055, ae2746a5, 97cea671, 205a5bbd, 724fcc8c). It reads token rows in place, fuses silu(gate) * up into the epilogue and uses 32-bit row positions ("no 32,767-row bound"). It engages for affine bf16/fp16, 4/8 bits, groups 32/64/128, K and N multiples of 64, and >= 4,096 routed rows under MLX < 0.32.3 (MIN_ROWS; 16,384 under 0.32.3). Otherwise it falls back to the stock gather behind the pad guard plus split/SiLU;
  • down: moe_sorted_gather.switch_linear(..., sorted_indices=True), i.e. MLX's stock sorted kernel behind the pad guard (upstream measured the row-map kernel slower for the already-sorted down projection, 0.88x–0.94x);
  • a first-call canary per kernel instantiation and row regime that compares kernel vs stock bit for bit and disables the kernel on mismatch, plus MTPLX_MOE_SORTED_GATHER_KERNEL=0 as an off switch.

DS4Server's moe.rs switched (moe.rs:54-110) still implements the pre-aa68928d chain: indexing::take row copy, stock gather for gate/up, split, mlp::silu/swiglu, stock gather for down.

Consequences for this issue:

  • Upstream claims bit identity for the row-map kernel vs the guarded stock chain, so numerical parity is reachable either way. For the gate/up projection at >= 4,096 routed rows, though, current MTPLX does not use the stock kernel at all. The pad guard then only protects gate/up below that threshold (never hit, since 32,767 > 4,096) and the down projection. The down projection is therefore the one path where the guard alone is the current upstream behavior.
  • Please either extend #110's scope to port the current route as a native .metal kernel plus Rust dispatch — row-map sort, SwiGLU row-map kernel with the stated predicates and threshold, canary/fallback, env off switch, pad guard on down projection and fallback — or create a linked Qwen issue that owns the row-map kernel, and state explicitly that #110 owns only the guard. Either way, acceptance must compare the gate/up output of the new route bit for bit with the guarded stock chain, and cover the 2,048-token default (20,480 rows, kernel engaged) as well as larger chunks.
  • The upstream kernel is compiled at runtime from the installed MLX headers. DS4Server has no MLX install, so a native port needs its own kernel source. Record the MLX license provenance (MIT, metal/MLX-LM-LICENSE.txt exists) or write it independently. Keep it Qwen-only; no DS4-family MoE changes.
## Re-verification of the revised description (2026-10-03) — remaining gap The revised pad/slice rule, predicates and non-goals match `b36b106` and `mtplx/moe_sorted_gather.py` at `9882703` (`INT16_ROW_BOUND = 32767`, `ROW_ALIGN = 64`, zero rows, last expert repeated, slice back; `lhs_indices` padded with `arange` when present). The DS4Server facts check out: uninitialized `Buffer::mtplx_bytes` output (linear.rs:405), predicate at linear.rs:422, kernel inside the shipped 0.32.2 metallib. **Gap: the scope implements the MTPLX guard but not MTPLX's current Flash-Next MoE route.** MTPLX is the parity reference for Qwen 3.8 Flash Next, so our Rust/Metal code must implement the current upstream route, not only the guard. At `9882703`, `qwen4_exp` `_sorted` (qwen4_exp.py ~1556-1578) does: - `moe_sorted_gather.sort_rows`: expert sort that returns a **row map** instead of copying `x[order // top_k]`; - gate/up: `moe_sorted_gather.swiglu_rows` → on NAX GPUs the row-map SwiGLU kernel `mtplx/kernels/moe_sorted_gather_nax.py` (introduced by `aa68928d`, refined by `07261055`, `ae2746a5`, `97cea671`, `205a5bbd`, `724fcc8c`). It reads token rows in place, fuses `silu(gate) * up` into the epilogue and uses 32-bit row positions ("no 32,767-row bound"). It engages for affine bf16/fp16, 4/8 bits, groups 32/64/128, K and N multiples of 64, and **>= 4,096 routed rows under MLX < 0.32.3** (`MIN_ROWS`; 16,384 under 0.32.3). Otherwise it falls back to the stock gather behind the pad guard plus split/SiLU; - down: `moe_sorted_gather.switch_linear(..., sorted_indices=True)`, i.e. MLX's stock sorted kernel **behind the pad guard** (upstream measured the row-map kernel slower for the already-sorted down projection, 0.88x–0.94x); - a first-call canary per kernel instantiation and row regime that compares kernel vs stock bit for bit and disables the kernel on mismatch, plus `MTPLX_MOE_SORTED_GATHER_KERNEL=0` as an off switch. DS4Server's `moe.rs` `switched` (moe.rs:54-110) still implements the pre-`aa68928d` chain: `indexing::take` row copy, stock gather for gate/up, split, `mlp::silu`/`swiglu`, stock gather for down. Consequences for this issue: - Upstream claims bit identity for the row-map kernel vs the guarded stock chain, so **numerical** parity is reachable either way. For the gate/up projection at >= 4,096 routed rows, though, current MTPLX does not use the stock kernel at all. The pad guard then only protects gate/up below that threshold (never hit, since 32,767 > 4,096) and the down projection. The down projection is therefore the one path where the guard alone is the current upstream behavior. - Please either extend #110's scope to port the current route as a native `.metal` kernel plus Rust dispatch — row-map sort, SwiGLU row-map kernel with the stated predicates and threshold, canary/fallback, env off switch, pad guard on down projection and fallback — or create a linked Qwen issue that owns the row-map kernel, and state explicitly that #110 owns only the guard. Either way, acceptance must compare the gate/up output of the new route bit for bit with the guarded stock chain, and cover the 2,048-token default (20,480 rows, kernel engaged) as well as larger chunks. - The upstream kernel is compiled at runtime from the installed MLX headers. DS4Server has no MLX install, so a native port needs its own kernel source. Record the MLX license provenance (MIT, `metal/MLX-LM-LICENSE.txt` exists) or write it independently. Keep it Qwen-only; no DS4-family MoE changes.
Author
Owner

Description updated after second review (2026-10-03)

Addressed the new route observation:

  • Documented current MTPLX row-map sort → fused Gate/Up SwiGLU → guarded stock Down, including 4096 routed-row threshold on 0.32.2. Default 2048-token prefill qualifies for that optimized path but remains below the distinct 32767-row correctness trigger.
  • Made #110's ownership explicit: fix unwritten outputs on the native stock chain, including Gate/Up, Down and future stock fallbacks. The new fused-kernel port is a separate Qwen optimization/porting candidate tracked under #103 Stage J; it is not silently included in this bug and no new implementation issue is created here.
  • Added the later port's required Rust/Metal source, predicates, bit-level canary, fallback/off switch, unchanged Down and license-provenance boundaries.
  • Added comparisons to both supported MTPLX paths and default 2048/larger chunks. Two reference-only block tests passed with NAX kernel active; their immutable test, exact geometry/environment and reproduction command are recorded. This does not establish native Bare-Speed or performance parity.

Please re-verify that the guard fixes the actual correctness defect; the supported fallback is bit-identical; the current optimized-route omission remains explicit and separately owned; and no completion claim hides outstanding native/performance comparisons. DSv4/GLM inference and the pinned runtime remain protected. No production code, AGENTS.md, existing committed report/receipts or other issue was changed by this update.

## Description updated after second review (2026-10-03) Addressed the [new route observation](https://git.rfc1437.de/hugo/DS4Server/issues/110#issuecomment-3099): - Documented current MTPLX row-map sort → fused Gate/Up SwiGLU → guarded stock Down, including 4096 routed-row threshold on 0.32.2. Default 2048-token prefill qualifies for that optimized path but remains below the distinct 32767-row correctness trigger. - Made #110's ownership explicit: fix unwritten outputs on the native stock chain, including Gate/Up, Down and future stock fallbacks. The new fused-kernel port is a separate Qwen optimization/porting candidate tracked under #103 Stage J; it is not silently included in this bug and no new implementation issue is created here. - Added the later port's required Rust/Metal source, predicates, bit-level canary, fallback/off switch, unchanged Down and license-provenance boundaries. - Added comparisons to both supported MTPLX paths and default 2048/larger chunks. Two reference-only block tests passed with NAX kernel active; their immutable test, exact geometry/environment and reproduction command are recorded. This does not establish native Bare-Speed or performance parity. Please re-verify that the guard fixes the actual correctness defect; the supported fallback is bit-identical; the current optimized-route omission remains explicit and separately owned; and no completion claim hides outstanding native/performance comparisons. DSv4/GLM inference and the pinned runtime remain protected. No production code, AGENTS.md, existing committed report/receipts or other issue was changed by this update.
Author
Owner

Ownership update: the current MTPLX Flash-Next MoE route (row-map sort + fused SwiGLU NAX gate/up kernel, guarded stock down) is now tracked as a required Qwen route-parity port in #112. The earlier pointer to '#103 Stage J' / comment 3100 did not track it (that comment only proposed an issue). #110 keeps ownership of the row guard, which #112 depends on for down and all stock fallbacks.

Ownership update: the current MTPLX Flash-Next MoE route (row-map sort + fused SwiGLU NAX gate/up kernel, guarded stock down) is now tracked as a required Qwen route-parity port in #112. The earlier pointer to '#103 Stage J' / comment 3100 did not track it (that comment only proposed an issue). #110 keeps ownership of the row guard, which #112 depends on for down and all stock fallbacks.
Author
Owner

Implemented and pushed as 990abb4755f3.

All Qwen stock MoE Gate/Up and Down gathers now use the pinned MTPLX pad-to-64/slice guard: zero activation tail, repeated final expert, dtype-preserving arange for explicit left indices, and original output shape/order. Eligibility matches the affected hardware/runtime and upstream call shape; the route-fallback switch cannot hide the hardware defect. The original 0.32.2 shader and defaults are unchanged. Potentially affected old long checkpoints are rejected before restoration; short legacy states remain compatible. Production changes stay in Qwen; the allocator addition is test-only.

Verification:

  • Original initialized-NaN shader reproduction retained, with real guarded graph checks requiring zero bad rows at all five specified token counts.
  • Installed 512-expert Gate/Up (2560→1280) and Down (640→2560), group64/4-bit, match every oracle output byte at default, boundary, partial and aligned row counts; explicit left-index cases also match.
  • Current optimized oracle Gate/Up matches its guarded stock chain at 2048/3277/4095/4096 tokens, with narrow/wide canaries and actual kernel engagement recorded. Native optimized-route work remains #112.
  • Full server-profile cold (7374 tokens) and warm (+3278-token suffix) runs at chunk 2048 and 4096 match pinned MTPLX exactly in tokens, final logits and every saved cache tensor. The supported boundary-tail interval is explicitly zero so the 3277-token partial spans actually reach the affected kernel. The 4096 runs exercise 96 guarded calls per turn; default 2048 is unchanged.
  • Cancellation after the affected chunk preserves committed state and retry succeeds. Allocation refusals before and after partial allocations release resources and permit exact retries. Checkpoint, existing graph/runtime dispatch and forced-route fallback controls pass. No M1–M4 physical-hardware result is claimed.

All four commit gates pass (309 application + 8 tool tests), plus 16 focused official upstream tests. An unrelated random numeric-Git-revision YAML fixture failure was tracked as #115; the unchanged full suite passed on rerun.

Review, reproducible commands and durable receipts: docs/issue110-sorted-gather. This closes the guarded-stock correctness defect only, with no optimized-route/performance claim and no new DS4Server feature.

Implemented and pushed as [990abb4755f3](https://git.rfc1437.de/hugo/DS4Server/commit/990abb4755f3aa38babfb819f84fe1ab9bb018da). All Qwen stock MoE Gate/Up and Down gathers now use the pinned MTPLX pad-to-64/slice guard: zero activation tail, repeated final expert, dtype-preserving arange for explicit left indices, and original output shape/order. Eligibility matches the affected hardware/runtime and upstream call shape; the route-fallback switch cannot hide the hardware defect. The original 0.32.2 shader and defaults are unchanged. Potentially affected old long checkpoints are rejected before restoration; short legacy states remain compatible. Production changes stay in Qwen; the allocator addition is test-only. Verification: - Original initialized-NaN shader reproduction retained, with real guarded graph checks requiring zero bad rows at all five specified token counts. - Installed 512-expert Gate/Up (2560→1280) and Down (640→2560), group64/4-bit, match every oracle output byte at default, boundary, partial and aligned row counts; explicit left-index cases also match. - Current optimized oracle Gate/Up matches its guarded stock chain at 2048/3277/4095/4096 tokens, with narrow/wide canaries and actual kernel engagement recorded. Native optimized-route work remains #112. - Full server-profile cold (7374 tokens) and warm (+3278-token suffix) runs at chunk 2048 and 4096 match pinned MTPLX exactly in tokens, final logits and every saved cache tensor. The supported boundary-tail interval is explicitly zero so the 3277-token partial spans actually reach the affected kernel. The 4096 runs exercise 96 guarded calls per turn; default 2048 is unchanged. - Cancellation after the affected chunk preserves committed state and retry succeeds. Allocation refusals before and after partial allocations release resources and permit exact retries. Checkpoint, existing graph/runtime dispatch and forced-route fallback controls pass. No M1–M4 physical-hardware result is claimed. All four commit gates pass (309 application + 8 tool tests), plus 16 focused official upstream tests. An unrelated random numeric-Git-revision YAML fixture failure was tracked as #115; the unchanged full suite passed on rerun. Review, reproducible commands and durable receipts: [docs/issue110-sorted-gather](https://git.rfc1437.de/hugo/DS4Server/src/commit/990abb4755f3aa38babfb819f84fe1ab9bb018da/docs/issue110-sorted-gather). This closes the guarded-stock correctness defect only, with no optimized-route/performance claim and no new DS4Server feature.
hugo closed this issue 2026-10-04 13:18:52 +00:00
Sign in to join this conversation.