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] bugfix for Flash Attention MLA with full cuda graph IMA following pr-25490 (#27128)
Signed-off-by: qqma <qqma@amazon.com> Co-authored-by: qqma <qqma@amazon.com>
Daisy-Ma-coder committed
Oct 22, 2025 at 12:36 UTC
5beacce2eae2cb36e44498e2ea80b2d3ee34fe61
1 file changed
+20
-13
vllm/v1/attention/backends/mla/flashattn_mla.py
+20
-13
@@ -89,10 +89,9 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
89
self.use_full_cuda_graph = (
90
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
91
)
92
+ self.max_cudagraph_size = self.compilation_config.max_capture_size
93
94
if self.use_full_cuda_graph and self.fa_aot_schedule:
94
- self.max_cudagraph_size = self.compilation_config.max_capture_size
95
-
95
if self.max_cudagraph_size > 992:
96
# This condition derives from FA3's internal heuristic.
97
# TODO(woosuk): Support larger cudagraph sizes.
@@ -114,7 +113,14 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
113
self.max_num_splits = 1
114
115
def _schedule_decode(
117
- self, num_reqs, cu_query_lens, max_query_len, seqlens, max_seq_len, causal
116
+ self,
117
+ num_reqs,
118
+ cu_query_lens,
119
+ max_query_len,
120
+ seqlens,
121
+ max_seq_len,
122
+ causal,
123
+ max_num_splits,
124
):
125
if self.fa_aot_schedule:
126
return get_scheduler_metadata(
@@ -130,7 +136,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
136
page_size=self.page_size,
137
cu_seqlens_q=cu_query_lens,
138
causal=causal,
133
- num_splits=self.max_num_splits,
139
+ num_splits=max_num_splits,
140
)
141
return None
142
@@ -148,6 +154,15 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
154
max_query_len = query_lens_cpu.max().item()
155
max_seq_len = seq_lens_device.max().item()
156
157
+ # For Flash Attention MLA + full cudagraph
158
+ max_num_splits = 0
159
+ if self.use_full_cuda_graph and num_decode_tokens <= self.max_cudagraph_size:
160
+ # NOTE(woosuk): Setting num_splits > 1 may increase the memory
161
+ # usage, because the intermediate buffers of size [num_splits,
162
+ # num_heads, num_tokens, head_size] are allocated. Therefore,
163
+ # we only set num_splits when using cuda graphs.
164
+ max_num_splits = self.max_num_splits
165
+
166
scheduler_metadata = self._schedule_decode(
167
num_reqs=seq_lens_cpu.numel(),
168
cu_query_lens=query_start_loc_device,
@@ -155,10 +170,9 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
170
seqlens=seq_lens_device,
171
max_seq_len=max_seq_len,
172
causal=True,
173
+ max_num_splits=max_num_splits,
174
)
175
160
- # For FA3 + full cudagraph
161
- max_num_splits = 0
176
if self.use_full_cuda_graph and scheduler_metadata is not None:
177
n = scheduler_metadata.shape[0]
178
# Ensure the persistent buffer is large enough
@@ -174,13 +188,6 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
188
self.scheduler_metadata[n:] = 0
189
scheduler_metadata = self.scheduler_metadata[:n]
190
177
- if num_decode_tokens <= self.max_cudagraph_size:
178
- # NOTE(woosuk): Setting num_splits > 1 may increase the memory
179
- # usage, because the intermediate buffers of size [num_splits,
180
- # num_heads, num_tokens, head_size] are allocated. Therefore,
181
- # we only set num_splits when using cuda graphs.
182
- max_num_splits = self.max_num_splits
183
-
191
if vllm_is_batch_invariant():
192
max_num_splits = 1
193