Skip to content

[TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x - #4563

Merged
cagrikymk merged 7 commits into
mainfrom
cagri/fp8_mqa_opt
Aug 5, 2026
Merged

[TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x#4563
cagrikymk merged 7 commits into
mainfrom
cagri/fp8_mqa_opt

Conversation

@cagrikymk

@cagrikymk cagrikymk commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Motivation

GLM 5 series use number of q-heads = 32, but this kernel was originally tuned for 64 heads.

This PR improves performance for GLM 5 via kernel updates + tuning.

Technical Details

Now the kernel supports handling 2 tokens together to increase the work per WG (since 32*2=64). While there is still some performance gap between GLM 5 config and dsv4 one, this PR closes that gap quite a bit.

Test Plan

Existing tests are extended to include GLM 5 case as well.

Test Result

Tests pass.

Performance

While DSV4 config (num_heads=64) shows some regression when # tokens~= 4k, specifically 1k/1k case, those cases appear rarely during inference as batching typically goes for >= 8k tokens.

Results are from MI355.

Baseline and new column format: {time} us / TFLOPs

num_heads = 32

ISL=OSL batch rows baseline new speedup
1024 1 1,024 15.0 µs / 287 10.7 µs / 401 1.40x
1024 4 4,096 27.2 µs / 633 27.6 µs / 624 0.99x
1024 8 8,192 50.4 µs / 682 39.1 µs / 879 1.29x
1024 16 16,384 93.3 µs / 738 71.3 µs / 965 1.31x
2048 1 2,048 31.8 µs / 541 25.4 µs / 678 1.25x
2048 2 4,096 46.3 µs / 743 44.8 µs / 767 1.03x
2048 4 8,192 85.6 µs / 803 63.0 µs / 1,091 1.36x
2048 8 16,384 162.7 µs / 845 117.8 µs / 1,167 1.38x
4096 1 4,096 83.2 µs / 827 81.2 µs / 847 1.02x
4096 2 8,192 152.5 µs / 902 106.5 µs / 1,291 1.43x
4096 4 16,384 293.3 µs / 937 200.2 µs / 1,373 1.47x
8192 1 8,192 288.4 µs / 953 194.1 µs / 1,417 1.49x
8192 2 16,384 571.7 µs / 962 376.6 µs / 1,460 1.52x
geomean 1.314x

num_heads = 64

ISL=OSL batch rows baseline new speedup
1024 1 1,024 18.9 µs / 456 18.6 µs / 462 1.02x
1024 4 4,096 34.0 µs / 1,012 40.5 µs / 850 0.84x
1024 8 8,192 61.2 µs / 1,123 62.6 µs / 1,099 0.98x
1024 16 16,384 115.6 µs / 1,190 110.5 µs / 1,245 1.05x
2048 1 2,048 39.3 µs / 876 38.9 µs / 885 1.01x
2048 2 4,096 56.5 µs / 1,217 59.5 µs / 1,155 0.95x
2048 4 8,192 104.8 µs / 1,312 100.9 µs / 1,363 1.04x
2048 8 16,384 193.4 µs / 1,422 187.7 µs / 1,465 1.03x
4096 1 4,096 99.6 µs / 1,380 97.6 µs / 1,408 1.02x
4096 2 8,192 178.3 µs / 1,542 178.6 µs / 1,539 1.00x
4096 4 16,384 340.1 µs / 1,617 329.0 µs / 1,672 1.03x
8192 1 8,192 330.9 µs / 1,662 313.7 µs / 1,753 1.05x
8192 2 16,384 651.6 µs / 1,688 613.9 µs / 1,791 1.06x
geomean 1.009x

@cagrikymk
cagrikymk requested a review from a team August 4, 2026 21:06
@github-actions

github-actions Bot commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

🏷️ 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: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 4563 --add-label <label>

vgokhale
vgokhale reviewed Aug 4, 2026
View reviewed changes
Comment thread aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py
@zufayu
zufayu requested a review from vgokhale August 5, 2026 02:48
@cagrikymk
cagrikymk merged commit aded0f8 into main Aug 5, 2026
56 checks passed
@cagrikymk
cagrikymk deleted the cagri/fp8_mqa_opt branch August 5, 2026 23:06
@vpietila-amd vpietila-amd mentioned this pull request Aug 6, 2026
Open
1 task
This was referenced Aug 8, 2026
Open
Open
Open
@amd-bot amd-bot mentioned this pull request Aug 16, 2026
Open
Ishiki-Iroha added a commit to Ishiki-Iroha/aiter that referenced this pull request Aug 21, 2026
@Ishiki-Iroha
perf(mi355x): backport GLM-5.x FP8 MQA logits tuning (ROCm#4563)
459480a
@jiejingzhangamd jiejingzhangamd mentioned this pull request Aug 28, 2026
Merged
xiaobochen-amd pushed a commit to xiaobochen-amd/aiter that referenced this pull request Aug 28, 2026
@jiejingzhangamd @cagrikymk
Jiejing Zhang @cursoragent
[Bugfix][Triton][gfx950] fp8_mqa_logits: gate buffer ops on the int32…
7bc2347 
… offset, not tensor bytes (#4)

* [TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (ROCm#4563)

(cherry picked from commit aded0f8)

Also carries the two-line `block_m = 1` for the non-gfx950 branch from ROCm#4246.
This commit routes the launch grid through block_m but only defines it under
`arch == "gfx950"`, so without those lines every other arch raises NameError at
the kernel launch. Upstream ran with that hole from ROCm#4563 until ROCm#4246.

* [Bugfix][Triton][gfx950] fp8_mqa_logits: gate buffer ops on the int32 offset, not tensor bytes

Prefill shapes whose fp32 logits pass 2 GiB abort the AMDGCN backend at
JIT time (Sequence.h:275 "Begin must be less or equal to End"). BLOCK_M=2
is selected for seq_len > 4096 but only compiles with buffer stores, and
the buffer-store gate switches off at 2 GiB, so the two combine into a
hard crash. For GLM-5.x that is any chunked prefill of 8192 tokens
against a context past 65,536.

The gate had the wrong unit. Buffer ops address through a 32-bit offset,
but the kernel re-bases the pointer per row and per KV tile before each
access, so that offset never has to span the tensor. What must fit in
int32 is the largest element offset the kernel forms, because the row
strides stay 32-bit on the buffer path and are only widened to int64 on
the fallback path. Counting bytes rather than elements made the limit 4x
too tight for an fp32 output.

Measured on MI355X, 32 heads x 128 head_dim, against the plain-store
path. Output is bit-identical in every case (max rel err 0.00e+00):

  s_q     s_k      logits    plain store   buffer store   speedup
  8192    65536    2.00 GiB      3.557 ms       2.941 ms     1.21x
  8192    95457    2.91 GiB      6.105 ms       4.577 ms     1.33x
  8192   131072    4.00 GiB      9.812 ms       6.489 ms     1.51x
  16384  131072    8.00 GiB     21.738 ms      12.879 ms     1.69x
  8192   262144    8.00 GiB     25.710 ms      14.184 ms     1.81x

The new boundary is exact rather than approximate: 8192x262144 and
16384x131072 both place the largest offset at exactly INT32_MAX and are
bit-correct, while 16384x139264 (1.06 x 2^31) and 32768x131072 (2^32)
fault.

End to end, a 21-layer GLM-5.x indexer prefill of 9,695 new tokens
against a 96,960-token context at chunk 8192 drops from 150.2 ms to
118.4 ms.

The existing cases top out at s_q=1024, s_k=1560, four orders of
magnitude below the gate, which is why nothing caught this. Added
8192x65664 and 8192x98304, which crash the process on current main.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Mehmet Cagri <mehmet.kaymak@amd.com>
Co-authored-by: Jiejing Zhang <jiejizha@smci355-ccs-aus-n02-29.prov.aus.ccs.cpe.ice.amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@zhuyuhua-v zhuyuhua-v mentioned this pull request Aug 28, 2026
Draft
1 task
zhuyuhua-v added a commit that referenced this pull request Aug 28, 2026
@zhuyuhua-v
Revert "[TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (#…
893479b 
…4563)"

This reverts commit aded0f8.
JohnQinAMD pushed a commit to xiaobochen-amd/aiter that referenced this pull request Aug 28, 2026
@jiejingzhangamd @cagrikymk
Jiejing Zhang @cursoragent yanyuan.qin@amd.com
[Bugfix][Triton][gfx950] fp8_mqa_logits: gate buffer ops on the int32…
da75436 
… offset, not tensor bytes (#4)

* [TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (ROCm#4563)

(cherry picked from commit aded0f8)

Also carries the two-line `block_m = 1` for the non-gfx950 branch from ROCm#4246.
This commit routes the launch grid through block_m but only defines it under
`arch == "gfx950"`, so without those lines every other arch raises NameError at
the kernel launch. Upstream ran with that hole from ROCm#4563 until ROCm#4246.

* [Bugfix][Triton][gfx950] fp8_mqa_logits: gate buffer ops on the int32 offset, not tensor bytes

Prefill shapes whose fp32 logits pass 2 GiB abort the AMDGCN backend at
JIT time (Sequence.h:275 "Begin must be less or equal to End"). BLOCK_M=2
is selected for seq_len > 4096 but only compiles with buffer stores, and
the buffer-store gate switches off at 2 GiB, so the two combine into a
hard crash. For GLM-5.x that is any chunked prefill of 8192 tokens
against a context past 65,536.

The gate had the wrong unit. Buffer ops address through a 32-bit offset,
but the kernel re-bases the pointer per row and per KV tile before each
access, so that offset never has to span the tensor. What must fit in
int32 is the largest element offset the kernel forms, because the row
strides stay 32-bit on the buffer path and are only widened to int64 on
the fallback path. Counting bytes rather than elements made the limit 4x
too tight for an fp32 output.

Measured on MI355X, 32 heads x 128 head_dim, against the plain-store
path. Output is bit-identical in every case (max rel err 0.00e+00):

  s_q     s_k      logits    plain store   buffer store   speedup
  8192    65536    2.00 GiB      3.557 ms       2.941 ms     1.21x
  8192    95457    2.91 GiB      6.105 ms       4.577 ms     1.33x
  8192   131072    4.00 GiB      9.812 ms       6.489 ms     1.51x
  16384  131072    8.00 GiB     21.738 ms      12.879 ms     1.69x
  8192   262144    8.00 GiB     25.710 ms      14.184 ms     1.81x

The new boundary is exact rather than approximate: 8192x262144 and
16384x131072 both place the largest offset at exactly INT32_MAX and are
bit-correct, while 16384x139264 (1.06 x 2^31) and 32768x131072 (2^32)
fault.

End to end, a 21-layer GLM-5.x indexer prefill of 9,695 new tokens
against a 96,960-token context at chunk 8192 drops from 150.2 ms to
118.4 ms.

The existing cases top out at s_q=1024, s_k=1560, four orders of
magnitude below the gate, which is why nothing caught this. Added
8192x65664 and 8192x98304, which crash the process on current main.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Mehmet Cagri <mehmet.kaymak@amd.com>
Co-authored-by: Jiejing Zhang <jiejizha@smci355-ccs-aus-n02-29.prov.aus.ccs.cpe.ice.amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
zhuyuhua-v added a commit that referenced this pull request Aug 31, 2026
@zhuyuhua-v
Revert "[TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (#…
150ebf3 
…4563)"

This reverts commit aded0f8.
zhuyuhua-v added a commit that referenced this pull request Aug 31, 2026
@zhuyuhua-v
Revert "[TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (#…
34d0afc 
…4563)"

This reverts commit aded0f8.
zhuyuhua-v added a commit that referenced this pull request Sep 1, 2026
@zhuyuhua-v
Revert "[TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (#…
142afca 
…4563)"

This reverts commit aded0f8.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@vgokhale vgokhale Awaiting requested review from vgokhale

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@cagrikymk @vgokhale