[Bugfix][GLM-5.3-Flash] SM90 sparse MLA: index_kpool mismatch leads to corruption via unread query token (#58704)
Signed-off-by: Matt Mastracci <matthew@mastracci.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Matt Mastracci committed
Sep 25, 2026 at 10:36 UTC
2edbb32e58aa719bd02ae568fe05ae3e87d2b38e
2 files changed
+56
-3
tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
+54
@@ -20,6 +20,7 @@ from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 import (
20
FlashInferMLASparseSM90Builder,
21
FlashInferMLASparseSM90Impl,
22
)
23
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
24
# isort: on
25
26
BLOCK_SIZE = 64
@@ -313,6 +314,59 @@ def test_kv_lens_host_empty():
314
assert num_rows == 0 and lens.numel() == 0
315
316
317
+_NO_KPOOL = object()
318
+
319
+
320
+@pytest.mark.parametrize(
321
+ "index_kpool,prefill_lens",
322
+ [
323
+ (4, [2048, 2049, 2050, 2051]),
324
+ (None, [2048] * 4),
325
+ (_NO_KPOOL, [2048] * 4),
326
+ ],
327
+ ids=["kpool4", "kpool_none", "no_kpool_attr"],
328
+)
329
+def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens):
330
+ """The builder took kpool from the KV cache spec, whose tokens_per_state is
331
+ 1, so the tail pool was never read."""
332
+ monkeypatch.setattr(
333
+ sm90_mod.FlashInferMLASparseMetadataBuilder,
334
+ "__init__",
335
+ lambda *_args, **_kwargs: None,
336
+ )
337
+ monkeypatch.setattr(sm90_mod, "_SM90State", lambda *_args, **_kwargs: None)
338
+ impl, _ = make_impl(64)
339
+ spec = MLAAttentionSpec(
340
+ block_size=BLOCK_SIZE, num_kv_heads=1, head_size=576, dtype=torch.bfloat16
341
+ )
342
+ assert spec.tokens_per_state == 1
343
+ hf_config = SimpleNamespace(index_topk=2048)
344
+ if index_kpool is not _NO_KPOOL:
345
+ hf_config.index_kpool = index_kpool
346
+ vllm_config = SimpleNamespace(
347
+ compilation_config=SimpleNamespace(
348
+ static_forward_context={"attn": SimpleNamespace(impl=impl)}
349
+ ),
350
+ scheduler_config=SimpleNamespace(
351
+ max_num_batched_tokens=64, async_scheduling=False
352
+ ),
353
+ model_config=SimpleNamespace(hf_text_config=hf_config),
354
+ )
355
+ builder = FlashInferMLASparseSM90Builder(
356
+ spec, ["attn"], vllm_config, torch.device("cpu")
357
+ )
358
+ # req0: prefill chunk ending at context 42295; req1: context <= topk.
359
+ cam = SimpleNamespace(
360
+ num_reqs=2,
361
+ query_start_loc_cpu=torch.tensor([0, 4, 6], dtype=torch.int32),
362
+ seq_lens_cpu_upper_bound=torch.tensor([42295, 100], dtype=torch.int32),
363
+ positions=None,
364
+ )
365
+ num_rows, lens = builder._kv_lens_host(cam)
366
+ assert num_rows == 6
367
+ assert lens.tolist() == prefill_lens + [99, 100]
368
+
369
+
370
def test_supports_combination_gates(monkeypatch, default_vllm_config):
371
monkeypatch.setattr(sm90_mod, "has_flashinfer_sm90_nope_mla", lambda: True)
372
call = lambda **kw: FlashInferMLASparseSM90Backend.supports_combination(
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
+2
-3
@@ -330,7 +330,7 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
330
hf_config = vllm_config.model_config.hf_text_config
331
assert hf_config.index_topk is not None
332
self._index_topk = int(hf_config.index_topk)
333
- self._index_kpool = int(kv_cache_spec.tokens_per_state)
333
+ self._index_kpool = int(getattr(hf_config, "index_kpool", 1) or 1)
334
335
def _kv_lens_host(self, cam: CommonAttentionMetadata) -> tuple[int, torch.Tensor]:
336
"""Exact per-row KV lengths, host-side (the flashinfer wrapper bakes
@@ -388,8 +388,7 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
388
+ 1
389
)
390
topk = self._index_topk
391
- kpool = max(self._index_kpool, 1)
392
- lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool)
391
+ lens = torch.where(ctx <= topk, ctx, topk + ctx % self._index_kpool)
392
return num_rows, lens.to(torch.int32)
393
394
def build(