[heslami]

Motivation

Stacked on #TBD (the indexer launch PR) (and the stack under it). Merge those first; this PR's diff is its last two commits.

At 8 rows (MTP at two requests) the MoE stages ran every routed slot once a row. An expert that two rows picked streamed its weights twice, and the shared expert, which every row uses, ran its tiles once a row in down. At 8 rows, 8 picks a row plus the shared expert give up to 65 (row, expert) uses a step, and rows of the same requests share many of them. The MoE phase of the launch is bound by the bytes each CU streams after its first weights, so every repeated expert is time.

• up/gate now runs one job per 8-intermediate tile of each expert the step uses: the shared one and every distinct routed expert, in order of first use. Each job takes every row that picked the expert as an MFMA column, and the jobs are dealt round-robin over the CTAs.
• down streams each of those experts once too. Its units multiply the expert's tile with every using row's mid chunk and keep the products in LDS. Each row's output is then summed over its units in the order and wave segments of the one-row form.
• The skip: a down weight batch that lies entirely past the step's expert groups skips its MFMAs (their buffers are already empty).

Up to 4 rows the repeats are few and the per-expert overheads (finding the groups, the products in LDS, the summing pass) cost about what they save, so 2 and 4 rows keep the one-row forms. Dense layers keep theirs too.

Every output and every write stays bitwise equal to the base, at TP8 and TP4, every width, on the bf16 and the E4M3 MLA cache (Accuracy Tests).

Result on MI355X, base → this PR (TP8, FP8 E4M3 KV, MTP 3-1-4 at the golden accept length 2.99, 8K):

• Per cycle at two requests (8 rows): TBD ms. At one request (4 rows): no change expected, since the 2- and 4-row paths keep their code and measured within noise.
• TP4, two requests: TBD ms a cycle.
• Kernel-only (time_layers.py, TP8): model order at 8 rows, 8K and 128K: TBD µs a launch (Speed Tests).

Modifications

Two commits, one per change. Each was checked against its parent (Accuracy Tests) and timed against it (Speed Tests).

#	Commit	Change	Numerics	Kernel-only (8 rows, µs a launch)
1	each expert once in up/gate and down	Above 4 rows (not dense layers): every wave finds the step's distinct routed experts from the top-8s (lane p holds pick p, a ballot marks first uses). up/gate jobs are the 8-intermediate tiles of the shared expert and of each distinct expert, each with every using row as an MFMA column, dealt round-robin over the CTAs. Every CTA routes, issues the activations' poll and its first job's weights, then checks the poll. down's units take one expert's tile against every using row and store the products in LDS; a row's output is summed over its units in the one-row form's order and wave segments, so a column's sums are a one-row tile's. The rows the multi-column MFMAs read (activations and FP8 mid blocks) are padded by 4 words against bank conflicts	bitwise	TBD
2	dead-batch skip	A wave's down weight batch whose first unit lies past the step's expert groups skips its MFMAs	bitwise	TBD

• Layout: no new mailboxes and no scratch growth. The group table (count, each group's expert, each row's slot of it) and down's products live in the LDS X region. Above 4 rows the up/gate stage's task count (timeline stamps only) covers up to 65 groups' tiles.
• poll takes an optional between() that runs after the first loads are issued and before they are checked; stage_xq passes it through. up/gate uses it to issue the first job's weights under the activations' round trip.
• Up to 4 rows and dense layers keep their code. Their objects only move by build jitter (Registers).

Accuracy Tests

All on MI355X in the image below, with the in-tree tools under torchrun at TP8 and TP4.

• Bitwise (check_layer.py --dump, --compare):

  • Kinds shared (3 launches), index (3), model (the first 8 launches) and dense (3).
  • Widths 2, 4 and 8 at TP8 and TP4.
  • Row sets p0, p1 and v4. bf16 and E4M3 MLA cache. 3 steps.
  • 6,588 digests a TP a tree.

  Against	TP8	TP4
  1 each expert once, against the base	TBD	TBD
  2 dead-batch skip, against commit 1	TBD	TBD

  • No error word was set.
  • Negative control: one input value negated changes TBD (TP8) and TBD (TP4) of the digests.

• NextN launches (check_nextn.py --only entry; NextN index at 4 and 8 rows, NextN shared at 2, 4 and 8 rows): TBD at TP8 and TP4 on both caches.

• E2E (test/e2e/test_glm53_server.py, TP8): TBD.

• GSM8K with MTP, two requests: TBD.

• Unit tests: TBD.

Speed Tests and Profiling

Setup as the base: image lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025 with SGLang's python/sglang overlaid at 5044779f, zai-org/GLM-5.3 at ba32f8c1, the e2e test's server flags with --kv-cache-dtype fp8_e4m3.

TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms:

Requests	base	this PR	Change
1 (verify 4 rows, draft extend 4)	TBD	TBD	TBD
2 (8 rows)	TBD	TBD	TBD

TP4, FP8 E4M3 KV, MTP at AL 2.99, two requests: TBD → TBD ms.

Kernel-only (time_layers.py, TP8, E4M3 KV, µs a launch; each commit against its parent, rounds in alternating order):

Launch	base	commit 1	commit 2	1 − base	2 − 1
Model order (75), 2 rows, 8K	TBD	TBD	TBD	TBD	TBD
Model order (75), 4 rows, 8K	TBD	TBD	TBD	TBD	TBD
Model order (75), 8 rows, 8K	TBD	TBD	TBD	TBD	TBD
Model order (75), 8 rows, 128K	TBD	TBD	TBD	TBD	TBD
Shared, 8 rows, 8K	TBD	TBD	TBD	TBD	TBD
Indexer, 8 rows, 8K	TBD	TBD	TBD	TBD	TBD

Registers

TBD.

Reviewer notes

• Why 8 rows only. At 2 and 4 rows the dedup measured within noise during development (Speed Tests), and the one-row forms keep their per-row prefetch.
• The skip needs the dedup's group count; it does not apply to the one-row down.
• The group table is per CTA. Every CTA routes all rows and finds the groups itself, so no cross-CTA hand-off is added; the routing and its tie order are unchanged.
• kernel.py grows by about 470 lines, most of it the 8-row up/gate and down paths; a candidate for the stage-builder split.

Checklist

• pre-commit run --all-files passes. (ruff, ruff-format and isort clean on the touched files; full run TBD)
• Tests under test/ cover the change.
• The SGLang commit or PR this was tested against is named above.
• Accuracy and speed results are attached when kernel or forward code changes. (Serving: TBD)
• A new model adds one line to .github/CODEOWNERS that names its owner. (No new model.)

[heslami]

Superseded by the regrouped preview #34 (and #35 for the long-context indexer).