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 6, 05:45 UTC
[ROCm] Enable DeepEP ROCm as all2allbackend for AMD GPUs. (#34692)
Signed-off-by: Tej Kiran <vpolamre@amd.com> Co-authored-by: Tej Kiran <vpolamre@amd.com>
Chaitanya Sri Krishna Lolla committed
Mar 21, 2026 at 13:02 UTC
3982bc2cd0bd9d633060b22e9ff683d8316a0f82
7 files changed
+68
-29
docker/Dockerfile.rocm_base
+1
-1
@@ -44,7 +44,7 @@ ENV DEBIAN_FRONTEND=noninteractive
44
45
# Install Python and other dependencies
46
RUN apt-get update -y \
47
- && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev \
47
+ && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \
48
&& for i in 1 2 3; do \
49
add-apt-repository -y ppa:deadsnakes/ppa && break || \
50
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
vllm/distributed/device_communicators/all2all.py
+19
-6
@@ -10,6 +10,7 @@ import vllm.envs as envs
10
from vllm.distributed import get_dp_group, get_ep_group
11
from vllm.forward_context import get_forward_context
12
from vllm.logger import init_logger
13
+from vllm.platforms import current_platform
14
from vllm.utils.flashinfer import (
15
has_flashinfer_nvlink_one_sided,
16
has_flashinfer_nvlink_two_sided,
@@ -325,14 +326,20 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
326
327
assert num_rdma_bytes is not None
328
assert num_qps_per_rank is not None
328
- return dict(
329
+ # TODO: remove platform-specific logic
330
+ # once ROCm DeepEP is updated with the latest APIs.
331
+ kwargs = dict(
332
group=self.cpu_group,
333
num_nvl_bytes=num_nvl_bytes,
334
num_rdma_bytes=num_rdma_bytes,
335
low_latency_mode=False,
336
num_qps_per_rank=num_qps_per_rank,
334
- explicitly_destroy=True,
337
)
338
+ if not current_platform.is_rocm():
339
+ kwargs.update(
340
+ explicitly_destroy=True,
341
+ )
342
+ return kwargs
343
344
def get_handle(self, kwargs):
345
assert len(kwargs) == 0, (
@@ -397,16 +404,22 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
404
)
405
406
assert num_rdma_bytes is not None
400
- return dict(
407
+ # TODO: remove platform-specific logic
408
+ # once ROCm DeepEP is updated with the latest APIs.
409
+ kwargs = dict(
410
group=self.cpu_group,
411
num_nvl_bytes=num_nvl_bytes,
412
num_rdma_bytes=num_rdma_bytes,
413
low_latency_mode=True,
414
num_qps_per_rank=num_qps_per_rank,
406
- allow_nvlink_for_low_latency_mode=True,
407
- allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
408
- explicitly_destroy=True,
415
)
416
+ if not current_platform.is_rocm():
417
+ kwargs.update(
418
+ allow_nvlink_for_low_latency_mode=True,
419
+ allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
420
+ explicitly_destroy=True,
421
+ )
422
+ return kwargs
423
424
def get_handle(self, kwargs):
425
"""
vllm/model_executor/layers/fused_moe/config.py
+2
-2
@@ -346,7 +346,7 @@ class FusedMoEQuantConfig:
346
347
@property
348
def use_fp8_w8a8(self) -> bool:
349
- return self.quant_dtype == torch.float8_e4m3fn
349
+ return self.quant_dtype == current_platform.fp8_dtype()
350
351
@property
352
def use_int8_w8a8(self) -> bool:
@@ -566,7 +566,7 @@ def fp8_w8a8_moe_quant_config(
566
Construct a quant config for fp8 activations and fp8 weights.
567
"""
568
return FusedMoEQuantConfig.make(
569
- torch.float8_e4m3fn,
569
+ current_platform.fp8_dtype(),
570
w1_scale=w1_scale,
571
g1_alphas=g1_alphas,
572
w2_scale=w2_scale,
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+41
-17
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
16
moe_kernel_quantize_input,
17
normalize_batched_scales_shape,
18
)
19
+from vllm.platforms import current_platform
20
from vllm.v1.worker.ubatching import (
21
dbo_current_ubatch_id,
22
dbo_enabled,
@@ -290,23 +291,46 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
291
292
# Dispatch
293
dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids)
293
- expert_x, expert_num_tokens, handle, _, hook = self.buffer.low_latency_dispatch(
294
- a1,
295
- dispatch_topk_ids,
296
- self.max_tokens_per_rank,
297
- num_experts,
298
- use_fp8=self.use_fp8_dispatch,
299
- round_scale=self.use_ue8m0_dispatch,
300
- use_ue8m0=self.use_ue8m0_dispatch,
301
- **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
302
- **(
303
- dict(x_global_scale=qc_a1_gscale_or_scale)
304
- if qc_a1_gscale_or_scale is not None
305
- else dict()
306
- ),
307
- async_finish=False,
308
- return_recv_hook=True,
309
- )
294
+ if current_platform.is_rocm():
295
+ (
296
+ expert_x,
297
+ expert_num_tokens,
298
+ handle,
299
+ _,
300
+ hook,
301
+ ) = self.buffer.low_latency_dispatch(
302
+ a1,
303
+ dispatch_topk_ids,
304
+ self.max_tokens_per_rank,
305
+ num_experts,
306
+ use_fp8=self.use_fp8_dispatch,
307
+ async_finish=False,
308
+ return_recv_hook=True,
309
+ )
310
+ else:
311
+ (
312
+ expert_x,
313
+ expert_num_tokens,
314
+ handle,
315
+ _,
316
+ hook,
317
+ ) = self.buffer.low_latency_dispatch(
318
+ a1,
319
+ dispatch_topk_ids,
320
+ self.max_tokens_per_rank,
321
+ num_experts,
322
+ use_fp8=self.use_fp8_dispatch,
323
+ round_scale=self.use_ue8m0_dispatch,
324
+ use_ue8m0=self.use_ue8m0_dispatch,
325
+ **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
326
+ **(
327
+ dict(x_global_scale=qc_a1_gscale_or_scale)
328
+ if qc_a1_gscale_or_scale is not None
329
+ else dict()
330
+ ),
331
+ async_finish=False,
332
+ return_recv_hook=True,
333
+ )
334
self.handles[a2a_idx] = handle
335
336
return (
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+2
-1
@@ -1017,6 +1017,7 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
1017
torch.float16,
1018
torch.bfloat16,
1019
torch.float8_e4m3fn,
1020
+ torch.float8_e4m3fnuz,
1021
]
1022
assert expert_tokens_meta is not None
1023
@@ -1046,7 +1047,7 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
1047
compute_type = tl.float16
1048
elif hidden_states.dtype == torch.float32:
1049
compute_type = tl.float32
1049
- elif hidden_states.dtype == torch.float8_e4m3fn:
1050
+ elif hidden_states.dtype == current_platform.fp8_dtype():
1051
compute_type = tl.bfloat16
1052
else:
1053
raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-1
@@ -1616,7 +1616,7 @@ def _get_config_quant_dtype(
1616
fused_experts_impl.
1617
"""
1618
if use_fp8_w8a8:
1619
- return torch.float8_e4m3fn
1619
+ return current_platform.fp8_dtype()
1620
elif use_int8_w8a8:
1621
return torch.int8
1622
elif ocp_mx_scheme == "w_mxfp4_a_mxfp4":
vllm/model_executor/layers/fused_moe/utils.py
+2
-1
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
25
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
26
per_tensor_dequantize,
27
)
28
+from vllm.platforms import current_platform
29
from vllm.triton_utils import tl, triton
30
from vllm.utils.math_utils import cdiv
31
from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -265,7 +266,7 @@ def moe_kernel_quantize_input(
266
# weights are already dequantized, and we proceed with normal
267
# activation quantization below.
268
268
- if quant_dtype == torch.float8_e4m3fn:
269
+ if quant_dtype == current_platform.fp8_dtype():
270
return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
271
elif quant_dtype == torch.int8:
272
return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)