Skip to content

npu docs: a MoE decode step is decided by down_proj's A shape, not the expert stream - #896

Merged
bong-water-water-bong merged 5 commits into
mainfrom
npu-moe-dispatch-cost
Nov 3, 2026
Merged

bong-water-water-bong merged 5 commits into
mainfrom
npu-moe-dispatch-cost

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

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 N chunk-list dispatch, measured
2026-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 — but down_proj's input is
each 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, then
gate|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 tensor
shapes), not recalled: routed experts 1,019,215,872 B (8 of 128 x 3 projections, pb16 NW=16) +
attention 509,607,936 + lm_head 175,030,272 = 1,703,854,080 B/token; --pad-n 1024 on gate/up
gives 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.

per-dispatch fixed cost down shares one A down per expert Qwen3-4B dense, same arithmetic
0.086 ms (fit) 11.05 tok/s (1.06x) 8.37 (0.80x) 10.45
0.116 ms (quoted floor) 9.80 (1.01x) 7.09 (0.73x) 9.68

(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-margin
finding. No code, no device, tools/copyright.py --check exits 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).

…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.
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

#896 ready for review at 47459ff (docs-only, docs/npu.md +28/-0, MERGEABLE, copyright --check exit 0; no new files, no device).

@github-actions

github-actions Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

PR Reviewer Guide 🔍

(Review updated until commit c016a2b)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪
🧪 No relevant tests
🔒 No security concerns identified
⚡ Recommended focus areas for review

Documentation Clarity

The documentation update describes a hardware-specific behavior regarding MoE dispatch costs, but it does not include any code changes or tests. While the PR is labeled as "docs-only", the content is quite technical and could benefit from clearer explanation of how the dispatch cost difference impacts performance or design decisions. It would be helpful to include a brief note on how this information should be used in future development or tuning efforts.

**A chunk-list dispatch shares ONE A across its chunks: exact for gate and up, impossible for down.**
At batch 1 every selected expert of a layer sees the same activation row, so gate and up are one
dispatch each — but `down_proj`'s input is each expert's own SwiGLU output, so it needs a per-chunk A
or one dispatch per expert. On this model (48 layers; per layer q, gate, k, v, o in the pb16i split,
then gate|up|down) that is **385 dispatches per token with the shared A and 721 without**, and that
gap decides whether the case meets its own bar:

| per-dispatch fixed cost | down shares one A | down per expert | Qwen3-4B dense, same arithmetic |
|---|---|---|---|
| 0.086 ms (fit below) | **11.05 tok/s (1.06x)** | 8.37 (0.80x) | 10.45 |
| 0.116 ms (quoted floor) | 9.80 (1.01x) | 7.09 (0.73x) | 9.68 |

That is at 34 GB/s in-loop over a **1,703,854,080 B/token** stream — routed experts 1,019,215,872
(8 of 128 x 3 projections, pb16 NW=16) + attention 509,607,936 + `lm_head` 175,030,272, with
`--pad-n 1024` on gate/up taking it to 1,930,346,496 — and the ratio holds across 34-43.4 GB/s:
**0.97-1.06x the 4B-dense bar when down shares the dispatch, 0.68-0.80x when it cannot.** So the
deciding term is the down_proj A shape, not the expert stream: sizing a MoE case on the expert
bytes alone understates the stream by 40 % (they are 60 % of it) and misses the term that decides
it.

The fixed cost above is measured, not assumed. Two runs of the same 50,135,040-B weight chunk in
the arena's chunk-list design — 1.051 ms for one chunk, 2.016 ms for two — fit a marginal rate of
51.95 GB/s and **0.086 ms paid once per dispatch**, i.e. a second weight BO on the same dispatch
costs nothing measurable (within the fit residual; the 2-chunk run is 4.1 % faster per chunk).
Pricing N chunk BOs as N launches overcharges the dispatch column by that factor: it would add 672
dispatches per token here (48 layers x gate and up x 8 chunks, against 96 real ones), +66.5 ms at
0.099 ms each.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@bong-water-water-bong bong-water-water-bong added the tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied label Nov 3, 2026
@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

Persistent review updated to latest commit 4734664

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

Persistent review updated to latest commit 86ef568

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

Persistent review updated to latest commit ba9ebe1

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

Persistent review updated to latest commit c016a2b

@bong-water-water-bong
bong-water-water-bong merged commit 03eb3c9 into main Nov 3, 2026
8 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the npu-moe-dispatch-cost branch November 3, 2026 14:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Review effort 2/5 tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant