ITADN

[Bug] Paged causal prefill returns finite LSE for fully masked rows when qo_len > kv_len

#4452Openanguyen8 创建于 18 天前
needs-triage
A
anguyen8commented
### 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 条评论