From d21601c0d93a9e4474a061905c6b845be2ca0bea Mon Sep 17 00:00:00 2001
From: xiaohong42 <940683523@qq.com>
Date: Tue, 9 Jun 2026 03:50:14 +0000
Subject: [PATCH] [rollout, sglang] feat: support sglang ROCm backend via aiter
 defaults and ray init env

On ROCm (torch.version.hip is not None), default the SGLang attention
backend to "aiter" when the user has not set one. The non-attention
kernels (RMSNorm/RoPE/MoE/quant) are routed through AITER via
SGLANG_USE_AITER, which is now injected through the platform abstraction
(PlatformROCm.rollout_env_vars()) instead of os.environ.setdefault: the
SGLang server actor's runtime_env consumes get_platform().rollout_env_vars()
and ray_noset_envvars(), mirroring the vLLM/trtllm rollout servers. Both
remain overridable: an explicit engine_kwargs.sglang.attention_backend or
an exported SGLANG_USE_AITER still win.

Also inject RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 into PPO_RAY_RUNTIME_ENV
on ROCm so Ray 2.x does not force-clear accelerator visibility for the
num_gpus=0 SGLang server actor (otherwise the actor sees no GPU). Scoped
to ROCm to keep Ray's default behavior on other platforms.

AI assistance (Cursor) was used; the change was reviewed line-by-line.

Co-authored-by: Cursor <cursor@cursor.com>
Signed-off-by: xiaohong42 <940683523@qq.com>
---
 verl/trainer/constants_ppo.py                    | 10 ++++++++++
 .../sglang_rollout/async_sglang_server.py        | 16 ++++++++++++----
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/verl/trainer/constants_ppo.py b/verl/trainer/constants_ppo.py
index 92f39710f03..ad77f4f2185 100644
--- a/verl/trainer/constants_ppo.py
+++ b/verl/trainer/constants_ppo.py
@@ -15,6 +15,7 @@
 import json
 import os
 
+import torch
 from ray._private.runtime_env.constants import RAY_JOB_CONFIG_JSON_ENV_VAR
 
 from verl.utils.device import get_device_capability
@@ -27,6 +28,14 @@
 if (_major or 0) >= 10 and os.environ.get("TLLM_DISABLE_NVLS_MNNVL", "0") == "1":
     _gb200_nccl_env = {"NCCL_NVLS_ENABLE": "0", "NCCL_MNNVL_ENABLE": "0"}
 
+# On ROCm, Ray 2.x force-clears accelerator visibility for num_gpus=0 actors
+# (e.g. the SGLang server actor), leaving them unable to see any GPU. Disable
+# that override so the actor keeps its HIP visibility. Scoped to ROCm to avoid
+# changing Ray's default behavior on other platforms.
+_rocm_ray_env = {}
+if torch.version.hip is not None:
+    _rocm_ray_env = {"RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO": "0"}
+
 PPO_RAY_RUNTIME_ENV = {
     "env_vars": {
         "TOKENIZERS_PARALLELISM": "true",
@@ -43,6 +52,7 @@
         "HCCL_NPU_SOCKET_PORT_RANGE": "auto",
         "HSA_NO_SCRATCH_RECLAIM": "1",
         **_gb200_nccl_env,
+        **_rocm_ray_env,
     },
 }
 
diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py
index 2bef78ab18b..171a9b3f2e7 100644
--- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py
+++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py
@@ -42,6 +42,7 @@
 )
 from sglang.srt.managers.tokenizer_manager import ServerStatus
 
+from verl.plugin.platform import get_platform
 from verl.utils.config import omega_conf_to_dataclass
 from verl.utils.device import get_visible_devices_keyword
 from verl.utils.net_utils import get_free_port, is_valid_ipv6_address
@@ -256,9 +257,11 @@ async def launch_server(self, master_address: str = None, master_port: int = Non
         attention_backend = engine_kwargs.pop("attention_backend", None)
         mm_attention_backend = engine_kwargs.pop("mm_attention_backend", None)
         if attention_backend is None:
-            # FA3 CUDA-graph capture is broken on sglang>=0.5.12 (#22800);
-            # default to flashinfer (users can opt into fa4 via engine_kwargs).
-            if version.parse(sglang.__version__) >= version.parse("0.5.12"):
+            if torch.version.hip is not None:
+                attention_backend = "aiter"
+            elif version.parse(sglang.__version__) >= version.parse("0.5.12"):
+                # FA3 CUDA-graph capture is broken on sglang>=0.5.12 (#22800);
+                # default to flashinfer (users can opt into fa4 via engine_kwargs).
                 attention_backend = "flashinfer"
             else:
                 attention_backend = "fa3"
@@ -791,7 +794,12 @@ async def launch_servers(self):
                     node_id=node_id,
                     soft=False,
                 ),
-                runtime_env={"env_vars": {f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}": "1"}},
+                runtime_env={
+                    "env_vars": {
+                        **{var: "1" for var in get_platform().ray_noset_envvars()},
+                        **get_platform().rollout_env_vars(),
+                    }
+                },
                 name=name,
                 max_concurrency=self.max_concurrency,
             ).remote(