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", 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, ) 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) 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, 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",