@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2427094152

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, 04:45 UTC

[Feature] Support EPLB for DeepSeek v4 Mega Moe (#43339)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Wei Zhao (Engrg-Hardware 1) <weizha@login-lyris01.lyris.clusters.nvidia.com>

Wei Zhao committed Jun 2, 2026 at 13:56 UTC 242709415287e2f7ebe6e3fcddea349c498307d9

4 files changed +232 -46

vllm/distributed/eplb/eplb_utils.py

+16 -7

@@ -61,25 +61,31 @@ class CpuGpuEvent:

61 self._recorded.set()

64 -def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:

64 +def override_envs_for_eplb(

65 + parallel_config: ParallelConfig,

66 + moe_backend: str | None = None,

67 +) -> None:

68 """

69 Override environment variables for EPLB when specific conditions are met.

71 Args:

72 parallel_config: The parallel configuration object.

73 + moe_backend: The configured MoE backend (e.g. ``deep_gemm_mega_moe``).

74 """

75 is_data_parallel = parallel_config.data_parallel_size > 1

76 is_eplb_enabled = parallel_config.enable_eplb

77 async_eplb = parallel_config.eplb_config.use_async

78 is_deepep_ll = parallel_config.all2all_backend == "deepep_low_latency"

79 + is_mega_moe = moe_backend == "deep_gemm_mega_moe"

80 is_nccl_based_eplb_communicator = parallel_config.eplb_config.communicator in (

81 "torch_nccl",

82 "pynccl",

83 )

80 - # Override NCCL_MAX_CTAS to avoid hangs when using async EPLB with the

81 - # DeepEP low-latency backend.

85 + # Override NCCL_MAX_CTAS to avoid hangs when EPLB's NCCL weight exchange

86 + # contends with MoE backend's cooperative-launch on GPU SMs.

87 #

88 + # DeepEP low-latency:

89 # The hang happens when two ranks interleave kernel launches differently

90 # between NCCL collectives (used by async EPLB weight exchange) and DeepEP

91 # low-latency (LL) kernels. DeepEP LL uses a cooperative launch and tries

@@ -94,12 +100,14 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:

100 # Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for the DeepEP

101 # cooperative kernel to launch and complete, breaking the deadlock.

102 # See: https://github.com/deepseek-ai/DeepEP/issues/496

103 + #

104 + # DeepGEMM Mega MoE also uses cooperative launch and will cause hang even

105 + # with sync EPLB.

106 if (

107 is_data_parallel

108 and is_eplb_enabled

100 - and is_deepep_ll

101 - and async_eplb

109 and is_nccl_based_eplb_communicator

110 + and ((is_deepep_ll and async_eplb) or is_mega_moe)

111 ):

112 current_value_str = os.getenv("NCCL_MAX_CTAS")

113

@@ -108,9 +116,10 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:

116

117 override_value = 8

118 os.environ["NCCL_MAX_CTAS"] = str(override_value)

119 + backend = "deepep_low_latency" if is_deepep_ll else "deep_gemm_mega_moe"

120 logger.info_once(

121 f"EPLB: Setting NCCL_MAX_CTAS={override_value} "

113 - "for expert parallel with NCCL-based EPLB communicator and "

114 - "deepep_low_latency backend",

122 + f"for expert parallel with NCCL-based EPLB communicator and "

123 + f"cooperative MoE backend ({backend})",

124 scope="global",

125 )

vllm/models/deepseek_v4/nvidia/model.py

+211 -38

@@ -1,7 +1,7 @@

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

4 +from collections.abc import Callable, Iterable, MutableSequence, Sequence

5 from itertools import islice

7 import regex as re

@@ -15,6 +15,7 @@ from vllm.distributed import (

15 get_tensor_model_parallel_rank,

16 get_tensor_model_parallel_world_size,

17 )

18 +from vllm.distributed.eplb.eplb_state import EplbLayerState

19 from vllm.model_executor.kernels.mhc.tilelang import (

20 hc_head_fused_kernel_tilelang,

21 mhc_fused_post_pre_tilelang,

@@ -23,6 +24,9 @@ from vllm.model_executor.kernels.mhc.tilelang import (

24 )

25 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp

26 from vllm.model_executor.layers.fused_moe import FusedMoE

27 +from vllm.model_executor.layers.fused_moe.router.base_router import (

28 + eplb_map_to_physical_and_record,

29 +)

30 from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (

31 fused_topk_bias,

32 )

@@ -40,7 +44,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

44 VocabParallelEmbedding,

45 )

46 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

43 -from vllm.model_executor.models.interfaces import SupportsPP

47 +from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP

48 from vllm.model_executor.models.utils import (

49 AutoWeightsLoader,

50 PPMissingLayer,

@@ -144,6 +148,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):

148 hidden_size: int,

149 intermediate_size: int,

150 prefix: str = "",

151 + num_logical_experts: int | None = None,

152 ):

153 super().__init__()

154 self.prefix = prefix

@@ -156,6 +161,12 @@ class DeepseekV4MegaMoEExperts(nn.Module):

161 self.intermediate_size = intermediate_size

162 self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

163

164 + self.num_logical_experts = (

165 + num_logical_experts if num_logical_experts is not None else num_experts

166 + )

167 +

168 + self.eplb_state = EplbLayerState()

169 +

170 weight_attrs = {"weight_loader": self.weight_loader}

171 self.w13_weight = nn.Parameter(

172 torch.zeros(

@@ -206,10 +217,22 @@ class DeepseekV4MegaMoEExperts(nn.Module):

217 self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None

218 self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None

219

209 - def _map_global_expert_id(self, expert_id: int) -> int:

210 - if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:

211 - return -1

212 - return expert_id - self.experts_start_idx

220 + # Register in the static forward context so the custom-op wrapper

221 + # can look up this module by name from within a torch.compile graph.

222 + compilation_config = vllm_config.compilation_config

223 + if prefix in compilation_config.static_forward_context:

224 + raise ValueError(f"Duplicate layer name: {prefix}")

225 + compilation_config.static_forward_context[prefix] = self

226 +

227 + def _map_global_expert_id(self, expert_id: int) -> list[int]:

228 + """Return local (per-rank) slot offsets where logical expert

229 + `expert_id` should land on this rank.

230 + """

231 + physical_ids: list[int] = []

232 + for p in range(self.experts_start_idx, self.experts_end_idx):

233 + if p % self.num_logical_experts == expert_id:

234 + physical_ids.append(p - self.experts_start_idx)

235 + return physical_ids

236

237 def weight_loader(

238 self,

@@ -220,30 +243,38 @@ class DeepseekV4MegaMoEExperts(nn.Module):

243 expert_id: int,

244 return_success: bool = False,

245 ) -> bool | None:

223 - local_expert_id = self._map_global_expert_id(expert_id)

224 - if local_expert_id == -1:

246 + local_expert_ids = self._map_global_expert_id(expert_id)

247 + if not local_expert_ids:

248 return False if return_success else None

249

227 - expert_data = param.data[local_expert_id]

228 - if shard_id in ("w1", "w3"):

229 - if "w13_" not in weight_name:

230 - return False if return_success else None

231 - shard_offset = 0 if shard_id == "w1" else self.intermediate_size

232 - expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)

233 - elif shard_id == "w2":

234 - if "w2_" not in weight_name:

235 - return False if return_success else None

236 - else:

237 - raise ValueError(f"Unsupported expert shard id: {shard_id}")

250 + loaded_any = False

251 + for local_expert_id in local_expert_ids:

252 + expert_data = param.data[local_expert_id]

253 + if shard_id in ("w1", "w3"):

254 + if "w13_" not in weight_name:

255 + continue

256 + shard_offset = 0 if shard_id == "w1" else self.intermediate_size

257 + expert_data = expert_data.narrow(

258 + 0, shard_offset, self.intermediate_size

259 + )

260 + elif shard_id == "w2":

261 + if "w2_" not in weight_name:

262 + continue

263 + else:

264 + raise ValueError(f"Unsupported expert shard id: {shard_id}")

265

239 - if expert_data.shape != loaded_weight.shape:

240 - raise ValueError(

241 - f"DeepSeek V4 MegaMoE expert weight shape mismatch for "

242 - f"{weight_name}: parameter shard {tuple(expert_data.shape)} "

243 - f"vs checkpoint {tuple(loaded_weight.shape)}"

244 - )

245 - expert_data.copy_(loaded_weight)

246 - return True if return_success else None

266 + if expert_data.shape != loaded_weight.shape:

267 + raise ValueError(

268 + f"DeepSeek V4 MegaMoE expert weight shape mismatch for "

269 + f"{weight_name}: parameter shard {tuple(expert_data.shape)} "

270 + f"vs checkpoint {tuple(loaded_weight.shape)}"

271 + )

272 + expert_data.copy_(loaded_weight)

273 + loaded_any = True

274 +

275 + if return_success:

276 + return loaded_any

277 + return None

278

279 @staticmethod

280 def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:

@@ -264,7 +295,9 @@ class DeepseekV4MegaMoEExperts(nn.Module):

295 return

296

297 self._check_runtime_supported()

267 - import vllm.third_party.deep_gemm as deep_gemm

298 + from vllm.utils.deep_gemm import _import_deep_gemm

299 +

300 + deep_gemm = _import_deep_gemm()

301

302 w13_scale = deep_gemm.transform_sf_into_required_layout(

303 self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),

@@ -298,7 +331,9 @@ class DeepseekV4MegaMoEExperts(nn.Module):

331 self.w2_weight_scale = None

332

333 def get_symm_buffer(self):

301 - import vllm.third_party.deep_gemm as deep_gemm

334 + from vllm.utils.deep_gemm import _import_deep_gemm

335 +

336 + deep_gemm = _import_deep_gemm()

337

338 group = get_ep_group().device_group

339 device = torch.accelerator.current_device_index()

@@ -324,6 +359,52 @@ class DeepseekV4MegaMoEExperts(nn.Module):

359 self._symm_buffer_cache[key] = symm_buffer

360 return symm_buffer

361

362 + def set_eplb_state(

363 + self,

364 + moe_layer_idx: int,

365 + expert_load_view: torch.Tensor,

366 + logical_to_physical_map: torch.Tensor,

367 + logical_replica_count: torch.Tensor,

368 + ) -> None:

369 + self.eplb_state.set_layer_state(

370 + moe_layer_idx,

371 + expert_load_view,

372 + logical_to_physical_map,

373 + logical_replica_count,

374 + )

375 +

376 + def get_expert_weights(self) -> list[torch.Tensor]:

377 + self.finalize_weights()

378 + assert self._transformed_l1_weights is not None

379 + assert self._transformed_l2_weights is not None

380 +

381 + def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor:

382 + """Return a (num_local_experts, -1) view with contiguous memory layout."""

383 + assert t.shape[0] == self.num_local_experts

384 + if t.is_contiguous():

385 + return t.view(self.num_local_experts, -1)

386 + elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]:

387 + # scales have shape (E, M, N) with memory layout (E, N, M)

388 + back = torch.transpose(t, 1, 2)

389 + assert back.is_contiguous()

390 + return back.view(self.num_local_experts, -1)

391 +

392 + raise AssertionError(

393 + f"DSv4 EPLB {name}: non-contiguous expert tensor with "

394 + f"unexpected layout shape={tuple(t.shape)} "

395 + f"stride={tuple(t.stride())} dtype={t.dtype}"

396 + )

397 +

398 + return [

399 + _to_eplb_view("l1_packed", self._transformed_l1_weights[0]),

400 + _to_eplb_view("l1_scale", self._transformed_l1_weights[1]),

401 + _to_eplb_view("l2_weight", self._transformed_l2_weights[0]),

402 + _to_eplb_view("l2_scale", self._transformed_l2_weights[1]),

403 + ]

404 +

405 + def update_expert_map(self) -> None:

406 + pass

407 +

408 def forward(

409 self,

410 hidden_states: torch.Tensor,

@@ -358,10 +439,27 @@ class DeepseekV4MegaMoEExperts(nn.Module):

439 activation_clamp: float | None,

440 fast_math: bool,

441 ) -> None:

361 - import vllm.third_party.deep_gemm as deep_gemm

442 + from vllm.utils.deep_gemm import _import_deep_gemm

443 +

444 + deep_gemm = _import_deep_gemm()

445

446 symm_buffer = self.get_symm_buffer()

447 num_tokens = hidden_states.shape[0]

448 +

449 + # EPLB: map logical expert IDs to physical replicas and record load.

450 + eplb_state = self.eplb_state

451 + if eplb_state.logical_to_physical_map is not None:

452 + assert eplb_state.expert_load_view is not None

453 + assert eplb_state.logical_replica_count is not None

454 + assert eplb_state.should_record_tensor is not None

455 + topk_ids = eplb_map_to_physical_and_record(

456 + topk_ids=topk_ids,

457 + expert_load_view=eplb_state.expert_load_view,

458 + logical_to_physical_map=eplb_state.logical_to_physical_map,

459 + logical_replica_count=eplb_state.logical_replica_count,

460 + record_enabled=eplb_state.should_record_tensor,

461 + )

462 +

463 prepare_megamoe_inputs(

464 hidden_states,

465 topk_weights,

@@ -493,17 +591,33 @@ class DeepseekV4MoE(nn.Module):

591 self.ep_group = get_ep_group()

592 self.ep_size = self.ep_group.world_size

593 self.ep_rank = self.ep_group.rank_in_group

496 - assert config.n_routed_experts % self.ep_size == 0

594

498 - self.n_local_experts = config.n_routed_experts // self.ep_size

499 - self.experts_start_idx = self.ep_rank * self.n_local_experts

500 - self.experts_end_idx = self.experts_start_idx + self.n_local_experts

595 + eplb_config = vllm_config.parallel_config.eplb_config

596 + self.n_redundant_experts = eplb_config.num_redundant_experts

597 + self.n_routed_experts = config.n_routed_experts

598 + self.n_shared_experts = config.n_shared_experts or 0

599 + self.n_logical_experts = self.n_routed_experts

600 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

601 + assert self.n_physical_experts % self.ep_size == 0, (

602 + f"n_physical_experts={self.n_physical_experts} must be divisible by "

603 + f"ep_size={self.ep_size}. Adjust num_redundant_experts."

604 + )

605 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size

606 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts

607 + self.physical_expert_end = (

608 + self.physical_expert_start + self.n_local_physical_experts

609 + )

610 +

611 + self.n_local_experts = self.n_local_physical_experts

612 + self.experts_start_idx = self.physical_expert_start

613 + self.experts_end_idx = self.physical_expert_end

614

615 self.experts = DeepseekV4MegaMoEExperts(

616 vllm_config,

504 - num_experts=config.n_routed_experts,

505 - num_local_experts=self.n_local_experts,

506 - experts_start_idx=self.experts_start_idx,

617 + num_experts=self.n_physical_experts,

618 + num_local_experts=self.n_local_physical_experts,

619 + experts_start_idx=self.physical_expert_start,

620 + num_logical_experts=self.n_logical_experts,

621 top_k=config.num_experts_per_tok,

622 hidden_size=config.hidden_size,

623 intermediate_size=config.moe_intermediate_size,

@@ -1242,7 +1356,44 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:

1356 )

1357

1358

1245 -class DeepseekV4ForCausalLM(nn.Module, SupportsPP):

1359 +class DeepseekV4MixtureOfExperts(MixtureOfExperts):

1360 + moe_mlp_layers: list["DeepseekV4MoE"]

1361 +

1362 + def extract_moe_parameters(self, example_moe: "DeepseekV4MoE | None") -> None:

1363 + if example_moe is None:

1364 + self.num_moe_layers = 0

1365 + self.num_expert_groups = 0

1366 + self.num_logical_experts = 0

1367 + self.num_physical_experts = 0

1368 + self.num_local_physical_experts = 0

1369 + self.num_routed_experts = 0

1370 + self.num_shared_experts = 0

1371 + self.num_redundant_experts = 0

1372 + return

1373 + self.num_logical_experts = example_moe.n_logical_experts

1374 + self.num_physical_experts = example_moe.n_physical_experts

1375 + self.num_local_physical_experts = example_moe.n_local_physical_experts

1376 + self.num_routed_experts = example_moe.n_routed_experts

1377 + self.num_shared_experts = example_moe.n_shared_experts

1378 + self.num_redundant_experts = example_moe.n_redundant_experts

1379 +

1380 + def update_physical_experts_metadata(

1381 + self,

1382 + num_physical_experts: int,

1383 + num_local_physical_experts: int,

1384 + ) -> None:

1385 + assert self.num_local_physical_experts == num_local_physical_experts

1386 + self.num_physical_experts = num_physical_experts

1387 + self.num_local_physical_experts = num_local_physical_experts

1388 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts

1389 + for moe in self.moe_mlp_layers:

1390 + moe.n_local_physical_experts = num_local_physical_experts

1391 + moe.n_physical_experts = num_physical_experts

1392 + moe.n_redundant_experts = self.num_redundant_experts

1393 + moe.experts.update_expert_map()

1394 +

1395 +

1396 +class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts):

1397 model_cls = DeepseekV4Model

1398

1399 # Default mapper assumes the original FP4-expert checkpoint layout.

@@ -1274,6 +1425,28 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP):

1425 self.model.make_empty_intermediate_tensors

1426 )

1427

1428 + self.set_moe_parameters()

1429 +

1430 + def set_moe_parameters(self) -> None:

1431 + self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = []

1432 + self.num_expert_groups = getattr(self.config, "n_group", 1)

1433 + self.num_moe_layers = self.config.num_hidden_layers

1434 + self.moe_layers: list[nn.Module] = []

1435 + self.moe_mlp_layers: list[DeepseekV4MoE] = []

1436 + example_moe: DeepseekV4MoE | None = None

1437 + for layer in self.model.layers:

1438 + if isinstance(layer, PPMissingLayer):

1439 + continue

1440 + if not isinstance(layer, DeepseekV4DecoderLayer):

1441 + continue

1442 + if isinstance(layer.ffn, DeepseekV4MoE):

1443 + example_moe = layer.ffn

1444 + self.moe_mlp_layers.append(layer.ffn)

1445 + self.moe_layers.append(layer.ffn.experts)

1446 +

1447 + self.num_moe_layers = len(self.moe_layers)

1448 + self.extract_moe_parameters(example_moe)

1449 +

1450 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1451 return self.model.embed_input_ids(input_ids)

1452

vllm/utils/deep_gemm.py

@@ -140,6 +140,7 @@ _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None

140 _transform_sf_into_required_layout_impl: Callable[..., Any] | None = None

141

142

143 +@functools.cache

144 def _import_deep_gemm():

145 """Import the deep_gemm module.

146

vllm/v1/worker/gpu_worker.py

+4 -1

@@ -1147,7 +1147,10 @@ def init_worker_distributed_environment(

1147 from vllm.model_executor.layers.batch_invariant import init_batch_invariance

1148

1149 init_batch_invariance()

1150 - override_envs_for_eplb(parallel_config)

1150 + override_envs_for_eplb(

1151 + parallel_config,

1152 + moe_backend=getattr(vllm_config.kernel_config, "moe_backend", None),

1153 + )

1154 set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)

1155

1156 init_method = distributed_init_method or "env://"