[BugFix][DP/EP] Fix CUTLASS MLA hang under load #26026
Merged
+32
−32
Conversation
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
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses a potential hang in the CUTLASS MLA kernel under load. The changes are mostly stylistic, involving code re-indentation. However, there is one critical bug fix that changes a cutlass::arch::NamedBarrier::arrive() call to arrive_and_wait(). This change correctly resolves a race condition that could lead to a deadlock between compute and load warps, which is the likely cause of the hang. My review confirms this fix is correct and critical.
csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp
Show resolved
Hide resolved
|
TODOs (follow-up)
|
LucasWilkinson
force-pushed
the
lwilkinson/fix-cutlass-mla-hang
branch
from
746ef1a to
a36b036
Compare
Co-authored-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
LucasWilkinson
force-pushed
the
lwilkinson/fix-cutlass-mla-hang
branch
from
a36b036 to
3a8a634
Compare
mgoin
approved these changes
Oct 1, 2025
mgoin
added
bug
| (kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, | ||
| kNamedBarrierEpilogue | ||
| ).arrive(); | ||
| ).arrive_and_wait(); |
There was a problem hiding this comment.
note for reviewer --- this is the line that fixes the deadlock
robertgshaw2-redhat
changed the title
Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
tlrmchlsmth
approved these changes
Oct 1, 2025
1 task
|
Failing tests are known on main and being resolved. Blackwell tests are green, merging |
soldni
pushed a commit
to soldni/vllm
that referenced
this pull request
Oct 1, 2025
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com>
simon-mo
pushed a commit
that referenced
this pull request
Oct 2, 2025
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com> Signed-off-by: simon-mo <simon.mo@hey.com>
pdasigi
pushed a commit
to pdasigi/vllm
that referenced
this pull request
Oct 2, 2025
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com>
yewentao256
pushed a commit
that referenced
this pull request
Oct 3, 2025
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com> Signed-off-by: yewentao256 <zhyanwentao@126.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
The early return in
compute(calls arrive:but didn't have any barrier before looping around and calling it again causing a deadlock when the load warps waits on:
Co-authored-by: Robert Shaw robshaw@redhat.com