diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index df1b0bed4..381cf7fec 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -251,7 +251,7 @@ jobs: fail-fast: false matrix: runner: [linux-mi325-gpu-1] - part: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + part: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] runs-on: ${{matrix.runner}} steps: - name: Checkout code @@ -273,7 +273,7 @@ jobs: - name: Run test timeout-minutes: 30 run: | - bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 13 --timeout-per-file 1800 + bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 14 --timeout-per-file 1800 stage-b-test-small-1-gpu-amd-mi35x: needs: [check-changes, stage-a-test-1-amd] @@ -484,7 +484,7 @@ jobs: docker exec ci_sglang rocm-smi --showmeminfo vram 2>/dev/null || echo "rocm-smi not available" - name: Run diffusion server tests (1-GPU) - timeout-minutes: 45 + timeout-minutes: 60 run: | # AMD CI: All 1-GPU tests except FLUX.2 (FLUX.1 covers same code path) # Tests: T2V, T2I, I2V, LoRA diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 1357db036..d2e47479d 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -21,7 +21,7 @@ ENV BUILD_TRITON="0" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" -ENV AITER_COMMIT="v0.1.9.post1" +ENV AITER_COMMIT="v0.1.10.post2" # =============================== # Base image 950 and args @@ -31,7 +31,7 @@ ENV BUILD_TRITON="0" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="0" ENV BUILD_MOONCAKE="1" -ENV AITER_COMMIT="v0.1.9.post1" +ENV AITER_COMMIT="v0.1.10.post2" # =============================== # Chosen arch and args FROM ${GPU_ARCH} diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index a99e230e2..cf867a6a1 100644 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -268,6 +268,7 @@ class AiterAttnBackend(AttentionBackend): self, qo_indptr, kv_indptr, + kv_last_page_len, work_metadata, work_info_set, work_indptr, @@ -287,6 +288,7 @@ class AiterAttnBackend(AttentionBackend): meta = get_mla_metadata_v1( qo_indptr, kv_indptr, + kv_last_page_len, self.num_head // nhead_kv, nhead_kv, True, @@ -367,6 +369,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( qo_indptr, kv_indptr, + kv_last_page_len, work_metadata, work_info_set, work_indptr, @@ -423,6 +426,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( qo_indptr, kv_indptr, + self.kv_last_page_len[:bs], work_metadata, work_info_set, work_indptr, @@ -518,6 +522,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( qo_indptr, kv_indptr, + self.kv_last_page_len[:bs], work_metadata, work_info_set, work_indptr, @@ -716,6 +721,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( qo_indptr, kv_indptr, + kv_last_page_len, self.work_metadata, self.work_info_set, self.work_indptr, @@ -786,6 +792,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( qo_indptr, kv_indptr, + kv_last_page_len, self.work_metadata, self.work_info_set, self.work_indptr, @@ -872,6 +879,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( qo_indptr, kv_indptr, + kv_last_page_len, self.work_metadata, self.work_info_set, self.work_indptr, @@ -1144,6 +1152,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( self.forward_metadata.qo_indptr, self.forward_metadata.kv_indptr, + self.forward_metadata.kv_last_page_len, work_metadata, work_info_set, work_indptr, @@ -1165,8 +1174,8 @@ class AiterAttnBackend(AttentionBackend): self.forward_metadata.kv_indices, self.forward_metadata.kv_last_page_len, self.forward_metadata.max_q_len, - layer.scaling, - layer.logit_cap, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, work_meta_data=work_metadata, work_indptr=work_indptr, work_info_set=work_info_set, @@ -1195,6 +1204,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( self.forward_metadata.qo_indptr, self.forward_metadata.kv_indptr, + self.forward_metadata.kv_last_page_len, work_metadata, work_info_set, work_indptr, @@ -1232,8 +1242,8 @@ class AiterAttnBackend(AttentionBackend): self.forward_metadata.kv_indices, self.forward_metadata.kv_last_page_len, self.forward_metadata.max_q_len, - layer.scaling, - layer.logit_cap, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, work_meta_data=work_metadata, work_indptr=work_indptr, work_info_set=work_info_set, @@ -1262,8 +1272,8 @@ class AiterAttnBackend(AttentionBackend): self.forward_metadata.kv_indices, self.forward_metadata.kv_last_page_len, self.forward_metadata.max_q_len, - layer.scaling, - layer.logit_cap, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, work_meta_data=work_metadata, work_indptr=work_indptr, work_info_set=work_info_set, @@ -1353,6 +1363,7 @@ class AiterAttnBackend(AttentionBackend): self.make_mla_meta_data( self.forward_metadata.qo_indptr, self.forward_metadata.kv_indptr, + self.forward_metadata.kv_last_page_len, work_metadata, work_info_set, work_indptr, @@ -1374,8 +1385,8 @@ class AiterAttnBackend(AttentionBackend): self.forward_metadata.kv_indices, self.forward_metadata.kv_last_page_len, self.forward_metadata.max_q_len, - layer.scaling, - layer.logit_cap, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, work_meta_data=work_metadata, work_indptr=work_indptr, work_info_set=work_info_set, diff --git a/scripts/ci/amd/amd_ci_install_dependency.sh b/scripts/ci/amd/amd_ci_install_dependency.sh index 52273f0f7..bc3215b2e 100755 --- a/scripts/ci/amd/amd_ci_install_dependency.sh +++ b/scripts/ci/amd/amd_ci_install_dependency.sh @@ -233,7 +233,7 @@ if [[ "${NEED_REBUILD}" == "true" ]]; then echo "[CI-AITER-CHECK] === AITER REBUILD START ===" # uninstall existing aiter - docker exec ci_sglang pip uninstall -y aiter || true + docker exec ci_sglang pip uninstall -y amd-aiter || true # delete old aiter directory docker exec ci_sglang rm -rf /sgl-workspace/aiter diff --git a/scripts/ci/amd/amd_ci_warmup_aiter.py b/scripts/ci/amd/amd_ci_warmup_aiter.py index b12d68717..461426013 100755 --- a/scripts/ci/amd/amd_ci_warmup_aiter.py +++ b/scripts/ci/amd/amd_ci_warmup_aiter.py @@ -32,10 +32,12 @@ def warmup_aiter_kernels(): device = torch.device("cuda:0") start_time = time.time() - # Warmup RMSNorm kernel (module_rmsnorm) - most commonly used - # SGLang uses rmsnorm2d_fwd and rmsnorm2d_fwd_with_add from aiter + # Warmup module_rmsnorm_quant (small module, ~2MB) + # Triggered by rmsnorm2d_fwd when hidden_size <= 8192 try: - print("\n[1/4] Warming up RMSNorm kernel (rmsnorm2d_fwd)...") + print( + "\n[1/5] Warming up module_rmsnorm_quant (rmsnorm2d_fwd, hidden<=8192)..." + ) from aiter import rmsnorm2d_fwd hidden_size = 4096 @@ -44,37 +46,62 @@ def warmup_aiter_kernels(): weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device) eps = 1e-6 - # This triggers JIT compilation + # hidden_size=4096 <= 8192 -> takes rmsnorm() path -> compiles module_rmsnorm_quant _ = rmsnorm2d_fwd(x, weight, eps) torch.cuda.synchronize() - print(f" RMSNorm kernel (rmsnorm2d_fwd) compiled successfully") + print(" module_rmsnorm_quant compiled successfully") except Exception as e: - print(f" RMSNorm warmup failed (may not be available): {e}") + print(f" module_rmsnorm_quant warmup failed: {e}") - # Warmup fused add RMSNorm kernel + # Warmup module_rmsnorm (large CK module, ~159MB) + # Triggered by rmsnorm2d_fwd_with_add (always uses CK path) + # NOTE: rmsnorm2d_fwd_with_add signature is: + # rmsnorm2d_fwd_with_add(out, input, residual_in, residual_out, weight, epsilon) try: - print("\n[2/4] Warming up fused add RMSNorm kernel (rmsnorm2d_fwd_with_add)...") + print("\n[2/5] Warming up module_rmsnorm (rmsnorm2d_fwd_with_add, CK path)...") from aiter import rmsnorm2d_fwd_with_add hidden_size = 4096 batch_size = 512 x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device=device) - residual = torch.randn( + residual_in = torch.randn( batch_size, hidden_size, dtype=torch.bfloat16, device=device ) + output = torch.empty_like(x) + residual_out = torch.empty_like(x) weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device) eps = 1e-6 - # This triggers JIT compilation - _ = rmsnorm2d_fwd_with_add(x, residual, weight, eps) + # This triggers JIT compilation of module_rmsnorm (CK kernels) + rmsnorm2d_fwd_with_add(output, x, residual_in, residual_out, weight, eps) torch.cuda.synchronize() - print(f" Fused add RMSNorm kernel compiled successfully") + print(" module_rmsnorm compiled successfully") except Exception as e: - print(f" Fused add RMSNorm warmup failed (may not be available): {e}") + print(f" module_rmsnorm warmup failed: {e}") + + # Warmup module_rmsnorm via rmsnorm2d_fwd with large hidden_size (CK path) + # When hidden_size > 8192, rmsnorm2d_fwd takes the rmsnorm2d_fwd_ck path + # which also uses module_rmsnorm (already compiled in step 2, but this + # ensures the CK rmsnorm2d_fwd path is exercised as well) + try: + print("\n[3/5] Warming up rmsnorm2d_fwd CK path (hidden>8192)...") + from aiter import rmsnorm2d_fwd + + hidden_size = 16384 # > 8192 to trigger rmsnorm2d_fwd_ck (module_rmsnorm) + batch_size = 32 + x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device=device) + weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device) + eps = 1e-6 + + _ = rmsnorm2d_fwd(x, weight, eps) + torch.cuda.synchronize() + print(" rmsnorm2d_fwd CK path compiled successfully") + except Exception as e: + print(f" rmsnorm2d_fwd CK path warmup skipped: {e}") # Warmup rotary embedding kernel if available try: - print("\n[3/4] Warming up rotary embedding kernel...") + print("\n[4/5] Warming up rotary embedding kernel...") from aiter import rotary_embedding head_size = 128 @@ -92,13 +119,13 @@ def warmup_aiter_kernels(): _ = rotary_embedding(positions, query, key, head_size, cos, sin, True) torch.cuda.synchronize() - print(f" Rotary embedding kernel compiled successfully") + print(" Rotary embedding kernel compiled successfully") except Exception as e: print(f" Rotary embedding warmup skipped (may not be available): {e}") # Warmup activation kernels if available try: - print("\n[4/4] Warming up activation kernels...") + print("\n[5/5] Warming up activation kernels...") from aiter import silu_and_mul hidden_size = 4096 @@ -110,7 +137,7 @@ def warmup_aiter_kernels(): silu_and_mul(out, x) torch.cuda.synchronize() - print(f" Activation kernel compiled successfully") + print(" Activation kernel compiled successfully") except Exception as e: print(f" Activation warmup skipped (may not be available): {e}") diff --git a/test/registered/attention/test_triton_attention_backend.py b/test/registered/attention/test_triton_attention_backend.py index dd5bd5b48..52757d383 100644 --- a/test/registered/attention/test_triton_attention_backend.py +++ b/test/registered/attention/test_triton_attention_backend.py @@ -21,7 +21,7 @@ from sglang.test.test_utils import ( # Triton attention backend integration test with latency benchmark and MMLU eval register_cuda_ci(est_time=200, suite="stage-b-test-large-1-gpu") -register_amd_ci(est_time=1110, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=1400, suite="stage-b-test-small-1-gpu-amd") class TestTritonAttnBackend(CustomTestCase): diff --git a/test/registered/dllm/test_llada2_mini_amd.py b/test/registered/dllm/test_llada2_mini_amd.py index 88309a48e..ad7aeae40 100644 --- a/test/registered/dllm/test_llada2_mini_amd.py +++ b/test/registered/dllm/test_llada2_mini_amd.py @@ -20,7 +20,7 @@ from sglang.test.test_utils import ( write_github_step_summary, ) -register_amd_ci(est_time=520, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=1000, suite="stage-b-test-small-1-gpu-amd") class TestLLaDA2MiniAMD(CustomTestCase): diff --git a/test/registered/eval/test_eval_accuracy_large.py b/test/registered/eval/test_eval_accuracy_large.py index e7328c29f..6a0a39581 100644 --- a/test/registered/eval/test_eval_accuracy_large.py +++ b/test/registered/eval/test_eval_accuracy_large.py @@ -21,7 +21,7 @@ from sglang.test.test_utils import ( ) register_cuda_ci(est_time=300, suite="stage-b-test-small-1-gpu") -register_amd_ci(est_time=300, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=420, suite="stage-b-test-small-1-gpu-amd") class TestEvalAccuracyLarge(CustomTestCase): diff --git a/test/registered/mla/test_mla_fp8.py b/test/registered/mla/test_mla_fp8.py index d55d64ee8..0135c7edb 100644 --- a/test/registered/mla/test_mla_fp8.py +++ b/test/registered/mla/test_mla_fp8.py @@ -14,7 +14,7 @@ from sglang.test.test_utils import ( # MLA FP8 KV cache test with MGSM evaluation register_cuda_ci(est_time=77, suite="stage-b-test-large-1-gpu") -register_amd_ci(est_time=360, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=800, suite="stage-b-test-small-1-gpu-amd") class TestMLA(CustomTestCase): diff --git a/test/registered/models/test_vlm_models.py b/test/registered/models/test_vlm_models.py index f0a9c76d9..ffcba0acb 100644 --- a/test/registered/models/test_vlm_models.py +++ b/test/registered/models/test_vlm_models.py @@ -17,7 +17,7 @@ from sglang.test.test_utils import is_in_ci register_cuda_ci(est_time=228, suite="stage-b-test-large-1-gpu") -register_amd_ci(est_time=420, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=850, suite="stage-b-test-small-1-gpu-amd") _is_hip = is_hip() # VLM models for testing diff --git a/test/registered/perf/test_bench_serving_2gpu.py b/test/registered/perf/test_bench_serving_2gpu.py index 3c8cc216a..00e8a058b 100644 --- a/test/registered/perf/test_bench_serving_2gpu.py +++ b/test/registered/perf/test_bench_serving_2gpu.py @@ -15,7 +15,7 @@ from sglang.test.test_utils import ( ) register_cuda_ci(est_time=600, suite="stage-b-test-large-2-gpu") -register_amd_ci(est_time=600, suite="stage-b-test-large-2-gpu-amd") +register_amd_ci(est_time=1100, suite="stage-b-test-large-2-gpu-amd") class TestBenchServing2GPU(CustomTestCase): diff --git a/test/registered/quant/test_eval_fp8_accuracy.py b/test/registered/quant/test_eval_fp8_accuracy.py index f4f02ccee..22b183785 100644 --- a/test/registered/quant/test_eval_fp8_accuracy.py +++ b/test/registered/quant/test_eval_fp8_accuracy.py @@ -15,7 +15,7 @@ from sglang.test.test_utils import ( ) register_cuda_ci(est_time=250, suite="stage-b-test-large-1-gpu") -register_amd_ci(est_time=303, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=600, suite="stage-b-test-small-1-gpu-amd") class TestEvalFP8Accuracy(CustomTestCase): diff --git a/test/registered/quant/test_torchao.py b/test/registered/quant/test_torchao.py index 9f1ed4133..f6e70ee06 100644 --- a/test/registered/quant/test_torchao.py +++ b/test/registered/quant/test_torchao.py @@ -7,7 +7,7 @@ from sglang import Engine from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=103, suite="stage-b-test-small-1-gpu") -register_amd_ci(est_time=106, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=230, suite="stage-b-test-small-1-gpu-amd") from sglang.lang.chat_template import get_chat_template_by_model_path from sglang.srt.utils import kill_process_tree from sglang.test.run_eval import run_eval diff --git a/test/registered/scheduler/test_abort.py b/test/registered/scheduler/test_abort.py index e57119efc..da4c70adb 100644 --- a/test/registered/scheduler/test_abort.py +++ b/test/registered/scheduler/test_abort.py @@ -18,7 +18,7 @@ from sglang.test.test_utils import ( ) register_cuda_ci(est_time=131, suite="stage-b-test-small-1-gpu") -register_amd_ci(est_time=51, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=300, suite="stage-b-test-small-1-gpu-amd") class TestAbort(CustomTestCase): diff --git a/test/registered/scheduler/test_retract_decode.py b/test/registered/scheduler/test_retract_decode.py index 80dc74cec..2c59ec247 100644 --- a/test/registered/scheduler/test_retract_decode.py +++ b/test/registered/scheduler/test_retract_decode.py @@ -18,7 +18,7 @@ from sglang.test.test_utils import ( from sglang.utils import is_in_ci register_cuda_ci(est_time=311, suite="stage-b-test-small-1-gpu") -register_amd_ci(est_time=450, suite="stage-b-test-small-1-gpu-amd") +register_amd_ci(est_time=600, suite="stage-b-test-small-1-gpu-amd") class TestRetractDecode(CustomTestCase):