fix: generate CK FMHA kernels for gfx1250

[tomjen12]

Summary

• Pass the runtime gfx target to CK dense forward and backward code generation.
• Pass the same target to CK varlen backward code generation.
• Enable gfx1250 CK dispatch instead of falling back with an invalid-argument error.

Manual validation

• Validated dense and varlen FlashAttention forward/backward on gfx1250 using downstream captured workloads.
• Compared outputs and Q/K/V gradients with PyTorch math SDPA.
• Confirmed CK forward and backward kernels launch in Torch Profiler.

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

[Copilot]

Pull request overview

This PR updates CK FMHA kernel generation command lines so that runtime codegen includes an explicit --targets <gfx> flag, aiming to generate kernels for gfx1250 instead of producing an empty set and failing dispatch.

Changes:

• Append a CK target-arch flag to CK generate.py invocations for dense MHA fwd and bwd codegen.
• Append the same CK target-arch flag to varlen MHA bwd codegen.
• Import and reuse a shared _ck_targets_flag() helper for these command-line updates.

Suppressed comments (2)

aiter/ops/mha.py:1144

• Same as in fwd: relying on _ck_targets_flag() (which uses get_gfx() rather than rocminfo-detected get_gfx_runtime()) can still produce a kernel set for a non-running arch when GPU_ARCHS is set, undermining the goal of enabling gfx1250 runtime dispatch. Updating the helper to use get_gfx_runtime() would make the --targets selection robust.

    blob_gen_cmd = [
        f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d bwd "
        "--receipt 300 --filter {} --output_dir {{}}{}".format(
            filter, _ck_targets_flag()
        ),

aiter/ops/mha.py:1404

• Same concern here: if _ck_targets_flag() resolves the arch via get_gfx() (GPU_ARCHS/build-time) instead of get_gfx_runtime(), varlen bwd codegen may still target the wrong ISA on multi-arch setups. Using get_gfx_runtime() for the --targets flag would better align with runtime kernel generation/dispatch for gfx1250.

    blob_gen_cmd = [
        f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d bwd "
        "--receipt 400 --filter {} --output_dir {{}}{}".format(
            filter, _ck_targets_flag()
        ),

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.