[CK] Add tuned CK a8w8 blockscale GEMM configs for Gemma-4-31B FP8-block

[mustafayildirim]

Summary

Add per-model tuned GEMM configurations for google/gemma-4-31B-it served with FP8 block-scale quantization (RedHatAI/gemma-4-31B-it-FP8-block).

The tuned configs cover 85 (M,N,K) shapes — 7 unique weight shapes × 13 batch sizes (M=1..16384) — including:

Weight shape	Role
N=5376, K∈{8192, 16384, 21504}	q/k/v/o projections
N∈{16384, 20480, 43008}, K=5376	gate/up/down MLP
N=262144, K=5376	lm_head — previously untuned entirely

Notable wins:

• Decode (M=1): 37–46% faster configs for the small projection GEMMs (N=5376) vs the default kernelId=8 splitK=0
• Prefill M=16384: previously untuned batch size

Tuning details

• Hardware: gfx942 (MI325X, 304 CUs)
• Script: csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py
• Args: --libtype ck --splitK --shape_grouped --warmup 5 --iters 101
• Backend: CK only (85 shapes). CKTile was tuned and benchmarked (see verification), but the kernels are not compiled into the stock vllm/vllm-openai-rocm:v0.28.0 image — including them causes RuntimeError at vLLM startup on that image. The 3 CKTile shapes fall back to default CK (no regression); they can be added back once a vLLM image ships an aiter build that includes this CSV in its prebuild manifest.

Correctness

All 85 entries have errRatio ≤ 0.01 (max = 0.0032, atol/rtol = 0.01).

Multi-target build fix

Also fixes a pre-existing bug in candidate_kernels_by_name (both gemm_a8w8_blockscale_cktile_instance.py and gemm_a8w8_bpreshuffle_cktile_common.py) where the name registry was built from only the current get_gfx() arch. In a multi-target build (GPU_ARCHS=gfx942;gfx950), this caused codegen to crash on cross-arch kernel names. Fix: make kernels_by_name the union of all arch kernel lists (using list(values) + list(values) to avoid integer-key collisions). Reported by @yzhou103.

Branch state

Rebased onto upstream main (Sep 2026). The diff is 3 files: the tuned CSV plus the multi-arch name-registry fix in the two cktile instance/common files. A previously included FlyDSL LDS override was dropped as obsolete after upstream #5116 removed aiter/ops/flydsl/utils.py.

Remaining uncovered shapes

6 shapes have no CK tuning and fall back to the default CK config (same as current behavior, no regression):

M	N	K	Note
4096, 8192, 16384	262144	5376	lm_head — no existing CK kernel supports these tile combos
1024	5376	16384	tuned successfully with CKTile (errRatio 0.0078), removed with the other CKTile rows (stock image lacks the kernels)
16384	5376	21504	tuned with CKTile (0.0), removed with CKTile rows
16384	43008	5376	tuned with CKTile (0.0), removed with CKTile rows

The 3 CKTile rows can be restored once a vLLM image ships an aiter build whose prebuild manifest includes this CSV. The 3 lm_head shapes need new CK kernel instances (C++ CK template work) — tracked as a follow-up.

How it works

The file is auto-discovered by aiter's model_configs/ glob — no env var or code change needed. Any deployment using aiter with a Gemma-4-31B FP8-block model picks up these tuned configs automatically.

Verification

• amd2 (MI325X / gfx942): vLLM (v0.28.0 image) + this CSV + attention fix: 32k-in/1k-out warm latency 16.26s → 16.20s vs 85-entry; tuned decode configs give the large win vs baseline (~83s with profiler overhead).
• amd4 (MI350X / gfx950): CSV loads, vLLM serves correctly. gfx942-tuned CK entries are not applicable to the gfx950 registry (expected); gfx950-specific tuning is a follow-up.
• CKTile: verified that rebuilding module_gemm_a8w8_blockscale_cktile.so from this branch compiles the previously-missing 128x128x128_1x4x1_16x16x64 BPC≥2 variants, after which the full 88-entry CSV (85 CK + 3 CKTile) loads and runs without errors.

---

AI assistance was used. All tuned entries were verified for numerical correctness by the tuner.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5062 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[Copilot]

Pull request overview

Adds a per-model tuned CK GEMM configuration table for running Gemma-4-31B FP8 block-scale (a8w8_blockscale) workloads on gfx942, enabling the runtime to auto-merge and use these tuned entries via the existing model_configs/ discovery mechanism.

Changes:

• Introduces a new tuned CSV for a8w8_blockscale GEMM covering key Gemma-4-31B shapes (including lm_head and large prefill/decode batches).
• Provides kernelId/splitK selections and measured perf metadata (us, tflops, bw, errRatio) for gfx942 MI325X-class GPUs.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

aiter/configs/model_configs/a8w8_blockscale_tuned_gemm_gemma4_31b.csv:87

• This row has errRatio=0.0242, which exceeds the PR description’s stated correctness bounds (atol=0.01, rtol=0.01). If errRatio is the metric used for that claim, either re-tune this shape with a stricter error threshold or adjust the PR description to reflect the actual error bound being enforced.

gfx942,304,1024,5376,16384,cktile,2,3,690.845,a8w8_blockscale_cktile_128x128x128_1x4x1_16x16x64_intrawave_0x1x0_2,261.11,167.72,0.0242

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.

[dwiddows]

The GEMM tunings look sensible, though I haven't tested them in detail. (I've used AI analysis to check the config shape coverage, lookup reachability, duplicate/collision safety, hardware/model dimensions, and description.)

Remaining questions include whether the CKTile errRatio=0.0242 is acceptable, and whether the performance summary is sufficient already. @valarLip @yzhou103 Please could you review these points, and advise whether anything else is needed for this PR?

[yzhou103]

build aiter wheels failed.

The 3 CKTile rows at the bottom trigger a pre-existing bug in gemm_a8w8_blockscale_cktile_instance.py:155-162: candidate_kernels_by_name is built from only the single arch that get_gfx() returns, but build_tune_dict filters CSV rows for all build targets via get_build_targets(). In a multi-target build (GPU_ARCHS=gfx942;gfx950), get_gfx() picks gfx950 (last entry), so the name registry only has gfx950 kernels (K_Warp_Tile=128). The gfx942 CKTile rows reference 16x16x64 kernels that don't exist in the gfx950 registry → codegen crashes → lookup.h never generated → compilation fails.

Same bug exists in cktile_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_cktile_common.py:415.

Suggested fix (either as a prerequisite PR or included here): make kernels_by_name the union of all arch kernel lists instead of only the get_gfx() arch. candidate_kernels_cktile_dict (used for codegen) stays single-arch — only the name lookup dict needs the union.

[yzhou103]

The GEMM tunings look sensible, though I haven't tested them in detail. (I've used AI analysis to check the config shape coverage, lookup reachability, duplicate/collision safety, hardware/model dimensions, and description.)

Remaining questions include whether the CKTile errRatio=0.0242 is acceptable, and whether the performance summary is sufficient already. @valarLip @yzhou103 Please could you review these points, and advise whether anything else is needed for this PR?

About the errRatio, errRatio=0.0242 is within the tuner's default threshold (0.05), so it's acceptable at the per-op level. Whether it's acceptable end-to-end depends on the model owner's accuracy requirements. If not sure, was a splitK=1 or splitK=2 alternative tried for this shape (M=1024, N=5376, K=16384)? If a lower-splitK config exists with comparable performance but lower error, it should be preferred.

[mustafayildirim]

Thanks @copilot for the review. Both points addressed:

1. errRatio=0.0242 on line 87 (CKTile entry): The PR description has been updated to document this explicitly. CK entries use atol=0.01/rtol=0.01 (max errRatio=0.0032). CKTile entries pass the tuner default tolerance (errRatio ≤ 0.05). The 0.0242 entry is expected for FP8 block-scale with splitK=3.

2. Description said --libtype ck but CSV includes cktile entries: Description updated. The 85 CK entries were tuned with --libtype ck, and the 3 CKTile entries (covering shapes CK could not support) were tuned separately with --libtype cktile. Total is 88 shapes (85 CK + 3 CKTile), which now matches the description.

[mustafayildirim]

Thanks @yzhou103 for the detailed review! Both issues addressed in the latest push:

1. Multi-target build fix: Fixed candidate_kernels_by_name in both gemm_a8w8_blockscale_cktile_instance.py and gemm_a8w8_bpreshuffle_cktile_common.py to be the union of all arch kernel lists instead of only the get_gfx() arch. The candidate_kernels_cktile_dict (used for codegen) stays single-arch — only the name lookup dict is now the union.

2. errRatio on (M=1024, N=5376, K=16384): Re-tuned with --profile_file to get all candidates. splitK=2 is only 5% slower than splitK=3 (675us vs 642us) with errRatio=0.0078 (within 0.01). Updated the CSV to use splitK=2. All 88 entries now have errRatio ≤ 0.01.

splitK	us	errRatio
0	871.7	0.0
1	765.5	0.0
2	675.2	0.0078 ← now used
3	642.3	0.0242

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (2)

csrc/cktile_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_cktile_common.py:415

• Merging the expanded kernel dicts with {**a, **b} drops kernels because both arches use overlapping integer IDs (e.g., 0, 1, ...). That means kernels_by_name may still miss names from the earlier arch and can reproduce the multi-target crash this change is trying to fix. Build the name registry from the concatenated values instead of merging dicts by key (and reuse the already-expanded kernels_list to avoid extra work).

_all_bpreshuffle_kernels = {**expand_blockpercu(kernels_list_942), **expand_blockpercu(kernels_list_950)}
kernels_by_name = {v.name: v for v in _all_bpreshuffle_kernels.values()}

csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_cktile_instance.py:153

• {**expand_blockpercu(kernels_list_942), **expand_blockpercu(kernels_list_95x)} merges by integer ID, so keys collide (both dicts start at 0) and one arch’s kernels are silently dropped. That defeats the purpose of including all arches for multi-target builds. Build candidate_kernels_by_name from the concatenated values instead (and reuse candidate_kernels_cktile_dict for the current arch).

_all_cktile_kernels = {**expand_blockpercu(kernels_list_942), **expand_blockpercu(kernels_list_95x)}
candidate_kernels_by_name = {v.name: v for v in _all_cktile_kernels.values()}

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

[yzhou103]

please fix code style check, so it can run ci tests.

[dwiddows]

Here's a merge that adds the linebreaks from the black formatter, which should fix the code style errors and unblock the ci tests: mustafayildirim/aiter@mustafa/gemma4-31b-blockscale-tuned-gemm...dwiddows:aiter:dwiddows/pr5062-black-wraps

[dwiddows]

Also I think the FlyDSL changes in #5089 are now unnecessary since #5116 was merged, and removed aiter/ops/flydsl/utils.py.

For this PR, it should be possible to remove aiter/ops/flydsl/utils.py and op_tests/flydsl_tests/test_flydsl_utils.py, and avoid these merge conflicts.

[mustafayildirim]

Fixed both items:

1. Black style: Cherry-picked @dwiddows's dwiddows/pr5062-black-wraps commit (aefedef) — wraps the two kernels_by_name dict comprehensions. black --check now passes on both files.

2. Dropped the FlyDSL changes: Removed aiter/ops/flydsl/utils.py and op_tests/flydsl_tests/test_flydsl_utils.py — correct call, upstream [Triton/Gluon] [HIP] [CK] Remove obsolete availability helpers #5116 removed that file entirely and replaced the LDS checks with chip_info.get_lds_capacity_bytes(). Also merged latest upstream main (99 commits), so conflicts are resolved (mergeable: MERGEABLE).

The PR diff is now minimal: the Gemma-4-31B tuned CSV (85 CK entries) + the multi-arch kernels_by_name build fix in the two cktile instance/common files. @dwiddows thanks for the wrap-up branch!

[Copilot]

🟢 Approval recommended

The changes are consistent with existing tuned-config discovery/format and the multi-target kernelName registry fix aligns with how build_tune_dict validates CSV rows across multiple GPU_ARCHS targets.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Lite

[yzhou103]

LGTM

[mustafayildirim]

Rebased onto latest main (39 new upstream commits, including #5280 which adds the rowcol_wp_v2 pipeline to the bpreshuffle cktile registry). Conflict resolved by keeping upstream's append_expanded composition for the gfx950 pipeline list while preserving the multi-arch kernels_by_name union fix — the union now also covers the new rowcol_wp_v2 kernels. black --check passes.

[Copilot]

🟡 Changes recommended

The shipped tuned CSV appears to be missing at least one claimed-covered (M,N,K) shape, creating a discrepancy with the PR description and forcing a fallback to default dispatch for that shape.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Architecture-specific kernel changes and extensive tuning data warrant final human validation.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Lite