[Triton/Gluon] [ASM] [HIP] MHA v4: support GQA, add gfx950 bf16, add gfx942 i8/fp8

[jcaraban]

⚠️ These MHA kernels are mainly tested & intended for Diffusion Inference workloads. However, mha_v4 gfx950 bf16 is virtually as accurate as mha_v3 at +100 TFlops 👌

Motivation

Extend the MHA v4 ASM path to support GQA without expanding K/V heads. This enables inputs such as Q [B, 32768, 64, 128] with K/V [B, 32768, 4, 128]. This PR also consolidates the related MHA v4 format, quantization, and gfx942/gfx950 kernel that were missing from the first PR --> #4627

Technical Details

• Support GQA ratios 1, 2, 4, 8, 16 in raw and packed MHA v4 APIs.
• Add BF16, native FP8, and MXFP8 MHA v4 dispatch paths.
  • Add new gfx942 FP8 and refresh INT8/FP8 code objects.
  • Add new gfx950 BF16 code object and manifest entry.
• Fix rotated FP8 preprocessing for singleton-head layouts.
• Improve ragged MXFP6/MXFP4 packing and layout validation.

The current scope remains dense, non-causal attention with head dimension 128 and BF16 output.

Test Plan

• Run the complete focused MHA v4 test suite: pytest -q op_tests/test_mha_v4.py
• Verify GQA output against an equivalent MHA invocation with repeated K/V heads.
• Exercise eager and torch.compile(fullgraph=True) paths across supported formats.
• Run the Sage benchmark with both a small smoke-test shape and the target long-context shape.
• Verify every gfx942/gfx950 manifest entry references an existing code object.
• Run Black, Ruff, and git diff --check.

Test Result

• MHA v4 suite: 72 passed ; Black and Ruff: passed
• Manifest code-object validation: passed
• GQA versus repeated-K/V MHA: bitwise equal

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[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
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 4967 --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.

[valarLip]

focus on you MHA v4 no changes to other part

[jcaraban]

sorry, I vibed myself out of my own files... reverted and added rotate_activation_hd128() into MHA v4. Note however aiter.ops.quant.rotate_activation may segfault for H=1. For example:

import torch
from aiter.ops.quant import rotate_activation

# BHSD permuted to BSHD before rotation
source = torch.empty((1, 1, 512, 128), device="cuda", dtype=torch.bfloat16)
input = source.permute(0, 2, 1, 3)
out = torch.empty_like(input)

print(f"shape={tuple(input.shape)}", flush=True)
print(f"stride={input.stride()}, contiguous={input.is_contiguous()}", flush=True)
print(f"stride(-2)={input.stride(-2)}, actual row spacing={input.size(-1)}", flush=True)

rotate_activation(out, input)
torch.cuda.synchronize()
print("unexpectedly survived", flush=True)

[valarLip]

thanks, @junhaha666 will take a look this one