@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0b73736a0d

Fork 0

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 15, 06:00 UTC

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

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

22quinn committed Jun 14, 2025 at 22:43 UTC 0b73736a0d86a2fd0c2548a52eb3877611fa5915

17 files changed +24 -19

tests/kernels/attention/test_attention.py

+16

@@ -10,6 +10,7 @@ import torch

10 from tests.kernels.allclose_default import get_default_atol, get_default_rtol

11 from tests.kernels.utils import opcheck

12 from vllm import _custom_ops as ops

13 +from vllm.attention.layer import Attention, MultiHeadAttention

14 from vllm.platforms import current_platform

15 from vllm.utils import get_max_shared_memory_bytes

@@ -506,3 +507,18 @@ def test_multi_query_kv_attention_with_alibi(

507 device,

508 use_alibi=True,

509 )

510 +

511 +

512 +@pytest.mark.parametrize("attention_cls", [Attention, MultiHeadAttention])

513 +def test_num_heads_not_divisble_by_num_kv_heads(attention_cls: type) -> None:

514 + head_size = 64

515 + scale = float(1.0 / (head_size**0.5))

516 + num_heads = 16

517 + num_kv_heads = 5

518 + with pytest.raises(AssertionError):

519 + _ = attention_cls(

520 + num_heads=num_heads,

521 + head_size=head_size,

522 + scale=scale,

523 + num_kv_heads=num_kv_heads,

524 + )

vllm/attention/backends/blocksparse_attn.py

+1 -3

@@ -65,7 +65,6 @@ class BlocksparseParams:

65 assert self.block_size > 0

66 assert self.local_blocks >= 0

67 assert self.vert_stride >= 1

68 - assert self.num_heads % self.num_kv_heads == 0

69 tp_size = get_tensor_model_parallel_world_size()

70 tp_rank = get_tensor_model_parallel_rank()

@@ -329,9 +328,8 @@ class BlocksparseFlashAttentionImpl(AttentionImpl):

328 self.head_size = head_size

329 self.scale = float(scale)

330 self.alibi_slopes = alibi_slopes

332 - self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads

331 + self.num_kv_heads = num_kv_heads

332

334 - assert self.num_heads % self.num_kv_heads == 0

333 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

334

335 self.local_blocks = self.blocksparse_params.local_blocks

vllm/attention/backends/dual_chunk_flash_attn.py

-1

@@ -307,7 +307,6 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):

307 if sliding_window is not None else (-1, -1))

308 self.kv_cache_dtype = kv_cache_dtype

309

310 - assert self.num_heads % self.num_kv_heads == 0

310 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

311 if sliding_window is not None:

312 # NOTE(woosuk): flash-attn's sliding window does not work with

vllm/attention/backends/flash_attn.py

-1

@@ -654,7 +654,6 @@ class FlashAttentionImpl(AttentionImpl):

654 logits_soft_cap = 0

655 self.logits_soft_cap = logits_soft_cap

656

657 - assert self.num_heads % self.num_kv_heads == 0

657 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

658

659 support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()

vllm/attention/backends/flashinfer.py

-1

@@ -957,7 +957,6 @@ class FlashInferImpl(AttentionImpl):

957 self.kv_cache_dtype = kv_cache_dtype

958 self.logits_soft_cap = logits_soft_cap

959

960 - assert self.num_heads % self.num_kv_heads == 0

960 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

961

962 if attn_type != AttentionType.DECODER:

vllm/attention/backends/hpu_attn.py

-1

@@ -148,7 +148,6 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

148 alibi_slopes_tensor = torch.tensor(alibi_slopes,

149 dtype=torch.bfloat16)

150 self.alibi_slopes = alibi_slopes_tensor

151 - assert self.num_heads % self.num_kv_heads == 0

151 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

152

153 if self.prefill_impl == 'fsdpa':

vllm/attention/backends/ipex_attn.py

-1

@@ -145,7 +145,6 @@ class IpexAttnBackendImpl(AttentionImpl[IpexAttnMetadata]):

145 self.sliding_window = sliding_window

146 self.kv_cache_dtype = kv_cache_dtype

147

148 - assert self.num_heads % self.num_kv_heads == 0

148 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

149 self.need_mask = (self.sliding_window is not None)

150 if logits_soft_cap is None:

vllm/attention/backends/pallas.py

+1 -2

@@ -121,9 +121,8 @@ class PallasAttentionBackendImpl(AttentionImpl):

121 self.num_heads = num_heads

122 self.head_size = head_size

123 self.scale = float(scale)

124 - self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads

124 + self.num_kv_heads = num_kv_heads

125

126 - assert self.num_heads % self.num_kv_heads == 0

126 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

127 self.logits_soft_cap = logits_soft_cap

128 if head_size % 128 != 0:

vllm/attention/backends/rocm_flash_attn.py

-1

@@ -528,7 +528,6 @@ class ROCmFlashAttentionImpl(AttentionImpl):

528 if sliding_window is not None else (-1, -1))

529 self.kv_cache_dtype = kv_cache_dtype

530

531 - assert self.num_heads % self.num_kv_heads == 0

531 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

532

533 self.paged_attn_module = _get_paged_attn_module()

vllm/attention/backends/torch_sdpa.py

-1

@@ -433,7 +433,6 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):

433 self.sliding_window = sliding_window

434 self.kv_cache_dtype = kv_cache_dtype

435

436 - assert self.num_heads % self.num_kv_heads == 0

436 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

437 self.need_mask = (self.alibi_slopes is not None

438 or self.sliding_window is not None)

vllm/attention/backends/xformers.py

-1

@@ -415,7 +415,6 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):

415 self.sliding_window = sliding_window

416 self.kv_cache_dtype = kv_cache_dtype

417

418 - assert self.num_heads % self.num_kv_heads == 0

418 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

419

420 supported_head_sizes = PagedAttention.get_supported_head_sizes()

vllm/attention/layer.py

+6 -1

@@ -80,6 +80,9 @@ class Attention(nn.Module):

80 calculate_kv_scales = False

81 if num_kv_heads is None:

82 num_kv_heads = num_heads

83 + assert num_heads % num_kv_heads == 0, \

84 + f"num_heads ({num_heads}) is not " \

85 + f"divisible by num_kv_heads ({num_kv_heads})"

87 # The default k/v_scale is set to 1.0. This is ignored

88 # when kv-cache is not fp8, and should be used with

@@ -291,7 +294,9 @@ class MultiHeadAttention(nn.Module):

294 self.scale = scale

295 self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads

296

294 - assert self.num_heads % self.num_kv_heads == 0

297 + assert self.num_heads % self.num_kv_heads == 0, \

298 + f"num_heads ({self.num_heads}) is not " \

299 + f"divisible by num_kv_heads ({self.num_kv_heads})"

300 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

301

302 dtype = torch.get_default_dtype()

vllm/v1/attention/backends/flash_attn.py

-1

@@ -545,7 +545,6 @@ class FlashAttentionImpl(AttentionImpl):

545 self.logits_soft_cap = logits_soft_cap

546 self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

547

548 - assert self.num_heads % self.num_kv_heads == 0

548 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

549

550 support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()

vllm/v1/attention/backends/flashinfer.py

-1

@@ -532,7 +532,6 @@ class FlashInferImpl(AttentionImpl):

532 self.logits_soft_cap = logits_soft_cap

533 self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

534

535 - assert self.num_heads % self.num_kv_heads == 0

535 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

536

537 if attn_type != AttentionType.DECODER:

vllm/v1/attention/backends/flex_attention.py

-1

@@ -376,7 +376,6 @@ class FlexAttentionImpl(AttentionImpl):

376 raise NotImplementedError(

377 "FlexAttention does not support logits soft cap yet.")

378

379 - assert self.num_heads % self.num_kv_heads == 0

379 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

380

381 if kv_sharing_target_layer_name is not None:

vllm/v1/attention/backends/pallas.py

-1

@@ -131,7 +131,6 @@ class PallasAttentionBackendImpl(AttentionImpl):

131 self.logits_soft_cap = logits_soft_cap

132 self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

133

134 - assert self.num_heads % self.num_kv_heads == 0

134 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

135 if head_size % 128 != 0:

136 raise NotImplementedError("Head size must be a multiple of 128.")

vllm/v1/attention/backends/triton_attn.py

-1

@@ -114,7 +114,6 @@ class TritonAttentionImpl(AttentionImpl):

114

115 self.use_irope = use_irope

116

117 - assert self.num_heads % self.num_kv_heads == 0

117 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

118

119 support_head_sizes = TritonAttentionBackend.get_supported_head_sizes()