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 13, 09:30 UTC
[EPLB][ROCm]: support EPBL for ROCm backend (#27731)
Signed-off-by: Perry Zhang <perzhang@amd.com> Co-authored-by: Perry Zhang <perzhang@amd.com>
PerryZhang01 committed
Nov 13, 2025 at 02:16 UTC
a1e7fa362a514589d4042f2287bde2a0653eb959
4 files changed
+23
-6
vllm/config/parallel.py
+2
-2
@@ -278,10 +278,10 @@ class ParallelConfig:
278
)
279
280
if self.enable_eplb:
281
- if not current_platform.is_cuda():
281
+ if not current_platform.is_cuda_alike():
282
raise ValueError(
283
"Expert parallelism load balancing is only supported on "
284
- "CUDA devices now."
284
+ "CUDA devices or ROCm devices now."
285
)
286
if not self.enable_expert_parallel:
287
raise ValueError("enable_expert_parallel must be True to use EPLB.")
vllm/model_executor/layers/fused_moe/layer.py
+5
-1
@@ -1218,7 +1218,11 @@ class FusedMoE(CustomOp):
1218
1219
def get_expert_weights(self) -> Iterable[torch.Tensor]:
1220
weights = list(self.named_parameters())
1221
- assert all(weight.is_contiguous() for _, weight in weights)
1221
+ assert all(
1222
+ weight.is_contiguous()
1223
+ for name, weight in weights
1224
+ if not name.startswith("_shared_experts.")
1225
+ )
1226
1227
# Filter out the non-expert weights.
1228
# `e_score_correction_bias` is a bias for each logical expert,
vllm/model_executor/layers/quantization/compressed_tensors/__init__.py
+3
@@ -0,0 +1,3 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+13
-3
@@ -1019,9 +1019,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1019
logical_replica_count: torch.Tensor | None = None,
1020
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1021
if enable_eplb:
1022
- raise NotImplementedError(
1023
- "EPLB not supported for `CompressedTensorsW8A8Fp8MoEMethod` yet."
1024
- )
1022
+ assert expert_load_view is not None
1023
+ assert logical_to_physical_map is not None
1024
+ assert logical_replica_count is not None
1025
+ assert isinstance(layer, FusedMoE)
1026
1027
topk_weights, topk_ids, _ = FusedMoE.select_experts(
1028
hidden_states=x,
@@ -1037,6 +1038,11 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1038
e_score_correction_bias=e_score_correction_bias,
1039
indices_type=self.topk_indices_dtype,
1040
num_fused_shared_experts=layer.num_fused_shared_experts,
1041
+ enable_eplb=enable_eplb,
1042
+ expert_map=expert_map,
1043
+ expert_load_view=expert_load_view,
1044
+ logical_to_physical_map=logical_to_physical_map,
1045
+ logical_replica_count=logical_replica_count,
1046
)
1047
1048
per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
@@ -1145,6 +1151,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1151
quant_config=self.moe_quant_config,
1152
)
1153
1154
+ @property
1155
+ def supports_eplb(self) -> bool:
1156
+ return True
1157
+
1158
1159
class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1160
def __init__(