| 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.config import VllmConfig |
| 12 | from vllm.distributed import ( |
| 13 | get_ep_group, |
| 14 | get_pp_group, |
| 15 | get_tensor_model_parallel_rank, |
| 16 | get_tensor_model_parallel_world_size, |
| 17 | ) |
| 18 | from vllm.forward_context import get_forward_context |
| 19 | from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp |
| 20 | from vllm.model_executor.layers.fused_moe import ( |
| 21 | FusedMoEFactory, |
| 22 | GateLinear, |
| 23 | fused_moe_make_expert_params_mapping, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 26 | fused_topk_bias, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 29 | from vllm.model_executor.layers.linear import ( |
| 30 | MergedColumnParallelLinear, |
| 31 | RowParallelLinear, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 34 | from vllm.model_executor.layers.mhc import ( |
| 35 | HCHeadOp, |
| 36 | MHCFusedPostPreOp, |
| 37 | MHCPostOp, |
| 38 | MHCPreOp, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 41 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 42 | ParallelLMHead, |
| 43 | VocabParallelEmbedding, |
| 44 | ) |
| 45 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 46 | from vllm.model_executor.models.interfaces import ( |
| 47 | EagleModelMixin, |
| 48 | SupportsEagle3, |
| 49 | SupportsLoRA, |
| 50 | SupportsPP, |
| 51 | ) |
| 52 | from vllm.model_executor.models.utils import ( |
| 53 | AutoWeightsLoader, |
| 54 | PPMissingLayer, |
| 55 | WeightsMapper, |
| 56 | extract_layer_index, |
| 57 | is_pp_missing_parameter, |
| 58 | make_layers, |
| 59 | maybe_prefix, |
| 60 | ) |
| 61 | from vllm.model_executor.utils import set_weight_attrs |
| 62 | from vllm.models.deepseek_v4.xpu.xpu_sparse import DeepseekV4XPUAttention |
| 63 | from vllm.platforms import current_platform |
| 64 | from vllm.sequence import IntermediateTensors |
| 65 | from vllm.triton_utils import tl, triton |
| 66 | from vllm.utils.torch_utils import direct_register_custom_op |
| 67 | |
| 68 | |
| 69 | class DeepseekV4MLP(nn.Module): |
| 70 | def __init__( |
| 71 | self, |
| 72 | hidden_size: int, |
| 73 | intermediate_size: int, |
| 74 | hidden_act: str, |
| 75 | swiglu_limit: float | None = None, |
| 76 | quant_config: QuantizationConfig | None = None, |
| 77 | reduce_results: bool = True, |
| 78 | is_sequence_parallel: bool = False, |
| 79 | prefix: str = "", |
| 80 | ) -> None: |
| 81 | super().__init__() |
| 82 | |
| 83 | # If is_sequence_parallel, the input and output tensors are sharded |
| 84 | # across the ranks within the tp_group. In this case the weights are |
| 85 | # replicated and no collective ops are needed. |
| 86 | # Otherwise we use standard TP with an allreduce at the end. |
| 87 | self.gate_up_proj = MergedColumnParallelLinear( |
| 88 | hidden_size, |
| 89 | [intermediate_size] * 2, |
| 90 | bias=False, |
| 91 | quant_config=quant_config, |
| 92 | disable_tp=is_sequence_parallel, |
| 93 | prefix=f"{prefix}.gate_up_proj", |
| 94 | ) |
| 95 | self.down_proj = RowParallelLinear( |
| 96 | intermediate_size, |
| 97 | hidden_size, |
| 98 | bias=False, |
| 99 | quant_config=quant_config, |
| 100 | reduce_results=reduce_results, |
| 101 | disable_tp=is_sequence_parallel, |
| 102 | prefix=f"{prefix}.down_proj", |
| 103 | ) |
| 104 | if hidden_act != "silu": |
| 105 | raise ValueError( |
| 106 | f"Unsupported activation: {hidden_act}. Only silu is supported for now." |
| 107 | ) |
| 108 | if swiglu_limit is not None: |
| 109 | self.act_fn = SiluAndMulWithClamp(swiglu_limit) |
| 110 | else: |
| 111 | self.act_fn = SiluAndMul() |
| 112 | |
| 113 | def forward(self, x): |
| 114 | gate_up, _ = self.gate_up_proj(x) |
| 115 | x = self.act_fn(gate_up) |
| 116 | x, _ = self.down_proj(x) |
| 117 | return x |
| 118 | |
| 119 | |
| 120 | @triton.jit |
| 121 | def _deepseek_v4_stage_mega_moe_inputs_kernel( |
| 122 | hidden_states, |
| 123 | x_fp8, |
| 124 | x_sf, |
| 125 | topk_ids, |
| 126 | topk_weights, |
| 127 | topk_idx_out, |
| 128 | topk_weights_out, |
| 129 | hidden_stride_m: tl.constexpr, |
| 130 | hidden_stride_k: tl.constexpr, |
| 131 | x_stride_m: tl.constexpr, |
| 132 | x_stride_k: tl.constexpr, |
| 133 | x_sf_stride_m: tl.constexpr, |
| 134 | x_sf_stride_k: tl.constexpr, |
| 135 | topk_ids_stride_m: tl.constexpr, |
| 136 | topk_ids_stride_k: tl.constexpr, |
| 137 | topk_weights_stride_m: tl.constexpr, |
| 138 | topk_weights_stride_k: tl.constexpr, |
| 139 | topk_idx_stride_m: tl.constexpr, |
| 140 | topk_idx_stride_k: tl.constexpr, |
| 141 | topk_weights_out_stride_m: tl.constexpr, |
| 142 | topk_weights_out_stride_k: tl.constexpr, |
| 143 | hidden_size: tl.constexpr, |
| 144 | top_k: tl.constexpr, |
| 145 | BLOCK_K: tl.constexpr, |
| 146 | GROUP_K: tl.constexpr, |
| 147 | BLOCK_TOPK: tl.constexpr, |
| 148 | ) -> None: |
| 149 | token_id = tl.program_id(0) |
| 150 | k_block_id = tl.program_id(1) |
| 151 | |
| 152 | k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) |
| 153 | k_mask = k_offsets < hidden_size |
| 154 | hidden = tl.load( |
| 155 | hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k, |
| 156 | mask=k_mask, |
| 157 | other=0.0, |
| 158 | ).to(tl.float32) |
| 159 | |
| 160 | num_groups: tl.constexpr = BLOCK_K // GROUP_K |
| 161 | hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K]) |
| 162 | amax = tl.max(hidden_groups, axis=1) |
| 163 | amax = tl.maximum(amax, 1.0e-4) |
| 164 | |
| 165 | scale = amax / 448.0 |
| 166 | scale_bits = scale.to(tl.uint32, bitcast=True) |
| 167 | scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( |
| 168 | tl.uint32 |
| 169 | ) |
| 170 | scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) |
| 171 | rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True) |
| 172 | |
| 173 | hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K]) |
| 174 | scaled = hidden_groups * (1.0 / rounded_scale)[:, None] |
| 175 | scaled = tl.reshape(scaled, [BLOCK_K]) |
| 176 | fp8 = scaled.to(tl.float8e4nv) |
| 177 | tl.store( |
| 178 | x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k, |
| 179 | fp8, |
| 180 | mask=k_mask, |
| 181 | ) |
| 182 | |
| 183 | scale_offsets = tl.arange(0, num_groups) |
| 184 | packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32) |
| 185 | tl.store( |
| 186 | x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, |
| 187 | packed_scale, |
| 188 | ) |
| 189 | |
| 190 | if k_block_id == 0: |
| 191 | topk_offsets = tl.arange(0, BLOCK_TOPK) |
| 192 | topk_mask = topk_offsets < top_k |
| 193 | |
| 194 | ids = tl.load( |
| 195 | topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k, |
| 196 | mask=topk_mask, |
| 197 | other=0, |
| 198 | ).to(tl.int64) |
| 199 | tl.store( |
| 200 | topk_idx_out |
| 201 | + token_id * topk_idx_stride_m |
| 202 | + topk_offsets * topk_idx_stride_k, |
| 203 | ids, |
| 204 | mask=topk_mask, |
| 205 | ) |
| 206 | |
| 207 | weights = tl.load( |
| 208 | topk_weights |
| 209 | + token_id * topk_weights_stride_m |
| 210 | + topk_offsets * topk_weights_stride_k, |
| 211 | mask=topk_mask, |
| 212 | other=0.0, |
| 213 | ) |
| 214 | tl.store( |
| 215 | topk_weights_out |
| 216 | + token_id * topk_weights_out_stride_m |
| 217 | + topk_offsets * topk_weights_out_stride_k, |
| 218 | weights, |
| 219 | mask=topk_mask, |
| 220 | ) |
| 221 | |
| 222 | |
| 223 | def _stage_deepseek_v4_mega_moe_inputs( |
| 224 | hidden_states: torch.Tensor, |
| 225 | topk_weights: torch.Tensor, |
| 226 | topk_ids: torch.Tensor, |
| 227 | x_fp8: torch.Tensor, |
| 228 | x_sf: torch.Tensor, |
| 229 | topk_idx_out: torch.Tensor, |
| 230 | topk_weights_out: torch.Tensor, |
| 231 | ) -> None: |
| 232 | num_tokens, hidden_size = hidden_states.shape |
| 233 | if num_tokens == 0: |
| 234 | return |
| 235 | if hidden_size % 128 != 0: |
| 236 | raise ValueError( |
| 237 | "DeepSeek V4 MegaMoE input staging requires hidden_size to be " |
| 238 | "a multiple of 128." |
| 239 | ) |
| 240 | top_k = topk_ids.shape[1] |
| 241 | if topk_weights.shape != topk_ids.shape: |
| 242 | raise ValueError( |
| 243 | "DeepSeek V4 MegaMoE input staging requires topk_weights and " |
| 244 | "topk_ids to have the same shape." |
| 245 | ) |
| 246 | |
| 247 | block_k = 128 |
| 248 | grid = (num_tokens, triton.cdiv(hidden_size, block_k)) |
| 249 | block_topk = triton.next_power_of_2(top_k) |
| 250 | _deepseek_v4_stage_mega_moe_inputs_kernel[grid]( |
| 251 | hidden_states, |
| 252 | x_fp8, |
| 253 | x_sf, |
| 254 | topk_ids, |
| 255 | topk_weights, |
| 256 | topk_idx_out, |
| 257 | topk_weights_out, |
| 258 | hidden_states.stride(0), |
| 259 | hidden_states.stride(1), |
| 260 | x_fp8.stride(0), |
| 261 | x_fp8.stride(1), |
| 262 | x_sf.stride(0), |
| 263 | x_sf.stride(1), |
| 264 | topk_ids.stride(0), |
| 265 | topk_ids.stride(1), |
| 266 | topk_weights.stride(0), |
| 267 | topk_weights.stride(1), |
| 268 | topk_idx_out.stride(0), |
| 269 | topk_idx_out.stride(1), |
| 270 | topk_weights_out.stride(0), |
| 271 | topk_weights_out.stride(1), |
| 272 | hidden_size, |
| 273 | top_k, |
| 274 | BLOCK_K=block_k, |
| 275 | GROUP_K=32, |
| 276 | BLOCK_TOPK=block_topk, |
| 277 | num_warps=4, |
| 278 | ) |
| 279 | |
| 280 | |
| 281 | def make_deepseek_v4_expert_params_mapping( |
| 282 | num_experts: int, |
| 283 | ) -> list[tuple[str, str, int, str]]: |
| 284 | return [ |
| 285 | ( |
| 286 | "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_", |
| 287 | f"experts.{expert_id}.{weight_name}.", |
| 288 | expert_id, |
| 289 | shard_id, |
| 290 | ) |
| 291 | for expert_id in range(num_experts) |
| 292 | for shard_id, weight_name in [ |
| 293 | ("w1", "w1"), |
| 294 | ("w2", "w2"), |
| 295 | ("w3", "w3"), |
| 296 | ] |
| 297 | ] |
| 298 | |
| 299 | |
| 300 | class DeepseekV4MegaMoEExperts(nn.Module): |
| 301 | _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} |
| 302 | |
| 303 | def __init__( |
| 304 | self, |
| 305 | vllm_config: VllmConfig, |
| 306 | *, |
| 307 | num_experts: int, |
| 308 | num_local_experts: int, |
| 309 | experts_start_idx: int, |
| 310 | top_k: int, |
| 311 | hidden_size: int, |
| 312 | intermediate_size: int, |
| 313 | prefix: str = "", |
| 314 | ): |
| 315 | super().__init__() |
| 316 | self.prefix = prefix |
| 317 | self.num_experts = num_experts |
| 318 | self.num_local_experts = num_local_experts |
| 319 | self.experts_start_idx = experts_start_idx |
| 320 | self.experts_end_idx = experts_start_idx + num_local_experts |
| 321 | self.top_k = top_k |
| 322 | self.hidden_size = hidden_size |
| 323 | self.intermediate_size = intermediate_size |
| 324 | self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 325 | |
| 326 | weight_attrs = {"weight_loader": self.weight_loader} |
| 327 | self.w13_weight = nn.Parameter( |
| 328 | torch.zeros( |
| 329 | num_local_experts, |
| 330 | 2 * intermediate_size, |
| 331 | hidden_size // 2, |
| 332 | dtype=torch.uint8, |
| 333 | ), |
| 334 | requires_grad=False, |
| 335 | ) |
| 336 | set_weight_attrs(self.w13_weight, weight_attrs) |
| 337 | |
| 338 | self.w13_weight_scale = nn.Parameter( |
| 339 | torch.zeros( |
| 340 | num_local_experts, |
| 341 | 2 * intermediate_size, |
| 342 | hidden_size // 32, |
| 343 | dtype=torch.uint8, |
| 344 | ), |
| 345 | requires_grad=False, |
| 346 | ) |
| 347 | set_weight_attrs(self.w13_weight_scale, weight_attrs) |
| 348 | self.w13_weight_scale.quant_method = "block" |
| 349 | |
| 350 | self.w2_weight = nn.Parameter( |
| 351 | torch.zeros( |
| 352 | num_local_experts, |
| 353 | hidden_size, |
| 354 | intermediate_size // 2, |
| 355 | dtype=torch.uint8, |
| 356 | ), |
| 357 | requires_grad=False, |
| 358 | ) |
| 359 | set_weight_attrs(self.w2_weight, weight_attrs) |
| 360 | |
| 361 | self.w2_weight_scale = nn.Parameter( |
| 362 | torch.zeros( |
| 363 | num_local_experts, |
| 364 | hidden_size, |
| 365 | intermediate_size // 32, |
| 366 | dtype=torch.uint8, |
| 367 | ), |
| 368 | requires_grad=False, |
| 369 | ) |
| 370 | set_weight_attrs(self.w2_weight_scale, weight_attrs) |
| 371 | self.w2_weight_scale.quant_method = "block" |
| 372 | |
| 373 | self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None |
| 374 | self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None |
| 375 | |
| 376 | # Register in the static forward context so the custom-op wrapper |
| 377 | # can look up this module by name from within a torch.compile graph. |
| 378 | compilation_config = vllm_config.compilation_config |
| 379 | if prefix in compilation_config.static_forward_context: |
| 380 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 381 | compilation_config.static_forward_context[prefix] = self |
| 382 | |
| 383 | def _map_global_expert_id(self, expert_id: int) -> int: |
| 384 | if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx: |
| 385 | return -1 |
| 386 | return expert_id - self.experts_start_idx |
| 387 | |
| 388 | def weight_loader( |
| 389 | self, |
| 390 | param: nn.Parameter, |
| 391 | loaded_weight: torch.Tensor, |
| 392 | weight_name: str, |
| 393 | shard_id: str, |
| 394 | expert_id: int, |
| 395 | return_success: bool = False, |
| 396 | ) -> bool | None: |
| 397 | local_expert_id = self._map_global_expert_id(expert_id) |
| 398 | if local_expert_id == -1: |
| 399 | return False if return_success else None |
| 400 | |
| 401 | expert_data = param.data[local_expert_id] |
| 402 | if shard_id in ("w1", "w3"): |
| 403 | if "w13_" not in weight_name: |
| 404 | return False if return_success else None |
| 405 | shard_offset = 0 if shard_id == "w1" else self.intermediate_size |
| 406 | expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size) |
| 407 | elif shard_id == "w2": |
| 408 | if "w2_" not in weight_name: |
| 409 | return False if return_success else None |
| 410 | else: |
| 411 | raise ValueError(f"Unsupported expert shard id: {shard_id}") |
| 412 | |
| 413 | if expert_data.shape != loaded_weight.shape: |
| 414 | raise ValueError( |
| 415 | f"DeepSeek V4 MegaMoE expert weight shape mismatch for " |
| 416 | f"{weight_name}: parameter shard {tuple(expert_data.shape)} " |
| 417 | f"vs checkpoint {tuple(loaded_weight.shape)}" |
| 418 | ) |
| 419 | expert_data.copy_(loaded_weight) |
| 420 | return True if return_success else None |
| 421 | |
| 422 | @staticmethod |
| 423 | def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: |
| 424 | return (sf.to(torch.int32) << 23).view(torch.float32) |
| 425 | |
| 426 | def _check_runtime_supported(self) -> None: |
| 427 | raise NotImplementedError("DeepSeek V4 MegaMoE is not supported on XPU.") |
| 428 | |
| 429 | def finalize_weights(self) -> None: |
| 430 | if self._transformed_l1_weights is not None: |
| 431 | return |
| 432 | |
| 433 | self._check_runtime_supported() |
| 434 | import vllm.third_party.deep_gemm as deep_gemm |
| 435 | |
| 436 | w13_scale = deep_gemm.transform_sf_into_required_layout( |
| 437 | self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), |
| 438 | 2 * self.intermediate_size, |
| 439 | self.hidden_size, |
| 440 | (1, 32), |
| 441 | self.num_local_experts, |
| 442 | ) |
| 443 | w2_scale = deep_gemm.transform_sf_into_required_layout( |
| 444 | self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), |
| 445 | self.hidden_size, |
| 446 | self.intermediate_size, |
| 447 | (1, 32), |
| 448 | self.num_local_experts, |
| 449 | ) |
| 450 | self._transformed_l1_weights, self._transformed_l2_weights = ( |
| 451 | deep_gemm.transform_weights_for_mega_moe( |
| 452 | (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), |
| 453 | (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), |
| 454 | ) |
| 455 | ) |
| 456 | # Drop the original loader-side parameters: the MegaMoE kernels only |
| 457 | # consume the transformed views above. transform_weights_for_mega_moe |
| 458 | # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights) |
| 459 | # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that |
| 460 | # aliases the original storage, and _transformed_l2_weights still holds |
| 461 | # it, so the storage stays live after we drop the Parameter. |
| 462 | self.w13_weight = None |
| 463 | self.w13_weight_scale = None |
| 464 | self.w2_weight = None |
| 465 | self.w2_weight_scale = None |
| 466 | |
| 467 | def get_symm_buffer(self): |
| 468 | import vllm.third_party.deep_gemm as deep_gemm |
| 469 | |
| 470 | group = get_ep_group().device_group |
| 471 | device = torch.accelerator.current_device_index() |
| 472 | key = ( |
| 473 | id(group), |
| 474 | device, |
| 475 | self.num_experts, |
| 476 | self.max_num_tokens, |
| 477 | self.top_k, |
| 478 | self.hidden_size, |
| 479 | self.intermediate_size, |
| 480 | ) |
| 481 | symm_buffer = self._symm_buffer_cache.get(key) |
| 482 | if symm_buffer is None: |
| 483 | symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( |
| 484 | group, |
| 485 | self.num_experts, |
| 486 | self.max_num_tokens, |
| 487 | self.top_k, |
| 488 | self.hidden_size, |
| 489 | self.intermediate_size, |
| 490 | ) |
| 491 | self._symm_buffer_cache[key] = symm_buffer |
| 492 | return symm_buffer |
| 493 | |
| 494 | def forward( |
| 495 | self, |
| 496 | hidden_states: torch.Tensor, |
| 497 | topk_weights: torch.Tensor, |
| 498 | topk_ids: torch.Tensor, |
| 499 | *, |
| 500 | activation_clamp: float | None, |
| 501 | fast_math: bool = True, |
| 502 | ) -> torch.Tensor: |
| 503 | if hidden_states.shape[0] > self.max_num_tokens: |
| 504 | raise ValueError( |
| 505 | f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " |
| 506 | f"but the symmetric buffer was sized for {self.max_num_tokens}." |
| 507 | ) |
| 508 | y = torch.empty_like(hidden_states, dtype=torch.bfloat16) |
| 509 | torch.ops.vllm.deepseek_v4_mega_moe_experts( |
| 510 | hidden_states, |
| 511 | topk_weights, |
| 512 | topk_ids, |
| 513 | y, |
| 514 | self.prefix, |
| 515 | activation_clamp, |
| 516 | fast_math, |
| 517 | ) |
| 518 | return y |
| 519 | |
| 520 | def _run_mega_moe( |
| 521 | self, |
| 522 | hidden_states: torch.Tensor, |
| 523 | topk_weights: torch.Tensor, |
| 524 | topk_ids: torch.Tensor, |
| 525 | y: torch.Tensor, |
| 526 | activation_clamp: float | None, |
| 527 | fast_math: bool, |
| 528 | ) -> None: |
| 529 | import vllm.third_party.deep_gemm as deep_gemm |
| 530 | |
| 531 | symm_buffer = self.get_symm_buffer() |
| 532 | num_tokens = hidden_states.shape[0] |
| 533 | _stage_deepseek_v4_mega_moe_inputs( |
| 534 | hidden_states, |
| 535 | topk_weights, |
| 536 | topk_ids, |
| 537 | symm_buffer.x[:num_tokens], |
| 538 | symm_buffer.x_sf[:num_tokens], |
| 539 | symm_buffer.topk_idx[:num_tokens], |
| 540 | symm_buffer.topk_weights[:num_tokens], |
| 541 | ) |
| 542 | |
| 543 | # This method must have been already called during the weight loading phase. |
| 544 | # We call it again here to cover the dummy weight loading case. |
| 545 | self.finalize_weights() |
| 546 | |
| 547 | assert self._transformed_l1_weights is not None |
| 548 | assert self._transformed_l2_weights is not None |
| 549 | deep_gemm.fp8_fp4_mega_moe( |
| 550 | y, |
| 551 | self._transformed_l1_weights, |
| 552 | self._transformed_l2_weights, |
| 553 | symm_buffer, |
| 554 | activation_clamp=activation_clamp, |
| 555 | fast_math=fast_math, |
| 556 | ) |
| 557 | |
| 558 | |
| 559 | DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] |
| 560 | |
| 561 | |
| 562 | def _deepseek_v4_mega_moe_experts_op( |
| 563 | hidden_states: torch.Tensor, |
| 564 | topk_weights: torch.Tensor, |
| 565 | topk_ids: torch.Tensor, |
| 566 | out: torch.Tensor, |
| 567 | layer_name: str, |
| 568 | activation_clamp: float | None, |
| 569 | fast_math: bool, |
| 570 | ) -> None: |
| 571 | self = get_forward_context().no_compile_layers[layer_name] |
| 572 | self._run_mega_moe( |
| 573 | hidden_states, |
| 574 | topk_weights, |
| 575 | topk_ids, |
| 576 | out, |
| 577 | activation_clamp, |
| 578 | fast_math, |
| 579 | ) |
| 580 | |
| 581 | |
| 582 | direct_register_custom_op( |
| 583 | op_name="deepseek_v4_mega_moe_experts", |
| 584 | op_func=_deepseek_v4_mega_moe_experts_op, |
| 585 | mutates_args=["out"], |
| 586 | ) |
| 587 | |
| 588 | |
| 589 | class DeepseekV4MoE(nn.Module): |
| 590 | def __init__( |
| 591 | self, |
| 592 | vllm_config: VllmConfig, |
| 593 | prefix: str = "", |
| 594 | ): |
| 595 | super().__init__() |
| 596 | |
| 597 | self.tp_size = get_tensor_model_parallel_world_size() |
| 598 | config = vllm_config.model_config.hf_config |
| 599 | quant_config = vllm_config.quant_config |
| 600 | self.prefix = prefix |
| 601 | self.use_mega_moe = ( |
| 602 | vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" |
| 603 | ) |
| 604 | if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: |
| 605 | raise NotImplementedError( |
| 606 | "DeepSeek V4 MegaMoE currently requires expert parallel. " |
| 607 | "Enable it with --enable-expert-parallel, or pick a different " |
| 608 | "moe backend." |
| 609 | ) |
| 610 | |
| 611 | self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) |
| 612 | self.hidden_size = config.hidden_size |
| 613 | |
| 614 | self.n_routed_experts = config.n_routed_experts |
| 615 | self.n_activated_experts = config.num_experts_per_tok |
| 616 | self.moe_intermediate_size = config.moe_intermediate_size |
| 617 | self.swiglu_limit = config.swiglu_limit |
| 618 | self.renormalize = config.norm_topk_prob |
| 619 | self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") |
| 620 | if self.use_mega_moe and self.scoring_func != "sqrtsoftplus": |
| 621 | raise NotImplementedError( |
| 622 | "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." |
| 623 | ) |
| 624 | if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": |
| 625 | raise NotImplementedError( |
| 626 | "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" |
| 627 | f"{config.expert_dtype!r}. Drop --kernel-config moe_backend=" |
| 628 | "deep_gemm_mega_moe for this checkpoint." |
| 629 | ) |
| 630 | |
| 631 | self.gate = GateLinear( |
| 632 | input_size=config.hidden_size, |
| 633 | output_size=config.n_routed_experts, |
| 634 | bias=False, |
| 635 | out_dtype=torch.float32, |
| 636 | prefix=f"{prefix}.gate", |
| 637 | ) |
| 638 | |
| 639 | self.gate.e_score_correction_bias = None |
| 640 | self.gate.tid2eid = None |
| 641 | is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers |
| 642 | self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 |
| 643 | if is_hash_moe: |
| 644 | # hash MoE doesn't use e_score_correction_bias |
| 645 | # Use randint instead of empty to avoid garbage values causing |
| 646 | # invalid memory access in dummy mode (--load-format="dummy") |
| 647 | self.gate.tid2eid = nn.Parameter( |
| 648 | torch.randint( |
| 649 | 0, |
| 650 | config.n_routed_experts, |
| 651 | (config.vocab_size, config.num_experts_per_tok), |
| 652 | dtype=self.hash_indices_dtype, |
| 653 | ), |
| 654 | requires_grad=False, |
| 655 | ) |
| 656 | elif getattr(config, "topk_method", None) == "noaux_tc": |
| 657 | self.gate.e_score_correction_bias = nn.Parameter( |
| 658 | torch.empty(config.n_routed_experts, dtype=torch.float32), |
| 659 | requires_grad=False, |
| 660 | ) |
| 661 | |
| 662 | if config.n_shared_experts is None: |
| 663 | self.shared_experts = None |
| 664 | else: |
| 665 | intermediate_size = config.moe_intermediate_size * config.n_shared_experts |
| 666 | |
| 667 | self.shared_experts = DeepseekV4MLP( |
| 668 | hidden_size=config.hidden_size, |
| 669 | intermediate_size=intermediate_size, |
| 670 | hidden_act=config.hidden_act, |
| 671 | swiglu_limit=self.swiglu_limit, |
| 672 | quant_config=quant_config, |
| 673 | reduce_results=self.use_mega_moe, |
| 674 | prefix=f"{prefix}.shared_experts", |
| 675 | ) |
| 676 | |
| 677 | if self.use_mega_moe: |
| 678 | self._init_mega_moe_experts(vllm_config, config, prefix) |
| 679 | else: |
| 680 | self._init_fused_moe_experts(config, quant_config, prefix) |
| 681 | |
| 682 | def _init_mega_moe_experts( |
| 683 | self, |
| 684 | vllm_config: VllmConfig, |
| 685 | config, |
| 686 | prefix: str, |
| 687 | ) -> None: |
| 688 | self.ep_group = get_ep_group() |
| 689 | self.ep_size = self.ep_group.world_size |
| 690 | self.ep_rank = self.ep_group.rank_in_group |
| 691 | assert config.n_routed_experts % self.ep_size == 0 |
| 692 | |
| 693 | self.n_local_experts = config.n_routed_experts // self.ep_size |
| 694 | self.experts_start_idx = self.ep_rank * self.n_local_experts |
| 695 | self.experts_end_idx = self.experts_start_idx + self.n_local_experts |
| 696 | |
| 697 | self.experts = DeepseekV4MegaMoEExperts( |
| 698 | vllm_config, |
| 699 | num_experts=config.n_routed_experts, |
| 700 | num_local_experts=self.n_local_experts, |
| 701 | experts_start_idx=self.experts_start_idx, |
| 702 | top_k=config.num_experts_per_tok, |
| 703 | hidden_size=config.hidden_size, |
| 704 | intermediate_size=config.moe_intermediate_size, |
| 705 | prefix=f"{prefix}.experts", |
| 706 | ) |
| 707 | |
| 708 | def _init_fused_moe_experts( |
| 709 | self, |
| 710 | config, |
| 711 | quant_config, |
| 712 | prefix: str, |
| 713 | ) -> None: |
| 714 | self.tp_rank = get_tensor_model_parallel_rank() |
| 715 | assert config.n_routed_experts % self.tp_size == 0 |
| 716 | |
| 717 | self.n_local_experts = config.n_routed_experts // self.tp_size |
| 718 | self.experts_start_idx = self.tp_rank * self.n_local_experts |
| 719 | self.experts_end_idx = self.experts_start_idx + self.n_local_experts |
| 720 | self.experts = FusedMoEFactory( |
| 721 | shared_experts=self.shared_experts, |
| 722 | gate=self.gate, |
| 723 | num_experts=config.n_routed_experts, |
| 724 | top_k=config.num_experts_per_tok, |
| 725 | hidden_size=config.hidden_size, |
| 726 | intermediate_size=config.moe_intermediate_size, |
| 727 | renormalize=config.norm_topk_prob, |
| 728 | quant_config=quant_config, |
| 729 | prefix=f"{prefix}.experts", |
| 730 | scoring_func=self.scoring_func, |
| 731 | routed_scaling_factor=self.routed_scaling_factor, |
| 732 | e_score_correction_bias=self.gate.e_score_correction_bias, |
| 733 | hash_indices_table=self.gate.tid2eid, |
| 734 | swiglu_limit=self.swiglu_limit, |
| 735 | router_logits_dtype=torch.float32, |
| 736 | ) |
| 737 | |
| 738 | def forward( |
| 739 | self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None |
| 740 | ) -> torch.Tensor: |
| 741 | if self.gate.tid2eid is not None and input_ids is None: |
| 742 | raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") |
| 743 | |
| 744 | if not self.use_mega_moe: |
| 745 | return self._forward_fused_moe(hidden_states, input_ids) |
| 746 | |
| 747 | org_shape = hidden_states.shape |
| 748 | router_logits, _ = self.gate(hidden_states) |
| 749 | topk_weights, topk_ids = fused_topk_bias( |
| 750 | hidden_states=hidden_states, |
| 751 | gating_output=router_logits, |
| 752 | scoring_func=self.scoring_func, |
| 753 | e_score_correction_bias=self.gate.e_score_correction_bias.data |
| 754 | if self.gate.e_score_correction_bias is not None |
| 755 | else None, |
| 756 | topk=self.n_activated_experts, |
| 757 | renormalize=self.renormalize, |
| 758 | indices_type=self.hash_indices_dtype, |
| 759 | input_tokens=input_ids, |
| 760 | hash_indices_table=self.gate.tid2eid, |
| 761 | routed_scaling_factor=self.routed_scaling_factor, |
| 762 | ) |
| 763 | activation_clamp = ( |
| 764 | float(self.swiglu_limit) if self.swiglu_limit is not None else None |
| 765 | ) |
| 766 | final_hidden_states = self.experts( |
| 767 | hidden_states, |
| 768 | topk_weights, |
| 769 | topk_ids, |
| 770 | activation_clamp=activation_clamp, |
| 771 | ) |
| 772 | |
| 773 | if self.shared_experts is not None: |
| 774 | shared_output = self.shared_experts(hidden_states) |
| 775 | final_hidden_states += shared_output |
| 776 | |
| 777 | return final_hidden_states.view(org_shape) |
| 778 | |
| 779 | def _forward_fused_moe( |
| 780 | self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None |
| 781 | ) -> torch.Tensor: |
| 782 | org_shape = hidden_states.shape |
| 783 | final_hidden_states = self.experts( |
| 784 | hidden_states=hidden_states, |
| 785 | router_logits=hidden_states, |
| 786 | input_ids=input_ids, |
| 787 | ) |
| 788 | |
| 789 | return final_hidden_states.view(org_shape) |
| 790 | |
| 791 | def finalize_mega_moe_weights(self) -> None: |
| 792 | if self.use_mega_moe: |
| 793 | self.experts.finalize_weights() |
| 794 | |
| 795 | |
| 796 | class DeepseekV4DecoderLayer(nn.Module): |
| 797 | def __init__( |
| 798 | self, |
| 799 | vllm_config, |
| 800 | prefix, |
| 801 | topk_indices_buffer: torch.Tensor | None = None, |
| 802 | aux_stream_list: list | None = None, |
| 803 | ): |
| 804 | super().__init__() |
| 805 | |
| 806 | # Lazy import to avoid top-level tilelang dependency. |
| 807 | # Registers both torch.ops.vllm.mhc_pre and mhc_post |
| 808 | import vllm.model_executor.layers.mhc # noqa: F401 |
| 809 | |
| 810 | config = vllm_config.model_config.hf_config |
| 811 | self.hidden_size = config.hidden_size |
| 812 | |
| 813 | self.rms_norm_eps = config.rms_norm_eps |
| 814 | self.attn = DeepseekV4XPUAttention( |
| 815 | vllm_config, |
| 816 | prefix=f"{prefix}.attn", |
| 817 | topk_indices_buffer=topk_indices_buffer, |
| 818 | aux_stream_list=aux_stream_list, |
| 819 | ) |
| 820 | self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") |
| 821 | |
| 822 | self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) |
| 823 | self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) |
| 824 | self.hc_mult = config.hc_mult |
| 825 | self.hc_sinkhorn_iters = config.hc_sinkhorn_iters |
| 826 | self.hc_eps = config.hc_eps |
| 827 | self.hc_post_alpha = 2.0 |
| 828 | mix_hc = (2 + self.hc_mult) * self.hc_mult |
| 829 | hc_dim = self.hc_mult * self.hidden_size |
| 830 | self.hc_attn_fn = nn.Parameter( |
| 831 | torch.empty( |
| 832 | (mix_hc, hc_dim), |
| 833 | dtype=torch.float32, |
| 834 | ), |
| 835 | requires_grad=False, |
| 836 | ) |
| 837 | self.hc_ffn_fn = nn.Parameter( |
| 838 | torch.empty( |
| 839 | (mix_hc, hc_dim), |
| 840 | dtype=torch.float32, |
| 841 | ), |
| 842 | requires_grad=False, |
| 843 | ) |
| 844 | self.hc_attn_base = nn.Parameter( |
| 845 | torch.empty( |
| 846 | mix_hc, |
| 847 | dtype=torch.float32, |
| 848 | ), |
| 849 | requires_grad=False, |
| 850 | ) |
| 851 | self.hc_ffn_base = nn.Parameter( |
| 852 | torch.empty( |
| 853 | mix_hc, |
| 854 | dtype=torch.float32, |
| 855 | ), |
| 856 | requires_grad=False, |
| 857 | ) |
| 858 | self.hc_attn_scale = nn.Parameter( |
| 859 | torch.empty( |
| 860 | 3, |
| 861 | dtype=torch.float32, |
| 862 | ), |
| 863 | requires_grad=False, |
| 864 | ) |
| 865 | self.hc_ffn_scale = nn.Parameter( |
| 866 | torch.empty( |
| 867 | 3, |
| 868 | dtype=torch.float32, |
| 869 | ), |
| 870 | requires_grad=False, |
| 871 | ) |
| 872 | self.mhc_pre = MHCPreOp() |
| 873 | self.mhc_post = MHCPostOp() |
| 874 | self.mhc_fused_post_pre = MHCFusedPostPreOp() |
| 875 | |
| 876 | def hc_pre( |
| 877 | self, |
| 878 | x: torch.Tensor, |
| 879 | hc_fn: torch.Tensor, |
| 880 | hc_scale: torch.Tensor, |
| 881 | hc_base: torch.Tensor, |
| 882 | ): |
| 883 | post_mix, res_mix, layer_input = self.mhc_pre( |
| 884 | residual=x, |
| 885 | fn=hc_fn, |
| 886 | hc_scale=hc_scale, |
| 887 | hc_base=hc_base, |
| 888 | rms_eps=self.rms_norm_eps, |
| 889 | hc_pre_eps=self.hc_eps, |
| 890 | hc_sinkhorn_eps=self.hc_eps, |
| 891 | hc_post_mult_value=self.hc_post_alpha, |
| 892 | sinkhorn_repeat=self.hc_sinkhorn_iters, |
| 893 | ) |
| 894 | return layer_input, post_mix, res_mix |
| 895 | |
| 896 | def hc_post( |
| 897 | self, |
| 898 | x: torch.Tensor, |
| 899 | residual: torch.Tensor, |
| 900 | post: torch.Tensor, |
| 901 | comb: torch.Tensor, |
| 902 | ): |
| 903 | return self.mhc_post(x, residual, post, comb) |
| 904 | |
| 905 | def forward( |
| 906 | self, |
| 907 | x: torch.Tensor, |
| 908 | positions: torch.Tensor, |
| 909 | input_ids: torch.Tensor | None, |
| 910 | post_mix: torch.Tensor | None = None, |
| 911 | res_mix: torch.Tensor | None = None, |
| 912 | residual: torch.Tensor | None = None, |
| 913 | ) -> tuple[ |
| 914 | torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None |
| 915 | ]: |
| 916 | if residual is None: |
| 917 | # First layer: run standalone hc_pre |
| 918 | residual = x |
| 919 | x, post_mix, res_mix = self.hc_pre( |
| 920 | x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base |
| 921 | ) |
| 922 | else: |
| 923 | residual, post_mix, res_mix, x = self.mhc_fused_post_pre( |
| 924 | x, |
| 925 | residual, |
| 926 | post_mix, |
| 927 | res_mix, |
| 928 | self.hc_attn_fn, |
| 929 | self.hc_attn_scale, |
| 930 | self.hc_attn_base, |
| 931 | self.rms_norm_eps, |
| 932 | self.hc_eps, |
| 933 | self.hc_eps, |
| 934 | self.hc_post_alpha, |
| 935 | self.hc_sinkhorn_iters, |
| 936 | ) |
| 937 | |
| 938 | x = self.attn_norm(x) |
| 939 | x = self.attn(positions, x, None) |
| 940 | |
| 941 | residual, post_mix, res_mix, x = self.mhc_fused_post_pre( |
| 942 | x, |
| 943 | residual, |
| 944 | post_mix, |
| 945 | res_mix, |
| 946 | self.hc_ffn_fn, |
| 947 | self.hc_ffn_scale, |
| 948 | self.hc_ffn_base, |
| 949 | self.rms_norm_eps, |
| 950 | self.hc_eps, |
| 951 | self.hc_eps, |
| 952 | self.hc_post_alpha, |
| 953 | self.hc_sinkhorn_iters, |
| 954 | ) |
| 955 | x = self.ffn_norm(x) |
| 956 | x = self.ffn(x, input_ids) |
| 957 | return x, residual, post_mix, res_mix |
| 958 | |
| 959 | |
| 960 | class DeepseekV4Model(nn.Module, EagleModelMixin): |
| 961 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 962 | super().__init__() |
| 963 | |
| 964 | config = vllm_config.model_config.hf_config |
| 965 | quant_config = vllm_config.quant_config |
| 966 | self.config = config |
| 967 | self.use_mega_moe = ( |
| 968 | vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" |
| 969 | ) |
| 970 | if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: |
| 971 | raise NotImplementedError( |
| 972 | "DeepSeek V4 MegaMoE currently requires expert parallel. " |
| 973 | "Enable it with --enable-expert-parallel, or pick a different " |
| 974 | "moe backend." |
| 975 | ) |
| 976 | self.vocab_size = config.vocab_size |
| 977 | self.hc_eps = config.hc_eps |
| 978 | self.hc_mult = config.hc_mult |
| 979 | self.hc_dim = self.hc_mult * config.hidden_size |
| 980 | self.rms_norm_eps = config.rms_norm_eps |
| 981 | |
| 982 | # Disable aux streams on XPU — no multi-stream overlap support. |
| 983 | aux_stream_list = None |
| 984 | |
| 985 | self.device = current_platform.device_type |
| 986 | # Reserved topk indices buffer for all Indexer layers to reuse. |
| 987 | self.topk_indices_buffer = torch.empty( |
| 988 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 989 | config.index_topk, |
| 990 | dtype=torch.int32, |
| 991 | device=self.device, |
| 992 | ) |
| 993 | |
| 994 | if get_pp_group().is_first_rank: |
| 995 | self.embed_tokens = VocabParallelEmbedding( |
| 996 | config.vocab_size, |
| 997 | config.hidden_size, |
| 998 | quant_config=quant_config, |
| 999 | prefix=f"{prefix}.embed_tokens", |
| 1000 | ) |
| 1001 | else: |
| 1002 | self.embed_tokens = PPMissingLayer() |
| 1003 | |
| 1004 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 1005 | config.num_hidden_layers, |
| 1006 | lambda prefix: DeepseekV4DecoderLayer( |
| 1007 | vllm_config, |
| 1008 | prefix=prefix, |
| 1009 | topk_indices_buffer=self.topk_indices_buffer, |
| 1010 | aux_stream_list=aux_stream_list, |
| 1011 | ), |
| 1012 | prefix=f"{prefix}.layers", |
| 1013 | ) |
| 1014 | |
| 1015 | if get_pp_group().is_last_rank: |
| 1016 | self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) |
| 1017 | else: |
| 1018 | self.norm = PPMissingLayer() |
| 1019 | |
| 1020 | self.hc_head_fn = nn.Parameter( |
| 1021 | torch.empty( |
| 1022 | self.hc_mult, |
| 1023 | self.hc_dim, |
| 1024 | dtype=torch.float32, |
| 1025 | ), |
| 1026 | requires_grad=False, |
| 1027 | ) |
| 1028 | self.hc_head_base = nn.Parameter( |
| 1029 | torch.empty( |
| 1030 | self.hc_mult, |
| 1031 | dtype=torch.float32, |
| 1032 | ), |
| 1033 | requires_grad=False, |
| 1034 | ) |
| 1035 | self.hc_head_scale = nn.Parameter( |
| 1036 | torch.empty(1, dtype=torch.float32), |
| 1037 | requires_grad=False, |
| 1038 | ) |
| 1039 | self.hc_head_op = HCHeadOp() |
| 1040 | spec_config = vllm_config.speculative_config |
| 1041 | needs_mtp_hidden_states = spec_config is not None and ( |
| 1042 | spec_config.use_eagle() or spec_config.uses_draft_model() |
| 1043 | ) |
| 1044 | if get_pp_group().is_last_rank and needs_mtp_hidden_states: |
| 1045 | self._mtp_hidden_buffer = torch.empty( |
| 1046 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 1047 | self.hc_dim, |
| 1048 | dtype=vllm_config.model_config.dtype, |
| 1049 | device=self.device, |
| 1050 | ) |
| 1051 | else: |
| 1052 | self._mtp_hidden_buffer = None |
| 1053 | |
| 1054 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 1055 | return self.embed_tokens(input_ids) |
| 1056 | |
| 1057 | def make_empty_intermediate_tensors( |
| 1058 | self, |
| 1059 | batch_size: int, |
| 1060 | dtype: torch.dtype, |
| 1061 | device: torch.device, |
| 1062 | ) -> IntermediateTensors: |
| 1063 | # PP intermediate tensors carry the multi-stream hidden_states |
| 1064 | # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the |
| 1065 | # token embedding to hc_mult streams before the first decoder |
| 1066 | # layer and keeps that shape until hc_head() collapses it. |
| 1067 | return IntermediateTensors( |
| 1068 | { |
| 1069 | "hidden_states": torch.zeros( |
| 1070 | (batch_size, self.hc_mult, self.config.hidden_size), |
| 1071 | dtype=dtype, |
| 1072 | device=device, |
| 1073 | ), |
| 1074 | } |
| 1075 | ) |
| 1076 | |
| 1077 | def forward( |
| 1078 | self, |
| 1079 | input_ids: torch.Tensor, |
| 1080 | positions: torch.Tensor, |
| 1081 | intermediate_tensors: IntermediateTensors | None, |
| 1082 | inputs_embeds: torch.Tensor | None = None, |
| 1083 | ) -> torch.Tensor | IntermediateTensors: |
| 1084 | if get_pp_group().is_first_rank: |
| 1085 | if inputs_embeds is not None: |
| 1086 | hidden_states = inputs_embeds |
| 1087 | else: |
| 1088 | hidden_states = self.embed_input_ids(input_ids) |
| 1089 | hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) |
| 1090 | else: |
| 1091 | assert intermediate_tensors is not None |
| 1092 | hidden_states = intermediate_tensors["hidden_states"] |
| 1093 | |
| 1094 | if self.use_mega_moe: |
| 1095 | input_ids = input_ids.to(torch.int64) |
| 1096 | |
| 1097 | residual, post_mix, res_mix = None, None, None |
| 1098 | aux_hidden_states: list[torch.Tensor] = [] |
| 1099 | for idx, layer in enumerate( |
| 1100 | islice(self.layers, self.start_layer, self.end_layer), |
| 1101 | start=self.start_layer, |
| 1102 | ): |
| 1103 | hidden_states, residual, post_mix, res_mix = layer( |
| 1104 | hidden_states, |
| 1105 | positions, |
| 1106 | input_ids, |
| 1107 | post_mix, |
| 1108 | res_mix, |
| 1109 | residual, |
| 1110 | ) |
| 1111 | if idx + 1 in self.aux_hidden_state_layers: |
| 1112 | aux_recon = layer.hc_post(hidden_states, residual, post_mix, res_mix) |
| 1113 | aux_hidden_states.append(aux_recon.mean(dim=1)) |
| 1114 | # The fused path defers the final hc_post to the next layer's |
| 1115 | # fused_post_pre. After the last layer we must apply it explicitly. |
| 1116 | if layer is not None: |
| 1117 | hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) |
| 1118 | |
| 1119 | if not get_pp_group().is_last_rank: |
| 1120 | return IntermediateTensors({"hidden_states": hidden_states}) |
| 1121 | |
| 1122 | if self._mtp_hidden_buffer is not None: |
| 1123 | num_tokens = hidden_states.shape[0] |
| 1124 | self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) |
| 1125 | |
| 1126 | hidden_states = self.hc_head_op( |
| 1127 | hidden_states, |
| 1128 | self.hc_head_fn, |
| 1129 | self.hc_head_scale, |
| 1130 | self.hc_head_base, |
| 1131 | self.rms_norm_eps, |
| 1132 | self.hc_eps, |
| 1133 | ) |
| 1134 | hidden_states = self.norm(hidden_states) |
| 1135 | if len(aux_hidden_states) > 0: |
| 1136 | return hidden_states, aux_hidden_states |
| 1137 | return hidden_states |
| 1138 | |
| 1139 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1140 | stacked_params_mapping = [ |
| 1141 | # (param_name, shard_name, shard_id) |
| 1142 | ("gate_up_proj", "w1", 0), |
| 1143 | ("gate_up_proj", "w3", 1), |
| 1144 | ("attn.fused_wqa_wkv", "attn.wq_a", 0), |
| 1145 | ("attn.fused_wqa_wkv", "attn.wkv", 1), |
| 1146 | ("compressor.fused_wkv_wgate", "compressor.wkv", 0), |
| 1147 | ("compressor.fused_wkv_wgate", "compressor.wgate", 1), |
| 1148 | ] |
| 1149 | params_dict = dict(self.named_parameters()) |
| 1150 | loaded_params: set[str] = set() |
| 1151 | |
| 1152 | # TP for attention |
| 1153 | tp_size = get_tensor_model_parallel_world_size() |
| 1154 | tp_rank = get_tensor_model_parallel_rank() |
| 1155 | n_head = self.config.num_attention_heads |
| 1156 | n_local_head = n_head // tp_size |
| 1157 | head_rank_start = n_local_head * tp_rank |
| 1158 | head_rank_end = n_local_head * (tp_rank + 1) |
| 1159 | |
| 1160 | # Pre-compute expert mapping ONCE. |
| 1161 | expert_mapping = self.get_expert_mapping() |
| 1162 | |
| 1163 | for name, loaded_weight in weights: |
| 1164 | for param_name, weight_name, shard_id in stacked_params_mapping: |
| 1165 | # Skip non-stacked layers and experts (experts handled below). |
| 1166 | if ".experts." in name: |
| 1167 | continue |
| 1168 | if weight_name not in name: |
| 1169 | continue |
| 1170 | name = name.replace(weight_name, param_name) |
| 1171 | |
| 1172 | if is_pp_missing_parameter(name, self): |
| 1173 | break |
| 1174 | if name not in params_dict: |
| 1175 | head, _, leaf = name.rpartition(".") |
| 1176 | suffixed = f"{head}.base_layer.{leaf}" |
| 1177 | if suffixed in params_dict: |
| 1178 | name = suffixed |
| 1179 | param = params_dict[name] |
| 1180 | weight_loader = param.weight_loader |
| 1181 | weight_loader(param, loaded_weight, shard_id) |
| 1182 | loaded_params.add(name) |
| 1183 | break |
| 1184 | else: |
| 1185 | if ".experts." in name: |
| 1186 | # E8M0 scales are stored as float8_e8m0fnu in |
| 1187 | # checkpoints but the MoE param is uint8. copy_() |
| 1188 | # would do a numeric conversion (e.g. 2^-7 → 0), |
| 1189 | # destroying the raw exponent bytes. |
| 1190 | if ( |
| 1191 | "weight_scale" in name |
| 1192 | and loaded_weight.dtype == torch.float8_e8m0fnu |
| 1193 | ): |
| 1194 | loaded_weight = loaded_weight.view(torch.uint8) |
| 1195 | for mapping in expert_mapping: |
| 1196 | param_name, weight_name, expert_id, expert_shard_id = mapping |
| 1197 | if weight_name not in name: |
| 1198 | continue |
| 1199 | name_mapped = name.replace(weight_name, param_name) |
| 1200 | if is_pp_missing_parameter(name_mapped, self): |
| 1201 | continue |
| 1202 | param = params_dict[name_mapped] |
| 1203 | # We should ask the weight loader to return success or not |
| 1204 | # here since otherwise we may skip experts with other |
| 1205 | # available replicas. |
| 1206 | weight_loader = typing.cast( |
| 1207 | Callable[..., bool], param.weight_loader |
| 1208 | ) |
| 1209 | success = weight_loader( |
| 1210 | param, |
| 1211 | loaded_weight, |
| 1212 | name_mapped, |
| 1213 | shard_id=expert_shard_id, |
| 1214 | expert_id=expert_id, |
| 1215 | return_success=True, |
| 1216 | ) |
| 1217 | if success: |
| 1218 | name = name_mapped |
| 1219 | break |
| 1220 | loaded_params.add(name_mapped) |
| 1221 | continue |
| 1222 | elif "attn_sink" in name: |
| 1223 | if is_pp_missing_parameter(name, self): |
| 1224 | continue |
| 1225 | narrow_weight = loaded_weight[head_rank_start:head_rank_end] |
| 1226 | n = narrow_weight.shape[0] |
| 1227 | params_dict[name][:n].copy_(narrow_weight) |
| 1228 | loaded_params.add(name) |
| 1229 | continue |
| 1230 | else: |
| 1231 | if is_pp_missing_parameter(name, self): |
| 1232 | continue |
| 1233 | # Non-LoRA params on a LoRA-wrapped module live at |
| 1234 | # ``<head>.base_layer.<leaf>``; the checkpoint is plain. |
| 1235 | if name not in params_dict: |
| 1236 | head, _, leaf = name.rpartition(".") |
| 1237 | suffixed = f"{head}.base_layer.{leaf}" |
| 1238 | if suffixed in params_dict: |
| 1239 | name = suffixed |
| 1240 | param = params_dict[name] |
| 1241 | weight_loader = getattr( |
| 1242 | param, "weight_loader", default_weight_loader |
| 1243 | ) |
| 1244 | weight_loader(param, loaded_weight) |
| 1245 | loaded_params.add(name) |
| 1246 | continue |
| 1247 | |
| 1248 | return loaded_params |
| 1249 | |
| 1250 | def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: |
| 1251 | first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) |
| 1252 | if first_layer.ffn.use_mega_moe: |
| 1253 | return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) |
| 1254 | # Params for weights, fp8 weight scales, fp8 activation scales |
| 1255 | # (param_name, weight_name, expert_id, shard_id) |
| 1256 | return fused_moe_make_expert_params_mapping( |
| 1257 | self, |
| 1258 | ckpt_gate_proj_name="w1", |
| 1259 | ckpt_down_proj_name="w2", |
| 1260 | ckpt_up_proj_name="w3", |
| 1261 | num_experts=self.config.n_routed_experts, |
| 1262 | ) |
| 1263 | |
| 1264 | def finalize_mega_moe_weights(self) -> None: |
| 1265 | for layer in islice(self.layers, self.start_layer, self.end_layer): |
| 1266 | layer.ffn.finalize_mega_moe_weights() |
| 1267 | |
| 1268 | |
| 1269 | def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: |
| 1270 | if expert_dtype == "fp4": |
| 1271 | # MXFP4 experts use Mxfp4MoEMethod, which registers scales as |
| 1272 | # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and |
| 1273 | # shared experts use Fp8LinearMethod's block scales, which |
| 1274 | # register as ``weight_scale_inv``. |
| 1275 | scale_regex = { |
| 1276 | # ``.base_layer.``-namespace variant (LoRA-wrapped experts). |
| 1277 | re.compile( |
| 1278 | r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$" |
| 1279 | ): r"\1.weight_scale", |
| 1280 | re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", |
| 1281 | re.compile(r"\.scale$"): ".weight_scale_inv", |
| 1282 | } |
| 1283 | else: |
| 1284 | # FP8 experts use Fp8MoEMethod (block_quant=True), which registers |
| 1285 | # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys |
| 1286 | # there. |
| 1287 | scale_regex = { |
| 1288 | # ``.base_layer.``-namespace variant of the above. |
| 1289 | re.compile( |
| 1290 | r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$" |
| 1291 | ): r"\1.weight_scale_inv", |
| 1292 | re.compile(r"\.scale$"): ".weight_scale_inv", |
| 1293 | } |
| 1294 | return WeightsMapper( |
| 1295 | orig_to_new_prefix={ |
| 1296 | "layers.": "model.layers.", |
| 1297 | "embed.": "model.embed.", |
| 1298 | "norm.": "model.norm.", |
| 1299 | "hc_head": "model.hc_head", |
| 1300 | "mtp.": "model.mtp.", |
| 1301 | }, |
| 1302 | orig_to_new_regex=scale_regex, |
| 1303 | orig_to_new_suffix={ |
| 1304 | "head.weight": "lm_head.weight", |
| 1305 | "embed.weight": "embed_tokens.weight", |
| 1306 | ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", |
| 1307 | }, |
| 1308 | orig_to_new_substr={ |
| 1309 | ".shared_experts.w2": ".shared_experts.down_proj", |
| 1310 | "mtp.": None, |
| 1311 | }, |
| 1312 | ) |
| 1313 | |
| 1314 | |
| 1315 | class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA): |
| 1316 | model_cls = DeepseekV4Model |
| 1317 | |
| 1318 | # Default mapper assumes the original FP4-expert checkpoint layout. |
| 1319 | # Overridden per-instance in __init__ when expert_dtype != "fp4". |
| 1320 | hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4") |
| 1321 | |
| 1322 | packed_modules_mapping = { |
| 1323 | "gate_up_proj": ["w1", "w3"], |
| 1324 | "fused_wqa_wkv": ["wq_a", "wkv"], |
| 1325 | "fused_wkv_wgate": ["wkv", "wgate"], |
| 1326 | } |
| 1327 | |
| 1328 | # The MTP draft head is not LoRA-adapted. |
| 1329 | lora_skip_prefixes = ["mtp."] |
| 1330 | |
| 1331 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1332 | super().__init__() |
| 1333 | |
| 1334 | config = vllm_config.model_config.hf_config |
| 1335 | self.config = config |
| 1336 | expert_dtype = getattr(config, "expert_dtype", "fp4") |
| 1337 | if expert_dtype != "fp4": |
| 1338 | self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype) |
| 1339 | |
| 1340 | self.model = self.model_cls( |
| 1341 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 1342 | ) |
| 1343 | if get_pp_group().is_last_rank: |
| 1344 | self.lm_head = ParallelLMHead( |
| 1345 | config.vocab_size, |
| 1346 | config.hidden_size, |
| 1347 | prefix=maybe_prefix(prefix, "lm_head"), |
| 1348 | ) |
| 1349 | else: |
| 1350 | self.lm_head = PPMissingLayer() |
| 1351 | self.logits_processor = LogitsProcessor(config.vocab_size) |
| 1352 | self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] |
| 1353 | self.model.make_empty_intermediate_tensors |
| 1354 | ) |
| 1355 | |
| 1356 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 1357 | return self.model.embed_input_ids(input_ids) |
| 1358 | |
| 1359 | def compute_logits( |
| 1360 | self, |
| 1361 | hidden_states: torch.Tensor, |
| 1362 | ) -> torch.Tensor | None: |
| 1363 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 1364 | return logits |
| 1365 | |
| 1366 | def forward( |
| 1367 | self, |
| 1368 | input_ids: torch.Tensor, |
| 1369 | positions: torch.Tensor, |
| 1370 | intermediate_tensors: IntermediateTensors | None = None, |
| 1371 | inputs_embeds: torch.Tensor | None = None, |
| 1372 | ) -> torch.Tensor | IntermediateTensors: |
| 1373 | hidden_states = self.model( |
| 1374 | input_ids, positions, intermediate_tensors, inputs_embeds |
| 1375 | ) |
| 1376 | return hidden_states |
| 1377 | |
| 1378 | def get_mtp_target_hidden_states(self) -> torch.Tensor | None: |
| 1379 | """Pre-hc_head residual stream buffer (max_num_batched_tokens, |
| 1380 | hc_mult * hidden_size) for the MTP draft model. Populated by |
| 1381 | forward(); valid after each target step.""" |
| 1382 | return getattr(self.model, "_mtp_hidden_buffer", None) |
| 1383 | |
| 1384 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1385 | loader = AutoWeightsLoader(self) |
| 1386 | loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 1387 | self.model.finalize_mega_moe_weights() |
| 1388 | return loaded_params |
| 1389 | |
| 1390 | def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: |
| 1391 | return self.model.get_expert_mapping() |