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

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

[Misc] Clean up cruft from previous FlashMLA sparse implementation (#26125)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Oct 7, 2025 at 22:09 UTC f80e7866c096c478021e04911fcdaedbd3d69930

9 files changed +79 -140

tests/kernels/attention/test_attention_selector.py

+2 -2

@@ -165,10 +165,10 @@ def test_env(

165 pytest.skip("FlashMLA only supports block_size 64")

166 else:

167 from vllm.v1.attention.backends.mla.flashmla import (

168 - is_flashmla_supported,

168 + is_flashmla_dense_supported,

169 )

170

171 - is_supported, _ = is_flashmla_supported()

171 + is_supported, _ = is_flashmla_dense_supported()

172 if not is_supported:

173 pytest.skip("FlashMLA not supported on this platform")

174 else:

tests/kernels/attention/test_flashmla.py

+6 -4

@@ -10,7 +10,7 @@ import torch

10 from vllm.attention.ops.flashmla import (

11 flash_mla_with_kvcache,

12 get_mla_metadata,

13 - is_flashmla_supported,

13 + is_flashmla_dense_supported,

14 )

15 from vllm.triton_utils import triton

@@ -27,13 +27,15 @@ def cal_diff(

29 FLASH_MLA_UNSUPPORTED_REASON = (

30 - is_flashmla_supported()[1]

31 - if not is_flashmla_supported()[0]

30 + is_flashmla_dense_supported()[1]

31 + if not is_flashmla_dense_supported()[0]

32 else "FlashMLA is supported"

33 )

36 -@pytest.mark.skipif(not is_flashmla_supported()[0], reason=FLASH_MLA_UNSUPPORTED_REASON)

36 +@pytest.mark.skipif(

37 + not is_flashmla_dense_supported()[0], reason=FLASH_MLA_UNSUPPORTED_REASON

38 +)

39 @pytest.mark.parametrize("b", [128])

40 @pytest.mark.parametrize("s_q", [1, 2])

41 @pytest.mark.parametrize("mean_sk", [4096, 8192, 16384])

tests/kernels/attention/test_flashmla_sparse.py

+9 -16

@@ -4,19 +4,12 @@ import pytest

4 import torch

7 -def _cuda_sm90_available() -> bool:

8 - if not torch.cuda.is_available():

9 - return False

10 - major, _ = torch.cuda.get_device_capability()

11 - return major == 9

12 -

13 -

7 def test_sparse_flashmla_metadata_smoke():

8 import vllm.attention.ops.flashmla as fm

17 - ok, reason = fm.is_flashmla_supported()

18 - if not ok or not _cuda_sm90_available():

19 - pytest.skip(reason or "SM90 not available")

10 + ok, reason = fm.is_flashmla_sparse_supported()

11 + if not ok:

12 + pytest.skip(reason)

14 device = torch.device("cuda")

15 batch_size = 1

@@ -43,9 +36,9 @@ def test_sparse_flashmla_metadata_smoke():

36 def test_sparse_flashmla_decode_smoke():

37 import vllm.attention.ops.flashmla as fm

46 - ok, reason = fm.is_flashmla_supported()

47 - if not ok or not _cuda_sm90_available():

48 - pytest.skip(reason or "SM90 not available")

39 + ok, reason = fm.is_flashmla_sparse_supported()

40 + if not ok:

41 + pytest.skip(reason)

43 device = torch.device("cuda")

44 batch_size = 1

@@ -106,9 +99,9 @@ def test_sparse_flashmla_decode_smoke():

99 def test_sparse_flashmla_prefill_smoke():

100 import vllm.attention.ops.flashmla as fm

101

109 - ok, reason = fm.is_flashmla_supported()

110 - if not ok or not _cuda_sm90_available():

111 - pytest.skip(reason or "SM90 not available")

102 + ok, reason = fm.is_flashmla_sparse_supported()

103 + if not ok:

104 + pytest.skip(reason)

105

106 device = torch.device("cuda")

107 s_q = 1

tests/v1/attention/test_sparse_mla_backends.py

+17 -66

@@ -24,12 +24,7 @@ from vllm import _custom_ops as ops

24 from vllm.attention.ops import flashmla

25 from vllm.model_executor.layers.linear import ColumnParallelLinear

26 from vllm.utils import cdiv

27 -from vllm.v1.attention.backends.mla.flashmla_sparse import (

28 - FlashMLASparseBackend,

29 - FlashMLASparseDecodeAndContextMetadata,

30 - FlashMLASparseImpl,

31 - FlashMLASparseMetadata,

32 -)

27 +from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend

28 from vllm.v1.attention.backends.mla.indexer import split_prefill_chunks

30 SPARSE_BACKEND_BATCH_SPECS = {

@@ -116,59 +111,6 @@ def _quantize_dequantize_fp8_ds_mla(

111 return dequant_kv_c, dequant_k_pe

112

113

119 -def test_sparse_backend_metadata_registration():

120 - backend = FlashMLASparseBackend

121 -

122 - assert backend.get_name() == "FLASHMLA_SPARSE_VLLM_V1"

123 - assert backend.get_metadata_cls() is FlashMLASparseMetadata

124 - assert backend.get_impl_cls() is FlashMLASparseImpl

125 -

126 - dtype_list = backend.get_supported_dtypes()

127 - assert torch.bfloat16 in dtype_list

128 -

129 - shape = backend.get_kv_cache_shape(

130 - num_blocks=2, block_size=64, num_kv_heads=1, head_size=576

131 - )

132 - assert shape == (2, 64, 576)

133 -

134 -

135 -def test_sparse_decode_metadata_filters_prefill_indices():

136 - prefill_context_lengths = torch.tensor([4, 2], dtype=torch.int32)

137 - metadata = FlashMLASparseDecodeAndContextMetadata(

138 - scheduler_metadata=torch.tensor([[0]], dtype=torch.int32),

139 - num_splits=torch.tensor([1, 1], dtype=torch.int32),

140 - cache_lens=torch.tensor([10, 12], dtype=torch.int32),

141 - prefill_context_lengths=prefill_context_lengths,

142 - )

143 -

144 - indices = torch.tensor([[0, 3, 5], [1, 2, 4]], dtype=torch.int32)

145 -

146 - context_indices, new_token_indices = metadata.filter_prefill_indices(indices)

147 -

148 - expected_context = torch.tensor([[-1, -1, 5], [-1, -1, 4]], dtype=torch.int32)

149 - expected_new_tokens = torch.tensor([[-1, -1, 1], [-1, 0, 2]], dtype=torch.int32)

150 -

151 - assert torch.equal(context_indices, expected_context)

152 - assert torch.equal(new_token_indices, expected_new_tokens)

153 -

154 -

155 -def test_sparse_impl_zero_fills_when_metadata_missing():

156 - impl = FlashMLASparseImpl.__new__(FlashMLASparseImpl)

157 - dummy_layer = object()

158 - q = torch.zeros((2, 1, 3))

159 - k_c = torch.zeros((2, 3))

160 - k_pe = torch.zeros((2, 1, 1))

161 - kv_cache = torch.zeros((1, 1, 1))

162 - output = torch.ones((2, 4))

163 -

164 - result = FlashMLASparseImpl.forward(

165 - impl, dummy_layer, q, k_c, k_pe, kv_cache, attn_metadata=None, output=output

166 - )

167 -

168 - assert result is output

169 - assert torch.all(result == 0)

170 -

171 -

114 @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))

115 @pytest.mark.parametrize("kv_cache_dtype", ["fp8_ds_mla", "auto"])

116 def test_sparse_backend_decode_correctness(dist_init, batch_name, kv_cache_dtype):

@@ -198,11 +140,12 @@ def test_sparse_backend_decode_correctness(dist_init, batch_name, kv_cache_dtype

140 max_model_len=max_seqlen,

141 num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),

142 block_size=block_size,

143 + hf_config_override={

144 + "index_topk": topk_tokens,

145 + "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],

146 + },

147 )

148 model_config = vllm_config.model_config

203 - model_config.hf_config = SimpleNamespace(

204 - attn_module_list_cfg=[{"topk_tokens": topk_tokens}]

205 - )

149 model_config.hf_text_config = SimpleNamespace(

150 q_lora_rank=None,

151 kv_lora_rank=kv_lora_rank,

@@ -301,6 +244,7 @@ def test_sparse_backend_decode_correctness(dist_init, batch_name, kv_cache_dtype

244 sdpa_reference = torch.cat(reference_outputs, dim=0)

245

246 vllm_config.cache_config.cache_dtype = kv_cache_dtype

247 + vllm_config.model_config.hf_config.index_topk = topk_tokens

248

249 common_attn_metadata = create_common_attn_metadata(

250 batch_spec,

@@ -352,7 +296,7 @@ def test_sparse_backend_decode_correctness(dist_init, batch_name, kv_cache_dtype

296 debug_indices = debug_indices.expand(metadata.num_actual_tokens, -1).clone()

297 mock_indexer = SimpleNamespace(topk_indices_buffer=debug_indices)

298

355 - ok, reason = flashmla.is_flashmla_supported()

299 + ok, reason = flashmla.is_flashmla_sparse_supported()

300 if not ok:

301 pytest.skip(reason)

302

@@ -397,9 +341,16 @@ def test_sparse_backend_decode_correctness(dist_init, batch_name, kv_cache_dtype

341 metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device

342 )

343

400 - backend_output = impl.forward(

401 - layer, query_vllm, kv_c_vllm, k_pe_vllm, kv_cache, metadata, output=out_buffer

402 - )

344 + with torch.inference_mode():

345 + backend_output = impl.forward(

346 + layer,

347 + query_vllm,

348 + kv_c_vllm,

349 + k_pe_vllm,

350 + kv_cache,

351 + metadata,

352 + output=out_buffer,

353 + )

354

355 assert backend_output.shape == sdpa_reference.shape

356 assert backend_output.dtype == sdpa_reference.dtype

tests/v1/attention/utils.py

@@ -181,6 +181,7 @@ def create_vllm_config(

181 max_num_batched_tokens: int = 8192,

182 enable_chunked_prefill: bool = True,

183 add_mock_model_methods: bool = True,

184 + hf_config_override: Optional[dict] = None,

185 ) -> VllmConfig:

186 """Create a VllmConfig for testing with reasonable defaults."""

187

@@ -235,6 +236,9 @@ def create_vllm_config(

236 lambda self, i: 1.0 / model_config.get_head_size() ** 0.5, model_config

237 )

238

239 + if hf_config_override:

240 + model_config.hf_config.update(hf_config_override)

241 +

242 return VllmConfig(

243 model_config=model_config,

244 cache_config=cache_config,

vllm/attention/ops/flashmla.py

+35 -9

@@ -31,21 +31,47 @@ else:

31 _flashmla_extension_C_AVAILABLE = False

34 -def is_flashmla_supported() -> tuple[bool, Optional[str]]:

34 +def _is_flashmla_available() -> tuple[bool, Optional[str]]:

35 + if not _flashmla_C_AVAILABLE:

36 + return (

37 + False,

38 + "vllm._flashmla_C is not available, likely was not "

39 + "compiled due to insufficient nvcc version or a supported arch "

40 + "was not in the list of target arches to compile for.",

41 + )

42 + if not _flashmla_extension_C_AVAILABLE:

43 + return (

44 + False,

45 + "vllm._flashmla_extension_C is not available, likely "

46 + "was not compiled due to a build error.",

47 + )

48 +

49 + return True, None

50 +

51 +

52 +def is_flashmla_dense_supported() -> tuple[bool, Optional[str]]:

53 """

54 Return: is_supported_flag, unsupported_reason (optional).

55 """

38 - if not current_platform.is_cuda():

39 - return False, "FlashMLA is only supported on CUDA devices."

56 + is_availble, maybe_reason = _is_flashmla_available()

57 + if not is_availble:

58 + return False, maybe_reason

59 if current_platform.get_device_capability()[0] != 9:

41 - return False, "FlashMLA is only supported on Hopper devices."

42 - if not _flashmla_C_AVAILABLE:

60 + return False, "FlashMLA Dense is only supported on Hopper devices."

61 + return True, None

62 +

63 +

64 +def is_flashmla_sparse_supported() -> tuple[bool, Optional[str]]:

65 + """

66 + Return: is_supported_flag, unsupported_reason (optional).

67 + """

68 + is_availble, maybe_reason = _is_flashmla_available()

69 + if not is_availble:

70 + return False, maybe_reason

71 + if current_platform.get_device_capability()[0] not in (9, 10):

72 return (

73 False,

45 - "vllm._flashmla_C is not available, likely was not "

46 - "compiled due to insufficient nvcc version or a supported arch "

47 - "(only sm90a currently) was not in the list of target arches to "

48 - "compile for.",

74 + "FlashMLA Sparse is only supported on Hopper and Blackwell devices.",

75 )

76 return True, None

vllm/platforms/cuda.py

+4 -4

@@ -146,11 +146,11 @@ class CudaPlatformBase(Platform):

146 use_cutlass_mla = envs.VLLM_ATTENTION_BACKEND == "CUTLASS_MLA"

147 use_flashinfer_mla = envs.VLLM_ATTENTION_BACKEND == "FLASHINFER_MLA"

148

149 - from vllm.attention.ops.flashmla import is_flashmla_supported

149 + from vllm.attention.ops.flashmla import is_flashmla_dense_supported

150

151 if (

152 use_flashmla

153 - and is_flashmla_supported()[0]

153 + and is_flashmla_dense_supported()[0]

154 and cache_config.block_size != 64

155 ):

156 cache_config.block_size = 64

@@ -256,7 +256,7 @@ class CudaPlatformBase(Platform):

256 "Set VLLM_USE_V1=1 to enable them."

257 )

258

259 - from vllm.attention.ops.flashmla import is_flashmla_supported

259 + from vllm.attention.ops.flashmla import is_flashmla_dense_supported

260 from vllm.attention.utils.fa_utils import flash_attn_supports_mla

261

262 if use_sparse:

@@ -277,7 +277,7 @@ class CudaPlatformBase(Platform):

277 and block_size in [32, 64]

278 )

279 use_flashmla = selected_backend == _Backend.FLASHMLA or (

280 - selected_backend is None and is_flashmla_supported()[0]

280 + selected_backend is None and is_flashmla_dense_supported()[0]

281 )

282 use_flashattn = selected_backend == _Backend.FLASH_ATTN_MLA or (

283 selected_backend is None and flash_attn_supports_mla()

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

+2 -2

@@ -10,7 +10,7 @@ from vllm.attention.backends.abstract import AttentionLayer, AttentionType

10 from vllm.attention.ops.flashmla import (

11 flash_mla_with_kvcache,

12 get_mla_metadata,

13 - is_flashmla_supported,

13 + is_flashmla_dense_supported,

14 )

15 from vllm.config import VllmConfig

16 from vllm.logger import init_logger

@@ -177,7 +177,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

177 **mla_args,

178 )

179

180 - is_supported, reason = is_flashmla_supported()

180 + is_supported, reason = is_flashmla_dense_supported()

181 assert is_supported, reason

182

183 unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]

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

-37

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -import math

3 from dataclasses import dataclass

4 from typing import TYPE_CHECKING, ClassVar, Optional

@@ -51,12 +50,6 @@ structured as:

50 """

54 -def _lse2_to_lse(lse_base2: torch.Tensor) -> torch.Tensor:

55 - # Convert base-2 LSE to natural-log LSE

56 - # Keep FP32 for numerical stability during the merge.

57 - return lse_base2.to(torch.float32) * math.log(2.0)

58 -

59 -

53 class FlashMLASparseBackend(AttentionBackend):

54 accept_output_buffer: bool = True

@@ -100,36 +93,6 @@ class FlashMLASparseBackend(AttentionBackend):

93 return [576]

103 -@dataclass

104 -class MLASparsePrefillMetadata:

105 - # NOTE(Chen): not call it "FlashMLASparsePrefillMetadata" because

106 - # the kernel is not from flashmla

107 - block_table: torch.Tensor

108 - has_context: bool = False

109 - context_lens: Optional[torch.Tensor] = None

110 -

111 -

112 -@dataclass

113 -class FlashMLASparseDecodeAndContextMetadata:

114 - scheduler_metadata: torch.Tensor = None

115 - num_splits: torch.Tensor = None

116 - cache_lens: torch.Tensor = None

117 - prefill_context_lengths: Optional[torch.Tensor] = None

118 - prefill_new_k_start_locs: Optional[torch.Tensor] = None

119 - dummy_block_table: torch.Tensor = None

120 -

121 - def filter_prefill_indices(

122 - self, indices: torch.Tensor

123 - ) -> tuple[torch.Tensor, torch.Tensor]:

124 - assert self.prefill_context_lengths is not None

125 - prefill_context_lengths = self.prefill_context_lengths.unsqueeze(-1)

126 - context_indices = torch.where(indices < prefill_context_lengths, indices, -1)

127 - new_token_indices = torch.where(

128 - indices >= prefill_context_lengths, indices - prefill_context_lengths, -1

129 - )

130 - return context_indices, new_token_indices

131 -

132 -

96 @dataclass

97 class FlashMLASparseMetadata:

98 num_reqs: int