Repository navigation
[Triton/Gluon] [Feature] MHA: add out_dtype to the Triton forward entry point - #7232
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
NOT READY FOR HUMAN REVIEWPR Self-ReviewRevision Author self-review accepted. Readiness report: Author answers
Show the author's self-review answersBehavior and intentValidation and unverified scopePR-specific questionChanged op testsPR state
Copilot findings
Code and op-test hygiene
Required CI
Shadow mode reports what would block but keeps the status successful. Set |
f6e7392 to
812a7a8
Compare
_flash_attn_forward allocates its output as float32 for FP8 inputs and as q.dtype otherwise, so a bfloat16 caller casts the float32 output on the host. out_dtype selects the output dtype; None keeps the current choice. The kernel accumulates in float32 and casts at the store, so the narrower output matches the float32 output to within one ULP: the backend can fuse the epilogue multiply into the store cast and round once instead of twice. Signed-off-by: Pham Binh <194774975+phambinhfin@users.noreply.github.com>
812a7a8 to
81e8b0f
Compare
github-actions
Bot
changed the title
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved issues include incorrect FP8 capability gating, validation that does not prove the claimed exact cast equality, and dao_ai failures for float16 or float32 outputs with bfloat16 inputs.
3 open issues
What changed in this PR
Adds optional output dtype selection to the Triton MHA forward entry point, allowing FP8 callers to request bfloat16 output without a separate cast.
Changes:
- Adds
out_dtypeto_flash_attn_forwardwhile preserving existing defaults. - Adds 24 test cases covering input dtypes, batch and varlen layouts, causal modes, and output dtypes.
- Leaves public wrapper signatures and the Gluon backend unchanged.
| File | Description |
|---|---|
op_tests/triton_tests/attention/test_mha.py |
Adds output dtype checks and tolerance-based comparisons with float32 output. |
aiter/ops/triton/attention/mha.py |
Allocates the forward output using the requested dtype. |
Review effort: Balanced
|
Static review (no GPU run). Head reviewed: 81e8b0f. No blocking defect: the [P2] The test checks a tolerance, while the description claims bit-exact equality. [P3] [P3] The FP8 skip gate uses a different predicate from the entry point. [P3] [P3] Coverage misses the consumer configuration and side outputs. Note: the vLLM consumer (vllm-project/vllm#62913) does not pass Not verified: |
…utput Reject an out_dtype other than float32, bfloat16 or float16. Where the implementation requires the output dtype to match q.dtype, reject a different out_dtype before dispatch instead of failing in fwd_prefill. The test now compares the output with the float32 output cast to out_dtype. An fp16 store can fuse the 1/l multiply into the cast and round once, so it allows a 1 ulp difference on at most 1 in 1000 values; a truncating store fails it. The cases are cut to one per behavior, including the vLLM caller's shape (FP8, varlen, head size 256). It also gates FP8 cases on is_fp8_avail(), checks that softmax_lse is unchanged, and covers the two rejections. Signed-off-by: Pham Binh <194774975+phambinhfin@users.noreply.github.com>
Hi @samremes, thanks for the review. All points are addressed in 9dd5742. For P2 (tolerance vs. the bit-exact claim): the test now compares out with ref.to(expected_dtype) at a 1-ulp bound, and allows at most 1 in 1000 values to differ. The 1 ulp is needed only for float16: with bfloat16 inputs and non-causal attention, the store compiles to v_fma_mixlo_f16, which rounds the 1/l multiply and the conversion once, so 13 of 524,288 values differ by 1 ulp from the double-rounded cast. bfloat16 outputs are bit-identical in every case. A store forced to round toward zero fails 5 of the 6 output cases (the sixth has a float32 output). The description and results are refreshed at the head. For P3 (fails late): a different out_dtype is now rejected in that branch before dispatch, next to the sink, PE and FP8 checks, with Output and input must have the same dtype, got out_dtype=... and q.dtype=.... _flash_attn_forward has no argument docs, so the limit is stated in the message and in the description. For P3 (FP8 gate): the skip now uses not is_fp8_avail(), the same check as the entry point. For P3 (validation): an out_dtype other than float32, bfloat16 or float16 raises ValueError. On the float16 range: the output is a convex combination of the rows of V, so it can overflow only if the dequantized V values exceed the float16 maximum, which is the caller's choice. For P3 (coverage): one case has the consumer's shape (FP8, varlen, causal, head size 256). Every case passes return_lse=True and checks that softmax_lse stays float32 and bit-identical. The rejection test covers the early check. I ran the consumer's explicit config separately at the head, with uneven sequence lengths: the bfloat16 and float16 outputs are bit-identical to the cast float32 output, causal and non-causal. I did not add it as a test case, because the config changes only the tile sizes. For the note: right, vllm-project/vllm#62913 still casts after the call. It will pass out_dtype once an aiter version with this change is available there. For "Not verified": now run at 9dd5742 on gfx942: all of test_mha.py (1202 passed, 1133 skipped; 1194 passed on main), the four other selected attention test files (same results as main), and the consumer configuration above. A non-float target (torch.int32) raises ValueError and is tested. An overflowing float16 target is not tested. gfx950 is covered by CI only. |
|
@azaidy Could you take look about this PR |
brunomazzottiamd
left a comment
brunomazzottiamd
left a comment
There was a problem hiding this comment.
Requesting minor changes that are easily addressable.
|
Added |
Name the dao_ai impl in its out_dtype check. In the test, split the softmax_lse assert in two, give each assert a message, and state how often the fp16 store differs from the double-rounded cast. Signed-off-by: Pham Binh <194774975+phambinhfin@users.noreply.github.com>
There was a problem hiding this comment.
🔵 Needs a closer look
The output dtype change affects Triton specialization and float16 rounding, so numerical correctness requires head-matched gfx942 and gfx950 CI validation unavailable in this source-only review.
No open issues 🎉
Review effort: Balanced
|
Flash Attention Integration / Flash Attention - Triton / RDNA3 (1 GPU) (pull_request) CI job seems to be stuck. I contacted AITER CI team and asked for help. |
|
Thanks @brunomazzottiamd a lot for contacting the CI team . It looks like the aiter-gfx1100 runners are offline: every RDNA3 job since yesterday 18:46 UTC is queued without a runner, |
@phambinhfin, the hang related to Flash Attention Integration / Flash Attention - Triton / RDNA3 (1 GPU) (pull_request) CI job was also happening in #7257. It got unstuck and I was able to merge #7257 a couple of minutes ago. AITER CI team told me the job requires a node type that's scarce in CI pool. Let's wait a bit more and hope for the best. |
…ry point (#7232) * [Triton/Gluon] MHA: add out_dtype to the Triton forward entry point _flash_attn_forward allocates its output as float32 for FP8 inputs and as q.dtype otherwise, so a bfloat16 caller casts the float32 output on the host. out_dtype selects the output dtype; None keeps the current choice. The kernel accumulates in float32 and casts at the store, so the narrower output matches the float32 output to within one ULP: the backend can fuse the epilogue multiply into the store cast and round once instead of twice. Signed-off-by: Pham Binh <194774975+phambinhfin@users.noreply.github.com>
Motivation
_flash_attn_forwardis the forward entry point of the Triton MHA backend. For FP8 inputs, it always allocates a float32 output. For other inputs, the output has the dtype ofq. A caller that needs a bfloat16 output from FP8 inputs must cast the float32 output after the call. This cast is one more kernel launch, and the float32 buffer is two times the size of the bfloat16 output.The draft FP8 prefill path for ROCm in vLLM (vllm-project/vllm#62913) calls
_flash_attn_forwardwith FP8 inputs and needs a bfloat16 output.Technical Details
_flash_attn_forwardgets a new argument,out_dtype: torch.dtype | None = None. It sets the dtype of the output tensor.Nonekeeps the current dtype: float32 for FP8 inputs,q.dtypefor other inputs.out_dtypemust be float32, bfloat16 or float16. Other dtypes raiseValueError.out_dtype, every implementation gets the same output dtype as before. An implementation that requires the output dtype to matchq.dtyperejects a differentout_dtypebefore dispatch, with a message that namesout_dtype.1/lmultiply into the conversion (v_fma_mixlo_f16), so the value is rounded once instead of twice.softmax_lsestays float32.flash_attn_func,flash_attn_varlen_func) and the Gluon backend do not change.Test Plan
Two new tests in
op_tests/triton_tests/attention/test_mha.py:test_mha_forward_out_dtype, 6 cases, one per behavior:out_dtype=None: the output stays float32.out_dtype=None: the output stays bfloat16.1/lmultiply into the conversion.Each case checks the output dtype, checks that
softmax_lseis bit-identical to the float32 run, and compares the output with the float32 output cast to the output dtype: at most 1 ulp apart, and at most 1 in 1000 values different.test_mha_forward_out_dtype_rejected, 2 cases:out_dtype=torch.int32raisesValueError, andout_dtype=torch.float16with bfloat16 inputs fails before dispatch in the implementation that requires matching dtypes.Also on this PR: the other attention test files that the Triton test selection picks for this change (
test_mha_dao_ai.py,test_mha_spatial_swizzle.py,test_mha_with_pe.py,test_mha_with_sink.py).Test Result
MI325X (gfx942), ROCm 7.2, Triton 3.8.0, torch 2.13.0. Baseline: main
bd4ff1c. PR:59028bd.TypeError, noout_dtypeargument)test_mha.pytest_mha_dao_ai.pytest_mha_spatial_swizzle.pytest_mha_with_pe.pytest_mha_with_sink.pyHow exact the narrow output is, measured on this PR over the full grid of these options (a separate script, not part of the test), against the float32 output cast to the output dtype:
v_fma_mixlo_f16.config): bit-identical for bfloat16 and float16 outputs, causal and non-causal, andsoftmax_lseis unchanged.The test fails when the store is wrong: with the store forced to round toward zero, 5 of the 6 cases fail (the other one has a float32 output, so the store does not narrow). With the two
out_dtypechecks removed, both rejection cases fail.CI ran the new tests on MI35X (gfx950) and, with the
ci:triton-300xlabel, on MI300X (gfx942). All 8 cases passed on both at9dd5742;59028bdonly adds assert messages.PR Self-Review
Reviewed revision
495d3b9
Behavior and intent
_flash_attn_forwardgets an optionalout_dtypeargument that sets the dtype of the output tensor. WithNone, the output dtype does not change: float32 for FP8 inputs andq.dtypefor other inputs. Values outside float32, bfloat16 and float16 are rejected, and where an implementation requires it,out_dtypemust equalq.dtype. The kernel casts from its float32 accumulator only when it stores the result. Thus a caller gets a bfloat16 output from FP8 inputs without a cast after the call.Validation and unverified scope
On MI325X (gfx942), the two new tests: 8 failed on main
bd4ff1c(TypeError) and 8 passed on this PR. All oftest_mha.py: 1194 passed on main and 1202 passed on this PR, with 1133 skipped on both. The four other selected attention test files have the same results on main and on this PR. gfx950 was not run locally; CI passed all 8 cases on MI35X (gfx950) and on MI300X (gfx942). Whether gfx950 fuses the float16 store the same way was not checked; the test allows it either way. The public wrappers do not takeout_dtype, so this PR does not cover them.PR-specific question
On main, all 8 new cases fail with
TypeError: _flash_attn_forward() got an unexpected keyword argument 'out_dtype'. On this PR they pass. The output check fails for a store that does not round to nearest: forced round-toward-zero fails 5 of the 6 output cases. No existing test can cover this, because main has no argument that sets the output dtype. The existing tests use the default dtype, and they pass on both trees.Changed op tests
op_tests/triton_tests/attention/test_mha.py: addstest_mha_forward_out_dtype(6 cases) andtest_mha_forward_out_dtype_rejected(2 cases). They check thatout_dtypesets the output dtype, that the output is the float32 output cast to that dtype, thatsoftmax_lsedoes not change, and that unsupported values are rejected before dispatch. They call_flash_attn_forwarddirectly, because the public wrappers do not takeout_dtype. No existing test sets an output dtype.