158 lines
9.4 KiB
Markdown
158 lines
9.4 KiB
Markdown
# 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.
|