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

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 1, 18:45 UTC

[DSV4] Decouple DS V4 Sparse MLA Metadata from DS V3.2 (#44699)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

Woosuk Kwon committed Jun 6, 2026 at 17:37 UTC 2a983c79acdbe78fa62e67eaa87ed4bbb7c8257e

7 files changed +449 -333

docs/design/attention_backends.md

+1 -1

@@ -240,5 +240,5 @@ default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.

241 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |

242 | `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

243 -| `FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | 256 | 512 | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

243 +| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |

244 | `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |

vllm/models/deepseek_v4/amd/rocm.py

+8 -10

@@ -9,17 +9,15 @@ import torch

9 from vllm.forward_context import get_forward_context

10 from vllm.models.deepseek_v4.attention import DeepseekV4Attention

11 from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache

12 -from vllm.models.deepseek_v4.nvidia.flashmla import (

13 - DeepseekV4FlashMLASparseBackend,

12 +from vllm.models.deepseek_v4.sparse_mla import (

13 + DeepseekV4FlashMLABackend,

14 + DeepseekV4FlashMLAMetadata,

15 + DeepseekV4FlashMLAMetadataBuilder,

16 )

17 from vllm.triton_utils import tl, triton

18 from vllm.v1.attention.backend import (

19 CommonAttentionMetadata,

20 )

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

20 - FlashMLASparseMetadata,

21 - FlashMLASparseMetadataBuilder,

22 -)

21 from vllm.v1.attention.backends.mla.sparse_swa import (

22 DeepseekSparseSWAMetadata,

23 DeepseekSparseSWAMetadataBuilder,

@@ -445,7 +443,7 @@ def _copy_ragged_to_graph_buffers(

443

444

445 @dataclass

448 -class DeepseekV4ROCMAiterMLASparseMetadata(FlashMLASparseMetadata):

446 +class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata):

447 """ROCm-specific DeepSeek V4 metadata carrying ragged decode topk."""

448

449 c128a_decode_topk_ragged_indices: torch.Tensor | None = None

@@ -458,12 +456,12 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):

456 decode_swa_ragged_indptr: torch.Tensor | None = None

457

458

461 -class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder):

459 +class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4FlashMLAMetadataBuilder):

460 def __init__(self, *args, **kwargs):

461 super().__init__(*args, **kwargs)

462 self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None

463 self.c128a_decode_topk_ragged_indptr_buffer: torch.Tensor | None = None

466 - if self.is_deepseek_v4 and self.compress_ratio == 128:

464 + if self.compress_ratio == 128:

465 max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens

466 self.c128a_decode_topk_ragged_indices_buffer = torch.empty(

467 max_tokens * self.c128a_max_compressed,

@@ -569,7 +567,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild

567 )

568

569

572 -class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend):

570 +class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLABackend):

571 @staticmethod

572 def get_name() -> str:

573 return "ROCM_FLASHMLA_SPARSE_DSV4"

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

+13 -10

@@ -18,13 +18,15 @@ from vllm.models.deepseek_v4.attention import DeepseekV4Attention

18 from vllm.models.deepseek_v4.common.ops import (

19 build_flashinfer_mixed_sparse_indices,

20 )

21 -from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLASparseBackend

21 from vllm.models.deepseek_v4.nvidia.ops.o_proj import (

22 compute_fp8_einsum_recipe,

23 deep_gemm_fp8_o_proj,

24 )

25 +from vllm.models.deepseek_v4.sparse_mla import (

26 + DeepseekV4FlashMLABackend,

27 + DeepseekV4FlashMLAMetadata,

28 +)

29 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4

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

31 if TYPE_CHECKING:

32 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata

@@ -47,13 +49,14 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:

49 return workspace

50 -class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLASparseBackend):

52 +class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):

53 """Shares the FlashMLA V4 metadata/cache pipeline; swaps the attention impl.

53 - Inheriting from the FlashMLA V4 backend reuses its ``FlashMLASparseMetadata``

54 - builder (which the V4 sparse-index pipeline needs — the V3.2 FlashInfer

55 - builder lacks the ``c128a_*`` fields), 256-token blocks, head_size 512, and

56 - the (num_blocks, block_size, 512) cache shape for non-``fp8_ds_mla`` dtypes.

55 + Inheriting from the FlashMLA V4 backend reuses its

56 + ``DeepseekV4FlashMLAMetadata`` builder (which the V4 sparse-index

57 + pipeline needs — the V3.2 FlashInfer builder lacks the ``c128a_*`` fields),

58 + 256-token blocks, head_size 512, and the (num_blocks, block_size, 512) cache

59 + shape for non-``fp8_ds_mla`` dtypes.

60 """

62 @staticmethod

@@ -162,7 +165,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

165

166 assert isinstance(attn_metadata, dict)

167 flashmla_metadata = cast(

165 - FlashMLASparseMetadata | None, attn_metadata.get(self.prefix)

168 + DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)

169 )

170 swa_metadata = cast(

171 "DeepseekSparseSWAMetadata | None",

@@ -190,7 +193,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

193 kv_cache: torch.Tensor | None,

194 swa_k_cache: torch.Tensor,

195 swa_metadata: "DeepseekSparseSWAMetadata",

193 - attn_metadata: FlashMLASparseMetadata | None,

196 + attn_metadata: DeepseekV4FlashMLAMetadata | None,

197 swa_only: bool,

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

199 """Build the combined sparse-index tensors for the mixed batch.

@@ -310,7 +313,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

313 kv_cache: torch.Tensor | None,

314 swa_k_cache: torch.Tensor,

315 swa_metadata: "DeepseekSparseSWAMetadata",

313 - attn_metadata: FlashMLASparseMetadata | None,

316 + attn_metadata: DeepseekV4FlashMLAMetadata | None,

317 swa_only: bool,

318 output: torch.Tensor,

319 ) -> None:

vllm/models/deepseek_v4/nvidia/flashmla.py

+7 -39

@@ -16,10 +16,9 @@ from vllm.models.deepseek_v4.nvidia.ops.o_proj import (

16 compute_fp8_einsum_recipe,

17 deep_gemm_fp8_o_proj,

18 )

19 -from vllm.v1.attention.backend import MultipleOf

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

21 - FlashMLASparseBackend,

22 - FlashMLASparseMetadata,

19 +from vllm.models.deepseek_v4.sparse_mla import (

20 + DeepseekV4FlashMLABackend,

21 + DeepseekV4FlashMLAMetadata,

22 )

23 from vllm.v1.attention.ops.flashmla import (

24 flash_mla_sparse_fwd,

@@ -31,41 +30,10 @@ if TYPE_CHECKING:

30 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata

34 -class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):

35 - @staticmethod

36 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

37 - return [256]

38 -

39 - @staticmethod

40 - def get_name() -> str:

41 - return "FLASHMLA_SPARSE_DSV4"

42 -

43 - @classmethod

44 - def get_supported_head_sizes(cls) -> list[int]:

45 - # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512 (overrides the

46 - # V3.2 default of 576 from FlashMLASparseBackend).

47 - return [512]

48 -

49 - @staticmethod

50 - def get_kv_cache_shape(

51 - num_blocks: int,

52 - block_size: int,

53 - num_kv_heads: int,

54 - head_size: int,

55 - cache_dtype_str: str = "auto",

56 - ) -> tuple[int, ...]:

57 - if cache_dtype_str == "fp8_ds_mla":

58 - # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).

59 - # head_size passed in is the semantic head_dim (512).

60 - return (num_blocks, block_size, 584)

61 - else:

62 - return (num_blocks, block_size, head_size)

63 -

64 -

33 class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

34 """FlashMLA sparse MLA attention layer for DeepSeek V4 (CUDA)."""

68 - backend_cls = DeepseekV4FlashMLASparseBackend

36 + backend_cls = DeepseekV4FlashMLABackend

38 def __init__(self, *args, **kwargs) -> None:

39 super().__init__(*args, **kwargs)

@@ -135,7 +103,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

103

104 assert isinstance(attn_metadata, dict)

105 flashmla_metadata = cast(

138 - FlashMLASparseMetadata | None, attn_metadata.get(self.prefix)

106 + DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)

107 )

108 swa_metadata = cast(

109 "DeepseekSparseSWAMetadata | None",

@@ -179,7 +147,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

147 q: torch.Tensor,

148 kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1

149 swa_metadata: "DeepseekSparseSWAMetadata",

182 - attn_metadata: FlashMLASparseMetadata | None,

150 + attn_metadata: DeepseekV4FlashMLAMetadata | None,

151 swa_only: bool,

152 output: torch.Tensor,

153 ) -> None:

@@ -273,7 +241,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

241 compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1

242 swa_k_cache: torch.Tensor,

243 output: torch.Tensor,

276 - attn_metadata: FlashMLASparseMetadata | None,

244 + attn_metadata: DeepseekV4FlashMLAMetadata | None,

245 swa_metadata: "DeepseekSparseSWAMetadata",

246 ) -> None:

247 swa_only = attn_metadata is None

vllm/models/deepseek_v4/sparse_mla.py new

+416

@@ -0,0 +1,416 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +"""DeepSeek-V4 FlashMLA sparse backend, metadata, and metadata builder."""

4 +

5 +from dataclasses import dataclass

6 +from typing import Any, ClassVar

7 +

8 +import numpy as np

9 +import torch

10 +

11 +from vllm.config import VllmConfig

12 +from vllm.config.cache import CacheDType

13 +from vllm.platforms.interface import DeviceCapability

14 +from vllm.triton_utils import tl, triton

15 +from vllm.utils.math_utils import cdiv

16 +from vllm.v1.attention.backend import (

17 + AttentionBackend,

18 + AttentionCGSupport,

19 + AttentionMetadata,

20 + AttentionMetadataBuilder,

21 + CommonAttentionMetadata,

22 + MultipleOf,

23 +)

24 +from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping

25 +from vllm.v1.attention.backends.utils import split_decodes_and_prefills

26 +from vllm.v1.kv_cache_interface import AttentionSpec

27 +

28 +# Pad C128A topk width to this alignment. 128 covers both h_q=64 (B_TOPK=64) and

29 +# h_q=128 (B_TOPK=128). FlashMLA decode asserts extra_topk % B_TOPK == 0;

30 +# unaligned widths (e.g. 17 = ceil(2136/128)) crash the sm100 head64 kernel.

31 +# Padded slots stay -1 and decode_lens caps them via topk_length, so the pad is a

32 +# no-op at kernel level. Mirrors _SPARSE_PREFILL_TOPK_ALIGNMENT in cache_utils.py.

33 +_C128A_TOPK_ALIGNMENT = 128

34 +

35 +

36 +class DeepseekV4FlashMLABackend(AttentionBackend):

37 + """DeepSeek-V4 sparse-MLA backend.

38 +

39 + Subclasses ``AttentionBackend`` directly (not the V3.2

40 + ``FlashMLASparseBackend``): DeepSeek-V4 runs its own attention layer

41 + (``DeepseekV4Attention``), so it does not reuse the V3.2 builder or impl, and

42 + only needs to declare its own metadata builder, KV-cache layout, and the

43 + sparse-MLA capability flags.

44 + """

45 +

46 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

47 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

48 + "auto",

49 + "bfloat16",

50 + "fp8_ds_mla",

51 + "fp8", # alias for fp8_ds_mla

52 + ]

53 +

54 + @staticmethod

55 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

56 + return [256]

57 +

58 + @staticmethod

59 + def get_name() -> str:

60 + return "FLASHMLA_SPARSE_DSV4"

61 +

62 + @staticmethod

63 + def get_builder_cls() -> type["DeepseekV4FlashMLAMetadataBuilder"]:

64 + return DeepseekV4FlashMLAMetadataBuilder

65 +

66 + @staticmethod

67 + def get_impl_cls() -> type[Any]:

68 + # DeepSeek-V4 runs its attention through ``DeepseekV4Attention.forward``,

69 + # not the generic ``Attention``/``MLAAttention`` layer, so the backend's

70 + # impl class is never instantiated.

71 + raise NotImplementedError(

72 + "DeepseekV4FlashMLABackend has no separate impl class; DeepSeek-V4 "

73 + "attention runs through DeepseekV4Attention."

74 + )

75 +

76 + @classmethod

77 + def get_supported_head_sizes(cls) -> list[int]:

78 + # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512.

79 + return [512]

80 +

81 + @classmethod

82 + def is_mla(cls) -> bool:

83 + return True

84 +

85 + @classmethod

86 + def is_sparse(cls) -> bool:

87 + return True

88 +

89 + @classmethod

90 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:

91 + return capability.major in [9, 10]

92 +

93 + @staticmethod

94 + def get_kv_cache_shape(

95 + num_blocks: int,

96 + block_size: int,

97 + num_kv_heads: int,

98 + head_size: int,

99 + cache_dtype_str: str = "auto",

100 + ) -> tuple[int, ...]:

101 + if cache_dtype_str == "fp8_ds_mla":

102 + # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).

103 + # head_size passed in is the semantic head_dim (512).

104 + return (num_blocks, block_size, 584)

105 + else:

106 + return (num_blocks, block_size, head_size)

107 +

108 +

109 +@dataclass

110 +class DeepseekV4FlashMLAMetadata(AttentionMetadata):

111 + num_reqs: int

112 + max_query_len: int

113 + max_seq_len: int

114 +

115 + num_actual_tokens: int # Number of tokens excluding padding.

116 + query_start_loc: torch.Tensor

117 + slot_mapping: torch.Tensor

118 +

119 + block_table: torch.Tensor

120 + req_id_per_token: torch.Tensor

121 + block_size: int

122 + topk_tokens: int

123 +

124 + # Pre-computed C128A metadata (compress_ratio == 128 only).

125 + # Decode: global slot ids + valid-entry counts (fused from positions).

126 + c128a_global_decode_topk_indices: torch.Tensor | None = None

127 + c128a_decode_topk_lens: torch.Tensor | None = None

128 + # Prefill: local topk indices (used by combine_topk_swa_indices).

129 + c128a_prefill_topk_indices: torch.Tensor | None = None

130 +

131 +

132 +class DeepseekV4FlashMLAMetadataBuilder(

133 + AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata]

134 +):

135 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

136 +

137 + def __init__(

138 + self,

139 + kv_cache_spec: AttentionSpec,

140 + layer_names: list[str],

141 + vllm_config: VllmConfig,

142 + device: torch.device,

143 + ) -> None:

144 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

145 + self.model_config = vllm_config.model_config

146 + # Classify single-token queries (plus num_speculative_tokens via

147 + # supports_spec_as_decode=True) as decodes; longer queries go to prefill.

148 + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

149 + self.topk_tokens = self.model_config.hf_config.index_topk

150 +

151 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens

152 + self.req_id_per_token_buffer = torch.empty(

153 + (max_num_batched_tokens,), dtype=torch.int32, device=device

154 + )

155 +

156 + assert hasattr(self.kv_cache_spec, "compress_ratio")

157 + self.compress_ratio = self.kv_cache_spec.compress_ratio

158 +

159 + # Pre-allocate compressed slot mapping buffer for CUDA graph address

160 + # stability when compress_ratio > 1.

161 + if self.compress_ratio > 1:

162 + self.compressed_slot_mapping_buffer = torch.empty(

163 + max_num_batched_tokens, dtype=torch.int64, device=device

164 + )

165 +

166 + # Pre-allocate C128A topk buffers for CUDA graph address stability.

167 + if self.compress_ratio == 128:

168 + c128a_max_compressed = cdiv(

169 + self.model_config.max_model_len, self.compress_ratio

170 + )

171 + c128a_max_compressed = (

172 + cdiv(c128a_max_compressed, _C128A_TOPK_ALIGNMENT)

173 + * _C128A_TOPK_ALIGNMENT

174 + )

175 + # Stored so _build_c128a_metadata passes it as the kernel's

176 + # max_compressed_tokens, matching the buffer stride. Otherwise the

177 + # kernel's default 8192 iterates past row width and spills writes

178 + # into adjacent rows (present in both decode and prefill branches of

179 + # _build_c128a_topk_metadata_kernel).

180 + self.c128a_max_compressed = c128a_max_compressed

181 + self.c128a_global_decode_buffer = torch.empty(

182 + (max_num_batched_tokens, c128a_max_compressed),

183 + dtype=torch.int32,

184 + device=device,

185 + )

186 + self.c128a_decode_lens_buffer = torch.empty(

187 + max_num_batched_tokens, dtype=torch.int32, device=device

188 + )

189 + self.c128a_prefill_buffer = torch.empty(

190 + (max_num_batched_tokens, c128a_max_compressed),

191 + dtype=torch.int32,

192 + device=device,

193 + )

194 +

195 + def build(

196 + self,

197 + common_prefix_len: int,

198 + common_attn_metadata: CommonAttentionMetadata,

199 + fast_build: bool = False,

200 + ) -> DeepseekV4FlashMLAMetadata:

201 + cm = common_attn_metadata

202 + num_tokens = cm.num_actual_tokens

203 + starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)

204 + seg_lengths = np.diff(starts)

205 + req_id_per_token = np.repeat(

206 + np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths

207 + )

208 + # Zero-fill for cudagraphs

209 + self.req_id_per_token_buffer.fill_(0)

210 + self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

211 + torch.from_numpy(req_id_per_token), non_blocking=True

212 + )

213 + req_id_per_token = self.req_id_per_token_buffer[:num_tokens]

214 +

215 + slot_mapping = cm.slot_mapping

216 + if self.compress_ratio > 1:

217 + slot_mapping = get_compressed_slot_mapping(

218 + cm.num_actual_tokens,

219 + cm.query_start_loc,

220 + cm.seq_lens,

221 + cm.block_table_tensor.clamp(min=0),

222 + int(self.kv_cache_spec.storage_block_size),

223 + self.compress_ratio,

224 + out=self.compressed_slot_mapping_buffer,

225 + )

226 +

227 + c128a_fields: dict[str, torch.Tensor | None] = {}

228 + if self.compress_ratio == 128:

229 + c128a_fields = self._build_c128a_metadata(cm, req_id_per_token)

230 +

231 + return DeepseekV4FlashMLAMetadata(

232 + num_reqs=cm.num_reqs,

233 + max_query_len=cm.max_query_len,

234 + max_seq_len=cm.max_seq_len,

235 + num_actual_tokens=cm.num_actual_tokens,

236 + query_start_loc=cm.query_start_loc,

237 + slot_mapping=slot_mapping,

238 + block_table=cm.block_table_tensor,

239 + req_id_per_token=req_id_per_token,

240 + block_size=self.kv_cache_spec.block_size,

241 + topk_tokens=self.topk_tokens,

242 + c128a_global_decode_topk_indices=c128a_fields.get(

243 + "c128a_global_decode_topk_indices"

244 + ),

245 + c128a_decode_topk_lens=c128a_fields.get("c128a_decode_topk_lens"),

246 + c128a_prefill_topk_indices=c128a_fields.get("c128a_prefill_topk_indices"),

247 + )

248 +

249 + def _build_c128a_metadata(

250 + self,

251 + cm: CommonAttentionMetadata,

252 + req_id_per_token: torch.Tensor,

253 + ) -> dict[str, torch.Tensor | None]:

254 + """Pre-compute C128A topk indices for DeepseekV4 (compress_ratio >= 128)."""

255 + # Must match SWA's decode split (no `require_uniform=True`) so

256 + # `c128a_global_decode_topk_indices.shape[0]` lines up with q in

257 + # `_forward_decode`. The per-token C128A kernel handles non-uniform

258 + # query lengths.

259 + (num_decodes, _, num_decode_tokens, num_prefill_tokens) = (

260 + split_decodes_and_prefills(

261 + cm,

262 + decode_threshold=self.reorder_batch_threshold or 1,

263 + )

264 + )

265 +

266 + num_total = num_decode_tokens + num_prefill_tokens

267 + if num_total == 0:

268 + return {}

269 +

270 + assert cm.positions is not None, (

271 + "positions is required for C128A metadata build"

272 + )

273 + block_size = self.kv_cache_spec.block_size // self.compress_ratio

274 + global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(

275 + cm.positions[:num_total],

276 + self.compress_ratio,

277 + num_decode_tokens,

278 + req_id_per_token,

279 + cm.block_table_tensor[:num_decodes],

280 + block_size,

281 + cm.slot_mapping,

282 + self.c128a_global_decode_buffer,

283 + self.c128a_decode_lens_buffer,

284 + self.c128a_prefill_buffer,

285 + max_compressed_tokens=self.c128a_max_compressed,

286 + )

287 +

288 + result: dict[str, torch.Tensor | None] = {}

289 + if num_decode_tokens > 0:

290 + result["c128a_global_decode_topk_indices"] = global_decode.view(

291 + num_decode_tokens, 1, -1

292 + )

293 + result["c128a_decode_topk_lens"] = decode_lens

294 + if num_prefill_tokens > 0:

295 + result["c128a_prefill_topk_indices"] = prefill_local

296 + return result

297 +

298 +

299 +def build_c128a_topk_metadata(

300 + positions: torch.Tensor,

301 + compress_ratio: int,

302 + num_decode_tokens: int,

303 + token_to_req_indices: torch.Tensor,

304 + block_table: torch.Tensor,

305 + block_size: int,

306 + slot_mapping: torch.Tensor,

307 + global_decode_buffer: torch.Tensor,

308 + decode_lens_buffer: torch.Tensor,

309 + prefill_buffer: torch.Tensor,

310 + max_compressed_tokens: int = 8192,

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

312 + """Single kernel for all C128A tokens (decode + prefill).

313 +

314 + Decode tokens: position → block_table lookup → global slot ids + topk_lens.

315 + Prefill tokens: position → local indices [0, ..., n-1, -1, ...].

316 +

317 + Writes into pre-allocated buffers for CUDA graph address stability.

318 + Returns slices of the buffers.

319 + """

320 + num_tokens = positions.shape[0]

321 + num_prefill_tokens = num_tokens - num_decode_tokens

322 +

323 + global_decode = global_decode_buffer[:num_decode_tokens]

324 + decode_lens = decode_lens_buffer[:num_decode_tokens]

325 + prefill_local = prefill_buffer[:num_prefill_tokens]

326 +

327 + if num_tokens == 0:

328 + return global_decode, decode_lens, prefill_local

329 +

330 + _build_c128a_topk_metadata_kernel[(num_tokens,)](

331 + global_decode_buffer,

332 + global_decode_buffer.stride(0),

333 + decode_lens_buffer,

334 + prefill_buffer,

335 + prefill_buffer.stride(0),

336 + positions,

337 + compress_ratio,

338 + max_compressed_tokens,

339 + num_decode_tokens,

340 + token_to_req_indices,

341 + block_table,

342 + block_table.stride(0),

343 + block_size,

344 + slot_mapping,

345 + BLOCK_SIZE=1024,

346 + )

347 + return global_decode, decode_lens, prefill_local

348 +

349 +

350 +@triton.jit

351 +def _build_c128a_topk_metadata_kernel(

352 + # Decode outputs

353 + global_decode_ptr,

354 + global_decode_stride,

355 + decode_lens_ptr,

356 + # Prefill output

357 + prefill_local_ptr,

358 + prefill_local_stride,

359 + # Inputs

360 + positions_ptr,

361 + compress_ratio,

362 + max_compressed_tokens,

363 + num_decode_tokens,

364 + token_to_req_indices_ptr,

365 + block_table_ptr,

366 + block_table_stride,

367 + block_size,

368 + slot_mapping_ptr,

369 + BLOCK_SIZE: tl.constexpr,

370 +):

371 + token_idx = tl.program_id(0)

372 + position = tl.load(positions_ptr + token_idx)

373 + num_compressed = (position + 1) // compress_ratio

374 + num_compressed = tl.minimum(num_compressed, max_compressed_tokens)

375 + is_decode = token_idx < num_decode_tokens

376 +

377 + if is_decode:

378 + # --- Decode: block-table lookup → global slot ids + count ---

379 + is_valid_token = tl.load(slot_mapping_ptr + token_idx) >= 0

380 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)

381 + count = tl.zeros((), dtype=tl.int32)

382 + for i in range(0, max_compressed_tokens, BLOCK_SIZE):

383 + offset = i + tl.arange(0, BLOCK_SIZE)

384 + mask = offset < max_compressed_tokens

385 + is_valid = offset < num_compressed

386 +

387 + block_indices = offset // block_size

388 + block_numbers = tl.load(

389 + block_table_ptr + req_idx * block_table_stride + block_indices,

390 + mask=mask & is_valid,

391 + )

392 + block_offsets = offset % block_size

393 + slot_ids = block_numbers * block_size + block_offsets

394 + slot_ids = tl.where(is_valid, slot_ids, -1)

395 + tl.store(

396 + global_decode_ptr + token_idx * global_decode_stride + offset,

397 + slot_ids,

398 + mask=mask,

399 + )

400 + count += tl.sum(is_valid.to(tl.int32), axis=0)

401 +

402 + tl.store(

403 + decode_lens_ptr + token_idx,

404 + tl.where(is_valid_token, count, 0),

405 + )

406 + else:

407 + # --- Prefill: write local indices ---

408 + pfx_idx = token_idx - num_decode_tokens

409 + for i in range(0, max_compressed_tokens, BLOCK_SIZE):

410 + offset = i + tl.arange(0, BLOCK_SIZE)

411 + mask = offset < max_compressed_tokens

412 + tl.store(

413 + prefill_local_ptr + pfx_idx * prefill_local_stride + offset,

414 + tl.where(offset < num_compressed, offset, -1),

415 + mask=mask,

416 + )

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

+3 -272

@@ -15,8 +15,6 @@ from vllm.model_executor.layers.attention.mla_attention import (

15 )

16 from vllm.platforms import current_platform

17 from vllm.platforms.interface import DeviceCapability

18 -from vllm.triton_utils import tl, triton

19 -from vllm.utils.math_utils import cdiv

18 from vllm.utils.platform_utils import num_compute_units

19 from vllm.utils.torch_utils import is_quantized_kv_cache

20 from vllm.v1.attention.backend import (

@@ -29,7 +27,6 @@ from vllm.v1.attention.backend import (

27 MultipleOf,

28 SparseMLAAttentionImpl,

29 )

32 -from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping

30 from vllm.v1.attention.backends.mla.sparse_utils import (

31 triton_convert_req_index_to_global_index,

32 )

@@ -118,9 +115,6 @@ class FlashMLASparseBackend(AttentionBackend):

115 @classmethod

116 def get_supported_head_sizes(cls) -> list[int]:

117 # DeepSeek V3.2 layout: 512 NoPE + 64 RoPE = 576.

121 - # DeepSeek V4 uses 448 NoPE + 64 RoPE = 512 and overrides this in

122 - # vllm/models/deepseek_v4/nvidia/flashmla.py:

123 - # DeepseekV4FlashMLASparseBackend.get_supported_head_sizes.

118 return [576]

119

120 @classmethod

@@ -223,13 +217,6 @@ class FlashMLASparseMetadata(AttentionMetadata):

217 fp8_extra_metadata: FP8SeparatePrefillDecode | FP8KernelMetadata | None = None

218 fp8_use_mixed_batch: bool = False

219

226 - # Pre-computed C128A metadata (DeepseekV4 only, compress_ratio == 128).

227 - # Decode: global slot ids + valid-entry counts (fused from positions).

228 - c128a_global_decode_topk_indices: torch.Tensor | None = None

229 - c128a_decode_topk_lens: torch.Tensor | None = None

230 - # Prefill: local topk indices (used by combine_topk_swa_indices).

231 - c128a_prefill_topk_indices: torch.Tensor | None = None

232 -

220

221 def get_prefill_workspace_size(max_model_len: int):

222 # NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size.

@@ -325,68 +312,6 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

312 device=device,

313 )

314

328 - # DeepseekV4: has compress_ratios in hf_config.

329 - hf_config = vllm_config.model_config.hf_config

330 - self.is_deepseek_v4 = (

331 - hasattr(hf_config, "compress_ratios") and len(hf_config.compress_ratios) > 0

332 - )

333 - self.compress_ratio = 1

334 - if self.is_deepseek_v4:

335 - assert hasattr(self.kv_cache_spec, "compress_ratio")

336 - self.compress_ratio = self.kv_cache_spec.compress_ratio

337 - # Pre-allocate compressed slot mapping buffer for CUDA graph

338 - # address stability when compress_ratio > 1.

339 - if self.compress_ratio > 1:

340 - max_num_batched_tokens = (

341 - vllm_config.scheduler_config.max_num_batched_tokens

342 - )

343 - self.compressed_slot_mapping_buffer = torch.empty(

344 - max_num_batched_tokens,

345 - dtype=torch.int64,

346 - device=self.device,

347 - )

348 -

349 - # Pre-allocate C128A topk buffers for CUDA graph address stability.

350 - if self.compress_ratio == 128:

351 - max_num_batched_tokens = (

352 - vllm_config.scheduler_config.max_num_batched_tokens

353 - )

354 - # Pad to B_TOPK alignment (128 covers both h_q=64 B_TOPK=64 and

355 - # h_q=128 B_TOPK=128). FlashMLA decode asserts extra_topk % B_TOPK

356 - # == 0; unaligned widths (e.g. 17 = ceil(2136/128)) crash the

357 - # sm100 head64 kernel. Padded slots stay -1 and decode_lens caps

358 - # them via topk_length, so the pad is a no-op at kernel level.

359 - # Mirrors _SPARSE_PREFILL_TOPK_ALIGNMENT in cache_utils.py.

360 - _C128A_TOPK_ALIGNMENT = 128

361 - c128a_max_compressed = cdiv(

362 - self.model_config.max_model_len, self.compress_ratio

363 - )

364 - c128a_max_compressed = (

365 - cdiv(c128a_max_compressed, _C128A_TOPK_ALIGNMENT)

366 - * _C128A_TOPK_ALIGNMENT

367 - )

368 - # Stored so _build_c128a_metadata passes it as the kernel's

369 - # max_compressed_tokens, matching the buffer stride. Otherwise

370 - # the kernel's default 8192 iterates past row width and spills

371 - # writes into adjacent rows (present in both decode and prefill

372 - # branches of _build_c128a_topk_metadata_kernel).

373 - self.c128a_max_compressed = c128a_max_compressed

374 - self.c128a_global_decode_buffer = torch.empty(

375 - (max_num_batched_tokens, c128a_max_compressed),

376 - dtype=torch.int32,

377 - device=self.device,

378 - )

379 - self.c128a_decode_lens_buffer = torch.empty(

380 - max_num_batched_tokens,

381 - dtype=torch.int32,

382 - device=self.device,

383 - )

384 - self.c128a_prefill_buffer = torch.empty(

385 - (max_num_batched_tokens, c128a_max_compressed),

386 - dtype=torch.int32,

387 - device=self.device,

388 - )

389 -

315 def _build_fp8_mixed_decode_prefill(

316 self,

317 common_attn_metadata: CommonAttentionMetadata,

@@ -582,109 +507,35 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

507 )

508 req_id_per_token = self.req_id_per_token_buffer[:num_tokens]

509

585 - slot_mapping = cm.slot_mapping

586 - if self.compress_ratio > 1:

587 - slot_mapping = get_compressed_slot_mapping(

588 - common_attn_metadata.num_actual_tokens,

589 - common_attn_metadata.query_start_loc,

590 - common_attn_metadata.seq_lens,

591 - common_attn_metadata.block_table_tensor.clamp(min=0),

592 - int(self.kv_cache_spec.storage_block_size),

593 - self.compress_ratio,

594 - out=self.compressed_slot_mapping_buffer,

595 - )

596 -

510 fp8_extra_metadata: (

511 FlashMLASparseMetadata.FP8SeparatePrefillDecode

512 | FlashMLASparseMetadata.FP8KernelMetadata

513 | None

514 ) = None

602 - fp8_use_mixed_batch = (

603 - self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.is_deepseek_v4

604 - )

605 - # DeepseekV4 has its own attention impl (DeepseekV4Attention) that does not

606 - # consume fp8_extra_metadata. Skipping the build here avoids a

607 - # forced D2H sync on seq_lens that would otherwise fire on every

608 - # prefill-bearing step, lifting GPU utilization on long-prefill

609 - # workloads (e.g. LongBench) from ~83% to ~100%.

610 - if self.use_fp8_kv_cache and not self.is_deepseek_v4:

515 + fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL

516 + if self.use_fp8_kv_cache:

517 if fp8_use_mixed_batch:

518 fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(cm)

519 else:

520 fp8_extra_metadata = self._build_fp8_separate_prefill_decode(cm)

521

616 - # Pre-compute C128A topk indices for DeepseekV4.

617 - c128a_fields = {}

618 - if self.is_deepseek_v4 and self.compress_ratio == 128:

619 - c128a_fields = self._build_c128a_metadata(cm, req_id_per_token)

620 -

522 metadata = FlashMLASparseMetadata(

523 num_reqs=cm.num_reqs,

524 max_query_len=cm.max_query_len,

525 max_seq_len=cm.max_seq_len,

526 num_actual_tokens=cm.num_actual_tokens,

527 query_start_loc=cm.query_start_loc,

627 - slot_mapping=slot_mapping,

528 + slot_mapping=cm.slot_mapping,

529 block_table=cm.block_table_tensor,

530 req_id_per_token=req_id_per_token,

531 block_size=self.kv_cache_spec.block_size,

532 topk_tokens=self.topk_tokens,

533 fp8_extra_metadata=fp8_extra_metadata,

534 fp8_use_mixed_batch=fp8_use_mixed_batch,

634 - **c128a_fields,

535 )

536

537 return metadata

538

639 - def _build_c128a_metadata(

640 - self,

641 - cm: CommonAttentionMetadata,

642 - req_id_per_token: torch.Tensor,

643 - ) -> dict[str, torch.Tensor | None]:

644 - """Pre-compute C128A topk indices for DeepseekV4 (compress_ratio >= 128)."""

645 - # Must match SWA's decode split (no `require_uniform=True`) so

646 - # `c128a_global_decode_topk_indices.shape[0]` lines up with q in

647 - # `_forward_decode`. The per-token C128A kernel handles non-uniform

648 - # query lengths.

649 - (num_decodes, _, num_decode_tokens, num_prefill_tokens) = (

650 - split_decodes_and_prefills(

651 - cm,

652 - decode_threshold=self.reorder_batch_threshold or 1,

653 - )

654 - )

655 -

656 - num_total = num_decode_tokens + num_prefill_tokens

657 - if num_total == 0:

658 - return {}

659 -

660 - assert cm.positions is not None, (

661 - "positions is required for C128A metadata build"

662 - )

663 - block_size = self.kv_cache_spec.block_size // self.compress_ratio

664 - global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(

665 - cm.positions[:num_total],

666 - self.compress_ratio,

667 - num_decode_tokens,

668 - req_id_per_token,

669 - cm.block_table_tensor[:num_decodes],

670 - block_size,

671 - cm.slot_mapping,

672 - self.c128a_global_decode_buffer,

673 - self.c128a_decode_lens_buffer,

674 - self.c128a_prefill_buffer,

675 - max_compressed_tokens=self.c128a_max_compressed,

676 - )

677 -

678 - result: dict[str, torch.Tensor | None] = {}

679 - if num_decode_tokens > 0:

680 - result["c128a_global_decode_topk_indices"] = global_decode.view(

681 - num_decode_tokens, 1, -1

682 - )

683 - result["c128a_decode_topk_lens"] = decode_lens

684 - if num_prefill_tokens > 0:

685 - result["c128a_prefill_topk_indices"] = prefill_local

686 - return result

687 -

539

540 class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):

541 @staticmethod

@@ -1027,123 +878,3 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):

878 )

879

880 return attn_out, None

1030 -

1031 -

1032 -def build_c128a_topk_metadata(

1033 - positions: torch.Tensor,

1034 - compress_ratio: int,

1035 - num_decode_tokens: int,

1036 - token_to_req_indices: torch.Tensor,

1037 - block_table: torch.Tensor,

1038 - block_size: int,

1039 - slot_mapping: torch.Tensor,

1040 - global_decode_buffer: torch.Tensor,

1041 - decode_lens_buffer: torch.Tensor,

1042 - prefill_buffer: torch.Tensor,

1043 - max_compressed_tokens: int = 8192,

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

1045 - """Single kernel for all C128A tokens (decode + prefill).

1046 -

1047 - Decode tokens: position → block_table lookup → global slot ids + topk_lens.

1048 - Prefill tokens: position → local indices [0, ..., n-1, -1, ...].

1049 -

1050 - Writes into pre-allocated buffers for CUDA graph address stability.

1051 - Returns slices of the buffers.

1052 - """

1053 - num_tokens = positions.shape[0]

1054 - num_prefill_tokens = num_tokens - num_decode_tokens

1055 -

1056 - global_decode = global_decode_buffer[:num_decode_tokens]

1057 - decode_lens = decode_lens_buffer[:num_decode_tokens]

1058 - prefill_local = prefill_buffer[:num_prefill_tokens]

1059 -

1060 - if num_tokens == 0:

1061 - return global_decode, decode_lens, prefill_local

1062 -

1063 - _build_c128a_topk_metadata_kernel[(num_tokens,)](

1064 - global_decode_buffer,

1065 - global_decode_buffer.stride(0),

1066 - decode_lens_buffer,

1067 - prefill_buffer,

1068 - prefill_buffer.stride(0),

1069 - positions,

1070 - compress_ratio,

1071 - max_compressed_tokens,

1072 - num_decode_tokens,

1073 - token_to_req_indices,

1074 - block_table,

1075 - block_table.stride(0),

1076 - block_size,

1077 - slot_mapping,

1078 - BLOCK_SIZE=1024,

1079 - )

1080 - return global_decode, decode_lens, prefill_local

1081 -

1082 -

1083 -@triton.jit

1084 -def _build_c128a_topk_metadata_kernel(

1085 - # Decode outputs

1086 - global_decode_ptr,

1087 - global_decode_stride,

1088 - decode_lens_ptr,

1089 - # Prefill output

1090 - prefill_local_ptr,

1091 - prefill_local_stride,

1092 - # Inputs

1093 - positions_ptr,

1094 - compress_ratio,

1095 - max_compressed_tokens,

1096 - num_decode_tokens,

1097 - token_to_req_indices_ptr,

1098 - block_table_ptr,

1099 - block_table_stride,

1100 - block_size,

1101 - slot_mapping_ptr,

1102 - BLOCK_SIZE: tl.constexpr,

1103 -):

1104 - token_idx = tl.program_id(0)

1105 - position = tl.load(positions_ptr + token_idx)

1106 - num_compressed = (position + 1) // compress_ratio

1107 - num_compressed = tl.minimum(num_compressed, max_compressed_tokens)

1108 - is_decode = token_idx < num_decode_tokens

1109 -

1110 - if is_decode:

1111 - # --- Decode: block-table lookup → global slot ids + count ---

1112 - is_valid_token = tl.load(slot_mapping_ptr + token_idx) >= 0

1113 - req_idx = tl.load(token_to_req_indices_ptr + token_idx)

1114 - count = tl.zeros((), dtype=tl.int32)

1115 - for i in range(0, max_compressed_tokens, BLOCK_SIZE):

1116 - offset = i + tl.arange(0, BLOCK_SIZE)

1117 - mask = offset < max_compressed_tokens

1118 - is_valid = offset < num_compressed

1119 -

1120 - block_indices = offset // block_size

1121 - block_numbers = tl.load(

1122 - block_table_ptr + req_idx * block_table_stride + block_indices,

1123 - mask=mask & is_valid,

1124 - )

1125 - block_offsets = offset % block_size

1126 - slot_ids = block_numbers * block_size + block_offsets

1127 - slot_ids = tl.where(is_valid, slot_ids, -1)

1128 - tl.store(

1129 - global_decode_ptr + token_idx * global_decode_stride + offset,

1130 - slot_ids,

1131 - mask=mask,

1132 - )

1133 - count += tl.sum(is_valid.to(tl.int32), axis=0)

1134 -

1135 - tl.store(

1136 - decode_lens_ptr + token_idx,

1137 - tl.where(is_valid_token, count, 0),

1138 - )

1139 - else:

1140 - # --- Prefill: write local indices ---

1141 - pfx_idx = token_idx - num_decode_tokens

1142 - for i in range(0, max_compressed_tokens, BLOCK_SIZE):

1143 - offset = i + tl.arange(0, BLOCK_SIZE)

1144 - mask = offset < max_compressed_tokens

1145 - tl.store(

1146 - prefill_local_ptr + pfx_idx * prefill_local_stride + offset,

1147 - tl.where(offset < num_compressed, offset, -1),

1148 - mask=mask,

1149 - )

vllm/v1/attention/backends/registry.py

+1 -1

@@ -78,7 +78,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):

78 )

79 # DeepSeek V4 sparse MLA backends (model-driven; selected via the V4 layer).

80 FLASHMLA_SPARSE_DSV4 = (

81 - "vllm.models.deepseek_v4.nvidia.flashmla.DeepseekV4FlashMLASparseBackend"

81 + "vllm.models.deepseek_v4.sparse_mla.DeepseekV4FlashMLABackend"

82 )

83 FLASHINFER_MLA_SPARSE_DSV4 = (

84 "vllm.models.deepseek_v4.nvidia.flashinfer_sparse."