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 8, 11:15 UTC
Reapply [Attention][FA3] Update FA3 to include new swizzle optimization (#34043)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Feb 11, 2026 at 08:07 UTC
c7914d30f90bc47f1c959d3330666885a0034f7d
6 files changed
+60
-44
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 188be16520ceefdc625fdf71365585d2ee348fe2
41
+ GIT_TAG 5824e6e2008271063c3229ab3e7032bd74abbbc6
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
tests/v1/cudagraph/test_cudagraph_dispatch.py
+13
-9
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import replace
4
from unittest.mock import MagicMock, patch
5
6
import pytest
@@ -132,36 +133,39 @@ class TestCudagraphDispatcher:
133
134
# Test dispatch logic
135
# 1. non-uniform batch, size in cudagraph size list
135
- desc_full_exact = BatchDescriptor(
136
- num_tokens=8,
137
- uniform=False,
138
- )
136
+ # FULL mode uses exact keys with num_reqs set
137
+ desc_full_with_reqs = BatchDescriptor(num_tokens=8, num_reqs=8, uniform=False)
138
+ # PIECEWISE mode uses relaxed keys with num_reqs=None
139
+ desc_piecewise = BatchDescriptor(num_tokens=8, num_reqs=None, uniform=False)
140
rt_mode, key = dispatcher.dispatch(
141
num_tokens=8, uniform_decode=False, has_lora=False
142
)
143
if cudagraph_mode_str == "FULL":
144
assert rt_mode == CUDAGraphMode.FULL
144
- assert key == desc_full_exact
145
+ assert key == desc_full_with_reqs
146
elif cudagraph_mode_str in ["FULL_AND_PIECEWISE", "PIECEWISE"]:
147
assert rt_mode == CUDAGraphMode.PIECEWISE
147
- assert key == desc_full_exact
148
+ assert key == desc_piecewise
149
else:
150
assert rt_mode == CUDAGraphMode.NONE
151
152
# 2. uniform decode batch, size in cudagraph size list
153
desc_uniform_exact = BatchDescriptor(num_tokens=8, num_reqs=8, uniform=True)
154
+ desc_non_uniform = BatchDescriptor(num_tokens=8, num_reqs=8, uniform=False)
155
rt_mode, key = dispatcher.dispatch(
156
num_tokens=8, uniform_decode=True, has_lora=False
157
)
158
if cudagraph_mode_str == "FULL":
159
+ # Pure FULL mode uses non-uniform keys for all batches
160
assert rt_mode == CUDAGraphMode.FULL
158
- assert key == desc_uniform_exact.relax_for_mixed_batch_cudagraphs()
161
+ assert key == desc_non_uniform
162
elif cudagraph_mode_str in ["FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"]:
163
+ # These modes have separate uniform decode keys
164
assert rt_mode == CUDAGraphMode.FULL
165
assert key == desc_uniform_exact
166
elif cudagraph_mode_str == "PIECEWISE":
167
assert rt_mode == CUDAGraphMode.PIECEWISE
164
- assert key == desc_uniform_exact.relax_for_mixed_batch_cudagraphs()
168
+ assert key == replace(desc_uniform_exact, num_reqs=None, uniform=False)
169
else:
170
assert rt_mode == CUDAGraphMode.NONE
171
@@ -180,7 +184,7 @@ class TestCudagraphDispatcher:
184
185
if "PIECEWISE" in cudagraph_mode_str: # string contains check
186
assert rt_mode == CUDAGraphMode.PIECEWISE
183
- assert key == desc_full_exact.relax_for_mixed_batch_cudagraphs()
187
+ assert key == replace(desc_full_exact, num_reqs=None, uniform=False)
188
else:
189
assert rt_mode == CUDAGraphMode.NONE
190
vllm/forward_context.py
+3
-15
@@ -5,7 +5,7 @@ import time
5
from collections import defaultdict
6
from contextlib import contextmanager
7
from dataclasses import dataclass, field
8
-from typing import Any, NamedTuple
8
+from typing import Any
9
10
import torch
11
@@ -26,7 +26,8 @@ batchsize_logging_interval: float = envs.VLLM_LOG_BATCHSIZE_INTERVAL
26
batchsize_forward_time: defaultdict = defaultdict(list)
27
28
29
-class BatchDescriptor(NamedTuple):
29
+@dataclass(frozen=True)
30
+class BatchDescriptor:
31
"""
32
Batch descriptor for cudagraph dispatching. We should keep the num of
33
items as minimal as possible to properly and uniquely describe the padded
@@ -56,19 +57,6 @@ class BatchDescriptor(NamedTuple):
57
to be properly captured.
58
"""
59
59
- def relax_for_mixed_batch_cudagraphs(self) -> "BatchDescriptor":
60
- """
61
- Return a relaxed version of current batch descriptor that is still compatible
62
- with PIECEWISE cudagraphs (or mixed prefill-decode FA cudagraphs).
63
- """
64
- return BatchDescriptor(
65
- self.num_tokens,
66
- num_reqs=None,
67
- uniform=False,
68
- has_lora=self.has_lora,
69
- num_active_loras=self.num_active_loras,
70
- )
71
-
60
61
def _compute_sp_num_tokens(
62
num_tokens_across_dp_cpu: torch.Tensor, sequence_parallel_size: int
vllm/v1/attention/backends/flash_attn.py
+11
-2
@@ -40,7 +40,7 @@ from vllm.model_executor.layers.batch_invariant import (
40
vllm_is_batch_invariant,
41
)
42
from vllm.platforms.interface import DeviceCapability
43
-from vllm.utils.math_utils import cdiv
43
+from vllm.utils.math_utils import cdiv, round_up
44
from vllm.v1.attention.backend import (
45
AttentionCGSupport,
46
AttentionMetadataBuilder,
@@ -310,8 +310,17 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
310
self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size
311
312
if self.use_full_cuda_graph and self.aot_schedule:
313
+ # FA3 scheduler_metadata size: 1 + round_up(batch_size, 4) * 4
314
+ # The +1 is for the tile_count_semaphore (synchronization).
315
+ # The 4 slots per batch element (num_prepare_batch_vectors) are:
316
+ # prepare_varlen + dynamic_split + sort_batches + head_swizzle
317
+ # See: https://github.com/vllm-project/flash-attention/blob/5824e6e/hopper/flash_api.cpp#L664-L671 # noqa: E501
318
+ max_batch_size = max(
319
+ vllm_config.scheduler_config.max_num_seqs,
320
+ self.max_cudagraph_size or 0,
321
+ )
322
self.scheduler_metadata = torch.zeros(
314
- vllm_config.scheduler_config.max_num_seqs + 1,
323
+ 1 + round_up(max_batch_size, 4) * 4,
324
dtype=torch.int32,
325
device=self.device,
326
)
vllm/v1/attention/backends/mla/flashattn_mla.py
+11
-1
@@ -21,6 +21,7 @@ from vllm.model_executor.layers.batch_invariant import (
21
vllm_is_batch_invariant,
22
)
23
from vllm.platforms.interface import DeviceCapability
24
+from vllm.utils.math_utils import round_up
25
from vllm.v1.attention.backend import (
26
AttentionCGSupport,
27
AttentionLayer,
@@ -129,8 +130,17 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
130
self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size
131
132
if self.use_full_cuda_graph and self.fa_aot_schedule:
133
+ # FA3 scheduler_metadata size: 1 + round_up(batch_size, 4) * 4
134
+ # The +1 is for the tile_count_semaphore (synchronization).
135
+ # The 4 slots per batch element (num_prepare_batch_vectors) are:
136
+ # prepare_varlen + dynamic_split + sort_batches + head_swizzle
137
+ # See: https://github.com/vllm-project/flash-attention/blob/5824e6e/hopper/flash_api.cpp#L664-L671 # noqa: E501
138
+ max_batch_size = max(
139
+ vllm_config.scheduler_config.max_num_seqs,
140
+ self.max_cudagraph_size or 0,
141
+ )
142
self.scheduler_metadata = torch.zeros(
133
- vllm_config.scheduler_config.max_num_seqs + 1,
143
+ 1 + round_up(max_batch_size, 4) * 4,
144
dtype=torch.int32,
145
device=self.device,
146
)
vllm/v1/cudagraph_dispatcher.py
+21
-16
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import replace
4
from itertools import product
5
6
from vllm.config import CUDAGraphMode, VllmConfig
@@ -180,12 +181,14 @@ class CudagraphDispatcher:
181
for bs, num_active_loras in product(
182
self.compilation_config.cudagraph_capture_sizes, lora_cases
183
):
183
- self.add_cudagraph_key(
184
- cudagraph_mode.mixed_mode(),
185
- self._create_padded_batch_descriptor(
186
- bs, False, num_active_loras > 0, num_active_loras
187
- ).relax_for_mixed_batch_cudagraphs(),
184
+ batch_desc = self._create_padded_batch_descriptor(
185
+ bs, False, num_active_loras > 0, num_active_loras
186
)
187
+ # Only relax for PIECEWISE mode. FULL mode needs exact num_reqs
188
+ # because FA3's scheduler_metadata computation depends on it.
189
+ if cudagraph_mode.mixed_mode() == CUDAGraphMode.PIECEWISE:
190
+ batch_desc = replace(batch_desc, num_reqs=None, uniform=False)
191
+ self.add_cudagraph_key(cudagraph_mode.mixed_mode(), batch_desc)
192
193
# if decode cudagraph mode is FULL, and we don't already have mixed
194
# mode full cudagraphs then add them here.
@@ -264,21 +267,23 @@ class CudagraphDispatcher:
267
batch_desc = self._create_padded_batch_descriptor(
268
num_tokens, uniform_decode, has_lora, effective_num_active_loras
269
)
267
- relaxed_batch_desc = batch_desc.relax_for_mixed_batch_cudagraphs()
268
-
269
- if not disable_full:
270
- # check if key exists for full cudagraph
271
- if batch_desc in self.cudagraph_keys[CUDAGraphMode.FULL]:
272
- return CUDAGraphMode.FULL, batch_desc
270
274
- # otherwise, check if the relaxed key exists
275
- if relaxed_batch_desc in self.cudagraph_keys[CUDAGraphMode.FULL]:
276
- return CUDAGraphMode.FULL, relaxed_batch_desc
271
+ # check if key exists for full cudagraph
272
+ # For pure FULL mode, keys are registered with uniform=False.
273
+ batch_desc_to_check = batch_desc
274
+ if self.cudagraph_mode == CUDAGraphMode.FULL:
275
+ batch_desc_to_check = replace(batch_desc, uniform=False)
276
+ if (
277
+ not disable_full
278
+ and batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.FULL]
279
+ ):
280
+ return CUDAGraphMode.FULL, batch_desc_to_check
281
282
# also check if the relaxed key exists for more "general"
283
# piecewise cudagraph
280
- if relaxed_batch_desc in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:
281
- return CUDAGraphMode.PIECEWISE, relaxed_batch_desc
284
+ batch_desc_to_check = replace(batch_desc, num_reqs=None, uniform=False)
285
+ if batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:
286
+ return CUDAGraphMode.PIECEWISE, batch_desc_to_check
287
288
# finally, just return no cudagraphs and a trivial batch descriptor
289
return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)