Skip to content

integrate new AITER mha v4 fmha entrypoint to xDiT - #752

Merged
jcaraban merged 5 commits into
mainfrom
mxfp_fmha_asm
Aug 24, 2026
Merged

jcaraban merged 5 commits into
mainfrom
mxfp_fmha_asm

Conversation

@jcaraban

Copy link
Copy Markdown
Collaborator

Motivation

Expose the low-precision MHA v4 kernels added by ROCm/aiter#4627 through xDiT.

Technical Details

  • Add I8FP8, MXFP4, MXFP6, F4F4, and F6F4 attention backends.
  • Delegate preprocessing and kernel dispatch to AITER's public mha_v4 API.
  • Support unequal query and KV sequence lengths.
  • Validate availability through the public AITER API.
  • Reject unsupported causal attention explicitly.
  • Document the new gfx950, head-dimension-128 backends.

Testing

  • Added eager correctness checks against SDPA.
  • Added torch.compile(fullgraph=True) coverage.
  • Covered unequal sequence lengths, cross-attention, partial tiles, and causal rejection.
  • All 20 focused tests pass on gfx950.

@jcaraban
jcaraban requested a review from avjves August 17, 2026 13:20
pass # Error is raised in runtime_state.py when an MHA v4 backend is selected.

try:
from aiter.ops.mha_v4 import (

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are these separate from the above check?

@jcaraban jcaraban Aug 24, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

because mha_v4_mxfp8 will fail to import until ROCm/aiter#4967 merges, so if it was with the above block it would make the other MHA v4 backends fail, while those import fine in AITER main

@jcaraban
jcaraban force-pushed the mxfp_fmha_asm branch 2 times, most recently from ce7864a to bc36ac7 Compare August 18, 2026 09:05
@jcaraban
jcaraban merged commit c11d0a0 into main Aug 24, 2026
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