[triton] Optimized Unified Attention for Gemma-4-31b

[a-sidorova]

Motivation

On AMD CDNA4 (MI355X / gfx950), the vLLM Triton unified-attention backend was
significantly faster than AITER's ROCM_AITER_UNIFIED_ATTN for google/gemma-4-31B-it,
and AITER's full-attention decode crashed outright with OutOfResources. Gemma-4
interleaves two large-head attention layers that fall outside the range the AITER configs
were tuned for (head_dim ≤ 128):

Layer	head_dim	GQA (q/kv)	block_size	Window
full (global)	512	32 / 4 (nqpkv=8)	64	global
sliding	256	32 / 16 (nqpkv=2)	32	1024

This PR retunes select_2d_config / select_3d_config for these shapes. Changes are
config-level only (BLOCK_M, TILE_SIZE, num_warps, num_stages, MIN_SEGMENTS) —
no compute-kernel logic is touched, so outputs stay bit-comparable. Result: AITER goes
from 2.3–5.6× slower (and crashing on decode due to reaching LDS limits) to parity-or-faster than the Triton
reference on every prefill/decode case, with full-attention decode fixed.

Technical Details

1. Full-attention prefill — eliminate accumulator spill (select_2d_config)

• BLOCK_M 128 → 16, num_warps=4, num_stages=1, TILE_SIZE=64.
• The f32 accumulator acc[BLOCK_M, 512] at BLOCK_M=128 needs 128×512×4/(4×64) = 256 VGPR/lane — the entire architectural register file — forcing everything else to spill to scratch and pinning occupancy at the 4-wave floor. BLOCK_M=16 cuts this 8× to 32 VGPR/lane and matches the MFMA M-tile exactly. TILE_SIZE=64 = one full KV page per tile (best gather locality on a scattered pool); a 2nd stage would need 256 KB > 160 KB LDS, so num_stages=1.

2. Full-attention decode — fix LDS overflow + raise memory-level parallelism (select_2d_config + select_3d_config)

• 2D: num_stages 2 → 1, num_warps 2 → 4. 3D split-KV: attn_warps 2 → 4, num_stages 2 → 1, MIN_SEGMENTS 8 → 16.
• With HEAD=512, TILE=64, one KV tile pair is 128 KB; double-buffering requested ~270 KB > 160 KB → the crash. num_stages=1 drops it to ~135 KB. MIN_SEGMENTS=16 splits the tiny decode grid (e.g. 64 WGs → 1024 WGs) so the ~8 TB/s HBM is saturated instead of latency-exposed.

3. Sliding-attention prefill — fill the grid (select_2d_config)

• BLOCK_M 128 → 16, num_warps 4 → 2, num_stages 1 → 2, TILE_SIZE 64 → 32.
• With nqpkv=2, BLOCK_M=128 → BLOCK_Q=64 launched only ~0.28 WG/CU (GPU mostly idle). BLOCK_M=16 → BLOCK_Q=8 raises this to ~2 WG/CU. The smaller HEAD=256 tiles (32 KB pair) leave room for num_stages=2 prefetch, and TILE_SIZE=32 = one page per tile.

Test Plan

• op_tests/triton_tests/attention/test_unified_attention.py

Test Result

Added tests for head_dim 256 and 512

Benchmark Result

Validated on model google/gemma-4-31B-it using vLLM v0.22, table below contains total token throughputs (t/s)

Configuration	TRITON_ATTN	AITER_UNIFIED_ATTN main branch	AITER_UNIFIED_ATTN this branch
tp=1, 1k->1k, conc=8	929.61	OutOfResources	1071.4
tp=1, 1k->1k, conc=16	1596.86	OutOfResources	1750.79
tp=1, 8k->1k, conc=16	4358.84	OutOfResources	5108.67
tp=4, 8k->1k, conc=16	9005.72	OutOfResources	9101.55
tp=4, 8k->1k, conc=32	13219.82	OutOfResources	13828.27

This optimization allows user to use AITER's Unified attention kernel for Gemma-4-31b model and increases total token throughput in x1.05-x1.17 times in comparison with using triton unified attention from vLLM.

SAFETENSORS_FAST_GPU=1 \
HIP_FORCE_DEV_KERNARG=1 \
TORCH_BLAS_PREFER_HIPBLASLT=1 \
NCCL_MIN_NCHANNELS=112 \
VLLM_ROCM_USE_AITER_RMSNORM=0 \
VLLM_ROCM_USE_AITER=1 \
VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 \
vllm serve google/gemma-4-31B-it \
  --dtype bfloat16 \
  --tensor-parallel-size X \
  --max-model-len 32768 \
  --max-num-seqs 2048 \
  --gpu-memory-utilization 0.90 \
  --block-size 32 \
  --enable-auto-tool-choice \
  --tool-call-parser gemma4 \
  --port 1616 \
  --host 0.0.0.0 \
  --no-enable-prefix-caching \
  --attention-backend ROCM_AITER_UNIFIED_ATTN

vllm bench serve \
  --backend openai-chat \
  --model google/gemma-4-31B-it \
  --base-url http://0.0.0.0:1616 \
  --endpoint /v1/chat/completions \
  --dataset-name random \
  --random-input-len X \
  --random-output-len X \
  --num-prompts X * 10 \
  --max-concurrency X

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: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 4044 --add-label <label>

[Copilot]

Pull request overview

Retunes Triton unified-attention config selection in AITER to better handle Gemma-4’s larger head sizes (256/512) on CDNA (gfx950), addressing performance regressions and avoiding LDS-related decode failures without changing kernel math.

Changes:

• Adjust select_2d_config prefill and decode heuristics for head_size >= 256/512 on non-RDNA targets (warps/stages/tile sizing).
• Adjust select_3d_config split-KV heuristics for head_size >= 512 (warps/stages and increased minimum segment count).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

op_tests/triton_tests/attention/test_unified_attention.py:580

• This parametrization doubles the already-large cartesian product of cases (head_size list grows from 2→4 values), which can make this test suite significantly slower and more resource-intensive. A straightforward way to keep coverage for 256/512 while limiting runtime is to couple (head_size, dtype combo) into a single parametrization so large head sizes run BF16-only, while smaller head sizes keep the broader dtype coverage.

@pytest.mark.parametrize("head_size", [64, 128, 256, 512])
@pytest.mark.parametrize("block_size", [16, 64])
@pytest.mark.parametrize("sliding_window", [None, 256])
@pytest.mark.parametrize(
    "soft_cap",
    [None, 50.0],
)
@pytest.mark.parametrize(
    "num_blocks",
    [2048, 32768],
)
@pytest.mark.parametrize(
    "q_dtype, kv_dtype, out_dtype, use_q_descale, use_kv_descale, use_out_scale",
    [
        (torch.bfloat16, torch.bfloat16, torch.bfloat16, False, False, False),
        (torch.bfloat16, e4m3_dtype, torch.bfloat16, False, True, False),
        (e4m3_dtype, e4m3_dtype, torch.bfloat16, True, True, False),
        (torch.float16, torch.float16, torch.float16, False, False, False),
    ],
)

[a-sidorova]

@ROCm/team_aiter may I ask you to take a look at these changes please? Thank you in advance