chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Update aiter to v0.1.10.post2 (#18423)

Browse Source 
Co-authored-by: kkHuang-amd <wunhuang@amd.com>
Co-authored-by: YC Tseng <yctseng@amd.com>
This commit is contained in:
Bingxu Chen
2026-02-09 14:08:24 +08:00
committed by GitHub
parent b564dcec61
commit 3f3c201243
15 changed files with 79 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
.github/workflows/pr-test-amd.yml vendored
View File

@@ -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

4
docker/rocm.Dockerfile
View File

@@ -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}

27
python/sglang/srt/layers/attention/aiter_backend.py
View File

@@ -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,

2
scripts/ci/amd/amd_ci_install_dependency.sh
View File

@@ -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

61
scripts/ci/amd/amd_ci_warmup_aiter.py
View File

@@ -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}")

2
test/registered/attention/test_triton_attention_backend.py
View File

@@ -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):

2
test/registered/dllm/test_llada2_mini_amd.py
View File

@@ -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):

2
test/registered/eval/test_eval_accuracy_large.py
View File

@@ -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):

2
test/registered/mla/test_mla_fp8.py
View File

@@ -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):

2
test/registered/models/test_vlm_models.py
View File

@@ -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

2
test/registered/perf/test_bench_serving_2gpu.py
View File

@@ -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):

2
test/registered/quant/test_eval_fp8_accuracy.py
View File

@@ -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):

2
test/registered/quant/test_torchao.py
View File

@@ -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

2
test/registered/scheduler/test_abort.py
View File

@@ -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):

2
test/registered/scheduler/test_retract_decode.py
View File

@@ -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):