[Bugfix][Triton] Avoid RDNA4 unified attention LDS overflow

[hogeheer499-commits]

Summary

Prevent the Triton 3D unified-attention kernel from selecting a two-stage
configuration whose staged K/V tiles already consume the full 64 KiB LDS
budget on gfx1200 and gfx1201.

Motivation

With a 256-wide BF16 KV head and TILE_SIZE=64, Triton 3.6 reports:

OutOfResources: shared memory, Required: 65792, Hardware limit: 65536

The K/V tiles alone require 64 KiB at two stages, leaving no room for Triton's
256-byte scratch allocation. This causes a hard startup failure in the
downstream vLLM report.

The Qwen3.6-27B model from that report uses a 256-wide full-attention head:
https://huggingface.co/Qwen/Qwen3.6-27B-FP8/blob/main/config.json

Changes

• Estimate the two-stage K/V tile footprint after the final tile size is known.
• Use one stage on gfx1200 and gfx1201 only when that footprint reaches the
  64 KiB LDS limit.
• Preserve two stages for smaller heads or tiles, FP8 KV caches, and gfx1151.
• Add focused selector regression coverage for both affected architectures and
  the configurations that must remain unchanged.

Testing

• 7 passed:
  pytest op_tests/triton_tests/attention/test_unified_attention.py -k test_select_3d_config_respects_rdna_lds_limit -q
• Triton 3.6 compile targeting gfx1200 and gfx1201:
  • before: two stages, 65,792 bytes, compile rejected
  • after: one stage, 32,768 bytes
• Real gfx1151 execution:
  • its existing two-stage 65,536-byte path remains selected and passes the
    numerical comparison
  • forcing the new gfx1201 selector path uses one stage and also passes the
    numerical comparison
• BF16 query with FP8 KV remains at two stages and compiles at 33,024 bytes.
• Black check, import-order check, Python byte compilation, and
  git diff --check pass.

I do not have native gfx1200/gfx1201 hardware, so the affected targets were
compiler-validated with Triton's architecture override. Native RDNA4 CI or
maintainer validation would still be valuable.

No new dependencies or breaking changes.

Fixes #4329.

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