Repository navigation
[Config] [Perf] Tune DeepSeek-V4-Pro small-M FP8 GEMMs on gfx1250 - #7644
Conversation
Add 15 gfx1250 mxfp8_128 bpreshuffle rows to the DeepSeek-V4 table for TP4 decode GEMMs that run the untuned fallback on main: - wqkv_a 2048x7168, wo_b 7168x4096, wq_b 16384x1536 and indexer wq_b 8192x1536 at M=4 and 16 (serve decode M 3, 4, 12, 16) - MTP draft e_proj / h_proj 7168x7168 at M=1, 2, 4, 8, 16, 32, 64 (serve every decode M up to 64) All rows are split-K 1 and bitwise equal to the fallback output. Add the untuned shape list for the in-tree mxfp8 tuner. No existing row changes. Signed-off-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: Reviews: |
READY FOR HUMAN REVIEWPR Self-ReviewRevision Author self-review accepted. Readiness report: Author answers
Show the author's self-review answersBehavior and intentValidation and unverified scopePR-specific questionChanged op testsPR state
Copilot findings
Code and op-test hygiene
Required CI
Shadow mode reports what would block but keeps the status successful. Set |
Summary
Add 15 gfx1250 rows to
aiter/configs/model_configs/dsv4_a8w8_blockscale_mxscale_bpreshuffle_tuned_gemm.csvfor DeepSeek-V4-Pro TP4 decode GEMMs that run the untuned FlyDSL fallback on main. Also add the matching untuned shape list for the in-tree tuner. Scope is TP4, with and without MTP (multi-token prediction, 3 draft tokens). On main these GEMMs take 25 to 101 us per call; the rows take 3.6 to 9.4 us. In a replay of one decode step's dense FP8 GEMMs, that removes 11 to 12 ms per step at batch 4 and 16 without MTP, and at batch 1 and 4 with MTP. Config only. No existing row changes.Motivation
A gfx1250 mxfp8_128 bpreshuffle GEMM with no tuned row goes to
_flydsl_mxfp8_fallback_kernel. That function sorts by-tile_n, so at small M it pickst16x1024x128(t32x/t64x1024x128above M=16). An N=2048 GEMM then launches 2 workgroups on 256 CUs, and N=7168 launches 7. At these M the GEMM is bound by reading the weight from HBM, and a handful of workgroups can't do that.On main, #7018 covers M=1/2/8/64 for four of these shapes, and the Kimi-K3 table covers the two shared-expert shapes. What still falls back at decode M (ATOM's CUDA-graph batch sizes are 1 to 8, 16, 32, 48, 64, …):
wqkv_a2048x7168,wo_b7168x4096,wq_b16384x1536 and the indexerwq_b8192x1536, at M=3, 4, 12 and 16. Without MTP that is batch 3, 4 and 16. With MTP the target model verifies 4 tokens per sequence, so batch 1 runs at M=4, batch 3 at M=12 and batch 4 at M=16.e_projandh_proj(both 7168x7168), at every M. Each of the 3 draft steps runse_projat M = T tokens andh_projat M = 4T (4 hyper-connection copies). The first draft step has T = 4 x batch and the next two have T = batch. That's 6 calls per decode step.Changes
w_scale_block128x128, split-K 1:dsv4_a8w8_blockscale_mxscale_bpreshuffle_untuned_gemm.csv: the 15 shapes, in the format of the Kimi-K3 untuned list. The table merge skips untuned files.kernelIdis the kernel's index in main'skernels_list. Dispatch readskernelName. Many older gfx1250 rows in this file have ids from an earlier list; this PR doesn't touch them.usis the timing method below. Older rows in the file use other methods, sousis not comparable across row groups.bwis in GB/s over A + B + C bytes.errRatioischeckAllclose(ref, out, rtol=0.01, atol=0.01)against an fp32 dequantized reference.Why not fix the fallback heuristic instead: that's the better long-term fix, and I'll send it separately. It would cover other models, TP sizes and the data-parallel-attention shapes without rows. The shape of it: among split-K 1 sync kernels that fit, prefer the widest
tile_nthat still launches at least aboutcu_num / 2workgroups, instead of the widesttile_noutright. I haven't measured that rule. These rows are still worth having afterwards, because they are the measured best on main.Performance (if applicable)
Hardware: one gfx1250 GPU (MI455X, 256 CUs), aiter main
0de43e662b+ this PR, FlyDSL 0.4.2. The branch is rebased ontoa0c4b69c14, whose new commits change only docs and review tooling.Timing method, used for every number here and for the
uscolumn: 50 back-to-back launches in one HIP graph, with weights cycling through more than 1 GiB of copies so each launch reads its weight from HBM, as in decode graph replay. Each number is the median of 5 replays divided by 50, then the median over 3 fresh processes. Across all candidates the spread between processes (max minus min, over median) was 1.2% at the median, 6% at p90 and 11% at most. For the picked kernels it was at most 3.7%.Per-step dense FP8 GEMM time. The replay sums every dense FP8 linear in one decode step, each call running whatever main's lookup picks for that table:
wqkv_a,wq_b,wo_b, shared expertgate_up1536x7168 andw27168x768, plus the indexerwq_bon its 30 layers.e_proj/h_proj, for each of the 3 draft steps.Each value is the median of 3 runs; the runs spread by at most 1.2%. At batch 8 and 16 with MTP, the remaining gain is the draft GEMMs.
h_projthere runs at M=128 and 256, which stay untuned and take 21 and 24 us.How much of a step that is: I couldn't serve with main's aiter (see unverified scope). The closest measurement is the older ATOM image's aiter with a superset table that tunes all of these GEMMs. Its dense FP8 GEMM time at M=1 is about the same as main + PR's (1.56 ms). There, a decode step took:
So on main the GEMM time these rows remove is roughly 45 to 55% of the decode step at those four points. It is about 2 to 3% at MTP batch 8 and 16. That's an estimate, not a serving measurement on main.
Per row, before (main's fallback) vs after, at every decode M the row serves:
wqkv_at16x1024x128100t32x32x1024_nb45.8-6.0wo_bt16x1024x12859t16x64x1024_nb35.0-5.2wq_bt16x1024x12825.4t16x64x256_nb43.9-4.0wq_bt16x1024x12825.2t32x32x512_nb33.6-3.7e_proj,h_projt16x1024x128101.5t16x64x1024_nb36.9-7.1e_proj,h_projt32x1024x12894.5t32x32x1024_nb47.2-7.5e_proj,h_projt64x1024x12899.4t64x32x512_nb49.3-9.4(
nbNis the kernel's number of LDS pipeline buffers.)How the kernels were picked:
Split-K:
t64x128x256with split-K 4 is 11% and 8% faster (8.25 vs 9.28 us, 8.71 vs 9.43 us).Testing
The rows are keyed to gfx1250 and cu_num 256, so MI250X and MI300X don't read them. No default CI job runs gfx1250 GEMMs. In CI, only the CSV checks (
test_csv_validation,test_config_shape_collision) cover this change.Documentation
Tuned-table rows and a shape list only.
Dependencies
Breaking Changes
None.
Merge order:
PR Self-Review
Reviewed revision
061877a8caa2668a84bee2534520bfd5f3e18e61Behavior and intent
On gfx1250,
gemm_a8w8_blockscale_bpreshufflewith e8m0 128x128 scales finds a tuned row for the 15 keys above and the padded M they serve. Before, it ran_flydsl_mxfp8_fallback_kernel, whose widest-N tile launches 2 to 16 workgroups at these M and takes 25 to 101 us. The new kernels take 3.6 to 9.4 us. They are from the same FlyDSL kernel family, use split-K 1, and their output is bitwise equal to the fallback's on every tested input. No existing row, other table or code path changes.Validation and unverified scope
All on one gfx1250 GPU, aiter main
0de43e662b+ this PR, FlyDSL 0.4.2, in a container separate from the serving image:python op_tests/test_gemm_a8w8_blockscale.py --ck_preshuffle True --flydsl -d bf16 -m 1 2 3 4 5 6 7 8 12 16 20 24 28 32 48 64 -nk 7168,7168 2048,7168 7168,4096 8192,1536 16384,1536passed 160 cases (constant and uniform/auto inputs), max error ratio 0.0024. Real model activations were not used.get_mxscale_bpreshuffle_configreturns the new row for all 15 keys (AITER_LOG_TUNED_CONFIG=1). The public op returns the same bytes as the row's kernel called directly, and no "falling back" warning appears.find_padded_m_rowfor every (gfx, cu_num, N, K, w_scale_block) in the merged table at M = 1 to 1024 plus 8 larger M, with main's table and with this PR. Of 104,232 lookups, the only 104 that changed went from no row to a new row. Removing any one of the 15 rows changes the kernel for some decode M.ds_storeof the output tile, then a workgroup barrier, thentensor_store_from_lds. A standalone Gluon repro of that sequence (no aiter kernel involved) returns stale LDS in about 7e-7 of tile stores. The open fixes for other kernels are #7202, #7605 and #7632. A matching fix for these a8w8 GEMMs exists on a branch of mine but is not upstream. Stress check on main: each row's kernel ran 1,000,000 times at 16 (shape, M) points covering the 15 rows (100 calls per HIP graph replay, 10,000 replays) with a bf16 4096^3 matmul on a second stream every replay. Every output was compared bitwise on the device with the first one. Result: 0 mismatches in 16M calls and 2.6e9 workgroup tile stores. At the repro's rate about 1,800 would be expected. The old fallback tiles also gave 0 in 10M calls. This bounds the rate for these kernels and shapes under this load; it doesn't show the race can't happen.python3 -m unittest op_tests.tuning_tests.test_csv_validation op_tests.tuning_tests.test_config_shape_collision: 51 passed.python3 csrc/gemm_a8w8_mxfp8_bpreshuffle/gemm_a8w8_mxfp8_bpreshuffle_tune.py --mxscale_table -i aiter/configs/model_configs/dsv4_a8w8_blockscale_mxscale_bpreshuffle_untuned_gemm.csv -o <out.csv>runs on main. It picks the same kernel for 7 of the 15 rows. Seven others are split-K 1 kernels within 2.2% of ours under the method above, which is inside the noise. For 7168x7168 M=64 it picks the split-K 4 kernel this PR declines. Itsuscolumn comes from a different timer and is about 6 to 10% higher.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_MXSCALE_BPRESHUFFLE. That aiter predates [Config] [Perf] Tune DeepSeek-V4 M1/M2/M8/M64 FP8 GEMMs on gfx1250 #7018, and many of that table's kernels differ from these rows.Unverified:
ImportError: pa_mqa_logits_mxfp4_block_table_width), so the step savings above are a GEMM replay, not a TPOT measurement.wo_b7168x2048), data-parallel-attention shapes (6144x7168, 65536x1536, 7168x16384, 7168x3072), thewo_abatched GEMM, and 7168x7168 above M=64 (h_projat MTP batch 8 and 16).PR-specific question
Consumers: gfx1250, cu_num 256, w_scale_block 128x128, (N, K) in {(2048, 7168), (7168, 4096), (8192, 1536), (16384, 1536)} at M = 3 to 4 and 9 to 16, and (7168, 7168) at M = 1 to 64. Of these, ATOM decode reaches M = 3, 4, 12 and 16 for the layer GEMMs, and 1 to 8, 12, 16, 20, 24, 28, 32, 48 and 64 for the draft GEMMs. Evidence that nothing else moved: the before/after comparison of all 104,232 (arch, N, K, block, M) lookups in the merged table changed only those 104 lookups, each from no row to a new row. The diff adds lines and removes none. The 29 keys that #7018 and the Kimi-K3 table already cover were left alone. Under this timing method some of #7018's rows are slower than alternatives, but under #7018's own single-launch method several are faster, so I didn't replace them. #7552 (open) sends M < 16 to the M=16 row only when no power-of-two row exists. Every M < 16 here hits an exact or power-of-two row first, so lookups with both PRs equal this PR alone. If #7552 merges, the M=1/2/4/8 rows of 7168x7168 and the M=4 rows would become redundant with the M=16 rows, because they use the same kernel.
Changed op tests
None.