[AMD] ci: bump ROCm images to SGLang v0.5.23 and expand MI350 nightly coverage

miles/backends/fsdp_utils/sglang_attn_bridge/hf_sglang_triton_patch.py

@@ -36,7 +36,7 @@ def forward(ctx, q, k, v, B, S):
         Returns:
             o: [total, num_heads, D], bf16
         """
-        from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd_unified
+        from sglang.kernels.ops.attention.extend_attention import extend_attention_fwd_unified
 
         total = B * S
         device = q.device
@@ -52,10 +52,12 @@ def forward(ctx, q, k, v, B, S):
             o,
             k,
             v,
-            qo_indptr,
-            kv_indptr,
-            kv_indices,
-            prefix_lens,
+            k_scale=1.0,
+            v_scale=1.0,
+            qo_indptr=qo_indptr,
+            kv_indptr=kv_indptr,
+            kv_indices=kv_indices,
+            prefix_lens=prefix_lens,
             max_len_extend=S,
             is_causal=True,
         )

tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py

@@ -13,7 +13,7 @@
     _prepare_download,
     _prepare_megatron_ckpt,
 )
-from tests.ci.ci_register import register_cuda_ci
+from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
 from tests.ci.metric_history import register_ci_gate
 
 from miles.utils.external_utils import command_utils
@@ -22,6 +22,8 @@
     est_time=1300, suite="stage-c-4-gpu-h200", labels=["megatron", "model-scripts"], hardware=["hopper"], num_gpus=4
 )
 
+register_rocm_ci(est_time=800, suite="nightly-stage-c-4-gpu-mi350", labels=["megatron", "model-scripts"])
+
 register_ci_gate(metric_key="train/grad_norm")
 register_ci_gate(metric_key="train/ppo_kl")
 register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff")
@@ -31,7 +33,7 @@
 
 
 def _args() -> ScriptArgs:
-    return ScriptArgs.from_env(
+    args = ScriptArgs.from_env(
         model_org="Pinaster",
         model_name="DeepSeek-V3.2-5layer",
         megatron_model_type="deepseek-v32-5layer",
@@ -55,6 +57,10 @@ def _args() -> ScriptArgs:
         ),
     )
 
+    if os.getenv("MILES_HARDWARE_PLATFORM") == "rocm":
+        args.extra_args += "--sglang-dsa-decode-backend triton --sglang-dsa-prefill-backend triton "
+    return args
+
 
 def prepare(args: ScriptArgs):
     _prepare_download(args)

tests/e2e/precision/test_dsv4_bshd_thd_rollout_parity.py

@@ -6,15 +6,13 @@
 from pathlib import Path
 
 import torch
-from scripts.run_deepseek_v4 import (
-    _DSV4_TE_PRECISION_CONFIG,
-    ScriptArgs,
-    _prepare_download,
-    _prepare_single,
-    _prepare_spmd,
-    _train,
-)
-from tests.ci.ci_register import register_cuda_ci
+from scripts.run_deepseek_v4 import _DSV4_TE_PRECISION_CONFIG
+
+if os.getenv("MILES_HARDWARE_PLATFORM") == "rocm":
+    from scripts.amd.run_deepseek_v4 import ScriptArgs, _prepare_download, _prepare_single, _prepare_spmd, _train
+else:
+    from scripts.run_deepseek_v4 import ScriptArgs, _prepare_download, _prepare_single, _prepare_spmd, _train
+from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
 
 from miles.utils.audit_utils.config_snapshot.enabled import SNAPSHOT_TESTING_ENABLED
 from miles.utils.audit_utils.config_snapshot.generated_value_registry import generated_value_registry
@@ -24,6 +22,7 @@
 register_cuda_ci(
     est_time=1100, suite="stage-c-4-gpu-h200", labels=["precision", "long"], hardware=["hopper"], num_gpus=4
 )
+register_rocm_ci(est_time=800, suite="nightly-stage-c-4-gpu-mi350", labels=["precision", "long"])
 
 _MODEL_NAME = "DeepSeek-V4-Flash-FP8-4layer"
 _NUM_GPUS = 4
@@ -64,7 +63,6 @@ def _prepare_args() -> ScriptArgs:
         enable_eval=False,
         num_nodes=1,
         num_gpus_per_node=_NUM_GPUS,
-        hardware="H200",
         skip_saving=True,
         use_fault_tolerance=False,
     )
@@ -106,7 +104,6 @@ def _run_args(
         enable_eval=False,
         num_nodes=1,
         num_gpus_per_node=_NUM_GPUS,
-        hardware="H200",
         hf_checkpoint=hf_checkpoint,
         skip_saving=True,
         use_fault_tolerance=False,

tests/e2e/sglang/test_session_v1_v2_parity.py

@@ -3,14 +3,15 @@
 import pytest
 import torch
 from huggingface_hub import snapshot_download
-from tests.ci.ci_register import register_cuda_ci
+from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
 from tests.e2e.sglang.utils.sglang_server import start_sglang_server
 from tests.session_parity_utils import V1, V2, assert_agentic_retry_trajectory_parity, run_agentic_retry_trajectories
 
 from miles.utils.test_utils.session_verify_agent import build_initial_messages
 from miles.utils.types import Sample
 
 register_cuda_ci(est_time=190, suite="stage-c-2-gpu-h200", labels=["sglang"], hardware=["hopper"], num_gpus=2)
+register_rocm_ci(est_time=200, suite="nightly-stage-c-2-gpu-mi350", labels=["sglang"])
 
 _MODEL_ID = "Qwen/Qwen3-8B"
 _MODEL_REVISION = "b968826d9c46dd6066d109eabc6255188de91218"
@@ -23,7 +24,6 @@
 @pytest.fixture(scope="module")
 def sglang_server():
     assert torch.cuda.is_available()
-    assert "H200" in torch.cuda.get_device_name(0)
     if _MODEL_PATH_OVERRIDE is None:
         snapshot_download(_MODEL_ID, revision=_MODEL_REVISION, local_dir=_DEFAULT_MODEL_PATH)
 
@@ -32,7 +32,7 @@ def sglang_server():
         enable_deterministic_inference=True,
         extra_args=[
             "--attention-backend",
-            "fa3",
+            "triton" if os.getenv("MILES_HARDWARE_PLATFORM") == "rocm" else "fa3",
             "--reasoning-parser",
             "qwen3",
             "--tool-call-parser",