[Triton/Gluon] [gfx950] add an optimized prefill fp8_mqa_logits for H64D128

[Dewei-Wang-sh]

Dispatched on gfx950 + seq_len >= 4096
Bench shows ~15% perf lift on 4Kx4K seqlen and dsv4 8K seqlen on mi350

config	live	FLOP	base	H64D128	d TFLOP/s
1x4096x4096, ratio 1	0.0084G (50.0%)	0.137 T	105.7 us / 1301 TF/s	92.8 us / 1482 TF/s	+13.9%
4x8192, ratio 4 (dsv4)	0.0335G (12.5%)	0.550 T	396.0 us / 1388 TF/s	349.7 us / 1572 TF/s	+13.3%

[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 5048 --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

This PR adds a gfx950-specific optimized Gluon FP8 MQA logits kernel variant for the H=64, D=128, seq_len≥4096 regime, and wires it into the existing fp8_mqa_logits dispatch to deliver the reported long-context throughput gains.

Changes:

• Adds a new _gluon_fp8_mqa_logits_kernel_H64D128 kernel optimized around BLOCK_M=4 and triple buffering.
• Extends the head-reduction planning logic to support an additional leading “row” axis needed by the tiled M dimension.
• Updates the Python wrapper dispatch to select the new kernel on gfx950 when the required features and shape constraints are met.

Reviewed changes

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

File	Description
aiter/ops/triton/attention/fp8_mqa_logits.py	Adds import + dispatch routing to the new gfx950 H64D128 optimized Gluon kernel for seq_len≥4096.
aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py	Implements the new H64D128 kernel and adapts the reduction plan to handle BLOCK_M>1.

---

💡 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 2 out of 2 changed files in this pull request and generated no new comments.

Suppressed comments (4)

aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py:971

• New Gluon entry kernel _gluon_fp8_mqa_logits_kernel_H64D128 is launchable but does not set a config-aware repr. The project’s kernel checklist requires make_kernel_repr(...) + @gluon.jit(repr=...) so compiled variants are identifiable in traces/logs (see aiter/ops/triton/README.md:304-305).

@gluon.jit
def _gluon_fp8_mqa_logits_kernel_H64D128(
    Q_ptr,  # fp8e4m3 [seq_len, NUM_HEADS, HEAD_SIZE]

aiter/ops/triton/attention/fp8_mqa_logits.py:217

• # FIXME: unify later leaves an explicitly unfinished dispatch path in a production wrapper. Please either (a) link to a tracking issue and clarify what remains to be unified, or (b) replace with a more specific TODO explaining the invariants that prevent unification today.

        # gfx950 H64D128 variant:
        # FIXME: unify later
        if (

aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py:967

• # FIXME: Unifying the two with full coverage is an unresolved note in a new kernel section. Please convert this to a TODO with a tracking issue (or remove it) so it’s clear what remaining correctness/perf gaps exist and how unification will be validated.

# Separate kernel rather than a BLOCK_M branch in the one above: the loop
# shapes genuinely differ (peeled double-buffer vs uniform triple-buffer with
# refill-at-top) and the store paths. Shares this module's reduction plan.
# FIXME: Unifying the two with full coverage

aiter/ops/triton/attention/fp8_mqa_logits.py:224

• This PR adds a new dispatch path for gfx950 H=64, D=128 with seq_len >= 4096, but the existing unit test matrix for fp8_mqa_logits only covers s_q <= 1024 (see op_tests/triton_tests/attention/test_fp8_mqa_logits.py). That means the new kernel can regress correctness/perf without being exercised in CI. Please add at least one correctness smoke test that triggers this branch (e.g., seq_len=4096, num_heads=64, head_dim=128, with a small seq_len_kv and a narrow window to keep the reference tractable).

        if (
            arch == "gfx950"
            and _gluon_fp8_mqa_logits_kernel_H64D128 is not None
            and ASYNC_COPY_SUPPORTS_DISTRIBUTED
            and num_heads == 64
            and head_size == 128
            and seq_len >= 4096
        ):

[Dewei-Wang-sh]

close since this has the same perf as #5216, yet pr5216 is more general