[AMD][DCP] Close two decode-LSE holes left by AITER DCP support - #54639
Closed
okorzh-amd wants to merge 2 commits into
Closed
okorzh-amd wants to merge 2 commits into
okorzh-amd wants to merge 2 commits into
Conversation
vllm-project#51705 gave the AITER MLA backend a decode LSE under decode context parallelism, so a DCP rank can hand its partial attention to the cross-shard merge. Two ways to reach that merge without an LSE remain. Gluon. `use_gluon_decode` takes no DCP argument, unlike its sibling `use_gluon_verify`, and the Gluon single-token branch in `forward_mqa` ends `return o, None`. Nothing structural keeps DCP off it -- only arithmetic: the builder passes the gathered head count, and Gluon bails at >= 16 heads, which every realistic DCP layout exceeds (Kimi-K3 at TP8/DCP8 gathers 12 x 8 = 96). At 6 local heads and DCP2 it gathers 12 and Gluon takes the batch, returning no LSE and tripping a bare `assert lse is not None` in the MLA layer. That shape is not hypothetical; it is the one the existing DCP decode test uses. Give the predicate a `dcp_world_size` and exclude DCP, matching verify. AITER capability. The DCP branch needs an AITER build whose `mla_decode_fwd` accepts `return_lse=True`. Nothing checks that at configuration time, so an older build loads the whole model and fails at the first decoded token on an assertion that names neither AITER nor the version. Probe the signature once and raise at construction instead. Neither changes any kernel selection for a supported configuration. Test: tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py tests/v1/attention/test_rocm_aiter_mla_mtp_split.py 565 passed, against 505 on the same commit without this change -- the delta is exactly the added cases, no regressions. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Oxana Korzh <okorzh@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Comments, docstrings and the error string only; no behavior change. 565 tests still pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Oxana Korzh <okorzh@amd.com>
Contributor
Author
|
PR #54546 resolves the issue. Closing this one |
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.
#51705 gave the AITER MLA backend a decode LSE under decode context parallelism, so a DCP rank can hand its partial attention to the cross-shard merge. Two ways to reach that merge without an LSE remain.
Gluon.
use_gluon_decodetakes no DCP argument, unlike its siblinguse_gluon_verify, and the Gluon single-token branch inforward_mqaendsreturn o, None. Nothing structural keeps DCP off it -- only arithmetic: the builder passes the gathered head count, and Gluon bails at >= 16 heads, which every realistic DCP layout exceeds (Kimi-K3 at TP8/DCP8 gathers 12 x 8 = 96). At 6 local heads and DCP2 it gathers 12 and Gluon takes the batch, returning no LSE and tripping a bareassert lse is not Nonein the MLA layer. That shape is not hypothetical; it is the one the existing DCP decode test uses. Give the predicate adcp_world_sizeand exclude DCP, matching verify.AITER capability. The DCP branch needs an AITER build whose
mla_decode_fwdacceptsreturn_lse=True. Nothing checks that at configuration time, so an older build loads the whole model and fails at the first decoded token on an assertion that names neither AITER nor the version. Probe the signature once and raise at construction instead.Neither changes any kernel selection for a supported configuration.
Test Plan
Test: tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
565 passed, against 505 on the same commit without this change --
the delta is exactly the added cases, no regressions.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.