minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[XPU][CI] key persistent JIT kernel cache by image content ID (#35337)

Browse Source
This commit is contained in:
ashwini rathi
2026-08-20 10:27:59 +08:00
committed by GitHub
parent 9db4ba8da1
commit 238ba40c27
3 changed files with 45 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-xpu.yml
+1 -1
View File
@@ -206,7 +206,7 @@ jobs:
docker exec ci_sglang_xpu /bin/bash -c '/opt/venv/bin/hf auth login --token ${HF_TOKEN}'
- name: Run stage-b tests
timeout-minutes: 60
timeout-minutes: 120
run: |
docker exec ci_sglang_xpu bash -c "source /opt/venv/bin/activate && cd /sglang-checkout/test && python3 run_suite.py --hw xpu --suite stage-b-test-1-gpu-xpu --enable-retry"
python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py
+12 -5
View File
@@ -112,6 +112,9 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop(
)
index = tl.load(initial_state_indices + i_n).to(tl.int32)
# Padded rows carry the -1 sentinel; without this guard the sentinel
# reaches pointer arithmetic and addresses before the state pool.
valid_state = index >= 0
h0 = initial_state + index * stride_h
ht = initial_state + index * stride_h
if USE_INITIAL_STATE:
@@ -128,18 +131,20 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop(
for k_blk in range(0, K, 64):
# Load h: from initial_state (i_t==0) or scratch (i_t>0)
if i_t == 0:
if USE_INITIAL_STATE:
if USE_INITIAL_STATE and valid_state:
p_hs = tl.make_block_ptr(
h0, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
)
b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
else:
b_h = tl.zeros([BV, 64], dtype=tl.float32)
else:
elif valid_state:
p_hs = tl.make_block_ptr(
ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
)
b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
else:
b_h = tl.zeros([BV, 64], dtype=tl.float32)
# Store pre-update h to output
p_ho = tl.make_block_ptr(
@@ -181,18 +186,20 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop(
for k_blk in range(0, K, 64):
# Reload h (same source as Phase 1)
if i_t == 0:
if USE_INITIAL_STATE:
if USE_INITIAL_STATE and valid_state:
p_hs = tl.make_block_ptr(
h0, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
)
b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
else:
b_h = tl.zeros([BV, 64], dtype=tl.float32)
else:
elif valid_state:
p_hs = tl.make_block_ptr(
ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
)
b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
else:
b_h = tl.zeros([BV, 64], dtype=tl.float32)
# Gate decay on h
if USE_G:
@@ -215,7 +222,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop(
b_h += tl.trans(tl.dot(b_k, b_v))
# Save updated h to scratch (initial_state) for next time step
if INPLACE_UPDATE:
if INPLACE_UPDATE and valid_state:
p_hs = tl.make_block_ptr(
ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
)
scripts/ci/xpu/xpu_ci_start_container.sh
+32 -5
View File
@@ -109,18 +109,45 @@ elif [[ -r "${HF_TOKEN_FILE}" ]]; then
HF_TOKEN_VALUE=$(cat "${HF_TOKEN_FILE}")
fi
# Persistent JIT kernel cache (Triton/Inductor/NEO/SYCL) keyed by GPU mask.
# Cold JIT compile can push test_xpu_basic past its 1200s timeout on B580.
XPU_KERNEL_CACHE_HOST="${XPU_KERNEL_CACHE_DIR:-${HOME}/.cache/sglang-xpu-ci/kernel-cache-gpu${ZE_AFFINITY_MASK:-shared}}"
# Persistent JIT kernel cache keyed by GPU mask + image ID (new image -> new
# cache; avoids dlopen of stale .so's like libsycl.so.8 after a torch bump).
if [[ -n "${XPU_KERNEL_CACHE_DIR:-}" ]]; then
XPU_KERNEL_CACHE_HOST="${XPU_KERNEL_CACHE_DIR}"
else
# `|| IMG_ID_SHORT=""` keeps pipefail from killing the script on inspect failure.
IMG_ID_SHORT=$(docker image inspect --format '{{.Id}}' "${IMAGE}" 2>/dev/null \
| sed 's/^sha256://' | cut -c1-12) || IMG_ID_SHORT=""
CACHE_ROOT="${HOME}/.cache/sglang-xpu-ci"
GPU_KEY="gpu${ZE_AFFINITY_MASK:-shared}"
if [[ -n "${IMG_ID_SHORT}" ]]; then
XPU_KERNEL_CACHE_HOST="${CACHE_ROOT}/kernel-cache-${GPU_KEY}-${IMG_ID_SHORT}"
# Prune caches for other image IDs + the legacy unversioned dir (root-owned).
shopt -s nullglob
stale_siblings=("${CACHE_ROOT}"/kernel-cache-"${GPU_KEY}"-* "${CACHE_ROOT}/kernel-cache-${GPU_KEY}")
shopt -u nullglob
for sibling in "${stale_siblings[@]}"; do
[[ -d "${sibling}" ]] || continue
[[ "${sibling}" == "${XPU_KERNEL_CACHE_HOST}" ]] && continue
echo "Pruning stale kernel cache: ${sibling}"
docker run --rm -v "${CACHE_ROOT}:/c" busybox:latest \
rm -rf "/c/$(basename "${sibling}")" || true
done
else
# Throwaway per-run dir; legacy path may be poisoned. Next good run prunes it.
echo "Warning: could not resolve image ID for ${IMAGE}; using throwaway cache." >&2
XPU_KERNEL_CACHE_HOST="${CACHE_ROOT}/kernel-cache-${GPU_KEY}-unversioned-$$"
fi
fi
mkdir -p "${XPU_KERNEL_CACHE_HOST}"/{triton,inductor,neo,sycl}
echo "Using persistent XPU kernel cache: ${XPU_KERNEL_CACHE_HOST}"
# Cap the cache (default 5 GiB); over-cap resets it (misses just recompile).
# Cap the cache (default 5 GiB); over-cap resets it via busybox (root-owned).
XPU_KERNEL_CACHE_MAX_MB="${XPU_KERNEL_CACHE_MAX_MB:-5120}"
cache_mb=$(du -sm "${XPU_KERNEL_CACHE_HOST}" 2>/dev/null | cut -f1)
if [[ -n "${cache_mb}" && "${cache_mb}" -gt "${XPU_KERNEL_CACHE_MAX_MB}" ]]; then
echo "XPU kernel cache is ${cache_mb} MiB (> ${XPU_KERNEL_CACHE_MAX_MB} MiB cap); resetting it."
rm -rf "${XPU_KERNEL_CACHE_HOST:?}"/{triton,inductor,neo,sycl}
docker run --rm -v "${XPU_KERNEL_CACHE_HOST}:/c" busybox:latest \
sh -c 'rm -rf /c/triton /c/inductor /c/neo /c/sycl' || true
mkdir -p "${XPU_KERNEL_CACHE_HOST}"/{triton,inductor,neo,sycl}
fi