minke.yu/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-08 22:08:24 -08:00
committed by GitHub
co-authored by kkHuang-amd YC Tseng
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
.github/workflows/pr-test-amd.yml
+3 -3
View File
@@ -251,7 +251,7 @@ jobs:
fail-fast: false fail-fast: false
matrix: matrix:
runner: [linux-mi325-gpu-1] 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}} runs-on: ${{matrix.runner}}
steps: steps:
- name: Checkout code - name: Checkout code
@@ -273,7 +273,7 @@ jobs:
- name: Run test - name: Run test
timeout-minutes: 30 timeout-minutes: 30
run: | 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: stage-b-test-small-1-gpu-amd-mi35x:
needs: [check-changes, stage-a-test-1-amd] 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" docker exec ci_sglang rocm-smi --showmeminfo vram 2>/dev/null || echo "rocm-smi not available"
- name: Run diffusion server tests (1-GPU) - name: Run diffusion server tests (1-GPU)
timeout-minutes: 45 timeout-minutes: 60
run: | run: |
# AMD CI: All 1-GPU tests except FLUX.2 (FLUX.1 covers same code path) # AMD CI: All 1-GPU tests except FLUX.2 (FLUX.1 covers same code path)
# Tests: T2V, T2I, I2V, LoRA # Tests: T2V, T2I, I2V, LoRA
docker/rocm.Dockerfile
+2 -2
View File
@@ -21,7 +21,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0" ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1" ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1" ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT="v0.1.9.post1" ENV AITER_COMMIT="v0.1.10.post2"
# =============================== # ===============================
# Base image 950 and args # Base image 950 and args
@@ -31,7 +31,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0" ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="0" ENV BUILD_AITER_ALL="0"
ENV BUILD_MOONCAKE="1" ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT="v0.1.9.post1" ENV AITER_COMMIT="v0.1.10.post2"
# =============================== # ===============================
# Chosen arch and args # Chosen arch and args
FROM ${GPU_ARCH} FROM ${GPU_ARCH}
python/sglang/srt/layers/attention/aiter_backend.py
+19 -8
View File
@@ -268,6 +268,7 @@ class AiterAttnBackend(AttentionBackend):
self, self,
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
kv_last_page_len,
work_metadata, work_metadata,
work_info_set, work_info_set,
work_indptr, work_indptr,
@@ -287,6 +288,7 @@ class AiterAttnBackend(AttentionBackend):
meta = get_mla_metadata_v1( meta = get_mla_metadata_v1(
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
kv_last_page_len,
self.num_head // nhead_kv, self.num_head // nhead_kv,
nhead_kv, nhead_kv,
True, True,
@@ -367,6 +369,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
kv_last_page_len,
work_metadata, work_metadata,
work_info_set, work_info_set,
work_indptr, work_indptr,
@@ -423,6 +426,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
self.kv_last_page_len[:bs],
work_metadata, work_metadata,
work_info_set, work_info_set,
work_indptr, work_indptr,
@@ -518,6 +522,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
self.kv_last_page_len[:bs],
work_metadata, work_metadata,
work_info_set, work_info_set,
work_indptr, work_indptr,
@@ -716,6 +721,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
kv_last_page_len,
self.work_metadata, self.work_metadata,
self.work_info_set, self.work_info_set,
self.work_indptr, self.work_indptr,
@@ -786,6 +792,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
kv_last_page_len,
self.work_metadata, self.work_metadata,
self.work_info_set, self.work_info_set,
self.work_indptr, self.work_indptr,
@@ -872,6 +879,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
qo_indptr, qo_indptr,
kv_indptr, kv_indptr,
kv_last_page_len,
self.work_metadata, self.work_metadata,
self.work_info_set, self.work_info_set,
self.work_indptr, self.work_indptr,
@@ -1144,6 +1152,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
self.forward_metadata.qo_indptr, self.forward_metadata.qo_indptr,
self.forward_metadata.kv_indptr, self.forward_metadata.kv_indptr,
self.forward_metadata.kv_last_page_len,
work_metadata, work_metadata,
work_info_set, work_info_set,
work_indptr, work_indptr,
@@ -1165,8 +1174,8 @@ class AiterAttnBackend(AttentionBackend):
self.forward_metadata.kv_indices, self.forward_metadata.kv_indices,
self.forward_metadata.kv_last_page_len, self.forward_metadata.kv_last_page_len,
self.forward_metadata.max_q_len, self.forward_metadata.max_q_len,
layer.scaling, sm_scale=layer.scaling,
layer.logit_cap, logit_cap=layer.logit_cap,
work_meta_data=work_metadata, work_meta_data=work_metadata,
work_indptr=work_indptr, work_indptr=work_indptr,
work_info_set=work_info_set, work_info_set=work_info_set,
@@ -1195,6 +1204,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
self.forward_metadata.qo_indptr, self.forward_metadata.qo_indptr,
self.forward_metadata.kv_indptr, self.forward_metadata.kv_indptr,
self.forward_metadata.kv_last_page_len,
work_metadata, work_metadata,
work_info_set, work_info_set,
work_indptr, work_indptr,
@@ -1232,8 +1242,8 @@ class AiterAttnBackend(AttentionBackend):
self.forward_metadata.kv_indices, self.forward_metadata.kv_indices,
self.forward_metadata.kv_last_page_len, self.forward_metadata.kv_last_page_len,
self.forward_metadata.max_q_len, self.forward_metadata.max_q_len,
layer.scaling, sm_scale=layer.scaling,
layer.logit_cap, logit_cap=layer.logit_cap,
work_meta_data=work_metadata, work_meta_data=work_metadata,
work_indptr=work_indptr, work_indptr=work_indptr,
work_info_set=work_info_set, work_info_set=work_info_set,
@@ -1262,8 +1272,8 @@ class AiterAttnBackend(AttentionBackend):
self.forward_metadata.kv_indices, self.forward_metadata.kv_indices,
self.forward_metadata.kv_last_page_len, self.forward_metadata.kv_last_page_len,
self.forward_metadata.max_q_len, self.forward_metadata.max_q_len,
layer.scaling, sm_scale=layer.scaling,
layer.logit_cap, logit_cap=layer.logit_cap,
work_meta_data=work_metadata, work_meta_data=work_metadata,
work_indptr=work_indptr, work_indptr=work_indptr,
work_info_set=work_info_set, work_info_set=work_info_set,
@@ -1353,6 +1363,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_meta_data( self.make_mla_meta_data(
self.forward_metadata.qo_indptr, self.forward_metadata.qo_indptr,
self.forward_metadata.kv_indptr, self.forward_metadata.kv_indptr,
self.forward_metadata.kv_last_page_len,
work_metadata, work_metadata,
work_info_set, work_info_set,
work_indptr, work_indptr,
@@ -1374,8 +1385,8 @@ class AiterAttnBackend(AttentionBackend):
self.forward_metadata.kv_indices, self.forward_metadata.kv_indices,
self.forward_metadata.kv_last_page_len, self.forward_metadata.kv_last_page_len,
self.forward_metadata.max_q_len, self.forward_metadata.max_q_len,
layer.scaling, sm_scale=layer.scaling,
layer.logit_cap, logit_cap=layer.logit_cap,
work_meta_data=work_metadata, work_meta_data=work_metadata,
work_indptr=work_indptr, work_indptr=work_indptr,
work_info_set=work_info_set, work_info_set=work_info_set,
scripts/ci/amd/amd_ci_install_dependency.sh
+1 -1
View File
@@ -233,7 +233,7 @@ if [[ "${NEED_REBUILD}" == "true" ]]; then
echo "[CI-AITER-CHECK] === AITER REBUILD START ===" echo "[CI-AITER-CHECK] === AITER REBUILD START ==="
# uninstall existing aiter # 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 # delete old aiter directory
docker exec ci_sglang rm -rf /sgl-workspace/aiter docker exec ci_sglang rm -rf /sgl-workspace/aiter
scripts/ci/amd/amd_ci_warmup_aiter.py
+44 -17
View File
@@ -32,10 +32,12 @@ def warmup_aiter_kernels():
device = torch.device("cuda:0") device = torch.device("cuda:0")
start_time = time.time() start_time = time.time()
# Warmup RMSNorm kernel (module_rmsnorm) - most commonly used # Warmup module_rmsnorm_quant (small module, ~2MB)
# SGLang uses rmsnorm2d_fwd and rmsnorm2d_fwd_with_add from aiter # Triggered by rmsnorm2d_fwd when hidden_size <= 8192
try: 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 from aiter import rmsnorm2d_fwd
hidden_size = 4096 hidden_size = 4096
@@ -44,37 +46,62 @@ def warmup_aiter_kernels():
weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device) weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device)
eps = 1e-6 eps = 1e-6
# This triggers JIT compilation # hidden_size=4096 <= 8192 -> takes rmsnorm() path -> compiles module_rmsnorm_quant
_ = rmsnorm2d_fwd(x, weight, eps) _ = rmsnorm2d_fwd(x, weight, eps)
torch.cuda.synchronize() torch.cuda.synchronize()
print(f" RMSNorm kernel (rmsnorm2d_fwd) compiled successfully") print(" module_rmsnorm_quant compiled successfully")
except Exception as e: 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: 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 from aiter import rmsnorm2d_fwd_with_add
hidden_size = 4096 hidden_size = 4096
batch_size = 512 batch_size = 512
x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device=device) 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 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) weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device)
eps = 1e-6 eps = 1e-6
# This triggers JIT compilation # This triggers JIT compilation of module_rmsnorm (CK kernels)
_ = rmsnorm2d_fwd_with_add(x, residual, weight, eps) rmsnorm2d_fwd_with_add(output, x, residual_in, residual_out, weight, eps)
torch.cuda.synchronize() torch.cuda.synchronize()
print(f" Fused add RMSNorm kernel compiled successfully") print(" module_rmsnorm compiled successfully")
except Exception as e: 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 # Warmup rotary embedding kernel if available
try: try:
print("\n[3/4] Warming up rotary embedding kernel...") print("\n[4/5] Warming up rotary embedding kernel...")
from aiter import rotary_embedding from aiter import rotary_embedding
head_size = 128 head_size = 128
@@ -92,13 +119,13 @@ def warmup_aiter_kernels():
_ = rotary_embedding(positions, query, key, head_size, cos, sin, True) _ = rotary_embedding(positions, query, key, head_size, cos, sin, True)
torch.cuda.synchronize() torch.cuda.synchronize()
print(f" Rotary embedding kernel compiled successfully") print(" Rotary embedding kernel compiled successfully")
except Exception as e: except Exception as e:
print(f" Rotary embedding warmup skipped (may not be available): {e}") print(f" Rotary embedding warmup skipped (may not be available): {e}")
# Warmup activation kernels if available # Warmup activation kernels if available
try: try:
print("\n[4/4] Warming up activation kernels...") print("\n[5/5] Warming up activation kernels...")
from aiter import silu_and_mul from aiter import silu_and_mul
hidden_size = 4096 hidden_size = 4096
@@ -110,7 +137,7 @@ def warmup_aiter_kernels():
silu_and_mul(out, x) silu_and_mul(out, x)
torch.cuda.synchronize() torch.cuda.synchronize()
print(f" Activation kernel compiled successfully") print(" Activation kernel compiled successfully")
except Exception as e: except Exception as e:
print(f" Activation warmup skipped (may not be available): {e}") print(f" Activation warmup skipped (may not be available): {e}")
test/registered/attention/test_triton_attention_backend.py
+1 -1
View File
@@ -21,7 +21,7 @@ from sglang.test.test_utils import (
# Triton attention backend integration test with latency benchmark and MMLU eval # 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_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): class TestTritonAttnBackend(CustomTestCase):
test/registered/dllm/test_llada2_mini_amd.py
+1 -1
View File
@@ -20,7 +20,7 @@ from sglang.test.test_utils import (
write_github_step_summary, 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): class TestLLaDA2MiniAMD(CustomTestCase):
test/registered/eval/test_eval_accuracy_large.py
+1 -1
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_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): class TestEvalAccuracyLarge(CustomTestCase):
test/registered/mla/test_mla_fp8.py
+1 -1
View File
@@ -14,7 +14,7 @@ from sglang.test.test_utils import (
# MLA FP8 KV cache test with MGSM evaluation # MLA FP8 KV cache test with MGSM evaluation
register_cuda_ci(est_time=77, suite="stage-b-test-large-1-gpu") 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): class TestMLA(CustomTestCase):
test/registered/models/test_vlm_models.py
+1 -1
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_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() _is_hip = is_hip()
# VLM models for testing # VLM models for testing
test/registered/perf/test_bench_serving_2gpu.py
+1 -1
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_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): class TestBenchServing2GPU(CustomTestCase):
test/registered/quant/test_eval_fp8_accuracy.py
+1 -1
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_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): class TestEvalFP8Accuracy(CustomTestCase):
test/registered/quant/test_torchao.py
+1 -1
View File
@@ -7,7 +7,7 @@ from sglang import Engine
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci 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_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.lang.chat_template import get_chat_template_by_model_path
from sglang.srt.utils import kill_process_tree from sglang.srt.utils import kill_process_tree
from sglang.test.run_eval import run_eval from sglang.test.run_eval import run_eval
test/registered/scheduler/test_abort.py
+1 -1
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_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): class TestAbort(CustomTestCase):
test/registered/scheduler/test_retract_decode.py
+1 -1
View File
@@ -18,7 +18,7 @@ from sglang.test.test_utils import (
from sglang.utils import is_in_ci from sglang.utils import is_in_ci
register_cuda_ci(est_time=311, suite="stage-b-test-small-1-gpu") 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): class TestRetractDecode(CustomTestCase):