@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 86173ad593

Fork 0

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

[Kernel] Support decode context parallelism on Blackwell with CUTLASS MLA (#24385)

Signed-off-by: Ming Yang <minos.future@gmail.com> Signed-off-by: youkaichao <youkaichao@gmail.com> Co-authored-by: youkaichao <youkaichao@gmail.com>

Ming Yang committed Sep 7, 2025 at 18:27 UTC 86173ad593ff2d0c699230c4c2cd3ee5715d6e5d

5 files changed +63 -32

csrc/attention/mla/sm100_cutlass_mla_kernel.cu

+12 -5

@@ -36,6 +36,7 @@ limitations under the License.

36 #if !defined(CUDA_VERSION) || CUDA_VERSION < 12040

37 void sm100_cutlass_mla_decode(

38 torch::Tensor const& out,

39 + torch::Tensor const& lse,

40 torch::Tensor const& q_nope,

41 torch::Tensor const& q_pe,

42 torch::Tensor const& kv_c_and_k_pe_cache,

@@ -99,6 +100,7 @@ struct MlaSm100 {

100 template <typename T>

101 typename T::Fmha::Arguments args_from_options(

102 at::Tensor const& out,

103 + at::Tensor const& lse,

104 at::Tensor const& q_nope,

105 at::Tensor const& q_pe,

106 at::Tensor const& kv_c_and_k_pe_cache,

@@ -162,7 +164,10 @@ typename T::Fmha::Arguments args_from_options(

164 stride_PT,

165 page_count_total,

166 page_size},

165 - {static_cast<ElementOut*>(out.data_ptr()), stride_O, static_cast<ElementAcc*>(nullptr), stride_LSE},

167 + {static_cast<ElementOut*>(out.data_ptr()),

168 + stride_O,

169 + static_cast<ElementAcc*>(lse.defined() ? lse.data_ptr() : nullptr),

170 + stride_LSE},

171 hw_info,

172 // TODO(trevor-m): Change split_kv back to -1 when

173 // https://github.com/NVIDIA/cutlass/issues/2274 is fixed. Split_kv=1 will

@@ -181,6 +186,7 @@ typename T::Fmha::Arguments args_from_options(

186 template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>

187 void runMla(

188 at::Tensor const& out,

189 + at::Tensor const& lse,

190 at::Tensor const& q_nope,

191 at::Tensor const& q_pe,

192 at::Tensor const& kv_c_and_k_pe_cache,

@@ -192,7 +198,7 @@ void runMla(

198 cudaStream_t stream) {

199 using MlaSm100Type = MlaSm100<Element, ElementOut, IsPaged128, PersistenceOption>;

200 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);

201 + auto arguments = args_from_options<MlaSm100Type>(out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, sm_scale, num_kv_splits);

202

203 CUTLASS_CHECK(fmha.can_implement(arguments));

204

@@ -214,6 +220,7 @@ void runMla(

220

221 void sm100_cutlass_mla_decode(

222 torch::Tensor const& out,

223 + torch::Tensor const& lse,

224 torch::Tensor const& q_nope,

225 torch::Tensor const& q_pe,

226 torch::Tensor const& kv_c_and_k_pe_cache,

@@ -234,13 +241,13 @@ void sm100_cutlass_mla_decode(

241 DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {

242 if (in_dtype == at::ScalarType::Half) {

243 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);

244 + out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);

245 } else if (in_dtype == at::ScalarType::BFloat16) {

246 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);

247 + out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);

248 } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {

249 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);

250 + out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);

251 } else {

252 TORCH_CHECK(false, "Unsupported input data type of MLA");

253 }

csrc/torch_bindings.cpp

+4 -4

@@ -516,10 +516,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

516

517 // SM100 CUTLASS MLA decode

518 ops.def(

519 - "sm100_cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe,"

520 - " Tensor kv_c_and_k_pe_cache, Tensor seq_lens,"

521 - " Tensor page_table, Tensor workspace, float "

522 - "scale,"

519 + "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"

520 + " Tensor q_pe, Tensor kv_c_and_k_pe_cache,"

521 + " Tensor seq_lens, Tensor page_table,"

522 + " Tensor workspace, float scale,"

523 " int num_kv_splits) -> ()");

524 // conditionally compiled so impl in source file

525

tests/kernels/test_cutlass_mla_decode.py

+22 -10

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 import random

5 +from typing import Optional

7 import pytest

8 import torch

@@ -14,14 +15,20 @@ from vllm.triton_utils import triton

15 def cal_diff(x: torch.Tensor,

16 y: torch.Tensor,

17 name: str,

17 - use_fp8: bool = False) -> None:

18 + use_fp8: bool = False,

19 + diff_threshold: Optional[float] = None) -> None:

20 x, y = x.double(), y.double()

21 cos_diff = 1 - 2 * (x * y).sum().item() / max(

22 (x * x + y * y).sum().item(), 1e-12)

21 - if (use_fp8):

22 - assert cos_diff < 1e-4

23 + if diff_threshold is not None:

24 + # directly compare the cos_diff with the threshold

25 + assert cos_diff < diff_threshold

26 else:

24 - assert cos_diff < 1e-5

27 + # use the default threshold

28 + if (use_fp8):

29 + assert cos_diff < 1e-4

30 + else:

31 + assert cos_diff < 1e-5

34 CUTLASS_MLA_UNSUPPORTED_REASON = \

@@ -118,11 +125,13 @@ def test_cutlass_mla_decode(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size,

125 dtype=torch.uint8)

126

127 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()

128 + output_lse = torch.empty((b, MAX_HEADS),

129 + dtype=torch.float32,

130 + device=q_nope.device)

131 + ops.sm100_cutlass_mla_decode(out_ans, output_lse, q_nope, q_pe,

132 + kv_cache_flat, cache_seqlens, block_table,

133 + workspace, scale, 1)

134 + return out_ans[:, :h_q].contiguous(), output_lse[:, :h_q].contiguous()

135

136 def scaled_dot_product_attention(query, key, value, is_causal=False):

137 query = query.float()

@@ -165,11 +174,14 @@ def test_cutlass_mla_decode(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size,

174 lse[i] = lse_i

175 return out, lse

176

168 - out_cutlass = cutlass_mla()

177 + out_cutlass, lse_cutlass = cutlass_mla()

178 out_torch, lse_torch = ref_mla()

179 # Extract the single token (s_q=1) slice to match cutlass output shape

180 out_torch_slice = out_torch[:, 0, :, :] # [b, h_q, dv]

181 + lse_torch_slice = lse_torch[:, 0, :] # [b, h_q]

182 cal_diff(out_cutlass, out_torch_slice, "out", use_fp8)

183 + # lse has larger numerical error, so use a larger threshold

184 + cal_diff(lse_cutlass, lse_torch_slice, "lse", use_fp8, diff_threshold=1e-3)

185

186 t = triton.testing.do_bench(cutlass_mla)

187 FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2

vllm/_custom_ops.py

+3 -3

@@ -1833,13 +1833,13 @@ def cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,

1833 return out

1834

1835

1836 -def sm100_cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,

1837 - q_pe: torch.Tensor,

1836 +def sm100_cutlass_mla_decode(out: torch.Tensor, lse: torch.Tensor,

1837 + q_nope: torch.Tensor, q_pe: torch.Tensor,

1838 kv_c_and_k_pe_cache: torch.Tensor,

1839 seq_lens: torch.Tensor, page_table: torch.Tensor,

1840 workspace: torch.Tensor, scale: float,

1841 num_kv_splits: int) -> torch.Tensor:

1842 - torch.ops._C.sm100_cutlass_mla_decode(out, q_nope, q_pe,

1842 + torch.ops._C.sm100_cutlass_mla_decode(out, lse, q_nope, q_pe,

1843 kv_c_and_k_pe_cache, seq_lens,

1844 page_table, workspace, scale,

1845 num_kv_splits)

vllm/v1/attention/backends/mla/cutlass_mla.py

+22 -10

@@ -76,6 +76,7 @@ g_sm100_workspace = SM100Workspace(128 * 1024 * 1024) # 128MB

78 class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

79 + can_return_lse_for_decode: bool = True

81 def __init__(

82 self,

@@ -138,7 +139,7 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

139 workspace: torch.Tensor,

140 sm_scale: float,

141 num_kv_splits: int,

141 - ) -> torch.Tensor:

142 + ) -> tuple[torch.Tensor, torch.Tensor]:

143 assert (q_nope.ndim == 3

144 ), f"q_nope must be a 3D tensor, but got {q_nope.ndim}"

145 assert (

@@ -193,9 +194,13 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

194 dtype = (torch.bfloat16 if is_quantized_kv_cache(self.kv_cache_dtype)

195 else q_nope.dtype)

196 out = q_nope.new_empty((B_q, MAX_HEADS, D_latent), dtype=dtype)

197 + lse = (torch.empty(

198 + (B_q, MAX_HEADS), dtype=torch.float32, device=q_nope.device)

199 + if self.need_to_return_lse_for_decode else torch.Tensor())

200

201 ops.sm100_cutlass_mla_decode(

202 out,

203 + lse,

204 q_nope,

205 q_pe,

206 kv_c_and_k_pe_cache,

@@ -205,7 +210,9 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

210 sm_scale,

211 num_kv_splits,

212 )

208 - return out[:, :H].contiguous()

213 + returned_lse = lse[:, :H].contiguous(

214 + ) if self.need_to_return_lse_for_decode else lse

215 + return out[:, :H].contiguous(), returned_lse

216

217 def _sm100_forward_decode(

218 self,

@@ -213,7 +220,7 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

220 q_pe: torch.Tensor,

221 kv_c_and_k_pe_cache: torch.Tensor,

222 attn_metadata: MLACommonMetadata,

216 - ) -> torch.Tensor:

223 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

224 assert kv_c_and_k_pe_cache.numel() > 0

225 assert attn_metadata.decode is not None

226

@@ -226,13 +233,18 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

233 q_nope = q_nope.clone()

234 q_pe = q_pe.clone()

235

229 - o = self._sm100_cutlass_mla_decode(q_nope, q_pe, kv_c_and_k_pe_cache,

230 - attn_metadata.decode.seq_lens,

231 - attn_metadata.decode.block_table,

232 - self._workspace.get_buf(),

233 - self.scale, self._num_kv_splits)

236 + o, lse = self._sm100_cutlass_mla_decode(

237 + q_nope,

238 + q_pe,

239 + kv_c_and_k_pe_cache,

240 + attn_metadata.decode.seq_lens,

241 + attn_metadata.decode.block_table,

242 + self._workspace.get_buf(),

243 + self.scale,

244 + self._num_kv_splits,

245 + )

246

235 - return o

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

@@ -286,4 +298,4 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

298 attn_metadata), None

299

300 return self._sm100_forward_decode(q_nope, q_pe, kv_c_and_k_pe_cache,

289 - attn_metadata), None

301 + attn_metadata)