[XPU][CI] key persistent JIT kernel cache by image content ID

.github/workflows/pr-test-xpu.yml

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

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

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