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
[Perf] Add decode full-graph support to FlashInfer-MLA backend (#26313)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Oct 6, 2025 at 19:03 UTC
f77df94647ca278575c2de4dada36d21d089e979
1 file changed
+12
-1
vllm/v1/attention/backends/mla/flashinfer_mla.py
+12
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Optional, Union
4
+from typing import ClassVar, Optional, Union
5
6
import torch
7
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
@@ -12,13 +12,20 @@ from vllm.v1.attention.backends.mla.common import (
12
MLACommonBackend,
13
MLACommonImpl,
14
MLACommonMetadata,
15
+ MLACommonMetadataBuilder,
16
)
17
+from vllm.v1.attention.backends.utils import AttentionCGSupport
18
19
logger = init_logger(__name__)
20
21
FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
22
23
24
+class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
25
+ # enable full CUDA Graph support for decode-only capture
26
+ cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
27
+
28
+
29
class FlashInferMLABackend(MLACommonBackend):
30
@staticmethod
31
def get_name() -> str:
@@ -28,6 +35,10 @@ class FlashInferMLABackend(MLACommonBackend):
35
def get_impl_cls() -> type["FlashInferMLAImpl"]:
36
return FlashInferMLAImpl
37
38
+ @staticmethod
39
+ def get_builder_cls() -> type["FlashInferMLAMetadataBuilder"]:
40
+ return FlashInferMLAMetadataBuilder
41
+
42
43
g_fi_workspace = torch.zeros(
44
FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE,