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, 19:30 UTC
[ROCm] Remove MegaMoE integration in deepseek v4 (#43629)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 26, 2026 at 08:56 UTC
c8414a82712bd775b7243b19a264d9623c3bb369
2 files changed
+10
-645
vllm/models/deepseek_v4/amd/model.py
+2
-623
@@ -11,17 +11,12 @@ import torch.nn as nn
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
from vllm.distributed import (
14
- get_ep_group,
14
get_pp_group,
15
get_tensor_model_parallel_rank,
16
get_tensor_model_parallel_world_size,
17
)
19
-from vllm.forward_context import get_forward_context
18
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
19
from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
22
-from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
23
- fused_topk_bias,
24
-)
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.linear import (
22
ColumnParallelLinear,
@@ -52,7 +47,6 @@ from vllm.model_executor.models.utils import (
47
make_layers,
48
maybe_prefix,
49
)
55
-from vllm.model_executor.utils import set_weight_attrs
50
from vllm.models.deepseek_v4.attention import (
51
DeepseekV4Indexer,
52
DeepseekV4MLAModules,
@@ -60,8 +54,6 @@ from vllm.models.deepseek_v4.attention import (
54
)
55
from vllm.platforms import current_platform
56
from vllm.sequence import IntermediateTensors
63
-from vllm.triton_utils import tl, triton
64
-from vllm.utils.torch_utils import direct_register_custom_op
57
58
59
class DeepseekV4MLP(nn.Module):
@@ -115,501 +107,6 @@ class DeepseekV4MLP(nn.Module):
107
return x
108
109
118
-@triton.jit
119
-def _deepseek_v4_stage_mega_moe_inputs_kernel(
120
- hidden_states,
121
- x_fp8,
122
- x_sf,
123
- topk_ids,
124
- topk_weights,
125
- topk_idx_out,
126
- topk_weights_out,
127
- hidden_stride_m: tl.constexpr,
128
- hidden_stride_k: tl.constexpr,
129
- x_stride_m: tl.constexpr,
130
- x_stride_k: tl.constexpr,
131
- x_sf_stride_m: tl.constexpr,
132
- x_sf_stride_k: tl.constexpr,
133
- topk_ids_stride_m: tl.constexpr,
134
- topk_ids_stride_k: tl.constexpr,
135
- topk_weights_stride_m: tl.constexpr,
136
- topk_weights_stride_k: tl.constexpr,
137
- topk_idx_stride_m: tl.constexpr,
138
- topk_idx_stride_k: tl.constexpr,
139
- topk_weights_out_stride_m: tl.constexpr,
140
- topk_weights_out_stride_k: tl.constexpr,
141
- hidden_size: tl.constexpr,
142
- top_k: tl.constexpr,
143
- BLOCK_K: tl.constexpr,
144
- GROUP_K: tl.constexpr,
145
- BLOCK_TOPK: tl.constexpr,
146
-) -> None:
147
- token_id = tl.program_id(0)
148
- k_block_id = tl.program_id(1)
149
-
150
- k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
151
- k_mask = k_offsets < hidden_size
152
- hidden = tl.load(
153
- hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
154
- mask=k_mask,
155
- other=0.0,
156
- ).to(tl.float32)
157
-
158
- num_groups: tl.constexpr = BLOCK_K // GROUP_K
159
- hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
160
- amax = tl.max(hidden_groups, axis=1)
161
- amax = tl.maximum(amax, 1.0e-4)
162
-
163
- scale = amax / 448.0
164
- scale_bits = scale.to(tl.uint32, bitcast=True)
165
- scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
166
- tl.uint32
167
- )
168
- scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
169
- rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
170
-
171
- hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
172
- scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
173
- scaled = tl.reshape(scaled, [BLOCK_K])
174
- fp8 = scaled.to(tl.float8e4nv)
175
- tl.store(
176
- x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
177
- fp8,
178
- mask=k_mask,
179
- )
180
-
181
- scale_offsets = tl.arange(0, num_groups)
182
- packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
183
- tl.store(
184
- x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
185
- packed_scale,
186
- )
187
-
188
- if k_block_id == 0:
189
- topk_offsets = tl.arange(0, BLOCK_TOPK)
190
- topk_mask = topk_offsets < top_k
191
-
192
- ids = tl.load(
193
- topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
194
- mask=topk_mask,
195
- other=0,
196
- ).to(tl.int64)
197
- tl.store(
198
- topk_idx_out
199
- + token_id * topk_idx_stride_m
200
- + topk_offsets * topk_idx_stride_k,
201
- ids,
202
- mask=topk_mask,
203
- )
204
-
205
- weights = tl.load(
206
- topk_weights
207
- + token_id * topk_weights_stride_m
208
- + topk_offsets * topk_weights_stride_k,
209
- mask=topk_mask,
210
- other=0.0,
211
- )
212
- tl.store(
213
- topk_weights_out
214
- + token_id * topk_weights_out_stride_m
215
- + topk_offsets * topk_weights_out_stride_k,
216
- weights,
217
- mask=topk_mask,
218
- )
219
-
220
-
221
-def _stage_deepseek_v4_mega_moe_inputs(
222
- hidden_states: torch.Tensor,
223
- topk_weights: torch.Tensor,
224
- topk_ids: torch.Tensor,
225
- x_fp8: torch.Tensor,
226
- x_sf: torch.Tensor,
227
- topk_idx_out: torch.Tensor,
228
- topk_weights_out: torch.Tensor,
229
-) -> None:
230
- num_tokens, hidden_size = hidden_states.shape
231
- if num_tokens == 0:
232
- return
233
- if hidden_size % 128 != 0:
234
- raise ValueError(
235
- "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
236
- "a multiple of 128."
237
- )
238
- top_k = topk_ids.shape[1]
239
- if topk_weights.shape != topk_ids.shape:
240
- raise ValueError(
241
- "DeepSeek V4 MegaMoE input staging requires topk_weights and "
242
- "topk_ids to have the same shape."
243
- )
244
-
245
- block_k = 128
246
- grid = (num_tokens, triton.cdiv(hidden_size, block_k))
247
- block_topk = triton.next_power_of_2(top_k)
248
- _deepseek_v4_stage_mega_moe_inputs_kernel[grid](
249
- hidden_states,
250
- x_fp8,
251
- x_sf,
252
- topk_ids,
253
- topk_weights,
254
- topk_idx_out,
255
- topk_weights_out,
256
- hidden_states.stride(0),
257
- hidden_states.stride(1),
258
- x_fp8.stride(0),
259
- x_fp8.stride(1),
260
- x_sf.stride(0),
261
- x_sf.stride(1),
262
- topk_ids.stride(0),
263
- topk_ids.stride(1),
264
- topk_weights.stride(0),
265
- topk_weights.stride(1),
266
- topk_idx_out.stride(0),
267
- topk_idx_out.stride(1),
268
- topk_weights_out.stride(0),
269
- topk_weights_out.stride(1),
270
- hidden_size,
271
- top_k,
272
- BLOCK_K=block_k,
273
- GROUP_K=32,
274
- BLOCK_TOPK=block_topk,
275
- num_warps=4,
276
- )
277
-
278
-
279
-def make_deepseek_v4_expert_params_mapping(
280
- num_experts: int,
281
-) -> list[tuple[str, str, int, str]]:
282
- return [
283
- (
284
- "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
285
- f"experts.{expert_id}.{weight_name}.",
286
- expert_id,
287
- shard_id,
288
- )
289
- for expert_id in range(num_experts)
290
- for shard_id, weight_name in [
291
- ("w1", "w1"),
292
- ("w2", "w2"),
293
- ("w3", "w3"),
294
- ]
295
- ]
296
-
297
-
298
-class DeepseekV4MegaMoEExperts(nn.Module):
299
- _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
300
-
301
- def __init__(
302
- self,
303
- vllm_config: VllmConfig,
304
- *,
305
- num_experts: int,
306
- num_local_experts: int,
307
- experts_start_idx: int,
308
- top_k: int,
309
- hidden_size: int,
310
- intermediate_size: int,
311
- prefix: str = "",
312
- ):
313
- super().__init__()
314
- self.prefix = prefix
315
- self.num_experts = num_experts
316
- self.num_local_experts = num_local_experts
317
- self.experts_start_idx = experts_start_idx
318
- self.experts_end_idx = experts_start_idx + num_local_experts
319
- self.top_k = top_k
320
- self.hidden_size = hidden_size
321
- self.intermediate_size = intermediate_size
322
- self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
323
-
324
- weight_attrs = {"weight_loader": self.weight_loader}
325
- self.w13_weight = nn.Parameter(
326
- torch.zeros(
327
- num_local_experts,
328
- 2 * intermediate_size,
329
- hidden_size // 2,
330
- dtype=torch.uint8,
331
- ),
332
- requires_grad=False,
333
- )
334
- set_weight_attrs(self.w13_weight, weight_attrs)
335
-
336
- self.w13_weight_scale = nn.Parameter(
337
- torch.zeros(
338
- num_local_experts,
339
- 2 * intermediate_size,
340
- hidden_size // 32,
341
- dtype=torch.uint8,
342
- ),
343
- requires_grad=False,
344
- )
345
- set_weight_attrs(self.w13_weight_scale, weight_attrs)
346
- self.w13_weight_scale.quant_method = "block"
347
-
348
- self.w2_weight = nn.Parameter(
349
- torch.zeros(
350
- num_local_experts,
351
- hidden_size,
352
- intermediate_size // 2,
353
- dtype=torch.uint8,
354
- ),
355
- requires_grad=False,
356
- )
357
- set_weight_attrs(self.w2_weight, weight_attrs)
358
-
359
- self.w2_weight_scale = nn.Parameter(
360
- torch.zeros(
361
- num_local_experts,
362
- hidden_size,
363
- intermediate_size // 32,
364
- dtype=torch.uint8,
365
- ),
366
- requires_grad=False,
367
- )
368
- set_weight_attrs(self.w2_weight_scale, weight_attrs)
369
- self.w2_weight_scale.quant_method = "block"
370
-
371
- self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
372
- self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
373
-
374
- # Register in the static forward context so the custom-op wrapper
375
- # can look up this module by name from within a torch.compile graph.
376
- compilation_config = vllm_config.compilation_config
377
- if prefix in compilation_config.static_forward_context:
378
- raise ValueError(f"Duplicate layer name: {prefix}")
379
- compilation_config.static_forward_context[prefix] = self
380
-
381
- def _map_global_expert_id(self, expert_id: int) -> int:
382
- if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
383
- return -1
384
- return expert_id - self.experts_start_idx
385
-
386
- def weight_loader(
387
- self,
388
- param: nn.Parameter,
389
- loaded_weight: torch.Tensor,
390
- weight_name: str,
391
- shard_id: str,
392
- expert_id: int,
393
- return_success: bool = False,
394
- ) -> bool | None:
395
- local_expert_id = self._map_global_expert_id(expert_id)
396
- if local_expert_id == -1:
397
- return False if return_success else None
398
-
399
- expert_data = param.data[local_expert_id]
400
- if shard_id in ("w1", "w3"):
401
- if "w13_" not in weight_name:
402
- return False if return_success else None
403
- shard_offset = 0 if shard_id == "w1" else self.intermediate_size
404
- expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)
405
- elif shard_id == "w2":
406
- if "w2_" not in weight_name:
407
- return False if return_success else None
408
- else:
409
- raise ValueError(f"Unsupported expert shard id: {shard_id}")
410
-
411
- if expert_data.shape != loaded_weight.shape:
412
- raise ValueError(
413
- f"DeepSeek V4 MegaMoE expert weight shape mismatch for "
414
- f"{weight_name}: parameter shard {tuple(expert_data.shape)} "
415
- f"vs checkpoint {tuple(loaded_weight.shape)}"
416
- )
417
- expert_data.copy_(loaded_weight)
418
- return True if return_success else None
419
-
420
- @staticmethod
421
- def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:
422
- return (sf.to(torch.int32) << 23).view(torch.float32)
423
-
424
- def _check_runtime_supported(self) -> None:
425
- if not torch.cuda.is_available():
426
- raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
427
- device = self.w13_weight.device
428
- if device.type != "cuda":
429
- raise NotImplementedError(
430
- "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA."
431
- )
432
- if torch.cuda.get_device_capability(device)[0] != 10:
433
- raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
434
- if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
435
- raise ValueError(
436
- "DeepGEMM MegaMoE requires hidden and intermediate sizes "
437
- "to be multiples of 128."
438
- )
439
-
440
- def finalize_weights(self) -> None:
441
- if self._transformed_l1_weights is not None:
442
- return
443
-
444
- self._check_runtime_supported()
445
- import vllm.third_party.deep_gemm as deep_gemm
446
-
447
- w13_scale = deep_gemm.transform_sf_into_required_layout(
448
- self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
449
- 2 * self.intermediate_size,
450
- self.hidden_size,
451
- (1, 32),
452
- self.num_local_experts,
453
- )
454
- w2_scale = deep_gemm.transform_sf_into_required_layout(
455
- self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
456
- self.hidden_size,
457
- self.intermediate_size,
458
- (1, 32),
459
- self.num_local_experts,
460
- )
461
- self._transformed_l1_weights, self._transformed_l2_weights = (
462
- deep_gemm.transform_weights_for_mega_moe(
463
- (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
464
- (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
465
- )
466
- )
467
- # Drop the original loader-side parameters: the MegaMoE kernels only
468
- # consume the transformed views above. transform_weights_for_mega_moe
469
- # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
470
- # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
471
- # aliases the original storage, and _transformed_l2_weights still holds
472
- # it, so the storage stays live after we drop the Parameter.
473
- self.w13_weight = None
474
- self.w13_weight_scale = None
475
- self.w2_weight = None
476
- self.w2_weight_scale = None
477
-
478
- def get_symm_buffer(self):
479
- import vllm.third_party.deep_gemm as deep_gemm
480
-
481
- group = get_ep_group().device_group
482
- device = torch.accelerator.current_device_index()
483
- key = (
484
- id(group),
485
- device,
486
- self.num_experts,
487
- self.max_num_tokens,
488
- self.top_k,
489
- self.hidden_size,
490
- self.intermediate_size,
491
- )
492
- symm_buffer = self._symm_buffer_cache.get(key)
493
- if symm_buffer is None:
494
- symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
495
- group,
496
- self.num_experts,
497
- self.max_num_tokens,
498
- self.top_k,
499
- self.hidden_size,
500
- self.intermediate_size,
501
- )
502
- self._symm_buffer_cache[key] = symm_buffer
503
- return symm_buffer
504
-
505
- def forward(
506
- self,
507
- hidden_states: torch.Tensor,
508
- topk_weights: torch.Tensor,
509
- topk_ids: torch.Tensor,
510
- *,
511
- activation_clamp: float | None,
512
- fast_math: bool = True,
513
- ) -> torch.Tensor:
514
- if hidden_states.shape[0] > self.max_num_tokens:
515
- raise ValueError(
516
- f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
517
- f"but the symmetric buffer was sized for {self.max_num_tokens}."
518
- )
519
- y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
520
- torch.ops.vllm.deepseek_v4_mega_moe_experts(
521
- hidden_states,
522
- topk_weights,
523
- topk_ids,
524
- y,
525
- self.prefix,
526
- activation_clamp,
527
- fast_math,
528
- )
529
- return y
530
-
531
- def _run_mega_moe(
532
- self,
533
- hidden_states: torch.Tensor,
534
- topk_weights: torch.Tensor,
535
- topk_ids: torch.Tensor,
536
- y: torch.Tensor,
537
- activation_clamp: float | None,
538
- fast_math: bool,
539
- ) -> None:
540
- import vllm.third_party.deep_gemm as deep_gemm
541
-
542
- symm_buffer = self.get_symm_buffer()
543
- num_tokens = hidden_states.shape[0]
544
- _stage_deepseek_v4_mega_moe_inputs(
545
- hidden_states,
546
- topk_weights,
547
- topk_ids,
548
- symm_buffer.x[:num_tokens],
549
- symm_buffer.x_sf[:num_tokens],
550
- symm_buffer.topk_idx[:num_tokens],
551
- symm_buffer.topk_weights[:num_tokens],
552
- )
553
-
554
- # This method must have been already called during the weight loading phase.
555
- # We call it again here to cover the dummy weight loading case.
556
- self.finalize_weights()
557
-
558
- assert self._transformed_l1_weights is not None
559
- assert self._transformed_l2_weights is not None
560
- deep_gemm.fp8_fp4_mega_moe(
561
- y,
562
- self._transformed_l1_weights,
563
- self._transformed_l2_weights,
564
- symm_buffer,
565
- activation_clamp=activation_clamp,
566
- fast_math=fast_math,
567
- )
568
-
569
-
570
-DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
571
-
572
-
573
-def _deepseek_v4_mega_moe_experts_op(
574
- hidden_states: torch.Tensor,
575
- topk_weights: torch.Tensor,
576
- topk_ids: torch.Tensor,
577
- out: torch.Tensor,
578
- layer_name: str,
579
- activation_clamp: float | None,
580
- fast_math: bool,
581
-) -> None:
582
- self = get_forward_context().no_compile_layers[layer_name]
583
- self._run_mega_moe(
584
- hidden_states,
585
- topk_weights,
586
- topk_ids,
587
- out,
588
- activation_clamp,
589
- fast_math,
590
- )
591
-
592
-
593
-def _deepseek_v4_mega_moe_experts_op_fake(
594
- hidden_states: torch.Tensor,
595
- topk_weights: torch.Tensor,
596
- topk_ids: torch.Tensor,
597
- out: torch.Tensor,
598
- layer_name: str,
599
- activation_clamp: float | None,
600
- fast_math: bool,
601
-) -> None:
602
- return None
603
-
604
-
605
-direct_register_custom_op(
606
- op_name="deepseek_v4_mega_moe_experts",
607
- op_func=_deepseek_v4_mega_moe_experts_op,
608
- mutates_args=["out"],
609
- fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
610
-)
611
-
612
-
110
class DeepseekV4MoE(nn.Module):
111
def __init__(
112
self,
@@ -622,15 +119,6 @@ class DeepseekV4MoE(nn.Module):
119
config = vllm_config.model_config.hf_config
120
quant_config = vllm_config.quant_config
121
self.prefix = prefix
625
- self.use_mega_moe = (
626
- vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
627
- )
628
- if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
629
- raise NotImplementedError(
630
- "DeepSeek V4 MegaMoE currently requires expert parallel. "
631
- "Enable it with --enable-expert-parallel, or pick a different "
632
- "moe backend."
633
- )
122
123
self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
124
self.hidden_size = config.hidden_size
@@ -641,16 +129,6 @@ class DeepseekV4MoE(nn.Module):
129
self.swiglu_limit = config.swiglu_limit
130
self.renormalize = config.norm_topk_prob
131
self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
644
- if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
645
- raise NotImplementedError(
646
- "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
647
- )
648
- if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
649
- raise NotImplementedError(
650
- "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
651
- f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
652
- "deep_gemm_mega_moe for this checkpoint."
653
- )
132
133
self.gate = GateLinear(
134
input_size=config.hidden_size,
@@ -663,7 +141,7 @@ class DeepseekV4MoE(nn.Module):
141
self.gate.e_score_correction_bias = None
142
self.gate.tid2eid = None
143
is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
666
- self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
144
+ self.hash_indices_dtype = torch.int32
145
if is_hash_moe:
146
# hash MoE doesn't use e_score_correction_bias
147
# Use randint instead of empty to avoid garbage values causing
@@ -694,47 +172,10 @@ class DeepseekV4MoE(nn.Module):
172
hidden_act=config.hidden_act,
173
swiglu_limit=self.swiglu_limit,
174
quant_config=quant_config,
697
- reduce_results=self.use_mega_moe,
175
+ reduce_results=False,
176
prefix=f"{prefix}.shared_experts",
177
)
178
701
- if self.use_mega_moe:
702
- self._init_mega_moe_experts(vllm_config, config, prefix)
703
- else:
704
- self._init_fused_moe_experts(config, quant_config, prefix)
705
-
706
- def _init_mega_moe_experts(
707
- self,
708
- vllm_config: VllmConfig,
709
- config,
710
- prefix: str,
711
- ) -> None:
712
- self.ep_group = get_ep_group()
713
- self.ep_size = self.ep_group.world_size
714
- self.ep_rank = self.ep_group.rank_in_group
715
- assert config.n_routed_experts % self.ep_size == 0
716
-
717
- self.n_local_experts = config.n_routed_experts // self.ep_size
718
- self.experts_start_idx = self.ep_rank * self.n_local_experts
719
- self.experts_end_idx = self.experts_start_idx + self.n_local_experts
720
-
721
- self.experts = DeepseekV4MegaMoEExperts(
722
- vllm_config,
723
- num_experts=config.n_routed_experts,
724
- num_local_experts=self.n_local_experts,
725
- experts_start_idx=self.experts_start_idx,
726
- top_k=config.num_experts_per_tok,
727
- hidden_size=config.hidden_size,
728
- intermediate_size=config.moe_intermediate_size,
729
- prefix=f"{prefix}.experts",
730
- )
731
-
732
- def _init_fused_moe_experts(
733
- self,
734
- config,
735
- quant_config,
736
- prefix: str,
737
- ) -> None:
179
self.tp_rank = get_tensor_model_parallel_rank()
180
assert config.n_routed_experts % self.tp_size == 0
181
@@ -766,44 +207,6 @@ class DeepseekV4MoE(nn.Module):
207
if self.gate.tid2eid is not None and input_ids is None:
208
raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
209
769
- if not self.use_mega_moe:
770
- return self._forward_fused_moe(hidden_states, input_ids)
771
-
772
- org_shape = hidden_states.shape
773
- router_logits, _ = self.gate(hidden_states)
774
- topk_weights, topk_ids = fused_topk_bias(
775
- hidden_states=hidden_states,
776
- gating_output=router_logits,
777
- scoring_func=self.scoring_func,
778
- e_score_correction_bias=self.gate.e_score_correction_bias.data
779
- if self.gate.e_score_correction_bias is not None
780
- else None,
781
- topk=self.n_activated_experts,
782
- renormalize=self.renormalize,
783
- indices_type=self.hash_indices_dtype,
784
- input_tokens=input_ids,
785
- hash_indices_table=self.gate.tid2eid,
786
- routed_scaling_factor=self.routed_scaling_factor,
787
- )
788
- activation_clamp = (
789
- float(self.swiglu_limit) if self.swiglu_limit is not None else None
790
- )
791
- final_hidden_states = self.experts(
792
- hidden_states,
793
- topk_weights,
794
- topk_ids,
795
- activation_clamp=activation_clamp,
796
- )
797
-
798
- if self.shared_experts is not None:
799
- shared_output = self.shared_experts(hidden_states)
800
- final_hidden_states += shared_output
801
-
802
- return final_hidden_states.view(org_shape)
803
-
804
- def _forward_fused_moe(
805
- self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
806
- ) -> torch.Tensor:
210
org_shape = hidden_states.shape
211
if self.experts.is_internal_router:
212
# In this case, the gate/router runs inside the FusedMoE class
@@ -822,10 +225,6 @@ class DeepseekV4MoE(nn.Module):
225
226
return final_hidden_states.view(org_shape)
227
825
- def finalize_mega_moe_weights(self) -> None:
826
- if self.use_mega_moe:
827
- self.experts.finalize_weights()
828
-
228
229
class DeepseekV4Attention(nn.Module):
230
def __init__(
@@ -1211,15 +610,6 @@ class DeepseekV4Model(nn.Module):
610
config = vllm_config.model_config.hf_config
611
quant_config = vllm_config.quant_config
612
self.config = config
1214
- self.use_mega_moe = (
1215
- vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
1216
- )
1217
- if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
1218
- raise NotImplementedError(
1219
- "DeepSeek V4 MegaMoE currently requires expert parallel. "
1220
- "Enable it with --enable-expert-parallel, or pick a different "
1221
- "moe backend."
1222
- )
613
self.vocab_size = config.vocab_size
614
self.hc_eps = config.hc_eps
615
self.hc_mult = config.hc_mult
@@ -1348,9 +738,6 @@ class DeepseekV4Model(nn.Module):
738
assert intermediate_tensors is not None
739
hidden_states = intermediate_tensors["hidden_states"]
740
1351
- if self.use_mega_moe:
1352
- input_ids = input_ids.to(torch.int64)
1353
-
741
residual, post_mix, res_mix = None, None, None
742
for layer in islice(self.layers, self.start_layer, self.end_layer):
743
hidden_states, residual, post_mix, res_mix = layer(
@@ -1482,9 +869,6 @@ class DeepseekV4Model(nn.Module):
869
return loaded_params
870
871
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1485
- first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
1486
- if first_layer.ffn.use_mega_moe:
1487
- return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
872
# Params for weights, fp8 weight scales, fp8 activation scales
873
# (param_name, weight_name, expert_id, shard_id)
874
return FusedMoE.make_expert_params_mapping(
@@ -1495,10 +879,6 @@ class DeepseekV4Model(nn.Module):
879
num_experts=self.config.n_routed_experts,
880
)
881
1498
- def finalize_mega_moe_weights(self) -> None:
1499
- for layer in islice(self.layers, self.start_layer, self.end_layer):
1500
- layer.ffn.finalize_mega_moe_weights()
1501
-
882
883
def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
884
if expert_dtype == "fp4":
@@ -1601,7 +981,6 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
981
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
982
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
983
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1604
- self.model.finalize_mega_moe_weights()
984
return loaded_params
985
986
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
vllm/models/deepseek_v4/amd/mtp.py
+8
-22
@@ -40,10 +40,7 @@ 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
-)
43
+from .model import DeepseekV4DecoderLayer
44
45
logger = init_logger(__name__)
46
@@ -330,19 +327,13 @@ class DeepSeekV4MTP(nn.Module):
327
head_rank_end = n_local_head * (tp_rank + 1)
328
329
# Pre-compute expert mapping ONCE.
333
- first_layer = next(iter(self.model.layers.values()))
334
- if first_layer.mtp_block.ffn.use_mega_moe:
335
- expert_mapping = make_deepseek_v4_expert_params_mapping(
336
- self.config.n_routed_experts
337
- )
338
- else:
339
- expert_mapping = FusedMoE.make_expert_params_mapping(
340
- self,
341
- ckpt_gate_proj_name="w1",
342
- ckpt_down_proj_name="w2",
343
- ckpt_up_proj_name="w3",
344
- num_experts=self.config.n_routed_experts,
345
- )
330
+ expert_mapping = FusedMoE.make_expert_params_mapping(
331
+ self,
332
+ ckpt_gate_proj_name="w1",
333
+ ckpt_down_proj_name="w2",
334
+ ckpt_up_proj_name="w3",
335
+ num_experts=self.config.n_routed_experts,
336
+ )
337
338
# FP8 experts register ``..._weight_scale_inv`` (block_quant) while
339
# FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
@@ -465,14 +456,9 @@ class DeepSeekV4MTP(nn.Module):
456
f"Use a checkpoint that includes MTP layer weights, "
457
f"or disable speculative decoding."
458
)
468
- self.finalize_mega_moe_weights()
459
logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
460
return loaded_params
461
472
- def finalize_mega_moe_weights(self) -> None:
473
- for layer in self.model.layers.values():
474
- layer.mtp_block.ffn.finalize_mega_moe_weights()
475
-
462
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
463
"""
464
Rewrite the weight name to match the format of the original model.