@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 165 lines 5.47 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5 from unittest.mock import Mock, patch
6
7 import numpy as np
8 import pytest
9 import torch
10
11 from vllm.config.mamba import MambaBackendEnum
12 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
13 from vllm.model_executor.warmup import replayssm_warmup as warmup
14 from vllm.platforms import current_platform
15 from vllm.utils.flashinfer import has_flashinfer
16
17 pytestmark = pytest.mark.skipif(
18 not current_platform.is_cuda() or not has_flashinfer(),
19 reason="FlashInfer ReplaySSM warmup tests require CUDA and FlashInfer",
20 )
21
22 PREFILL_KWARGS = {
23 "num_tokens": 128,
24 "skip_eplb": True,
25 "is_profile": True,
26 "randomize_inputs": True,
27 }
28
29
30 def _autotune_runner(
31 *,
32 use_v2_model_runner: bool = False,
33 query_len: int = 6,
34 max_num_seqs: int = 32,
35 num_blocks: int = 17,
36 max_num_tokens: int = 100,
37 use_replayssm: bool = True,
38 backend: MambaBackendEnum = MambaBackendEnum.FLASHINFER,
39 ) -> SimpleNamespace:
40 return SimpleNamespace(
41 vllm_config=SimpleNamespace(
42 cache_config=SimpleNamespace(use_replayssm=use_replayssm),
43 mamba_config=SimpleNamespace(backend=backend),
44 use_v2_model_runner=use_v2_model_runner,
45 ),
46 uniform_decode_query_len=query_len,
47 decode_query_len=query_len,
48 max_num_tokens=max_num_tokens,
49 scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
50 kv_cache_config=SimpleNamespace(num_blocks=num_blocks),
51 )
52
53
54 @pytest.mark.parametrize(
55 ("runner_kwargs", "expected_num_reqs"),
56 [
57 # max_num_seqs (32) vs max_num_tokens // query_len (16) vs blocks-1 (16).
58 (dict(query_len=6, use_v2_model_runner=False), 16),
59 (dict(query_len=6, use_v2_model_runner=True), 16),
60 # num_blocks - 1 is the binding constraint.
61 (dict(query_len=1, max_num_tokens=128, max_num_seqs=64, num_blocks=5), 4),
62 ],
63 ids=["v1", "v2", "clamped_to_state_capacity"],
64 )
65 def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs):
66 query_len = runner_kwargs["query_len"]
67 with patch.object(
68 warmup, "flashinfer_replayssm_autotune_supported", return_value=True
69 ):
70 result = warmup._replayssm_autotune_kwargs(
71 _autotune_runner(**runner_kwargs), PREFILL_KWARGS
72 )
73
74 assert result is not None
75 max_num_reqs, decode_kwargs = result
76 assert max_num_reqs == expected_num_reqs
77 expected_kwargs = {
78 **PREFILL_KWARGS,
79 "num_tokens": expected_num_reqs * query_len,
80 "uniform_decode": True,
81 }
82 if runner_kwargs.get("use_v2_model_runner"):
83 expected_kwargs["valid_dummy_state_slots"] = True
84 else:
85 expected_kwargs.update(
86 allow_microbatching=False,
87 force_attention=True,
88 profile_seq_lens=query_len + 1,
89 )
90 assert decode_kwargs == expected_kwargs
91
92
93 @pytest.mark.parametrize(
94 ("runner_kwargs", "flashinfer_supported"),
95 [
96 (dict(use_replayssm=False), True),
97 (dict(backend=MambaBackendEnum.TRITON), True),
98 ({}, False),
99 (dict(num_blocks=1), True),
100 ],
101 ids=[
102 "replayssm_disabled",
103 "non_flashinfer_backend",
104 "kernel_unavailable",
105 "zero_non_padding_slots",
106 ],
107 )
108 def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported):
109 with patch.object(
110 warmup,
111 "flashinfer_replayssm_autotune_supported",
112 return_value=flashinfer_supported,
113 ):
114 result = warmup._replayssm_autotune_kwargs(
115 _autotune_runner(**runner_kwargs), PREFILL_KWARGS
116 )
117 assert result is None
118
119
120 def test_replayssm_autotune_slots_restore_state_and_trackers():
121 mixer = MambaMixer2.__new__(MambaMixer2)
122 torch.nn.Module.__init__(mixer)
123 mixer.use_replayssm = True
124 mixer.replayssm_buffer_len = 16
125 mixer.kv_cache = (
126 torch.full((4, 2), 3.0),
127 torch.full((4, 2), 3.0),
128 *(torch.full((4, 2, 17), 3.0) for _ in range(3)),
129 )
130 mixer._replayssm_ring_start = torch.full((4,), 3, dtype=torch.int32)
131 mixer._replayssm_prev_num_accepted = torch.full((4,), 3, dtype=torch.int32)
132 mixer._replayssm_prev_query_len = torch.full((4,), 3, dtype=torch.int32)
133 tracked = (
134 *mixer.kv_cache,
135 mixer._replayssm_ring_start,
136 mixer._replayssm_prev_num_accepted,
137 mixer._replayssm_prev_query_len,
138 )
139
140 block_ids = np.arange(10, 14, dtype=np.int32).reshape(4, 1)
141 original_block_ids = block_ids.copy()
142 block_table = SimpleNamespace(block_table=SimpleNamespace(np=block_ids))
143 multi_group_block_table = SimpleNamespace(
144 block_tables=[block_table], commit_block_table=Mock()
145 )
146 runner = SimpleNamespace(
147 vllm_config=SimpleNamespace(use_v2_model_runner=False),
148 input_batch=SimpleNamespace(block_table=multi_group_block_table),
149 get_model=lambda: SimpleNamespace(modules=lambda: (mixer,)),
150 )
151
152 with warmup._temporary_replayssm_autotune_state(runner, 2):
153 assert block_ids[:2, 0].tolist() == [1, 2]
154 for tensor in tracked:
155 tensor[1:3].fill_(9)
156
157 assert np.array_equal(block_ids, original_block_ids)
158 assert multi_group_block_table.commit_block_table.call_args_list == [
159 ((2,), {}),
160 ((2,), {}),
161 ]
162 for tensor in tracked:
163 assert torch.count_nonzero(tensor[1:3]) == 0
164 assert torch.all(tensor[0] == 3)
165 assert torch.all(tensor[3] == 3)