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][CI] Fix v1 attention backend tests and add CI coverage (#26597)
Signed-off-by: Mohammad Miadh Angkad <MAngkad.BSDSBA2027@aim.edu> Signed-off-by: Mohammad Miadh Angkad <mangkad.bsdsba2027@aim.edu> Co-authored-by: Ye (Charlotte) Qi <yeq@meta.com>
Mohammad Miadh Angkad committed
Oct 28, 2025 at 23:42 UTC
a8c02fb5bf2e3191ec93af7c671697583a2e0c56
4 files changed
+67
-26
.buildkite/test-pipeline.yaml
+9
@@ -313,6 +313,15 @@ steps:
313
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
314
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
315
316
+- label: V1 Test attention (H100) # 10min
317
+ timeout_in_minutes: 30
318
+ gpu: h100
319
+ source_file_dependencies:
320
+ - vllm/v1/attention
321
+ - tests/v1/attention
322
+ commands:
323
+ - pytest -v -s v1/attention
324
+
325
- label: V1 Test others (CPU) # 5 mins
326
source_file_dependencies:
327
- vllm/
tests/v1/attention/test_attention_backends.py
+22
-14
@@ -428,6 +428,7 @@ def _test_backend_correctness(
428
# [num_blocks, 2, block_size, num_kv_heads, head_size]
429
# Select the appropriate KV cache format for each backend
430
kv_cache_for_backend = kv_cache
431
+ reset_kv_cache_layout = False
432
if backend_name in (_Backend.FLASHINFER, _Backend.TRITON_ATTN):
433
kv_cache_for_backend = kv_cache.transpose(0, 1)
434
@@ -437,20 +438,27 @@ def _test_backend_correctness(
438
kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)
439
)
440
set_kv_cache_layout("HND")
440
-
441
- backend_output = run_attention_backend(
442
- backend_name,
443
- kv_cache_spec,
444
- ["placeholder"],
445
- vllm_config,
446
- device,
447
- common_attn_metadata,
448
- query_vllm,
449
- key_vllm,
450
- value_vllm,
451
- kv_cache_for_backend,
452
- sliding_window=sliding_window,
453
- )
441
+ reset_kv_cache_layout = True
442
+ elif backend_name == _Backend.TRITON_ATTN:
443
+ kv_cache_for_backend = kv_cache_for_backend.contiguous()
444
+
445
+ try:
446
+ backend_output = run_attention_backend(
447
+ backend_name,
448
+ kv_cache_spec,
449
+ ["placeholder"],
450
+ vllm_config,
451
+ device,
452
+ common_attn_metadata,
453
+ query_vllm,
454
+ key_vllm,
455
+ value_vllm,
456
+ kv_cache_for_backend,
457
+ sliding_window=sliding_window,
458
+ )
459
+ finally:
460
+ if reset_kv_cache_layout:
461
+ set_kv_cache_layout(None)
462
463
# Check shape and dtype consistency
464
assert backend_output.shape == sdpa_output.shape, (
tests/v1/attention/test_mla_backends.py
+2
-1
@@ -155,7 +155,7 @@ def create_and_prepopulate_kv_cache(
155
scale_tensor = scale_tensor.to(device=device, dtype=torch.float32)
156
else:
157
# Create MLA KV cache: (num_blocks, block_size, head_size)
158
- kv_cache = torch.empty(
158
+ kv_cache = torch.zeros(
159
num_blocks, block_size, head_size, dtype=dtype, device=device
160
)
161
kv_cache_flat = kv_cache.view(-1, head_size)
@@ -212,6 +212,7 @@ def create_and_prepopulate_kv_cache(
212
start = start_block_idx
213
end = start + num_blocks_for_seq
214
block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
215
+ block_table[i, num_blocks_for_seq:] = 0
216
start_block_idx += num_blocks_for_seq
217
218
# Create a realistic slot mapping that corresponds to the block table
vllm/v1/attention/backends/flex_attention.py
+34
-11
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with FlexAttention."""
4
5
+import math
6
from dataclasses import dataclass
7
8
import torch
@@ -592,9 +593,10 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
593
self.headdim = self.model_config.get_head_size()
594
self.block_size = kv_cache_spec.block_size
595
self.kv_cache_spec = kv_cache_spec
595
- self.direct_build: bool = is_torch_equal_or_newer("2.9.0.dev0")
596
- self.q_block_size: int = 16 if is_torch_equal_or_newer("2.9.0.dev0") else 128
597
- self.kv_block_size: int = 16 if is_torch_equal_or_newer("2.9.0.dev0") else 128
596
+ supports_small_blocks = is_torch_equal_or_newer("2.9.0.dev0")
597
+ self.direct_build: bool = supports_small_blocks
598
+ self.q_block_size: int = 16 if supports_small_blocks else 128
599
+ self.kv_block_size: int = self.block_size if supports_small_blocks else 128
600
601
def build(
602
self,
@@ -867,6 +869,22 @@ def get_kernel_options(
869
kernel_options: dict[str, int | bool] = {
870
"FORCE_USE_FLEX_ATTENTION": True,
871
}
872
+
873
+ def ensure_divisible(candidate: int, block_size: int) -> int:
874
+ """Pick a kernel block size that divides the logical block."""
875
+ if block_size <= 0:
876
+ return candidate
877
+ candidate = min(candidate, block_size)
878
+ if candidate <= 0:
879
+ return block_size
880
+ if block_size % candidate == 0:
881
+ return candidate
882
+
883
+ candidate = math.gcd(candidate, block_size)
884
+ if candidate <= 1:
885
+ return block_size
886
+ return candidate
887
+
888
if vllm_is_batch_invariant():
889
kernel_options["BLOCK_M"] = 16
890
kernel_options["BLOCK_N"] = 16
@@ -877,17 +895,22 @@ def get_kernel_options(
895
kernel_options["BLOCK_N"] = block_n
896
return kernel_options
897
else:
880
- kernel_options["BLOCK_M"] = 64
881
- kernel_options["BLOCK_N"] = 64
882
- if query.dtype == torch.float32:
883
- kernel_options["BLOCK_M"] = 32
884
- kernel_options["BLOCK_N"] = 32
885
- # if current_platform.is_cuda():
898
+ preferred_block = 32 if query.dtype == torch.float32 else 64
899
+ block_m_candidate = ensure_divisible(preferred_block, block_m)
900
+ block_n_candidate = ensure_divisible(preferred_block, block_n)
901
+
902
if torch.cuda.is_available():
903
device_props = torch.cuda.get_device_properties()
904
max_shared_memory = device_props.shared_memory_per_block_optin
905
if max_shared_memory < 144 * 1024:
890
- kernel_options["BLOCK_M"] = kernel_options["BLOCK_M"] // 2
891
- kernel_options["BLOCK_N"] = kernel_options["BLOCK_N"] // 2
906
+ block_m_candidate = ensure_divisible(
907
+ max(1, block_m_candidate // 2), block_m
908
+ )
909
+ block_n_candidate = ensure_divisible(
910
+ max(1, block_n_candidate // 2), block_n
911
+ )
912
+
913
+ kernel_options["BLOCK_M"] = block_m_candidate
914
+ kernel_options["BLOCK_N"] = block_n_candidate
915
916
return kernel_options