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 3, 09:15 UTC
[Perf][MLA] Enable FULL cudagraph capture for TRITON_MLA decode (#42885)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
May 19, 2026 at 05:29 UTC
0191354827560fe38f68b4e7207f8824d6152ca3
1 file changed
+10
vllm/v1/attention/backends/mla/triton_mla.py
+10
@@ -12,12 +12,14 @@ from vllm.model_executor.layers.attention.mla_attention import (
12
MLACommonBackend,
13
MLACommonImpl,
14
MLACommonMetadata,
15
+ MLACommonMetadataBuilder,
16
)
17
from vllm.platforms import current_platform
18
from vllm.platforms.interface import DeviceCapability
19
from vllm.triton_utils import triton
20
from vllm.utils.torch_utils import is_quantized_kv_cache
21
from vllm.v1.attention.backend import (
22
+ AttentionCGSupport,
23
AttentionLayer,
24
AttentionType,
25
MultipleOf,
@@ -27,6 +29,10 @@ from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
29
logger = init_logger(__name__)
30
31
32
+class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
33
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
34
+
35
+
36
class TritonMLABackend(MLACommonBackend):
37
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
38
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -63,6 +69,10 @@ class TritonMLABackend(MLACommonBackend):
69
def get_impl_cls() -> type["TritonMLAImpl"]:
70
return TritonMLAImpl
71
72
+ @staticmethod
73
+ def get_builder_cls() -> type["TritonMLAMetadataBuilder"]:
74
+ return TritonMLAMetadataBuilder
75
+
76
@classmethod
77
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
78
return True