Skip to content

[BugFix] Potential bug fix for test_async_tp_pass_correctness #33854

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
LucasWilkinson wants to merge 1 commit into from
Closed

[BugFix] Potential bug fix for test_async_tp_pass_correctness #33854

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions vllm/v1/attention/backends/flash_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,11 +487,13 @@ def schedule(
if self.use_full_cuda_graph and scheduler_metadata is not None:
n = scheduler_metadata.shape[0]
self.scheduler_metadata[:n] = scheduler_metadata
# NOTE(woosuk): We should zero out the rest of the scheduler
# metadata to guarantee the correctness. Otherwise, some thread
# blocks may use the invalid scheduler metadata and overwrite the
# output buffer.
self.scheduler_metadata[n:] = 0
# NOTE(woosuk, lucas): Zero from n-1 onwards. Positions >= n must be
# zeroed to prevent invalid metadata from being used. The
# semaphore at position n-1 must also be zeroed before each
# forward pass because when num_splits == 1, FA3's internal
# semaphore reset uses PyTorch zero_() which isn't captured in
# CUDA graphs.
self.scheduler_metadata[n - 1 :] = 0
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 5, 2026

Choose a reason for hiding this comment

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

high

This change seems correct, but it might introduce a bug if n can be 0. If n=0, n-1 becomes -1, and self.scheduler_metadata[-1:] = 0 will only zero out the last element of the buffer. The previous behavior for n=0 was to zero out the entire buffer (self.scheduler_metadata[0:] = 0), which seems safer for resetting state when there are no requests.

While n is likely always >= 1 (since scheduler_metadata size is batch_size * 4 + 1), it's safer to handle the n=0 case explicitly to prevent potential issues. A more robust implementation would be:

Suggested change
self.scheduler_metadata[n - 1 :] = 0
self.scheduler_metadata[max(0, n - 1):] = 0

scheduler_metadata = self.scheduler_metadata[:n]

attn_metadata = FlashAttentionMetadata(
Expand Down
12 changes: 7 additions & 5 deletions vllm/v1/attention/backends/mla/flashattn_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,11 +224,13 @@ def _build_decode(
f"{self.scheduler_metadata.shape[0]}"
)
self.scheduler_metadata[:n] = scheduler_metadata
# NOTE(woosuk): We should zero out the rest of the scheduler
# metadata to guarantee the correctness. Otherwise, some thread
# blocks may use the invalid scheduler metadata and overwrite the
# output buffer.
self.scheduler_metadata[n:] = 0
# NOTE(woosuk, lucas): Zero from n-1 onwards. Positions >= n must be
# zeroed to prevent invalid metadata from being used. The
# semaphore at position n-1 must also be zeroed before each
# forward pass because when num_splits == 1, FA3's internal
# semaphore reset uses PyTorch zero_() which isn't captured in
# CUDA graphs.
self.scheduler_metadata[n - 1 :] = 0
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 5, 2026

Choose a reason for hiding this comment

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

high

This change seems correct, but it might introduce a bug if n can be 0. If n=0, n-1 becomes -1, and self.scheduler_metadata[-1:] = 0 will only zero out the last element of the buffer. The previous behavior for n=0 was to zero out the entire buffer (self.scheduler_metadata[0:] = 0), which seems safer for resetting state when there are no requests.

While n is likely always >= 1 (since scheduler_metadata size is batch_size * 4 + 1), it's safer to handle the n=0 case explicitly to prevent potential issues. A more robust implementation would be:

Suggested change
self.scheduler_metadata[n - 1 :] = 0
self.scheduler_metadata[max(0, n - 1):] = 0

scheduler_metadata = self.scheduler_metadata[:n]

metadata = FlashAttnMLADecodeMetadata(
Expand Down
Loading