[Bug] Paged causal prefill returns finite LSE for fully masked rows when qo_len > kv_len
needs-triage
### Before submitting
- [x] I searched open and closed issues for this wrapper/shape case.
- [x] This is related to #4267 and the proposed fix in #4401, but the trigger is a genuinely empty causal row rather than an extremely negative valid row.
### Environment
The failure was observed with:
- GPU: NVIDIA B200 (SM100)
- FlashInfer: `0.6.8.post1`
- PyTorch: `2.9.0`
- CUDA runtime: `13.0`
- dtype: BF16
The script below prints the environment it actually runs in. No model weights or external data are required.
### Bug description
`BatchPrefillWithPagedKVCacheWrapper(causal=True)` returns finite LSE values near `-5e4` for fully masked leading query rows when `qo_len > kv_len`. The mathematically correct empty-attention state is output zero and LSE `-inf`.
FlashInfer's documented bottom-right causal alignment permits this shape: with `qo_len=34` and `kv_len=1`, the first 33 query rows have no valid key and only the final query row can attend the one KV token.
This differs from #4267: that issue has valid finite logits below the sentinel, whereas this case uses ordinary random values and has no valid logits at all. PR #4401 says fully masked behavior is unchanged, but its current test file does not construct a fully masked row and does not request LSE.
This case is also important for validating the proposed true-infinity change. In a scheduled empty row, online-softmax updates may encounter `-inf - -inf` before later empty-row normalization. The reproduction below can determine whether the one-line change is sufficient or an explicit empty-row state is required.
### Minimal reproduction
```python
import math
import flashinfer
import torch
device = torch.device("cuda")
dtype = torch.bfloat16
qo_len = 34
kv_len = 1
num_qo_heads = 32
num_kv_heads = 8
head_dim = 128
page_size = 1
num_pages = 2
sm_scale = 1.0 / math.sqrt(head_dim)
print("flashinfer:", flashinfer.__version__)
print("torch:", torch.__version__)
print("CUDA runtime:", torch.version.cuda)
print("GPU:", torch.cuda.get_device_name(device))
torch.manual_seed(200)
q = torch.randn(
(qo_len, num_qo_heads, head_dim), dtype=dtype, device=device
)
k_cache = torch.randn(
(num_pages, page_size, num_kv_heads, head_dim),
dtype=dtype,
device=device,
)
v_cache = torch.randn_like(k_cache)
qo_indptr = torch.tensor([0, qo_len], dtype=torch.int32, device=device)
paged_kv_indptr = torch.tensor([0, 1], dtype=torch.int32, device=device)
paged_kv_indices = torch.tensor([1], dtype=torch.int32, device=device)
paged_kv_last_page_len = torch.tensor([1], dtype=torch.int32, device=device)
workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device)
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
workspace,
kv_layout="NHD",
backend="fa2",
)
wrapper.plan(
qo_indptr=qo_indptr,
paged_kv_indptr=paged_kv_indptr,
paged_kv_indices=paged_kv_indices,
paged_kv_last_page_len=paged_kv_last_page_len,
num_qo_heads=num_qo_heads,
num_kv_heads=num_kv_heads,
head_dim_qk=head_dim,
page_size=page_size,
causal=True,
sm_scale=sm_scale,
q_data_type=dtype,
kv_data_type=dtype,
)
out, lse = wrapper.run(q, (k_cache, v_cache), return_lse=True)
# Bottom-right causal alignment with qo_len=34 and kv_len=1 makes rows
# [0, 33) fully masked. Their exact reference state is output=0, LSE=-inf.
num_fully_masked_rows = qo_len - kv_len
empty_out = out[:num_fully_masked_rows]
empty_lse = lse[:num_fully_masked_rows]
print("fully masked head-rows:", empty_lse.numel())
print("finite LSE values among them:", torch.isfinite(empty_lse).sum().item())
print("first fully masked output norm:", empty_out[0, 0].float().norm().item())
print("first fully masked LSE:", empty_lse[0, 0].item())
torch.testing.assert_close(
empty_out,
torch.zeros_like(empty_out),
rtol=0,
atol=0,
)
assert torch.isneginf(empty_lse).all(), "fully masked rows must have LSE=-inf"
```
### Expected behavior
All `33 * 32 = 1,056` fully masked head-rows should return output zero and LSE `-inf`.
### Actual behavior
On the affected build, all 1,056 fully masked LSE values were finite; `lse[0, 0]` was approximately `-49993.98`. The final assertion fails.
### Related upstream work
- #4267 identifies the finite sentinel but its runnable reproduction always has valid keys.
- #4401 proposes true infinity and claims fully masked behavior is unchanged, but currently has no truly fully-masked or LSE test.
- #3047 establishes `output=0, LSE=-inf` as the intended empty-attention state in another FlashInfer attention path.
- #4051 concerns merging already-empty attention states and is related semantically, but occurs after the attention kernel rather than in paged GQA prefill.
0 条评论