Repository navigation
npu docs: a MoE decode step is decided by down_proj's A shape, not the expert stream - #896
Conversation
…e expert stream The chunk-list dispatch shares one A across its N weight chunks. That is exact for gate/up (every selected expert of a layer sees the same activation row) and impossible for down_proj (each expert's input is its own SwiGLU output), so down costs one dispatch per expert: 385 dispatches/token with a shared A, 721 without. Priced with the same arithmetic that reproduces the 4B-dense bar, MoE-on-NPU lands at 0.97-1.06x that bar when down shares the dispatch and 0.68-0.80x when it cannot, over a 1,703,854,080 B/token stream measured from the container's own tile table (experts 60 %, attention 30 %, lm_head 10 %). The per-dispatch fixed cost is measured, not assumed: the arena's 1-chunk and 2-chunk runs of the same 50,135,040-B chunk (1.051 / 2.016 ms) fit 51.95 GB/s marginal and 0.086 ms paid once per dispatch, so an extra weight BO on the same dispatch is free within the fit residual, and pricing N chunk BOs as N launches overcharges by that factor.
PR Reviewer Guide 🔍(Review updated until commit c016a2b)Here are some key observations to aid the review process:
|
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY
Docs-only, 28 lines in docs/npu.md, and I re-derived the arithmetic instead of taking the table.
Dispatch counts. 48 layers, 8 dispatches per layer without the down term (q, gate, k, v, o, then gate|up|down), plus one lm_head: 48 x 8 + 1 = 385 with down_proj sharing the dispatch's single A. Per-expert down adds 7 more per layer: 48 x 15 + 1 = 721. Both match the text, and the 336-dispatch difference is exactly 48 x 7.
The stream. 1,019,215,872 + 509,607,936 + 175,030,272 = 1,703,854,080 B/token — exact. The routed-expert share is 59.8 %, so "sizing a MoE case on the expert bytes alone understates the stream by 40 %" is right.
The fixed cost, from the two runs quoted. 1.051 ms for one 50,135,040 B chunk and 2.016 ms for two: marginal (2.016-1.051)/50,135,040 B = 51.94 GB/s, and 1.051 - 50,135,040/51.94e9 = 0.0861 ms paid once per dispatch. The second chunk costs 4.1 % less per chunk than the first ((2.016/2)/1.051), which is what the text says. The overcharge figure also closes: 48 layers x gate and up x 7 extra chunks = 672 extra dispatches beside the 96 real ones, and 672 x 0.099 ms = 66.5 ms.
The ratios are the quotients they claim: 11.05/10.45 = 1.057, 8.37/10.45 = 0.801, 9.80/9.68 = 1.012, 7.09/9.68 = 0.732.
Nothing here needs a device to re-check, which is the point of putting it in docs/npu.md — the next person sizing a MoE case on this route reads the deciding term before spending a ticket on it. The one thing I would keep in mind is that the tok/s rows come from a cost model fed by those two measurements, not from an end-to-end run; the text says "fit below" and "quoted floor", so that is disclosed rather than implied.
|
Persistent review updated to latest commit 4734664 |
|
Persistent review updated to latest commit 86ef568 |
|
Persistent review updated to latest commit ba9ebe1 |
|
Persistent review updated to latest commit c016a2b |
Docs-only. Records one measured fact about the NPU's chunk-list dispatch before someone reopens a
MoE inference case on this route (the parked TMOE-NPU lane's brief), because the missing term is
larger than the term the case was sized on.
Root cause of the gap in the record
@runtime_sequence(A, W0..W{N-1}, C)(the arena's--wchunks Nchunk-list dispatch, measured2026-11-01 12:43) shares one A across its N weight chunks. At batch 1 every selected expert of
a layer sees the same activation row, so this is exact for
gate/up— butdown_proj's input iseach expert's own SwiGLU output, so down cannot use it and costs one dispatch per expert.
Per layer, pb16i: q, gate, k, v (four launches,
npu/forward/src/q4nx_forward.cpp:771), o, thengate|up|down. That is 385 dispatches/token with a shared A and 721 without on Qwen3-Coder-30B-A3B
(48 layers). 336 extra dispatches is the whole difference.
Measured before/after
Stream per decode step, read out of the container's own tile table (
__metadata__.moe+ the tensorshapes), not recalled: routed experts 1,019,215,872 B (8 of 128 x 3 projections, pb16 NW=16) +
attention 509,607,936 +
lm_head175,030,272 = 1,703,854,080 B/token;--pad-n 1024on gate/upgives 1,930,346,496.
Per-dispatch fixed cost, from the arena's own two runs of the same 50,135,040-B chunk:
1.051 ms (1 chunk) and 2.016 ms (2 chunks) fit 51.95 GB/s marginal and 0.086 ms paid once per
dispatch — a second weight BO on the same dispatch costs nothing measurable. So pricing N chunk
BOs as N launches (which an earlier entry in this lane's state file did) overcharges by that factor:
+672 dispatches, +66.5 ms/token.
(at 34 GB/s in-loop; the ratio holds over 34-43.4 GB/s: 0.97-1.06x with a shared A, 0.68-0.80x without.)
So the deciding term for a MoE case here is the down_proj A shape, not the expert bytes — which
are only 60 % of a decode step's stream. A case sized on expert bytes alone understates the stream by
40 % and misses the term that decides whether it clears its bar.
Docs-only:
docs/npu.md, +27 lines, inside "Routed expert stacks" next to the existing top-8-marginfinding. No code, no device,
tools/copyright.py --checkexits 0 (no new files).Not part of this PR, and stated so it is not misread: this is arithmetic on measured rates, not a
device measurement of a MoE kernel — no MoE-shaped ELF exists (the lane is parked by the owner).