[Bugfix] Fix model run _npu_flash_attention hang issue #4410
Conversation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
![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
The pull request effectively addresses a reported hang issue in _npu_flash_attention by refining the handling of attention mask dtypes. The key change involves refactoring the logic for retrieving the chunked_prefill_attn_mask into a dedicated method, get_chunked_prefill_attn_mask. This new method explicitly ensures the mask is converted to torch.bool, which is crucial for the correct operation of the attention mechanism. This refactoring not only fixes the bug but also enhances code clarity and maintainability by centralizing the dtype conversion for this specific mask.
| def get_chunked_prefill_attn_mask(self): | ||
| return self.chunked_prefill_attn_mask.to(torch.bool) | ||
|
|
||
| def get_attn_mask(self, max_seq_len: int, dtype: torch.dtype, |
There was a problem hiding this comment.
The introduction of get_chunked_prefill_attn_mask and the removal of the conditional logic from get_attn_mask is a significant improvement. This refactoring clearly separates the responsibility of providing the chunked prefill attention mask and explicitly ensures its torch.bool dtype. This directly addresses the reported "wrong attention mask dtype" issue, which was causing hangs in _npu_flash_attention, by enforcing the correct data type for this specific mask. It also makes the get_attn_mask method more focused on its general purpose.
| elif attn_state == AscendAttentionState.PrefillCacheHit: | ||
| return self.attn_mask_builder.get_attn_mask( | ||
| 2048, self.dtype, self.device) | ||
| return self.attn_mask_builder.get_chunked_prefill_attn_mask() |
There was a problem hiding this comment.
Updating the call to use the new get_chunked_prefill_attn_mask() method is a correct and consistent application of the refactored logic. This change ensures that the attention mask used in the PrefillCacheHit state consistently has the torch.bool dtype, which is essential for preventing the _npu_flash_attention hang issue as described in the PR.
Semmer2
force-pushed
the
FixFAHangIssue
branch
3 times, most recently
from
f5f2eac to
36d1b5e
Compare
Fix model run _npu_flash_attention in _forward_prefill_no_cache hang issue, it was caused by wrong attention mask dtype. Signed-off-by: Ting FU <[email protected]>
Semmer2
force-pushed
the
FixFAHangIssue
branch
from
36d1b5e to
5b7ef34
Compare
1 participant
Fix model run _npu_flash_attention in _forward_prefill_no_cache hang issue, it was caused by wrong attention mask dtype.
What this PR does / why we need it?
Does this PR introduce any user-facing change?
No
How was this patch tested?
Yes, tesed on Qwen2.5-VL and Qwen2.5-Omni