[Kernel] Raise verbose error and consolidate num_heads/num_kv_heads divisibility check (#19339)

Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>

Author: 22quinn

tests/kernels/attention/test_attention.py

@@ -10,6 +10,7 @@
 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
 from tests.kernels.utils import opcheck
 from vllm import _custom_ops as ops
+from vllm.attention.layer import Attention, MultiHeadAttention
 from vllm.platforms import current_platform
 from vllm.utils import get_max_shared_memory_bytes
 
@@ -506,3 +507,18 @@ def test_multi_query_kv_attention_with_alibi(
         device,
         use_alibi=True,
     )
+
+
+@pytest.mark.parametrize("attention_cls", [Attention, MultiHeadAttention])
+def test_num_heads_not_divisble_by_num_kv_heads(attention_cls: type) -> None:
+    head_size = 64
+    scale = float(1.0 / (head_size**0.5))
+    num_heads = 16
+    num_kv_heads = 5
+    with pytest.raises(AssertionError):
+        _ = attention_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=scale,
+            num_kv_heads=num_kv_heads,
+        )

vllm/attention/backends/blocksparse_attn.py

@@ -65,7 +65,6 @@ def __post_init__(self):
         assert self.block_size > 0
         assert self.local_blocks >= 0
         assert self.vert_stride >= 1
-        assert self.num_heads % self.num_kv_heads == 0
 
         tp_size = get_tensor_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank()
@@ -329,9 +328,8 @@ def __init__(
         self.head_size = head_size
         self.scale = float(scale)
         self.alibi_slopes = alibi_slopes
-        self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
+        self.num_kv_heads = num_kv_heads
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
         self.local_blocks = self.blocksparse_params.local_blocks

vllm/attention/backends/dual_chunk_flash_attn.py

@@ -307,7 +307,6 @@ def __init__(
                                if sliding_window is not None else (-1, -1))
         self.kv_cache_dtype = kv_cache_dtype
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         if sliding_window is not None:
             # NOTE(woosuk): flash-attn's sliding window does not work with

vllm/attention/backends/flash_attn.py

@@ -654,7 +654,6 @@ def __init__(
             logits_soft_cap = 0
         self.logits_soft_cap = logits_soft_cap
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
         support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()

vllm/attention/backends/flashinfer.py

@@ -957,7 +957,6 @@ def __init__(
         self.kv_cache_dtype = kv_cache_dtype
         self.logits_soft_cap = logits_soft_cap
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
         if attn_type != AttentionType.DECODER:

vllm/attention/backends/hpu_attn.py

@@ -148,7 +148,6 @@ def __init__(
             alibi_slopes_tensor = torch.tensor(alibi_slopes,
                                                dtype=torch.bfloat16)
             self.alibi_slopes = alibi_slopes_tensor
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
         if self.prefill_impl == 'fsdpa':

vllm/attention/backends/ipex_attn.py

@@ -145,7 +145,6 @@ def __init__(
         self.sliding_window = sliding_window
         self.kv_cache_dtype = kv_cache_dtype
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.need_mask = (self.sliding_window is not None)
         if logits_soft_cap is None:

vllm/attention/backends/pallas.py

@@ -121,9 +121,8 @@ def __init__(
         self.num_heads = num_heads
         self.head_size = head_size
         self.scale = float(scale)
-        self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
+        self.num_kv_heads = num_kv_heads
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.logits_soft_cap = logits_soft_cap
         if head_size % 128 != 0:

vllm/attention/backends/rocm_flash_attn.py

@@ -528,7 +528,6 @@ def __init__(
                                if sliding_window is not None else (-1, -1))
         self.kv_cache_dtype = kv_cache_dtype
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
         self.paged_attn_module = _get_paged_attn_module()

vllm/attention/backends/torch_sdpa.py

@@ -433,7 +433,6 @@ def __init__(
         self.sliding_window = sliding_window
         self.kv_cache_dtype = kv_cache_dtype
 
-        assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.need_mask = (self.alibi_slopes is not None
                           or self.sliding_window is not None)