From e53382cc2ec863b2c990bca920198ec31d963b6f Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Tue, 24 Jun 2025 11:21:52 -0400 Subject: [PATCH] Sage Moore fixes for full cuda graph support for DeepEP+DeepGEMM LL Signed-off-by: Tyler Michael Smith --- vllm/v1/attention/backends/mla/flashmla.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py index be26e0060db5e..46da28540d110 100644 --- a/vllm/v1/attention/backends/mla/flashmla.py +++ b/vllm/v1/attention/backends/mla/flashmla.py @@ -75,7 +75,8 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]): 1, # MQA for the decode path ) - if self.runner.full_cuda_graph: + n = num_splits.size(0) + if self.runner.full_cuda_graph and (n-1) <= self.runner.cudagraph_batch_sizes[-1]: # First time around (CUDAGraph capture), allocate the static buffer if self.cg_buf_tile_scheduler_metadata is None: self.cg_buf_tile_scheduler_metadata = tile_scheduler_metadata