Repository navigation
Conversation
- Accept caller-owned outputs in gemm_a8w8_bpreshuffle and gemm_a4w4 through an optional out= argument. - Above the bpreshuffle table, reuse the largest tuned row up to 3x its M; chunk FlyDSL launches at their row limit. - Add large-M A8W8 bpreshuffle and KDA prefill tuning with dispatch tests. - Replace 9 DSv3 N=2112 K=7168 bpreshuffle rows with faster Kimi-K3 rows that have the same keys. - Add Kimi-K3 BF16 tuning, keeping upstream rows and avoiding duplicate untuned keys. - Add gfx950 shared-expert FP4 GEMM configs; M buckets that were slower than DEFAULT.json use the default config.
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
github-actions
Bot
changed the title
There was a problem hiding this comment.
🟡 Changes recommended
The A8W8 caller-owned output accepts unsupported noncontiguous buffers, and new dispatch/API paths lack required coverage.
5 open findings
What changed in this PR
Adds Kimi-K3 GEMM tuning, caller-owned outputs, and large-M FP8 configuration reuse.
Changes:
- Adds gfx950 BF16, FP8, and FP4 tuning data.
- Adds optional
out=buffers to A8W8 and A4W4 GEMMs. - Reuses and chunks large-M FlyDSL configurations.
| File | Description |
|---|---|
op_tests/test_gemm_a8w8_bpreshuffle_large_m.py |
Tests large-M reuse and chunking. |
aiter/ops/triton/configs/gfx950/triton/gemm/gemm_afp4wfp4/GEMM-AFP4WFP4-N=7168-K=768.json |
Adds AFP4/WFP4 tuning. |
aiter/ops/triton/configs/gfx950/triton/gemm/gemm_afp4wfp4/GEMM-AFP4WFP4-N=1536-K=7168.json |
Adds AFP4/WFP4 tuning. |
aiter/ops/triton/configs/gfx950/triton/gemm/gemm_a16wfp4/GEMM-A16WFP4-N=7168-K=768.json |
Adds A16/WFP4 tuning. |
aiter/ops/triton/configs/gfx950/triton/gemm/gemm_a16wfp4/GEMM-A16WFP4-N=1536-K=7168.json |
Adds A16/WFP4 tuning. |
aiter/ops/gemm_op_a8w8.py |
Adds output buffers and large-M reuse. |
aiter/ops/gemm_op_a4w4.py |
Adds padded caller-owned outputs. |
aiter/configs/model_configs/kimik3_bf16_untuned_gemm.csv |
Adds BF16 tuning inputs. |
aiter/configs/model_configs/kimik3_bf16_tuned_gemm.csv |
Adds BF16 selected kernels. |
aiter/configs/model_configs/dsv3_a8w8_bpreshuffle_tuned_gemm.csv |
Removes superseded rows. |
aiter/configs/model_configs/a8w8_bpreshuffle_untuned_gemm_kimik3.csv |
Adds FP8 tuning inputs. |
aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv |
Adds FP8 selected kernels. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if out.shape != (m, n) or out.dtype != dtype or out.device != XQ.device: | ||
| raise ValueError( | ||
| f"out must be shape {(m, n)}, dtype {dtype}, device {XQ.device}; " | ||
| f"got shape {tuple(out.shape)}, dtype {out.dtype}, device {out.device}" |
| m_padded = (m + 31) // 32 * 32 | ||
| if out is None: | ||
| out = torch.empty((m_padded, n), dtype=dtype, device=A.device) | ||
| else: | ||
| out = _check_gemm_a4w4_out(out, m_padded, n, dtype, A.device) |
| found = rows.get( | ||
| _bpreshuffle_group(has_gfx, get_gfx(), get_cu_num(), n, k, q_dtype_w) | ||
| ) |
| bias: Tensor | None = None, | ||
| dtype: torch.dtype = dtypes.bf16, | ||
| check: bool = False, | ||
| out: Tensor | None = None, |
| apreshuffle: bool | None = False, | ||
| global_A_scale: Tensor | None = None, # NVFP4 per-tensor | ||
| global_B_scale: Tensor | None = None, # NVFP4 per-tensor | ||
| out: Tensor | None = None, # [ceil(M / 32) * 32, N] caller-owned output |
|
Closing in favor of a split, one concern per PR, as suggested in review:
|
|
Update on the large-M bpreshuffle reuse: it's dropped, so there won't be a fifth PR. I swept every gfx950 (N, K) group in the merged table where reuse would run: 43 groups, at 1.5x, 2x and 3x the largest tuned M. Reuse was slower than the default kernel in 9 groups, at every ratio. Examples: DSv3 N=2112 K=7168 at 0.17x, GLM-4.7 N=3072 K=5120 at 0.69x, DSv3 N=3072 K=1536 at 0.83x. Limiting it to gfx950 doesn't help, because the losses are on gfx950 too. Without reuse, the FlyDSL chunking isn't needed either. Instead, #6410 now has tuned rows for Kimi-K3 N=7168 K=3584 at M=288..1024. Against main, that's 1.05x-1.72x (geomean 1.19x) across M=264..1024, compared with 0.95x-1.46x from the reuse path. The sweep also turned up an existing issue on main: the FlyDSL 8-wave preshuffle kernel ( |
Motivation
Kimi-K3 runs FP8 bpreshuffle, BF16 and MXFP4 GEMM shapes on gfx950 that have no tuned config, so they use default kernels. FP8 calls with M above the tuned table also use the default CK kernel. Some callers need the GEMM to write into a buffer they own.
Technical Details
gemm_a8w8_bpreshuffleandgemm_a4w4accept an optionalout=tensor. The shape, dtype and device are checked.gemm_a4w4,outmust be[ceil(M/32)*32, N], because the kernels write whole 32-row tiles. The result is the first M rows.out=is not supported on gfx1250.gemm_a8w8_bpreshufflereuses the row with the largest tuned M. This only applies up to 3x that M; larger M use the default kernel. A row tuned for a small M became slower than the default kernel at much larger M.PRESHUFFLE_M_MAX).a8w8_bpreshuffle_tuned_gemm_kimik3.csv: 87 rows for N=7168 K=3584, N=6336 K=7168, N=2112 K=7168 and N=6144 K=7168.kimik3_bf16_tuned_gemm.csv: 217 rows.gemm_afp4wfp4andgemm_a16wfp4at N=1536 K=7168 and N=7168 K=768. M buckets that were slower thanDEFAULT.jsonuse the default config.dsv3_a8w8_bpreshuffle_tuned_gemm.csv. The new Kimi-K3 rows have the same keys and a lower tuned time. Keeping both would make the config merge fail with a duplicate-shape error.op_tests/test_gemm_a8w8_bpreshuffle_large_m.pytests the reuse path, the chunking, the 3x limit, and that M between tuned rows uses the default kernel.Test Plan
pytest op_tests/test_gemm_a8w8_bpreshuffle_large_m.pypython3 -m unittest op_tests.tuning_tests.test_config_shape_collisionupstream/main(80f1165) on one MI355X. The same scripts run on both trees, back to back, for 2 or 3 rounds. Timing usesrun_perftest, and each output is checked against a torch reference withcheckAllclose. Shapes:gemm_a8w8_bpreshuffle: new rows, M above the table, and the DSv3 N=2112 shapes.tgemm.mm: new rows.out=for both GEMMs.Test Result
Setup: MI355X (gfx950), ROCm 7.2, torch 2.11.0+rocm7.2, Triton 3.7.0, FlyDSL 0.3.4.1. Base is
mainat 80f1165. Base and head ran back to back on the same GPU with the same script.Summary, this PR vs main (GPU kernel time, best of 2-3 rounds):
tgemm.mm, new rowsgemm_a16wfp4, M <= 256gemm_a16wfp4, M > 256gemm_afp4wfp4, M <= 256gemm_afp4wfp4, M > 256Shapes between 0.97x and 0.99x run the same kernel or config as main, so those gaps are run-to-run noise.
gemm_a8w8_bpreshuffle(head vs base):tgemm.mm: geomean 1.06x (0.98x to 1.27x). One tuned row (M=6043 N=7168 K=1536) was removed. M=6043 now matches the M=6086 row, which uses the same kernel.gemm_a16wfp41.01x to 3.07x at M<=256 and 1.00x above.gemm_afp4wfp40.98x to 2.30x. Shapes below 1.00x use the same config as base.gemm_a8w8_bpreshufflewith and withoutout=matches eager.gemm_a4w4without=matches eager underaot_eager. Under inductor,gemm_a4w4fails onupstream/maintoo, even withoutout=. That is not caused by this change.Submission Checklist