@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e2efe79695

[ROCm]Migrating Deepseek V3.2 to vllm/models/deepseek_v32/ (#47207)

Signed-off-by: sroberts-amd <sroberts@amd.com> Signed-off-by: Stacy Roberts <sroberts@amd.com> Co-authored-by: Elvir Crncevic <elvircrn@gmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

sroberts-amd committed Jul 30, 2026 at 08:37 UTC e2efe79695e0f92395c029a6a7907c02a0374678
12 files changed +1031 -70
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+1 -1
@@ -27,7 +27,7 @@ rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses).
27 import pytest
28 import torch
29
30 -from vllm.models.deepseek_v32.nvidia import kernels as K
30 +from vllm.models.deepseek_v32.common import kernels as K
31 from vllm.platforms import current_platform
32
33 FP8 = torch.float8_e4m3fn
vllm/models/deepseek_v32/__init__.py
+9 -5
@@ -10,11 +10,15 @@ attend. The same model code serves any DSA checkpoint, including GLM-5.2
10
11 from vllm.platforms import current_platform
12
13 -if current_platform.is_rocm() or current_platform.is_xpu():
14 - raise NotImplementedError("deepseek_v32 currently supports NVIDIA SM100 only.")
15 -
16 -from .nvidia.model import DeepseekV32ForCausalLM
17 -from .nvidia.mtp import DeepseekV32MTP
13 +if current_platform.is_rocm():
14 + from .amd.model import DeepseekV32ForCausalLM
15 + from .amd.mtp import DeepseekV32MTP
16 +elif current_platform.is_xpu():
17 + raise NotImplementedError("deepseek_v32 does not yet support XPU.")
18 +else:
19 + # Covers Blackwell (sm100) and all other CUDA devices.
20 + from .nvidia.model import DeepseekV32ForCausalLM
21 + from .nvidia.mtp import DeepseekV32MTP
22
23 __all__ = [
24 "DeepseekV32ForCausalLM",
vllm/models/deepseek_v32/amd/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v32/amd/model.py new
+327
@@ -0,0 +1,327 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import typing
5 +from collections.abc import Callable, Iterable
6 +from itertools import islice
7 +
8 +import torch
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.distributed import get_pp_group
12 +from vllm.model_executor.layers.fused_moe import (
13 + fused_moe_make_expert_params_mapping,
14 +)
15 +from vllm.model_executor.layers.layernorm import RMSNorm
16 +from vllm.model_executor.layers.vocab_parallel_embedding import (
17 + VocabParallelEmbedding,
18 +)
19 +from vllm.model_executor.model_loader.weight_utils import (
20 + default_weight_loader,
21 + maybe_remap_kv_scale_name,
22 +)
23 +from vllm.model_executor.models.deepseek_v2 import (
24 + DeepseekV2ForCausalLM,
25 + DeepseekV2MLP,
26 + DeepseekV2MoE,
27 + _try_load_fp8_indexer_wk,
28 + get_spec_layer_idx_from_weight_name,
29 +)
30 +from vllm.model_executor.models.utils import (
31 + PPMissingLayer,
32 + get_pp_missing_layer_names,
33 + is_pp_missing_parameter,
34 + make_empty_intermediate_tensors_factory,
35 + make_layers,
36 +)
37 +from vllm.models.deepseek_v32.common.fused_ops import fused_allreduce_rms_norm
38 +from vllm.sequence import IntermediateTensors
39 +
40 +from .rocm import DeepseekV32MLAAttention
41 +
42 +
43 +class DeepseekV32DecoderLayer(torch.nn.Module):
44 + def __init__(
45 + self,
46 + vllm_config: VllmConfig,
47 + prefix: str,
48 + config=None,
49 + topk_indices_buffer: torch.Tensor | None = None,
50 + ) -> None:
51 + super().__init__()
52 +
53 + if config is None:
54 + config = vllm_config.model_config.hf_config
55 + quant_config = vllm_config.quant_config
56 + parallel_config = vllm_config.parallel_config
57 +
58 + self.hidden_size = config.hidden_size
59 + moe_layer_freq = getattr(config, "moe_layer_freq", 1)
60 + layer_idx = int(prefix.split(sep=".")[-1])
61 + self.layer_idx = layer_idx
62 + self.use_mha = False
63 +
64 + self.self_attn = DeepseekV32MLAAttention(
65 + vllm_config=vllm_config,
66 + config=config,
67 + prefix=f"{prefix}.self_attn",
68 + topk_indices_buffer=topk_indices_buffer,
69 + )
70 +
71 + if (
72 + config.n_routed_experts is not None
73 + and layer_idx >= config.first_k_dense_replace
74 + and layer_idx % moe_layer_freq == 0
75 + ):
76 + self.mlp = DeepseekV2MoE(
77 + config=config,
78 + parallel_config=parallel_config,
79 + quant_config=quant_config,
80 + prefix=f"{prefix}.mlp",
81 + )
82 + self.mlp.experts.moe_config.skip_final_all_reduce = True
83 + else:
84 + self.mlp = DeepseekV2MLP(
85 + hidden_size=config.hidden_size,
86 + intermediate_size=config.intermediate_size,
87 + hidden_act=config.hidden_act,
88 + quant_config=quant_config,
89 + prefix=f"{prefix}.mlp",
90 + reduce_results=False,
91 + )
92 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
93 + self.post_attention_layernorm = RMSNorm(
94 + config.hidden_size, eps=config.rms_norm_eps
95 + )
96 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
97 +
98 + def forward(
99 + self,
100 + positions: torch.Tensor,
101 + hidden_states: torch.Tensor,
102 + residual: torch.Tensor | None,
103 + ) -> tuple[torch.Tensor, torch.Tensor]:
104 + if residual is None:
105 + residual = hidden_states
106 + hidden_states = self.input_layernorm(hidden_states)
107 + else:
108 + hidden_states, residual = fused_allreduce_rms_norm(
109 + hidden_states, residual, self.input_layernorm
110 + )
111 + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
112 + hidden_states, residual = fused_allreduce_rms_norm(
113 + hidden_states, residual, self.post_attention_layernorm
114 + )
115 + hidden_states = self.mlp(hidden_states)
116 + return hidden_states, residual
117 +
118 +
119 +class DeepseekV32Model(torch.nn.Module):
120 + fall_back_to_pt_during_load = False
121 +
122 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
123 + super().__init__()
124 +
125 + config = vllm_config.model_config.hf_config
126 + quant_config = vllm_config.quant_config
127 + self.config = config
128 + from vllm.platforms import current_platform
129 +
130 + self.device = current_platform.device_type
131 +
132 + self.vocab_size = config.vocab_size
133 + self.is_v32 = True
134 + topk_indices_buffer = torch.empty(
135 + vllm_config.scheduler_config.max_num_batched_tokens,
136 + config.index_topk,
137 + dtype=torch.int32,
138 + device=self.device,
139 + )
140 +
141 + if get_pp_group().is_first_rank:
142 + self.embed_tokens = VocabParallelEmbedding(
143 + config.vocab_size,
144 + config.hidden_size,
145 + quant_config=quant_config,
146 + prefix=f"{prefix}.embed_tokens",
147 + )
148 + else:
149 + self.embed_tokens = PPMissingLayer()
150 +
151 + self.start_layer, self.end_layer, self.layers = make_layers(
152 + config.num_hidden_layers,
153 + lambda prefix: DeepseekV32DecoderLayer(
154 + vllm_config=vllm_config,
155 + prefix=prefix,
156 + topk_indices_buffer=topk_indices_buffer,
157 + ),
158 + prefix=f"{prefix}.layers",
159 + )
160 +
161 + if get_pp_group().is_last_rank:
162 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
163 + else:
164 + self.norm = PPMissingLayer()
165 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
166 + ["hidden_states", "residual"], config.hidden_size
167 + )
168 +
169 + self.aux_hidden_state_layers = tuple[int, ...]()
170 + self.num_redundant_experts = (
171 + vllm_config.parallel_config.eplb_config.num_redundant_experts
172 + )
173 +
174 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
175 + return self.embed_tokens(input_ids)
176 +
177 + def forward(
178 + self,
179 + input_ids: torch.Tensor | None,
180 + positions: torch.Tensor,
181 + intermediate_tensors: IntermediateTensors | None = None,
182 + inputs_embeds: torch.Tensor | None = None,
183 + ) -> torch.Tensor | IntermediateTensors:
184 + if get_pp_group().is_first_rank:
185 + if inputs_embeds is not None:
186 + hidden_states = inputs_embeds
187 + else:
188 + assert input_ids is not None
189 + hidden_states = self.embed_input_ids(input_ids)
190 + residual = None
191 + else:
192 + assert intermediate_tensors is not None
193 + hidden_states = intermediate_tensors["hidden_states"]
194 + residual = intermediate_tensors["residual"]
195 +
196 + aux_hidden_states = []
197 + for idx, layer in enumerate(
198 + islice(self.layers, self.start_layer, self.end_layer),
199 + start=self.start_layer,
200 + ):
201 + if idx in self.aux_hidden_state_layers:
202 + aux_hidden_states.append(hidden_states + residual)
203 + hidden_states, residual = layer(positions, hidden_states, residual)
204 +
205 + if not get_pp_group().is_last_rank:
206 + return IntermediateTensors(
207 + {"hidden_states": hidden_states, "residual": residual}
208 + )
209 +
210 + hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm)
211 + if len(aux_hidden_states) > 0:
212 + return hidden_states, aux_hidden_states
213 + return hidden_states
214 +
215 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
216 + stacked_params_mapping = [
217 + ("gate_up_proj", "gate_proj", 0),
218 + ("gate_up_proj", "up_proj", 1),
219 + ("fused_qkv_a_proj", "q_a_proj", 0),
220 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
221 + ("wk_weights_proj", "wk", 0),
222 + ("wk_weights_proj", "weights_proj", 1),
223 + ]
224 + expert_params_mapping = fused_moe_make_expert_params_mapping(
225 + self,
226 + ckpt_gate_proj_name="gate_proj",
227 + ckpt_down_proj_name="down_proj",
228 + ckpt_up_proj_name="up_proj",
229 + num_experts=self.config.n_routed_experts,
230 + num_redundant_experts=self.num_redundant_experts,
231 + )
232 +
233 + pp_missing_layer_names = get_pp_missing_layer_names(self)
234 + params_dict = dict(self.named_parameters())
235 + loaded_params: set[str] = set()
236 + _pending_wk_fp8: dict = {}
237 + for name, loaded_weight in weights:
238 + if "rotary_emb.inv_freq" in name:
239 + continue
240 + if get_spec_layer_idx_from_weight_name(self.config, name) is not None:
241 + continue
242 + if _try_load_fp8_indexer_wk(
243 + name,
244 + loaded_weight,
245 + _pending_wk_fp8,
246 + params_dict,
247 + loaded_params,
248 + pp_missing_layer_names,
249 + ):
250 + continue
251 +
252 + for param_name, weight_name, shard_id in stacked_params_mapping:
253 + if weight_name not in name:
254 + continue
255 + if ("mlp.experts." in name) and name not in params_dict:
256 + continue
257 + name_mapped = name.replace(weight_name, param_name)
258 + if (
259 + param_name == "fused_qkv_a_proj"
260 + ) and name_mapped not in params_dict:
261 + continue
262 + name = name_mapped
263 + if name.endswith(".bias") and name not in params_dict:
264 + continue
265 + if is_pp_missing_parameter(name, self):
266 + continue
267 + param = params_dict[name]
268 + param.weight_loader(param, loaded_weight, shard_id)
269 + break
270 + else:
271 + is_expert_weight = False
272 + for mapping in expert_params_mapping:
273 + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
274 + if weight_name not in name:
275 + continue
276 + is_expert_weight = True
277 + name_mapped = name.replace(weight_name, param_name)
278 + if is_pp_missing_parameter(name_mapped, self):
279 + continue
280 + param = params_dict[name_mapped]
281 + weight_loader = typing.cast(
282 + Callable[..., bool], param.weight_loader
283 + )
284 + success = weight_loader(
285 + param,
286 + loaded_weight,
287 + name_mapped,
288 + shard_id=shard_id,
289 + expert_id=expert_id,
290 + return_success=True,
291 + )
292 + if success:
293 + name = name_mapped
294 + break
295 + else:
296 + if is_expert_weight:
297 + continue
298 + if name.endswith(".bias") and name not in params_dict:
299 + continue
300 + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
301 + if name is None:
302 + continue
303 + if is_pp_missing_parameter(name, self):
304 + continue
305 + param = params_dict[name]
306 + loader = getattr(param, "weight_loader", default_weight_loader)
307 + loader(param, loaded_weight)
308 + loaded_params.add(name)
309 + return loaded_params
310 +
311 +
312 +class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
313 + model_cls = DeepseekV32Model
314 +
315 + def set_moe_parameters(self):
316 + self.num_expert_groups = getattr(self.config, "n_group", 1)
317 + self.moe_layers = []
318 + self.moe_mlp_layers = []
319 + example_moe = None
320 + for layer in self.model.layers:
321 + if isinstance(layer, PPMissingLayer):
322 + continue
323 + if isinstance(layer.mlp, DeepseekV2MoE):
324 + example_moe = layer.mlp
325 + self.moe_mlp_layers.append(layer.mlp)
326 + self.moe_layers.append(layer.mlp.experts)
327 + self.extract_moe_parameters(example_moe)
vllm/models/deepseek_v32/amd/mtp.py new
+392
@@ -0,0 +1,392 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import typing
6 +from collections.abc import Callable, Iterable
7 +
8 +import torch
9 +import torch.nn as nn
10 +
11 +from vllm._aiter_ops import rocm_aiter_ops
12 +from vllm.config import VllmConfig
13 +from vllm.distributed import tensor_model_parallel_all_reduce
14 +from vllm.model_executor.layers.fused_moe import (
15 + fused_moe_make_expert_params_mapping,
16 +)
17 +from vllm.model_executor.layers.layernorm import RMSNorm
18 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
19 +from vllm.model_executor.layers.vocab_parallel_embedding import (
20 + VocabParallelEmbedding,
21 +)
22 +from vllm.model_executor.model_loader.weight_utils import (
23 + default_weight_loader,
24 + maybe_remap_kv_scale_name,
25 +)
26 +from vllm.model_executor.models.deepseek_mtp import SharedHead
27 +from vllm.model_executor.models.deepseek_v2 import (
28 + DeepseekV2MixtureOfExperts,
29 + DeepseekV2MoE,
30 + _try_load_fp8_indexer_wk,
31 + get_spec_layer_idx_from_weight_name,
32 +)
33 +from vllm.model_executor.models.utils import (
34 + get_pp_missing_layer_names,
35 + maybe_prefix,
36 +)
37 +from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
38 +from vllm.platforms import current_platform
39 +from vllm.sequence import IntermediateTensors
40 +
41 +from .model import DeepseekV32DecoderLayer
42 +
43 +
44 +class DeepseekV32MultiTokenPredictorLayer(nn.Module):
45 + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
46 + super().__init__()
47 + assert vllm_config.speculative_config is not None
48 + config = vllm_config.speculative_config.draft_model_config.hf_config
49 + self.config = config
50 + quant_config = vllm_config.quant_config
51 +
52 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
53 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
54 + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
55 +
56 + topk_indices_buffer = torch.empty(
57 + vllm_config.scheduler_config.max_num_batched_tokens,
58 + config.index_topk,
59 + dtype=torch.int32,
60 + device=current_platform.device_type,
61 + )
62 + self.shared_head = SharedHead(
63 + config=config, prefix=prefix, quant_config=quant_config
64 + )
65 + self.mtp_block = DeepseekV32DecoderLayer(
66 + vllm_config,
67 + prefix,
68 + config=config,
69 + topk_indices_buffer=topk_indices_buffer,
70 + )
71 +
72 + def forward(
73 + self,
74 + input_ids: torch.Tensor,
75 + positions: torch.Tensor,
76 + previous_hidden_states: torch.Tensor,
77 + inputs_embeds: torch.Tensor | None = None,
78 + spec_step_index: int = 0,
79 + ) -> torch.Tensor:
80 + assert inputs_embeds is not None
81 + eh_input = fused_eh_norm(
82 + positions,
83 + inputs_embeds,
84 + previous_hidden_states,
85 + self.enorm.weight,
86 + self.hnorm.weight,
87 + self.enorm.variance_epsilon,
88 + )
89 + hidden_states = self.eh_proj(eh_input)
90 + hidden_states, residual = self.mtp_block(
91 + positions=positions, hidden_states=hidden_states, residual=None
92 + )
93 + hidden_states = tensor_model_parallel_all_reduce(hidden_states)
94 + hidden_states = residual + hidden_states
95 + hidden_states = self.shared_head.norm(hidden_states)
96 + return hidden_states, hidden_states
97 +
98 +
99 +class DeepseekV32MultiTokenPredictor(nn.Module):
100 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
101 + super().__init__()
102 + config = vllm_config.model_config.hf_config
103 + self.mtp_start_layer_idx = config.num_hidden_layers
104 + self.num_mtp_layers = config.num_nextn_predict_layers
105 + self.layers = torch.nn.ModuleDict(
106 + {
107 + str(idx): DeepseekV32MultiTokenPredictorLayer(
108 + vllm_config, f"{prefix}.layers.{idx}"
109 + )
110 + for idx in range(
111 + self.mtp_start_layer_idx,
112 + self.mtp_start_layer_idx + self.num_mtp_layers,
113 + )
114 + }
115 + )
116 + self.embed_tokens = VocabParallelEmbedding(
117 + config.vocab_size,
118 + config.hidden_size,
119 + prefix=maybe_prefix(prefix, "embed_tokens"),
120 + )
121 + self.logits_processor = LogitsProcessor(config.vocab_size)
122 +
123 + def set_skip_topk(self, skip: bool):
124 + for layer in self.layers.values():
125 + self_attn = getattr(layer.mtp_block, "self_attn", None)
126 + if self_attn is not None and hasattr(self_attn, "skip_topk"):
127 + self_attn.skip_topk = skip
128 +
129 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
130 + return self.embed_tokens(input_ids)
131 +
132 + def forward(
133 + self,
134 + input_ids: torch.Tensor,
135 + positions: torch.Tensor,
136 + previous_hidden_states: torch.Tensor,
137 + inputs_embeds: torch.Tensor | None = None,
138 + spec_step_idx: int = 0,
139 + ) -> torch.Tensor:
140 + if inputs_embeds is None:
141 + inputs_embeds = self.embed_tokens(input_ids)
142 + current_step_idx = spec_step_idx % self.num_mtp_layers
143 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
144 + input_ids,
145 + positions,
146 + previous_hidden_states,
147 + inputs_embeds,
148 + current_step_idx,
149 + )
150 +
151 + def compute_logits(
152 + self,
153 + hidden_states: torch.Tensor,
154 + spec_step_idx: int = 0,
155 + ) -> torch.Tensor:
156 + current_step_idx = spec_step_idx % self.num_mtp_layers
157 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
158 + return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
159 +
160 +
161 +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
162 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
163 + super().__init__()
164 + self.config = vllm_config.model_config.hf_config
165 + self.quant_config = vllm_config.quant_config
166 + self.model = DeepseekV32MultiTokenPredictor(
167 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
168 + )
169 + self.set_moe_parameters()
170 +
171 + def set_moe_parameters(self):
172 + self.num_moe_layers = self.config.num_nextn_predict_layers
173 + self.num_expert_groups = self.config.n_group
174 + self.moe_layers = []
175 + self.moe_mlp_layers = []
176 + example_moe = None
177 + for layer in self.model.layers.values():
178 + mlp = layer.mtp_block.mlp
179 + if isinstance(mlp, DeepseekV2MoE):
180 + example_moe = mlp
181 + self.moe_mlp_layers.append(mlp)
182 + self.moe_layers.append(mlp.experts)
183 + self.extract_moe_parameters(example_moe)
184 +
185 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
186 + return self.model.embed_input_ids(input_ids)
187 +
188 + def forward(
189 + self,
190 + input_ids: torch.Tensor | None,
191 + positions: torch.Tensor,
192 + hidden_states: torch.Tensor,
193 + intermediate_tensors: IntermediateTensors | None = None,
194 + inputs_embeds: torch.Tensor | None = None,
195 + spec_step_idx: int = 0,
196 + ) -> torch.Tensor:
197 + return self.model(
198 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
199 + )
200 +
201 + def compute_logits(
202 + self,
203 + hidden_states: torch.Tensor,
204 + spec_step_idx: int = 0,
205 + ) -> torch.Tensor | None:
206 + return self.model.compute_logits(hidden_states, spec_step_idx)
207 +
208 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
209 + spec_layer_weight_names = [
210 + "embed_tokens",
211 + "enorm",
212 + "hnorm",
213 + "eh_proj",
214 + "shared_head",
215 + ]
216 + shared_weight_names = ["embed_tokens"]
217 + spec_layer_weight = False
218 + shared_weight = False
219 + for weight_name in spec_layer_weight_names:
220 + if weight_name in name:
221 + spec_layer_weight = True
222 + if weight_name in shared_weight_names:
223 + shared_weight = True
224 + break
225 + if not spec_layer_weight:
226 + name = name.replace(
227 + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
228 + )
229 + elif shared_weight:
230 + name = name.replace(f"model.layers.{spec_layer}.", "model.")
231 + return name
232 +
233 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
234 + rocm_aiter_moe_shared_expert_enabled = (
235 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
236 + )
237 + stacked_params_mapping = [
238 + ("gate_up_proj", "gate_proj", 0),
239 + ("gate_up_proj", "up_proj", 1),
240 + ("fused_qkv_a_proj", "q_a_proj", 0),
241 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
242 + ("wk_weights_proj", "wk", 0),
243 + ("wk_weights_proj", "weights_proj", 1),
244 + ]
245 + expert_params_mapping = fused_moe_make_expert_params_mapping(
246 + self,
247 + ckpt_gate_proj_name="gate_proj",
248 + ckpt_down_proj_name="down_proj",
249 + ckpt_up_proj_name="up_proj",
250 + num_experts=self.config.n_routed_experts
251 + + (
252 + self.config.n_shared_experts
253 + if rocm_aiter_moe_shared_expert_enabled
254 + else 0
255 + ),
256 + )
257 +
258 + pp_missing_layer_names = get_pp_missing_layer_names(self)
259 + params_dict = dict(self.named_parameters())
260 + loaded_params: set[str] = set()
261 + _pending_wk_fp8: dict = {}
262 + for name, loaded_weight in weights:
263 + if "rotary_emb.inv_freq" in name:
264 + continue
265 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
266 + if spec_layer is None:
267 + continue
268 + is_fusion_moe_shared_experts_layer = (
269 + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
270 + )
271 + name = self._rewrite_spec_layer_name(spec_layer, name)
272 +
273 + if _try_load_fp8_indexer_wk(
274 + name,
275 + loaded_weight,
276 + _pending_wk_fp8,
277 + params_dict,
278 + loaded_params,
279 + pp_missing_layer_names,
280 + ):
281 + continue
282 +
283 + for param_name, weight_name, shard_id in stacked_params_mapping:
284 + if weight_name not in name:
285 + continue
286 + if ("mlp.experts." in name) and name not in params_dict:
287 + continue
288 + if is_fusion_moe_shared_experts_layer:
289 + continue
290 + name_mapped = name.replace(weight_name, param_name)
291 + if (
292 + param_name == "fused_qkv_a_proj"
293 + ) and name_mapped not in params_dict:
294 + continue
295 + else:
296 + name = name_mapped
297 + if name.endswith(".bias") and name not in params_dict:
298 + continue
299 + param = params_dict[name]
300 + weight_loader = param.weight_loader
301 + weight_loader(param, loaded_weight, shard_id)
302 + break
303 + else:
304 + num_chunks = 1
305 + if is_fusion_moe_shared_experts_layer:
306 + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
307 + split_dim = (
308 + 1
309 + if ("down_proj.weight" in name and loaded_weight.ndim > 1)
310 + else 0
311 + )
312 + total = loaded_weight.shape[split_dim]
313 + assert total % num_chunks == 0
314 + chunk_size = total // num_chunks
315 +
316 + for j in range(num_chunks):
317 + chunk_name = name
318 + weight_to_load = loaded_weight
319 + if is_fusion_moe_shared_experts_layer:
320 + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
321 + if loaded_weight.ndim == 1:
322 + weight_to_load = loaded_weight[chunk_slice]
323 + elif split_dim == 0:
324 + weight_to_load = loaded_weight[chunk_slice, :]
325 + else:
326 + weight_to_load = loaded_weight[:, chunk_slice]
327 + chunk_name = name.replace(
328 + "mlp.shared_experts",
329 + f"mlp.experts.{self.config.n_routed_experts + j}",
330 + )
331 +
332 + is_expert_weight = False
333 + for mapping in expert_params_mapping:
334 + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
335 + if weight_name not in chunk_name:
336 + continue
337 + is_expert_weight = True
338 + name_mapped = chunk_name.replace(weight_name, param_name)
339 + param = params_dict[name_mapped]
340 + weight_loader = typing.cast(
341 + Callable[..., bool], param.weight_loader
342 + )
343 + success = weight_loader(
344 + param,
345 + weight_to_load,
346 + name_mapped,
347 + shard_id=shard_id,
348 + expert_id=expert_id,
349 + return_success=True,
350 + )
351 + if success:
352 + if not is_fusion_moe_shared_experts_layer:
353 + name = name_mapped
354 + else:
355 + loaded_params.add(name_mapped)
356 + break
357 + else:
358 + if is_expert_weight:
359 + continue
360 + if name.endswith(".bias") and name not in params_dict:
361 + continue
362 + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
363 + if name is None:
364 + continue
365 + if (
366 + spec_layer != self.model.mtp_start_layer_idx
367 + and ".layers" not in name
368 + ):
369 + continue
370 + param = params_dict[name]
371 + weight_loader = getattr(
372 + param, "weight_loader", default_weight_loader
373 + )
374 + weight_loader(param, loaded_weight)
375 + if not is_fusion_moe_shared_experts_layer:
376 + loaded_params.add(name)
377 +
378 + loaded_layers: set[int] = set()
379 + for param_name in loaded_params:
380 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
381 + if spec_layer is not None:
382 + loaded_layers.add(spec_layer)
383 + for layer_idx in range(
384 + self.model.mtp_start_layer_idx,
385 + self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
386 + ):
387 + if layer_idx not in loaded_layers:
388 + raise ValueError(
389 + f"MTP speculative decoding layer {layer_idx} weights "
390 + f"missing from checkpoint."
391 + )
392 + return loaded_params
vllm/models/deepseek_v32/amd/rocm.py new
+263
@@ -0,0 +1,263 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
7 +from vllm.forward_context import get_forward_context
8 +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
9 +from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
10 +from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer
11 +from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
12 +from vllm.utils.torch_utils import is_quantized_kv_cache
13 +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
14 +from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
15 + ROCMAiterMLASparseBackend,
16 +)
17 +
18 +
19 +class DeepseekV32MLASparseBackend(ROCMAiterMLASparseBackend):
20 + @staticmethod
21 + def get_supported_kernel_block_sizes() -> list:
22 + return [16, 32]
23 +
24 +
25 +class DeepseekV32ROCmIndexerBackend(DeepseekV32IndexerBackend):
26 + @staticmethod
27 + def get_supported_kernel_block_sizes() -> list:
28 + return [16, 32]
29 +
30 +
31 +class DeepseekV32ROCmIndexerCache(DeepseekV32IndexerCache):
32 + def get_attn_backend(self):
33 + return DeepseekV32ROCmIndexerBackend
34 +
35 +
36 +class DeepseekV32ROCmIndexer(DeepseekV32Indexer):
37 + indexer_cache_cls = DeepseekV32ROCmIndexerCache
38 +
39 +
40 +class DeepseekV32MLAAttention(DeepseekV32Attention):
41 + require_fp8_kv_cache: bool = False
42 + indexer_cls = DeepseekV32ROCmIndexer
43 +
44 + def __init__(self, vllm_config, config, prefix, topk_indices_buffer=None):
45 + super().__init__(
46 + vllm_config,
47 + config,
48 + prefix,
49 + topk_indices_buffer,
50 + attn_backend=DeepseekV32MLASparseBackend,
51 + )
52 +
53 + self.indexer_op: SparseAttnIndexer | None = None
54 + if self.indexer is not None:
55 + self.indexer_op = SparseAttnIndexer(
56 + self.indexer.k_cache,
57 + self.indexer.quant_block_size,
58 + self.indexer.scale_fmt,
59 + self.indexer.topk_tokens,
60 + self.indexer.head_dim,
61 + self.indexer.max_model_len,
62 + self.indexer.max_total_seq_len,
63 + topk_indices_buffer,
64 + skip_k_cache_insert=True,
65 + )
66 + self._fp8_kv = is_quantized_kv_cache(self.kv_cache_dtype)
67 + self._fp8_kv_needs_view = self._fp8_kv and self.kv_cache_dtype != "fp8_ds_mla"
68 +
69 + def _compute_ql_nope(self, q_c: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
70 + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
71 + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
72 + q_nope = q_nope.transpose(0, 1) # (N, tokens, P)
73 +
74 + if self.is_aiter_triton_fp4_bmm_enabled:
75 + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
76 +
77 + ql_nope = batched_gemm_a16wfp4(
78 + q_nope, self.W_K, self.W_K_scale, transpose_bm=True, prequant=True
79 + )
80 + elif self.is_aiter_triton_fp8_bmm_enabled:
81 + from vllm._aiter_ops import rocm_aiter_ops
82 +
83 + ql_nope = rocm_aiter_ops.triton_fp8_bmm(
84 + q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
85 + )
86 + else:
87 + ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1)
88 +
89 + return ql_nope, q_pe
90 +
91 + def _run_indexer(
92 + self,
93 + q_c: torch.Tensor,
94 + index_q_fp8: torch.Tensor | None,
95 + index_weights_out: torch.Tensor | None,
96 + ) -> None:
97 + """Run the ROCm sparse indexer (forward_hip) if this layer has an indexer."""
98 + if self.indexer_op is not None:
99 + self.indexer_op.forward_hip(q_c, index_q_fp8, None, index_weights_out)
100 +
101 + def _build_q_for_attn(
102 + self,
103 + ql_nope: torch.Tensor,
104 + mqa_q: torch.Tensor,
105 + num_actual: int,
106 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
107 + if self._fp8_kv:
108 + # fp8 KV: mqa_q is the full [ql_nope; q_pe] packed as fp8.
109 + return mqa_q[:num_actual]
110 +
111 + return (ql_nope[:num_actual], mqa_q[:num_actual])
112 +
113 + def _compute_uv_out(
114 + self,
115 + attn_out: torch.Tensor,
116 + output: torch.Tensor,
117 + num_actual: int,
118 + ) -> None:
119 + x = attn_out.view(
120 + num_actual, self.num_local_heads, self.kv_lora_rank
121 + ).transpose(0, 1) # (N, tokens, L)
122 + out_view = output[:num_actual].view(
123 + num_actual, self.num_local_heads, self.v_head_dim
124 + )
125 +
126 + if self.is_aiter_triton_fp4_bmm_enabled:
127 + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
128 +
129 + batched_gemm_a16wfp4(
130 + x, self.W_V, self.W_V_scale, out_view, transpose_bm=True, prequant=True
131 + )
132 + elif self.is_aiter_triton_fp8_bmm_enabled:
133 + from vllm._aiter_ops import rocm_aiter_ops
134 +
135 + rocm_aiter_ops.triton_fp8_bmm(
136 + x,
137 + self.W_V,
138 + self.W_V_scale,
139 + group_size=128,
140 + transpose_bm=True,
141 + YQ=out_view,
142 + )
143 + else:
144 + torch.bmm(x, self.W_UV, out=out_view.transpose(0, 1))
145 +
146 + @eager_break_during_capture
147 + def _fused_attention(
148 + self,
149 + positions: torch.Tensor,
150 + q_c: torch.Tensor,
151 + kv_c: torch.Tensor,
152 + k_pe: torch.Tensor,
153 + index_k: torch.Tensor | None,
154 + index_weights: torch.Tensor | None,
155 + output: torch.Tensor,
156 + ) -> None:
157 + forward_context = get_forward_context()
158 + attn_metadata_raw = forward_context.attn_metadata
159 + if isinstance(attn_metadata_raw, dict):
160 + attn_metadata = attn_metadata_raw.get(self.layer_name)
161 + elif isinstance(attn_metadata_raw, list):
162 + attn_metadata = attn_metadata_raw[0].get(self.layer_name)
163 + else:
164 + attn_metadata = attn_metadata_raw
165 +
166 + slot_mapping = forward_context.slot_mapping
167 + assert isinstance(slot_mapping, dict)
168 + mla_slot = slot_mapping.get(self.layer_name)
169 +
170 + if self.indexer is not None:
171 + has_indexer = True
172 + indexer_k_norm_w = self.indexer.k_norm.weight
173 + indexer_k_norm_bias = self.indexer.k_norm.bias
174 + indexer_k_norm_eps = self.indexer.k_norm.eps
175 + indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
176 + indexer_k_cache = self.indexer.k_cache.kv_cache
177 + indexer_softmax_scale = self.indexer.softmax_scale
178 + indexer_n_head_scale = self.indexer.n_head**-0.5
179 + else:
180 + has_indexer = False
181 + indexer_k_norm_w = None
182 + indexer_k_norm_bias = None
183 + indexer_k_norm_eps = 1e-6
184 + indexer_k_rope_cos_sin_cache = None
185 + indexer_k_cache = None
186 + indexer_softmax_scale = 0.0
187 + indexer_n_head_scale = 0.0
188 +
189 + if attn_metadata is None:
190 + mla_kv_cache = None
191 + mla_k_scale = None
192 + indexer_k_cache = None
193 + mla_slot = None
194 + else:
195 + mla_kv_cache = self.kv_cache
196 + mla_k_scale = self._k_scale
197 +
198 + q_c = fused_norm_rope(
199 + positions,
200 + q_c,
201 + self.q_a_layernorm.weight,
202 + self.q_a_layernorm.variance_epsilon,
203 + kv_c,
204 + self.kv_a_layernorm.weight,
205 + self.kv_a_layernorm.variance_epsilon,
206 + k_pe,
207 + self.rotary_emb.cos_sin_cache,
208 + index_k,
209 + indexer_k_norm_w,
210 + indexer_k_norm_bias,
211 + indexer_k_norm_eps,
212 + indexer_k_rope_cos_sin_cache,
213 + self.topk_indices_buffer,
214 + slot_mapping=mla_slot,
215 + indexer_k_cache=indexer_k_cache,
216 + mla_kv_cache=mla_kv_cache,
217 + mla_kv_cache_dtype=self.kv_cache_dtype,
218 + mla_k_scale=mla_k_scale,
219 + has_indexer=has_indexer,
220 + index_rope_interleave=self._index_rope_interleave,
221 + )
222 +
223 + ql_nope, q_pe = self._compute_ql_nope(q_c)
224 +
225 + if self.indexer is not None:
226 + index_q = self.indexer.wq_b(q_c)[0]
227 + index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
228 + else:
229 + index_q = None
230 +
231 + index_q_fp8, index_weights_out, mqa_q = fused_q(
232 + positions,
233 + q_pe,
234 + self.rotary_emb.cos_sin_cache,
235 + index_q,
236 + self.indexer_rope_emb.cos_sin_cache if has_indexer else None,
237 + ql_nope,
238 + self._q_scale,
239 + index_weights,
240 + indexer_softmax_scale,
241 + indexer_n_head_scale,
242 + has_indexer=has_indexer,
243 + index_rope_interleave=self._index_rope_interleave,
244 + quantize_mqa=self._fp8_kv,
245 + )
246 +
247 + self._run_indexer(q_c, index_q_fp8, index_weights_out)
248 +
249 + if attn_metadata is None:
250 + output.zero_()
251 + return
252 +
253 + num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined]
254 + kv_cache = self.kv_cache
255 + if self._fp8_kv_needs_view:
256 + kv_cache = kv_cache.view(torch.float8_e4m3fn)
257 +
258 + q_for_attn = self._build_q_for_attn(ql_nope, mqa_q, num_actual)
259 + attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined]
260 + q_for_attn, kv_cache, attn_metadata, self
261 + )
262 +
263 + self._compute_uv_out(attn_out, output, num_actual)
vllm/models/deepseek_v32/attention.py renamed
+32 -60
@@ -31,12 +31,13 @@ from vllm.model_executor.models.deepseek_v2 import (
31 yarn_get_mscale,
32 )
33 from vllm.model_executor.models.utils import extract_layer_index
34 +from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
35 from vllm.utils.torch_utils import is_quantized_kv_cache
36
36 -from .kernels import fused_norm_rope, fused_q
37 -
37
38 class DeepseekV32Indexer(nn.Module):
39 + indexer_cache_cls = DeepseekV32IndexerCache
40 +
41 def __init__(
42 self,
43 vllm_config: VllmConfig,
@@ -55,7 +56,6 @@ class DeepseekV32Indexer(nn.Module):
56 self.rope_dim = config.qk_rope_head_dim
57 self.q_lora_rank = q_lora_rank
58
58 - # No tensor parallel, just replicated.
59 self.wq_b = ReplicatedLinear(
60 self.q_lora_rank,
61 self.head_dim * self.n_head,
@@ -63,8 +63,7 @@ class DeepseekV32Indexer(nn.Module):
63 quant_config=quant_config,
64 prefix=f"{prefix}.wq_b",
65 )
66 - # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
67 - # FP8 wk weights are upcasted to BF16 during loading to keep this fused.
66 +
67 self.wk_weights_proj = MergedColumnParallelLinear(
68 hidden_size,
69 [self.head_dim, self.n_head],
@@ -80,9 +79,8 @@ class DeepseekV32Indexer(nn.Module):
79 self.quant_block_size = 128
80 self.topk_indices_buffer = topk_indices_buffer
81
83 - # fp8 naive cache: value in fp8 + fp32 scale per quant_block_size element.
82 assert cache_config is not None, "DeepSeek V3.2 indexer requires cache_config"
85 - self.k_cache = DeepseekV32IndexerCache(
83 + self.k_cache = type(self).indexer_cache_cls(
84 head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
85 dtype=torch.uint8,
86 prefix=f"{prefix}.k_cache",
@@ -120,7 +118,7 @@ class DeepseekV32Indexer(nn.Module):
118 q_pe, q_nope = torch.split(
119 q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
120 )
123 - # Fused wk + weights_proj: one GEMM, then split.
121 +
122 kw, _ = self.wk_weights_proj(hidden_states)
123 k = kw[:, : self.head_dim]
124 weights = kw[:, self.head_dim :]
@@ -131,14 +129,13 @@ class DeepseekV32Indexer(nn.Module):
129 )
130
131 q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
134 - # RoPE (NeoX) can introduce extra leading dims; reshape back to flat.
132 +
133 q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
134 k_pe = k_pe.reshape(-1, 1, self.rope_dim)
135
136 q = torch.cat([q_pe, q_nope], dim=-1)
137 k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
138
141 - # Only quant q here; k quant is fused with cache insertion.
139 q = q.view(-1, self.head_dim)
140 q_fp8, q_scale = per_token_group_quant_fp8(
141 q,
@@ -158,9 +155,10 @@ class DeepseekV32Indexer(nn.Module):
155
156
157 class DeepseekV32Attention(MLAAttention):
161 - # Narrow the base's broadly-typed `indexer` to the concrete type so the
162 - # `if self.indexer is not None` guards below type-check its attributes.
158 indexer: "DeepseekV32Indexer | None"
159 + indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer
160 +
161 + require_fp8_kv_cache: bool = True
162
163 def __init__(
164 self,
@@ -168,6 +166,7 @@ class DeepseekV32Attention(MLAAttention):
166 config: DeepseekV2Config | DeepseekV3Config,
167 prefix: str,
168 topk_indices_buffer: torch.Tensor | None = None,
169 + attn_backend: "type | None" = None,
170 ) -> None:
171 quant_config = vllm_config.quant_config
172 cache_config = vllm_config.cache_config
@@ -200,11 +199,6 @@ class DeepseekV32Attention(MLAAttention):
199 mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
200 scaling = scaling * mscale * mscale
201
203 - # DSA "shared indexer" pattern: only some layers carry an indexer; the
204 - # rest reuse the top-k written by the previous indexer layer into the
205 - # shared topk_indices_buffer. DeepSeek-V3.2 builds it on every layer
206 - # (index_topk_freq defaults to 1); GLM-5.2 uses index_topk_freq=4 so
207 - # only layers [0,1,2,6,10,...] (+ MTP) carry one.
202 layer_id = extract_layer_index(prefix)
203 index_topk_freq = getattr(config, "index_topk_freq", 1)
204 index_topk_pattern = getattr(config, "index_topk_pattern", None)
@@ -217,12 +211,10 @@ class DeepseekV32Attention(MLAAttention):
211 skip_topk = index_topk_pattern[layer_id] == "S"
212 else:
213 skip_topk = False
220 - # MTP/nextn layers always build a full indexer (they toggle at runtime).
214 +
215 num_hidden_layers = getattr(config, "num_hidden_layers", None)
216 is_mtp_layer = num_hidden_layers is not None and layer_id >= num_hidden_layers
217
224 - # Build kv_b_proj + indexer first; they are passed to MLAAttention.__init__
225 - # (which runs nn.Module.__init__ and registers them).
218 kv_b_proj = ColumnParallelLinear(
219 kv_lora_rank,
220 num_heads * (qk_nope_head_dim + v_head_dim),
@@ -232,7 +224,7 @@ class DeepseekV32Attention(MLAAttention):
224 )
225 indexer = None
226 if not skip_topk or is_mtp_layer:
235 - indexer = DeepseekV32Indexer(
227 + indexer = type(self).indexer_cls(
228 vllm_config,
229 config,
230 hidden_size,
@@ -243,8 +235,6 @@ class DeepseekV32Attention(MLAAttention):
235 prefix=f"{prefix}.indexer",
236 )
237
246 - # Set up the MLA engine (impl, KV cache, scales, backend, registration,
247 - # and process_weights_after_loading) via the MLAAttention base.
238 super().__init__(
239 num_heads=num_local_heads,
240 scale=scaling,
@@ -260,41 +250,32 @@ class DeepseekV32Attention(MLAAttention):
250 use_sparse=True,
251 indexer=indexer,
252 topk_indices_buffer=topk_indices_buffer,
253 + attn_backend=attn_backend,
254 )
255
256 self.num_local_heads = num_local_heads
257 self.qk_head_dim = qk_head_dim
258 self.indexer = indexer
259 self.topk_indices_buffer = topk_indices_buffer
269 - # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0
270 - # computes the top-k, steps 1+ set this True to reuse it.
260 +
261 self.skip_topk = False
272 - # Fused fp8 paths: Triton fused norm/rope/cache + fused-q. Two layouts,
273 - # picked by the sparse MLA backend's query support:
274 - # * supports_quant_query_input (FlashInfer sparse, SM100): per-tensor
275 - # fp8 cache + a single packed fp8 MQA query.
276 - # * not supported (FlashMLA sparse, SM90/SM100): fp8_ds_mla cache
277 - # (per-128 block-scaled fp8 NoPE + unquantized bf16 RoPE) + a bf16
278 - # (ql_nope, q_pe) query tuple. FA3 cannot mix a bf16 query with an
279 - # fp8 KV cache, so FlashMLA (which dequantizes internally) is used.
280 - # FlashMLA sparse runs on both Hopper and Blackwell, so this is the
281 - # only DSA path on SM90 and an opt-in alternative on SM100.
282 - assert is_quantized_kv_cache(self.kv_cache_dtype), (
283 - "deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse "
284 - "MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) "
285 - "or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)."
286 - )
287 - self._fp8_query = self.impl.supports_quant_query_input
288 - if not self._fp8_query:
289 - assert self.kv_cache_dtype == "fp8_ds_mla", (
290 - "deepseek_v32 (nvidia) on a bf16-query sparse MLA backend "
291 - "(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. "
292 - "Launch with --kv-cache-dtype fp8_ds_mla."
262 +
263 + if self.require_fp8_kv_cache:
264 + assert is_quantized_kv_cache(self.kv_cache_dtype), (
265 + "deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse "
266 + "MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) "
267 + "or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)."
268 )
294 - # The paged KV cache is stored as uint8 and viewed as fp8 for the decode
295 - # (per-tensor fp8). The fp8_ds_mla layout is consumed as raw bytes.
296 - self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla"
297 - # GLM-5.2 uses interleaved indexer RoPE; DeepSeek-V3.2 uses NeoX.
269 + self._fp8_query = self.impl.supports_quant_query_input
270 + if not self._fp8_query:
271 + assert self.kv_cache_dtype == "fp8_ds_mla", (
272 + "deepseek_v32 (nvidia) on a bf16-query sparse MLA backend "
273 + "(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. "
274 + "Launch with --kv-cache-dtype fp8_ds_mla."
275 + )
276 +
277 + self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla"
278 +
279 self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False)
280
281 # Remaining MLA projections (registered on this module).
@@ -313,9 +294,7 @@ class DeepseekV32Attention(MLAAttention):
294 prefix=f"{prefix}.q_b_proj",
295 )
296 self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps)
316 - # reduce_results=False: the attention all-reduce is fused with the
317 - # following post_attention_layernorm in the decoder layer via
318 - # fused_allreduce_rms_norm.
297 +
298 self.o_proj = RowParallelLinear(
299 num_heads * v_head_dim,
300 hidden_size,
@@ -380,11 +359,6 @@ class DeepseekV32Attention(MLAAttention):
359 index_weights: torch.Tensor | None,
360 output: torch.Tensor,
361 ) -> None:
383 - # One eager break for the whole attention. In FULL cudagraph mode (pure
384 - # decode) this decorator is a no-op, so everything here is captured; in
385 - # PIECEWISE (prefill) it runs eagerly. The cache writes, sparse indexer,
386 - # and forward_mqa all depend on per-step metadata and must not be split
387 - # out (PIECEWISE capture would otherwise miss them).
362 forward_context = get_forward_context()
363 attn_metadata_raw = forward_context.attn_metadata
364 if isinstance(attn_metadata_raw, dict):
@@ -516,8 +490,6 @@ class DeepseekV32Attention(MLAAttention):
490 :num_actual
491 ]
492 else:
519 - # FlashMLA sparse: bf16 (ql_nope, q_pe) tuple. mqa_q is the RoPE'd
520 - # q_pe; ql_nope is consumed directly.
493 mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual])
494 attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined]
495 mqa_q_arg, kv_cache, attn_metadata, self
vllm/models/deepseek_v32/common/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v32/common/fused_ops.py renamed
vllm/models/deepseek_v32/common/kernels.py renamed
vllm/models/deepseek_v32/nvidia/model.py
+2 -3
@@ -38,11 +38,10 @@ from vllm.model_executor.models.utils import (
38 make_layers,
39 sequence_parallel_chunk,
40 )
41 +from vllm.models.deepseek_v32.attention import DeepseekV32Attention
42 +from vllm.models.deepseek_v32.common.fused_ops import fused_allreduce_rms_norm
43 from vllm.sequence import IntermediateTensors
44
43 -from .attention import DeepseekV32Attention
44 -from .fused_ops import fused_allreduce_rms_norm
45 -
45
46 def _all_gather_sp_states(
47 hidden_states: torch.Tensor,
vllm/models/deepseek_v32/nvidia/mtp.py
+1 -1
@@ -38,10 +38,10 @@ from vllm.model_executor.models.utils import (
38 get_pp_missing_layer_names,
39 maybe_prefix,
40 )
41 +from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
42 from vllm.platforms import current_platform
43 from vllm.sequence import IntermediateTensors
44
44 -from .kernels import fused_eh_norm
45 from .model import DeepseekV32DecoderLayer
46
47