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
[Attention] Blackwell FP8 MLA support with CUTLASS_MLA backend (#23289)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Sep 3, 2025 at 14:05 UTC
a742322092379a25d83300a4d44abc7cf225a700
4 files changed
+186
-107
csrc/attention/mla/sm100_cutlass_mla_kernel.cu
+8
-8
@@ -64,11 +64,11 @@ struct IsPersistent {
64
static const bool value = v;
65
};
66
67
-template <typename T, bool IsPaged128, typename PersistenceOption = IsPersistent<true>>
67
+template <typename T, typename TOut, bool IsPaged128, typename PersistenceOption = IsPersistent<true>>
68
struct MlaSm100 {
69
using Element = T;
70
using ElementAcc = float;
71
- using ElementOut = T;
71
+ using ElementOut = TOut;
72
73
using TileShape = Shape<_128, _128, Shape<_512, _64>>;
74
using TileShapeH = cute::tuple_element_t<0, TileShape>;
@@ -178,7 +178,7 @@ typename T::Fmha::Arguments args_from_options(
178
return arguments;
179
}
180
181
-template <typename Element, bool IsPaged128, typename PersistenceOption>
181
+template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>
182
void runMla(
183
at::Tensor const& out,
184
at::Tensor const& q_nope,
@@ -190,7 +190,7 @@ void runMla(
190
double sm_scale,
191
int64_t num_kv_splits,
192
cudaStream_t stream) {
193
- using MlaSm100Type = MlaSm100<Element, IsPaged128, PersistenceOption>;
193
+ using MlaSm100Type = MlaSm100<Element, ElementOut, IsPaged128, PersistenceOption>;
194
typename MlaSm100Type::Fmha fmha;
195
auto arguments = args_from_options<MlaSm100Type>(out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, sm_scale, num_kv_splits);
196
@@ -233,13 +233,13 @@ void sm100_cutlass_mla_decode(
233
DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
234
DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
235
if (in_dtype == at::ScalarType::Half) {
236
- runMla<cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
236
+ runMla<cutlass::half_t, cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
237
out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
238
} else if (in_dtype == at::ScalarType::BFloat16) {
239
- runMla<cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
239
+ runMla<cutlass::bfloat16_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
240
out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
241
} else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
242
- runMla<cutlass::float_e4m3_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
242
+ runMla<cutlass::float_e4m3_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
243
out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
244
} else {
245
TORCH_CHECK(false, "Unsupported input data type of MLA");
@@ -253,7 +253,7 @@ void sm100_cutlass_mla_decode(
253
int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
254
// Workspace size depends on ElementAcc and ElementLSE (same as ElementAcc)
255
// which are float, so Element type here doesn't matter.
256
- using MlaSm100Type = MlaSm100<cutlass::half_t, true>;
256
+ using MlaSm100Type = MlaSm100<cutlass::half_t, cutlass::half_t, true>;
257
258
// Get split kv. Requires problem shape and sm_count only.
259
typename MlaSm100Type::Fmha::Arguments arguments;
tests/kernels/test_cutlass_mla_decode.py
+167
-83
@@ -1,96 +1,180 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+import random
5
+
6
import pytest
7
import torch
5
-import torch.nn.functional as F
6
-from torch import Tensor
8
9
import vllm._custom_ops as ops
10
from vllm.platforms import current_platform
10
-
11
-if not current_platform.has_device_capability(100):
12
- pytest.skip(
13
- reason="Cutlass MLA Requires compute capability of 10 or above.",
14
- allow_module_level=True)
15
-
16
-
17
-def ref_mla(
18
- out: Tensor, # (bs, num_heads, v_head_dim)
19
- query: Tensor, # (bs, num_heads, head_dim)
20
- kv_cache: Tensor, # (num_blocks, block_size, head_dim)
21
- scale: float,
22
- block_tables: Tensor, # (bs, max_num_blocks)
23
- seq_lens: Tensor, # (bs,)
24
-):
25
- bs, num_heads, v_head_dim = out.shape
26
- head_dim = query.shape[2]
27
-
28
- for i in range(bs):
29
- # gather and flatten KV-cache
30
- kv = kv_cache[
31
- block_tables[i]] # (max_num_blocks, block_size, head_dim)
32
- kv = kv.view(1, -1,
33
- head_dim)[:, :seq_lens[i]] # (1, seq_len, head_dim)
34
- v = kv[:, :, :v_head_dim]
35
-
36
- q = query[i].view(num_heads, 1, head_dim)
37
- o = F.scaled_dot_product_attention(q,
38
- kv,
39
- v,
40
- scale=scale,
41
- enable_gqa=True)
42
- out[i] = o.view(num_heads, v_head_dim)
43
-
44
- return out
45
-
46
-
47
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
48
-@pytest.mark.parametrize("mean_seq_len", [128, 1024, 4096])
49
-@pytest.mark.parametrize("bs", [1, 2, 4])
11
+from vllm.triton_utils import triton
12
+
13
+
14
+def cal_diff(x: torch.Tensor,
15
+ y: torch.Tensor,
16
+ name: str,
17
+ use_fp8: bool = False) -> None:
18
+ x, y = x.double(), y.double()
19
+ cos_diff = 1 - 2 * (x * y).sum().item() / max(
20
+ (x * x + y * y).sum().item(), 1e-12)
21
+ if (use_fp8):
22
+ assert cos_diff < 1e-4
23
+ else:
24
+ assert cos_diff < 1e-5
25
+
26
+
27
+CUTLASS_MLA_UNSUPPORTED_REASON = \
28
+ "Cutlass MLA Requires compute capability of 10 or above." \
29
+ if not current_platform.is_device_capability(100) \
30
+ else "Cutlass MLA is supported"
31
+
32
+
33
+@pytest.mark.skipif(not current_platform.has_device_capability(100),
34
+ reason=CUTLASS_MLA_UNSUPPORTED_REASON)
35
+@pytest.mark.parametrize("b", [128])
36
+@pytest.mark.parametrize("s_q", [1])
37
+@pytest.mark.parametrize("mean_sk", [4096, 8192, 16384])
38
+@pytest.mark.parametrize("h_q", [16, 32, 64, 128])
39
+@pytest.mark.parametrize("h_kv", [1])
40
+@pytest.mark.parametrize("d", [576])
41
+@pytest.mark.parametrize("dv", [512])
42
+@pytest.mark.parametrize("block_size", [64])
43
+@pytest.mark.parametrize("causal", [True])
44
@pytest.mark.parametrize("varlen", [False, True])
51
-@pytest.mark.parametrize("block_size", [16, 64, 128])
52
-def test_cutlass_mla_decode(dtype: torch.dtype, mean_seq_len: int, bs: int,
53
- varlen: bool, block_size: int):
54
- torch.set_default_dtype(dtype)
55
- torch.set_default_device('cuda')
45
+@pytest.mark.parametrize("torch_dtype", [torch.bfloat16, torch.float8_e4m3fn])
46
+@torch.inference_mode()
47
+def test_cutlass_mla_decode(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size,
48
+ causal, varlen, torch_dtype):
49
+ device = torch.device("cuda:0")
50
+ if torch_dtype == torch.float8_e4m3fn:
51
+ init_dtype = torch.bfloat16
52
+ else:
53
+ init_dtype = torch_dtype
54
+ torch.set_default_dtype(init_dtype)
55
+ torch.set_default_device(device)
56
+ torch.cuda.set_device(device)
57
torch.manual_seed(42)
58
+ random.seed(42)
59
58
- d = 576
59
- h_q = 128
60
- dv = 512
60
+ print(f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, "
61
+ f"{d=}, {dv=}, {causal=}, {varlen=}, {torch_dtype=}")
62
62
- q_nope_dim = 128
63
- q_pe_dim = 64
64
- scale = (q_nope_dim + q_pe_dim)**(-0.5)
63
+ use_fp8 = torch_dtype == torch.float8_e4m3fn
64
+ scale = math.sqrt(d)**(-1)
65
+ cache_seqlens = torch.full((b, ), mean_sk, dtype=torch.int32)
66
if varlen:
66
- seq_lens = torch.empty(bs).normal_(mean_seq_len, mean_seq_len / 2)
67
- seq_lens = seq_lens.clip(2).to(torch.int32)
67
+ for i in range(b):
68
+ cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2),
69
+ s_q)
70
+ total_seqlens = cache_seqlens.sum().item()
71
+ max_seqlen = cache_seqlens.max().item()
72
+ max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256
73
+
74
+ q = torch.randn(b, s_q, h_q, d)
75
+ block_table = torch.arange(b * max_seqlen_pad // block_size,
76
+ dtype=torch.int32).view(
77
+ b, max_seqlen_pad // block_size)
78
+ blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d)
79
+ blocked_v = blocked_k[..., :dv]
80
+
81
+ init_dtype = q.dtype
82
+ if use_fp8:
83
+ fp8_dtype = torch.float8_e4m3fn
84
+ descale_q = torch.ones((1), dtype=torch.float32)
85
+ descale_k = torch.ones((1), dtype=torch.float32)
86
+
87
+ q = q.to(fp8_dtype)
88
+ blocked_k = blocked_k.to(fp8_dtype)
89
+ blocked_v = blocked_v.to(fp8_dtype)
90
else:
69
- seq_lens = torch.full((bs, ), mean_seq_len, dtype=torch.int32)
70
- max_seq_len = seq_lens.max().item()
71
- block_num = (max_seq_len + block_size - 1) // block_size
72
-
73
- # Pad block_num so that small blocks can be packed into full 128-sized
74
- # CUTLASS tiles. One 128-wide tile can hold (128 // block_size) small
75
- # blocks.
76
- pack_factor = 128 // block_size
77
- block_num = ((block_num + pack_factor - 1) // pack_factor) * pack_factor
78
-
79
- # Amplify input values to ensure test coverage of edge cases where CUTLASS
80
- # kernel errors occur with split_k settings.
81
- q = torch.randn(bs, h_q, d) * 100
82
- block_table = torch.randint(0,
83
- bs * block_num, (bs, block_num),
84
- dtype=torch.int32)
85
-
86
- kv_cache = torch.randn(block_table.numel(), block_size, d)
87
-
88
- out_ref = q.new_zeros(bs, h_q, dv)
89
- ref_mla(out_ref, q, kv_cache, scale, block_table, seq_lens)
90
- out_ans = torch.zeros_like(out_ref)
91
- q_nope = q[:, :, :dv].clone()
92
- q_pe = q[:, :, dv:].clone()
93
- ops.cutlass_mla_decode(out_ans, q_nope, q_pe, kv_cache, seq_lens,
94
- block_table, scale)
95
-
96
- torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
91
+ descale_q = None
92
+ descale_k = None
93
+
94
+ def cutlass_mla():
95
+ MAX_HEADS = 128
96
+
97
+ q_reshaped = q.squeeze(1)
98
+ q_nope = q_reshaped[:, :, :dv].clone()
99
+ q_pe = q_reshaped[:, :, dv:].clone()
100
+
101
+ if h_q < MAX_HEADS:
102
+ q_nope_padded = q_nope.new_empty((b, MAX_HEADS, dv))
103
+ q_nope_padded[:, :h_q] = q_nope
104
+ q_nope = q_nope_padded
105
+
106
+ q_pe_padded = q_pe.new_empty((b, MAX_HEADS, d - dv))
107
+ q_pe_padded[:, :h_q] = q_pe
108
+ q_pe = q_pe_padded
109
+
110
+ kv_cache_flat = blocked_k.squeeze(2)
111
+ device_properties = torch.cuda.get_device_properties(
112
+ torch.device("cuda:0"))
113
+ sm_count = device_properties.multi_processor_count
114
+ workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
115
+ max_seqlen * block_size, b, sm_count, num_kv_splits=1)
116
+ workspace = torch.empty(workspace_size,
117
+ device="cuda",
118
+ dtype=torch.uint8)
119
+
120
+ out_ans = torch.empty(b, MAX_HEADS, dv, dtype=init_dtype)
121
+
122
+ ops.sm100_cutlass_mla_decode(out_ans, q_nope, q_pe, kv_cache_flat,
123
+ cache_seqlens, block_table, workspace,
124
+ scale, 1)
125
+ return out_ans[:, :h_q].contiguous()
126
+
127
+ def scaled_dot_product_attention(query, key, value, is_causal=False):
128
+ query = query.float()
129
+ key = key.float()
130
+ value = value.float()
131
+ key = key.repeat_interleave(h_q // h_kv, dim=0)
132
+ value = value.repeat_interleave(h_q // h_kv, dim=0)
133
+ attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))
134
+ if is_causal:
135
+ s_q = query.shape[-2]
136
+ s_k = key.shape[-2]
137
+ attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype)
138
+ temp_mask = torch.ones(s_q, s_k,
139
+ dtype=torch.bool).tril(diagonal=s_k - s_q)
140
+ attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
141
+ attn_bias.to(query.dtype)
142
+ attn_weight += attn_bias
143
+ lse = attn_weight.logsumexp(dim=-1)
144
+ attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
145
+ return attn_weight @ value, lse
146
+
147
+ def ref_mla():
148
+ q_ = (q.to(torch.float) * descale_q).to(init_dtype) if use_fp8 else q
149
+ blocked_k_ = (blocked_k.to(torch.float) *
150
+ descale_k).to(init_dtype) if use_fp8 else blocked_k
151
+ blocked_v_ = (blocked_v.to(torch.float) *
152
+ descale_k).to(init_dtype) if use_fp8 else blocked_v
153
+ out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32)
154
+ lse = torch.empty(b, h_q, s_q, dtype=torch.float32)
155
+ for i in range(b):
156
+ begin = i * max_seqlen_pad
157
+ end = begin + cache_seqlens[i]
158
+ out_i, lse_i = scaled_dot_product_attention(
159
+ q_[i].transpose(0, 1),
160
+ blocked_k_.view(-1, h_kv, d)[begin:end].transpose(0, 1),
161
+ blocked_v_.view(-1, h_kv, dv)[begin:end].transpose(0, 1),
162
+ is_causal=causal,
163
+ )
164
+ out[i] = out_i.transpose(0, 1)
165
+ lse[i] = lse_i
166
+ return out, lse
167
+
168
+ out_cutlass = cutlass_mla()
169
+ out_torch, lse_torch = ref_mla()
170
+ # Extract the single token (s_q=1) slice to match cutlass output shape
171
+ out_torch_slice = out_torch[:, 0, :, :] # [b, h_q, dv]
172
+ cal_diff(out_cutlass, out_torch_slice, "out", use_fp8)
173
+
174
+ t = triton.testing.do_bench(cutlass_mla)
175
+ FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2
176
+ bytes = (total_seqlens * h_kv * d +
177
+ b * s_q * h_q * d) * (torch.finfo(torch_dtype).bits // 8) + (
178
+ b * s_q * h_q * dv) * (torch.finfo(init_dtype).bits // 8)
179
+ print(f"{t:.3f} ms, {FLOPS / 10 ** 9 / t:.0f} TFLOPS,",
180
+ f"{bytes / 10 ** 6 / t:.0f} GB/s")
vllm/platforms/cuda.py
+2
-2
@@ -500,8 +500,8 @@ class CudaPlatformBase(Platform):
500
else:
501
attention_backend = "FLASHMLA"
502
503
- # Only FlashMLA supports fp8
504
- if attention_backend == "FLASHMLA":
503
+ # Only FlashMLA and CUTLASS_MLA support fp8
504
+ if attention_backend in ["FLASHMLA", "CUTLASS_MLA"]:
505
supported = True
506
else:
507
supported = (not fp8_attention)
vllm/v1/attention/backends/mla/cutlass_mla.py
+9
-14
@@ -108,10 +108,6 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
108
"are not implemented for "
109
"CutlassMLAImpl")
110
111
- if is_quantized_kv_cache(self.kv_cache_dtype):
112
- raise NotImplementedError(
113
- "CutlassMLA V1 with FP8 KV cache not yet supported")
114
-
111
self._use_old_cutlass_mla = False
112
force_old_cutlass = os.environ.get("FORCE_OLD_CUTLASS_MLA", None)
113
if force_old_cutlass:
@@ -182,11 +178,10 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
178
> 0), f"block num must be greater than 0, got {block_num}"
179
assert block_num % (128 / PAGE_SIZE) == 0
180
185
- # TODO(kaixih@nvidia): support fp8
181
assert q_nope.dtype in (
187
- torch.float16,
188
- torch.bfloat16,
189
- ), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}."
182
+ torch.float16, torch.bfloat16, torch.float8_e4m3fn), (
183
+ f"q_nope.dtype needs to be fp16 or bf16 or e4m3 but got "
184
+ f"{q_nope.dtype}.")
185
assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype
186
assert (
187
seq_lens.dtype == torch.int32
@@ -195,7 +190,9 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
190
page_table.dtype == torch.int32
191
), f"page_table.dtype needs to be int32 but got {page_table.dtype}."
192
198
- out = q_nope.new_empty((B_q, MAX_HEADS, D_latent))
193
+ dtype = (torch.bfloat16 if is_quantized_kv_cache(self.kv_cache_dtype)
194
+ else q_nope.dtype)
195
+ out = q_nope.new_empty((B_q, MAX_HEADS, D_latent), dtype=dtype)
196
197
ops.sm100_cutlass_mla_decode(
198
out,
@@ -220,9 +217,6 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
217
assert kv_c_and_k_pe_cache.numel() > 0
218
assert attn_metadata.decode is not None
219
223
- if self.kv_cache_dtype.startswith("fp8"):
224
- raise NotImplementedError("FP8 Cutlass MLA not yet supported")
225
-
220
# Adjust workspace size (if necessary)
221
self._workspace.ensure_size(attn_metadata, self._num_kv_splits)
222
@@ -252,8 +246,9 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
246
assert kv_c_and_k_pe_cache.numel() > 0
247
assert attn_metadata.decode is not None
248
255
- if self.kv_cache_dtype.startswith("fp8"):
256
- raise NotImplementedError("FP8 Cutlass MLA not yet supported")
249
+ if is_quantized_kv_cache(self.kv_cache_dtype):
250
+ raise NotImplementedError(
251
+ "FP8 Cutlass MLA not supported with FORCE_OLD_CUTLASS_MLA")
252
253
B = q_nope.shape[0]
254