From 836bb3839ffefcda8283ea7d41671a89e1a613df Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 15 Sep 2026 23:03:58 +0800 Subject: [PATCH] [Bugfix] Prevent out-of-bounds access in FlashInfer SM90 sparse MLA mixed batches (#56969) Signed-off-by: chaunceyjiang --- .../test_flashinfer_mla_sparse_sm90.py | 26 ++++++++++++++----- .../mla/flashinfer_mla_sparse_sm90.py | 8 ++++++ 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index b20c163306..8dc082b08b 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -151,32 +151,44 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): assert kwargs == {} -def test_builder_attaches_its_state(monkeypatch): +@pytest.mark.parametrize("use_mha", [False, True]) +@pytest.mark.parametrize("num_decodes", [0, 1]) +def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes): + """MHA prefill rows must not make the MQA kernel read beyond its query.""" builder = object.__new__(FlashInferMLASparseSM90Builder) builder._index_topk = 2048 builder._index_kpool = 4 builder._async_scheduling = False builder.state = FakeState(TOPK) + builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha) metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata) metadata.state = None + metadata.num_prefills = 1 + metadata.num_decode_tokens = num_decodes monkeypatch.setattr( sm90_mod.FlashInferMLASparseMetadataBuilder, "build", lambda *_args, **_kwargs: metadata, ) cam = SimpleNamespace( - num_reqs=1, - query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), - seq_lens=torch.tensor([1], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32), + num_reqs=num_decodes + 1, + query_start_loc_cpu=torch.tensor( + [0, 1, 6] if num_decodes else [0, 5], dtype=torch.int32 + ), + seq_lens_cpu_upper_bound=torch.tensor( + [1402, 5] if num_decodes else [5], dtype=torch.int32 + ), positions=None, ) result = builder.build(0, cam) assert result.state is builder.state - assert builder.state.plan_calls[0][0] == 1 - assert builder.state.plan_calls[0][1].tolist() == [1] + expected_lens = [1402] if num_decodes else [] + if not use_mha: + expected_lens += [1, 2, 3, 4, 5] + assert builder.state.plan_calls[0][0] == len(expected_lens) + assert builder.state.plan_calls[0][1].tolist() == expected_lens def test_plan_uses_state_params(monkeypatch): diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 90dc4cd947..830bc52ab9 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py @@ -277,6 +277,7 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): attention_layer = vllm_config.compilation_config.static_forward_context[ layer_names[0] ] + self._attention_layer = attention_layer impl = attention_layer.impl if not isinstance(impl, FlashInferMLASparseSM90Impl): raise TypeError( @@ -376,6 +377,13 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): # Replan every step outside any CUDA graph capture with this step's # exact per-row lengths; captured runs read the refreshed buffers. num_rows, kv_lens = self._kv_lens_host(common_attn_metadata) + # MHA prefills are removed from q before forward_mqa. Plan only the + # decode prefix in that case, using the same routing decision as MLA. + if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha( + metadata + ): + num_rows = metadata.num_decode_tokens + kv_lens = kv_lens[:num_rows] self.state.plan(num_rows, kv_lens) metadata.state = self.state return metadata