From 86faf6a806439595602cb9b2b1a3bc5cbd4cb952 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 17 Jun 2026 04:59:09 +0000 Subject: [PATCH 1/2] fix(hip,aiter): bootstrap non-causal varlen .so for single-prefill with logits cap The single-prefill AITER path routes through mha_varlen_fwd when logits_soft_cap > 0, and AITER lazy-JIT-builds that .so split by causality (mask vs nmask). The bootstrap helper hardcoded is_causal=True, so a non-causal request hit a dlopen failure for the never-built nmask variant. Thread causal through the bootstrap so it builds the variant matching the request. Co-Authored-By: Claude Opus 4.7 --- flashinfer/prefill_rocm.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index b4485e2ff0..b879eecdc1 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -369,15 +369,17 @@ def _auto_select_prefill_backend( @functools.lru_cache(maxsize=None) def _aiter_bootstrap_single_prefill_varlen( dtype: torch.dtype, + causal: bool, head_dim: int, device_idx: int, ) -> None: - """Force AITER's lazy JIT to compile mha_varlen_fwd_*.so for logits+causal variants. + """Force AITER's lazy JIT to compile mha_varlen_fwd_*.so logits variants. - Non-logits and non-causal variants are pre-shipped with the AITER package. - The logits+causal combination is missing from the pre-built set and must be - bootstrapped on first use. Used by single-prefill when logits_soft_cap > 0 - (which forces the varlen .so because mha_fwd has no _logits arm). + Non-logits variants are pre-shipped with the AITER package; the logits + combination is missing from the pre-built set and must be bootstrapped on + first use. Used by single-prefill when logits_soft_cap > 0 (which forces the + varlen .so because mha_fwd has no _logits arm). The .so is split by causal + (mask vs nmask), so build the variant matching the actual request. """ device = torch.device("cuda", device_idx) q = torch.zeros(2, 2, head_dim, dtype=dtype, device=device) @@ -394,7 +396,7 @@ def _aiter_bootstrap_single_prefill_varlen( softmax_scale=scale, logits_soft_cap=0.5, zero_tensors=False, - is_causal=True, + is_causal=causal, window_size_left=-1, window_size_right=-1, sink_size=0, @@ -1351,7 +1353,7 @@ def single_prefill_with_kv_cache( with _aiter_bootstrap_lock: if logits_soft_cap > 0: _aiter_bootstrap_single_prefill_varlen( - q.dtype, q.shape[-1], q.device.index or 0 + q.dtype, causal, q.shape[-1], q.device.index or 0 ) else: _aiter_bootstrap_single_prefill_mha_fwd( From 9f81d192f73dfb70de918cb2313dd81b7454bed8 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 17 Jun 2026 05:04:18 +0000 Subject: [PATCH 2/2] docs(hip,aiter): clarify varlen .so causality split in bootstrap comments Address Copilot review: the logits varlen .so is split by causality (mask vs nmask) and neither is pre-shipped, so both the docstring and the call-site comment now reflect that we bootstrap the requested variant rather than only the (logits, causal) combo. Co-Authored-By: Claude Opus 4.7 --- flashinfer/prefill_rocm.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index b879eecdc1..db50584537 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -376,7 +376,7 @@ def _aiter_bootstrap_single_prefill_varlen( """Force AITER's lazy JIT to compile mha_varlen_fwd_*.so logits variants. Non-logits variants are pre-shipped with the AITER package; the logits - combination is missing from the pre-built set and must be bootstrapped on + variants are missing from the pre-built set and must be bootstrapped on first use. Used by single-prefill when logits_soft_cap > 0 (which forces the varlen .so because mha_fwd has no _logits arm). The .so is split by causal (mask vs nmask), so build the variant matching the actual request. @@ -1347,7 +1347,8 @@ def single_prefill_with_kv_cache( "use backend='fa2' or backend='auto' instead." ) # logits_soft_cap > 0 forces the varlen .so (mha_fwd template has no _logits - # arm); only the (logits, causal) combo isn't pre-shipped by AITER. + # arm); the logits .so is split by causality (mask vs nmask) and neither is + # pre-shipped by AITER, so bootstrap the variant matching the request. # logits_soft_cap == 0 takes the mha_fwd .so, none of whose variants are # pre-shipped — JIT-build the exact (dtype, causal, has_lse) we need. with _aiter_bootstrap_lock: