[Bugfix] Prevent out-of-bounds access in FlashInfer SM90 sparse MLA mixed batches (#56969)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Sep 15, 2026 at 23:03 UTC
836bb3839ffefcda8283ea7d41671a89e1a613df
2 files changed
+27
-7
tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
+19
-7
@@ -151,32 +151,44 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
151
assert kwargs == {}
152
153
154
-def test_builder_attaches_its_state(monkeypatch):
154
+@pytest.mark.parametrize("use_mha", [False, True])
155
+@pytest.mark.parametrize("num_decodes", [0, 1])
156
+def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes):
157
+ """MHA prefill rows must not make the MQA kernel read beyond its query."""
158
builder = object.__new__(FlashInferMLASparseSM90Builder)
159
builder._index_topk = 2048
160
builder._index_kpool = 4
161
builder._async_scheduling = False
162
builder.state = FakeState(TOPK)
163
+ builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha)
164
metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata)
165
metadata.state = None
166
+ metadata.num_prefills = 1
167
+ metadata.num_decode_tokens = num_decodes
168
monkeypatch.setattr(
169
sm90_mod.FlashInferMLASparseMetadataBuilder,
170
"build",
171
lambda *_args, **_kwargs: metadata,
172
)
173
cam = SimpleNamespace(
168
- num_reqs=1,
169
- query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32),
170
- seq_lens=torch.tensor([1], dtype=torch.int32),
171
- seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32),
174
+ num_reqs=num_decodes + 1,
175
+ query_start_loc_cpu=torch.tensor(
176
+ [0, 1, 6] if num_decodes else [0, 5], dtype=torch.int32
177
+ ),
178
+ seq_lens_cpu_upper_bound=torch.tensor(
179
+ [1402, 5] if num_decodes else [5], dtype=torch.int32
180
+ ),
181
positions=None,
182
)
183
184
result = builder.build(0, cam)
185
186
assert result.state is builder.state
178
- assert builder.state.plan_calls[0][0] == 1
179
- assert builder.state.plan_calls[0][1].tolist() == [1]
187
+ expected_lens = [1402] if num_decodes else []
188
+ if not use_mha:
189
+ expected_lens += [1, 2, 3, 4, 5]
190
+ assert builder.state.plan_calls[0][0] == len(expected_lens)
191
+ assert builder.state.plan_calls[0][1].tolist() == expected_lens
192
193
194
def test_plan_uses_state_params(monkeypatch):
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
+8
@@ -277,6 +277,7 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
277
attention_layer = vllm_config.compilation_config.static_forward_context[
278
layer_names[0]
279
]
280
+ self._attention_layer = attention_layer
281
impl = attention_layer.impl
282
if not isinstance(impl, FlashInferMLASparseSM90Impl):
283
raise TypeError(
@@ -376,6 +377,13 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
377
# Replan every step outside any CUDA graph capture with this step's
378
# exact per-row lengths; captured runs read the refreshed buffers.
379
num_rows, kv_lens = self._kv_lens_host(common_attn_metadata)
380
+ # MHA prefills are removed from q before forward_mqa. Plan only the
381
+ # decode prefix in that case, using the same routing decision as MLA.
382
+ if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha(
383
+ metadata
384
+ ):
385
+ num_rows = metadata.num_decode_tokens
386
+ kv_lens = kv_lens[:num_rows]
387
self.state.plan(num_rows, kv_lens)
388
metadata.state = self.state
389
return metadata