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
[Feature] Full Cuda Graph Support for Cutlass MLA and 6% E2E Throughput Improvement (#22763)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 15, 2025 at 02:27 UTC
5c3fbfe46bf62d339a42476120d9bf268fedfa24
2 files changed
+88
-2
tests/compile/piecewise/test_full_cudagraph.py
+74
@@ -66,6 +66,80 @@ def llm_pair(request):
66
)
67
68
69
+@pytest.fixture(scope="class")
70
+def cutlass_mla_llm_pair(request):
71
+ model = request.param
72
+
73
+ # force V1 engine and Cutlass MLA backend
74
+ with temporary_environ({
75
+ "VLLM_USE_V1": "1",
76
+ "VLLM_ATTENTION_BACKEND": "CUTLASS_MLA",
77
+ "FORCE_NUM_KV_SPLITS":
78
+ "1", # TODO: remove this when hang issue is fixed
79
+ }):
80
+ full = LLM(
81
+ model=model,
82
+ gpu_memory_utilization=0.45,
83
+ trust_remote_code=True,
84
+ max_model_len=1024,
85
+ compilation_config=CompilationConfig(
86
+ full_cuda_graph=True,
87
+ cudagraph_capture_sizes=[16, 32, 64, 128, 256, 512],
88
+ ),
89
+ )
90
+ piecewise = LLM(
91
+ model=model,
92
+ gpu_memory_utilization=0.45,
93
+ trust_remote_code=True,
94
+ max_model_len=1024,
95
+ compilation_config=CompilationConfig(),
96
+ )
97
+
98
+ yield weakref.proxy(full), weakref.proxy(piecewise)
99
+ del full
100
+ del piecewise
101
+
102
+ wait_for_gpu_memory_to_clear(
103
+ devices=[0],
104
+ threshold_ratio=0.1,
105
+ )
106
+
107
+
108
+@pytest.mark.parametrize(
109
+ "cutlass_mla_llm_pair",
110
+ [
111
+ # use an MLA model
112
+ "deepseek-ai/DeepSeek-V2-Lite",
113
+ ],
114
+ indirect=True)
115
+@pytest.mark.skipif(current_platform.get_device_capability() != (10, 0),
116
+ reason="Only Blackwell GPUs support Cutlass MLA")
117
+class TestFullCUDAGraphCutlassMLA:
118
+ """
119
+ Validate full CUDA Graph with Cutlass MLA (decode-only capture).
120
+ """
121
+
122
+ @pytest.mark.parametrize(("batch_size", "max_tokens"), [
123
+ (8, 8),
124
+ ])
125
+ def test_full_cudagraph_sm100_cutlass_mla(
126
+ self, batch_size, max_tokens, cutlass_mla_llm_pair: tuple[LLM,
127
+ LLM]):
128
+ piecewise_llm, full_cudagraph_llm = cutlass_mla_llm_pair
129
+
130
+ prompts = ["Hello, my name is"] * batch_size
131
+ sampling_params = SamplingParams(temperature=0.0,
132
+ max_tokens=max_tokens,
133
+ top_p=0.95)
134
+
135
+ piecewise_responses = piecewise_llm.generate(prompts, sampling_params)
136
+ full_responses = full_cudagraph_llm.generate(prompts, sampling_params)
137
+
138
+ for piecewise_res, full_res in zip(piecewise_responses,
139
+ full_responses):
140
+ assert piecewise_res.outputs[0].text == full_res.outputs[0].text
141
+
142
+
143
@pytest.mark.parametrize(
144
"llm_pair",
145
[
vllm/v1/attention/backends/mla/cutlass_mla.py
+14
-2
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
-from typing import Optional
5
+from typing import ClassVar, Optional
6
7
import torch
8
@@ -12,11 +12,19 @@ from vllm.attention.backends.abstract import (AttentionType,
12
from vllm.logger import init_logger
13
from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
14
MLACommonImpl,
15
- MLACommonMetadata)
15
+ MLACommonMetadata,
16
+ MLACommonMetadataBuilder)
17
+from vllm.v1.attention.backends.utils import AttentionCGSupport
18
19
logger = init_logger(__name__)
20
21
22
+class CutlassMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
23
+ # enable full CUDA Graph support for decode-only capture
24
+ attn_cudagraph_support: ClassVar[
25
+ AttentionCGSupport] = AttentionCGSupport.PURE_DECODE_ONLY
26
+
27
+
28
class CutlassMLABackend(MLACommonBackend):
29
30
@staticmethod
@@ -27,6 +35,10 @@ class CutlassMLABackend(MLACommonBackend):
35
def get_impl_cls() -> type["CutlassMLAImpl"]:
36
return CutlassMLAImpl
37
38
+ @staticmethod
39
+ def get_builder_cls() -> type["CutlassMLAMetadataBuilder"]:
40
+ return CutlassMLAMetadataBuilder
41
+
42
43
class SM100Workspace:
44