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 4, 10:00 UTC
[ROCm] Enable DBO (Dynamic Batch Optimization) on ROCm (#34726)
Signed-off-by: raviguptaamd <ravi.gupta@amd.com>
raviguptaamd committed
May 1, 2026 at 02:18 UTC
7075df79b3094bb6f6d28021c4df8631af10b2b8
2 files changed
+14
-5
vllm/envs.py
+12
-3
@@ -1627,9 +1627,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
1627
"VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool(
1628
int(os.getenv("VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL", "0"))
1629
),
1630
- # The number of SMs to allocate for communication kernels when running DBO
1631
- # the rest of the SMs on the device will be allocated to compute
1632
- "VLLM_DBO_COMM_SMS": lambda: int(os.getenv("VLLM_DBO_COMM_SMS", "20")),
1630
+ # The number of SMs/CUs to allocate for communication kernels when
1631
+ # running DBO; the rest will be allocated to compute.
1632
+ # Default: 20 on CUDA (SMs), 64 on ROCm (CUs).
1633
+ "VLLM_DBO_COMM_SMS": lambda: int(
1634
+ os.getenv(
1635
+ "VLLM_DBO_COMM_SMS",
1636
+ "64"
1637
+ if hasattr(__import__("torch").version, "hip")
1638
+ and __import__("torch").version.hip is not None
1639
+ else "20",
1640
+ )
1641
+ ),
1642
# Enable max_autotune & coordinate_descent_tuning in inductor_config
1643
# to compile static shapes passed from compile_sizes in compilation_config
1644
# If set to 1, enable max_autotune; By default, this is enabled (1)
vllm/v1/worker/gpu_ubatch_wrapper.py
+2
-2
@@ -71,8 +71,8 @@ class SMControlContextManager:
71
A function that sets the number of SMs for computation.
72
"""
73
74
- assert current_platform.is_cuda(), (
75
- "SM control is currently only supported on CUDA"
74
+ assert current_platform.is_cuda() or current_platform.is_rocm(), (
75
+ "SM/CU control is supported on CUDA and ROCm platforms"
76
)
77
device = torch.accelerator.current_device_index()
78
total_sms = num_compute_units(device)