Repository navigation
[AMD] ci: bump ROCm images to SGLang v0.5.23 and expand MI350 nightly coverage - #4599
Merged
guapisolo merged 5 commits intoNov 2, 2026
Merged
Conversation
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 <jessicajiang324@gmail.com>
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 <jessicajiang324@gmail.com>
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 <jessicajiang324@gmail.com>
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 <jessicajiang324@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
XinyuJiangCMU
requested review from
Rockdu,
Zhichenzzz and
guapisolo
as code owners
November 2, 2026 03:19
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 <jessicajiang324@gmail.com>
XinyuJiangCMU
changed the title
guapisolo
approved these changes
Nov 2, 2026
This was referenced Nov 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Register three more existing Miles tests in the AMD MI350 nightly, fix the ROCm FSDP Triton attention bridge for the current SGLang, and move the ROCm images to the SGLang v0.5.23 bases.
tests/e2e/precision/test_dsv4_bshd_thd_rollout_parity.py: DeepSeek-V4 rollout followed by one training step in both the BSHD (padded) and THD (packed) layouts, comparing log-probabilities, gradient norm, loss and PPO KL. On ROCm it uses the existing AMD DeepSeek-V4 launcher. The explicithardware="H200"is dropped; the launcher defaultautoresolves to H200 on the registered CUDA runner. Four-GPU nightly.tests/e2e/sglang/test_session_v1_v2_parity.py: v1/v2 session payload parity on Qwen3-8B. On ROCm, SGLang uses Triton attention instead of FlashAttention 3 (FA3). The H200-only device assertion is removed, since CUDA CI already runs this file on Hopper. Two-GPU nightly.tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py: five-layer DeepSeek-V3.2 FP8 smoke. On ROCm, SGLang uses Triton for DeepSeek sparse attention (DSA) prefill and decode, because the default FlashMLA sparse kernels are CUDA-only. Four-GPU nightly.miles/backends/fsdp_utils/sglang_attn_bridge/hf_sglang_triton_patch.py: on ROCm, FSDP training with Triton attention calls SGLang's Triton attention kernel through this bridge, which still imported the kernel from a module SGLang has removed. It now imports the kernel from its current module and passes the unit KV scales that the current kernel expects for BF16 attention.docker/build.py,docker/Dockerfile.rocm: move every ROCm variant to its SGLangv0.5.23release base, following build(docker): bump SGLang to v0.5.23 #4548. That change moved the CUDA image tov0.5.23but kept thev0.5.22ROCm bases, while the ROCm image checks out the samesglang-milesstack, so the kernels prebuilt in the base and the SGLang source came from different releases. Thev0.5.23base contains thesglang-milesv0.5.23fork point with no ahead-of-time kernel changes in between, and keeps torch2.11.0+rocm10.0.0, so the existing ROCm wheels still apply.Validation
All three tests passed locally with ROCm 10 on MI350X/MI355X. With the bridge fix, forward and backward through the bridge complete on MI355X with finite, nonzero q/k/v gradients; the old import fails.
For item 5, an image built on the new ROCm 10 MI350 base with the current
sglang-milespasses the DSv4 BSHD/THD parity and session v1/v2 parity tests above, andtests/e2e/megatron/test_mimo_7B_mtp_only_grad.py, which exercises speculative decoding. The DeepSeek V3.2 smoke has not been rerun on the new base yet.Co-author-with: @JessicaJiang-123