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, 03:30 UTC
[DSV4] Minor cleanup for DeepseekV4MegaMoEExperts (#44367)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 2, 2026 at 17:54 UTC
b254e0456c9851414cf9aada98a9129235614429
1 file changed
+1
-18
vllm/models/deepseek_v4/nvidia/model.py
+1
-18
@@ -420,25 +420,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
420
f"but the symmetric buffer was sized for {self.max_num_tokens}."
421
)
422
y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
423
- self._run_mega_moe(
424
- hidden_states,
425
- topk_weights,
426
- topk_ids,
427
- y,
428
- activation_clamp,
429
- fast_math,
430
- )
431
- return y
423
433
- def _run_mega_moe(
434
- self,
435
- hidden_states: torch.Tensor,
436
- topk_weights: torch.Tensor,
437
- topk_ids: torch.Tensor,
438
- y: torch.Tensor,
439
- activation_clamp: float | None,
440
- fast_math: bool,
441
- ) -> None:
424
from vllm.utils.deep_gemm import _import_deep_gemm
425
426
deep_gemm = _import_deep_gemm()
@@ -484,6 +466,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
466
activation_clamp=activation_clamp,
467
fast_math=fast_math,
468
)
469
+ return y
470
471
472
DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]