fix(hip,aiter): bootstrap non-causal varlen .so for single-prefill with logits cap

[demandal25]

Summary

Fixes a dlopen failure in the AITER single-prefill path when logits_soft_cap > 0 and causal=False. The bootstrap that triggers AITER's lazy JIT only ever built the causal variant of the mha_varlen_fwd .so, so a non-causal request with a logits soft cap failed to load the never-built nmask variant.

What changed

• flashinfer/prefill_rocm.py — _aiter_bootstrap_single_prefill_varlen now takes a causal argument and threads it into is_causal, so the bootstrap compiles the .so variant matching the actual request instead of hardcoding is_causal=True. AITER splits the varlen .so by causality (mask vs nmask); the caller in single_prefill_with_kv_cache now passes causal through. Added @functools.lru_cache to match the sibling bootstrap helpers and avoid re-running the warmup on every call.

Root cause

When logits_soft_cap > 0, single-prefill routes through mha_varlen_fwd (the mha_fwd template has no _logits arm). AITER ships non-logits variants prebuilt but lazy-JIT-builds the logits ones, keyed by causality. The bootstrap warmup only exercised is_causal=True, leaving the non-causal logits .so absent — so backend="aiter", causal=False, logits_soft_cap=8.0 raised RuntimeError: AITER .so not found: ...mha_varlen_fwd_fp16_logits_nbias_nmask_....

Test plan

• pytest tests/rocm_tests/test_single_prefill_kernels_hip.py — previously failing aiter non-causal logits cases now pass
• pre-commit run -a

[Copilot]

Pull request overview

Fixes an AITER dlopen failure in ROCm single-prefill when logits_soft_cap > 0 and causal=False by ensuring the bootstrap warmup compiles the correct varlen .so variant (mask vs nmask) instead of always bootstrapping the causal one.

Changes:

• Thread causal into _aiter_bootstrap_single_prefill_varlen(..., is_causal=causal) so the lazy-JIT builds the correct logits varlen artifact.
• Add @functools.lru_cache to the varlen bootstrap helper to avoid repeating warmups.
• Update the single-prefill caller to pass causal into the varlen bootstrap.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.