| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for BaseMambaAttentionMetadataBuilder batch classification.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | from tests.v1.attention.utils import MockMambaBuilder |
| 8 | from vllm.config.compilation import CUDAGraphMode |
| 9 | |
| 10 | |
| 11 | def _make_vllm_config( |
| 12 | max_model_len: int, |
| 13 | max_num_seqs: int, |
| 14 | num_speculative_tokens: int = 0, |
| 15 | block_size: int | None = None, |
| 16 | ): |
| 17 | """Create a minimal mock VllmConfig with only the fields the builder |
| 18 | accesses, avoiding any model download / HF config inspection.""" |
| 19 | speculative_config = ( |
| 20 | SimpleNamespace( |
| 21 | num_speculative_tokens=num_speculative_tokens, |
| 22 | parallel_drafting=False, |
| 23 | ) |
| 24 | if num_speculative_tokens > 0 |
| 25 | else None |
| 26 | ) |
| 27 | return SimpleNamespace( |
| 28 | cache_config=SimpleNamespace( |
| 29 | block_size=block_size, |
| 30 | mamba_cache_mode="none", |
| 31 | use_replayssm=False, |
| 32 | replayssm_buffer_len=16, |
| 33 | ), |
| 34 | compilation_config=SimpleNamespace( |
| 35 | cudagraph_mode=CUDAGraphMode.FULL, |
| 36 | max_cudagraph_capture_size=None, |
| 37 | ), |
| 38 | speculative_config=speculative_config, |
| 39 | num_speculative_tokens=num_speculative_tokens, |
| 40 | parallel_config=SimpleNamespace(decode_context_parallel_size=1), |
| 41 | scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs), |
| 42 | model_config=SimpleNamespace(max_model_len=max_model_len), |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | def test_mamba_single_token_prompt_runs_as_prefill(): |
| 47 | seq_lens = [8, 9, 1] |
| 48 | config = _make_vllm_config(256, len(seq_lens), block_size=16) |
| 49 | metadata = MockMambaBuilder.build_mamba_metadata( |
| 50 | config, |
| 51 | seq_lens=seq_lens, |
| 52 | query_lens=[1] * len(seq_lens), |
| 53 | is_prefilling=[False, False, True], |
| 54 | ) |
| 55 | |
| 56 | assert metadata.num_decodes == 2 |
| 57 | assert metadata.num_prefills == 1 |
| 58 | assert metadata.has_initial_states_p.tolist() == [False] |