@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 402759d472
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 14, 05:30 UTC

[Attention] FlashAttn MLA (#14258)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com> Co-authored-by: Matthew Bonanni <mbonanni001@gmail.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Lucas Wilkinson committed Sep 4, 2025 at 05:47 UTC 402759d4727d9a377598a09d06770770d4e184c6
22 files changed +491 -211
.buildkite/check-wheel-size.py
+4 -4
@@ -5,11 +5,11 @@ import os
5 import sys
6 import zipfile
7
8 -# Read the VLLM_MAX_SIZE_MB environment variable, defaulting to 400 MiB
9 -# Note that we have 400 MiB quota, please use it wisely.
10 -# See https://github.com/pypi/support/issues/3792 .
8 +# Read the VLLM_MAX_SIZE_MB environment variable, defaulting to 450 MiB
9 +# Note that we have 800 MiB quota, please use it wisely.
10 +# See https://github.com/pypi/support/issues/6326 .
11 # Please also sync the value with the one in Dockerfile.
12 -VLLM_MAX_SIZE_MB = int(os.environ.get("VLLM_MAX_SIZE_MB", 400))
12 +VLLM_MAX_SIZE_MB = int(os.environ.get("VLLM_MAX_SIZE_MB", 450))
13
14
15 def print_top_10_largest_files(zip_file):
cmake/external_projects/vllm_flash_attn.cmake
+1 -1
@@ -38,7 +38,7 @@ else()
38 FetchContent_Declare(
39 vllm-flash-attn
40 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
41 - GIT_TAG 57b4e68b9f9d94750b46de8f8dbd2bfcc86edd4f
41 + GIT_TAG ee4d25bd84e0cbc7e0b9b9685085fd5db2dcb62a
42 GIT_PROGRESS TRUE
43 # Don't share the vllm-flash-attn build between build types
44 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
docker/Dockerfile
+1 -1
@@ -237,7 +237,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
237 # Check the size of the wheel if RUN_WHEEL_CHECK is true
238 COPY .buildkite/check-wheel-size.py check-wheel-size.py
239 # sync the default value with .buildkite/check-wheel-size.py
240 -ARG VLLM_MAX_SIZE_MB=400
240 +ARG VLLM_MAX_SIZE_MB=450
241 ENV VLLM_MAX_SIZE_MB=$VLLM_MAX_SIZE_MB
242 ARG RUN_WHEEL_CHECK=true
243 RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \
tests/kernels/attention/test_attention_selector.py
+83 -27
@@ -22,7 +22,7 @@ def clear_cache():
22
23 # Define MLA and non-MLA backends separately
24 DEVICE_MLA_BACKENDS = {
25 - "cuda": ["TRITON_MLA", "FLASHMLA"],
25 + "cuda": ["TRITON_MLA", "FLASHMLA", "FLASH_ATTN_MLA", "CUTLASS_MLA"],
26 "hip": ["TRITON_MLA", "ROCM_AITER_MLA"],
27 "cpu": [],
28 }
@@ -98,21 +98,14 @@ def test_env(
98 with patch("vllm.attention.selector.current_platform",
99 RocmPlatform()):
100 if use_mla:
101 - # Validate HIP MLA backend-block_size combinations
102 - valid_combination = (
103 - (name == "TRITON_MLA" and block_size != 1)
104 - or (name == "ROCM_AITER_MLA" and block_size == 1))
105 -
106 - if valid_combination:
107 - backend = get_attn_backend(16,
108 - torch.float16,
109 - torch.float16,
110 - block_size,
111 - False,
112 - use_mla=use_mla)
113 - expected = f"{name}_VLLM_V1" if use_v1 else name
114 - assert backend.get_name() == expected
115 - else:
101 + # ROCm MLA backend logic:
102 + # - TRITON_MLA: supported when block_size != 1
103 + # - ROCM_AITER_MLA: supported when block_size == 1
104 + # If backend is forced but doesn't match block_size,
105 + # should raise ValueError
106 +
107 + if name == "TRITON_MLA" and block_size == 1:
108 + # TRITON_MLA doesn't support block_size == 1
109 with pytest.raises(ValueError) as exc_info:
110 get_attn_backend(16,
111 torch.float16,
@@ -122,6 +115,27 @@ def test_env(
115 use_mla=use_mla)
116 assert f"The selected backend, {name}" in str(
117 exc_info.value)
118 + elif name == "ROCM_AITER_MLA" and block_size != 1:
119 + # ROCM_AITER_MLA only supports block_size == 1
120 + with pytest.raises(ValueError) as exc_info:
121 + get_attn_backend(16,
122 + torch.float16,
123 + torch.float16,
124 + block_size,
125 + False,
126 + use_mla=use_mla)
127 + assert f"The selected backend, {name}" in str(
128 + exc_info.value)
129 + else:
130 + # Valid backend-block_size combination
131 + backend = get_attn_backend(16,
132 + torch.float16,
133 + torch.float16,
134 + block_size,
135 + False,
136 + use_mla=use_mla)
137 + expected = f"{name}_VLLM_V1" if use_v1 else name
138 + assert backend.get_name() == expected
139 else:
140 backend = get_attn_backend(16,
141 torch.float16,
@@ -136,16 +150,57 @@ def test_env(
150 with patch("vllm.attention.selector.current_platform",
151 CudaPlatform()):
152 if use_mla:
139 - if name == "FLASHMLA" and block_size == 64:
140 - from vllm.attention.backends.flashmla import (
141 - is_flashmla_supported)
142 -
143 - # only on cuda platforms with specific capability.
144 - is_supported, _ = is_flashmla_supported()
145 -
146 - if not is_supported:
147 - # if platform is not supported then skip this case.
148 - pytest.skip()
153 + # CUDA MLA backend logic:
154 + # - CUTLASS_MLA: only supported with block_size == 128
155 + # and Blackwell GPUs (SM 10.0), V1 only
156 + # - FLASHMLA: only supported with block_size == 64
157 + # - FLASH_ATTN_MLA: V1 only
158 + # - TRITON_MLA: fallback for other cases
159 +
160 + if name == "CUTLASS_MLA":
161 + if not use_v1:
162 + # CUTLASS_MLA only supported on V1 engine
163 + pytest.skip(
164 + "CUTLASS_MLA only supported on V1 engine")
165 + elif block_size != 128:
166 + # CUTLASS_MLA only supports block_size == 128
167 + pytest.skip(
168 + "CUTLASS_MLA only supports block_size 128")
169 + else:
170 + backend = get_attn_backend(16,
171 + torch.float16,
172 + torch.float16,
173 + block_size,
174 + False,
175 + use_mla=use_mla)
176 + expected = "CUTLASS_MLA_VLLM_V1"
177 + assert backend.get_name() == expected
178 + elif name == "FLASHMLA":
179 + if block_size != 64:
180 + # FlashMLA only supports block_size == 64
181 + pytest.skip("FlashMLA only supports block_size 64")
182 + else:
183 + from vllm.attention.backends.flashmla import (
184 + is_flashmla_supported)
185 + is_supported, _ = is_flashmla_supported()
186 + if not is_supported:
187 + pytest.skip(
188 + "FlashMLA not supported on this platform")
189 + else:
190 + backend = get_attn_backend(16,
191 + torch.float16,
192 + torch.float16,
193 + block_size,
194 + False,
195 + use_mla=use_mla)
196 + expected = f"{name}_VLLM_V1" if use_v1 else name
197 + assert backend.get_name() == expected
198 + elif name == "FLASH_ATTN_MLA":
199 + if not use_v1:
200 + # FlashAttention MLA only supported on V1 engine
201 + pytest.skip(
202 + "FlashAttention MLA only supported on V1 engine"
203 + )
204 else:
205 backend = get_attn_backend(16,
206 torch.float16,
@@ -153,9 +208,10 @@ def test_env(
208 block_size,
209 False,
210 use_mla=use_mla)
156 - expected = f"{name}_VLLM_V1" if use_v1 else name
211 + expected = "FLASH_ATTN_MLA"
212 assert backend.get_name() == expected
213 else:
214 + # TRITON_MLA or other fallback
215 backend = get_attn_backend(16,
216 torch.float16,
217 torch.float16,
tests/v1/attention/test_attention_backends.py
-16
@@ -70,22 +70,6 @@ BATCH_SPECS = {
70 }
71
72
73 -def create_dummy_kv_cache(kv_cache_spec: FullAttentionSpec,
74 - device: torch.device,
75 - num_blocks: int = 100) -> torch.Tensor:
76 - """Create a dummy KV cache tensor for testing."""
77 - kv_cache = torch.randn(
78 - 2, # K and V
79 - num_blocks,
80 - kv_cache_spec.block_size,
81 - kv_cache_spec.num_kv_heads,
82 - kv_cache_spec.head_size,
83 - dtype=_convert_dtype_to_torch(kv_cache_spec.dtype),
84 - device=device,
85 - )
86 - return kv_cache
87 -
88 -
73 def create_and_prepopulate_kv_cache(
74 k_contexts: list[torch.Tensor],
75 v_contexts: list[torch.Tensor],
tests/v1/attention/test_mla_backends.py
+101 -101
@@ -15,7 +15,7 @@ from vllm.v1.attention.backends.utils import CommonAttentionMetadata
15 from vllm.v1.kv_cache_interface import FullAttentionSpec
16
17 BACKENDS_TO_TEST = [
18 - _Backend.CUTLASS_MLA, _Backend.FLASHMLA_VLLM_V1,
18 + _Backend.CUTLASS_MLA, _Backend.FLASHMLA_VLLM_V1, _Backend.FLASH_ATTN_MLA,
19 _Backend.TRITON_MLA_VLLM_V1
20 ]
21
@@ -69,20 +69,6 @@ BATCH_SPECS = {
69 }
70
71
72 -def create_dummy_kv_cache(kv_cache_spec: FullAttentionSpec,
73 - device: torch.device,
74 - num_blocks: int = 100) -> torch.Tensor:
75 - """Create a dummy KV cache tensor for testing."""
76 - kv_cache = torch.randn(
77 - num_blocks,
78 - kv_cache_spec.block_size,
79 - kv_cache_spec.head_size, # latent dimension
80 - dtype=_convert_dtype_to_torch(kv_cache_spec.dtype),
81 - device=device,
82 - )
83 - return kv_cache
84 -
85 -
72 def create_and_prepopulate_kv_cache(
73 kv_c_contexts: list[torch.Tensor],
74 k_pe_contexts: list[torch.Tensor],
@@ -315,7 +301,7 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
301
302 # 2. Generate data and compute SDPA reference output for MLA
303 all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
318 - all_sdpa_outputs = []
304 + all_sdpa_outputs: list[list[torch.Tensor]] = []
305 kv_c_contexts, k_pe_contexts = [], []
306
307 # Create shared MLA weight matrices for consistency across all sequences
@@ -331,6 +317,9 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
317 device=device)
318 kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)
319
320 + for i, backend in enumerate(BACKENDS_TO_TEST):
321 + all_sdpa_outputs.append([])
322 +
323 for i in range(batch_size):
324 s_len = seq_lens[i]
325 q_len = query_lens[i]
@@ -358,85 +347,93 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
347 dtype=dtype,
348 device=device)
349
361 - # Determine if this is decode (single token)
362 - # or prefill (multiple tokens)
363 - is_decode = q_len == 1
350 + # Determine if this is decode or prefill
351 + is_decode = []
352 + for i, backend in enumerate(BACKENDS_TO_TEST):
353 + builder_cls, _ = get_attention_backend(backend)
354 + is_decode.append(q_len <= builder_cls.reorder_batch_threshold)
355
356 # Split q into nope and rope components
357 q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
358
368 - if is_decode:
369 - # Decode path: MQA-style attention in latent space
370 - # Transform q_nope to latent space: q_nope @ W_UK
371 - # q_nope: [1, num_heads, qk_nope_head_dim]
372 - # W_UK: [kv_lora_rank, num_heads, qk_nope_head_dim]
373 - ql_nope = torch.einsum("qnh,lnh->qnl", q_nope,
374 - W_UK) # [1, num_heads, kv_lora_rank]
375 -
376 - # Build MQA attention inputs
377 - # Q: [1, num_heads, kv_lora_rank + qk_rope_head_dim]
378 - q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
379 - # K: [s_len, kv_lora_rank + qk_rope_head_dim]
380 - # (broadcasted to all heads)
381 - k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
382 - k_mqa = k_mqa.unsqueeze(1).expand(-1, num_q_heads, -1)
383 - # V: [s_len, kv_lora_rank] (broadcasted to all heads)
384 - v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_q_heads, -1)
385 -
386 - # SDPA expects (N, H, L, D)
387 - q_sdpa_in = q_mqa.unsqueeze(0).transpose(1, 2)
388 - k_sdpa_in = k_mqa.unsqueeze(0).transpose(1, 2)
389 - v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)
390 -
391 - sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(
392 - q_sdpa_in, k_sdpa_in, v_sdpa_in, is_causal=False, scale=scale)
393 - sdpa_out_i = sdpa_out_i.transpose(1, 2).squeeze(
394 - 0) # [1, num_heads, kv_lora_rank]
395 -
396 - # Project back to output space: sdpa_out @ W_UV
397 - sdpa_out_i = torch.einsum("qnl,lnv->qnv", sdpa_out_i, W_UV)
398 - sdpa_out_i = sdpa_out_i.flatten(start_dim=-2)
399 - else:
400 - # Prefill path: MHA-style attention with full sequence
401 - # Apply kv_b_proj to the full kv_c tensor
402 - kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_full,
403 - kv_b_proj_weight)
404 - k_nope_full, v_full = kv_nope_full.split(
405 - [qk_nope_head_dim, v_head_dim], dim=-1)
406 -
407 - # Build attention inputs for full sequence
408 - q_mha = torch.cat([q_nope, q_pe],
409 - dim=-1) # [q_len, num_heads, total_dim]
410 - k_pe_full_expanded = k_pe_full.expand(-1, num_q_heads, -1)
411 - k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1)
412 -
413 - # Create custom attention mask:
414 - # - Query tokens can attend to all context tokens
415 - # - Query tokens can only attend to query tokens up to their pos
416 - attn_mask = torch.ones(q_len,
417 - s_len,
418 - dtype=torch.bool,
419 - device=device)
420 - # Apply causal mask only to the query portion (context_len onwards)
421 - causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
422 - attn_mask[:, context_len:] = causal_mask
423 -
424 - # SDPA expects (N, H, L, D)
425 - q_sdpa_in = q_mha.unsqueeze(0).transpose(1, 2)
426 - k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
427 - v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
428 -
429 - # Single attention call with custom mask
430 - sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(
431 - q_sdpa_in,
432 - k_sdpa_in,
433 - v_sdpa_in,
434 - attn_mask=attn_mask,
435 - scale=scale)
436 - sdpa_out_i = sdpa_out_i.transpose(1, 2).squeeze(0)
437 - sdpa_out_i = sdpa_out_i.flatten(start_dim=-2)
438 -
439 - all_sdpa_outputs.append(sdpa_out_i)
359 + #######################################################
360 + # Decode path: MQA-style attention in latent space
361 + # Transform q_nope to latent space: q_nope @ W_UK
362 + # q_nope: [1, num_heads, qk_nope_head_dim]
363 + # W_UK: [kv_lora_rank, num_heads, qk_nope_head_dim]
364 + ql_nope = torch.einsum("qnh,lnh->qnl", q_nope,
365 + W_UK) # [1, num_heads, kv_lora_rank]
366 +
367 + # Build MQA attention inputs
368 + # Q: [1, num_heads, kv_lora_rank + qk_rope_head_dim]
369 + q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
370 + # K: [s_len, kv_lora_rank + qk_rope_head_dim]
371 + # (broadcasted to all heads)
372 + k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
373 + k_mqa = k_mqa.unsqueeze(1).expand(-1, num_q_heads, -1)
374 + # V: [s_len, kv_lora_rank] (broadcasted to all heads)
375 + v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_q_heads, -1)
376 +
377 + # Create custom attention mask for decode path:
378 + # - Query tokens can attend to all context tokens
379 + # - Query tokens can only attend to query tokens up to their position
380 + attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)
381 + # Apply causal mask only to the query portion (context_len onwards)
382 + causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
383 + attn_mask[:, context_len:] = causal_mask
384 +
385 + # SDPA expects (N, H, L, D)
386 + q_sdpa_in = q_mqa.unsqueeze(0).transpose(1, 2)
387 + k_sdpa_in = k_mqa.unsqueeze(0).transpose(1, 2)
388 + v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)
389 +
390 + sdpa_out_i_decode = torch.nn.functional.scaled_dot_product_attention(
391 + q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale)
392 + sdpa_out_i_decode = sdpa_out_i_decode.transpose(1, 2).squeeze(
393 + 0) # [1, num_heads, kv_lora_rank]
394 +
395 + # Project back to output space: sdpa_out @ W_UV
396 + sdpa_out_i_decode = torch.einsum("qnl,lnv->qnv", sdpa_out_i_decode,
397 + W_UV)
398 + sdpa_out_i_decode = sdpa_out_i_decode.flatten(start_dim=-2)
399 +
400 + #######################################################
401 + # Prefill path: MHA-style attention with full sequence
402 + # Apply kv_b_proj to the full kv_c tensor
403 + kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_full, kv_b_proj_weight)
404 + k_nope_full, v_full = kv_nope_full.split(
405 + [qk_nope_head_dim, v_head_dim], dim=-1)
406 +
407 + # Build attention inputs for full sequence
408 + q_mha = torch.cat([q_nope, q_pe],
409 + dim=-1) # [q_len, num_heads, total_dim]
410 + k_pe_full_expanded = k_pe_full.expand(-1, num_q_heads, -1)
411 + k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1)
412 +
413 + # Create custom attention mask:
414 + # - Query tokens can attend to all context tokens
415 + # - Query tokens can only attend to query tokens up to their pos
416 + attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)
417 + # Apply causal mask only to the query portion (context_len onwards)
418 + causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
419 + attn_mask[:, context_len:] = causal_mask
420 +
421 + # SDPA expects (N, H, L, D)
422 + q_sdpa_in = q_mha.unsqueeze(0).transpose(1, 2)
423 + k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
424 + v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
425 +
426 + # Single attention call with custom mask
427 + sdpa_out_i_prefill = torch.nn.functional.scaled_dot_product_attention(
428 + q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale)
429 + sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0)
430 + sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2)
431 +
432 + for i, backend in enumerate(BACKENDS_TO_TEST):
433 + if is_decode[i]:
434 + all_sdpa_outputs[i].append(sdpa_out_i_decode)
435 + else:
436 + all_sdpa_outputs[i].append(sdpa_out_i_prefill)
437
438 # Inputs for vLLM MLA backends are just the new tokens
439 all_q_vllm.append(q_c)
@@ -451,7 +448,9 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
448 query_vllm = torch.cat(all_q_vllm, dim=0)
449 kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
450 k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
454 - sdpa_output = torch.cat(all_sdpa_outputs, dim=0)
451 + sdpa_outputs = []
452 + for i, backend in enumerate(BACKENDS_TO_TEST):
453 + sdpa_outputs.append(torch.cat(all_sdpa_outputs[i], dim=0))
454
455 # Create mock kv_b_proj using the same weights as reference implementation
456 from vllm.model_executor.layers.linear import ColumnParallelLinear
@@ -486,7 +485,7 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
485 randomize_blocks=True)
486
487 # 4. Run vLLM backends and compare
489 - for backend_name in BACKENDS_TO_TEST:
488 + for i, backend_name in enumerate(BACKENDS_TO_TEST):
489 backend_output = run_attention_backend(
490 backend_name, kv_cache_spec, ["placeholder"], vllm_config, device,
491 common_attn_metadata, query_vllm, kv_c_vllm, k_pe_vllm, kv_cache,
@@ -494,12 +493,12 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
493 mock_kv_b_proj)
494
495 # Check shape and dtype consistency
497 - assert backend_output.shape == sdpa_output.shape, (
496 + assert backend_output.shape == sdpa_outputs[i].shape, (
497 f"[{backend_name}] shape {backend_output.shape} != "
499 - f"SDPA shape {sdpa_output.shape}")
500 - assert backend_output.dtype == sdpa_output.dtype, (
498 + f"SDPA shape {sdpa_outputs[i].shape}")
499 + assert backend_output.dtype == sdpa_outputs[i].dtype, (
500 f"[{backend_name}] dtype {backend_output.dtype} != "
502 - f"SDPA dtype {sdpa_output.dtype}")
501 + f"SDPA dtype {sdpa_outputs[i].dtype}")
502
503 assert torch.isfinite(backend_output).all(), (
504 f"[{backend_name}] produced non-finite values")
@@ -508,12 +507,13 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
507 rtol = 1e-2
508 atol = 5e-1
509
511 - max_diff = torch.max(torch.abs(backend_output - sdpa_output)).item()
510 + max_diff = torch.max(torch.abs(backend_output -
511 + sdpa_outputs[i])).item()
512 max_rel_diff = torch.max(
513 - torch.abs(backend_output - sdpa_output) /
514 - torch.abs(sdpa_output)).item()
513 + torch.abs(backend_output - sdpa_outputs[i]) /
514 + torch.abs(sdpa_outputs[i])).item()
515 all_close = torch.allclose(backend_output,
516 - sdpa_output,
516 + sdpa_outputs[i],
517 rtol=rtol,
518 atol=atol)
519
tests/v1/attention/utils.py
+2
@@ -139,6 +139,8 @@ def get_attention_backend(backend_name: _Backend):
139 "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
140 _Backend.FLASHMLA_VLLM_V1:
141 "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
142 + _Backend.FLASH_ATTN_MLA:
143 + "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
144 _Backend.TRITON_MLA_VLLM_V1:
145 "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
146 }
vllm/attention/utils/fa_utils.py
+13
@@ -68,5 +68,18 @@ def flash_attn_supports_fp8() -> bool:
68 current_platform.get_device_capability().major == 9
69
70
71 +def flash_attn_supports_mla():
72 + from vllm.platforms import current_platform
73 + if current_platform.is_cuda():
74 + try:
75 + from vllm.vllm_flash_attn.flash_attn_interface import (
76 + is_fa_version_supported)
77 + return is_fa_version_supported(3) \
78 + and current_platform.get_device_capability()[0] == 9
79 + except (ImportError, AssertionError):
80 + pass
81 + return False
82 +
83 +
84 def is_flash_attn_varlen_func_available() -> bool:
85 return current_platform.is_cuda() or current_platform.is_xpu()
vllm/engine/arg_utils.py
+2
@@ -1488,6 +1488,8 @@ class EngineArgs:
1488 "TRITON_MLA",
1489 "CUTLASS_MLA",
1490 "FLASHMLA",
1491 + "FLASHMLA_VLLM_V1",
1492 + "FLASH_ATTN_MLA",
1493 "FLASHINFER",
1494 "FLASHINFER_VLLM_V1",
1495 "ROCM_AITER_MLA",
vllm/envs.py
+1
@@ -463,6 +463,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
463 # - "ROCM_FLASH": use ROCmFlashAttention
464 # - "FLASHINFER": use flashinfer
465 # - "FLASHMLA": use FlashMLA
466 + # - "FLASH_ATTN_MLA": use FlashAttention for MLA
467 "VLLM_ATTENTION_BACKEND":
468 lambda: os.getenv("VLLM_ATTENTION_BACKEND", None),
469
vllm/platforms/cuda.py
+40 -28
@@ -223,9 +223,30 @@ class CudaPlatformBase(Platform):
223 if use_mla:
224 # TODO(lucas): refactor to be more concise
225 # we should probably consider factoring out V1 here
226 - if selected_backend == _Backend.CUTLASS_MLA or (
227 - cls.is_device_capability(100) and selected_backend is None
228 - and block_size == 128):
226 +
227 + from vllm.attention.ops.flashmla import is_flashmla_supported
228 + from vllm.attention.utils.fa_utils import flash_attn_supports_mla
229 +
230 + use_cutlassmla = selected_backend == _Backend.CUTLASS_MLA or (
231 + selected_backend is None and cls.is_device_capability(100)
232 + and block_size == 128)
233 + use_flashmla = selected_backend in [
234 + _Backend.FLASHMLA, _Backend.FLASHMLA_VLLM_V1
235 + ] or (selected_backend is None and is_flashmla_supported()[0])
236 + use_flashattn = selected_backend == _Backend.FLASH_ATTN_MLA or (
237 + selected_backend is None and flash_attn_supports_mla())
238 + use_triton = selected_backend == _Backend.TRITON_MLA or (
239 + selected_backend is None)
240 +
241 + def _get_version(name, import_suffix) -> str:
242 + if use_v1:
243 + logger.info_once(f"Using {name} backend on V1 engine.")
244 + return f"vllm.v1.attention.backends.mla.{import_suffix}"
245 + else:
246 + logger.info_once(f"Using {name} backend.")
247 + return f"vllm.attention.backends.{import_suffix}"
248 +
249 + if use_cutlassmla:
250 if use_v1:
251 logger.info_once("Using Cutlass MLA backend on V1 engine.")
252 return ("vllm.v1.attention.backends.mla."
@@ -233,36 +254,27 @@ class CudaPlatformBase(Platform):
254 else:
255 logger.warning(
256 "Cutlass MLA backend is only supported on V1 engine")
236 - if selected_backend == _Backend.TRITON_MLA or block_size != 64:
237 - if use_v1:
238 - logger.info_once("Using Triton MLA backend on V1 engine.")
239 - return ("vllm.v1.attention.backends.mla."
240 - "triton_mla.TritonMLABackend")
241 - else:
242 - logger.info("Using Triton MLA backend.")
243 - return "vllm.attention.backends.triton_mla.TritonMLABackend"
244 - else:
245 - from vllm.attention.backends.flashmla import (
246 - is_flashmla_supported)
247 - if not is_flashmla_supported()[0]:
248 - logger.warning(
249 - "FlashMLA backend is not supported due to %s",
250 - is_flashmla_supported()[1])
251 - elif block_size != 64:
257 + if use_flashmla:
258 + if block_size != 64:
259 logger.warning(
260 "FlashMLA backend is not supported for block size %d"
261 " (currently only supports block size 64).",
262 block_size)
263 else:
257 - if use_v1:
258 - logger.info_once(
259 - "Using FlashMLA backend on V1 engine.")
260 - return ("vllm.v1.attention.backends.mla."
261 - "flashmla.FlashMLABackend")
262 - else:
263 - logger.info("Using FlashMLA backend.")
264 - return ("vllm.attention.backends."
265 - "flashmla.FlashMLABackend")
264 + return _get_version("FlashMLA", "flashmla.FlashMLABackend")
265 + if use_flashattn:
266 + if use_v1:
267 + logger.info_once(
268 + "Using FlashAttention MLA backend on V1 engine.")
269 + return ("vllm.v1.attention.backends.mla."
270 + "flashattn_mla.FlashAttnMLABackend")
271 + else:
272 + logger.warning(
273 + "FlashAttention MLA backend is only supported on V1 "
274 + "engine.")
275 + if use_triton:
276 + return _get_version("Triton MLA",
277 + "triton_mla.TritonMLABackend")
278 if use_v1:
279 FLASHINFER_V1 = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" # noqa: E501
280 FLEX_ATTENTION_V1 = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501
vllm/platforms/interface.py
+3 -2
@@ -52,9 +52,10 @@ class _Backend(enum.Enum):
52 FLASHINFER_VLLM_V1 = enum.auto()
53 TRITON_MLA = enum.auto() # Supported by V1
54 TRITON_MLA_VLLM_V1 = enum.auto()
55 - FLASHMLA_VLLM_V1 = enum.auto()
56 - FLASHMLA = enum.auto() # Supported by V1
55 CUTLASS_MLA = enum.auto()
56 + FLASHMLA = enum.auto() # Supported by V1
57 + FLASHMLA_VLLM_V1 = enum.auto()
58 + FLASH_ATTN_MLA = enum.auto() # Supported by V1
59 PALLAS = enum.auto()
60 PALLAS_VLLM_V1 = enum.auto()
61 IPEX = enum.auto()
vllm/v1/attention/backends/flashinfer.py
+2 -1
@@ -317,7 +317,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
317 num_reqs = common_attn_metadata.num_reqs
318 num_actual_tokens = common_attn_metadata.num_actual_tokens
319 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens =\
320 - split_decodes_and_prefills(common_attn_metadata)
320 + split_decodes_and_prefills(common_attn_metadata,
321 + decode_threshold=self.reorder_batch_threshold)
322
323 page_size = self.page_size
324 max_q_len = common_attn_metadata.max_query_len
vllm/v1/attention/backends/linear_attn.py
+3 -2
@@ -52,8 +52,9 @@ class LinearAttentionMetadataBuilder(
52 state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
53
54 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
55 - split_decodes_and_prefills(common_attn_metadata,
56 - decode_threshold=1))
55 + split_decodes_and_prefills(
56 + common_attn_metadata,
57 + decode_threshold=self.reorder_batch_threshold))
58
59 attn_metadata = LinearAttentionMetadata(
60 num_prefills=num_prefills,
vllm/v1/attention/backends/mamba1_attn.py
+3 -2
@@ -50,8 +50,9 @@ class Mamba1AttentionMetadataBuilder(
50 query_start_loc.device)
51
52 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
53 - split_decodes_and_prefills(common_attn_metadata,
54 - decode_threshold=1))
53 + split_decodes_and_prefills(
54 + common_attn_metadata,
55 + decode_threshold=self.reorder_batch_threshold))
56
57 has_initial_states = None
58 padded_decodes = num_decodes
vllm/v1/attention/backends/mamba2_attn.py
+3 -2
@@ -115,8 +115,9 @@ class Mamba2AttentionMetadataBuilder(
115 state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
116
117 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
118 - split_decodes_and_prefills(common_attn_metadata,
119 - decode_threshold=1))
118 + split_decodes_and_prefills(
119 + common_attn_metadata,
120 + decode_threshold=self.reorder_batch_threshold))
121
122 # Compute seq_idx, chunk_indices and chunk_offsets for prefill only
123 if num_prefills > 0:
vllm/v1/attention/backends/mla/common.py
+12 -5
@@ -578,11 +578,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
578 prefill.prefill_main = self._fi_prefill_main
579 prefill.prefill_chunks = self._fi_prefill_chunks
580
581 - def _build_decode(self, block_table_tensor: torch.Tensor,
582 - seq_lens: torch.Tensor):
581 + def _build_decode(
582 + self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,
583 + seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,
584 + query_start_loc_device: torch.Tensor) -> MLACommonDecodeMetadata:
585 return MLACommonDecodeMetadata(
586 block_table=block_table_tensor,
585 - seq_lens=seq_lens,
587 + seq_lens=seq_lens_device,
588 )
589
590 def build_for_cudagraph_capture(
@@ -618,6 +620,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
620 query_start_loc = common_attn_metadata.query_start_loc
621 query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
622 seq_lens = common_attn_metadata.seq_lens
623 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu
624
625 query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
626
@@ -625,7 +628,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
628 query_seq_lens_cpu)
629
630 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = \
628 - split_decodes_and_prefills(common_attn_metadata)
631 + split_decodes_and_prefills(common_attn_metadata,
632 + decode_threshold=self.reorder_batch_threshold)
633
634 assert num_decodes + num_prefills == num_reqs
635 assert num_decode_tokens + num_prefill_tokens == num_tokens
@@ -725,7 +729,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
729 if num_decodes > 0:
730 decode_metadata = self._build_decode(
731 block_table_tensor=block_table_tensor[:num_decodes, ...],
728 - seq_lens=seq_lens[:num_decodes],
732 + seq_lens_cpu=seq_lens_cpu[:num_decodes],
733 + seq_lens_device=seq_lens[:num_decodes],
734 + query_start_loc_cpu=query_start_loc_cpu[:num_decodes + 1],
735 + query_start_loc_device=query_start_loc[:num_decodes + 1],
736 )
737
738 attn_metadata = self.metadata_cls(
vllm/v1/attention/backends/mla/flashattn_mla.py new
+189
@@ -0,0 +1,189 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass
5 +from typing import ClassVar, Optional
6 +
7 +import torch
8 +
9 +from vllm.attention.backends.abstract import (AttentionLayer, AttentionType,
10 + is_quantized_kv_cache)
11 +from vllm.attention.utils.fa_utils import (flash_attn_supports_mla,
12 + get_flash_attn_version)
13 +from vllm.config import VllmConfig
14 +from vllm.logger import init_logger
15 +from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
16 + MLACommonDecodeMetadata,
17 + MLACommonImpl,
18 + MLACommonMetadata,
19 + MLACommonMetadataBuilder)
20 +from vllm.v1.kv_cache_interface import AttentionSpec
21 +from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata
22 +
23 +logger = init_logger(__name__)
24 +
25 +
26 +class FlashAttnMLABackend(MLACommonBackend):
27 +
28 + @staticmethod
29 + def get_name() -> str:
30 + return "FLASH_ATTN_MLA"
31 +
32 + @staticmethod
33 + def get_metadata_cls() -> type["FlashAttnMLAMetadata"]:
34 + return FlashAttnMLAMetadata
35 +
36 + @staticmethod
37 + def get_builder_cls() -> type["FlashAttnMLAMetadataBuilder"]:
38 + return FlashAttnMLAMetadataBuilder
39 +
40 + @staticmethod
41 + def get_impl_cls() -> type["FlashAttnMLAImpl"]:
42 + return FlashAttnMLAImpl
43 +
44 +
45 +@dataclass
46 +class FlashAttnMLADecodeMetadata(MLACommonDecodeMetadata):
47 + query_start_loc: torch.Tensor
48 + max_query_len: int
49 + max_seq_len: int
50 + scheduler_metadata: Optional[torch.Tensor] = None
51 +
52 +
53 +@dataclass
54 +class FlashAttnMLAMetadata(MLACommonMetadata[FlashAttnMLADecodeMetadata]):
55 + pass
56 +
57 +
58 +class FlashAttnMLAMetadataBuilder(
59 + MLACommonMetadataBuilder[FlashAttnMLAMetadata]):
60 + reorder_batch_threshold: ClassVar[int] = 512
61 +
62 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
63 + vllm_config: VllmConfig, device: torch.device):
64 + super().__init__(kv_cache_spec, layer_names, vllm_config, device,
65 + FlashAttnMLAMetadata)
66 + self.fa_aot_schedule = (get_flash_attn_version() == 3)
67 +
68 + def _schedule_decode(self, num_reqs, cu_query_lens, max_query_len, seqlens,
69 + max_seq_len, causal):
70 + if self.fa_aot_schedule:
71 + return get_scheduler_metadata(
72 + batch_size=num_reqs,
73 + max_seqlen_q=max_query_len,
74 + max_seqlen_k=max_seq_len,
75 + num_heads_q=self.num_heads,
76 + num_heads_kv=1,
77 + headdim=self.mla_dims.qk_rope_head_dim,
78 + cache_seqlens=seqlens,
79 + qkv_dtype=self.kv_cache_spec.dtype,
80 + headdim_v=self.mla_dims.kv_lora_rank,
81 + page_size=self.page_size,
82 + cu_seqlens_q=cu_query_lens,
83 + causal=causal,
84 + )
85 + return None
86 +
87 + def _build_decode(
88 + self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,
89 + seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,
90 + query_start_loc_device: torch.Tensor
91 + ) -> FlashAttnMLADecodeMetadata:
92 + query_lens_cpu = (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1])
93 + max_query_len = query_lens_cpu.max().item()
94 + max_seq_len = seq_lens_cpu.max().item()
95 +
96 + scheduler_metadata = self._schedule_decode(
97 + num_reqs=seq_lens_cpu.numel(),
98 + cu_query_lens=query_start_loc_device,
99 + max_query_len=max_query_len,
100 + seqlens=seq_lens_device,
101 + max_seq_len=max_seq_len,
102 + causal=True,
103 + )
104 +
105 + return FlashAttnMLADecodeMetadata(
106 + block_table=block_table_tensor,
107 + seq_lens=seq_lens_device,
108 + query_start_loc=query_start_loc_device,
109 + max_query_len=max_query_len,
110 + max_seq_len=max_seq_len,
111 + scheduler_metadata=scheduler_metadata,
112 + )
113 +
114 +
115 +class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
116 +
117 + def __init__(
118 + self,
119 + num_heads: int,
120 + head_size: int,
121 + scale: float,
122 + num_kv_heads: int,
123 + alibi_slopes: Optional[list[float]],
124 + sliding_window: Optional[int],
125 + kv_cache_dtype: str,
126 + logits_soft_cap: Optional[float],
127 + attn_type: str,
128 + kv_sharing_target_layer_name: Optional[str],
129 + # MLA Specific Arguments
130 + **mla_args) -> None:
131 + super().__init__(num_heads, head_size, scale, num_kv_heads,
132 + alibi_slopes, sliding_window, kv_cache_dtype,
133 + logits_soft_cap, attn_type,
134 + kv_sharing_target_layer_name, **mla_args)
135 +
136 + assert flash_attn_supports_mla(), \
137 + "FlashAttnMLA is not supported on this device"
138 +
139 + unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
140 + if any(unsupported_features):
141 + raise NotImplementedError(
142 + "FlashAttnMLAImpl does not support one of the following: "
143 + "alibi_slopes, sliding_window, logits_soft_cap")
144 +
145 + if attn_type != AttentionType.DECODER:
146 + raise NotImplementedError("Encoder self-attention and "
147 + "encoder/decoder cross-attention "
148 + "are not implemented for "
149 + "FlashAttnMLAImpl")
150 +
151 + if is_quantized_kv_cache(self.kv_cache_dtype):
152 + raise NotImplementedError(
153 + "FlashAttnMLA V1 with FP8 KV cache not yet supported")
154 +
155 + def _forward_decode(
156 + self,
157 + q_nope: torch.Tensor,
158 + q_pe: torch.Tensor,
159 + kv_c_and_k_pe_cache: torch.Tensor,
160 + attn_metadata: FlashAttnMLAMetadata,
161 + layer: AttentionLayer,
162 + ) -> torch.Tensor:
163 + assert kv_c_and_k_pe_cache.numel() > 0
164 + assert attn_metadata.decode is not None
165 +
166 + if self.kv_cache_dtype.startswith("fp8"):
167 + raise NotImplementedError(
168 + "FP8 FlashAttention MLA not yet supported")
169 +
170 + kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]
171 + k_pe_cache = kv_c_and_k_pe_cache[..., self.kv_lora_rank:]
172 +
173 + o = flash_attn_varlen_func(
174 + q=q_pe,
175 + k=k_pe_cache.unsqueeze(-2), # Add head dim of 1
176 + v=kv_c_cache.unsqueeze(-2), # Add head dim of 1
177 + q_v=q_nope,
178 + max_seqlen_q=attn_metadata.decode.max_query_len,
179 + cu_seqlens_q=attn_metadata.decode.query_start_loc,
180 + max_seqlen_k=attn_metadata.decode.max_seq_len,
181 + seqused_k=attn_metadata.decode.seq_lens,
182 + block_table=attn_metadata.decode.block_table,
183 + softmax_scale=self.scale,
184 + causal=True,
185 + fa_version=3, # only version 3 is supported
186 + scheduler_metadata=attn_metadata.decode.scheduler_metadata,
187 + )
188 +
189 + return self._v_up_proj(o)
vllm/v1/attention/backends/mla/flashmla.py
+6 -4
@@ -85,11 +85,13 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
85 device=self.device,
86 dtype=torch.int32)
87
88 - def _build_decode(self, block_table_tensor: torch.Tensor,
89 - seq_lens: torch.Tensor) -> FlashMLADecodeMetadata:
88 + def _build_decode(
89 + self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,
90 + seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,
91 + query_start_loc_device: torch.Tensor) -> FlashMLADecodeMetadata:
92 tile_scheduler_metadata, num_splits = \
93 get_mla_metadata(
92 - seq_lens,
94 + seq_lens_device,
95 self.num_q_heads,
96 1, # MQA for the decode path
97 )
@@ -123,7 +125,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
125
126 return FlashMLADecodeMetadata(
127 block_table=block_table_tensor,
126 - seq_lens=seq_lens,
128 + seq_lens=seq_lens_device,
129 tile_scheduler_metadata=tile_scheduler_metadata,
130 num_splits=num_splits,
131 )
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+8 -6
@@ -104,12 +104,14 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
104 dtype=torch.int32,
105 device=device)
106
107 - def _build_decode(self, block_table_tensor: torch.Tensor,
108 - seq_lens: torch.Tensor) -> AiterMLADecodeMetadata:
107 + def _build_decode(
108 + self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,
109 + seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,
110 + query_start_loc_device: torch.Tensor) -> AiterMLADecodeMetadata:
111 page_size = self.kv_cache_spec.block_size
110 - block_table_bounds = (seq_lens + page_size - 1) // page_size
112 + block_table_bounds = (seq_lens_device + page_size - 1) // page_size
113 device = self.device
112 - num_reqs = seq_lens.size(0)
114 + num_reqs = seq_lens_device.size(0)
115
116 mask = (torch.arange(block_table_tensor.size(1),
117 dtype=block_table_tensor.dtype,
@@ -117,7 +119,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
119 < block_table_bounds.unsqueeze(1))
120 paged_kv_indices = block_table_tensor[mask]
121
120 - paged_kv_last_page_len = seq_lens % page_size
122 + paged_kv_last_page_len = seq_lens_device % page_size
123 paged_kv_last_page_len = torch.where(paged_kv_last_page_len == 0,
124 page_size, paged_kv_last_page_len)
125
@@ -156,7 +158,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
158
159 attn_metadata = AiterMLADecodeMetadata(
160 block_table=block_table_tensor,
159 - seq_lens=seq_lens,
161 + seq_lens=seq_lens_device,
162 paged_kv_indptr=paged_kv_indptr,
163 paged_kv_indices=paged_kv_indices,
164 paged_kv_last_page_len=paged_kv_last_page_len,
vllm/v1/attention/backends/short_conv_attn.py
+4 -3
@@ -58,8 +58,9 @@ class ShortConvAttentionMetadataBuilder(
58 state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
59
60 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
61 - split_decodes_and_prefills(common_attn_metadata,
62 - decode_threshold=1))
61 + split_decodes_and_prefills(
62 + common_attn_metadata,
63 + decode_threshold=self.reorder_batch_threshold))
64 has_initial_states = None
65 if num_prefills > 0:
66 #[batch,]
@@ -78,4 +79,4 @@ class ShortConvAttentionMetadataBuilder(
79 has_initial_states=has_initial_states,
80 state_indices_tensor=state_indices_tensor,
81 )
81 - return attn_metadata
\ No newline at end of file
82 + return attn_metadata
vllm/v1/attention/backends/xformers.py
+10 -6
@@ -3,7 +3,7 @@
3 """Attention layer with XFormersAttention."""
4
5 from dataclasses import dataclass
6 -from typing import TYPE_CHECKING, Optional
6 +from typing import TYPE_CHECKING, ClassVar, Optional
7
8 import torch
9
@@ -197,6 +197,8 @@ class XFormersAttentionMetadata:
197 class XFormersAttentionMetadataBuilder(
198 AttentionMetadataBuilder[XFormersAttentionMetadata]):
199
200 + reorder_batch_threshold: ClassVar[int] = 1
201 +
202 def __init__(
203 self,
204 kv_cache_spec: AttentionSpec,
@@ -212,9 +214,10 @@ class XFormersAttentionMetadataBuilder(
214
215 def reorder_batch(self, input_batch: "InputBatch",
216 scheduler_output: "SchedulerOutput") -> bool:
215 - return reorder_batch_to_split_decodes_and_prefills(input_batch,
216 - scheduler_output,
217 - decode_threshold=1)
217 + return reorder_batch_to_split_decodes_and_prefills(
218 + input_batch,
219 + scheduler_output,
220 + decode_threshold=self.reorder_batch_threshold)
221
222 def build(
223 self,
@@ -223,8 +226,9 @@ class XFormersAttentionMetadataBuilder(
226 fast_build: bool = False,
227 ) -> XFormersAttentionMetadata:
228 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
226 - split_decodes_and_prefills(common_attn_metadata,
227 - decode_threshold=1))
229 + split_decodes_and_prefills(
230 + common_attn_metadata,
231 + decode_threshold=self.reorder_batch_threshold))
232
233 num_actual_tokens = common_attn_metadata.num_actual_tokens
234 q_start_loc = common_attn_metadata.query_start_loc