fix: support non-128-aligned element sizes in HiCache JIT

[Chronostasys]

Motivation

HiCache JIT currently assumes each transferred element size is 128-byte aligned. Models with MLA KV layouts such as GLM-5.2 can use an element_size of 576 bytes, which is not divisible by 128. In that case the JIT kernel is rejected before it can be used, even though the data can still be safely transferred with smaller vectorized memory packages.

I checked upstream main and searched existing issues/PRs for related HiCache JIT fixes (HiCache JIT element_size 576, hicache multiple of 128, GLM-5.2 hicache) and did not find an existing fix for this case.

Changes

• Generalize the HiCache JIT load/store package selection from a fixed 128-byte package group to the largest supported package size (16B, 8B, or 4B) that evenly covers each worker thread's byte range.
• Preserve the previous behavior for 128-byte-aligned sizes, so existing common paths keep the same package size and loop count.
• Add a Python-side eligibility check matching the C++ template constraints, so unsupported non-aligned sizes fall back cleanly instead of failing during JIT compilation.

Validation

• python3 -m py_compile python/sglang/jit_kernel/hicache.py
• python3 -m ruff check --select=F401,F821,UP037 python/sglang/jit_kernel/hicache.py
• python3 -m isort --check-only python/sglang/jit_kernel/hicache.py
• git diff --check upstream/main..HEAD
• Verified the package-selection math keeps the old result for 128-byte-aligned sizes and supports element_size=576 with the default unroll configuration.

---

CI States

Latest PR Test (Base): ❌ Run #33533584580
Latest PR Test (Extra): ❌ Run #33533584383
Latest PR Test (AMD ROCm 7.2): ❌ Run #33533584427

_{✨ Presented to you with Mind Lab — A Lab for Experiential Intelligence.}​

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[hzh0425]

@DarkSharpness could you take a review?

[DarkSharpness]

LGTM

[hzh0425]

/rerun-group hicache

[github-actions]

Results for /rerun-group hicache:

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/hicache/test_hicache_spec_file_storage.py
cd test/ && python3 registered/hicache/test_hicache_variants.py

🚀 2-gpu-h100 (4 tests): ✅ View workflow run

cd test/ && python3 registered/hicache/test_hicache_spec_mooncake_storage.py
cd test/ && python3 registered/hicache/test_hicache_storage_file_backend.py
cd test/ && python3 registered/hicache/test_hicache_storage_mooncake_backend.py
cd test/ && python3 registered/hicache/test_hicache_storage_runtime_attach_detach.py

🚀 1-gpu-5090 (1 test): ❌ View workflow run

cd test/ && python3 registered/hicache/test_hicache_storage.py

🚀 4-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/hicache/test_hicache_storage_3fs_backend.py
cd test/ && python3 registered/hicache/test_qwen35_hicache.py

[alilxxey]

One more real-world element size for this PR: 656 bytes, from the W4AFP8 GLM-5.2 checkpoint
(PhalaCloud/GLM-5.2-W4AFP8, sglang v0.5.18, tp8 H200, --kv-cache-dtype fp8_e4m3,
hicache ratio 2.5, write_through). The JIT fallback is visible in our startup logs, so every
L2 transfer runs on the non-aligned path.

Correctness. We enabled Mooncake L3 in production on v0.5.17, i.e. with #31443 already in,
with the timeout prefetch policy. Within hours users got silently corrupted generations, zero
errors in the logs. Rolled L3 back, clean for three weeks on L2 only. We cannot pin the
corruption on the JIT fallback with certainty, but the non-128-aligned element size is the one
dimension where our setup differs from the configs reported as fixed. Relevant because #30057
was just closed as "fixed by #31443", and that fix did not cover our case.

Performance. Under load, per replica (8 ranks aggregated): rate(load_back_bytes_total)
7.5 to 10.3 GiB/s, rate(load_back_duration_seconds_sum) 0.86 to 0.94, i.e. the replica spends
86 to 94 percent of wall time loading back at about 1.4 GiB/s per rank. If the 4B package
fallback is part of why we are that far from link speed, merging this helps L2-only deployments
too.

[Chronostasys]

Rebased onto the new sglang.kernels layout — #32072 retired the jit_kernel directory, so the patch now lives in kernels/jit/csrc/kvcacheio/hicache.cuh and kernels/ops/kvcache/hicache.py. Logic is unchanged from the original commits; only the paths moved.

The report from @alilxxey above is worth extra weight IMO, since it covers both sides of this fix:

• Correctness: silent generation corruption on W4AFP8 GLM-5.2 (656B element size) with Mooncake L3 + timeout prefetch on v0.5.17 — the non-128-aligned path is the one dimension their setup differs from the configs reported fixed by [HiCache]: Optimize hybrid/DSA L3 prefetch result sync and usable-prefix clamping #31443.
• Performance: replicas spend 86-94% of wall time in load_back at ~1.4 GiB/s per rank with the 4B fallback — so this helps L2-only deployments too, not just the non-aligned case.

@DarkSharpness could you take another look when you get a chance? Happy to address any feedback on the new layout.

[alilxxey]

One clarification on my 656 case:

_default_unroll(656) returns 2 (656 falls in the 512 to 1024 bucket), so num_threads = 32 // 2 = 16. The new check then gives 656 % 16 == 0 but 656 // 16 == 41, and 41 is not a multiple of
4, so can_use_hicache_jit_kernel returns False. 656 = 2^4 x 41 only divides cleanly at
num_threads <= 4, i.e. unroll >= 8, and no caller passes unroll explicitly
(memory_pool_host.py calls without it), so there is no reachable configuration for this size.
576 works, 656 does not.

[Chronostasys]

Good catch, thanks for walking through the arithmetic — 656 = 2^4 × 41 indeed has no viable config at the default unroll.

Fixed in 3487820: the non-staged paths now resolve the unroll by walking the ladder from the performance default toward fewer threads per package (num_threads 16 → 8 → 4 → 2 → 1), so 656 lands on unroll=8 (num_threads=4, 164 bytes/thread, 4B packages). The resolver is shared by can_use_hicache_jit_kernel and the transfer helpers, so the check and the actual kernel construction can't disagree on the chosen unroll anymore. Sizes that divide cleanly at the default (576, 1024, ...) keep their original unroll, and explicitly passing an unroll still overrides the search.

[san-tian]

We hit exactly this in production: a GLM-5.3 fine-tuned FP8 model (MLA layout, element_size=576, 8×B200, sglang v0.5.15.post1 + fusionfix container) crashed with Fatal Python error: Segmentation fault under long-context task load. Crash frame:

void (anonymous namespace)::HiCacheStagedWriteBackKernel<576l, 2u, 2u, 1024u>::run_staged_impl<true>(...)
jit_kernel/hicache.py:341 in transfer_hicache_all_layer_mla_staged_lf_pf

Startup log showed Unsupported element_size = 576 for JIT HiCache kernel — so the JIT path rejected 576 and we fell back to the prebuilt C++ staged-writeback kernel, which segfaulted at runtime during long-context KV write-back. Same root cause as described here.

Impact: production data plane crashed mid-task, systemd restart loop amplified the outage (concurrent long-context requests hit every restarted instance).

Temporary mitigation: disabled HiCache entirely on single-box deployments (host offload is not critical there), throughput cost ~15% (933 vs 1099 tok/s at concurrency 8 on 8×B200).

We'd love to see this land — after merge we'll re-enable HiCache and run the long-context concurrency regression to confirm. Thanks for the fix!

[Chronostasys]

Thanks for the detailed report — and sorry it took a production crash to surface this. 576 is exactly the primary case this fix targets: after it lands, the JIT path accepts 576 (num_threads=16, 36 B/thread, 4B packages) instead of rejecting it, so your load path stays on the JIT kernels rather than the prebuilt fallback that segfaulted for you. Your post-merge long-context regression will be the real confirmation — if anything in the staged write-back path still misbehaves for 576, happy to follow up on it here.