[PERF] Optimize Triton unified attention prefill and decode

[vorapolsiloai]

Summary

• skip per-element causal and bounds masking for prefill tiles that are fully visible to every row in a query block
• add a tuned FP8 prefill launch configuration for the validated gfx950 shape family
• use a tuned single-launch 2D path for conservative small-batch decode shapes instead of split-KV plus reduction

This rebases the validated optimization from 389fb3d3 onto current main while preserving Mohsen Saffari's authorship.

Performance validation

Tested end-to-end with Qwen3-VL-235B-A22B-Instruct-MXFP4 by applying only this AITER change to the same pinned vLLM Image C base. Results use same-device/pair crossover placement; one thermally throttled GPU was detected and excluded.

• TP1, 60k input / 600 output, concurrency 32: +14.68% median output throughput; 12.80% lower duration; 14.10% lower mean TTFT; 13.30% lower mean TPOT across three healthy same-GPU comparisons
• TP2, 60k input / 600 output, concurrency 32: +11.57% median output throughput; 10.37% lower duration; 12.81% lower mean TTFT; 9.97% lower mean TPOT across three same-pair comparisons
• TP1, 8k input / 1k output, concurrency 128: +1.35% median output throughput across three healthy same-GPU comparisons
• TP2, 8k input / 1k output, concurrency 256: two independent Latin-square passes produced -1.41% and +1.77% median output-throughput changes; the combined six comparisons are approximately neutral, so no repeatable short-context regression was observed

Test plan

• clean cherry-pick onto AITER main at 1b741c0
• git diff --check upstream/main...HEAD
• Python byte-compilation of both changed modules
• TP1 end-to-end serving crossover: 8k/1k and 60k/600
• TP2 three-pair Latin-square serving crossover: 8k/1k and 60k/600
• run the upstream unified-attention correctness suite on the rebased branch
• upstream CI

Review notes

The PR is intentionally draft while the upstream correctness suite and CI are pending. The kernel fast loop is prefill-only; sliding-window and decode calls retain the general masked loop. The launch overrides are restricted by dtype, head size, GQA ratio, architecture, sequence length, and batch conditions documented inline.

Made with Cursor

[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 4761 --add-label <label>