[Triton/Gluon] Move gluon gemm_a8w8_blockscale kernel into _gluon_kernels by vgokhale · Pull Request #4917 · ROCm/aiter

vgokhale · 2026-08-21T17:39:41Z

Move A8W8 blockscale gemm to the correct location.

Also deleted invalid configs as part of this move.

github-actions · 2026-08-21T17:40:11Z

🏷️ 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
`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 4917 --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.

The gfx950 gluon block-scale kernels move to aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8_blockscale.py, joining the gfx1250 variant that already lives under _gluon_kernels. The module holds only device code -- the two kernels, the gluon.jit helpers they call, the layout constexpr helpers and _SUPPORTED_TILES -- and its import path is hardcoded, so there are no arch branches inside it. All 15 device functions are unchanged. The host side merges into aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py, which already dispatches gluon for gfx1250, so gfx950 becomes a second gluon arch of the same entry point. The gfx950 gluon path no longer carries its own config loader. It used to build gemm/gluon/ paths by hand and skip config buckets whose tile the kernel cannot run; the shared get_gemm_config resolver already reads those same files, so the unrunnable buckets are dropped from the seven gfx950-GEMM-A8W8_BLOCKSCALE*.json files (25 in total, all at small M) and the path calls _get_config with backend="gluon" like every other path. 24 of those buckets have a tile outside _SUPPORTED_TILES, which is what the old picker skipped at runtime. The 25th, M_LEQ_128 in gfx950-GEMM-A8W8_BLOCKSCALE-N=4608-K=7168.json, has num_warps=2 while the kernel static-asserts NUM_WARPS == 4. Its tile is supported, so the old picker selected it and every M <= 128 at that shape failed to compile; dropping it lets those M fall through to M_LEQ_2048, which is tuned for the same BLOCK_SIZE_K with num_warps=4, and they now run. No unit test covers N=4608 K=7168, which is why this went unnoticed. Verified equivalent: the old picker on the old files and the resolver on the trimmed files pick the same config for 149 of 161 (M, file) pairs, and the 12 differences are exactly the N=4608 K=7168 M <= 128 cases above, where the old choice did not compile. The kernel keeps its supported-tile static_assert, so a config with an unrunnable tile still fails loudly. It also no longer carries its own split-K reduce kernel. The triton and gfx1250 gluon paths already reduce with the shared _gemm_splitk_reduce_kernel, and the two kernels compute the same sum, so gfx950 uses the shared one and the gluon reduce is deleted. Measured over 66 shapes, the shared kernel is within +1.6% worst case and faster on average: mean -1.0% on the shapes the tuned split-K buckets serve, -3.5% on larger kernel-bound shapes. gfx950 and gfx1250 share one launch site. Their kernels already agreed on the first 27 arguments, so the arch only selects the kernel and a small dict of arch-specific constexprs -- NUM_WARPS/NUM_STAGES for gfx950, warp_bases and NUM_BUFFERS for gfx1250 -- the same shape attention/fp8_mqa_logits uses. That also folds the previously duplicated gluon and triton launch bodies into one call, and puts gfx950 on the shared config, split-K and reduce path, so compute_splitk_params now derives its SPLITK_BLOCK_SIZE (verified to produce the same value for every reachable config and K) and skip_reduce works there. The arch dispatch is exhaustive rather than correct-by-elimination, so adding an arch to _GLUON_SUPPORTED_ARCHS without a launch branch raises instead of silently loading another arch's kernels. gfx950 has one blockscale kernel, which serves both kernel_type values; those two are currently identical code on gfx1250 as well, so kernel_type is validated once against the shared vocabulary. Arch handling is now three explicit tables in place of _is_gluon_available: _GLUON_SUPPORTED_ARCHS for gemm_a8w8_blockscale, _GLUON_PRESHUFFLE_ARCHS for gemm_a8w8_blockscale_preshuffle, which has no gfx950 gluon kernel, and _GLUON_DEFAULT_ARCHS for where gluon is the default backend. gfx1250 defaults to gluon and gfx950 to triton, which is what _is_gluon_available already resolved to -- it was gated on _GLUON_SUPPORTED_ARCHS, gfx1250-only at the time. gfx950 stays on triton by default because gluon was previously unreachable from this entry point there, so defaulting to it would move every existing caller off the tuned triton kernel; callers opt in with backend="gluon". Each entry point now defaults off its own table, so preshuffle keys on _GLUON_PRESHUFFLE_ARCHS rather than borrowing _GLUON_DEFAULT_ARCHS. Both tuples are gfx1250 today, so this is not a behaviour change. Removing _is_gluon_available also drops two dead pieces: a call that passed it an unsupported preshuffle= argument, and a second backend resolution later in the preshuffle body that could never run, since the first one leaves backend non-None. The old aiter.ops.triton.gluon.gemm_a8w8_blockscale import path keeps working through a _BACKWARD_COMPAT_MAP entry, and resolves to the shared entry point, so callers wanting gluon through it must pass backend="gluon". The unit test covers both import paths, picking one per case from a stable checksum of the case parameters; its gfx950 special case is gone, since gfx950 gluon is now reached the same way as gfx1250.

rahulbatra85 · 2026-08-25T17:39:10Z

 def parse_args(args: list[str] | None = None):
    parser = get_parser(kernel_name="A8W8 GEMM Blockscale")
    parser = add_argparse_ff(parser)
    parser.add_argument(


Instead of adding an arg gluon, I think it's better to add an arg "backend" which can be set to "gluon" or "triton" for now. Default can be triton or gluon depending on whatever makes sense for the op.

Benefit of this is if in future another backend is added, then no change would be needed in the benchmark. The backend arg from command line will be passed along to the op wrapper.

rahulbatra85 · 2026-08-25T17:42:37Z

        for shape in get_x_vals()
    ],
 )
 @pytest.mark.parametrize("backend", ["gluon", "triton"])


Also, add None to test the default backend which will be gluon or triton. I think this can be a separate test function because no one point is going through all shapes for this

)

vgokhale changed the base branch from main to vinayak/gemm_a8w8_to_gluon_kernels August 21, 2026 21:22

vgokhale force-pushed the vinayak/gemm_a8w8_blockscale_to_gluon_kernels branch from 68d159f to c751da9 Compare August 21, 2026 21:24

Boss2002n reviewed Aug 22, 2026

View reviewed changes

Comment thread aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8_blockscale.py

Boss2002n requested changes Aug 22, 2026

View reviewed changes

vgokhale force-pushed the vinayak/gemm_a8w8_blockscale_to_gluon_kernels branch 2 times, most recently from 026b755 to 3c74c1c Compare August 24, 2026 20:44

vgokhale force-pushed the vinayak/gemm_a8w8_to_gluon_kernels branch from 3e56e0d to 8cd7f4b Compare August 24, 2026 20:44

vgokhale force-pushed the vinayak/gemm_a8w8_blockscale_to_gluon_kernels branch from 3c74c1c to 9ebdd56 Compare August 24, 2026 20:57

vgokhale force-pushed the vinayak/gemm_a8w8_to_gluon_kernels branch from 8cd7f4b to e92e0ab Compare August 24, 2026 20:58

vgokhale marked this pull request as ready for review August 25, 2026 14:54

Boss2002n approved these changes Aug 25, 2026

View reviewed changes

vgokhale force-pushed the vinayak/gemm_a8w8_blockscale_to_gluon_kernels branch from 9ebdd56 to d1c841a Compare August 25, 2026 17:19

vgokhale force-pushed the vinayak/gemm_a8w8_blockscale_to_gluon_kernels branch from d1c841a to 59ccc1c Compare August 25, 2026 17:25

vgokhale force-pushed the vinayak/gemm_a8w8_to_gluon_kernels branch from e92e0ab to f59cdbb Compare August 25, 2026 17:26

rahulbatra85 reviewed Aug 25, 2026

View reviewed changes

vgokhale merged commit 66666f9 into vinayak/gemm_a8w8_to_gluon_kernels Aug 25, 2026

vgokhale deleted the vinayak/gemm_a8w8_blockscale_to_gluon_kernels branch August 25, 2026 18:58

vgokhale added a commit that referenced this pull request Aug 26, 2026

Move gluon gemm_a8w8_blockscale kernel into _gluon_kernels/gfx950 (#4917

52323b3

)

Navigation Menu

Uh oh!

[Triton/Gluon] Move gluon gemm_a8w8_blockscale kernel into _gluon_kernels - #4917

[Triton/Gluon] Move gluon gemm_a8w8_blockscale kernel into _gluon_kernels#4917
vgokhale merged 1 commit into
vinayak/gemm_a8w8_to_gluon_kernelsfrom
vinayak/gemm_a8w8_blockscale_to_gluon_kernels

vgokhale commented Aug 21, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 21, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rahulbatra85 Aug 25, 2026 •

edited

Loading

Uh oh!

rahulbatra85 Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

vgokhale commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Aug 21, 2026

🏷️ CI Guide

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rahulbatra85 Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rahulbatra85 Aug 25, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vgokhale commented Aug 21, 2026 •

edited

Loading

rahulbatra85 Aug 25, 2026 •

edited

Loading