From d9f3eb5b33710f858b0bb8e9525d45eacaa24a92 Mon Sep 17 00:00:00 2001 From: Xinyu Jiang Date: Mon, 2 Nov 2026 03:09:42 +0000 Subject: [PATCH 1/5] [AMD] ci: run DSv4 BSHD/THD parity in the MI350 nightly Use the existing AMD DeepSeek-V4 launcher on ROCm and register the original four-GPU BSHD/THD parity test. Keep optimizer and training offload at the launcher defaults. Drop the explicit H200 constructor argument. The generic launcher defaults to auto and resolves H200 on the registered H200 runner. Preserve the original workload and all numerical comparisons. Co-authored-by: Zhiyao Jiang --- .../test_dsv4_bshd_thd_rollout_parity.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tests/e2e/precision/test_dsv4_bshd_thd_rollout_parity.py b/tests/e2e/precision/test_dsv4_bshd_thd_rollout_parity.py index 4b94968054e..1ea59f9b3bf 100644 --- a/tests/e2e/precision/test_dsv4_bshd_thd_rollout_parity.py +++ b/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, From abba40c0632a1e32da72074bc04af572de1bdf87 Mon Sep 17 00:00:00 2001 From: Xinyu Jiang Date: Mon, 2 Nov 2026 03:09:42 +0000 Subject: [PATCH 2/5] [AMD] ci: run session v1/v2 parity in the MI350 nightly Use the existing deterministic Triton backend on ROCm and register the original session parity file in the MI350 nightly. Select the backend with MILES_HARDWARE_PLATFORM. The H200-only assertion is redundant with CUDA CI hardware selection and excludes H100 despite its FA3 support; the test compares v1/v2 payloads on the same server and GPU. Set the ROCm estimate to 200 seconds from the recorded 136.119-second run, using a 30% margin rounded up to the next hundred seconds. Co-authored-by: Zhiyao Jiang --- tests/e2e/sglang/test_session_v1_v2_parity.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/sglang/test_session_v1_v2_parity.py b/tests/e2e/sglang/test_session_v1_v2_parity.py index c7387fdf47c..e912e5acc78 100644 --- a/tests/e2e/sglang/test_session_v1_v2_parity.py +++ b/tests/e2e/sglang/test_session_v1_v2_parity.py @@ -3,7 +3,7 @@ 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 @@ -11,6 +11,7 @@ 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", From c026f496d91a34962ef500c0aa9d29ea911ee09f Mon Sep 17 00:00:00 2001 From: Xinyu Jiang Date: Mon, 2 Nov 2026 03:09:42 +0000 Subject: [PATCH 3/5] [AMD] ci: run DeepSeek V3.2 FP8 smoke nightly on MI350 Select Triton DSA prefill and decode only on ROCm; keep the CUDA launch and the original five-layer, four-GPU, two-rollout test intact. Register the measured full-file smoke with est_time=800. Co-authored-by: Zhiyao Jiang --- .../model_scripts/test_deepseek_v32_5layer_fp8.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py b/tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py index 7baa33cb7e8..25b713a0f45 100644 --- a/tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py +++ b/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) From 78f7ad318bb079646d5bd6c2cf5ebde587a16ee5 Mon Sep 17 00:00:00 2001 From: Xinyu Jiang Date: Mon, 2 Nov 2026 03:18:35 +0000 Subject: [PATCH 4/5] [AMD] fsdp: call the current SGLang attention kernel directly Import the pinned SGLang kernel from its current module and pass unit KV scales explicitly for the BF16 attention bridge. Co-authored-by: Zhiyao Jiang --- .../sglang_attn_bridge/hf_sglang_triton_patch.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/sglang_attn_bridge/hf_sglang_triton_patch.py b/miles/backends/fsdp_utils/sglang_attn_bridge/hf_sglang_triton_patch.py index 690aa35f18f..0084748ac15 100644 --- a/miles/backends/fsdp_utils/sglang_attn_bridge/hf_sglang_triton_patch.py +++ b/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, ) From afb3bbaabaaf4415ac53622047d6520c14d1e4a3 Mon Sep 17 00:00:00 2001 From: Xinyu Jiang Date: Mon, 2 Nov 2026 19:21:47 +0000 Subject: [PATCH 5/5] [AMD] docker: bump ROCm bases to SGLang v0.5.23 Use the matching release base for each ROCm variant alongside the current sglang-miles source stack. Update the default and variant examples. Co-authored-by: Zhiyao Jiang --- docker/Dockerfile.rocm | 8 ++++---- docker/build.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index a25cfcdb358..627a8d69cca 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -2,20 +2,20 @@ # ROCm image for Miles (built on a prebuilt sglang base): # # 1. rocm724-mi35x (gfx950, ROCm 7.2.4 / MI350 / MI355X, Python 3.12 base) -# build-arg:SGLANG_IMAGE_REPO=lmsysorg/sglang SGLANG_IMAGE_TAG=v0.5.22-rocm724-mi35x +# build-arg:SGLANG_IMAGE_REPO=lmsysorg/sglang SGLANG_IMAGE_TAG=v0.5.23-rocm724-mi35x # WHEELS_TAG_ROCM=rocm724-gfx950-v0.5.20 TE_USE_WHEEL=1 APPLY_ROCR_VMMFIX=1 # # 2. rocm10-mi35x (gfx950, ROCm 10 / MI350 / MI355X, Python 3.12 base) -# build-arg:SGLANG_IMAGE_REPO=lmsysorg/sglang SGLANG_IMAGE_TAG=v0.5.22-rocm10-mi35x +# build-arg:SGLANG_IMAGE_REPO=lmsysorg/sglang SGLANG_IMAGE_TAG=v0.5.23-rocm10-mi35x # WHEELS_TAG_ROCM=rocm10-gfx950-v0.5.18 TE_USE_WHEEL=1 APEX_USE_PREBUILT=1 (cp312 shelf) # # 3. rocm10-mi30x (gfx942, ROCm 10 / MI300X / MI325X, Python 3.12 base) -# build-arg:GPU_ARCH=gfx942 SGLANG_IMAGE_REPO=lmsysorg/sglang SGLANG_IMAGE_TAG=v0.5.22-rocm10-mi30x +# build-arg:GPU_ARCH=gfx942 SGLANG_IMAGE_REPO=lmsysorg/sglang SGLANG_IMAGE_TAG=v0.5.23-rocm10-mi30x # WHEELS_TAG_ROCM=rocm10-gfx942-v0.5.20 TE_USE_WHEEL=1 APEX_USE_PREBUILT=1 (cp312 shelf) # AITER_PREBUILD_JIT=1 ARG SGLANG_IMAGE_REPO=lmsysorg/sglang -ARG SGLANG_IMAGE_TAG=v0.5.22-rocm724-mi35x +ARG SGLANG_IMAGE_TAG=v0.5.23-rocm724-mi35x FROM ${SGLANG_IMAGE_REPO}:${SGLANG_IMAGE_TAG} AS sglang SHELL ["/bin/bash", "-ceuxo", "pipefail"] diff --git a/docker/build.py b/docker/build.py index c891d966563..ca392897f2f 100644 --- a/docker/build.py +++ b/docker/build.py @@ -61,7 +61,7 @@ "build_args": { "GPU_ARCH": "gfx950", "SGLANG_IMAGE_REPO": "lmsysorg/sglang", - "SGLANG_IMAGE_TAG": "v0.5.22-rocm724-mi35x", + "SGLANG_IMAGE_TAG": "v0.5.23-rocm724-mi35x", "WHEELS_TAG_ROCM": "rocm724-gfx950-v0.5.20", "APPLY_ROCR_VMMFIX": "1", "TE_USE_WHEEL": "1", @@ -75,7 +75,7 @@ "build_args": { "GPU_ARCH": "gfx950", "SGLANG_IMAGE_REPO": "lmsysorg/sglang", - "SGLANG_IMAGE_TAG": "v0.5.22-rocm10-mi35x", + "SGLANG_IMAGE_TAG": "v0.5.23-rocm10-mi35x", "WHEELS_TAG_ROCM": "rocm10-gfx950-v0.5.18", "APEX_USE_PREBUILT": "1", "TE_USE_WHEEL": "1", @@ -89,7 +89,7 @@ "build_args": { "GPU_ARCH": "gfx942", "SGLANG_IMAGE_REPO": "lmsysorg/sglang", - "SGLANG_IMAGE_TAG": "v0.5.22-rocm10-mi30x", + "SGLANG_IMAGE_TAG": "v0.5.23-rocm10-mi30x", "WHEELS_TAG_ROCM": "rocm10-gfx942-v0.5.20", "APEX_USE_PREBUILT": "1", "TE_USE_WHEEL": "1",