[Bugfix] Prevent NaN poisoning in xpu_mla_sparse for fully-masked index chunks (#48366)
Signed-off-by: Nick Iusiumbeli <nickuspro@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Nick Iusiumbeli committed
Jul 27, 2026 at 04:06 UTC
f0553889c04ae5bf3df48245302ba5ac19fe3432
2 files changed
+57
-2
tests/kernels/attention/test_xpu_mla_sparse.py
+52
@@ -116,3 +116,55 @@ def test_bf16_triton_sparse_mla(device_str, dtype):
116
assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2)
117
assert torch.allclose(max_logits, ref_max_logits, atol=1e-3, rtol=1e-3)
118
assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3)
119
+
120
+
121
+@pytest.mark.parametrize("device_str", ["xpu"])
122
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
123
+@pytest.mark.skipif(
124
+ not torch.xpu.is_available(),
125
+ reason="XPU is required",
126
+)
127
+def test_bf16_triton_sparse_mla_masked_chunks(device_str, dtype):
128
+ """Rows whose leading BLOCK_N index entries are all masked must not NaN.
129
+
130
+ Regression test: with an -inf running max, a fully-masked leading chunk
131
+ produced re_scale = exp2(-inf - -inf) = NaN, permanently poisoning the
132
+ accumulator even though valid keys followed in later chunks.
133
+ """
134
+ device = torch.device(device_str)
135
+ s_q = 3
136
+ s_kv = 256
137
+ h_q = 64
138
+ h_kv = 1
139
+ d_qk = 576
140
+ d_v = 512
141
+ topk = 128 # 8 chunks of BLOCK_N=16
142
+
143
+ torch.random.manual_seed(1234)
144
+
145
+ q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device)
146
+ kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device)
147
+ indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device)
148
+ # row 0: valid keys only in chunks 1-2 -> leading AND trailing masked chunks
149
+ indices[0, 0, 16:48] = torch.arange(32, dtype=torch.int32, device=device)
150
+ # row 1: fully valid
151
+ indices[1, 0, :] = torch.arange(topk, dtype=torch.int32, device=device)
152
+ # row 2: no valid key at all
153
+
154
+ sm_scale = d_qk**-0.5
155
+
156
+ out, max_logits, lse = triton_bf16_mla_sparse_interface(
157
+ q, kv, indices, sm_scale, d_v
158
+ )
159
+ assert out.isfinite().all()
160
+
161
+ ref_out, _, ref_max_logits, ref_lse = reference_mla_sparse_prefill(
162
+ q, kv, indices, sm_scale, d_v
163
+ )
164
+ assert torch.allclose(out[:2], ref_out[:2], atol=1e-2, rtol=1e-2)
165
+ assert torch.allclose(max_logits[:2], ref_max_logits[:2], atol=1e-3, rtol=1e-3)
166
+ assert torch.allclose(lse[:2], ref_lse[:2], atol=1e-3, rtol=1e-3)
167
+ # A row with no valid key yields zeros (the reference's convention); its
168
+ # lse/max_logits are large-negative finite rather than the reference's
169
+ # +inf/-inf placeholders, so only the output is compared here.
170
+ assert torch.allclose(out[2], torch.zeros_like(out[2]))
vllm/v1/attention/ops/xpu_mla_sparse.py
+5
-2
@@ -73,7 +73,10 @@ def _bf16_mla_sparse_kernel(
73
q_buffer + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
74
)
75
76
- e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
76
+ # Use a finite sentinel rather than -inf: if every key in a chunk is
77
+ # masked (e.g. leading -1 padding in `indices`) an -inf running max gives
78
+ # re_scale = exp2(-inf - -inf) = NaN, permanently poisoning acc / e_sum.
79
+ e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - 1.0e30
80
e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
81
acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
82
@@ -122,7 +125,7 @@ def _bf16_mla_sparse_kernel(
125
126
# apply scaling
127
qk *= sm_scale
125
- qk = tl.where((mask_h[:, None]) & (mask_kv[None, :]), qk, -float("inf"))
128
+ qk = tl.where((mask_h[:, None]) & (mask_kv[None, :]), qk, -1.0e30)
129
130
# load v
131
mask_v_d = offs_dv < dim_v