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 5, 23:30 UTC
[Bugfix] Restore CUDA graph persistent buffers for FP8 FlashMLA decode (#35175)
Signed-off-by: haosdent <haosdent@gmail.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
haosdent committed
Mar 27, 2026 at 00:13 UTC
0aac2048bf3a7e60eaddf1ebcb4165ed777eb8ff
1 file changed
+15
vllm/v1/attention/backends/mla/flashmla.py
+15
@@ -172,6 +172,21 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
172
num_q_tokens_per_head_k,
173
1, # MQA for the decode path
174
)
175
+
176
+ # Copy FP8 metadata into persistent CUDA graph buffers
177
+ if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
178
+ assert self.cg_buf_tile_scheduler_metadata is not None
179
+ assert self.cg_buf_num_splits is not None
180
+ n = tile_scheduler_metadata.size(0)
181
+ assert n <= self.cg_buf_tile_scheduler_metadata.size(0)
182
+ self.cg_buf_tile_scheduler_metadata[:n].copy_(tile_scheduler_metadata)
183
+ tile_scheduler_metadata = self.cg_buf_tile_scheduler_metadata[:n]
184
+
185
+ n = num_splits.size(0)
186
+ assert n <= self.cg_buf_num_splits.size(0)
187
+ self.cg_buf_num_splits[:n].copy_(num_splits)
188
+ num_splits = self.cg_buf_num_splits[:n]
189
+
190
scheduler_metadata.tile_scheduler_metadata = tile_scheduler_metadata
191
scheduler_metadata.num_splits = num_splits
192