Skip to content

bump aiter to 859e56f98 for gfx1201/gfx1200 FP8 support#186

Open
johnodin99 wants to merge 1 commit into
ROCm:tridaofrom
johnodin99:johnodin99/bump-aiter-gfx1201-fp8
Open

bump aiter to 859e56f98 for gfx1201/gfx1200 FP8 support#186
johnodin99 wants to merge 1 commit into
ROCm:tridaofrom
johnodin99:johnodin99/bump-aiter-gfx1201-fp8

Conversation

@johnodin99

Copy link
Copy Markdown

Motivation

This PR addresses the FP8 architecture gate for RDNA4 (gfx1200 / gfx1201) in the AMD Triton flash-attention path.

Currently, FP8 support can be blocked by the architecture support table even when running on an FP8-capable RDNA4 platform. This PR updates the bundled aiter submodule to a newer revision that includes RDNA4 FP8 architecture support.

Related issue:

Technical Details

This PR bumps the third_party/aiter submodule:

third_party/aiter: 428e8e761 -> 859e56f98

The newer aiter revision includes RDNA4 FP8 support for gfx1200 / gfx1201, including updates to the FP8 architecture capability checks used by the AMD Triton flash-attention path.

The nested CK submodule under aiter is also updated through the aiter submodule state:

third_party/aiter/3rdparty/composable_kernel: 020b6f435e787b6f7f40aabf638a3ef9ab6a9002

This PR is intentionally limited to the aiter submodule bump. It does not directly modify flash-attention kernels or change the AMD Triton FA2 direct-FP8 tensor path.

Note:

Direct FP8 tensors through flash_attn_func currently hit the AMD Triton FA2 interface, which raises NotImplementedError for FP8 tensors. The validated FP8 path in this PR is the AITER FA3 FP8 path.

Test Plan

Tested locally on gfx1201 with:

ROCm: 7.2
PyTorch: 2.11.0+rocm7.2
FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
GPU_ARCHS=gfx1201

The following checks were run:

  1. Build flash-attention with the updated aiter submodule.

  2. Verify the FP8 architecture gate:

arch: gfx1201
supports_fp8: True
PR-A gate check: PASS
  1. Run FP16 forward smoke test using flash_attn_func.

  2. Run FP16 backward smoke test using flash_attn_func.

  3. Run AITER FA3 FP8 forward numerical correctness check against an FP32 reference attention implementation.

The FP8 numerical check covers:

causal=False, head_dim=128
causal=False, head_dim=64
causal=True,  head_dim=128
causal=True,  head_dim=64

Test Result

All local validation checks passed on gfx1201.

Summary:

FP8 arch gate:
  arch=gfx1201
  supports_fp8=True
  PASS

FP16 forward smoke test:
  PASS

FP16 backward smoke test:
  PASS

AITER FA3 FP8 forward numerical check:
  ALL FP8 FWD NUMERICAL CHECKS PASS

FP8 numerical accuracy versus FP32 reference:

causal=False, head_dim=128:
  rel_fro=5.1931e-02
  cosine=0.998652
  PASS

causal=False, head_dim=64:
  rel_fro=5.4826e-02
  cosine=0.998498
  PASS

causal=True, head_dim=128:
  rel_fro=4.6094e-02
  cosine=0.998937
  PASS

causal=True, head_dim=64:
  rel_fro=4.8156e-02
  cosine=0.998840
  PASS

Submission Checklist

@johnodin99

Copy link
Copy Markdown
Author

check record

check_record.txt

@rocking5566

Copy link
Copy Markdown
Collaborator

Thanks for the fix! tridao is a read-only mirror we re-sync from Dao-AILab upstream, so PRs merged there get overwritten — please don't target it directly.

Two options:

  • Retarget to ck_improve_main — lands in the ROCm fork right away. We batch PRs from it back to upstream, but the timing of that isn't fixed.
  • If you're in a hurry, open the PR directly against Dao-AILab upstream — once it merges there, we can sync it into the fork at any time.

Either works; pick whichever gets you unblocked fastest.

@johnodin99

Copy link
Copy Markdown
Author

Thanks for the fix! tridao is a read-only mirror we re-sync from Dao-AILab upstream, so PRs merged there get overwritten — please don't target it directly.

Two options:

  • Retarget to ck_improve_main — lands in the ROCm fork right away. We batch PRs from it back to upstream, but the timing of that isn't fixed.
  • If you're in a hurry, open the PR directly against Dao-AILab upstream — once it merges there, we can sync it into the fork at any time.

Either works; pick whichever gets you unblocked fastest.

Thanks for the guidance. Before retargeting, I checked the current ck_improve_main state. It appears to pin third_party/aiter at 315fb08e948855ab3a5be5555ff83892d1be0895 through commit b02b07e, and that aiter revision already includes gfx1200 / gfx1201 in FP8_ARCHS, is_fp8_avail(), and the FP8 dtype selection logic.

I also found that 315fb08e and 859e56f98 have diverged histories, so retargeting this PR as-is would replace the current aiter pin with a different history rather than perform a straightforward forward bump.

It therefore looks like the architecture-gate portion of this PR may already be superseded on ck_improve_main and current Dao-AILab upstream. Would you prefer that I close #186 as already covered, or is there another aiter revision/branch that you would like this PR to target?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants