| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Equivalence test for ``precopy_mamba_align_fused_kernel``. |
| 4 | |
| 5 | The fused "align" pre-copy must migrate mamba state across block boundaries |
| 6 | with byte-identical semantics to the scalar V1 copy specs |
| 7 | (``get_conv_copy_spec`` / ``get_temporal_copy_spec``): |
| 8 | |
| 9 | * conv state (SD layout, conv_width > 0): shift the sliding window by |
| 10 | ``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` -> |
| 11 | ``state[bt[dst_col], :conv_width - token_bias]``. |
| 12 | * temporal state (conv_width == 0): ``token_bias`` selects the accepted |
| 13 | speculative column -- ``state[bt[src_col + token_bias]]`` -> |
| 14 | ``state[bt[dst_col]]``. |
| 15 | |
| 16 | The kernel must also no-op when ``src_col < 0`` (fresh request) or |
| 17 | ``src_col == dst_col`` (no boundary crossed). V2 callers pass an explicit |
| 18 | ``idx_mapping``; V1 align preprocessing launches in batch order with |
| 19 | ``idx_mapping=None``. |
| 20 | """ |
| 21 | |
| 22 | from __future__ import annotations |
| 23 | |
| 24 | from collections.abc import Callable |
| 25 | from types import SimpleNamespace |
| 26 | from typing import Any |
| 27 | |
| 28 | import numpy as np |
| 29 | import torch |
| 30 | |
| 31 | from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils |
| 32 | from vllm.platforms import current_platform |
| 33 | from vllm.v1.worker import mamba_utils as worker_mamba_utils |
| 34 | from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, precopy_mamba_align_fused_kernel |
| 35 | |
| 36 | _parametrize: Callable[..., Callable[[Any], Any]] |
| 37 | |
| 38 | try: |
| 39 | import pytest |
| 40 | |
| 41 | _cuda_required = pytest.mark.skipif( |
| 42 | not current_platform.is_cuda(), |
| 43 | reason="precopy_mamba_align_fused_kernel needs CUDA/Triton", |
| 44 | ) |
| 45 | _parametrize = pytest.mark.parametrize |
| 46 | except ModuleNotFoundError: # allow running directly as ``python <thisfile>`` |
| 47 | |
| 48 | def _cuda_required(fn): |
| 49 | return fn |
| 50 | |
| 51 | def _no_parametrize(_name, _values): |
| 52 | def _deco(fn): |
| 53 | return fn |
| 54 | |
| 55 | return _deco |
| 56 | |
| 57 | _parametrize = _no_parametrize |
| 58 | |
| 59 | |
| 60 | NUM_LAYERS = 3 |
| 61 | CONV_WIDTH = 4 # conv_kernel - 1 + num_spec |
| 62 | CONV_DIM = 96 |
| 63 | SSM_SHAPE = (4, 16, 16) |
| 64 | MAX_COLS = 8 |
| 65 | |
| 66 | |
| 67 | def _build_state(num_blocks, device, conv_state_dim_first): |
| 68 | """Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools.""" |
| 69 | convs, ssms = [], [] |
| 70 | for _ in range(NUM_LAYERS): |
| 71 | conv_shape = ( |
| 72 | (num_blocks, CONV_DIM, CONV_WIDTH) |
| 73 | if conv_state_dim_first |
| 74 | else (num_blocks, CONV_WIDTH, CONV_DIM) |
| 75 | ) |
| 76 | convs.append( |
| 77 | torch.randn( |
| 78 | *conv_shape, |
| 79 | dtype=torch.bfloat16, |
| 80 | device=device, |
| 81 | ) |
| 82 | ) |
| 83 | ssms.append( |
| 84 | torch.randn(num_blocks, *SSM_SHAPE, dtype=torch.float32, device=device) |
| 85 | ) |
| 86 | return convs, ssms |
| 87 | |
| 88 | |
| 89 | def _build_meta(convs, ssms, device, conv_state_dim_first): |
| 90 | """Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer.""" |
| 91 | n = NUM_LAYERS * 2 |
| 92 | base = torch.zeros(n, dtype=torch.int64, device=device) |
| 93 | blk_stride = torch.zeros(n, dtype=torch.int64, device=device) |
| 94 | elem = torch.zeros(n, dtype=torch.int32, device=device) |
| 95 | inner = torch.zeros(n, dtype=torch.int64, device=device) |
| 96 | width = torch.zeros(n, dtype=torch.int32, device=device) |
| 97 | group = torch.zeros(n, dtype=torch.int32, device=device) |
| 98 | drc = torch.zeros(n, dtype=torch.int32, device=device) # DS rows (unused, SD) |
| 99 | drs = torch.zeros(n, dtype=torch.int64, device=device) |
| 100 | i = 0 |
| 101 | for layer in range(NUM_LAYERS): |
| 102 | conv, ssm = convs[layer], ssms[layer] |
| 103 | # conv (SD): width = size(1), inner = stride(1) |
| 104 | base[i] = conv.data_ptr() |
| 105 | blk_stride[i] = conv.stride(0) * conv.element_size() |
| 106 | elem[i] = conv.element_size() |
| 107 | if conv_state_dim_first: |
| 108 | width[i] = conv.size(2) |
| 109 | inner[i] = 1 |
| 110 | drc[i] = conv.size(1) |
| 111 | drs[i] = conv.stride(1) * conv.element_size() |
| 112 | else: |
| 113 | width[i] = conv.size(1) |
| 114 | inner[i] = conv.stride(1) |
| 115 | i += 1 |
| 116 | # ssm (temporal): width = 0, inner = elems per block |
| 117 | base[i] = ssm.data_ptr() |
| 118 | blk_stride[i] = ssm.stride(0) * ssm.element_size() |
| 119 | elem[i] = ssm.element_size() |
| 120 | width[i] = 0 |
| 121 | inner[i] = ssm[0].numel() |
| 122 | i += 1 |
| 123 | return base, blk_stride, elem, inner, width, group, drc, drs |
| 124 | |
| 125 | |
| 126 | def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs, conv_dim_first): |
| 127 | """Apply the V1 copy semantics on clones, reading from the pre-copy state.""" |
| 128 | conv_pre = [c.clone() for c in convs] |
| 129 | ssm_pre = [s.clone() for s in ssms] |
| 130 | conv_ref = [c.clone() for c in convs] |
| 131 | ssm_ref = [s.clone() for s in ssms] |
| 132 | for r in range(num_reqs): |
| 133 | sc, dc, tb = int(src_col[r]), int(dst_col[r]), int(bias[r]) |
| 134 | if sc < 0 or sc == dc: |
| 135 | continue |
| 136 | sblk, dblk = int(bt[r, sc]), int(bt[r, dc]) |
| 137 | tblk = int(bt[r, sc + tb]) # temporal src column shifted by bias |
| 138 | for layer in range(NUM_LAYERS): |
| 139 | if conv_dim_first: |
| 140 | conv_ref[layer][dblk, :, : CONV_WIDTH - tb] = conv_pre[layer][ |
| 141 | sblk, :, tb: |
| 142 | ] |
| 143 | else: |
| 144 | conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:] |
| 145 | ssm_ref[layer][dblk] = ssm_pre[layer][tblk] |
| 146 | return conv_ref, ssm_ref |
| 147 | |
| 148 | |
| 149 | @_parametrize("conv_state_dim_first", [False, True]) |
| 150 | @_parametrize("num_reqs", [1, 4, 16]) |
| 151 | @_parametrize("token_bias", [0, 1, 2]) |
| 152 | @_parametrize("has_idx_mapping", [True, False]) |
| 153 | @_parametrize("temporal_tiles", [1, _TEMPORAL_TILES]) |
| 154 | @_cuda_required |
| 155 | def test_precopy_matches_v1_copy_specs( |
| 156 | num_reqs, |
| 157 | token_bias, |
| 158 | has_idx_mapping, |
| 159 | conv_state_dim_first, |
| 160 | temporal_tiles, |
| 161 | ): |
| 162 | device = torch.device("cuda") |
| 163 | torch.manual_seed(0) |
| 164 | # Distinct physical block per (req, col) so copies never alias. |
| 165 | num_blocks = num_reqs * MAX_COLS + 1 |
| 166 | bt = torch.empty(num_reqs, MAX_COLS, dtype=torch.int32, device=device) |
| 167 | for r in range(num_reqs): |
| 168 | bt[r] = torch.arange( |
| 169 | 1 + r * MAX_COLS, 1 + (r + 1) * MAX_COLS, dtype=torch.int32, device=device |
| 170 | ) |
| 171 | |
| 172 | # Per-req columns: req 0 fresh (src=-1, skip), req 1 same block (skip), |
| 173 | # the rest cross from col 1 -> col 0 with the given spec token bias. |
| 174 | src_col = torch.full((num_reqs,), 1, dtype=torch.int32, device=device) |
| 175 | dst_col = torch.zeros(num_reqs, dtype=torch.int32, device=device) |
| 176 | bias = torch.full((num_reqs,), token_bias, dtype=torch.int32, device=device) |
| 177 | if num_reqs >= 1: |
| 178 | src_col[0] = -1 # fresh -> no copy |
| 179 | if num_reqs >= 2: |
| 180 | dst_col[1] = 1 # src_col == dst_col -> no copy |
| 181 | |
| 182 | convs, ssms = _build_state(num_blocks, device, conv_state_dim_first) |
| 183 | conv_ref, ssm_ref = _reference( |
| 184 | convs, |
| 185 | ssms, |
| 186 | bt.cpu(), |
| 187 | src_col.cpu(), |
| 188 | dst_col.cpu(), |
| 189 | bias.cpu(), |
| 190 | num_reqs, |
| 191 | conv_state_dim_first, |
| 192 | ) |
| 193 | |
| 194 | base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( |
| 195 | convs, ssms, device, conv_state_dim_first |
| 196 | ) |
| 197 | bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) |
| 198 | idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) |
| 199 | grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles) |
| 200 | precopy_mamba_align_fused_kernel[grid]( |
| 201 | dst_col, |
| 202 | src_col, |
| 203 | bias, |
| 204 | bt_ptrs, |
| 205 | bt.stride(0), |
| 206 | base, |
| 207 | blk_stride, |
| 208 | elem, |
| 209 | inner, |
| 210 | width, |
| 211 | group, |
| 212 | drc, |
| 213 | drs, |
| 214 | idx_mapping if has_idx_mapping else None, |
| 215 | num_reqs, |
| 216 | COPY_BLOCK_SIZE=1024, |
| 217 | CONV_STATE_DIM_FIRST=conv_state_dim_first, |
| 218 | HAS_IDX_MAPPING=has_idx_mapping, |
| 219 | TEMPORAL_TILES=temporal_tiles, |
| 220 | ) |
| 221 | torch.accelerator.synchronize() |
| 222 | |
| 223 | for layer in range(NUM_LAYERS): |
| 224 | torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0) |
| 225 | torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) |
| 226 | |
| 227 | |
| 228 | def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch): |
| 229 | monkeypatch.setattr( |
| 230 | layer_mamba_utils, |
| 231 | "is_conv_state_dim_first", |
| 232 | lambda: True, |
| 233 | ) |
| 234 | state = torch.empty((2, CONV_DIM, CONV_WIDTH), dtype=torch.bfloat16) |
| 235 | |
| 236 | with pytest.raises(AssertionError, match="num_accepted_tokens > 1"): |
| 237 | layer_mamba_utils.get_conv_copy_spec( |
| 238 | state=state, |
| 239 | block_ids=[0, 1], |
| 240 | cur_block_idx=0, |
| 241 | num_accepted_tokens=3, |
| 242 | ) |
| 243 | |
| 244 | |
| 245 | class _FakeCpuGpuBuffer: |
| 246 | def __init__(self, n): |
| 247 | self.np = np.zeros(n, dtype=np.int32) |
| 248 | self.gpu = object() |
| 249 | self.copy_sizes = [] |
| 250 | |
| 251 | def copy_to_gpu(self, n=None): |
| 252 | self.copy_sizes.append(n) |
| 253 | return self.gpu |
| 254 | |
| 255 | |
| 256 | class _FakePrecopyContext: |
| 257 | def __init__(self, n): |
| 258 | self.is_initialized = True |
| 259 | self.mamba_group_ids = [0] |
| 260 | self.mamba_state_idx_buf = _FakeCpuGpuBuffer(n) |
| 261 | self.precopy_src_col_buf = _FakeCpuGpuBuffer(n) |
| 262 | self.precopy_token_bias_buf = _FakeCpuGpuBuffer(n) |
| 263 | self.calls = [] |
| 264 | |
| 265 | def initialize_from_forward_context(self, *args, **kwargs): |
| 266 | raise AssertionError("test context is pre-initialized") |
| 267 | |
| 268 | def run_fused_precopy( |
| 269 | self, |
| 270 | *, |
| 271 | num_reqs, |
| 272 | state_idx_gpu, |
| 273 | src_col_gpu, |
| 274 | token_bias_gpu, |
| 275 | idx_mapping, |
| 276 | ): |
| 277 | self.calls.append( |
| 278 | { |
| 279 | "num_reqs": num_reqs, |
| 280 | "state_idx": self.mamba_state_idx_buf.np[:num_reqs].copy(), |
| 281 | "src_col": self.precopy_src_col_buf.np[:num_reqs].copy(), |
| 282 | "token_bias": self.precopy_token_bias_buf.np[:num_reqs].copy(), |
| 283 | "idx_mapping": idx_mapping, |
| 284 | } |
| 285 | ) |
| 286 | |
| 287 | |
| 288 | def _make_preprocess_case(token_bias): |
| 289 | req_ids = ["fresh", "same", "cross_a", "cross_b"] |
| 290 | scheduler_output = SimpleNamespace( |
| 291 | finished_req_ids=set(), |
| 292 | preempted_req_ids=set(), |
| 293 | scheduled_cached_reqs=SimpleNamespace(resumed_req_ids=set()), |
| 294 | num_scheduled_tokens={ |
| 295 | "fresh": 1, |
| 296 | "same": 1, |
| 297 | "cross_a": 1, |
| 298 | "cross_b": 2, |
| 299 | }, |
| 300 | ) |
| 301 | input_batch = SimpleNamespace( |
| 302 | req_ids=req_ids, |
| 303 | num_accepted_tokens_cpu=np.array( |
| 304 | [token_bias + 1, token_bias + 1, token_bias + 1, 2], |
| 305 | dtype=np.int32, |
| 306 | ), |
| 307 | ) |
| 308 | requests = { |
| 309 | "fresh": SimpleNamespace(req_id="fresh", num_computed_tokens=0), |
| 310 | "same": SimpleNamespace(req_id="same", num_computed_tokens=5), |
| 311 | "cross_a": SimpleNamespace(req_id="cross_a", num_computed_tokens=8), |
| 312 | "cross_b": SimpleNamespace(req_id="cross_b", num_computed_tokens=7), |
| 313 | } |
| 314 | mamba_state_idx = {"same": 1, "cross_a": 0, "cross_b": 1} |
| 315 | return scheduler_output, input_batch, requests, mamba_state_idx |
| 316 | |
| 317 | |
| 318 | @_parametrize("token_bias", [1, 2]) |
| 319 | def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bias): |
| 320 | block_size = 4 |
| 321 | mamba_spec = SimpleNamespace(block_size=block_size, num_speculative_blocks=1) |
| 322 | # Fakes stand in for real config/context objects. The fused path tests |
| 323 | # below only touch the attributes preprocess_mamba actually reads. |
| 324 | cache_config: Any = SimpleNamespace(enable_prefix_caching=True) |
| 325 | kv_cache_config: Any = SimpleNamespace() |
| 326 | scalar_copy_calls = [] |
| 327 | |
| 328 | def fake_collect( |
| 329 | copy_bufs, |
| 330 | kv_cache_config, |
| 331 | mamba_state_copy_funcs, |
| 332 | mamba_group_ids, |
| 333 | src_block_idx, |
| 334 | dest_block_idx, |
| 335 | accept_token_bias, |
| 336 | req_state, |
| 337 | forward_context, |
| 338 | ): |
| 339 | scalar_copy_calls.append( |
| 340 | ( |
| 341 | req_state.req_id, |
| 342 | src_block_idx, |
| 343 | dest_block_idx, |
| 344 | accept_token_bias, |
| 345 | ) |
| 346 | ) |
| 347 | |
| 348 | monkeypatch.setattr(worker_mamba_utils, "collect_mamba_copy_meta", fake_collect) |
| 349 | monkeypatch.setattr( |
| 350 | worker_mamba_utils, "do_mamba_copy_block", lambda copy_bufs: None |
| 351 | ) |
| 352 | |
| 353 | scalar_case = _make_preprocess_case(token_bias) |
| 354 | fused_case = _make_preprocess_case(token_bias) |
| 355 | |
| 356 | scalar_copy_bufs: Any = SimpleNamespace( |
| 357 | mamba_group_ids=[0], |
| 358 | mamba_spec=mamba_spec, |
| 359 | offset=0, |
| 360 | ) |
| 361 | worker_mamba_utils.preprocess_mamba( |
| 362 | scheduler_output=scalar_case[0], |
| 363 | kv_cache_config=kv_cache_config, |
| 364 | cache_config=cache_config, |
| 365 | mamba_state_idx=scalar_case[3], |
| 366 | input_batch=scalar_case[1], |
| 367 | requests=scalar_case[2], |
| 368 | forward_context={}, |
| 369 | mamba_state_copy_funcs={}, |
| 370 | copy_bufs=scalar_copy_bufs, |
| 371 | ) |
| 372 | |
| 373 | ctx: Any = _FakePrecopyContext(len(fused_case[1].req_ids)) |
| 374 | fused_copy_bufs: Any = SimpleNamespace( |
| 375 | mamba_group_ids=[0], |
| 376 | mamba_spec=mamba_spec, |
| 377 | offset=0, |
| 378 | ) |
| 379 | worker_mamba_utils.preprocess_mamba( |
| 380 | scheduler_output=fused_case[0], |
| 381 | kv_cache_config=kv_cache_config, |
| 382 | cache_config=cache_config, |
| 383 | mamba_state_idx=fused_case[3], |
| 384 | input_batch=fused_case[1], |
| 385 | requests=fused_case[2], |
| 386 | forward_context={}, |
| 387 | mamba_state_copy_funcs={}, |
| 388 | copy_bufs=fused_copy_bufs, |
| 389 | align_ctx=ctx, |
| 390 | ) |
| 391 | |
| 392 | assert fused_case[3] == scalar_case[3] |
| 393 | np.testing.assert_array_equal( |
| 394 | fused_case[1].num_accepted_tokens_cpu, |
| 395 | scalar_case[1].num_accepted_tokens_cpu, |
| 396 | ) |
| 397 | assert scalar_copy_calls == [ |
| 398 | ("cross_a", 0, 2, token_bias), |
| 399 | ("cross_b", 1, 2, 1), |
| 400 | ] |
| 401 | assert len(ctx.calls) == 1 |
| 402 | call = ctx.calls[0] |
| 403 | assert call["num_reqs"] == len(fused_case[1].req_ids) |
| 404 | assert call["idx_mapping"] is None |
| 405 | np.testing.assert_array_equal(call["state_idx"], np.array([0, 1, 2, 2])) |
| 406 | np.testing.assert_array_equal(call["src_col"], np.array([-1, -1, 0, 1])) |
| 407 | np.testing.assert_array_equal(call["token_bias"], np.array([0, 0, token_bias, 1])) |
| 408 | fused_copy_calls = [ |
| 409 | (req_id, int(src), int(dst), int(bias)) |
| 410 | for req_id, src, dst, bias in zip( |
| 411 | fused_case[1].req_ids, |
| 412 | call["src_col"], |
| 413 | call["state_idx"], |
| 414 | call["token_bias"], |
| 415 | ) |
| 416 | if int(src) != -1 and int(src) != int(dst) |
| 417 | ] |
| 418 | assert fused_copy_calls == scalar_copy_calls |
| 419 | |
| 420 | |
| 421 | if __name__ == "__main__": |
| 422 | from itertools import product |
| 423 | |
| 424 | _CASES = product( |
| 425 | (1, 4, 16), # num_reqs |
| 426 | (0, 1, 2), # token_bias |
| 427 | (True, False), # has_idx_mapping |
| 428 | (False, True), # conv_state_dim_first |
| 429 | (1, _TEMPORAL_TILES), # temporal_tiles |
| 430 | ) |
| 431 | for nr, tb, mapping, dim_first, tt in _CASES: |
| 432 | test_precopy_matches_v1_copy_specs(nr, tb, mapping, dim_first, tt) |
| 433 | print( |
| 434 | f"OK num_reqs={nr} token_bias={tb} has_idx_mapping={mapping} " |
| 435 | f"conv_dim_first={dim_first} temporal_tiles={tt}" |
| 436 | ) |