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, 05:30 UTC
[Bugfix] Fixing division by zero in triton_attn if query_heads/kv_heads > 16 (#23424)
Signed-off-by: Burkhard Ringlein <ngl@zurich.ibm.com>
Burkhard Ringlein committed
Sep 3, 2025 at 17:01 UTC
6d80ae83e1455cb0e47196cea557398fde0f03d1
1 file changed
+2
-1
vllm/attention/ops/triton_unified_attention.py
+2
-1
@@ -674,7 +674,8 @@ def unified_attention(
674
num_queries_per_kv = num_query_heads // num_kv_heads
675
head_size = q.shape[2]
676
677
- BLOCK_M = 16
677
+ BLOCK_M = 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(
678
+ num_queries_per_kv)
679
BLOCK_Q = BLOCK_M // num_queries_per_kv
680
681
# Ideally we would launch with kernel with: