[ROCm][Bugfix][MLA] Fix mla_reduce_v1 num_kv_splits arg for FP8 MLA prefill - #46810
Closed
xaguilar-amd wants to merge 1 commit into
Closed
xaguilar-amd wants to merge 1 commit into
xaguilar-amd wants to merge 1 commit into
Conversation
…s arg to mla_reduce_v1; pass it (the compute-unit count) when the installed AITER expects it, detected via capability probe. Fixes FP8 MLA prefill crashing on gfx950 with "Expected int for num_kv_splits but found Tensor". Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
Author
|
cc: @ChuanLi1101 |
Contributor
Author
|
Bug fixed as well with #46760 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
AITER PR #3391 (released in v0.1.16) added a required positional
num_kv_splits: inttomla_reduce_v1, inserted betweenmax_seqlen_qandfinal_output. The AITER MLA backend's FP8 prefill call site still used the old signature, so the output tensor landed in thenum_kv_splitsslot and FP8 MLA prefill crashed on gfx950 with:Expected a value of type 'int' for argument 'num_kv_splits' but instead found type 'Tensor'This updates the call site to match the new signature. Validated on
amd/Kimi-K2.5-MXFP4-AttnFP8(gfx950).Changes
num_kv_splitstomla_reduce_v1for the FP8 MLA prefill path inrocm_aiter_mla.py.get_cu_num()) as the split count — the tight per-reduce-tile upper bound for the persistent-scheduling prefill kernel, equal to the kernel's own internalmax_splitsfloor (not a placeholder).inspect.signature→ torch op schema fallback, matching the pattern invllm/_aiter_ops.py), so builds on the older signature keep working.Test plan
amd/Kimi-K2.5-MXFP4-AttnFP8).num_kv_splitspassed).