chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

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