Skip to content

One decode GEMM table on MI325X, each size's kernel picked by measurement - #1104

Draft
merceod wants to merge 9 commits into
users/atindra/pr-36-t6-fused-draftfrom
users/atindra/pr-38-f4-gemm-arbitration
Draft

merceod wants to merge 9 commits into
users/atindra/pr-36-t6-fused-draftfrom
users/atindra/pr-38-f4-gemm-arbitration

Conversation

@merceod

@merceod merceod commented Nov 1, 2026 •

Copy link
Copy Markdown
Collaborator

Stacked on #1102 (pr-36-t6-fused-draft).

Why

Decode linears on ROCm were routed by three configs (decode_gemv, decode_gemv_silu_mul, decode_linear) and a chain of passes, each deciding for its own kernel. Nothing put AITER's skinny GEMM, the Triton GEMV and tuned hipBLASLt side by side for one size, and moving a size to another kernel meant editing a different file per kernel.

What changed

Stacked on #1102, nothing picked first: the table replaces the decode GEMV configs #1102 added rows 3, 6, 12 and the fc to, and the last commit puts those rows into the table.

  • KernelChoice in mstar/kernels/registry.py reads one per-device table, configs/gfx942/mi325x/decode_gemm.json, which replaces the three configs. Each entry keeps every candidate's time and parameters beside the choice, so moving a size is a one-word edit. A fused gate/up has its own gated1- entry.
  • mstar/platforms/rocm_gemm.py: one pass route for AITER, Triton, tuned hipBLASLt and plain hipBLASLt at rows 1-128 in decode buckets. DecodeLinearPass and the pass chain are gone. The route keys on a bucket's rows, so Speculative decoding on Qwen3.5: a fused draft step and an fp8 draft head #1102's verify buckets (8 rows, 24 tokens) keep their passes. Tuned solutions apply only on the stack they were tuned on.
  • python -m benchmark.kernels.promote decode_gemm --store DIR writes the table from the harness store. A size leaves what the tree serves only for a 2 % gain and one graph node (a fused gate/up for the 2 % alone): without the node floor a first table moved 2B's rows-1 out_proj to Triton for 0.17 us a call and broke an Inductor fold (247 to 265 kernels a step).
  • 185 entries: 159 from the search (144 sizes over 9B, 2B and 0.8B, 15 gate/up), Speculative decoding on Qwen3.5: a fused draft step and an fp8 draft head #1102's 16 rows with the kernel and knobs it shipped, and 10 naming the class default. The search store alone gives back the 159 unchanged.
  • Off ROCm skinny_gemm_passes() returns None before any table read, and products without a table keep today's defaults. Docs in docs/kernels.rst.

Verification

model c=1 c=4 c=8 c=16 c=32 c=48 c=64
9B 224 / 224 766 / 767 1473 / 1467 2586 / 2657 4313 / 4636 5793 / 6060 7139 / 7273
2B 608 / 619 2035 / 2026 3824 / 3765 6760 / 6816 10585 / 11193 14585 / 15006 17408 / 18453
0.8B 880 / 839 2601 / 2639 4881 / 5020 8819 / 8923 13974 / 14243 19167 / 19254 23107 / 22923
  • 9B +2.7 % at c=16, +7.5 % at c=32, +4.6 % at c=48, +1.9 % at c=64, ITL p50 7.10 to 6.56 ms at c=32. 2B +5.7 % at c=32 and +6.0 % at c=64 (ITL 3.27 to 3.11 ms), -1.5 % at c=8. 0.8B +1 to +3 % at c=4-32, flat at c=48-64, and -2 to -5 % at c=1 on the median, where both arms carry stalled runs (best run -1.9 %).
  • Stalled runs are as common on the tip as here: 41 of 144 against 40 of 138. Two losses a shared node showed (2B c=1 -20 %, 9B c=4 -34 %) did not come back on the exclusive node.
  • On the week-3 tip with Speculative decoding on Qwen3.5: a fused draft step and an fp8 draft head #1102's rows in, one job at a time on an empty queue, direction only: 9B 225.9 against 223.8 tok/s at c=1 and 7287 against 7147 at c=64 (+0.6 %, +1.8 %). 9B mtp2 reads 409 against 410 at 2.45 tokens a chunk, so Speculative decoding on Qwen3.5: a fused draft step and an fp8 draft head #1102's path keeps its numbers.
  • Correctness: 25 greedy passes on 9B and 2B with 0 unexplained divergences, and greedy PASS on every window I arm. 1416 GPU cases check every kernel the table holds against the parity bar, and each served one from a graph replay.
  • Tests, from the branch and integration trees, not this cut branch: CPU Core failing set equal to the base's (34) on the branch and on the week-3 tip, test/rocm 1965 passed and 9 skipped on the tip with the merged table, ruff clean.

Not in this PR

  • The 0.8B c=1 -2 % is unexplained. Its rows-1 entries come from one search pass, and a rerun of those 7 sizes on a quiet node changed no pick. The two passes disagree on 9B rows-16 gate/up, which rests on the pooled time.
  • Rows 96 and 128 keep hipBLASLt's pick, and the Triton kernels weren't tried past 16 rows. With Speculative decoding on Qwen3.5: a fused draft step and an fp8 draft head #1102's rows in, only 9B was served (c=1, c=64, mtp2 c=1): no 2B, 0.8B, 9B c=16 or c=32 yet.
  • Only MI325X has a table. Norm has nothing to pick (Triton wins at every decode row), and attention needs a bench that times a candidate inside a step first. NVIDIA not run, and nothing there reads the table.

Rebased onto main

Rebased onto main (ff7819d, 2026-11-05); the same 9 commits. The MI325X numbers here were measured on the AMD line before its rebase; the rebased line read level on the Qwen3-TTS, Qwen3-Omni, Qwen3.5 mtp2 and FLUX.2 klein smoke rows, and -8.7 % on the Qwen3.5-0.8B text c=64 row, which waits on the twin-slot re-port (phase 2).

@merceod
merceod force-pushed the users/atindra/pr-36-t6-fused-draft branch from 3707ec6 to d077088 Compare November 5, 2026 17:52
@merceod
merceod force-pushed the users/atindra/pr-38-f4-gemm-arbitration branch from bb62fe4 to 3d41607 Compare November 5, 2026 17:52

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant