@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1806d1adfc
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 2, 22:00 UTC

[ROCm] [DSv4] [Perf] Support DeepSeek v4 MTP (#43385)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

TJian committed May 24, 2026 at 18:43 UTC 1806d1adfc9b598bc6eb94de38a330aaad04c291
8 files changed +2340 -52
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+22 -11
@@ -81,17 +81,28 @@ def _patch_make_bitmatrix_metadata() -> None:
81 import triton.language as tl
82
83 try:
84 - from vllm.third_party.triton_kernels.tensor_details import (
85 - bitmatrix as _bm,
86 - )
87 - from vllm.third_party.triton_kernels.tensor_details.bitmatrix import (
88 - BitmatrixMetadata,
89 - _keyed_add,
90 - cdiv,
91 - )
92 - from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501
93 - sum_bitmatrix_rows,
94 - )
84 + if current_platform.is_rocm():
85 + from triton_kernels.tensor_details import bitmatrix as _bm
86 + from triton_kernels.tensor_details.bitmatrix import (
87 + BitmatrixMetadata,
88 + _keyed_add,
89 + cdiv,
90 + )
91 + from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501
92 + sum_bitmatrix_rows,
93 + )
94 + else:
95 + from vllm.third_party.triton_kernels.tensor_details import (
96 + bitmatrix as _bm,
97 + )
98 + from vllm.third_party.triton_kernels.tensor_details.bitmatrix import (
99 + BitmatrixMetadata,
100 + _keyed_add,
101 + cdiv,
102 + )
103 + from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501
104 + sum_bitmatrix_rows,
105 + )
106 except ImportError:
107 return
108
vllm/models/deepseek_v4/amd/model.py deleted
-1
@@ -1 +0,0 @@
1 -../nvidia/model.py
\ No newline at end of file
vllm/models/deepseek_v4/amd/model.py new
+1612
@@ -0,0 +1,1612 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import typing
4 +from collections.abc import Callable, Iterable
5 +from itertools import islice
6 +
7 +import regex as re
8 +import torch
9 +import torch.nn as nn
10 +
11 +from vllm.compilation.decorators import support_torch_compile
12 +from vllm.config import VllmConfig
13 +from vllm.distributed import (
14 + get_ep_group,
15 + get_pp_group,
16 + get_tensor_model_parallel_rank,
17 + get_tensor_model_parallel_world_size,
18 +)
19 +from vllm.forward_context import get_forward_context
20 +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
21 +from vllm.model_executor.layers.fused_moe import FusedMoE
22 +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
23 + fused_topk_bias,
24 +)
25 +from vllm.model_executor.layers.fused_moe.router.norm_gate_linear import (
26 + NormGateLinear,
27 +)
28 +from vllm.model_executor.layers.layernorm import RMSNorm
29 +from vllm.model_executor.layers.linear import (
30 + ColumnParallelLinear,
31 + MergedColumnParallelLinear,
32 + RowParallelLinear,
33 +)
34 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
35 +from vllm.model_executor.layers.mhc import (
36 + HCHeadOp,
37 + MHCFusedPostPreOp,
38 + MHCPostOp,
39 + MHCPreOp,
40 +)
41 +from vllm.model_executor.layers.quantization import QuantizationConfig
42 +from vllm.model_executor.layers.rotary_embedding import get_rope
43 +from vllm.model_executor.layers.vocab_parallel_embedding import (
44 + ParallelLMHead,
45 + VocabParallelEmbedding,
46 +)
47 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
48 +from vllm.model_executor.models.interfaces import SupportsPP
49 +from vllm.model_executor.models.utils import (
50 + AutoWeightsLoader,
51 + PPMissingLayer,
52 + WeightsMapper,
53 + extract_layer_index,
54 + is_pp_missing_parameter,
55 + make_layers,
56 + maybe_prefix,
57 +)
58 +from vllm.model_executor.utils import set_weight_attrs
59 +from vllm.models.deepseek_v4.nvidia.ops.attention import (
60 + DeepseekV4Indexer,
61 + DeepseekV4MLAModules,
62 + DeepseekV4MultiHeadLatentAttentionWrapper,
63 +)
64 +from vllm.platforms import current_platform
65 +from vllm.sequence import IntermediateTensors
66 +from vllm.triton_utils import tl, triton
67 +from vllm.utils.torch_utils import direct_register_custom_op
68 +
69 +
70 +class DeepseekV4MLP(nn.Module):
71 + def __init__(
72 + self,
73 + hidden_size: int,
74 + intermediate_size: int,
75 + hidden_act: str,
76 + swiglu_limit: float | None = None,
77 + quant_config: QuantizationConfig | None = None,
78 + reduce_results: bool = True,
79 + is_sequence_parallel: bool = False,
80 + prefix: str = "",
81 + ) -> None:
82 + super().__init__()
83 +
84 + # If is_sequence_parallel, the input and output tensors are sharded
85 + # across the ranks within the tp_group. In this case the weights are
86 + # replicated and no collective ops are needed.
87 + # Otherwise we use standard TP with an allreduce at the end.
88 + self.gate_up_proj = MergedColumnParallelLinear(
89 + hidden_size,
90 + [intermediate_size] * 2,
91 + bias=False,
92 + quant_config=quant_config,
93 + disable_tp=is_sequence_parallel,
94 + prefix=f"{prefix}.gate_up_proj",
95 + )
96 + self.down_proj = RowParallelLinear(
97 + intermediate_size,
98 + hidden_size,
99 + bias=False,
100 + quant_config=quant_config,
101 + reduce_results=reduce_results,
102 + disable_tp=is_sequence_parallel,
103 + prefix=f"{prefix}.down_proj",
104 + )
105 + if hidden_act != "silu":
106 + raise ValueError(
107 + f"Unsupported activation: {hidden_act}. Only silu is supported for now."
108 + )
109 + if swiglu_limit is not None:
110 + self.act_fn = SiluAndMulWithClamp(swiglu_limit)
111 + else:
112 + self.act_fn = SiluAndMul()
113 +
114 + def forward(self, x):
115 + gate_up, _ = self.gate_up_proj(x)
116 + x = self.act_fn(gate_up)
117 + x, _ = self.down_proj(x)
118 + return x
119 +
120 +
121 +@triton.jit
122 +def _deepseek_v4_stage_mega_moe_inputs_kernel(
123 + hidden_states,
124 + x_fp8,
125 + x_sf,
126 + topk_ids,
127 + topk_weights,
128 + topk_idx_out,
129 + topk_weights_out,
130 + hidden_stride_m: tl.constexpr,
131 + hidden_stride_k: tl.constexpr,
132 + x_stride_m: tl.constexpr,
133 + x_stride_k: tl.constexpr,
134 + x_sf_stride_m: tl.constexpr,
135 + x_sf_stride_k: tl.constexpr,
136 + topk_ids_stride_m: tl.constexpr,
137 + topk_ids_stride_k: tl.constexpr,
138 + topk_weights_stride_m: tl.constexpr,
139 + topk_weights_stride_k: tl.constexpr,
140 + topk_idx_stride_m: tl.constexpr,
141 + topk_idx_stride_k: tl.constexpr,
142 + topk_weights_out_stride_m: tl.constexpr,
143 + topk_weights_out_stride_k: tl.constexpr,
144 + hidden_size: tl.constexpr,
145 + top_k: tl.constexpr,
146 + BLOCK_K: tl.constexpr,
147 + GROUP_K: tl.constexpr,
148 + BLOCK_TOPK: tl.constexpr,
149 +) -> None:
150 + token_id = tl.program_id(0)
151 + k_block_id = tl.program_id(1)
152 +
153 + k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
154 + k_mask = k_offsets < hidden_size
155 + hidden = tl.load(
156 + hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
157 + mask=k_mask,
158 + other=0.0,
159 + ).to(tl.float32)
160 +
161 + num_groups: tl.constexpr = BLOCK_K // GROUP_K
162 + hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
163 + amax = tl.max(hidden_groups, axis=1)
164 + amax = tl.maximum(amax, 1.0e-4)
165 +
166 + scale = amax / 448.0
167 + scale_bits = scale.to(tl.uint32, bitcast=True)
168 + scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
169 + tl.uint32
170 + )
171 + scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
172 + rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
173 +
174 + hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
175 + scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
176 + scaled = tl.reshape(scaled, [BLOCK_K])
177 + fp8 = scaled.to(tl.float8e4nv)
178 + tl.store(
179 + x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
180 + fp8,
181 + mask=k_mask,
182 + )
183 +
184 + scale_offsets = tl.arange(0, num_groups)
185 + packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
186 + tl.store(
187 + x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
188 + packed_scale,
189 + )
190 +
191 + if k_block_id == 0:
192 + topk_offsets = tl.arange(0, BLOCK_TOPK)
193 + topk_mask = topk_offsets < top_k
194 +
195 + ids = tl.load(
196 + topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
197 + mask=topk_mask,
198 + other=0,
199 + ).to(tl.int64)
200 + tl.store(
201 + topk_idx_out
202 + + token_id * topk_idx_stride_m
203 + + topk_offsets * topk_idx_stride_k,
204 + ids,
205 + mask=topk_mask,
206 + )
207 +
208 + weights = tl.load(
209 + topk_weights
210 + + token_id * topk_weights_stride_m
211 + + topk_offsets * topk_weights_stride_k,
212 + mask=topk_mask,
213 + other=0.0,
214 + )
215 + tl.store(
216 + topk_weights_out
217 + + token_id * topk_weights_out_stride_m
218 + + topk_offsets * topk_weights_out_stride_k,
219 + weights,
220 + mask=topk_mask,
221 + )
222 +
223 +
224 +def _stage_deepseek_v4_mega_moe_inputs(
225 + hidden_states: torch.Tensor,
226 + topk_weights: torch.Tensor,
227 + topk_ids: torch.Tensor,
228 + x_fp8: torch.Tensor,
229 + x_sf: torch.Tensor,
230 + topk_idx_out: torch.Tensor,
231 + topk_weights_out: torch.Tensor,
232 +) -> None:
233 + num_tokens, hidden_size = hidden_states.shape
234 + if num_tokens == 0:
235 + return
236 + if hidden_size % 128 != 0:
237 + raise ValueError(
238 + "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
239 + "a multiple of 128."
240 + )
241 + top_k = topk_ids.shape[1]
242 + if topk_weights.shape != topk_ids.shape:
243 + raise ValueError(
244 + "DeepSeek V4 MegaMoE input staging requires topk_weights and "
245 + "topk_ids to have the same shape."
246 + )
247 +
248 + block_k = 128
249 + grid = (num_tokens, triton.cdiv(hidden_size, block_k))
250 + block_topk = triton.next_power_of_2(top_k)
251 + _deepseek_v4_stage_mega_moe_inputs_kernel[grid](
252 + hidden_states,
253 + x_fp8,
254 + x_sf,
255 + topk_ids,
256 + topk_weights,
257 + topk_idx_out,
258 + topk_weights_out,
259 + hidden_states.stride(0),
260 + hidden_states.stride(1),
261 + x_fp8.stride(0),
262 + x_fp8.stride(1),
263 + x_sf.stride(0),
264 + x_sf.stride(1),
265 + topk_ids.stride(0),
266 + topk_ids.stride(1),
267 + topk_weights.stride(0),
268 + topk_weights.stride(1),
269 + topk_idx_out.stride(0),
270 + topk_idx_out.stride(1),
271 + topk_weights_out.stride(0),
272 + topk_weights_out.stride(1),
273 + hidden_size,
274 + top_k,
275 + BLOCK_K=block_k,
276 + GROUP_K=32,
277 + BLOCK_TOPK=block_topk,
278 + num_warps=4,
279 + )
280 +
281 +
282 +def make_deepseek_v4_expert_params_mapping(
283 + num_experts: int,
284 +) -> list[tuple[str, str, int, str]]:
285 + return [
286 + (
287 + "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
288 + f"experts.{expert_id}.{weight_name}.",
289 + expert_id,
290 + shard_id,
291 + )
292 + for expert_id in range(num_experts)
293 + for shard_id, weight_name in [
294 + ("w1", "w1"),
295 + ("w2", "w2"),
296 + ("w3", "w3"),
297 + ]
298 + ]
299 +
300 +
301 +class DeepseekV4MegaMoEExperts(nn.Module):
302 + _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
303 +
304 + def __init__(
305 + self,
306 + vllm_config: VllmConfig,
307 + *,
308 + num_experts: int,
309 + num_local_experts: int,
310 + experts_start_idx: int,
311 + top_k: int,
312 + hidden_size: int,
313 + intermediate_size: int,
314 + prefix: str = "",
315 + ):
316 + super().__init__()
317 + self.prefix = prefix
318 + self.num_experts = num_experts
319 + self.num_local_experts = num_local_experts
320 + self.experts_start_idx = experts_start_idx
321 + self.experts_end_idx = experts_start_idx + num_local_experts
322 + self.top_k = top_k
323 + self.hidden_size = hidden_size
324 + self.intermediate_size = intermediate_size
325 + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
326 +
327 + weight_attrs = {"weight_loader": self.weight_loader}
328 + self.w13_weight = nn.Parameter(
329 + torch.zeros(
330 + num_local_experts,
331 + 2 * intermediate_size,
332 + hidden_size // 2,
333 + dtype=torch.uint8,
334 + ),
335 + requires_grad=False,
336 + )
337 + set_weight_attrs(self.w13_weight, weight_attrs)
338 +
339 + self.w13_weight_scale = nn.Parameter(
340 + torch.zeros(
341 + num_local_experts,
342 + 2 * intermediate_size,
343 + hidden_size // 32,
344 + dtype=torch.uint8,
345 + ),
346 + requires_grad=False,
347 + )
348 + set_weight_attrs(self.w13_weight_scale, weight_attrs)
349 + self.w13_weight_scale.quant_method = "block"
350 +
351 + self.w2_weight = nn.Parameter(
352 + torch.zeros(
353 + num_local_experts,
354 + hidden_size,
355 + intermediate_size // 2,
356 + dtype=torch.uint8,
357 + ),
358 + requires_grad=False,
359 + )
360 + set_weight_attrs(self.w2_weight, weight_attrs)
361 +
362 + self.w2_weight_scale = nn.Parameter(
363 + torch.zeros(
364 + num_local_experts,
365 + hidden_size,
366 + intermediate_size // 32,
367 + dtype=torch.uint8,
368 + ),
369 + requires_grad=False,
370 + )
371 + set_weight_attrs(self.w2_weight_scale, weight_attrs)
372 + self.w2_weight_scale.quant_method = "block"
373 +
374 + self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
375 + self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
376 +
377 + # Register in the static forward context so the custom-op wrapper
378 + # can look up this module by name from within a torch.compile graph.
379 + compilation_config = vllm_config.compilation_config
380 + if prefix in compilation_config.static_forward_context:
381 + raise ValueError(f"Duplicate layer name: {prefix}")
382 + compilation_config.static_forward_context[prefix] = self
383 +
384 + def _map_global_expert_id(self, expert_id: int) -> int:
385 + if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
386 + return -1
387 + return expert_id - self.experts_start_idx
388 +
389 + def weight_loader(
390 + self,
391 + param: nn.Parameter,
392 + loaded_weight: torch.Tensor,
393 + weight_name: str,
394 + shard_id: str,
395 + expert_id: int,
396 + return_success: bool = False,
397 + ) -> bool | None:
398 + local_expert_id = self._map_global_expert_id(expert_id)
399 + if local_expert_id == -1:
400 + return False if return_success else None
401 +
402 + expert_data = param.data[local_expert_id]
403 + if shard_id in ("w1", "w3"):
404 + if "w13_" not in weight_name:
405 + return False if return_success else None
406 + shard_offset = 0 if shard_id == "w1" else self.intermediate_size
407 + expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)
408 + elif shard_id == "w2":
409 + if "w2_" not in weight_name:
410 + return False if return_success else None
411 + else:
412 + raise ValueError(f"Unsupported expert shard id: {shard_id}")
413 +
414 + if expert_data.shape != loaded_weight.shape:
415 + raise ValueError(
416 + f"DeepSeek V4 MegaMoE expert weight shape mismatch for "
417 + f"{weight_name}: parameter shard {tuple(expert_data.shape)} "
418 + f"vs checkpoint {tuple(loaded_weight.shape)}"
419 + )
420 + expert_data.copy_(loaded_weight)
421 + return True if return_success else None
422 +
423 + @staticmethod
424 + def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:
425 + return (sf.to(torch.int32) << 23).view(torch.float32)
426 +
427 + def _check_runtime_supported(self) -> None:
428 + if not torch.cuda.is_available():
429 + raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
430 + device = self.w13_weight.device
431 + if device.type != "cuda":
432 + raise NotImplementedError(
433 + "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA."
434 + )
435 + if torch.cuda.get_device_capability(device)[0] != 10:
436 + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
437 + if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
438 + raise ValueError(
439 + "DeepGEMM MegaMoE requires hidden and intermediate sizes "
440 + "to be multiples of 128."
441 + )
442 +
443 + def finalize_weights(self) -> None:
444 + if self._transformed_l1_weights is not None:
445 + return
446 +
447 + self._check_runtime_supported()
448 + import vllm.third_party.deep_gemm as deep_gemm
449 +
450 + w13_scale = deep_gemm.transform_sf_into_required_layout(
451 + self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
452 + 2 * self.intermediate_size,
453 + self.hidden_size,
454 + (1, 32),
455 + self.num_local_experts,
456 + )
457 + w2_scale = deep_gemm.transform_sf_into_required_layout(
458 + self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
459 + self.hidden_size,
460 + self.intermediate_size,
461 + (1, 32),
462 + self.num_local_experts,
463 + )
464 + self._transformed_l1_weights, self._transformed_l2_weights = (
465 + deep_gemm.transform_weights_for_mega_moe(
466 + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
467 + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
468 + )
469 + )
470 + # Drop the original loader-side parameters: the MegaMoE kernels only
471 + # consume the transformed views above. transform_weights_for_mega_moe
472 + # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
473 + # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
474 + # aliases the original storage, and _transformed_l2_weights still holds
475 + # it, so the storage stays live after we drop the Parameter.
476 + self.w13_weight = None
477 + self.w13_weight_scale = None
478 + self.w2_weight = None
479 + self.w2_weight_scale = None
480 +
481 + def get_symm_buffer(self):
482 + import vllm.third_party.deep_gemm as deep_gemm
483 +
484 + group = get_ep_group().device_group
485 + device = torch.accelerator.current_device_index()
486 + key = (
487 + id(group),
488 + device,
489 + self.num_experts,
490 + self.max_num_tokens,
491 + self.top_k,
492 + self.hidden_size,
493 + self.intermediate_size,
494 + )
495 + symm_buffer = self._symm_buffer_cache.get(key)
496 + if symm_buffer is None:
497 + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
498 + group,
499 + self.num_experts,
500 + self.max_num_tokens,
501 + self.top_k,
502 + self.hidden_size,
503 + self.intermediate_size,
504 + )
505 + self._symm_buffer_cache[key] = symm_buffer
506 + return symm_buffer
507 +
508 + def forward(
509 + self,
510 + hidden_states: torch.Tensor,
511 + topk_weights: torch.Tensor,
512 + topk_ids: torch.Tensor,
513 + *,
514 + activation_clamp: float | None,
515 + fast_math: bool = True,
516 + ) -> torch.Tensor:
517 + if hidden_states.shape[0] > self.max_num_tokens:
518 + raise ValueError(
519 + f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
520 + f"but the symmetric buffer was sized for {self.max_num_tokens}."
521 + )
522 + y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
523 + torch.ops.vllm.deepseek_v4_mega_moe_experts(
524 + hidden_states,
525 + topk_weights,
526 + topk_ids,
527 + y,
528 + self.prefix,
529 + activation_clamp,
530 + fast_math,
531 + )
532 + return y
533 +
534 + def _run_mega_moe(
535 + self,
536 + hidden_states: torch.Tensor,
537 + topk_weights: torch.Tensor,
538 + topk_ids: torch.Tensor,
539 + y: torch.Tensor,
540 + activation_clamp: float | None,
541 + fast_math: bool,
542 + ) -> None:
543 + import vllm.third_party.deep_gemm as deep_gemm
544 +
545 + symm_buffer = self.get_symm_buffer()
546 + num_tokens = hidden_states.shape[0]
547 + _stage_deepseek_v4_mega_moe_inputs(
548 + hidden_states,
549 + topk_weights,
550 + topk_ids,
551 + symm_buffer.x[:num_tokens],
552 + symm_buffer.x_sf[:num_tokens],
553 + symm_buffer.topk_idx[:num_tokens],
554 + symm_buffer.topk_weights[:num_tokens],
555 + )
556 +
557 + # This method must have been already called during the weight loading phase.
558 + # We call it again here to cover the dummy weight loading case.
559 + self.finalize_weights()
560 +
561 + assert self._transformed_l1_weights is not None
562 + assert self._transformed_l2_weights is not None
563 + deep_gemm.fp8_fp4_mega_moe(
564 + y,
565 + self._transformed_l1_weights,
566 + self._transformed_l2_weights,
567 + symm_buffer,
568 + activation_clamp=activation_clamp,
569 + fast_math=fast_math,
570 + )
571 +
572 +
573 +DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
574 +
575 +
576 +def _deepseek_v4_mega_moe_experts_op(
577 + hidden_states: torch.Tensor,
578 + topk_weights: torch.Tensor,
579 + topk_ids: torch.Tensor,
580 + out: torch.Tensor,
581 + layer_name: str,
582 + activation_clamp: float | None,
583 + fast_math: bool,
584 +) -> None:
585 + self = get_forward_context().no_compile_layers[layer_name]
586 + self._run_mega_moe(
587 + hidden_states,
588 + topk_weights,
589 + topk_ids,
590 + out,
591 + activation_clamp,
592 + fast_math,
593 + )
594 +
595 +
596 +def _deepseek_v4_mega_moe_experts_op_fake(
597 + hidden_states: torch.Tensor,
598 + topk_weights: torch.Tensor,
599 + topk_ids: torch.Tensor,
600 + out: torch.Tensor,
601 + layer_name: str,
602 + activation_clamp: float | None,
603 + fast_math: bool,
604 +) -> None:
605 + return None
606 +
607 +
608 +direct_register_custom_op(
609 + op_name="deepseek_v4_mega_moe_experts",
610 + op_func=_deepseek_v4_mega_moe_experts_op,
611 + mutates_args=["out"],
612 + fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
613 +)
614 +
615 +
616 +class DeepseekV4MoE(nn.Module):
617 + def __init__(
618 + self,
619 + vllm_config: VllmConfig,
620 + prefix: str = "",
621 + ):
622 + super().__init__()
623 +
624 + self.tp_size = get_tensor_model_parallel_world_size()
625 + config = vllm_config.model_config.hf_config
626 + quant_config = vllm_config.quant_config
627 + self.prefix = prefix
628 + self.use_mega_moe = (
629 + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
630 + )
631 + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
632 + raise NotImplementedError(
633 + "DeepSeek V4 MegaMoE currently requires expert parallel. "
634 + "Enable it with --enable-expert-parallel, or pick a different "
635 + "moe backend."
636 + )
637 +
638 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
639 + self.hidden_size = config.hidden_size
640 +
641 + self.n_routed_experts = config.n_routed_experts
642 + self.n_activated_experts = config.num_experts_per_tok
643 + self.moe_intermediate_size = config.moe_intermediate_size
644 + self.swiglu_limit = config.swiglu_limit
645 + self.renormalize = config.norm_topk_prob
646 + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
647 + if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
648 + raise NotImplementedError(
649 + "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
650 + )
651 + if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
652 + raise NotImplementedError(
653 + "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
654 + f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
655 + "deep_gemm_mega_moe for this checkpoint."
656 + )
657 +
658 + # Fused RMSNorm + gate: owns both ffn_norm and the gate matmul.
659 + self.norm_gate = NormGateLinear(
660 + hidden_size=config.hidden_size,
661 + num_experts=config.n_routed_experts,
662 + rms_eps=config.rms_norm_eps,
663 + prefix=f"{prefix}.norm_gate",
664 + )
665 + # Routing-side tensors live on ``norm_gate`` directly (not on the
666 + # inner gate); they are initialized to None in NormGatedLinear and
667 + # populated below depending on the MoE variant.
668 + is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
669 + self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
670 + if is_hash_moe:
671 + # hash MoE doesn't use e_score_correction_bias
672 + # Use randint instead of empty to avoid garbage values causing
673 + # invalid memory access in dummy mode (--load-format="dummy")
674 + self.norm_gate.tid2eid = nn.Parameter(
675 + torch.randint(
676 + 0,
677 + config.n_routed_experts,
678 + (config.vocab_size, config.num_experts_per_tok),
679 + dtype=self.hash_indices_dtype,
680 + ),
681 + requires_grad=False,
682 + )
683 + elif getattr(config, "topk_method", None) == "noaux_tc":
684 + self.norm_gate.e_score_correction_bias = nn.Parameter(
685 + torch.empty(config.n_routed_experts, dtype=torch.float32),
686 + requires_grad=False,
687 + )
688 +
689 + if config.n_shared_experts is None:
690 + self.shared_experts = None
691 + else:
692 + intermediate_size = config.moe_intermediate_size * config.n_shared_experts
693 +
694 + self.shared_experts = DeepseekV4MLP(
695 + hidden_size=config.hidden_size,
696 + intermediate_size=intermediate_size,
697 + hidden_act=config.hidden_act,
698 + swiglu_limit=self.swiglu_limit,
699 + quant_config=quant_config,
700 + reduce_results=self.use_mega_moe,
701 + prefix=f"{prefix}.shared_experts",
702 + )
703 +
704 + if self.use_mega_moe:
705 + self._init_mega_moe_experts(vllm_config, config, prefix)
706 + else:
707 + self._init_fused_moe_experts(config, quant_config, prefix)
708 +
709 + def _init_mega_moe_experts(
710 + self,
711 + vllm_config: VllmConfig,
712 + config,
713 + prefix: str,
714 + ) -> None:
715 + self.ep_group = get_ep_group()
716 + self.ep_size = self.ep_group.world_size
717 + self.ep_rank = self.ep_group.rank_in_group
718 + assert config.n_routed_experts % self.ep_size == 0
719 +
720 + self.n_local_experts = config.n_routed_experts // self.ep_size
721 + self.experts_start_idx = self.ep_rank * self.n_local_experts
722 + self.experts_end_idx = self.experts_start_idx + self.n_local_experts
723 +
724 + self.experts = DeepseekV4MegaMoEExperts(
725 + vllm_config,
726 + num_experts=config.n_routed_experts,
727 + num_local_experts=self.n_local_experts,
728 + experts_start_idx=self.experts_start_idx,
729 + top_k=config.num_experts_per_tok,
730 + hidden_size=config.hidden_size,
731 + intermediate_size=config.moe_intermediate_size,
732 + prefix=f"{prefix}.experts",
733 + )
734 +
735 + def _init_fused_moe_experts(
736 + self,
737 + config,
738 + quant_config,
739 + prefix: str,
740 + ) -> None:
741 + self.tp_rank = get_tensor_model_parallel_rank()
742 + assert config.n_routed_experts % self.tp_size == 0
743 +
744 + self.n_local_experts = config.n_routed_experts // self.tp_size
745 + self.experts_start_idx = self.tp_rank * self.n_local_experts
746 + self.experts_end_idx = self.experts_start_idx + self.n_local_experts
747 + # We don't pass `gate` into FusedMoE
748 + self.experts = FusedMoE(
749 + shared_experts=self.shared_experts,
750 + num_experts=config.n_routed_experts,
751 + top_k=config.num_experts_per_tok,
752 + hidden_size=config.hidden_size,
753 + intermediate_size=config.moe_intermediate_size,
754 + renormalize=config.norm_topk_prob,
755 + quant_config=quant_config,
756 + prefix=f"{prefix}.experts",
757 + scoring_func=self.scoring_func,
758 + routed_scaling_factor=self.routed_scaling_factor,
759 + e_score_correction_bias=self.norm_gate.e_score_correction_bias,
760 + hash_indices_table=self.norm_gate.tid2eid,
761 + swiglu_limit=self.swiglu_limit,
762 + router_logits_dtype=torch.float32,
763 + )
764 +
765 + def forward(
766 + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
767 + ) -> torch.Tensor:
768 + if self.norm_gate.tid2eid is not None and input_ids is None:
769 + raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
770 +
771 + if not self.use_mega_moe:
772 + return self._forward_fused_moe(hidden_states, input_ids)
773 +
774 + org_shape = hidden_states.shape
775 + normed_x, router_logits = self.norm_gate(hidden_states)
776 + topk_weights, topk_ids = fused_topk_bias(
777 + hidden_states=normed_x,
778 + gating_output=router_logits,
779 + scoring_func=self.scoring_func,
780 + e_score_correction_bias=self.norm_gate.e_score_correction_bias.data
781 + if self.norm_gate.e_score_correction_bias is not None
782 + else None,
783 + topk=self.n_activated_experts,
784 + renormalize=self.renormalize,
785 + indices_type=self.hash_indices_dtype,
786 + input_tokens=input_ids,
787 + hash_indices_table=self.norm_gate.tid2eid,
788 + routed_scaling_factor=self.routed_scaling_factor,
789 + )
790 + activation_clamp = (
791 + float(self.swiglu_limit) if self.swiglu_limit is not None else None
792 + )
793 + final_hidden_states = self.experts(
794 + normed_x,
795 + topk_weights,
796 + topk_ids,
797 + activation_clamp=activation_clamp,
798 + )
799 +
800 + if self.shared_experts is not None:
801 + shared_output = self.shared_experts(normed_x)
802 + final_hidden_states += shared_output
803 +
804 + return final_hidden_states.view(org_shape)
805 +
806 + def _forward_fused_moe(
807 + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
808 + ) -> torch.Tensor:
809 + assert not self.experts.is_internal_router
810 + org_shape = hidden_states.shape
811 + normed_x, router_logits = self.norm_gate(hidden_states)
812 + final_hidden_states = self.experts(
813 + hidden_states=normed_x,
814 + router_logits=router_logits,
815 + input_ids=input_ids,
816 + )
817 +
818 + return final_hidden_states.view(org_shape)
819 +
820 + def finalize_mega_moe_weights(self) -> None:
821 + if self.use_mega_moe:
822 + self.experts.finalize_weights()
823 +
824 +
825 +class DeepseekV4Attention(nn.Module):
826 + def __init__(
827 + self,
828 + vllm_config: VllmConfig,
829 + prefix: str,
830 + topk_indices_buffer: torch.Tensor | None = None,
831 + aux_stream_list: list[torch.cuda.Stream] | None = None,
832 + ):
833 + super().__init__()
834 + config = vllm_config.model_config.hf_config
835 + quant_config = vllm_config.quant_config
836 + layer_id = extract_layer_index(prefix)
837 +
838 + self.layer_id = layer_id
839 + self.hidden_size = config.hidden_size
840 + self.n_heads = config.num_attention_heads
841 + tp_size = get_tensor_model_parallel_world_size()
842 + assert self.n_heads % tp_size == 0
843 +
844 + self.n_local_heads = self.n_heads // tp_size
845 + self.q_lora_rank = config.q_lora_rank
846 + self.o_lora_rank = config.o_lora_rank
847 + self.head_dim = config.head_dim
848 + self.rope_head_dim = config.qk_rope_head_dim
849 + self.nope_head_dim = self.head_dim - self.rope_head_dim
850 + self.n_groups = config.o_groups
851 + self.n_local_groups = self.n_groups // tp_size
852 + self.window_size = config.sliding_window
853 + # NOTE(zyongye) Compress ratio can't be 0
854 + # we do this for because MTP layer is not included
855 + # in the compress ratio list
856 + if layer_id < config.num_hidden_layers:
857 + self.compress_ratio = max(1, config.compress_ratios[layer_id])
858 + else:
859 + self.compress_ratio = 1
860 + self.eps = config.rms_norm_eps
861 + self.max_position_embeddings = config.max_position_embeddings
862 +
863 + # Padded to min 64 heads for FlashMLA, initialized to -inf
864 + # (no sink effect). Weight loading fills the first n_local_heads slots.
865 + padded_heads = max(self.n_local_heads, 64)
866 + self.attn_sink = nn.Parameter(
867 + torch.full((padded_heads,), -float("inf"), dtype=torch.float32),
868 + requires_grad=False,
869 + )
870 +
871 + self.fused_wqa_wkv = MergedColumnParallelLinear(
872 + self.hidden_size,
873 + [self.q_lora_rank, self.head_dim],
874 + bias=False,
875 + quant_config=quant_config,
876 + prefix=f"{prefix}.fused_wqa_wkv",
877 + disable_tp=True, # fused ReplicatedLinear
878 + )
879 + self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
880 + self.wq_b = ColumnParallelLinear(
881 + self.q_lora_rank,
882 + self.n_heads * self.head_dim,
883 + bias=False,
884 + quant_config=quant_config,
885 + return_bias=False,
886 + prefix=f"{prefix}.wq_b",
887 + )
888 +
889 + self.kv_norm = RMSNorm(self.head_dim, self.eps)
890 + self.wo_a = ColumnParallelLinear(
891 + self.n_heads * self.head_dim // self.n_groups,
892 + self.n_groups * self.o_lora_rank,
893 + bias=False,
894 + quant_config=quant_config,
895 + return_bias=False,
896 + prefix=f"{prefix}.wo_a",
897 + )
898 + self.wo_a.is_bmm = True
899 + self.wo_a.bmm_batch_size = self.n_local_groups
900 + self.wo_b = RowParallelLinear(
901 + self.n_groups * self.o_lora_rank,
902 + self.hidden_size,
903 + bias=False,
904 + quant_config=quant_config,
905 + return_bias=False,
906 + prefix=f"{prefix}.wo_b",
907 + )
908 + self.softmax_scale = self.head_dim**-0.5
909 + self.scale_fmt = config.quantization_config["scale_fmt"]
910 +
911 + self.rope_parameters = config.rope_scaling
912 +
913 + # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it)
914 + rope_parameters = config.rope_parameters
915 + rope_parameters["rope_theta"] = (
916 + config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta
917 + )
918 + if config.rope_parameters["rope_type"] != "default":
919 + config.rope_parameters["rope_type"] = (
920 + "deepseek_yarn"
921 + if config.rope_parameters.get("apply_yarn_scaling", True)
922 + else "deepseek_llama_scaling"
923 + )
924 + rope_parameters["mscale"] = 0 # Disable mscale
925 + rope_parameters["mscale_all_dim"] = 0 # Disable mscale
926 + rope_parameters["is_deepseek_v4"] = True
927 + rope_parameters["rope_dim"] = self.rope_head_dim
928 + self.rotary_emb = get_rope(
929 + self.head_dim,
930 + max_position=self.max_position_embeddings,
931 + rope_parameters=rope_parameters,
932 + is_neox_style=False,
933 + )
934 +
935 + self.indexer = None
936 + if self.compress_ratio == 4:
937 + # Only C4A uses sparse attention and hence has indexer.
938 + self.indexer = DeepseekV4Indexer(
939 + vllm_config,
940 + config=config,
941 + hidden_size=self.hidden_size,
942 + q_lora_rank=self.q_lora_rank,
943 + quant_config=quant_config,
944 + cache_config=vllm_config.cache_config,
945 + topk_indices_buffer=topk_indices_buffer,
946 + compress_ratio=self.compress_ratio,
947 + prefix=f"{prefix}.indexer",
948 + )
949 +
950 + mla_modules = DeepseekV4MLAModules(
951 + vllm_config=vllm_config,
952 + fused_wqa_wkv=self.fused_wqa_wkv,
953 + q_norm=self.q_norm,
954 + wq_b=self.wq_b,
955 + kv_norm=self.kv_norm,
956 + wo_a=self.wo_a,
957 + wo_b=self.wo_b,
958 + attn_sink=self.attn_sink,
959 + rotary_emb=self.rotary_emb,
960 + indexer=self.indexer,
961 + indexer_rotary_emb=self.rotary_emb,
962 + topk_indices_buffer=topk_indices_buffer,
963 + aux_stream_list=aux_stream_list,
964 + )
965 + self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper(
966 + hidden_size=self.hidden_size,
967 + num_heads=self.n_local_heads,
968 + head_dim=self.head_dim,
969 + scale=self.softmax_scale,
970 + qk_nope_head_dim=self.nope_head_dim,
971 + qk_rope_head_dim=self.rope_head_dim,
972 + v_head_dim=self.head_dim,
973 + q_lora_rank=self.q_lora_rank,
974 + kv_lora_rank=self.head_dim,
975 + o_lora_rank=self.o_lora_rank,
976 + mla_modules=mla_modules,
977 + window_size=self.window_size,
978 + compress_ratio=self.compress_ratio,
979 + cache_config=vllm_config.cache_config,
980 + quant_config=quant_config,
981 + prefix=prefix,
982 + )
983 +
984 + def forward(
985 + self,
986 + positions: torch.Tensor,
987 + hidden_states: torch.Tensor,
988 + llama_4_scaling: torch.Tensor | None,
989 + ):
990 + return self.mla_attn(positions, hidden_states, llama_4_scaling)
991 +
992 +
993 +class DeepseekV4DecoderLayer(nn.Module):
994 + def __init__(
995 + self,
996 + vllm_config,
997 + prefix,
998 + topk_indices_buffer: torch.Tensor | None = None,
999 + aux_stream_list: list[torch.cuda.Stream] | None = None,
1000 + ):
1001 + super().__init__()
1002 +
1003 + # Lazy import to avoid top-level tilelang dependency.
1004 + # Registers both torch.ops.vllm.mhc_pre and mhc_post
1005 + import vllm.model_executor.layers.mhc # noqa: F401
1006 +
1007 + config = vllm_config.model_config.hf_config
1008 + self.hidden_size = config.hidden_size
1009 +
1010 + self.rms_norm_eps = config.rms_norm_eps
1011 + self.attn = DeepseekV4Attention(
1012 + vllm_config,
1013 + prefix=f"{prefix}.attn",
1014 + topk_indices_buffer=topk_indices_buffer,
1015 + aux_stream_list=aux_stream_list,
1016 + )
1017 + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
1018 +
1019 + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1020 + # ``ffn_norm`` is owned by ``self.ffn.norm_gate`` (fused with the
1021 + # router gate matmul); see ``NormGatedLinear``.
1022 + self.hc_mult = config.hc_mult
1023 + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
1024 + self.hc_eps = config.hc_eps
1025 + self.hc_post_alpha = 2.0
1026 + mix_hc = (2 + self.hc_mult) * self.hc_mult
1027 + hc_dim = self.hc_mult * self.hidden_size
1028 + self.hc_attn_fn = nn.Parameter(
1029 + torch.empty(
1030 + (mix_hc, hc_dim),
1031 + dtype=torch.float32,
1032 + ),
1033 + requires_grad=False,
1034 + )
1035 + self.hc_ffn_fn = nn.Parameter(
1036 + torch.empty(
1037 + (mix_hc, hc_dim),
1038 + dtype=torch.float32,
1039 + ),
1040 + requires_grad=False,
1041 + )
1042 + self.hc_attn_base = nn.Parameter(
1043 + torch.empty(
1044 + mix_hc,
1045 + dtype=torch.float32,
1046 + ),
1047 + requires_grad=False,
1048 + )
1049 + self.hc_ffn_base = nn.Parameter(
1050 + torch.empty(
1051 + mix_hc,
1052 + dtype=torch.float32,
1053 + ),
1054 + requires_grad=False,
1055 + )
1056 + self.hc_attn_scale = nn.Parameter(
1057 + torch.empty(
1058 + 3,
1059 + dtype=torch.float32,
1060 + ),
1061 + requires_grad=False,
1062 + )
1063 + self.hc_ffn_scale = nn.Parameter(
1064 + torch.empty(
1065 + 3,
1066 + dtype=torch.float32,
1067 + ),
1068 + requires_grad=False,
1069 + )
1070 + self.mhc_pre = MHCPreOp()
1071 + self.mhc_post = MHCPostOp()
1072 + self.mhc_fused_post_pre = MHCFusedPostPreOp()
1073 +
1074 + def hc_pre(
1075 + self,
1076 + x: torch.Tensor,
1077 + hc_fn: torch.Tensor,
1078 + hc_scale: torch.Tensor,
1079 + hc_base: torch.Tensor,
1080 + ):
1081 + post_mix, res_mix, layer_input = self.mhc_pre(
1082 + residual=x,
1083 + fn=hc_fn,
1084 + hc_scale=hc_scale,
1085 + hc_base=hc_base,
1086 + rms_eps=self.rms_norm_eps,
1087 + hc_pre_eps=self.hc_eps,
1088 + hc_sinkhorn_eps=self.hc_eps,
1089 + hc_post_mult_value=self.hc_post_alpha,
1090 + sinkhorn_repeat=self.hc_sinkhorn_iters,
1091 + )
1092 + return layer_input, post_mix, res_mix
1093 +
1094 + def hc_post(
1095 + self,
1096 + x: torch.Tensor,
1097 + residual: torch.Tensor,
1098 + post: torch.Tensor,
1099 + comb: torch.Tensor,
1100 + ):
1101 + return self.mhc_post(x, residual, post, comb)
1102 +
1103 + def _forward_cuda(
1104 + self,
1105 + x: torch.Tensor,
1106 + positions: torch.Tensor,
1107 + input_ids: torch.Tensor | None,
1108 + post_mix: torch.Tensor | None = None,
1109 + res_mix: torch.Tensor | None = None,
1110 + residual: torch.Tensor | None = None,
1111 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1112 + if residual is None:
1113 + # Run standalone hc_pre on first layer
1114 + residual = x
1115 + x, post_mix, res_mix = self.hc_pre(
1116 + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1117 + )
1118 + else:
1119 + residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
1120 + x,
1121 + residual,
1122 + post_mix,
1123 + res_mix,
1124 + self.hc_attn_fn,
1125 + self.hc_attn_scale,
1126 + self.hc_attn_base,
1127 + self.rms_norm_eps,
1128 + self.hc_eps,
1129 + self.hc_eps,
1130 + self.hc_post_alpha,
1131 + self.hc_sinkhorn_iters,
1132 + )
1133 +
1134 + x = self.attn_norm(x)
1135 + x = self.attn(positions, x, None)
1136 +
1137 + residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
1138 + x,
1139 + residual,
1140 + post_mix,
1141 + res_mix,
1142 + self.hc_ffn_fn,
1143 + self.hc_ffn_scale,
1144 + self.hc_ffn_base,
1145 + self.rms_norm_eps,
1146 + self.hc_eps,
1147 + self.hc_eps,
1148 + self.hc_post_alpha,
1149 + self.hc_sinkhorn_iters,
1150 + )
1151 + # ffn_norm is now folded into self.ffn.norm_gate; ffn() takes
1152 + # the pre-norm activation directly.
1153 + x = self.ffn(x, input_ids)
1154 + return x, residual, post_mix, res_mix
1155 +
1156 + def _forward_rocm(
1157 + self,
1158 + x: torch.Tensor,
1159 + positions: torch.Tensor,
1160 + input_ids: torch.Tensor | None,
1161 + post_mix: torch.Tensor | None = None,
1162 + res_mix: torch.Tensor | None = None,
1163 + residual: torch.Tensor | None = None,
1164 + ) -> tuple[
1165 + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
1166 + ]:
1167 + residual = x
1168 + x, post, comb = self.hc_pre(
1169 + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1170 + )
1171 + x = self.attn_norm(x)
1172 + x = self.attn(positions, x, None)
1173 + x = self.hc_post(x, residual, post, comb)
1174 +
1175 + residual = x
1176 + x, post, comb = self.hc_pre(
1177 + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
1178 + )
1179 + # ffn_norm is now folded into self.ffn.norm_gate; ffn() takes
1180 + # the pre-norm activation directly.
1181 + x = self.ffn(x, input_ids)
1182 + x = self.hc_post(x, residual, post, comb)
1183 + return x, None, None, None
1184 +
1185 + def forward(
1186 + self,
1187 + x: torch.Tensor,
1188 + positions: torch.Tensor,
1189 + input_ids: torch.Tensor | None,
1190 + post_mix: torch.Tensor | None = None,
1191 + res_mix: torch.Tensor | None = None,
1192 + residual: torch.Tensor | None = None,
1193 + ) -> tuple[
1194 + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
1195 + ]:
1196 + if current_platform.is_rocm():
1197 + return self._forward_rocm(
1198 + x, positions, input_ids, post_mix, res_mix, residual
1199 + )
1200 +
1201 + return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)
1202 +
1203 +
1204 +@support_torch_compile
1205 +class DeepseekV4Model(nn.Module):
1206 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1207 + super().__init__()
1208 +
1209 + config = vllm_config.model_config.hf_config
1210 + quant_config = vllm_config.quant_config
1211 + self.config = config
1212 + self.use_mega_moe = (
1213 + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
1214 + )
1215 + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
1216 + raise NotImplementedError(
1217 + "DeepSeek V4 MegaMoE currently requires expert parallel. "
1218 + "Enable it with --enable-expert-parallel, or pick a different "
1219 + "moe backend."
1220 + )
1221 + self.vocab_size = config.vocab_size
1222 + self.hc_eps = config.hc_eps
1223 + self.hc_mult = config.hc_mult
1224 + self.hc_dim = self.hc_mult * config.hidden_size
1225 + self.rms_norm_eps = config.rms_norm_eps
1226 +
1227 + # Three aux streams: one per non-default input GEMM in
1228 + # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
1229 + # (compressor kv_score, indexer.weights_proj, indexer.compressor
1230 + # kv_score). fused_wqa_wkv stays on the default stream.
1231 + # Disable them on ROCm because of hang issues.
1232 + aux_stream_list = (
1233 + None
1234 + if current_platform.is_rocm()
1235 + else [torch.cuda.Stream() for _ in range(3)]
1236 + )
1237 +
1238 + self.device = current_platform.device_type
1239 + # Reserved topk indices buffer for all Indexer layers to reuse.
1240 + self.topk_indices_buffer = torch.empty(
1241 + vllm_config.scheduler_config.max_num_batched_tokens,
1242 + config.index_topk,
1243 + dtype=torch.int32,
1244 + device=self.device,
1245 + )
1246 +
1247 + if get_pp_group().is_first_rank:
1248 + self.embed_tokens = VocabParallelEmbedding(
1249 + config.vocab_size,
1250 + config.hidden_size,
1251 + quant_config=quant_config,
1252 + prefix=f"{prefix}.embed_tokens",
1253 + )
1254 + else:
1255 + self.embed_tokens = PPMissingLayer()
1256 +
1257 + self.start_layer, self.end_layer, self.layers = make_layers(
1258 + config.num_hidden_layers,
1259 + lambda prefix: DeepseekV4DecoderLayer(
1260 + vllm_config,
1261 + prefix=prefix,
1262 + topk_indices_buffer=self.topk_indices_buffer,
1263 + aux_stream_list=aux_stream_list,
1264 + ),
1265 + prefix=f"{prefix}.layers",
1266 + )
1267 +
1268 + if get_pp_group().is_last_rank:
1269 + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
1270 + else:
1271 + self.norm = PPMissingLayer()
1272 +
1273 + self.hc_head_fn = nn.Parameter(
1274 + torch.empty(
1275 + self.hc_mult,
1276 + self.hc_dim,
1277 + dtype=torch.float32,
1278 + ),
1279 + requires_grad=False,
1280 + )
1281 + self.hc_head_base = nn.Parameter(
1282 + torch.empty(
1283 + self.hc_mult,
1284 + dtype=torch.float32,
1285 + ),
1286 + requires_grad=False,
1287 + )
1288 + self.hc_head_scale = nn.Parameter(
1289 + torch.empty(1, dtype=torch.float32),
1290 + requires_grad=False,
1291 + )
1292 + self.hc_head_op = HCHeadOp()
1293 + # Pre-hc_head residual stream buffer for the MTP draft. Stable
1294 + # address (outside the cudagraph pool) so the copy_ in forward()
1295 + # refreshes it correctly across captured shapes.
1296 + # refreshes it correctly across captured shapes. Only allocated on
1297 + # the last PP rank — that's where MTP target hidden states are
1298 + # produced.
1299 + if get_pp_group().is_last_rank:
1300 + self._mtp_hidden_buffer = torch.empty(
1301 + vllm_config.scheduler_config.max_num_batched_tokens,
1302 + self.hc_dim,
1303 + dtype=vllm_config.model_config.dtype,
1304 + device=self.device,
1305 + )
1306 + else:
1307 + self._mtp_hidden_buffer = None
1308 +
1309 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1310 + return self.embed_tokens(input_ids)
1311 +
1312 + def make_empty_intermediate_tensors(
1313 + self,
1314 + batch_size: int,
1315 + dtype: torch.dtype,
1316 + device: torch.device,
1317 + ) -> IntermediateTensors:
1318 + # PP intermediate tensors carry the multi-stream hidden_states
1319 + # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the
1320 + # token embedding to hc_mult streams before the first decoder
1321 + # layer and keeps that shape until hc_head() collapses it.
1322 + return IntermediateTensors(
1323 + {
1324 + "hidden_states": torch.zeros(
1325 + (batch_size, self.hc_mult, self.config.hidden_size),
1326 + dtype=dtype,
1327 + device=device,
1328 + ),
1329 + }
1330 + )
1331 +
1332 + def forward(
1333 + self,
1334 + input_ids: torch.Tensor,
1335 + positions: torch.Tensor,
1336 + intermediate_tensors: IntermediateTensors | None,
1337 + inputs_embeds: torch.Tensor | None = None,
1338 + ) -> torch.Tensor | IntermediateTensors:
1339 + if get_pp_group().is_first_rank:
1340 + if inputs_embeds is not None:
1341 + hidden_states = inputs_embeds
1342 + else:
1343 + hidden_states = self.embed_input_ids(input_ids)
1344 + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
1345 + else:
1346 + assert intermediate_tensors is not None
1347 + hidden_states = intermediate_tensors["hidden_states"]
1348 +
1349 + if self.use_mega_moe:
1350 + input_ids = input_ids.to(torch.int64)
1351 +
1352 + residual, post_mix, res_mix = None, None, None
1353 + for layer in islice(self.layers, self.start_layer, self.end_layer):
1354 + hidden_states, residual, post_mix, res_mix = layer(
1355 + hidden_states,
1356 + positions,
1357 + input_ids,
1358 + post_mix,
1359 + res_mix,
1360 + residual,
1361 + )
1362 + if layer is not None and current_platform.is_cuda():
1363 + hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
1364 +
1365 + if not get_pp_group().is_last_rank:
1366 + return IntermediateTensors({"hidden_states": hidden_states})
1367 +
1368 + # Stash pre-hc_head residual for the MTP draft (captured copy_).
1369 + num_tokens = hidden_states.shape[0]
1370 + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1371 +
1372 + hidden_states = self.hc_head_op(
1373 + hidden_states,
1374 + self.hc_head_fn,
1375 + self.hc_head_scale,
1376 + self.hc_head_base,
1377 + self.rms_norm_eps,
1378 + self.hc_eps,
1379 + )
1380 + hidden_states = self.norm(hidden_states)
1381 + return hidden_states
1382 +
1383 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1384 + stacked_params_mapping = [
1385 + # (param_name, shard_name, shard_id)
1386 + ("gate_up_proj", "w1", 0),
1387 + ("gate_up_proj", "w3", 1),
1388 + ("attn.fused_wqa_wkv", "attn.wq_a", 0),
1389 + ("attn.fused_wqa_wkv", "attn.wkv", 1),
1390 + ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
1391 + ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
1392 + ]
1393 + params_dict = dict(self.named_parameters())
1394 + loaded_params: set[str] = set()
1395 +
1396 + # TP for attention
1397 + tp_size = get_tensor_model_parallel_world_size()
1398 + tp_rank = get_tensor_model_parallel_rank()
1399 + n_head = self.config.num_attention_heads
1400 + n_local_head = n_head // tp_size
1401 + head_rank_start = n_local_head * tp_rank
1402 + head_rank_end = n_local_head * (tp_rank + 1)
1403 +
1404 + # Pre-compute expert mapping ONCE.
1405 + expert_mapping = self.get_expert_mapping()
1406 +
1407 + for name, loaded_weight in weights:
1408 + for param_name, weight_name, shard_id in stacked_params_mapping:
1409 + # Skip non-stacked layers and experts (experts handled below).
1410 + if ".experts." in name:
1411 + continue
1412 + if weight_name not in name:
1413 + continue
1414 + name = name.replace(weight_name, param_name)
1415 +
1416 + if is_pp_missing_parameter(name, self):
1417 + break
1418 + param = params_dict[name]
1419 + weight_loader = param.weight_loader
1420 + weight_loader(param, loaded_weight, shard_id)
1421 + loaded_params.add(name)
1422 + break
1423 + else:
1424 + if ".experts." in name:
1425 + # E8M0 scales are stored as float8_e8m0fnu in
1426 + # checkpoints but the MoE param is uint8. copy_()
1427 + # would do a numeric conversion (e.g. 2^-7 → 0),
1428 + # destroying the raw exponent bytes.
1429 + if (
1430 + "weight_scale" in name
1431 + and loaded_weight.dtype == torch.float8_e8m0fnu
1432 + ):
1433 + loaded_weight = loaded_weight.view(torch.uint8)
1434 + for mapping in expert_mapping:
1435 + param_name, weight_name, expert_id, expert_shard_id = mapping
1436 + if weight_name not in name:
1437 + continue
1438 + name_mapped = name.replace(weight_name, param_name)
1439 + if is_pp_missing_parameter(name_mapped, self):
1440 + continue
1441 + param = params_dict[name_mapped]
1442 + # We should ask the weight loader to return success or not
1443 + # here since otherwise we may skip experts with other
1444 + # available replicas.
1445 + weight_loader = typing.cast(
1446 + Callable[..., bool], param.weight_loader
1447 + )
1448 + success = weight_loader(
1449 + param,
1450 + loaded_weight,
1451 + name_mapped,
1452 + shard_id=expert_shard_id,
1453 + expert_id=expert_id,
1454 + return_success=True,
1455 + )
1456 + if success:
1457 + name = name_mapped
1458 + break
1459 + loaded_params.add(name_mapped)
1460 + continue
1461 + elif "attn_sink" in name:
1462 + if is_pp_missing_parameter(name, self):
1463 + continue
1464 + narrow_weight = loaded_weight[head_rank_start:head_rank_end]
1465 + n = narrow_weight.shape[0]
1466 + params_dict[name][:n].copy_(narrow_weight)
1467 + loaded_params.add(name)
1468 + continue
1469 + else:
1470 + if is_pp_missing_parameter(name, self):
1471 + continue
1472 + param = params_dict[name]
1473 + weight_loader = getattr(
1474 + param, "weight_loader", default_weight_loader
1475 + )
1476 + weight_loader(param, loaded_weight)
1477 + loaded_params.add(name)
1478 + continue
1479 +
1480 + return loaded_params
1481 +
1482 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1483 + first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
1484 + if first_layer.ffn.use_mega_moe:
1485 + return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
1486 + # Params for weights, fp8 weight scales, fp8 activation scales
1487 + # (param_name, weight_name, expert_id, shard_id)
1488 + return FusedMoE.make_expert_params_mapping(
1489 + self,
1490 + ckpt_gate_proj_name="w1",
1491 + ckpt_down_proj_name="w2",
1492 + ckpt_up_proj_name="w3",
1493 + num_experts=self.config.n_routed_experts,
1494 + )
1495 +
1496 + def finalize_mega_moe_weights(self) -> None:
1497 + for layer in islice(self.layers, self.start_layer, self.end_layer):
1498 + layer.ffn.finalize_mega_moe_weights()
1499 +
1500 +
1501 +def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1502 + if expert_dtype == "fp4":
1503 + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as
1504 + # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and
1505 + # shared experts use Fp8LinearMethod's block scales, which
1506 + # register as ``weight_scale_inv``.
1507 + scale_regex = {
1508 + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
1509 + re.compile(r"\.scale$"): ".weight_scale_inv",
1510 + }
1511 + else:
1512 + # FP8 experts use Fp8MoEMethod (block_quant=True), which registers
1513 + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
1514 + # there.
1515 + scale_regex = {
1516 + re.compile(r"\.scale$"): ".weight_scale_inv",
1517 + }
1518 + return WeightsMapper(
1519 + orig_to_new_prefix={
1520 + "layers.": "model.layers.",
1521 + "embed.": "model.embed.",
1522 + "norm.": "model.norm.",
1523 + "hc_head": "model.hc_head",
1524 + "mtp.": "model.mtp.",
1525 + },
1526 + orig_to_new_regex=scale_regex,
1527 + orig_to_new_suffix={
1528 + "head.weight": "lm_head.weight",
1529 + "embed.weight": "embed_tokens.weight",
1530 + # Pre-MoE norm + gate are now owned by ``DeepseekV4MoE.norm_gate``
1531 + # (see NormGatedLinear).
1532 + ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
1533 + ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
1534 + ".ffn.gate.bias": ".ffn.norm_gate.e_score_correction_bias",
1535 + # Hash MoE table also moved off the inner gate.
1536 + ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
1537 + },
1538 + orig_to_new_substr={
1539 + ".attn.compressor.": ".attn.mla_attn.compressor.",
1540 + ".shared_experts.w2": ".shared_experts.down_proj",
1541 + },
1542 + )
1543 +
1544 +
1545 +class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
1546 + model_cls = DeepseekV4Model
1547 +
1548 + # Default mapper assumes the original FP4-expert checkpoint layout.
1549 + # Overridden per-instance in __init__ when expert_dtype != "fp4".
1550 + hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
1551 +
1552 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1553 + super().__init__()
1554 +
1555 + config = vllm_config.model_config.hf_config
1556 + self.config = config
1557 + expert_dtype = getattr(config, "expert_dtype", "fp4")
1558 + if expert_dtype != "fp4":
1559 + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype)
1560 +
1561 + self.model = self.model_cls(
1562 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1563 + )
1564 + if get_pp_group().is_last_rank:
1565 + self.lm_head = ParallelLMHead(
1566 + config.vocab_size,
1567 + config.hidden_size,
1568 + prefix=maybe_prefix(prefix, "lm_head"),
1569 + )
1570 + else:
1571 + self.lm_head = PPMissingLayer()
1572 + self.logits_processor = LogitsProcessor(config.vocab_size)
1573 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1574 + self.model.make_empty_intermediate_tensors
1575 + )
1576 +
1577 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1578 + return self.model.embed_input_ids(input_ids)
1579 +
1580 + def compute_logits(
1581 + self,
1582 + hidden_states: torch.Tensor,
1583 + ) -> torch.Tensor | None:
1584 + logits = self.logits_processor(self.lm_head, hidden_states)
1585 + return logits
1586 +
1587 + def forward(
1588 + self,
1589 + input_ids: torch.Tensor,
1590 + positions: torch.Tensor,
1591 + intermediate_tensors: IntermediateTensors | None = None,
1592 + inputs_embeds: torch.Tensor | None = None,
1593 + ) -> torch.Tensor | IntermediateTensors:
1594 + hidden_states = self.model(
1595 + input_ids, positions, intermediate_tensors, inputs_embeds
1596 + )
1597 + return hidden_states
1598 +
1599 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
1600 + """Pre-hc_head residual stream buffer (max_num_batched_tokens,
1601 + hc_mult * hidden_size) for the MTP draft model. Populated by
1602 + forward(); valid after each target step."""
1603 + return getattr(self.model, "_mtp_hidden_buffer", None)
1604 +
1605 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1606 + loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
1607 + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1608 + self.model.finalize_mega_moe_weights()
1609 + return loaded_params
1610 +
1611 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1612 + return self.model.get_expert_mapping()
vllm/models/deepseek_v4/amd/mtp.py deleted
-1
@@ -1 +0,0 @@
1 -../nvidia/mtp.py
\ No newline at end of file
vllm/models/deepseek_v4/amd/mtp.py new
+520
@@ -0,0 +1,520 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4).
4 +
5 +Split from ``deepseek_mtp.py`` because the V4 architecture introduces several
6 +pieces that have no analogue in V3/V32:
7 + * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of
8 + the fused ``eh_proj``);
9 + * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``;
10 + * ``DeepseekV4DecoderLayer`` with its own aux-stream management;
11 + * V4-specific checkpoint weight-name remapping in ``load_weights``.
12 +"""
13 +
14 +import typing
15 +from collections.abc import Callable, Iterable
16 +
17 +import regex as re
18 +import torch
19 +import torch.nn as nn
20 +
21 +from vllm.compilation.decorators import support_torch_compile
22 +from vllm.config import VllmConfig
23 +from vllm.distributed import (
24 + get_tensor_model_parallel_rank,
25 + get_tensor_model_parallel_world_size,
26 +)
27 +from vllm.logger import init_logger
28 +from vllm.model_executor.layers.fused_moe import FusedMoE
29 +from vllm.model_executor.layers.layernorm import RMSNorm
30 +from vllm.model_executor.layers.linear import ReplicatedLinear
31 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
32 +from vllm.model_executor.layers.mhc import HCHeadOp
33 +from vllm.model_executor.layers.vocab_parallel_embedding import (
34 + VocabParallelEmbedding,
35 +)
36 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37 +from vllm.model_executor.models.deepseek_mtp import SharedHead
38 +from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
39 +from vllm.model_executor.models.utils import maybe_prefix
40 +from vllm.platforms import current_platform
41 +from vllm.sequence import IntermediateTensors
42 +
43 +from .model import (
44 + DeepseekV4DecoderLayer,
45 + make_deepseek_v4_expert_params_mapping,
46 +)
47 +
48 +logger = init_logger(__name__)
49 +
50 +# MoE expert scales are fused into per-layer w13/w2 tensors. The exact
51 +# parameter suffix depends on which FusedMoE method handles the experts:
52 +# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;
53 +# - fp8 experts (Fp8MoEMethod with block_quant=True) register
54 +# ``w{1,2,3}_weight_scale_inv``.
55 +# Other FP8 linear scales (including shared experts) always use
56 +# ``.weight_scale_inv``. Mirrors the per-instance mapper built by
57 +# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py.
58 +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
59 +
60 +
61 +class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
62 + def __init__(
63 + self,
64 + vllm_config: VllmConfig,
65 + topk_indices_buffer: torch.Tensor,
66 + prefix: str,
67 + aux_stream_list: list[torch.cuda.Stream] | None = None,
68 + ) -> None:
69 + super().__init__()
70 +
71 + assert vllm_config.speculative_config is not None
72 + config = vllm_config.speculative_config.draft_model_config.hf_config
73 + self.config = config
74 + quant_config = vllm_config.quant_config
75 + self.rms_norm_eps = config.rms_norm_eps
76 +
77 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
78 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
79 +
80 + # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than
81 + # fusing them the way V3 does with eh_proj.
82 + self.e_proj = ReplicatedLinear(
83 + config.hidden_size,
84 + config.hidden_size,
85 + bias=False,
86 + return_bias=False,
87 + quant_config=quant_config,
88 + )
89 + self.h_proj = ReplicatedLinear(
90 + config.hidden_size,
91 + config.hidden_size,
92 + bias=False,
93 + return_bias=False,
94 + quant_config=quant_config,
95 + )
96 +
97 + self.hc_eps = config.hc_eps
98 + self.hc_mult = config.hc_mult
99 + self.hc_dim = self.hc_mult * config.hidden_size
100 + self.hc_head_fn = nn.Parameter(
101 + torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32),
102 + requires_grad=False,
103 + )
104 + self.hc_head_base = nn.Parameter(
105 + torch.empty(self.hc_mult, dtype=torch.float32),
106 + requires_grad=False,
107 + )
108 + self.hc_head_scale = nn.Parameter(
109 + torch.empty(1, dtype=torch.float32),
110 + requires_grad=False,
111 + )
112 +
113 + self.shared_head = SharedHead(
114 + config=config, prefix=prefix, quant_config=quant_config
115 + )
116 + self.mtp_block = DeepseekV4DecoderLayer(
117 + vllm_config,
118 + prefix,
119 + topk_indices_buffer=topk_indices_buffer,
120 + aux_stream_list=aux_stream_list,
121 + )
122 +
123 + self.hc_head_op = HCHeadOp()
124 +
125 + def forward(
126 + self,
127 + input_ids: torch.Tensor,
128 + positions: torch.Tensor,
129 + previous_hidden_states: torch.Tensor,
130 + inputs_embeds: torch.Tensor | None = None,
131 + spec_step_index: int = 0,
132 + ) -> torch.Tensor:
133 + assert inputs_embeds is not None
134 + # masking inputs at position 0, as not needed by MTP
135 + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
136 + inputs_embeds = self.enorm(inputs_embeds)
137 +
138 + # Target stashes pre-hc_head residual as flat (T, hc_mult * D);
139 + # reshape to (T, hc_mult, D) — the training-time layout.
140 + previous_hidden_states = previous_hidden_states.view(
141 + -1, self.hc_mult, self.config.hidden_size
142 + )
143 + previous_hidden_states = self.hnorm(previous_hidden_states)
144 + hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
145 + inputs_embeds
146 + ).unsqueeze(-2)
147 + hidden_states, residual, post_mix, res_mix = self.mtp_block(
148 + positions=positions, x=hidden_states, input_ids=None
149 + )
150 + if current_platform.is_cuda():
151 + hidden_states = self.mtp_block.hc_post(
152 + hidden_states, residual, post_mix, res_mix
153 + )
154 + # Return the flat pre-hc_head residual so it can be re-fed as the
155 + # next spec step's `previous_hidden_states` when
156 + # num_speculative_tokens > 1. hc_head is deferred to compute_logits.
157 + return hidden_states.flatten(1)
158 +
159 +
160 +class DeepSeekV4MultiTokenPredictor(nn.Module):
161 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
162 + super().__init__()
163 + config = vllm_config.model_config.hf_config
164 + self.mtp_start_layer_idx = config.num_hidden_layers
165 + self.num_mtp_layers = config.num_nextn_predict_layers
166 + self.device = current_platform.device_type
167 +
168 + topk_tokens = config.index_topk
169 + self.topk_indices_buffer = torch.empty(
170 + vllm_config.scheduler_config.max_num_batched_tokens,
171 + topk_tokens,
172 + dtype=torch.int32,
173 + device=self.device,
174 + )
175 +
176 + # Three aux streams shared across all MTP layers, mirroring
177 + # DeepseekV4Model. ROCm runs the same work serially for now.
178 + aux_stream_list = (
179 + None
180 + if current_platform.is_rocm()
181 + else [torch.cuda.Stream() for _ in range(3)]
182 + )
183 +
184 + # to map the exact layer index from weights
185 + self.layers = torch.nn.ModuleDict(
186 + {
187 + str(idx): DeepSeekV4MultiTokenPredictorLayer(
188 + vllm_config,
189 + self.topk_indices_buffer,
190 + f"{prefix}.layers.{idx}",
191 + aux_stream_list=aux_stream_list,
192 + )
193 + for idx in range(
194 + self.mtp_start_layer_idx,
195 + self.mtp_start_layer_idx + self.num_mtp_layers,
196 + )
197 + }
198 + )
199 + self.embed_tokens = VocabParallelEmbedding(
200 + config.vocab_size,
201 + config.hidden_size,
202 + prefix=maybe_prefix(prefix, "embed_tokens"),
203 + )
204 + self.logits_processor = LogitsProcessor(config.vocab_size)
205 +
206 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
207 + return self.embed_tokens(input_ids)
208 +
209 + def forward(
210 + self,
211 + input_ids: torch.Tensor,
212 + positions: torch.Tensor,
213 + previous_hidden_states: torch.Tensor,
214 + inputs_embeds: torch.Tensor | None = None,
215 + spec_step_idx: int = 0,
216 + ) -> torch.Tensor:
217 + if inputs_embeds is None:
218 + inputs_embeds = self.embed_tokens(input_ids)
219 + current_step_idx = spec_step_idx % self.num_mtp_layers
220 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
221 + input_ids,
222 + positions,
223 + previous_hidden_states,
224 + inputs_embeds,
225 + current_step_idx,
226 + )
227 +
228 + def compute_logits(
229 + self,
230 + hidden_states: torch.Tensor,
231 + spec_step_idx: int = 0,
232 + ) -> torch.Tensor:
233 + current_step_idx = spec_step_idx % self.num_mtp_layers
234 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
235 + # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply
236 + # hc_head here so logits are computed from the dense hidden state.
237 + hidden_states = hidden_states.view(
238 + -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size
239 + )
240 + hidden_states = mtp_layer.hc_head_op(
241 + hidden_states,
242 + mtp_layer.hc_head_fn,
243 + mtp_layer.hc_head_scale,
244 + mtp_layer.hc_head_base,
245 + mtp_layer.rms_norm_eps,
246 + mtp_layer.hc_eps,
247 + )
248 + logits = self.logits_processor(
249 + mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
250 + )
251 + return logits
252 +
253 +
254 +@support_torch_compile
255 +class DeepSeekV4MTP(nn.Module):
256 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
257 + super().__init__()
258 + self.config = vllm_config.model_config.hf_config
259 + self.quant_config = vllm_config.quant_config
260 + self.model = DeepSeekV4MultiTokenPredictor(
261 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
262 + )
263 +
264 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
265 + return self.model.embed_input_ids(input_ids)
266 +
267 + def forward(
268 + self,
269 + input_ids: torch.Tensor | None,
270 + positions: torch.Tensor,
271 + hidden_states: torch.Tensor,
272 + intermediate_tensors: IntermediateTensors | None = None,
273 + inputs_embeds: torch.Tensor | None = None,
274 + spec_step_idx: int = 0,
275 + ) -> torch.Tensor:
276 + hidden_states = self.model(
277 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
278 + )
279 + return hidden_states
280 +
281 + def compute_logits(
282 + self,
283 + hidden_states: torch.Tensor,
284 + spec_step_idx: int = 0,
285 + ) -> torch.Tensor | None:
286 + return self.model.compute_logits(hidden_states, spec_step_idx)
287 +
288 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
289 + # Weight name remapping for checkpoint compatibility.
290 + # Maps checkpoint weight paths to model parameter paths.
291 + WEIGHT_NAME_REMAPPING: dict[str, str] = {
292 + ".emb.tok_emb.weight": ".embed_tokens.weight",
293 + ".head.weight": ".shared_head.head.weight",
294 + ".norm.weight": ".shared_head.norm.weight",
295 + # Pre-MoE norm + gate are now owned by
296 + # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear).
297 + ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
298 + ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
299 + ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
300 + }
301 +
302 + def _remap_weight_name(name: str) -> str:
303 + """Remap checkpoint weight names to model parameter names."""
304 + for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items():
305 + if old_pattern in name:
306 + name = name.replace(old_pattern, new_pattern)
307 + return name
308 +
309 + def _find_mtp_layer_idx(name: str) -> int:
310 + subnames = name.split(".")
311 + for subname in subnames:
312 + try:
313 + # we return the first encountered integer
314 + return int(subname)
315 + except ValueError:
316 + continue
317 + return 0
318 +
319 + stacked_params_mapping = [
320 + # (param_name, shard_name, shard_id)
321 + ("gate_up_proj", "w1", 0),
322 + ("gate_up_proj", "w3", 1),
323 + ("attn.fused_wqa_wkv", "attn.wq_a", 0),
324 + ("attn.fused_wqa_wkv", "attn.wkv", 1),
325 + ]
326 + params_dict = dict(self.named_parameters())
327 + loaded_params: set[str] = set()
328 +
329 + # TP for attention
330 + tp_size = get_tensor_model_parallel_world_size()
331 + tp_rank = get_tensor_model_parallel_rank()
332 + n_head = self.config.num_attention_heads
333 + n_local_head = n_head // tp_size
334 + head_rank_start = n_local_head * tp_rank
335 + head_rank_end = n_local_head * (tp_rank + 1)
336 +
337 + # Pre-compute expert mapping ONCE.
338 + first_layer = next(iter(self.model.layers.values()))
339 + if first_layer.mtp_block.ffn.use_mega_moe:
340 + expert_mapping = make_deepseek_v4_expert_params_mapping(
341 + self.config.n_routed_experts
342 + )
343 + else:
344 + expert_mapping = FusedMoE.make_expert_params_mapping(
345 + self,
346 + ckpt_gate_proj_name="w1",
347 + ckpt_down_proj_name="w2",
348 + ckpt_up_proj_name="w3",
349 + num_experts=self.config.n_routed_experts,
350 + )
351 +
352 + # FP8 experts register ``..._weight_scale_inv`` (block_quant) while
353 + # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
354 + # for the rename below based on the model's expert dtype.
355 + expert_scale_suffix = (
356 + ".weight_scale"
357 + if getattr(self.config, "expert_dtype", "fp4") == "fp4"
358 + else ".weight_scale_inv"
359 + )
360 +
361 + for name, loaded_weight in weights:
362 + mtp_layer_idx = _find_mtp_layer_idx(name)
363 + # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to
364 + # `model.layers.{num_hidden_layers + i}.*` so that
365 + # get_spec_layer_idx_from_weight_name can identify them.
366 + name = name.replace(
367 + f"mtp.{mtp_layer_idx}.",
368 + f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.",
369 + )
370 +
371 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
372 + if spec_layer is None:
373 + continue
374 +
375 + name = _remap_weight_name(name)
376 + name = self._rewrite_spec_layer_name(spec_layer, name)
377 +
378 + if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name:
379 + continue
380 + if name.endswith(".scale"):
381 + suffix = (
382 + expert_scale_suffix
383 + if _EXPERT_SCALE_RE.search(name)
384 + else ".weight_scale_inv"
385 + )
386 + name = name.removesuffix(".scale") + suffix
387 + for param_name, weight_name, shard_id in stacked_params_mapping:
388 + # Skip non-stacked layers and experts (experts handled below).
389 + if ".experts." in name:
390 + continue
391 + if weight_name not in name:
392 + continue
393 + name = name.replace(weight_name, param_name)
394 +
395 + param = params_dict[name]
396 + weight_loader = param.weight_loader
397 + weight_loader(param, loaded_weight, shard_id)
398 + loaded_params.add(name)
399 + break
400 + else:
401 + if ".experts." in name:
402 + # Reinterpret E8M0 scales as uint8 to preserve raw
403 + # exponent bytes; numeric copy_() would zero them.
404 + # Mirrors the main DeepseekV4 loader.
405 + if (
406 + "weight_scale" in name
407 + and loaded_weight.dtype == torch.float8_e8m0fnu
408 + ):
409 + loaded_weight = loaded_weight.view(torch.uint8)
410 + for mapping in expert_mapping:
411 + param_name, weight_name, expert_id, expert_shard_id = mapping
412 + if weight_name not in name:
413 + continue
414 + name_mapped = name.replace(weight_name, param_name)
415 + param = params_dict[name_mapped]
416 + # We should ask the weight loader to return success or not
417 + # here since otherwise we may skip experts with other
418 + # available replicas.
419 + weight_loader = typing.cast(
420 + Callable[..., bool], param.weight_loader
421 + )
422 + success = weight_loader(
423 + param,
424 + loaded_weight,
425 + name_mapped,
426 + shard_id=expert_shard_id,
427 + expert_id=expert_id,
428 + return_success=True,
429 + )
430 + if success:
431 + name = name_mapped
432 + loaded_params.add(name_mapped)
433 + break
434 + continue
435 + elif "attn_sink" in name:
436 + narrow_weight = loaded_weight[head_rank_start:head_rank_end]
437 + n = narrow_weight.shape[0]
438 + params_dict[name][:n].copy_(narrow_weight)
439 + loaded_params.add(name)
440 + continue
441 + else:
442 + if ".shared_experts.w2" in name:
443 + name = name.replace(
444 + ".shared_experts.w2", ".shared_experts.down_proj"
445 + )
446 + if name.endswith(".ffn.gate.bias"):
447 + # ``e_score_correction_bias`` lives on
448 + # ``norm_gate`` directly (not on the inner gate).
449 + name = name.replace(
450 + ".ffn.gate.bias",
451 + ".ffn.norm_gate.e_score_correction_bias",
452 + )
453 + param = params_dict[name]
454 + weight_loader = getattr(
455 + param, "weight_loader", default_weight_loader
456 + )
457 + weight_loader(param, loaded_weight)
458 + loaded_params.add(name)
459 + continue
460 +
461 + loaded_layers: set[int] = set()
462 + for param_name in loaded_params:
463 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
464 + if spec_layer is not None:
465 + loaded_layers.add(spec_layer)
466 + for layer_idx in range(
467 + self.model.mtp_start_layer_idx,
468 + self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
469 + ):
470 + if layer_idx not in loaded_layers:
471 + raise ValueError(
472 + f"MTP speculative decoding layer {layer_idx} weights "
473 + f"missing from checkpoint. The checkpoint may have "
474 + f"been quantized without including the MTP layers. "
475 + f"Use a checkpoint that includes MTP layer weights, "
476 + f"or disable speculative decoding."
477 + )
478 + self.finalize_mega_moe_weights()
479 + logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
480 + return loaded_params
481 +
482 + def finalize_mega_moe_weights(self) -> None:
483 + for layer in self.model.layers.values():
484 + layer.mtp_block.ffn.finalize_mega_moe_weights()
485 +
486 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
487 + """
488 + Rewrite the weight name to match the format of the original model.
489 + Add .mtp_block for modules in transformer layer block for spec layer
490 + and rename shared layer weights to be top level.
491 + """
492 + spec_layer_weight_names = [
493 + "embed_tokens",
494 + "enorm",
495 + "hnorm",
496 + "h_proj",
497 + "e_proj",
498 + "shared_head",
499 + "hc_head_fn",
500 + "hc_head_base",
501 + "hc_head_scale",
502 + ]
503 + shared_weight_names = ["embed_tokens"]
504 + spec_layer_weight = False
505 + shared_weight = False
506 + for weight_name in spec_layer_weight_names:
507 + if weight_name in name:
508 + spec_layer_weight = True
509 + if weight_name in shared_weight_names:
510 + shared_weight = True
511 + break
512 + if not spec_layer_weight:
513 + # treat rest weights as weights for transformer layer block
514 + name = name.replace(
515 + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
516 + )
517 + elif shared_weight:
518 + # treat shared weights as top level weights
519 + name = name.replace(f"model.layers.{spec_layer}.", "model.")
520 + return name
vllm/models/deepseek_v4/amd/rocm.py
+134 -23
@@ -44,6 +44,127 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
44 return indptr
45
46
47 +# ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes
48 +# do not touch the shared DeepSeek V4 cache utilities.
49 +_SPARSE_PREFILL_TOPK_ALIGNMENT = 128
50 +
51 +
52 +@triton.jit
53 +def _combine_topk_swa_indices_kernel(
54 + combined_indices_ptr,
55 + combined_indices_stride,
56 + combined_lens_ptr,
57 + topk_indices_ptr,
58 + topk_indices_stride,
59 + query_start_loc_ptr,
60 + seq_lens_ptr,
61 + gather_lens_ptr,
62 + M,
63 + N,
64 + TOP_K: tl.constexpr,
65 + COMPRESS_RATIO: tl.constexpr,
66 + WINDOW_SIZE: tl.constexpr,
67 + TOPK_WIDTH: tl.constexpr,
68 + PADDED_TOP_K: tl.constexpr,
69 +):
70 + batch_idx = tl.program_id(0)
71 + worker_id = tl.program_id(1)
72 + num_workers = tl.num_programs(1)
73 +
74 + base = tl.load(query_start_loc_ptr)
75 + query_start = tl.load(query_start_loc_ptr + batch_idx) - base
76 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
77 + query_len = query_end - query_start
78 + seq_len = tl.load(seq_lens_ptr + batch_idx)
79 + gather_len = tl.load(gather_lens_ptr + batch_idx)
80 + start_pos = seq_len - query_len
81 + gather_start = seq_len - gather_len
82 +
83 + for token_idx in range(query_start + worker_id, query_end, num_workers):
84 + token_idx_in_query = token_idx - query_start
85 + pos = start_pos + token_idx_in_query
86 + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
87 + swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
88 +
89 + topk_offset = tl.arange(0, PADDED_TOP_K)
90 + topk_mask = topk_offset < topk_len
91 + safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0)
92 + topk_indices = tl.load(
93 + topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset,
94 + mask=topk_mask,
95 + other=-1,
96 + )
97 + valid_topk = (topk_indices >= 0) & (topk_indices < N)
98 + topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1)
99 + tl.store(
100 + combined_indices_ptr + token_idx * combined_indices_stride + topk_offset,
101 + topk_indices,
102 + mask=topk_mask,
103 + )
104 +
105 + swa_offset = tl.arange(0, WINDOW_SIZE)
106 + tl.store(
107 + combined_indices_ptr
108 + + token_idx * combined_indices_stride
109 + + topk_len
110 + + swa_offset,
111 + M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start,
112 + mask=swa_offset < swa_len,
113 + )
114 +
115 + tl.store(combined_lens_ptr + token_idx, topk_len + swa_len)
116 +
117 +
118 +def combine_topk_swa_indices(
119 + topk_indices: torch.Tensor,
120 + query_start_loc: torch.Tensor,
121 + seq_lens: torch.Tensor,
122 + gather_lens: torch.Tensor,
123 + window_size: int,
124 + compress_ratio: int,
125 + topk: int,
126 + M: int,
127 + N: int,
128 +) -> tuple[torch.Tensor, torch.Tensor]:
129 + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
130 + num_tokens = topk_indices.shape[0]
131 + num_reqs = seq_lens.shape[0]
132 + combined_topk = (
133 + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
134 + // _SPARSE_PREFILL_TOPK_ALIGNMENT
135 + * _SPARSE_PREFILL_TOPK_ALIGNMENT
136 + )
137 + combined_indices = torch.full(
138 + (num_tokens, combined_topk),
139 + fill_value=-1,
140 + dtype=torch.int32,
141 + device=topk_indices.device,
142 + )
143 + combined_lens = torch.empty(
144 + num_tokens, dtype=torch.int32, device=topk_indices.device
145 + )
146 +
147 + num_workers = 128
148 + _combine_topk_swa_indices_kernel[(num_reqs, num_workers)](
149 + combined_indices,
150 + combined_indices.stride(0),
151 + combined_lens,
152 + topk_indices,
153 + topk_indices.stride(0),
154 + query_start_loc,
155 + seq_lens,
156 + gather_lens,
157 + M,
158 + N,
159 + TOP_K=topk,
160 + COMPRESS_RATIO=compress_ratio,
161 + WINDOW_SIZE=window_size,
162 + TOPK_WIDTH=topk_indices.shape[-1],
163 + PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
164 + )
165 + return combined_indices, combined_lens
166 +
167 +
168 @triton.jit
169 def _compute_topk_lens_kernel(
170 topk_lens_ptr,
@@ -704,31 +825,23 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
825 query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
826 )
827
707 - combined_ragged_indices, combined_ragged_indptr, combined_lens = (
708 - combine_topk_swa_indices_ragged(
709 - topk_indices[query_start:query_end],
710 - query_start_loc[
711 - num_decodes + chunk_start : num_decodes + chunk_end + 1
712 - ],
713 - seq_lens[chunk_start:chunk_end],
714 - gather_lens[chunk_start:chunk_end],
715 - layer.window_size,
716 - layer.compress_ratio,
717 - top_k,
718 - M,
719 - N,
720 - )
828 + combined_indices, combined_lens = combine_topk_swa_indices(
829 + topk_indices[query_start:query_end],
830 + query_start_loc[
831 + num_decodes + chunk_start : num_decodes + chunk_end + 1
832 + ],
833 + seq_lens[chunk_start:chunk_end],
834 + gather_lens[chunk_start:chunk_end],
835 + layer.window_size,
836 + layer.compress_ratio,
837 + top_k,
838 + M,
839 + N,
840 )
841 rocm_sparse_attn_prefill(
842 q=q[query_start:query_end],
843 kv=kv.view(-1, 1, q.shape[-1]),
725 - indices=torch.empty(
726 - q[query_start:query_end].shape[0],
727 - 1,
728 - 0,
729 - dtype=torch.int32,
730 - device=q.device,
731 - ),
844 + indices=combined_indices,
845 topk_length=combined_lens,
846 scale=layer.scale,
847 head_dim=layer.head_dim,
@@ -736,6 +849,4 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
849 rope_head_dim=layer.rope_head_dim,
850 attn_sink=layer.attn_sink,
851 output=output[query_start:query_end],
739 - ragged_indices=combined_ragged_indices,
740 - ragged_indptr=combined_ragged_indptr,
852 )
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+46 -16
@@ -144,35 +144,57 @@ def _cp_gather_indexer_quant_cache_kernel(
144 HEAD_DIM: tl.constexpr,
145 BLOCK_TILE_SIZE: tl.constexpr,
146 HEAD_TILE_SIZE: tl.constexpr,
147 + NUM_TOKENS: tl.constexpr,
148 + NUM_BATCHES: tl.constexpr,
149 + BLOCK_TABLE_WIDTH: tl.constexpr,
150 + NUM_BLOCKS: tl.constexpr,
151 ):
152 tid = tl.program_id(0)
153 offset = tl.arange(0, HEAD_DIM)
150 - batch_id = tl.load(token_to_seq_ptr + tid)
151 - batch_start = tl.load(cu_seqlen_ptr + batch_id)
152 - batch_end = tl.load(cu_seqlen_ptr + batch_id + 1)
154 + valid_tid = tid < NUM_TOKENS
155 + batch_id = tl.load(token_to_seq_ptr + tid, mask=valid_tid, other=-1)
156 + valid_batch = (batch_id >= 0) & (batch_id < NUM_BATCHES)
157 + safe_batch_id = tl.where(valid_batch, batch_id, 0)
158 + batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0)
159 + batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0)
160 batch_offset = tid - batch_start
154 - if tid >= batch_end:
161 + valid_token = valid_tid & valid_batch & (tid >= batch_start) & (tid < batch_end)
162 + if not valid_token:
163 return
164 block_table_id = batch_offset // block_size
165 block_offset = batch_offset % block_size
158 - block_table_offset = batch_id * block_table_stride + block_table_id
159 - block_id = tl.load(block_table_ptr + block_table_offset)
160 - tiled_block_id = block_offset // BLOCK_TILE_SIZE
161 - tiled_block_offset = block_offset % BLOCK_TILE_SIZE
166 + valid_block_table = (
167 + valid_token
168 + & (block_table_id >= 0)
169 + & (block_table_id < BLOCK_TABLE_WIDTH)
170 + & (block_offset >= 0)
171 + & (block_offset < block_size)
172 + )
173 + safe_block_table_id = tl.where(valid_block_table, block_table_id, 0)
174 + block_table_offset = safe_batch_id * block_table_stride + safe_block_table_id
175 + block_id = tl.load(
176 + block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
177 + )
178 + valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS)
179 + safe_block_id = tl.where(valid_block, block_id, 0)
180 + safe_block_offset = tl.where(valid_block, block_offset, 0)
181 + tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE
182 if LAYOUT == "SHUFFLE":
183 src_cache_offset = (
164 - block_id * kv_cache_stride
165 - + tiled_block_id * HEAD_DIM * BLOCK_TILE_SIZE
184 + safe_block_id * kv_cache_stride
185 + + (safe_block_offset // BLOCK_TILE_SIZE) * HEAD_DIM * BLOCK_TILE_SIZE
186 + tiled_block_offset * HEAD_TILE_SIZE
187 )
188 else:
169 - src_cache_offset = block_id * kv_cache_stride + block_offset * HEAD_DIM
170 - src_scale_offset = block_id * kv_cache_scale_stride + block_offset
189 + src_cache_offset = (
190 + safe_block_id * kv_cache_stride + safe_block_offset * HEAD_DIM
191 + )
192 + src_scale_offset = safe_block_id * kv_cache_scale_stride + safe_block_offset
193 dst_offset = tid * HEAD_DIM
194 src_scale_ptr = kv_cache_scale_ptr + src_scale_offset
195 src_cache_ptr = kv_cache_ptr + src_cache_offset
196 dst_k_ptr = k_fp8_ptr + dst_offset
175 - scale_val = tl.load(src_scale_ptr)
197 + scale_val = tl.load(src_scale_ptr, mask=valid_block, other=0.0)
198 tl.store(k_scale_ptr + tid, scale_val)
199 if LAYOUT == "SHUFFLE":
200 tiled_src_offset = (
@@ -182,7 +204,7 @@ def _cp_gather_indexer_quant_cache_kernel(
204 else:
205 tiled_src_offset = offset
206 val = tl.load(src_cache_ptr + tiled_src_offset)
185 - tl.store(dst_k_ptr + offset, val)
207 + tl.store(dst_k_ptr + offset, val, mask=valid_block)
208
209
210 def cp_gather_indexer_k_quant_cache_triton(
@@ -223,6 +245,10 @@ def cp_gather_indexer_k_quant_cache_triton(
245 head_dim,
246 block_tile_size,
247 head_tile_size,
248 + num_tokens,
249 + cu_seqlen.shape[0] - 1,
250 + block_table.shape[1],
251 + num_blocks,
252 )
253
254
@@ -943,8 +969,9 @@ def _pack_dense_prefix_to_ragged_kernel(
969 return
970
971 mask = offsets < row_len
972 + safe_offsets = tl.where(offsets < row_width, offsets, 0)
973 vals = tl.load(
947 - indices_ptr + row_idx * indices_stride0 + offsets,
974 + indices_ptr + row_idx * indices_stride0 + safe_offsets,
975 mask=mask & (offsets < row_width),
976 other=-1,
977 ).to(tl.int32)
@@ -1060,9 +1087,12 @@ def _sparse_attn_prefill_ragged_kernel(
1087 in_range = k_pos < kv_len
1088 slot = tl.load(kv_indices_ptr + kv_start + k_pos, mask=in_range, other=-1)
1089 valid = in_range & (slot >= 0) & (slot < num_kv)
1090 + safe_slot = tl.where(valid, slot, 0)
1091
1092 kv = tl.load(
1065 - kv_ptr + slot[:, None] * kv_stride_n + dim_offsets[None, :] * kv_stride_d,
1093 + kv_ptr
1094 + + safe_slot[:, None] * kv_stride_n
1095 + + dim_offsets[None, :] * kv_stride_d,
1096 mask=valid[:, None] & dim_mask[None, :],
1097 other=0.0,
1098 )
vllm/v1/spec_decode/llm_base_proposer.py
+6
@@ -240,6 +240,10 @@ class SpecDecodeBaseProposer:
240 # Determine allowed attention backends once during initialization.
241 self.allowed_attn_types: tuple | None = None
242 if current_platform.is_rocm():
243 + from vllm.models.deepseek_v4.amd.rocm import (
244 + DeepseekV4ROCMAiterMLASparseMetadata,
245 + DeepseekV4ROCMAiterSparseSWAMetadata,
246 + )
247 from vllm.v1.attention.backends.mla.indexer import (
248 DeepseekV32IndexerMetadata,
249 )
@@ -252,6 +256,8 @@ class SpecDecodeBaseProposer:
256 TritonAttentionMetadata,
257 RocmAttentionMetadata,
258 ROCMAiterMLASparseMetadata,
259 + DeepseekV4ROCMAiterMLASparseMetadata,
260 + DeepseekV4ROCMAiterSparseSWAMetadata,
261 DeepseekV32IndexerMetadata,
262 ]
263 # ROCM_AITER_FA is an optional backend