A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 10:30 UTC
[Bugfix] mamba: run single-token extends as decodes (#42430)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Netanel Haber committed
May 18, 2026 at 18:26 UTC
47829b1159335a010521ea3e5361d51744a36b0a
4 files changed
+120
-21
tests/v1/attention/test_mamba_update_block_table.py
+33
-20
@@ -16,23 +16,18 @@ from types import SimpleNamespace
16
17
import torch
18
19
+from tests.v1.attention.utils import MockMambaBuilder
20
from vllm.config.compilation import CUDAGraphMode
20
-from vllm.v1.attention.backends.mamba_attn import (
21
- BaseMambaAttentionMetadata,
22
- BaseMambaAttentionMetadataBuilder,
23
-)
21
+from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadata
22
from vllm.v1.kv_cache_interface import MambaSpec
23
24
27
-class _ConcreteMambaBuilder(
28
- BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetadata]
25
+def _make_vllm_config(
26
+ max_model_len: int,
27
+ max_num_seqs: int,
28
+ num_speculative_tokens: int = 0,
29
+ block_size: int | None = None,
30
):
30
- """Minimal concrete subclass for testing (base class is ABC)."""
31
-
32
- metadata_cls = BaseMambaAttentionMetadata
33
-
34
-
35
-def _make_vllm_config(max_model_len, max_num_seqs, num_speculative_tokens=0):
31
"""Create a minimal mock VllmConfig with only the fields the builder
32
accesses, avoiding any model download / HF config inspection."""
33
speculative_config = (
@@ -44,7 +39,10 @@ def _make_vllm_config(max_model_len, max_num_seqs, num_speculative_tokens=0):
39
else None
40
)
41
return SimpleNamespace(
47
- cache_config=SimpleNamespace(mamba_cache_mode="all"),
42
+ cache_config=SimpleNamespace(
43
+ block_size=block_size,
44
+ mamba_cache_mode="all",
45
+ ),
46
compilation_config=SimpleNamespace(
47
cudagraph_mode=CUDAGraphMode.FULL,
48
max_cudagraph_capture_size=None,
@@ -57,6 +55,21 @@ def _make_vllm_config(max_model_len, max_num_seqs, num_speculative_tokens=0):
55
)
56
57
58
+def test_mamba_single_token_prompt_runs_as_prefill():
59
+ seq_lens = [8, 9, 1]
60
+ config = _make_vllm_config(256, len(seq_lens), block_size=16)
61
+ metadata = MockMambaBuilder.build_mamba_metadata(
62
+ config,
63
+ seq_lens=seq_lens,
64
+ query_lens=[1] * len(seq_lens),
65
+ is_prefilling=[False, False, True],
66
+ )
67
+
68
+ assert metadata.num_decodes == 2
69
+ assert metadata.num_prefills == 1
70
+ assert metadata.has_initial_states_p.tolist() == [False]
71
+
72
+
73
def test_update_block_table_copies_block_idx_to_persistent_buffers():
74
"""update_block_table() must write block_idx tensors to the current
75
builder's persistent buffers, not leave them pointing to a different
@@ -77,8 +90,8 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
90
)
91
92
# Two builders simulating two KV cache groups with the same MambaSpec.
80
- builder_a = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
81
- builder_b = _ConcreteMambaBuilder(spec, ["layer1"], vllm_config, device)
93
+ builder_a = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
94
+ builder_b = MockMambaBuilder(spec, ["layer1"], vllm_config, device)
95
96
# Sanity: each builder has its own persistent buffer.
97
assert (
@@ -188,7 +201,7 @@ def test_state_indices_tensor_d_includes_num_speculative_blocks():
201
num_speculative_blocks=num_speculative_blocks,
202
)
203
191
- builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
204
+ builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
205
206
expected_cols = (max_model_len // block_size) + num_speculative_blocks
207
assert builder.state_indices_tensor_d.shape == (max_num_seqs, expected_cols)
@@ -221,7 +234,7 @@ def test_block_idx_cudagraph_capture_padded_by_num_reqs():
234
num_speculative_blocks=2,
235
)
236
224
- builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
237
+ builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
238
239
builder.block_idx_last_scheduled_token.fill_(-1)
240
builder.block_idx_last_computed_token.fill_(-1)
@@ -299,7 +312,7 @@ def test_block_idx_prev_step_persistent_buffer_allocated():
312
mamba_cache_mode="all",
313
num_speculative_blocks=2,
314
)
302
- builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
315
+ builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
316
317
assert hasattr(builder, "block_idx_last_scheduled_token_prev_step")
318
assert builder.block_idx_last_scheduled_token_prev_step.shape == (max_num_seqs,)
@@ -323,7 +336,7 @@ def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode():
336
dtypes=(torch.float32,),
337
mamba_cache_mode="all",
338
)
326
- builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
339
+ builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
340
341
assert not hasattr(builder, "block_idx_last_scheduled_token_prev_step")
342
@@ -351,7 +364,7 @@ def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer():
364
mamba_cache_mode="all",
365
num_speculative_blocks=2,
366
)
354
- builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
367
+ builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
368
builder.block_idx_last_scheduled_token.fill_(-1)
369
builder.block_idx_last_computed_token.fill_(-1)
370
builder.block_idx_last_scheduled_token_prev_step.fill_(-1)
tests/v1/attention/utils.py
+40
-1
@@ -24,8 +24,16 @@ from vllm.v1.attention.backend import (
24
AttentionType,
25
CommonAttentionMetadata,
26
)
27
+from vllm.v1.attention.backends.mamba_attn import (
28
+ BaseMambaAttentionMetadata,
29
+ BaseMambaAttentionMetadataBuilder,
30
+)
31
from vllm.v1.attention.backends.registry import AttentionBackendEnum
28
-from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, FullAttentionSpec
32
+from vllm.v1.kv_cache_interface import (
33
+ EncoderOnlyAttentionSpec,
34
+ FullAttentionSpec,
35
+ MambaSpec,
36
+)
37
38
39
@dataclass
@@ -376,3 +384,34 @@ full_cg_backend_configs = {
384
},
385
),
386
}
387
+
388
+
389
+class MockMambaBuilder(BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetadata]):
390
+ """Minimal concrete subclass for testing (base class is ABC)."""
391
+
392
+ metadata_cls = BaseMambaAttentionMetadata
393
+
394
+ @classmethod
395
+ def build_mamba_metadata(
396
+ cls,
397
+ vllm_config: VllmConfig,
398
+ seq_lens: list[int],
399
+ query_lens: list[int],
400
+ is_prefilling: list[bool],
401
+ *,
402
+ device: torch.device | None = None,
403
+ ) -> BaseMambaAttentionMetadata:
404
+ block_size = vllm_config.cache_config.block_size
405
+ device = device or torch.device("cpu")
406
+ mamba_spec = MambaSpec(
407
+ block_size=block_size, shapes=((1,), (1,)), dtypes=(torch.float32,)
408
+ )
409
+ builder = cls(mamba_spec, ["layer0"], vllm_config, device)
410
+ batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
411
+ common_metadata = create_common_attn_metadata(
412
+ batch_spec, block_size=block_size, device=device, arange_block_indices=True
413
+ )
414
+ common_metadata = common_metadata.replace(
415
+ is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool)
416
+ )
417
+ return builder.build(0, common_metadata)
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+25
@@ -8,6 +8,7 @@ from unittest.mock import patch
8
import pytest
9
import torch
10
11
+from tests.v1.attention.utils import MockMambaBuilder
12
from vllm import LLM, SamplingParams
13
from vllm.config import KVTransferConfig
14
from vllm.v1.core.single_type_kv_cache_manager import (
@@ -636,6 +637,30 @@ def test_mamba_n1_d_side(has_mamba, is_hma_required, expected_count):
637
assert is_async is True
638
639
640
+@pytest.mark.cpu_test
641
+def test_mamba_n1_d_side_builds_decode_metadata():
642
+ req = create_request(num_tokens=10, do_remote_prefill=True)
643
+ sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
644
+
645
+ num_computed_tokens, is_async = sched.get_num_new_matched_tokens(
646
+ req, num_computed_tokens=0
647
+ )
648
+
649
+ assert num_computed_tokens == req.num_prompt_tokens - 1
650
+ assert is_async is True
651
+
652
+ vllm_config = create_vllm_config()
653
+ metadata = MockMambaBuilder.build_mamba_metadata(
654
+ vllm_config,
655
+ seq_lens=[req.num_prompt_tokens],
656
+ query_lens=[1],
657
+ is_prefilling=[True],
658
+ )
659
+
660
+ assert metadata.num_decodes == 1
661
+ assert metadata.num_prefills == 0
662
+
663
+
664
@pytest.mark.cpu_test
665
def test_mamba_n1_p_side_truncation():
666
"""P-side: Mamba truncates prompt to N-1, sets max_tokens=1.
vllm/v1/attention/backends/mamba_attn.py
+22
@@ -385,6 +385,28 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
385
self.reorder_batch_threshold if num_accepted_tokens is not None else 1
386
)
387
388
+ # FULL-CG dispatch is shape-based, so one-token prefills with
389
+ # prior Mamba state can replay a decode graph while `is_prefilling`
390
+ # is still true. Treat them as decode/update rows. This is required
391
+ # for NIXL disagg's h(N-1)->N recompute path and for sporadic
392
+ # final single-token prefill chunks that land in a `uniform` FULL-CG
393
+ # batch. Relies on `reorder` putting short extends before pure prefills.
394
+ is_prefilling = common_attn_metadata.is_prefilling
395
+ assert is_prefilling is not None
396
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
397
+ assert seq_lens_cpu is not None
398
+ query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu)
399
+ single_token_prefill_rows = is_prefilling & (query_lens_cpu == 1)
400
+ # First-token prefills have no prior Mamba state and must stay prefills.
401
+ has_prior_state = seq_lens_cpu > 1
402
+ prefill_to_decode = single_token_prefill_rows & has_prior_state
403
+ if torch.any(prefill_to_decode).item():
404
+ is_prefilling = is_prefilling.clone()
405
+ is_prefilling[prefill_to_decode] = False
406
+ common_attn_metadata = common_attn_metadata.replace(
407
+ is_prefilling=is_prefilling
408
+ )
409
+
410
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
411
split_decodes_and_prefills(
412
common_attn_metadata,