xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-26 16:54:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Avoid too small block m/n for FlexAttention kernel option (#27853)

Browse Source 
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
This commit is contained in:
Isotr0py 2025-10-31 19:33:12 +08:00 committed by GitHub
parent e5ef4dfc11
commit 3933f18a5e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/v1/attention/backends/flex_attention.py
View File

@ -896,6 +896,8 @@ def get_kernel_options(
return kernel_options return kernel_options
else: else:
preferred_block = 32 if query.dtype == torch.float32 else 64 preferred_block = 32 if query.dtype == torch.float32 else 64
block_lower_bound = 16
block_m_candidate = ensure_divisible(preferred_block, block_m) block_m_candidate = ensure_divisible(preferred_block, block_m)
block_n_candidate = ensure_divisible(preferred_block, block_n) block_n_candidate = ensure_divisible(preferred_block, block_n)
@ -910,6 +912,9 @@ def get_kernel_options(
max(1, block_n_candidate // 2), block_n max(1, block_n_candidate // 2), block_n
) )
block_m_candidate = max(block_m_candidate, block_lower_bound)
block_n_candidate = max(block_n_candidate, block_lower_bound)
kernel_options["BLOCK_M"] = block_m_candidate kernel_options["BLOCK_M"] = block_m_candidate
kernel_options["BLOCK_N"] = block_n_candidate kernel_options["BLOCK_N"] = block_n_candidate