A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 00:30 UTC
[Bugfix] Avoid too small block m/n for FlexAttention kernel option (#27853)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Oct 31, 2025 at 19:33 UTC
3933f18a5e7b69b096d4b8f700dfa496e6716d86
1 file changed
+5
vllm/v1/attention/backends/flex_attention.py
+5
@@ -896,6 +896,8 @@ def get_kernel_options(
896
return kernel_options
897
else:
898
preferred_block = 32 if query.dtype == torch.float32 else 64
899
+ block_lower_bound = 16
900
+
901
block_m_candidate = ensure_divisible(preferred_block, block_m)
902
block_n_candidate = ensure_divisible(preferred_block, block_n)
903
@@ -910,6 +912,9 @@ def get_kernel_options(
912
max(1, block_n_candidate // 2), block_n
913
)
914
915
+ block_m_candidate = max(block_m_candidate, block_lower_bound)
916
+ block_n_candidate = max(block_n_candidate, block_lower_bound)
917
+
918
kernel_options["BLOCK_M"] = block_m_candidate
919
kernel_options["BLOCK_N"] = block_n_candidate
920