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
Remove old cutlass mla (#23961)
Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Sep 17, 2025 at 10:31 UTC
8f3616f422e34ccb0e79f1f00d72366c4dab24f1
6 files changed
+10
-345
CMakeLists.txt
-2
@@ -308,7 +308,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
308
"csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
309
"csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
310
"csrc/cutlass_extensions/common.cpp"
311
- "csrc/attention/mla/cutlass_mla_entry.cu"
311
"csrc/quantization/fp8/per_token_group_quant.cu")
312
313
set_gencode_flags_for_srcs(
@@ -595,7 +594,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
594
cuda_archs_loose_intersection(MLA_ARCHS "10.0a" "${CUDA_ARCHS}")
595
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
596
set(SRCS
598
- "csrc/attention/mla/cutlass_mla_kernels.cu"
597
"csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
598
set_gencode_flags_for_srcs(
599
SRCS "${SRCS}"
csrc/attention/mla/cutlass_mla_entry.cu
deleted
-38
@@ -1,38 +0,0 @@
1
-/*
2
- * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3
- *
4
- * Licensed under the Apache License, Version 2.0 (the "License");
5
- * you may not use this file except in compliance with the License.
6
- * You may obtain a copy of the License at
7
- *
8
- * http://www.apache.org/licenses/LICENSE-2.0
9
- *
10
- * Unless required by applicable law or agreed to in writing, software
11
- * distributed under the License is distributed on an "AS IS" BASIS,
12
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
- * See the License for the specific language governing permissions and
14
- * limitations under the License.
15
- */
16
-
17
-#include <torch/all.h>
18
-
19
-#if defined ENABLE_CUTLASS_MLA && ENABLE_CUTLASS_MLA
20
-void cutlass_mla_decode_sm100a(torch::Tensor const& out,
21
- torch::Tensor const& q_nope,
22
- torch::Tensor const& q_pe,
23
- torch::Tensor const& kv_c_and_k_pe_cache,
24
- torch::Tensor const& seq_lens,
25
- torch::Tensor const& page_table, double scale);
26
-#endif
27
-
28
-void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
29
- torch::Tensor const& q_pe,
30
- torch::Tensor const& kv_c_and_k_pe_cache,
31
- torch::Tensor const& seq_lens,
32
- torch::Tensor const& page_table, double scale) {
33
-#if defined ENABLE_CUTLASS_MLA && ENABLE_CUTLASS_MLA
34
- return cutlass_mla_decode_sm100a(out, q_nope, q_pe, kv_c_and_k_pe_cache,
35
- seq_lens, page_table, scale);
36
-#endif
37
- TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled cutlass MLA");
38
-}
csrc/attention/mla/cutlass_mla_kernels.cu
deleted
-225
@@ -1,225 +0,0 @@
1
-/*
2
- * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3
- *
4
- * Licensed under the Apache License, Version 2.0 (the "License");
5
- * you may not use this file except in compliance with the License.
6
- * You may obtain a copy of the License at
7
- *
8
- * http://www.apache.org/licenses/LICENSE-2.0
9
- *
10
- * Unless required by applicable law or agreed to in writing, software
11
- * distributed under the License is distributed on an "AS IS" BASIS,
12
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
- * See the License for the specific language governing permissions and
14
- * limitations under the License.
15
- */
16
-
17
-#include <torch/all.h>
18
-
19
-#include <ATen/cuda/CUDAContext.h>
20
-#include <c10/cuda/CUDAGuard.h>
21
-
22
-#include "cute/tensor.hpp"
23
-
24
-#include "cutlass/cutlass.h"
25
-#include "cutlass/kernel_hardware_info.h"
26
-
27
-#include "cutlass_extensions/common.hpp"
28
-
29
-#include "device/sm100_mla.hpp"
30
-#include "kernel/sm100_mla_tile_scheduler.hpp"
31
-
32
-using namespace cute;
33
-using namespace cutlass::fmha::kernel;
34
-
35
-template <typename T, bool PersistenceOption = true>
36
-struct MlaSm100 {
37
- using Element = T;
38
- using ElementAcc = float;
39
- using ElementOut = T;
40
-
41
- using TileShape = Shape<_128, _128, Shape<_512, _64>>;
42
- using TileShapeH = cute::tuple_element_t<0, TileShape>;
43
- using TileShapeD = cute::tuple_element_t<2, TileShape>;
44
-
45
- // H K (D_latent D_rope) B
46
- using ProblemShape = cute::tuple<TileShapeH, int, TileShapeD, int>;
47
-
48
- using StrideQ = cute::tuple<int64_t, _1, int64_t>; // H D B
49
- using StrideK = cute::tuple<int64_t, _1, int64_t>; // K D B
50
- using StrideO = StrideK; // H D B
51
- using StrideLSE = cute::tuple<_1, int>; // H B
52
-
53
- using TileScheduler =
54
- std::conditional_t<PersistenceOption, Sm100MlaPersistentTileScheduler,
55
- Sm100MlaIndividualTileScheduler>;
56
-
57
- using FmhaKernel =
58
- cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized<
59
- TileShape, Element, ElementAcc, ElementOut, ElementAcc, TileScheduler,
60
- /*kIsCpAsync=*/true>;
61
- using Fmha = cutlass::fmha::device::MLA<FmhaKernel>;
62
-};
63
-
64
-template <typename T>
65
-typename T::Fmha::Arguments args_from_options(
66
- at::Tensor const& out, at::Tensor const& q_nope, at::Tensor const& q_pe,
67
- at::Tensor const& kv_c_and_k_pe_cache, at::Tensor const& seq_lens,
68
- at::Tensor const& page_table, double scale) {
69
- cutlass::KernelHardwareInfo hw_info;
70
- hw_info.device_id = q_nope.device().index();
71
- hw_info.sm_count =
72
- cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
73
- hw_info.device_id);
74
-
75
- int batches = q_nope.sizes()[0];
76
- int page_count_per_seq = page_table.sizes()[1];
77
- int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
78
- int page_size = kv_c_and_k_pe_cache.sizes()[1];
79
- int max_seq_len = page_size * page_count_per_seq;
80
- using TileShapeH = typename T::TileShapeH;
81
- using TileShapeD = typename T::TileShapeD;
82
- auto problem_shape =
83
- cute::make_tuple(TileShapeH{}, max_seq_len, TileShapeD{}, batches);
84
-
85
- auto [H, K, D, B] = problem_shape;
86
- auto [D_latent, D_rope] = D;
87
-
88
- using StrideQ = typename T::StrideQ;
89
- using StrideK = typename T::StrideK;
90
- using StrideO = typename T::StrideO;
91
- using StrideLSE = typename T::StrideLSE;
92
-
93
- StrideQ stride_Q_latent = cute::make_tuple(
94
- static_cast<int64_t>(D_latent), _1{}, static_cast<int64_t>(H * D_latent));
95
- StrideQ stride_Q_rope = cute::make_tuple(static_cast<int64_t>(D_rope), _1{},
96
- static_cast<int64_t>(H * D_rope));
97
- StrideK stride_C =
98
- cute::make_tuple(static_cast<int64_t>(D_latent + D_rope), _1{},
99
- static_cast<int64_t>(page_size * (D_latent + D_rope)));
100
- StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
101
- StrideLSE stride_LSE = cute::make_tuple(_1{}, static_cast<int>(H));
102
- StrideO stride_O = cute::make_tuple(static_cast<int64_t>(D_latent), _1{},
103
- static_cast<int64_t>(H * D_latent));
104
-
105
- using Element = typename T::Element;
106
- using ElementOut = typename T::ElementOut;
107
- using ElementAcc = typename T::ElementAcc;
108
- auto Q_latent_ptr = static_cast<Element*>(q_nope.data_ptr());
109
- auto Q_rope_ptr = static_cast<Element*>(q_pe.data_ptr());
110
- auto C_ptr = static_cast<Element*>(kv_c_and_k_pe_cache.data_ptr());
111
- auto scale_f = static_cast<float>(scale);
112
- typename T::Fmha::Arguments arguments{
113
- problem_shape,
114
- {scale_f, Q_latent_ptr, stride_Q_latent, Q_rope_ptr, stride_Q_rope, C_ptr,
115
- stride_C, C_ptr + D_latent, stride_C,
116
- static_cast<int*>(seq_lens.data_ptr()),
117
- static_cast<int*>(page_table.data_ptr()), stride_PT, page_count_total,
118
- page_size},
119
- {static_cast<ElementOut*>(out.data_ptr()), stride_O,
120
- static_cast<ElementAcc*>(nullptr), stride_LSE},
121
- hw_info,
122
- 1, // split_kv
123
- nullptr, // is_var_split_kv
124
- };
125
- // TODO(kaixih@nvidia): When split_kv=-1 and is_var_split_kv=false, we compute
126
- // split_kv automatically based on batch size and sequence length to balance
127
- // workload across available SMs. Consider using var_split_kv for manual
128
- // control if needed.
129
- T::Fmha::set_split_kv(arguments);
130
- return arguments;
131
-}
132
-
133
-template <typename Element>
134
-void runMla(at::Tensor const& out, at::Tensor const& q_nope,
135
- at::Tensor const& q_pe, at::Tensor const& kv_c_and_k_pe_cache,
136
- at::Tensor const& seq_lens, at::Tensor const& page_table,
137
- float scale, cudaStream_t stream) {
138
- using MlaSm100Type = MlaSm100<Element>;
139
- typename MlaSm100Type::Fmha fmha;
140
- auto arguments = args_from_options<MlaSm100Type>(
141
- out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, scale);
142
- size_t workspace_size = MlaSm100Type::Fmha::get_workspace_size(arguments);
143
- auto const workspace_options =
144
- torch::TensorOptions().dtype(torch::kUInt8).device(q_nope.device());
145
- auto workspace = torch::empty(workspace_size, workspace_options);
146
-
147
- CUTLASS_CHECK(fmha.can_implement(arguments));
148
-
149
- CUTLASS_CHECK(fmha.initialize(arguments, workspace.data_ptr(), stream));
150
-
151
- CUTLASS_CHECK(fmha.run(arguments, workspace.data_ptr(), stream));
152
-}
153
-
154
-void cutlass_mla_decode_sm100a(torch::Tensor const& out,
155
- torch::Tensor const& q_nope,
156
- torch::Tensor const& q_pe,
157
- torch::Tensor const& kv_c_and_k_pe_cache,
158
- torch::Tensor const& seq_lens,
159
- torch::Tensor const& page_table, double scale) {
160
- TORCH_CHECK(q_nope.device().is_cuda(), "q_nope must be on CUDA");
161
- TORCH_CHECK(q_nope.dim() == 3, "q_nope must be a 3D tensor");
162
- TORCH_CHECK(q_pe.dim() == 3, "q_pe must be a 3D tensor");
163
- TORCH_CHECK(kv_c_and_k_pe_cache.dim() == 3,
164
- "kv_c_and_k_pe_cache must be a 3D tensor");
165
- TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must be a 1D tensor");
166
- TORCH_CHECK(page_table.dim() == 2, "page_table must be a 2D tensor");
167
- TORCH_CHECK(out.dim() == 3, "out must be a 3D tensor");
168
-
169
- auto B_q_nope = q_nope.size(0);
170
- auto H_q_nope = q_nope.size(1);
171
- auto D_q_nope = q_nope.size(2);
172
- auto B_q_pe = q_pe.size(0);
173
- auto H_q_pe = q_pe.size(1);
174
- auto D_q_pe = q_pe.size(2);
175
- auto B_pt = page_table.size(0);
176
- auto PAGE_NUM = page_table.size(1);
177
- auto PAGE_SIZE = kv_c_and_k_pe_cache.size(1);
178
- auto D_ckv = kv_c_and_k_pe_cache.size(2);
179
- auto B_o = out.size(0);
180
- auto H_o = out.size(1);
181
- auto D_o = out.size(2);
182
-
183
- TORCH_CHECK(D_q_nope == 512, "D_q_nope must be equal to 512");
184
- TORCH_CHECK(D_q_pe == 64, "D_q_pe must be equal to 64");
185
- TORCH_CHECK(D_ckv == 576, "D_ckv must be equal to 576");
186
- TORCH_CHECK(H_q_nope == H_q_pe && H_q_nope == H_o && H_o == 128,
187
- "H_q_nope, H_q_pe, and H_o must be equal to 128");
188
- TORCH_CHECK(PAGE_SIZE > 0 && (PAGE_SIZE & (PAGE_SIZE - 1)) == 0,
189
- "PAGE_SIZE must be a power of 2");
190
- TORCH_CHECK(
191
- B_q_nope == B_q_pe && B_q_nope == B_pt && B_q_nope == B_o,
192
- "Batch dims must be same for page_table, q_nope and q_pe, and out");
193
- TORCH_CHECK(PAGE_NUM % (128 / PAGE_SIZE) == 0,
194
- "PAGE_NUM must be divisible by 128 / PAGE_SIZE");
195
- TORCH_CHECK(D_o == 512, "D_o must be equal to 512");
196
-
197
- TORCH_CHECK(q_nope.dtype() == at::ScalarType::Half ||
198
- q_nope.dtype() == at::ScalarType::BFloat16 ||
199
- q_nope.dtype() == at::ScalarType::Float8_e4m3fn,
200
- "q_nope must be a half, bfloat16, or float8_e4m3fn tensor");
201
- TORCH_CHECK(kv_c_and_k_pe_cache.dtype() == q_nope.dtype() &&
202
- q_nope.dtype() == q_pe.dtype(),
203
- "kv_c_and_k_pe_cache, q_nope, and q_pe must be the same type");
204
- TORCH_CHECK(seq_lens.dtype() == torch::kInt32,
205
- "seq_lens must be a 32-bit integer tensor");
206
- TORCH_CHECK(page_table.dtype() == torch::kInt32,
207
- "page_table must be a 32-bit integer tensor");
208
-
209
- auto in_dtype = q_nope.dtype();
210
- const at::cuda::OptionalCUDAGuard device_guard(device_of(q_nope));
211
- const cudaStream_t stream =
212
- at::cuda::getCurrentCUDAStream(q_nope.get_device());
213
- if (in_dtype == at::ScalarType::Half) {
214
- runMla<cutlass::half_t>(out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens,
215
- page_table, scale, stream);
216
- } else if (in_dtype == at::ScalarType::BFloat16) {
217
- runMla<cutlass::bfloat16_t>(out, q_nope, q_pe, kv_c_and_k_pe_cache,
218
- seq_lens, page_table, scale, stream);
219
- } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
220
- runMla<cutlass::float_e4m3_t>(out, q_nope, q_pe, kv_c_and_k_pe_cache,
221
- seq_lens, page_table, scale, stream);
222
- } else {
223
- TORCH_CHECK(false, "Unsupported input data type of MLA");
224
- }
225
-}
csrc/torch_bindings.cpp
-7
@@ -510,13 +510,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
510
ops.def("cutlass_sparse_compress(Tensor a) -> Tensor[]");
511
ops.impl("cutlass_sparse_compress", &cutlass_sparse_compress);
512
513
- // CUTLASS MLA decode
514
- ops.def(
515
- "cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe,"
516
- " Tensor kv_c_and_k_pe_cache, Tensor seq_lens,"
517
- " Tensor page_table, float scale) -> ()");
518
- ops.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);
519
-
513
// SM100 CUTLASS MLA decode
514
ops.def(
515
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
vllm/_custom_ops.py
-9
@@ -1823,15 +1823,6 @@ def flash_mla_with_kvcache(
1823
return out, softmax_lse
1824
1825
1826
-def cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,
1827
- q_pe: torch.Tensor, kv_c_and_k_pe_cache: torch.Tensor,
1828
- seq_lens: torch.Tensor, page_table: torch.Tensor,
1829
- scale: float) -> torch.Tensor:
1830
- torch.ops._C.cutlass_mla_decode(out, q_nope, q_pe, kv_c_and_k_pe_cache,
1831
- seq_lens, page_table, scale)
1832
- return out
1833
-
1834
-
1826
def sm100_cutlass_mla_decode(out: torch.Tensor, lse: torch.Tensor,
1827
q_nope: torch.Tensor, q_pe: torch.Tensor,
1828
kv_c_and_k_pe_cache: torch.Tensor,
vllm/v1/attention/backends/mla/cutlass_mla.py
+10
-64
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
-from typing import ClassVar, Optional
5
+from typing import ClassVar, Optional, Union
6
7
import torch
8
@@ -109,12 +109,6 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
109
"are not implemented for "
110
"CutlassMLAImpl")
111
112
- self._use_old_cutlass_mla = False
113
- force_old_cutlass = os.environ.get("FORCE_OLD_CUTLASS_MLA", None)
114
- if force_old_cutlass:
115
- logger.warning_once("Forcing old cutlass mla kernel")
116
- self._use_old_cutlass_mla = True
117
-
112
# TODO: Currently, num_kv_splits is limited to 16 to avoid hanging
113
# issues. In case the code hangs, use:
114
# FORCE_NUM_KV_SPLITS=1
@@ -219,16 +213,22 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
213
214
return out, returned_lse
215
222
- def _sm100_forward_decode(
216
+ def _forward_decode(
217
self,
224
- q_nope: torch.Tensor,
225
- q_pe: torch.Tensor,
218
+ q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
219
kv_c_and_k_pe_cache: torch.Tensor,
220
attn_metadata: MLACommonMetadata,
221
+ layer: AttentionLayer,
222
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
223
assert kv_c_and_k_pe_cache.numel() > 0
224
assert attn_metadata.decode is not None
225
226
+ if type(q) is tuple:
227
+ q_nope, q_pe = q
228
+ else:
229
+ q_nope, q_pe = torch.split(
230
+ q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
231
+
232
# Adjust workspace size (if necessary)
233
self._workspace.ensure_size(attn_metadata, self._num_kv_splits)
234
@@ -245,57 +245,3 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
245
)
246
247
return o, (lse if self.need_to_return_lse_for_decode else None)
248
-
249
- # TODO: Currently we leave it here only for backup in case something is
250
- # wrong with the new SM100 CUTLASS MLA kernel
251
- def _old_forward_decode(
252
- self,
253
- q_nope: torch.Tensor,
254
- q_pe: torch.Tensor,
255
- kv_c_and_k_pe_cache: torch.Tensor,
256
- attn_metadata: MLACommonMetadata,
257
- ) -> torch.Tensor:
258
- assert kv_c_and_k_pe_cache.numel() > 0
259
- assert attn_metadata.decode is not None
260
-
261
- if is_quantized_kv_cache(self.kv_cache_dtype):
262
- raise NotImplementedError(
263
- "FP8 Cutlass MLA not supported with FORCE_OLD_CUTLASS_MLA")
264
-
265
- B = q_nope.shape[0]
266
-
267
- o = torch.empty((B, self.num_heads, self.kv_lora_rank),
268
- dtype=q_nope.dtype,
269
- device=q_nope.device)
270
-
271
- # Run MLA
272
- # Clone q_nope and q_pe to make sure strides computation is correct.
273
- q_nope = q_nope.clone()
274
- q_pe = q_pe.clone()
275
-
276
- ops.cutlass_mla_decode(o, q_nope, q_pe, kv_c_and_k_pe_cache,
277
- attn_metadata.decode.seq_lens,
278
- attn_metadata.decode.block_table, self.scale)
279
-
280
- return o
281
-
282
- def _forward_decode(
283
- self,
284
- q: torch.Tensor,
285
- kv_c_and_k_pe_cache: torch.Tensor,
286
- attn_metadata: MLACommonMetadata,
287
- layer: AttentionLayer,
288
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
289
- if type(q) is tuple:
290
- q_nope, q_pe = q
291
- else:
292
- q_nope, q_pe = torch.split(
293
- q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
294
- if self._use_old_cutlass_mla:
295
- # TODO: Remove the old cutlass MLA kernel after more extensive
296
- # testing
297
- return self._old_forward_decode(q_nope, q_pe, kv_c_and_k_pe_cache,
298
- attn_metadata), None
299
-
300
- return self._sm100_forward_decode(q_nope, q_pe, kv_c_and_k_pe_cache,
301
- attn_metadata)