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, 03:00 UTC
[Perf] [Hybrid] Fused Triton kernel for GPU-side Mamba state postprocessing (#40172)
Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Francesco Fusco committed
May 21, 2026 at 13:50 UTC
b730c4635288d75da4788bc28d8d26b5e5c3726c
4 files changed
+2786
-167
tests/v1/e2e/general/test_mamba_prefix_cache.py
+4
-48
@@ -293,7 +293,6 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable):
293
294
def get_fake_process_mamba_fn(
295
original_preprocess_mamba_fn: Callable,
296
- original_post_process_mamba_fn: Callable,
296
original_copy_fn: Callable,
297
):
298
copy_info: tuple[list[int], list[int], list[int]] | None = None
@@ -361,45 +360,6 @@ def get_fake_process_mamba_fn(
360
)
361
return ret
362
364
- def fake_post_process_mamba_fn(
365
- scheduler_output: SchedulerOutput,
366
- kv_cache_config: KVCacheConfig,
367
- cache_config: CacheConfig,
368
- input_batch: GPUInputBatch,
369
- requests: dict[str, CachedRequestState],
370
- mamba_state_idx: dict[str, int],
371
- num_spec_tokens: int,
372
- num_reqs: int,
373
- *,
374
- forward_context: dict[str, Any] | None = None,
375
- mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...] | None = None,
376
- copy_bufs: mamba_utils.MambaCopyBuffers | None = None,
377
- ):
378
- nonlocal copy_info
379
- copy_info = None
380
- ret = original_post_process_mamba_fn(
381
- scheduler_output,
382
- kv_cache_config,
383
- cache_config,
384
- input_batch,
385
- requests,
386
- mamba_state_idx,
387
- num_spec_tokens,
388
- num_reqs,
389
- forward_context=forward_context,
390
- mamba_state_copy_funcs=mamba_state_copy_funcs,
391
- copy_bufs=copy_bufs,
392
- )
393
- if cur_step_action is not None:
394
- assert forward_context is not None
395
- check_copy_info(
396
- cur_step_action.postprocess_copy_idx,
397
- kv_cache_config,
398
- forward_context,
399
- input_batch,
400
- )
401
- return ret
402
-
363
def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):
364
nonlocal copy_info
365
assert copy_info is None
@@ -410,7 +370,7 @@ def get_fake_process_mamba_fn(
370
copy_info = (src_state_list, dest_state_list, num_elements_list)
371
return original_copy_fn(copy_bufs)
372
413
- return fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn
373
+ return fake_preprocess_mamba_fn, fake_copy_fn
374
375
376
def run_ref_mamba_state_in_subprocess() -> None:
@@ -522,15 +482,11 @@ def apply_patch(monkeypatch: pytest.MonkeyPatch):
482
fake_allocate_slots_fn = get_fake_allocate_slots_fn(KVCacheManager.allocate_slots)
483
monkeypatch.setattr(KVCacheManager, "allocate_slots", fake_allocate_slots_fn)
484
525
- fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn = (
526
- get_fake_process_mamba_fn(
527
- mamba_utils.preprocess_mamba,
528
- mamba_utils.postprocess_mamba,
529
- mamba_utils.do_mamba_copy_block,
530
- )
485
+ fake_preprocess_mamba_fn, fake_copy_fn = get_fake_process_mamba_fn(
486
+ mamba_utils.preprocess_mamba,
487
+ mamba_utils.do_mamba_copy_block,
488
)
489
monkeypatch.setattr(mamba_utils, "preprocess_mamba", fake_preprocess_mamba_fn)
533
- monkeypatch.setattr(mamba_utils, "postprocess_mamba", fake_post_process_mamba_fn)
490
monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
491
492
tests/v1/worker/test_mamba_utils.py
+2071
-6
@@ -1,9 +1,91 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+from dataclasses import dataclass
5
+from typing import Any
6
from unittest.mock import MagicMock, patch
7
8
+import numpy as np
9
+import pytest
10
+import torch
11
+
12
+from vllm.model_executor.layers.mamba.mamba_utils import (
13
+ get_conv_copy_spec,
14
+ get_temporal_copy_spec,
15
+)
16
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
6
-from vllm.v1.worker.mamba_utils import preprocess_mamba
17
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec
18
+from vllm.v1.worker.mamba_utils import (
19
+ MambaCopyBuffers,
20
+ MambaSpecDecodeGPUContext,
21
+ collect_mamba_copy_meta,
22
+ do_mamba_copy_block,
23
+ preprocess_mamba,
24
+)
25
+
26
+MambaStateCopyFunc = Callable[..., Any]
27
+
28
+# Conv + temporal copy specs, in the order the tests' MambaSpec shapes expect.
29
+_COPY_FUNCS: tuple[MambaStateCopyFunc, ...] = (
30
+ get_conv_copy_spec,
31
+ get_temporal_copy_spec,
32
+)
33
+
34
+
35
+def postprocess_mamba(
36
+ scheduler_output: "SchedulerOutput",
37
+ kv_cache_config: "KVCacheConfig",
38
+ input_batch: Any,
39
+ requests: dict[str, Any],
40
+ forward_context: dict[str, Any],
41
+ mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
42
+ copy_bufs: "MambaCopyBuffers",
43
+):
44
+ """CPU reference for the align-mode postprocess.
45
+
46
+ Used as a golden against the GPU fused kernel (``postprocess_mamba_align_gpu``).
47
+ Mirrors what the production code did before the fused kernel replaced it;
48
+ kept here because production no longer has a CPU implementation.
49
+ """
50
+ assert input_batch.mamba_state_idx_cpu is not None
51
+ num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens
52
+ scheduled_spec_decode_tokens_dict = scheduler_output.scheduled_spec_decode_tokens
53
+ num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu
54
+ mamba_state_idx_cpu = input_batch.mamba_state_idx_cpu
55
+ mamba_group_ids = copy_bufs.mamba_group_ids
56
+ mamba_spec = copy_bufs.mamba_spec
57
+ copy_bufs.offset = 0
58
+ for i, req_id in enumerate(input_batch.req_ids):
59
+ req_state = requests[req_id]
60
+ num_computed_tokens = req_state.num_computed_tokens
61
+ num_draft_tokens = len(scheduled_spec_decode_tokens_dict.get(req_id, []))
62
+ num_scheduled_tokens = num_scheduled_tokens_dict[req_id]
63
+ num_accepted_tokens = num_accepted_tokens_cpu[i]
64
+ num_tokens_running_state = (
65
+ num_computed_tokens + num_scheduled_tokens - num_draft_tokens
66
+ )
67
+ new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens - 1
68
+ aligned_new_computed_tokens = (
69
+ new_num_computed_tokens // mamba_spec.block_size * mamba_spec.block_size
70
+ )
71
+ if aligned_new_computed_tokens >= num_tokens_running_state:
72
+ accept_token_bias = aligned_new_computed_tokens - num_tokens_running_state
73
+ src_block_idx = mamba_state_idx_cpu[i]
74
+ dest_block_idx = aligned_new_computed_tokens // mamba_spec.block_size - 1
75
+ collect_mamba_copy_meta(
76
+ copy_bufs,
77
+ kv_cache_config,
78
+ mamba_state_copy_funcs,
79
+ mamba_group_ids,
80
+ src_block_idx,
81
+ dest_block_idx,
82
+ accept_token_bias,
83
+ req_state,
84
+ forward_context,
85
+ )
86
+ if src_block_idx == dest_block_idx:
87
+ num_accepted_tokens_cpu[i] = 1
88
+ do_mamba_copy_block(copy_bufs)
89
90
91
def _make_scheduler_output(
@@ -38,7 +120,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
120
input_batch = MagicMock(req_ids=[])
121
copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec)
122
41
- mamba_state_idx = {
123
+ mamba_state_idx: dict[str, int] = {
124
"finished": 1,
125
"preempted": 2,
126
"resumed": 3, # only in resumed_req_ids, NOT in preempted
@@ -56,14 +138,1997 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
138
):
139
preprocess_mamba(
140
sched,
59
- MagicMock(),
141
+ MagicMock(), # kv_cache_config
142
cache_config,
143
mamba_state_idx,
144
input_batch,
63
- {},
64
- {},
65
- (),
145
+ {}, # requests
146
+ {}, # forward_context
147
+ (), # mamba_state_copy_funcs
148
copy_bufs,
149
)
150
151
assert mamba_state_idx == {"keep": 99}
152
+
153
+
154
+# -----------------------------------------------------------------------------
155
+# Golden tests for postprocess_mamba_fused_kernel
156
+# -----------------------------------------------------------------------------
157
+
158
+
159
+@dataclass
160
+class _TestConfig:
161
+ """Common test configuration for fused kernel tests."""
162
+
163
+ block_size: int = 16
164
+ num_blocks: int = 32
165
+ num_layers: int = 2
166
+ num_reqs: int = 4
167
+ max_num_reqs: int = 8
168
+ # Conv state shape: [num_blocks, conv_width, inner_dim]
169
+ conv_width: int = 4
170
+ conv_inner_dim: int = 64
171
+ # Temporal state shape: [num_blocks, state_dim]
172
+ temporal_state_dim: int = 128
173
+ dtype: torch.dtype = torch.float16
174
+
175
+
176
+class _MockCpuGpuBuffer:
177
+ """Mock CpuGpuBuffer for testing without pinned memory."""
178
+
179
+ def __init__(self, size: int, dtype: torch.dtype, device: torch.device):
180
+ self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
181
+ self.gpu = torch.zeros(size, dtype=dtype, device=device)
182
+ self.np = self.cpu.numpy()
183
+
184
+ def copy_to_gpu(self, n: int | None = None) -> torch.Tensor:
185
+ if n is None:
186
+ return self.gpu.copy_(self.cpu, non_blocking=True)
187
+ return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)
188
+
189
+
190
+def _make_postprocess_scheduler_output(
191
+ req_ids: list[str],
192
+ num_scheduled_tokens: dict[str, int],
193
+ scheduled_spec_decode_tokens: dict[str, list] | None = None,
194
+) -> SchedulerOutput:
195
+ """Create a minimal SchedulerOutput for postprocess testing."""
196
+ cached = CachedRequestData.make_empty()
197
+ return SchedulerOutput(
198
+ scheduled_new_reqs=[],
199
+ scheduled_cached_reqs=cached,
200
+ num_scheduled_tokens=num_scheduled_tokens,
201
+ total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
202
+ scheduled_spec_decode_tokens=scheduled_spec_decode_tokens or {},
203
+ scheduled_encoder_inputs={},
204
+ num_common_prefix_blocks=[],
205
+ finished_req_ids=set(),
206
+ free_encoder_mm_hashes=[],
207
+ preempted_req_ids=set(),
208
+ )
209
+
210
+
211
+def _make_mock_attention(
212
+ conv_state: torch.Tensor, temporal_state: torch.Tensor
213
+) -> MagicMock:
214
+ """Create a mock attention object with kv_cache."""
215
+ attention = MagicMock()
216
+ attention.kv_cache = [conv_state, temporal_state]
217
+ return attention
218
+
219
+
220
+def _make_dual_states(
221
+ cfg: "_TestConfig",
222
+ layer_names: list[str],
223
+ device: torch.device,
224
+ *,
225
+ num_blocks: int | None = None,
226
+) -> tuple[
227
+ list[torch.Tensor],
228
+ list[torch.Tensor],
229
+ list[torch.Tensor],
230
+ list[torch.Tensor],
231
+ dict[str, MagicMock],
232
+ dict[str, MagicMock],
233
+]:
234
+ """Allocate conv+temporal state tensors for the Python path, clone them for
235
+ the GPU path, and build matching ``forward_context`` dicts for both.
236
+
237
+ Returns ``(conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu)``
238
+ where the four state lists are parallel to ``layer_names``.
239
+ """
240
+ n_blocks = num_blocks if num_blocks is not None else cfg.num_blocks
241
+ conv_py = [
242
+ torch.randn(
243
+ n_blocks,
244
+ cfg.conv_width,
245
+ cfg.conv_inner_dim,
246
+ dtype=cfg.dtype,
247
+ device=device,
248
+ )
249
+ for _ in layer_names
250
+ ]
251
+ temporal_py = [
252
+ torch.randn(n_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device)
253
+ for _ in layer_names
254
+ ]
255
+ conv_gpu = [s.clone() for s in conv_py]
256
+ temporal_gpu = [s.clone() for s in temporal_py]
257
+ fwd_py = {
258
+ name: _make_mock_attention(c, t)
259
+ for name, c, t in zip(layer_names, conv_py, temporal_py)
260
+ }
261
+ fwd_gpu = {
262
+ name: _make_mock_attention(c, t)
263
+ for name, c, t in zip(layer_names, conv_gpu, temporal_gpu)
264
+ }
265
+ return conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu
266
+
267
+
268
+def _make_dual_layer_state(
269
+ cfg: "_TestConfig",
270
+ device: torch.device,
271
+ *,
272
+ num_blocks: int | None = None,
273
+ layer_name: str = "layer_0",
274
+) -> tuple[
275
+ torch.Tensor,
276
+ torch.Tensor,
277
+ torch.Tensor,
278
+ torch.Tensor,
279
+ dict[str, MagicMock],
280
+ dict[str, MagicMock],
281
+]:
282
+ """Single-layer convenience form of ``_make_dual_states``."""
283
+ conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu = _make_dual_states(
284
+ cfg, [layer_name], device, num_blocks=num_blocks
285
+ )
286
+ return conv_py[0], temporal_py[0], conv_gpu[0], temporal_gpu[0], fwd_py, fwd_gpu
287
+
288
+
289
+def _make_kv_cache_config(cfg: _TestConfig, layer_names: list[str]) -> KVCacheConfig:
290
+ """Create a KVCacheConfig with mamba groups."""
291
+ mamba_spec = MambaSpec(
292
+ block_size=cfg.block_size,
293
+ shapes=(
294
+ (cfg.conv_width, cfg.conv_inner_dim),
295
+ (cfg.temporal_state_dim,),
296
+ ),
297
+ dtypes=(cfg.dtype, cfg.dtype),
298
+ mamba_cache_mode="all",
299
+ )
300
+ group = KVCacheGroupSpec(
301
+ layer_names=layer_names,
302
+ kv_cache_spec=mamba_spec,
303
+ )
304
+ return KVCacheConfig(
305
+ num_blocks=cfg.num_blocks,
306
+ kv_cache_tensors=[],
307
+ kv_cache_groups=[group],
308
+ )
309
+
310
+
311
+def _make_input_batch(
312
+ req_ids: list[str],
313
+ num_accepted_tokens: list[int],
314
+ mamba_state_idx: list[int],
315
+) -> MagicMock:
316
+ """Create a mock GPUInputBatch."""
317
+ batch = MagicMock()
318
+ batch.req_ids = req_ids
319
+ batch.req_id_to_index = {rid: i for i, rid in enumerate(req_ids)}
320
+ # Use numpy arrays so modifications persist
321
+ batch.num_accepted_tokens_cpu = np.array(num_accepted_tokens, dtype=np.int32)
322
+ batch.mamba_state_idx_cpu = np.array(mamba_state_idx, dtype=np.int32)
323
+ return batch
324
+
325
+
326
+def _make_requests(
327
+ req_ids: list[str],
328
+ num_computed_tokens: list[int],
329
+ block_ids_per_req: list[list[int]],
330
+) -> dict[str, MagicMock]:
331
+ """Create mock CachedRequestState objects."""
332
+ requests = {}
333
+ for i, req_id in enumerate(req_ids):
334
+ req = MagicMock()
335
+ req.num_computed_tokens = num_computed_tokens[i]
336
+ req.block_ids = {0: block_ids_per_req[i]} # group_id=0
337
+ requests[req_id] = req
338
+ return requests
339
+
340
+
341
+def _make_copy_bufs(
342
+ cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device
343
+) -> MambaCopyBuffers:
344
+ """Create MambaCopyBuffers for the Python path."""
345
+
346
+ def make_buffer(n, dtype):
347
+ return _MockCpuGpuBuffer(n, dtype, device)
348
+
349
+ return MambaCopyBuffers.create(
350
+ max_num_reqs=cfg.max_num_reqs,
351
+ kv_cache_config=kv_cache_config,
352
+ copy_funcs=(get_conv_copy_spec, get_temporal_copy_spec),
353
+ make_buffer=make_buffer,
354
+ )
355
+
356
+
357
+def _make_gpu_ctx(
358
+ cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device
359
+) -> MambaSpecDecodeGPUContext:
360
+ """Create MambaSpecDecodeGPUContext for the GPU path."""
361
+
362
+ def make_buffer(n, dtype):
363
+ return _MockCpuGpuBuffer(n, dtype, device)
364
+
365
+ return MambaSpecDecodeGPUContext.create(
366
+ max_num_reqs=cfg.max_num_reqs,
367
+ kv_cache_config=kv_cache_config,
368
+ num_state_types=2,
369
+ device=device,
370
+ make_buffer=make_buffer,
371
+ )
372
+
373
+
374
+def _run_gpu_postprocess(
375
+ gpu_ctx: MambaSpecDecodeGPUContext,
376
+ *,
377
+ kv_cache_config: KVCacheConfig,
378
+ forward_context: dict[str, Any],
379
+ copy_funcs: tuple,
380
+ block_table: torch.Tensor,
381
+ req_ids: list[str],
382
+ num_accepted_tokens: list[int],
383
+ mamba_state_idx: list[int],
384
+ num_scheduled_tokens: dict[str, int],
385
+ num_computed_tokens: list[int],
386
+ num_draft_tokens: dict[str, int],
387
+ device: torch.device,
388
+) -> None:
389
+ """Initialize the GPU context against `block_table`, run the fused
390
+ postprocess kernel for `req_ids`, and synchronize."""
391
+
392
+ def t(values):
393
+ return torch.tensor(values, dtype=torch.int32, device=device)
394
+
395
+ gpu_ctx.initialize_from_forward_context(
396
+ kv_cache_config, forward_context, copy_funcs, [block_table]
397
+ )
398
+ gpu_ctx.run_fused_postprocess(
399
+ num_reqs=len(req_ids),
400
+ num_accepted_tokens_gpu=t(num_accepted_tokens),
401
+ mamba_state_idx_gpu=t(mamba_state_idx),
402
+ num_scheduled_tokens_gpu=t([num_scheduled_tokens[r] for r in req_ids]),
403
+ num_computed_tokens_gpu=t(num_computed_tokens),
404
+ num_draft_tokens_gpu=t([num_draft_tokens.get(r, 0) for r in req_ids]),
405
+ )
406
+ torch.accelerator.synchronize()
407
+
408
+
409
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
410
+class TestPostprocessMambaFusedKernel:
411
+ """Tests for postprocess_mamba_fused_kernel comparing GPU vs CPU paths."""
412
+
413
+ @pytest.fixture
414
+ def device(self):
415
+ return torch.device("cuda:0")
416
+
417
+ @pytest.fixture
418
+ def test_config(self):
419
+ return _TestConfig()
420
+
421
+ def test_matches_python_postprocess_mamba(self, device, test_config):
422
+ """
423
+ Golden test: GPU kernel produces identical results to Python impl.
424
+
425
+ This test:
426
+ 1. Sets up identical initial state for both paths
427
+ 2. Runs Python postprocess_mamba (modifies states via batch_memcpy)
428
+ 3. Runs GPU fused kernel (modifies states directly)
429
+ 4. Compares resulting state tensors and num_accepted_tokens
430
+ """
431
+ cfg = test_config
432
+ torch.manual_seed(42)
433
+
434
+ # Test scenario: 4 requests with different copy conditions
435
+ # Copy needed when: aligned_new_computed >= num_tokens_running_state
436
+ # where: num_tokens_running_state = num_computed + num_scheduled - num_draft
437
+ # new_num_computed = num_tokens_running_state + num_accepted - 1
438
+ # aligned_new_computed = (new_num_computed // block_size) * block_size
439
+ req_ids = ["req_0", "req_1", "req_2", "req_3"]
440
+
441
+ # Configure requests so some need copies, some don't
442
+ # block_size = 16
443
+ # req_0: running=60+5-2=63, new=63+3-1=65, aligned=64 >= 63 -> COPY
444
+ # req_1: running=30+3-0=33, new=33+2-1=34, aligned=32 < 33 -> NO COPY
445
+ # req_2: running=45+8-3=50, new=50+4-1=53, aligned=48 < 50 -> NO COPY
446
+ # req_3: running=10+6-0=16, new=16+2-1=17, aligned=16 >= 16 -> COPY
447
+ num_computed_tokens = [60, 30, 45, 10]
448
+ num_scheduled_tokens = {"req_0": 5, "req_1": 3, "req_2": 8, "req_3": 6}
449
+ num_draft_tokens = {"req_0": 2, "req_1": 0, "req_2": 3, "req_3": 0}
450
+ num_accepted_tokens = [3, 2, 4, 2]
451
+ mamba_state_idx = [3, 1, 2, 0] # source block indices
452
+
453
+ # Block IDs for each request (simulate block table)
454
+ block_ids_per_req = [
455
+ list(range(8)), # req_0: blocks 0-7
456
+ list(range(8, 16)), # req_1: blocks 8-15
457
+ list(range(16, 24)), # req_2: blocks 16-23
458
+ list(range(24, 32)), # req_3: blocks 24-31
459
+ ]
460
+
461
+ layer_names = [f"layer_{i}" for i in range(cfg.num_layers)]
462
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
463
+
464
+ (
465
+ conv_states_py,
466
+ temporal_states_py,
467
+ conv_states_gpu,
468
+ temporal_states_gpu,
469
+ forward_context_py,
470
+ forward_context_gpu,
471
+ ) = _make_dual_states(cfg, layer_names, device)
472
+
473
+ # --- Run Python path ---
474
+ scheduler_output = _make_postprocess_scheduler_output(
475
+ req_ids,
476
+ num_scheduled_tokens,
477
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
478
+ )
479
+ input_batch_py = _make_input_batch(
480
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
481
+ )
482
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
483
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
484
+
485
+ postprocess_mamba(
486
+ scheduler_output,
487
+ kv_cache_config,
488
+ input_batch_py,
489
+ requests,
490
+ forward_context_py,
491
+ _COPY_FUNCS,
492
+ copy_bufs,
493
+ )
494
+ torch.accelerator.synchronize()
495
+
496
+ # --- Run GPU path ---
497
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
498
+
499
+ num_reqs = len(req_ids)
500
+ max_blocks = max(len(b) for b in block_ids_per_req)
501
+ block_table_gpu = torch.zeros(
502
+ num_reqs, max_blocks, dtype=torch.int32, device=device
503
+ )
504
+ for i, block_ids in enumerate(block_ids_per_req):
505
+ block_table_gpu[i, : len(block_ids)] = torch.tensor(
506
+ block_ids, dtype=torch.int32
507
+ )
508
+ _run_gpu_postprocess(
509
+ gpu_ctx,
510
+ kv_cache_config=kv_cache_config,
511
+ forward_context=forward_context_gpu,
512
+ copy_funcs=_COPY_FUNCS,
513
+ block_table=block_table_gpu,
514
+ req_ids=req_ids,
515
+ num_accepted_tokens=num_accepted_tokens,
516
+ mamba_state_idx=mamba_state_idx,
517
+ num_scheduled_tokens=num_scheduled_tokens,
518
+ num_computed_tokens=num_computed_tokens,
519
+ num_draft_tokens=num_draft_tokens,
520
+ device=device,
521
+ )
522
+
523
+ # --- Compare results ---
524
+ # 1. Compare state tensors
525
+ for i in range(cfg.num_layers):
526
+ torch.testing.assert_close(
527
+ conv_states_gpu[i],
528
+ conv_states_py[i],
529
+ msg=f"Conv state mismatch at layer {i}",
530
+ )
531
+ torch.testing.assert_close(
532
+ temporal_states_gpu[i],
533
+ temporal_states_py[i],
534
+ msg=f"Temporal state mismatch at layer {i}",
535
+ )
536
+
537
+ # 2. Compare num_accepted_tokens updates
538
+ expected_accepted = torch.tensor(
539
+ input_batch_py.num_accepted_tokens_cpu[:num_reqs],
540
+ dtype=torch.int32,
541
+ device=device,
542
+ )
543
+ torch.testing.assert_close(
544
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
545
+ expected_accepted,
546
+ msg="num_accepted_tokens mismatch",
547
+ )
548
+
549
+ def test_no_copy_when_not_needed(self, device, test_config):
550
+ """Kernel should not modify state when no copy is needed."""
551
+ cfg = test_config
552
+ torch.manual_seed(123)
553
+
554
+ # Single request where no copy is needed:
555
+ # running = 30 + 3 = 33, new = 33 + 1 - 1 = 33, aligned = 32 < 33
556
+ req_ids = ["req_0"]
557
+ num_computed_tokens = [30]
558
+ num_scheduled_tokens = {"req_0": 3}
559
+ num_draft_tokens: dict[str, int] = {}
560
+ num_accepted_tokens = [1]
561
+ mamba_state_idx = [1]
562
+ block_ids_per_req = [list(range(8))]
563
+
564
+ layer_names = ["layer_0"]
565
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
566
+
567
+ # Create state tensor
568
+ conv_state = torch.randn(
569
+ cfg.num_blocks,
570
+ cfg.conv_width,
571
+ cfg.conv_inner_dim,
572
+ dtype=cfg.dtype,
573
+ device=device,
574
+ )
575
+ temporal_state = torch.randn(
576
+ cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device
577
+ )
578
+
579
+ # Clone to verify no modification
580
+ conv_state_orig = conv_state.clone()
581
+ temporal_state_orig = temporal_state.clone()
582
+
583
+ forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)}
584
+
585
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
586
+
587
+ num_reqs = len(req_ids)
588
+ block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
589
+ block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
590
+
591
+ _run_gpu_postprocess(
592
+ gpu_ctx,
593
+ kv_cache_config=kv_cache_config,
594
+ forward_context=forward_context,
595
+ copy_funcs=_COPY_FUNCS,
596
+ block_table=block_table_gpu,
597
+ req_ids=req_ids,
598
+ num_accepted_tokens=num_accepted_tokens,
599
+ mamba_state_idx=mamba_state_idx,
600
+ num_scheduled_tokens=num_scheduled_tokens,
601
+ num_computed_tokens=num_computed_tokens,
602
+ num_draft_tokens=num_draft_tokens,
603
+ device=device,
604
+ )
605
+
606
+ # State should be unchanged
607
+ torch.testing.assert_close(conv_state, conv_state_orig)
608
+ torch.testing.assert_close(temporal_state, temporal_state_orig)
609
+
610
+ @pytest.mark.parametrize("num_reqs", [1, 2, 8, 16])
611
+ def test_various_batch_sizes(self, device, test_config, num_reqs):
612
+ """Verify kernel works correctly with different batch sizes."""
613
+ cfg = _TestConfig(max_num_reqs=max(16, num_reqs))
614
+ torch.manual_seed(456)
615
+
616
+ req_ids = [f"req_{i}" for i in range(num_reqs)]
617
+ # All requests will trigger a copy
618
+ num_computed_tokens = [60] * num_reqs
619
+ num_scheduled_tokens = {r: 5 for r in req_ids}
620
+ num_draft_tokens = {r: 0 for r in req_ids}
621
+ num_accepted_tokens = [3] * num_reqs
622
+ mamba_state_idx = [3] * num_reqs
623
+ # Each request gets unique blocks
624
+ block_ids_per_req = [list(range(i * 8, (i + 1) * 8)) for i in range(num_reqs)]
625
+
626
+ # Ensure we have enough blocks
627
+ total_blocks = num_reqs * 8
628
+ cfg = _TestConfig(num_blocks=total_blocks, max_num_reqs=max(16, num_reqs))
629
+
630
+ layer_names = ["layer_0"]
631
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
632
+
633
+ (
634
+ conv_state_py,
635
+ temporal_state_py,
636
+ conv_state_gpu,
637
+ temporal_state_gpu,
638
+ forward_context_py,
639
+ forward_context_gpu,
640
+ ) = _make_dual_layer_state(cfg, device)
641
+
642
+ # Run Python path
643
+ scheduler_output = _make_postprocess_scheduler_output(
644
+ req_ids,
645
+ num_scheduled_tokens,
646
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
647
+ )
648
+ input_batch_py = _make_input_batch(
649
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
650
+ )
651
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
652
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
653
+
654
+ postprocess_mamba(
655
+ scheduler_output,
656
+ kv_cache_config,
657
+ input_batch_py,
658
+ requests,
659
+ forward_context_py,
660
+ _COPY_FUNCS,
661
+ copy_bufs,
662
+ )
663
+ torch.accelerator.synchronize()
664
+
665
+ # Run GPU path
666
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
667
+
668
+ max_blocks_per_req = 8
669
+ block_table_gpu = torch.zeros(
670
+ num_reqs, max_blocks_per_req, dtype=torch.int32, device=device
671
+ )
672
+ for i, block_ids in enumerate(block_ids_per_req):
673
+ block_table_gpu[i, : len(block_ids)] = torch.tensor(
674
+ block_ids, dtype=torch.int32
675
+ )
676
+
677
+ _run_gpu_postprocess(
678
+ gpu_ctx,
679
+ kv_cache_config=kv_cache_config,
680
+ forward_context=forward_context_gpu,
681
+ copy_funcs=_COPY_FUNCS,
682
+ block_table=block_table_gpu,
683
+ req_ids=req_ids,
684
+ num_accepted_tokens=num_accepted_tokens,
685
+ mamba_state_idx=mamba_state_idx,
686
+ num_scheduled_tokens=num_scheduled_tokens,
687
+ num_computed_tokens=num_computed_tokens,
688
+ num_draft_tokens=num_draft_tokens,
689
+ device=device,
690
+ )
691
+
692
+ # Compare results
693
+ torch.testing.assert_close(
694
+ conv_state_gpu, conv_state_py, msg="Conv state mismatch"
695
+ )
696
+ torch.testing.assert_close(
697
+ temporal_state_gpu, temporal_state_py, msg="Temporal state mismatch"
698
+ )
699
+
700
+ def test_block_table_with_realistic_stride(self, device, test_config):
701
+ """
702
+ Test kernel with realistic block table strides.
703
+
704
+ In real usage, the block table is pre-allocated with shape
705
+ [max_num_reqs, max_num_blocks_per_req] and then sliced to
706
+ [:num_reqs]. This means stride(0) = max_num_blocks_per_req,
707
+ which is typically much larger than the actual blocks used.
708
+
709
+ This test verifies the kernel handles non-tight strides correctly,
710
+ catching bugs where stride is incorrectly treated as bytes vs elements.
711
+ """
712
+ cfg = test_config
713
+ torch.manual_seed(789)
714
+
715
+ # Use multiple requests to exercise stride-based indexing
716
+ num_reqs = 4
717
+ req_ids = [f"req_{i}" for i in range(num_reqs)]
718
+
719
+ # All requests trigger copies (same setup as test_various_batch_sizes)
720
+ num_computed_tokens = [60] * num_reqs
721
+ num_scheduled_tokens = {r: 5 for r in req_ids}
722
+ num_draft_tokens: dict[str, int] = {}
723
+ num_accepted_tokens = [3] * num_reqs
724
+ mamba_state_idx = [3] * num_reqs
725
+
726
+ # Each request uses only 8 blocks, but we allocate much more
727
+ blocks_used_per_req = 8
728
+ block_ids_per_req = [
729
+ list(range(i * blocks_used_per_req, (i + 1) * blocks_used_per_req))
730
+ for i in range(num_reqs)
731
+ ]
732
+
733
+ total_blocks = num_reqs * blocks_used_per_req
734
+ cfg = _TestConfig(num_blocks=total_blocks, max_num_reqs=max(16, num_reqs))
735
+
736
+ layer_names = ["layer_0"]
737
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
738
+
739
+ (
740
+ conv_state_py,
741
+ temporal_state_py,
742
+ conv_state_gpu,
743
+ temporal_state_gpu,
744
+ forward_context_py,
745
+ forward_context_gpu,
746
+ ) = _make_dual_layer_state(cfg, device)
747
+
748
+ # Run Python path
749
+ scheduler_output = _make_postprocess_scheduler_output(
750
+ req_ids,
751
+ num_scheduled_tokens,
752
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
753
+ )
754
+ input_batch_py = _make_input_batch(
755
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
756
+ )
757
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
758
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
759
+
760
+ postprocess_mamba(
761
+ scheduler_output,
762
+ kv_cache_config,
763
+ input_batch_py,
764
+ requests,
765
+ forward_context_py,
766
+ _COPY_FUNCS,
767
+ copy_bufs,
768
+ )
769
+ torch.accelerator.synchronize()
770
+
771
+ # Run GPU path with REALISTIC block table stride
772
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
773
+
774
+ # KEY DIFFERENCE: Create a large block table like real code does
775
+ # Real system has max_num_blocks_per_req >> blocks actually used
776
+ max_num_reqs_full = 16
777
+ max_blocks_per_req_full = 512 # Much larger than blocks_used_per_req=8
778
+
779
+ # Allocate full-size table (simulates pre-allocated CpuGpuBuffer)
780
+ block_table_full = torch.zeros(
781
+ max_num_reqs_full, max_blocks_per_req_full, dtype=torch.int32, device=device
782
+ )
783
+
784
+ # Fill in actual block IDs (only first few columns used)
785
+ for i, block_ids in enumerate(block_ids_per_req):
786
+ block_table_full[i, : len(block_ids)] = torch.tensor(
787
+ block_ids, dtype=torch.int32
788
+ )
789
+
790
+ # Slice like real code: block_table.gpu[:num_reqs]
791
+ # This preserves stride(0) = 512, not 8!
792
+ block_table_gpu = block_table_full[:num_reqs]
793
+
794
+ # Verify stride is large (the key property we're testing)
795
+ assert block_table_gpu.stride(0) == max_blocks_per_req_full, (
796
+ f"Expected stride {max_blocks_per_req_full}, "
797
+ f"got {block_table_gpu.stride(0)}"
798
+ )
799
+
800
+ _run_gpu_postprocess(
801
+ gpu_ctx,
802
+ kv_cache_config=kv_cache_config,
803
+ forward_context=forward_context_gpu,
804
+ copy_funcs=_COPY_FUNCS,
805
+ block_table=block_table_gpu,
806
+ req_ids=req_ids,
807
+ num_accepted_tokens=num_accepted_tokens,
808
+ mamba_state_idx=mamba_state_idx,
809
+ num_scheduled_tokens=num_scheduled_tokens,
810
+ num_computed_tokens=num_computed_tokens,
811
+ num_draft_tokens=num_draft_tokens,
812
+ device=device,
813
+ )
814
+
815
+ # Compare results - this will fail if stride handling is incorrect
816
+ torch.testing.assert_close(
817
+ conv_state_gpu,
818
+ conv_state_py,
819
+ msg="Conv state mismatch - possible stride bug in kernel",
820
+ )
821
+ torch.testing.assert_close(
822
+ temporal_state_gpu,
823
+ temporal_state_py,
824
+ msg="Temporal state mismatch - possible stride bug in kernel",
825
+ )
826
+
827
+ def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1(
828
+ self, device, test_config
829
+ ):
830
+ """
831
+ Test the ``src_addr == dst_addr`` early-return path in
832
+ postprocess_mamba_fused_kernel matches Python behavior.
833
+
834
+ When src_addr == dst_addr (source and destination memory addresses are
835
+ identical), both implementations should:
836
+ 1. Skip the copy (state unchanged)
837
+ 2. Set num_accepted_tokens to 1
838
+
839
+ This condition occurs when:
840
+ - src_block_idx == dest_block_idx (same logical block)
841
+ - accept_token_bias == 0 (no offset within the block)
842
+
843
+ Python reference (collect_mamba_copy_meta):
844
+ if src_block_idx == dest_block_idx and accept_token_bias == 0:
845
+ return # No copy added
846
+
847
+ Python reference (postprocess_mamba):
848
+ if src_block_idx == dest_block_idx:
849
+ num_accepted_tokens_cpu[i] = 1
850
+
851
+ Test setup (block_size=16):
852
+ - num_tokens_running_state = 30 + 2 - 0 = 32
853
+ - new_num_computed = 32 + 1 - 1 = 32
854
+ - aligned_new_computed = 32
855
+ - accept_token_bias = 32 - 32 = 0
856
+ - dest_block_idx = 32 // 16 - 1 = 1
857
+ - src_block_idx = 1 (set explicitly)
858
+ """
859
+ cfg = test_config
860
+ torch.manual_seed(1001)
861
+
862
+ req_ids = ["req_0"]
863
+ num_computed_tokens = [30]
864
+ num_scheduled_tokens = {"req_0": 2}
865
+ num_draft_tokens: dict[str, int] = {}
866
+ num_accepted_tokens = [5] # Initial value, should become 1
867
+ mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx
868
+ block_ids_per_req = [list(range(8))]
869
+
870
+ layer_names = ["layer_0"]
871
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
872
+
873
+ (
874
+ conv_state_py,
875
+ temporal_state_py,
876
+ conv_state_gpu,
877
+ temporal_state_gpu,
878
+ forward_context_py,
879
+ forward_context_gpu,
880
+ ) = _make_dual_layer_state(cfg, device)
881
+
882
+ # Also clone to verify no modification
883
+ conv_state_orig = conv_state_py.clone()
884
+ temporal_state_orig = temporal_state_py.clone()
885
+
886
+ # --- Run Python path ---
887
+ scheduler_output = _make_postprocess_scheduler_output(
888
+ req_ids,
889
+ num_scheduled_tokens,
890
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
891
+ )
892
+ input_batch_py = _make_input_batch(
893
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
894
+ )
895
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
896
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
897
+
898
+ postprocess_mamba(
899
+ scheduler_output,
900
+ kv_cache_config,
901
+ input_batch_py,
902
+ requests,
903
+ forward_context_py,
904
+ _COPY_FUNCS,
905
+ copy_bufs,
906
+ )
907
+ torch.accelerator.synchronize()
908
+
909
+ # --- Run GPU path ---
910
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
911
+
912
+ num_reqs = len(req_ids)
913
+ block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
914
+ block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
915
+
916
+ _run_gpu_postprocess(
917
+ gpu_ctx,
918
+ kv_cache_config=kv_cache_config,
919
+ forward_context=forward_context_gpu,
920
+ copy_funcs=_COPY_FUNCS,
921
+ block_table=block_table_gpu,
922
+ req_ids=req_ids,
923
+ num_accepted_tokens=num_accepted_tokens,
924
+ mamba_state_idx=mamba_state_idx,
925
+ num_scheduled_tokens=num_scheduled_tokens,
926
+ num_computed_tokens=num_computed_tokens,
927
+ num_draft_tokens=num_draft_tokens,
928
+ device=device,
929
+ )
930
+
931
+ # --- Verify Python behavior (ground truth) ---
932
+ # State should be unchanged (no copy when src_addr == dst_addr)
933
+ torch.testing.assert_close(
934
+ conv_state_py,
935
+ conv_state_orig,
936
+ msg="Python: Conv state should be unchanged when src==dst",
937
+ )
938
+ torch.testing.assert_close(
939
+ temporal_state_py,
940
+ temporal_state_orig,
941
+ msg="Python: Temporal state should be unchanged when src==dst",
942
+ )
943
+ # num_accepted_tokens should be 1
944
+ assert input_batch_py.num_accepted_tokens_cpu[0] == 1, (
945
+ f"Python: num_accepted_tokens should be 1, "
946
+ f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
947
+ )
948
+
949
+ # --- Verify GPU matches Python ---
950
+ torch.testing.assert_close(
951
+ conv_state_gpu,
952
+ conv_state_py,
953
+ msg="GPU conv state should match Python",
954
+ )
955
+ torch.testing.assert_close(
956
+ temporal_state_gpu,
957
+ temporal_state_py,
958
+ msg="GPU temporal state should match Python",
959
+ )
960
+ expected_accepted = torch.tensor(
961
+ input_batch_py.num_accepted_tokens_cpu[:num_reqs],
962
+ dtype=torch.int32,
963
+ device=device,
964
+ )
965
+ torch.testing.assert_close(
966
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
967
+ expected_accepted,
968
+ msg="GPU num_accepted_tokens should match Python",
969
+ )
970
+
971
+ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1(
972
+ self, device, test_config
973
+ ):
974
+ """
975
+ Test the ``src_block_idx == dest_block_idx`` post-copy update in
976
+ postprocess_mamba_fused_kernel matches Python behavior.
977
+
978
+ When src_block_idx == dest_block_idx but accept_token_bias > 0, both
979
+ implementations should:
980
+ 1. Perform the copy (src_addr != dst_addr due to offset)
981
+ 2. Set num_accepted_tokens to 1 AFTER the copy
982
+
983
+ Python reference (postprocess_mamba):
984
+ if src_block_idx == dest_block_idx:
985
+ num_accepted_tokens_cpu[i] = 1
986
+
987
+ For conv states: copies state[block, offset:] to
988
+ state[block, :] (shifted window)
989
+ For temporal states: copies state[block_ids[src_idx + offset]] to
990
+ state[block_ids[dest_idx]]
991
+
992
+ Test setup (block_size=16):
993
+ - num_tokens_running_state = 30 + 1 - 0 = 31
994
+ - new_num_computed = 31 + 2 - 1 = 32
995
+ - aligned_new_computed = 32
996
+ - accept_token_bias = 32 - 31 = 1 (> 0, so copy happens)
997
+ - dest_block_idx = 32 // 16 - 1 = 1
998
+ - src_block_idx = 1 (set explicitly, == dest_block_idx)
999
+ """
1000
+ cfg = test_config
1001
+ torch.manual_seed(1002)
1002
+
1003
+ req_ids = ["req_0"]
1004
+ num_computed_tokens = [30]
1005
+ num_scheduled_tokens = {"req_0": 1}
1006
+ num_draft_tokens: dict[str, int] = {}
1007
+ num_accepted_tokens = [2] # Results in accept_token_bias = 1
1008
+ mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx
1009
+ block_ids_per_req = [list(range(8))]
1010
+
1011
+ layer_names = ["layer_0"]
1012
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
1013
+
1014
+ (
1015
+ conv_state_py,
1016
+ temporal_state_py,
1017
+ conv_state_gpu,
1018
+ temporal_state_gpu,
1019
+ forward_context_py,
1020
+ forward_context_gpu,
1021
+ ) = _make_dual_layer_state(cfg, device)
1022
+
1023
+ # Clone to verify modification
1024
+ conv_state_orig = conv_state_py.clone()
1025
+ temporal_state_orig = temporal_state_py.clone()
1026
+
1027
+ # --- Run Python path ---
1028
+ scheduler_output = _make_postprocess_scheduler_output(
1029
+ req_ids,
1030
+ num_scheduled_tokens,
1031
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
1032
+ )
1033
+ input_batch_py = _make_input_batch(
1034
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
1035
+ )
1036
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
1037
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
1038
+
1039
+ postprocess_mamba(
1040
+ scheduler_output,
1041
+ kv_cache_config,
1042
+ input_batch_py,
1043
+ requests,
1044
+ forward_context_py,
1045
+ _COPY_FUNCS,
1046
+ copy_bufs,
1047
+ )
1048
+ torch.accelerator.synchronize()
1049
+
1050
+ # --- Run GPU path ---
1051
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
1052
+
1053
+ num_reqs = len(req_ids)
1054
+ block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
1055
+ block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
1056
+
1057
+ _run_gpu_postprocess(
1058
+ gpu_ctx,
1059
+ kv_cache_config=kv_cache_config,
1060
+ forward_context=forward_context_gpu,
1061
+ copy_funcs=_COPY_FUNCS,
1062
+ block_table=block_table_gpu,
1063
+ req_ids=req_ids,
1064
+ num_accepted_tokens=num_accepted_tokens,
1065
+ mamba_state_idx=mamba_state_idx,
1066
+ num_scheduled_tokens=num_scheduled_tokens,
1067
+ num_computed_tokens=num_computed_tokens,
1068
+ num_draft_tokens=num_draft_tokens,
1069
+ device=device,
1070
+ )
1071
+
1072
+ # --- Verify Python behavior (ground truth) ---
1073
+ dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1
1074
+
1075
+ # Conv state should be modified (shifted copy within block)
1076
+ conv_changed = not torch.allclose(
1077
+ conv_state_py[dest_block_id], conv_state_orig[dest_block_id]
1078
+ )
1079
+ assert conv_changed, (
1080
+ "Python: Conv state should be modified when accept_token_bias > 0"
1081
+ )
1082
+
1083
+ # Temporal state should be modified (copy from different block)
1084
+ src_block_id_temporal = block_ids_per_req[0][2] # actual_src_block_idx = 2
1085
+ dest_block_id_temporal = block_ids_per_req[0][1] # dest_block_idx = 1
1086
+ torch.testing.assert_close(
1087
+ temporal_state_py[dest_block_id_temporal],
1088
+ temporal_state_orig[src_block_id_temporal],
1089
+ msg="Python: Temporal state copy should have happened",
1090
+ )
1091
+
1092
+ # num_accepted_tokens should be 1
1093
+ assert input_batch_py.num_accepted_tokens_cpu[0] == 1, (
1094
+ f"Python: num_accepted_tokens should be 1, "
1095
+ f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
1096
+ )
1097
+
1098
+ # --- Verify GPU matches Python ---
1099
+ torch.testing.assert_close(
1100
+ conv_state_gpu,
1101
+ conv_state_py,
1102
+ msg="GPU conv state should match Python",
1103
+ )
1104
+ torch.testing.assert_close(
1105
+ temporal_state_gpu,
1106
+ temporal_state_py,
1107
+ msg="GPU temporal state should match Python",
1108
+ )
1109
+ expected_accepted = torch.tensor(
1110
+ input_batch_py.num_accepted_tokens_cpu[:num_reqs],
1111
+ dtype=torch.int32,
1112
+ device=device,
1113
+ )
1114
+ torch.testing.assert_close(
1115
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
1116
+ expected_accepted,
1117
+ msg="GPU num_accepted_tokens should match Python",
1118
+ )
1119
+
1120
+ def test_different_block_idx_copies_without_setting_accepted_to_1(
1121
+ self, device, test_config
1122
+ ):
1123
+ """
1124
+ Test that neither special-case path triggers when
1125
+ src_block_idx != dest_block_idx, and GPU matches Python behavior.
1126
+
1127
+ When copying between different blocks:
1128
+ 1. src_addr != dst_addr (different blocks = different addresses)
1129
+ 2. src_block_idx != dest_block_idx
1130
+
1131
+ Therefore:
1132
+ - The ``src_addr == dst_addr`` early-return does NOT trigger
1133
+ - The ``src_block_idx == dest_block_idx`` post-copy update does NOT trigger
1134
+ - Copy happens normally
1135
+ - num_accepted_tokens remains UNCHANGED
1136
+
1137
+ Test setup (block_size=16):
1138
+ - num_tokens_running_state = 60 + 3 - 0 = 63
1139
+ - new_num_computed = 63 + 3 - 1 = 65
1140
+ - aligned_new_computed = 64
1141
+ - accept_token_bias = 64 - 63 = 1
1142
+ - dest_block_idx = 64 // 16 - 1 = 3
1143
+ - src_block_idx = 2 (set explicitly, != dest_block_idx)
1144
+ """
1145
+ cfg = test_config
1146
+ torch.manual_seed(1003)
1147
+
1148
+ req_ids = ["req_0"]
1149
+ num_computed_tokens = [60]
1150
+ num_scheduled_tokens = {"req_0": 3}
1151
+ num_draft_tokens: dict[str, int] = {}
1152
+ num_accepted_tokens = [3] # Should remain 3, NOT set to 1
1153
+ mamba_state_idx = [2] # src_block_idx = 2, dest_block_idx will be 3
1154
+ block_ids_per_req = [list(range(8))]
1155
+
1156
+ layer_names = ["layer_0"]
1157
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
1158
+
1159
+ (
1160
+ conv_state_py,
1161
+ temporal_state_py,
1162
+ conv_state_gpu,
1163
+ temporal_state_gpu,
1164
+ forward_context_py,
1165
+ forward_context_gpu,
1166
+ ) = _make_dual_layer_state(cfg, device)
1167
+
1168
+ # Clone to verify modification
1169
+ conv_state_orig = conv_state_py.clone()
1170
+
1171
+ # --- Run Python path ---
1172
+ scheduler_output = _make_postprocess_scheduler_output(
1173
+ req_ids,
1174
+ num_scheduled_tokens,
1175
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
1176
+ )
1177
+ input_batch_py = _make_input_batch(
1178
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
1179
+ )
1180
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
1181
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
1182
+
1183
+ postprocess_mamba(
1184
+ scheduler_output,
1185
+ kv_cache_config,
1186
+ input_batch_py,
1187
+ requests,
1188
+ forward_context_py,
1189
+ _COPY_FUNCS,
1190
+ copy_bufs,
1191
+ )
1192
+ torch.accelerator.synchronize()
1193
+
1194
+ # --- Run GPU path ---
1195
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
1196
+
1197
+ num_reqs = len(req_ids)
1198
+ block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
1199
+ block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
1200
+
1201
+ _run_gpu_postprocess(
1202
+ gpu_ctx,
1203
+ kv_cache_config=kv_cache_config,
1204
+ forward_context=forward_context_gpu,
1205
+ copy_funcs=_COPY_FUNCS,
1206
+ block_table=block_table_gpu,
1207
+ req_ids=req_ids,
1208
+ num_accepted_tokens=num_accepted_tokens,
1209
+ mamba_state_idx=mamba_state_idx,
1210
+ num_scheduled_tokens=num_scheduled_tokens,
1211
+ num_computed_tokens=num_computed_tokens,
1212
+ num_draft_tokens=num_draft_tokens,
1213
+ device=device,
1214
+ )
1215
+
1216
+ # --- Verify Python behavior (ground truth) ---
1217
+ dest_block_id = block_ids_per_req[0][3] # dest_block_idx = 3
1218
+
1219
+ # Copy DID happen (dest block should be modified)
1220
+ conv_changed = not torch.allclose(
1221
+ conv_state_py[dest_block_id], conv_state_orig[dest_block_id]
1222
+ )
1223
+ assert conv_changed, "Python: Conv state copy should have happened"
1224
+
1225
+ # num_accepted_tokens should NOT be changed to 1
1226
+ assert input_batch_py.num_accepted_tokens_cpu[0] == num_accepted_tokens[0], (
1227
+ f"Python: num_accepted_tokens should remain {num_accepted_tokens[0]}, "
1228
+ f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
1229
+ )
1230
+
1231
+ # --- Verify GPU matches Python ---
1232
+ torch.testing.assert_close(
1233
+ conv_state_gpu,
1234
+ conv_state_py,
1235
+ msg="GPU conv state should match Python",
1236
+ )
1237
+ torch.testing.assert_close(
1238
+ temporal_state_gpu,
1239
+ temporal_state_py,
1240
+ msg="GPU temporal state should match Python",
1241
+ )
1242
+ expected_accepted = torch.tensor(
1243
+ input_batch_py.num_accepted_tokens_cpu[:num_reqs],
1244
+ dtype=torch.int32,
1245
+ device=device,
1246
+ )
1247
+ torch.testing.assert_close(
1248
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
1249
+ expected_accepted,
1250
+ msg="GPU num_accepted_tokens should match Python",
1251
+ )
1252
+
1253
+ def test_prefix_caching_shared_block_does_not_set_accepted_to_1(
1254
+ self, device, test_config
1255
+ ):
1256
+ """
1257
+ Regression test: with prefix caching, different logical block indices
1258
+ can map to the same physical block. The kernel must NOT set
1259
+ num_accepted_tokens to 1 in that case.
1260
+
1261
+ When src_block_idx != dest_block_idx but block_table maps both to the
1262
+ same physical block ID, src_addr == dst_addr. The copy is correctly
1263
+ skipped (self-copy is a no-op), but num_accepted_tokens must be
1264
+ preserved — only logical-index equality justifies setting it to 1.
1265
+
1266
+ Test setup (block_size=16):
1267
+ - num_tokens_running_state = 30 + 2 - 0 = 32
1268
+ - new_num_computed = 32 + 3 - 1 = 34
1269
+ - aligned_new_computed = 32
1270
+ - accept_token_bias = 32 - 32 = 0
1271
+ - dest_block_idx = 32 // 16 - 1 = 1
1272
+ - src_block_idx = 0 (set explicitly, != dest_block_idx)
1273
+ - block_ids = [5, 5, ...] — prefix caching: both logical indices
1274
+ map to the same physical block 5
1275
+ """
1276
+ cfg = test_config
1277
+ torch.manual_seed(2001)
1278
+
1279
+ req_ids = ["req_0"]
1280
+ num_computed_tokens = [30]
1281
+ num_scheduled_tokens = {"req_0": 2}
1282
+ num_draft_tokens: dict[str, int] = {}
1283
+ num_accepted_tokens = [3] # Must stay 3, NOT become 1
1284
+ mamba_state_idx = [0] # src_block_idx = 0, dest_block_idx will be 1
1285
+
1286
+ # Prefix caching: logical blocks 0 and 1 share physical block 5
1287
+ block_ids_per_req = [[5, 5, 2, 3, 4, 6, 7, 8]]
1288
+
1289
+ layer_names = ["layer_0"]
1290
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
1291
+
1292
+ (
1293
+ conv_state_py,
1294
+ temporal_state_py,
1295
+ conv_state_gpu,
1296
+ temporal_state_gpu,
1297
+ forward_context_py,
1298
+ forward_context_gpu,
1299
+ ) = _make_dual_layer_state(cfg, device)
1300
+
1301
+ conv_state_orig = conv_state_py.clone()
1302
+ temporal_state_orig = temporal_state_py.clone()
1303
+
1304
+ # --- Run Python path ---
1305
+ scheduler_output = _make_postprocess_scheduler_output(
1306
+ req_ids,
1307
+ num_scheduled_tokens,
1308
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
1309
+ )
1310
+ input_batch_py = _make_input_batch(
1311
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
1312
+ )
1313
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
1314
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
1315
+
1316
+ postprocess_mamba(
1317
+ scheduler_output,
1318
+ kv_cache_config,
1319
+ input_batch_py,
1320
+ requests,
1321
+ forward_context_py,
1322
+ _COPY_FUNCS,
1323
+ copy_bufs,
1324
+ )
1325
+ torch.accelerator.synchronize()
1326
+
1327
+ # --- Run GPU path ---
1328
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
1329
+ num_reqs = len(req_ids)
1330
+ block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
1331
+ block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
1332
+
1333
+ _run_gpu_postprocess(
1334
+ gpu_ctx,
1335
+ kv_cache_config=kv_cache_config,
1336
+ forward_context=forward_context_gpu,
1337
+ copy_funcs=_COPY_FUNCS,
1338
+ block_table=block_table_gpu,
1339
+ req_ids=req_ids,
1340
+ num_accepted_tokens=num_accepted_tokens,
1341
+ mamba_state_idx=mamba_state_idx,
1342
+ num_scheduled_tokens=num_scheduled_tokens,
1343
+ num_computed_tokens=num_computed_tokens,
1344
+ num_draft_tokens=num_draft_tokens,
1345
+ device=device,
1346
+ )
1347
+
1348
+ # --- Verify Python behavior (ground truth) ---
1349
+ # Copy is self-to-self (same physical block), state unchanged
1350
+ torch.testing.assert_close(
1351
+ conv_state_py,
1352
+ conv_state_orig,
1353
+ msg="Python: Conv state should be unchanged (self-copy)",
1354
+ )
1355
+ torch.testing.assert_close(
1356
+ temporal_state_py,
1357
+ temporal_state_orig,
1358
+ msg="Python: Temporal state should be unchanged (self-copy)",
1359
+ )
1360
+ # num_accepted_tokens must NOT be set to 1 (src_block_idx != dest_block_idx)
1361
+ assert input_batch_py.num_accepted_tokens_cpu[0] == num_accepted_tokens[0], (
1362
+ f"Python: num_accepted_tokens should remain {num_accepted_tokens[0]}, "
1363
+ f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
1364
+ )
1365
+
1366
+ # --- Verify GPU matches Python ---
1367
+ torch.testing.assert_close(
1368
+ conv_state_gpu,
1369
+ conv_state_py,
1370
+ msg="GPU conv state should match Python",
1371
+ )
1372
+ torch.testing.assert_close(
1373
+ temporal_state_gpu,
1374
+ temporal_state_py,
1375
+ msg="GPU temporal state should match Python",
1376
+ )
1377
+ expected_accepted = torch.tensor(
1378
+ input_batch_py.num_accepted_tokens_cpu[:num_reqs],
1379
+ dtype=torch.int32,
1380
+ device=device,
1381
+ )
1382
+ torch.testing.assert_close(
1383
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
1384
+ expected_accepted,
1385
+ msg="GPU num_accepted_tokens should match Python (must NOT be 1)",
1386
+ )
1387
+
1388
+ def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_config):
1389
+ """
1390
+ Regression test: non-sequential physical block IDs under prefix caching
1391
+ with the needs_copy boundary at exact equality.
1392
+
1393
+ Under PC, the block allocator assigns physical block IDs in arbitrary
1394
+ order (e.g., [17, 3, 42, 9] instead of [0, 1, 2, 3]). The needs_copy
1395
+ condition is purely token-count based and must evaluate identically
1396
+ regardless of the physical block IDs assigned. This test verifies that
1397
+ the kernel's address arithmetic (block_table lookup, stride computation)
1398
+ produces correct copies when physical IDs are non-sequential.
1399
+
1400
+ Two requests exercise different boundary behaviors:
1401
+ - req_0: aligned_new_computed == num_tokens_running_state (exact boundary)
1402
+ This is the tightest edge: one fewer accepted token and no copy needed.
1403
+ - req_1: aligned_new_computed == num_tokens_running_state (exact boundary)
1404
+ Different block layout, src!=dest, real copy happens.
1405
+
1406
+ Both use non-sequential block IDs typical of PC reuse patterns.
1407
+
1408
+ Test setup (block_size=16):
1409
+ req_0:
1410
+ - num_tokens_running_state = 48 + 0 - 0 = 48
1411
+ - new_num_computed = 48 + 1 - 1 = 48
1412
+ - aligned_new_computed = 48
1413
+ - needs_copy = (48 >= 48) = True (exact boundary!)
1414
+ - accept_token_bias = 48 - 48 = 0
1415
+ - dest_block_idx = 48 // 16 - 1 = 2
1416
+ - src_block_idx = 2 (same as dest -> num_accepted = 1)
1417
+
1418
+ req_1:
1419
+ - num_tokens_running_state = 31 + 1 - 0 = 32
1420
+ - new_num_computed = 32 + 3 - 1 = 34
1421
+ - aligned_new_computed = 32
1422
+ - needs_copy = (32 >= 32) = True (exact boundary!)
1423
+ - accept_token_bias = 32 - 32 = 0
1424
+ - dest_block_idx = 32 // 16 - 1 = 1
1425
+ - src_block_idx = 0 (diff from dest -> num_accepted unchanged)
1426
+ """
1427
+ cfg = test_config
1428
+ torch.manual_seed(4001)
1429
+
1430
+ req_ids = ["req_0", "req_1"]
1431
+ num_computed_tokens = [48, 31]
1432
+ num_scheduled_tokens = {"req_0": 0, "req_1": 1}
1433
+ num_draft_tokens: dict[str, int] = {}
1434
+ num_accepted_tokens = [1, 3]
1435
+ mamba_state_idx = [2, 0]
1436
+
1437
+ # Non-sequential block IDs typical of prefix caching allocation
1438
+ block_ids_per_req = [
1439
+ [17, 3, 42, 9, 25, 11, 30, 2], # req_0: scattered physical blocks
1440
+ [41, 7, 22, 15, 38, 19, 4, 28], # req_1: different scattered blocks
1441
+ ]
1442
+
1443
+ layer_names = [f"layer_{i}" for i in range(cfg.num_layers)]
1444
+ # Need enough physical blocks for the scattered IDs
1445
+ num_blocks = 50
1446
+ local_cfg = _TestConfig(num_blocks=num_blocks, max_num_reqs=cfg.max_num_reqs)
1447
+ kv_cache_config = _make_kv_cache_config(local_cfg, layer_names)
1448
+
1449
+ (
1450
+ conv_states_py,
1451
+ temporal_states_py,
1452
+ conv_states_gpu,
1453
+ temporal_states_gpu,
1454
+ forward_context_py,
1455
+ forward_context_gpu,
1456
+ ) = _make_dual_states(local_cfg, layer_names, device)
1457
+
1458
+ # --- Run Python path ---
1459
+ scheduler_output = _make_postprocess_scheduler_output(
1460
+ req_ids,
1461
+ num_scheduled_tokens,
1462
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
1463
+ )
1464
+ input_batch_py = _make_input_batch(
1465
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
1466
+ )
1467
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
1468
+ copy_bufs = _make_copy_bufs(local_cfg, kv_cache_config, device)
1469
+
1470
+ postprocess_mamba(
1471
+ scheduler_output,
1472
+ kv_cache_config,
1473
+ input_batch_py,
1474
+ requests,
1475
+ forward_context_py,
1476
+ _COPY_FUNCS,
1477
+ copy_bufs,
1478
+ )
1479
+ torch.accelerator.synchronize()
1480
+
1481
+ # --- Run GPU path ---
1482
+ gpu_ctx = _make_gpu_ctx(local_cfg, kv_cache_config, device)
1483
+ num_reqs = len(req_ids)
1484
+ max_blocks = max(len(b) for b in block_ids_per_req)
1485
+ block_table_gpu = torch.zeros(
1486
+ num_reqs, max_blocks, dtype=torch.int32, device=device
1487
+ )
1488
+ for i, block_ids in enumerate(block_ids_per_req):
1489
+ block_table_gpu[i, : len(block_ids)] = torch.tensor(
1490
+ block_ids, dtype=torch.int32
1491
+ )
1492
+
1493
+ _run_gpu_postprocess(
1494
+ gpu_ctx,
1495
+ kv_cache_config=kv_cache_config,
1496
+ forward_context=forward_context_gpu,
1497
+ copy_funcs=_COPY_FUNCS,
1498
+ block_table=block_table_gpu,
1499
+ req_ids=req_ids,
1500
+ num_accepted_tokens=num_accepted_tokens,
1501
+ mamba_state_idx=mamba_state_idx,
1502
+ num_scheduled_tokens=num_scheduled_tokens,
1503
+ num_computed_tokens=num_computed_tokens,
1504
+ num_draft_tokens=num_draft_tokens,
1505
+ device=device,
1506
+ )
1507
+
1508
+ # --- Compare results ---
1509
+ for i in range(cfg.num_layers):
1510
+ torch.testing.assert_close(
1511
+ conv_states_gpu[i],
1512
+ conv_states_py[i],
1513
+ msg=f"Conv state mismatch at layer {i} with non-sequential block IDs",
1514
+ )
1515
+ torch.testing.assert_close(
1516
+ temporal_states_gpu[i],
1517
+ temporal_states_py[i],
1518
+ msg=(
1519
+ f"Temporal state mismatch at layer {i} "
1520
+ f"with non-sequential block IDs"
1521
+ ),
1522
+ )
1523
+
1524
+ expected_accepted = torch.tensor(
1525
+ input_batch_py.num_accepted_tokens_cpu[:num_reqs],
1526
+ dtype=torch.int32,
1527
+ device=device,
1528
+ )
1529
+ torch.testing.assert_close(
1530
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
1531
+ expected_accepted,
1532
+ msg="num_accepted_tokens mismatch with non-sequential block IDs",
1533
+ )
1534
+
1535
+ # Verify req_0 had num_accepted set to 1 (src==dest) and req_1 unchanged
1536
+ assert input_batch_py.num_accepted_tokens_cpu[0] == 1
1537
+ assert input_batch_py.num_accepted_tokens_cpu[1] == num_accepted_tokens[1]
1538
+
1539
+ def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_config):
1540
+ """
1541
+ Regression test: mixed batch under prefix caching where some requests
1542
+ have shared physical blocks (aliased) and others have distinct blocks,
1543
+ with the needs_copy boundary at various positions.
1544
+
1545
+ This tests the interaction between:
1546
+ 1. PC block aliasing (src and dest map to same physical block)
1547
+ 2. The needs_copy boundary (exact equality vs well-past vs no-copy)
1548
+ 3. Non-sequential physical block IDs
1549
+
1550
+ Batch of 4 requests:
1551
+ - req_0: needs_copy=True, src!=dest, shared physical block (PC aliased)
1552
+ -> copy skipped (src_addr==dst_addr), num_accepted PRESERVED
1553
+ - req_1: needs_copy=True, src!=dest, distinct blocks, non-sequential IDs
1554
+ -> real copy happens, num_accepted PRESERVED
1555
+ - req_2: needs_copy=False (below boundary)
1556
+ -> no action at all
1557
+ - req_3: needs_copy=True, src==dest (exact boundary, zero bias)
1558
+ -> copy skipped (self-copy), num_accepted SET TO 1
1559
+
1560
+ Test setup (block_size=16):
1561
+ req_0: running=30+2-0=32, new=32+3-1=34, aligned=32, 32>=32 -> COPY
1562
+ bias=0, dest=32//16-1=1, src=0 (!=dest)
1563
+ block_ids=[5,5,...] -> same physical -> skip, keep accepted=3
1564
+
1565
+ req_1: running=60+5-2=63, new=63+3-1=65, aligned=64, 64>=63 -> COPY
1566
+ bias=1, dest=64//16-1=3, src=2 (!=dest)
1567
+ block_ids=[41,7,22,15,...] -> distinct -> real copy, keep accepted=3
1568
+
1569
+ req_2: running=30+3-0=33, new=33+1-1=33, aligned=32, 32<33 -> NO COPY
1570
+
1571
+ req_3: running=48+0-0=48, new=48+1-1=48, aligned=48, 48>=48 -> COPY
1572
+ bias=0, dest=48//16-1=2, src=2 (==dest)
1573
+ block_ids=[10,20,30,...] -> distinct IDs, same logical idx
1574
+ -> self-copy (src_addr==dst_addr), set accepted=1
1575
+ """
1576
+ cfg = test_config
1577
+ torch.manual_seed(5001)
1578
+
1579
+ req_ids = ["req_0", "req_1", "req_2", "req_3"]
1580
+ num_computed_tokens = [30, 60, 30, 48]
1581
+ num_scheduled_tokens = {"req_0": 2, "req_1": 5, "req_2": 3, "req_3": 0}
1582
+ num_draft_tokens = {"req_1": 2}
1583
+ num_accepted_tokens = [3, 3, 1, 1]
1584
+ mamba_state_idx = [0, 2, 1, 2]
1585
+
1586
+ # Block IDs with various PC patterns:
1587
+ # req_0: shared blocks (PC alias: logical 0 and 1 -> physical 5)
1588
+ # req_1: distinct non-sequential blocks
1589
+ # req_2: doesn't matter (no copy)
1590
+ # req_3: distinct sequential blocks (no aliasing)
1591
+ block_ids_per_req = [
1592
+ [5, 5, 12, 18, 23, 31, 44, 2], # req_0: blocks 0,1 share phys 5
1593
+ [41, 7, 22, 15, 38, 19, 4, 28], # req_1: all distinct
1594
+ [10, 20, 30, 40, 1, 6, 8, 14], # req_2: irrelevant
1595
+ [10, 20, 30, 40, 1, 6, 8, 14], # req_3: distinct, dest=src=idx 2
1596
+ ]
1597
+
1598
+ layer_names = [f"layer_{i}" for i in range(cfg.num_layers)]
1599
+ num_blocks = 50
1600
+ local_cfg = _TestConfig(num_blocks=num_blocks, max_num_reqs=cfg.max_num_reqs)
1601
+ kv_cache_config = _make_kv_cache_config(local_cfg, layer_names)
1602
+
1603
+ (
1604
+ conv_states_py,
1605
+ temporal_states_py,
1606
+ conv_states_gpu,
1607
+ temporal_states_gpu,
1608
+ forward_context_py,
1609
+ forward_context_gpu,
1610
+ ) = _make_dual_states(local_cfg, layer_names, device)
1611
+
1612
+ # --- Run Python path ---
1613
+ scheduler_output = _make_postprocess_scheduler_output(
1614
+ req_ids,
1615
+ num_scheduled_tokens,
1616
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
1617
+ )
1618
+ input_batch_py = _make_input_batch(
1619
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
1620
+ )
1621
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
1622
+ copy_bufs = _make_copy_bufs(local_cfg, kv_cache_config, device)
1623
+
1624
+ postprocess_mamba(
1625
+ scheduler_output,
1626
+ kv_cache_config,
1627
+ input_batch_py,
1628
+ requests,
1629
+ forward_context_py,
1630
+ _COPY_FUNCS,
1631
+ copy_bufs,
1632
+ )
1633
+ torch.accelerator.synchronize()
1634
+
1635
+ # --- Run GPU path ---
1636
+ gpu_ctx = _make_gpu_ctx(local_cfg, kv_cache_config, device)
1637
+ num_reqs = len(req_ids)
1638
+ max_blocks = max(len(b) for b in block_ids_per_req)
1639
+ block_table_gpu = torch.zeros(
1640
+ num_reqs, max_blocks, dtype=torch.int32, device=device
1641
+ )
1642
+ for i, block_ids in enumerate(block_ids_per_req):
1643
+ block_table_gpu[i, : len(block_ids)] = torch.tensor(
1644
+ block_ids, dtype=torch.int32
1645
+ )
1646
+
1647
+ _run_gpu_postprocess(
1648
+ gpu_ctx,
1649
+ kv_cache_config=kv_cache_config,
1650
+ forward_context=forward_context_gpu,
1651
+ copy_funcs=_COPY_FUNCS,
1652
+ block_table=block_table_gpu,
1653
+ req_ids=req_ids,
1654
+ num_accepted_tokens=num_accepted_tokens,
1655
+ mamba_state_idx=mamba_state_idx,
1656
+ num_scheduled_tokens=num_scheduled_tokens,
1657
+ num_computed_tokens=num_computed_tokens,
1658
+ num_draft_tokens=num_draft_tokens,
1659
+ device=device,
1660
+ )
1661
+
1662
+ # --- Compare all state tensors ---
1663
+ for i in range(cfg.num_layers):
1664
+ torch.testing.assert_close(
1665
+ conv_states_gpu[i],
1666
+ conv_states_py[i],
1667
+ msg=(
1668
+ f"Conv state mismatch at layer {i} — "
1669
+ f"mixed PC batch with shared/distinct blocks"
1670
+ ),
1671
+ )
1672
+ torch.testing.assert_close(
1673
+ temporal_states_gpu[i],
1674
+ temporal_states_py[i],
1675
+ msg=(
1676
+ f"Temporal state mismatch at layer {i} — "
1677
+ f"mixed PC batch with shared/distinct blocks"
1678
+ ),
1679
+ )
1680
+
1681
+ # --- Compare num_accepted_tokens ---
1682
+ expected_accepted = torch.tensor(
1683
+ input_batch_py.num_accepted_tokens_cpu[:num_reqs],
1684
+ dtype=torch.int32,
1685
+ device=device,
1686
+ )
1687
+ torch.testing.assert_close(
1688
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
1689
+ expected_accepted,
1690
+ msg="num_accepted_tokens mismatch in mixed PC batch",
1691
+ )
1692
+
1693
+ # Verify per-request expectations:
1694
+ # req_0: src!=dest, shared block -> preserved (3)
1695
+ assert input_batch_py.num_accepted_tokens_cpu[0] == 3
1696
+ # req_1: src!=dest, distinct blocks -> preserved (3)
1697
+ assert input_batch_py.num_accepted_tokens_cpu[1] == 3
1698
+ # req_2: no copy -> preserved (1)
1699
+ assert input_batch_py.num_accepted_tokens_cpu[2] == 1
1700
+ # req_3: src==dest -> set to 1
1701
+ assert input_batch_py.num_accepted_tokens_cpu[3] == 1
1702
+
1703
+ def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr(
1704
+ self, device, test_config
1705
+ ):
1706
+ """
1707
+ Regression test for 6466ce0d vs 959ca0fd: the kernel's early-return
1708
+ guard must compare logical block indices, not physical addresses.
1709
+
1710
+ Under prefix caching, different logical blocks (src_block_idx=0,
1711
+ dest_block_idx=1) can map to the same physical block. When
1712
+ accept_token_bias=0, this makes src_addr == dst_addr for BOTH conv
1713
+ and temporal states. A buggy guard `if src_addr == dst_addr` would
1714
+ incorrectly set num_accepted_tokens=1; the correct guard is
1715
+ `if src_block_idx == dest_block_idx and accept_token_bias == 0`.
1716
+
1717
+ The Python reference only sets num_accepted_tokens=1 when
1718
+ src_block_idx == dest_block_idx (line 79 of postprocess_mamba).
1719
+ With src_block_idx=0, dest_block_idx=1, num_accepted_tokens must
1720
+ be preserved even though the physical addresses match.
1721
+
1722
+ Test setup (block_size=16):
1723
+ - num_tokens_running_state = 30 + 2 - 0 = 32
1724
+ - new_num_computed = 32 + 3 - 1 = 34
1725
+ - aligned_new_computed = 32
1726
+ - needs_copy = (32 >= 32) = True
1727
+ - accept_token_bias = 32 - 32 = 0
1728
+ - dest_block_idx = 32 // 16 - 1 = 1
1729
+ - src_block_idx = 0 (explicitly, != dest_block_idx)
1730
+ - block_ids = [7, 7, ...] -> physical aliasing via prefix caching
1731
+
1732
+ Expected: num_accepted_tokens stays 3 (not set to 1).
1733
+ Bug (959ca0fd): kernel saw src_addr == dst_addr, set it to 1.
1734
+ """
1735
+ cfg = test_config
1736
+ torch.manual_seed(6001)
1737
+
1738
+ req_ids = ["req_0"]
1739
+ num_computed_tokens = [30]
1740
+ num_scheduled_tokens = {"req_0": 2}
1741
+ num_draft_tokens: dict[str, int] = {}
1742
+ num_accepted_tokens = [3]
1743
+ mamba_state_idx = [0] # src_block_idx = 0
1744
+
1745
+ # Prefix caching: logical blocks 0 and 1 both map to physical block 7.
1746
+ block_ids_per_req = [[7, 7, 10, 11, 12, 13, 14, 15]]
1747
+
1748
+ layer_names = ["layer_0"]
1749
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
1750
+
1751
+ (
1752
+ conv_state_py,
1753
+ temporal_state_py,
1754
+ conv_state_gpu,
1755
+ temporal_state_gpu,
1756
+ forward_context_py,
1757
+ forward_context_gpu,
1758
+ ) = _make_dual_layer_state(cfg, device)
1759
+
1760
+ # --- Run Python path ---
1761
+ scheduler_output = _make_postprocess_scheduler_output(
1762
+ req_ids,
1763
+ num_scheduled_tokens,
1764
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
1765
+ )
1766
+ input_batch_py = _make_input_batch(
1767
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
1768
+ )
1769
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
1770
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
1771
+
1772
+ postprocess_mamba(
1773
+ scheduler_output,
1774
+ kv_cache_config,
1775
+ input_batch_py,
1776
+ requests,
1777
+ forward_context_py,
1778
+ _COPY_FUNCS,
1779
+ copy_bufs,
1780
+ )
1781
+ torch.accelerator.synchronize()
1782
+
1783
+ # Python reference: src_block_idx(0) != dest_block_idx(1) -> no change
1784
+ assert input_batch_py.num_accepted_tokens_cpu[0] == 3, (
1785
+ f"Python: num_accepted_tokens should remain 3, "
1786
+ f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
1787
+ )
1788
+
1789
+ # --- Run GPU path ---
1790
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
1791
+ num_reqs = len(req_ids)
1792
+ block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
1793
+ block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
1794
+
1795
+ _run_gpu_postprocess(
1796
+ gpu_ctx,
1797
+ kv_cache_config=kv_cache_config,
1798
+ forward_context=forward_context_gpu,
1799
+ copy_funcs=_COPY_FUNCS,
1800
+ block_table=block_table_gpu,
1801
+ req_ids=req_ids,
1802
+ num_accepted_tokens=num_accepted_tokens,
1803
+ mamba_state_idx=mamba_state_idx,
1804
+ num_scheduled_tokens=num_scheduled_tokens,
1805
+ num_computed_tokens=num_computed_tokens,
1806
+ num_draft_tokens=num_draft_tokens,
1807
+ device=device,
1808
+ )
1809
+
1810
+ # The critical assertion: kernel must NOT set num_accepted_tokens to 1
1811
+ # when src_block_idx != dest_block_idx, even though src_addr == dst_addr
1812
+ # due to prefix caching aliasing.
1813
+ #
1814
+ # Old kernel (959ca0fd): `if src_addr == dst_addr` -> FAILS here (sets 1)
1815
+ # Fixed kernel (6466ce0d): `if src_block_idx == dest_block_idx and
1816
+ # accept_token_bias == 0` -> PASSES (preserves 3)
1817
+ kernel_accepted = gpu_ctx.num_accepted_tokens_out[0].item()
1818
+ assert kernel_accepted == 3, (
1819
+ f"Kernel set num_accepted_tokens to {kernel_accepted} but expected 3. "
1820
+ f"The early-return guard likely compared physical addresses "
1821
+ f"(src_addr == dst_addr) instead of logical block indices "
1822
+ f"(src_block_idx == dest_block_idx). Under prefix caching, "
1823
+ f"different logical blocks can share the same physical block."
1824
+ )
1825
+
1826
+ # Also verify state tensors match Python
1827
+ torch.testing.assert_close(
1828
+ conv_state_gpu,
1829
+ conv_state_py,
1830
+ msg="GPU conv state should match Python",
1831
+ )
1832
+ torch.testing.assert_close(
1833
+ temporal_state_gpu,
1834
+ temporal_state_py,
1835
+ msg="GPU temporal state should match Python",
1836
+ )
1837
+
1838
+ def test_as_strided_temporal_copy_size(self, device, test_config):
1839
+ """
1840
+ Regression test for 240723d46: temporal copy_size must be
1841
+ inner_size * elem_size, not state_block_stride.
1842
+
1843
+ In production (gpu_model_runner.py), conv and temporal states share
1844
+ a raw buffer via torch.as_strided where stride(0) equals
1845
+ page_size_bytes / elem_size — larger than either state's natural
1846
+ element count. Using stride(0) as copy_size for temporal states
1847
+ overwrites into the next block's conv region.
1848
+
1849
+ Layout per page (384 float16 elements = 768 bytes):
1850
+ [conv: 256 elems | temporal: 128 elems]
1851
+
1852
+ The test triggers a temporal copy from block 4 to block 3. With the
1853
+ bug the kernel copies 768 bytes (page stride) instead of 256 bytes
1854
+ (128 * 2), overwriting conv_state[4] with conv_state[5]'s data.
1855
+
1856
+ Test setup (block_size=16):
1857
+ - running = 60 + 5 - 2 = 63
1858
+ - new = 63 + 3 - 1 = 65
1859
+ - aligned = 64 >= 63 -> COPY needed
1860
+ - accept_token_bias = 64 - 63 = 1
1861
+ - dest_block_idx = 64 // 16 - 1 = 3
1862
+ - temporal: actual_src_block_idx = 3 + 1 = 4 (block_ids[4] = 4)
1863
+ """
1864
+ cfg = test_config
1865
+ torch.manual_seed(7001)
1866
+
1867
+ req_ids = ["req_0"]
1868
+ num_computed_tokens = [60]
1869
+ num_scheduled_tokens = {"req_0": 5}
1870
+ num_draft_tokens = {"req_0": 2}
1871
+ num_accepted_tokens = [3]
1872
+ mamba_state_idx = [3]
1873
+ block_ids_per_req = [list(range(8))]
1874
+
1875
+ layer_names = ["layer_0"]
1876
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
1877
+
1878
+ # --- Production-like packed layout (mirrors gpu_model_runner.py) ---
1879
+ conv_shape = (cfg.conv_width, cfg.conv_inner_dim)
1880
+ temporal_shape = (cfg.temporal_state_dim,)
1881
+ dtype = cfg.dtype
1882
+ elem_size = torch.tensor([], dtype=dtype).element_size()
1883
+
1884
+ conv_natural_elems = cfg.conv_width * cfg.conv_inner_dim
1885
+ temporal_natural_elems = cfg.temporal_state_dim
1886
+ page_size_bytes = (conv_natural_elems + temporal_natural_elems) * elem_size
1887
+ num_element_per_page = page_size_bytes // elem_size
1888
+
1889
+ assert num_element_per_page > temporal_natural_elems, (
1890
+ "Test requires padded stride; page must be larger than one state"
1891
+ )
1892
+
1893
+ raw_py = torch.randn(
1894
+ cfg.num_blocks * num_element_per_page, dtype=dtype, device=device
1895
+ )
1896
+ raw_gpu = raw_py.clone()
1897
+
1898
+ def make_views(raw):
1899
+ conv_tgt = (cfg.num_blocks, *conv_shape)
1900
+ conv_nat_stride = torch.empty(conv_tgt).stride()
1901
+ conv = torch.as_strided(
1902
+ raw,
1903
+ size=conv_tgt,
1904
+ stride=(num_element_per_page, *conv_nat_stride[1:]),
1905
+ storage_offset=0,
1906
+ )
1907
+
1908
+ temp_tgt = (cfg.num_blocks, *temporal_shape)
1909
+ temp_nat_stride = torch.empty(temp_tgt).stride()
1910
+ temp = torch.as_strided(
1911
+ raw,
1912
+ size=temp_tgt,
1913
+ stride=(num_element_per_page, *temp_nat_stride[1:]),
1914
+ storage_offset=conv_natural_elems,
1915
+ )
1916
+ return conv, temp
1917
+
1918
+ conv_py, temp_py = make_views(raw_py)
1919
+ conv_gpu, temp_gpu = make_views(raw_gpu)
1920
+
1921
+ fwd_py = {"layer_0": _make_mock_attention(conv_py, temp_py)}
1922
+ fwd_gpu = {"layer_0": _make_mock_attention(conv_gpu, temp_gpu)}
1923
+
1924
+ # --- Python reference ---
1925
+ sched = _make_postprocess_scheduler_output(
1926
+ req_ids,
1927
+ num_scheduled_tokens,
1928
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
1929
+ )
1930
+ batch_py = _make_input_batch(
1931
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
1932
+ )
1933
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
1934
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
1935
+
1936
+ postprocess_mamba(
1937
+ sched,
1938
+ kv_cache_config,
1939
+ batch_py,
1940
+ requests,
1941
+ fwd_py,
1942
+ _COPY_FUNCS,
1943
+ copy_bufs,
1944
+ )
1945
+ torch.accelerator.synchronize()
1946
+
1947
+ # --- GPU fused kernel ---
1948
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
1949
+ num_reqs = 1
1950
+ block_table = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
1951
+ block_table[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
1952
+
1953
+ _run_gpu_postprocess(
1954
+ gpu_ctx,
1955
+ kv_cache_config=kv_cache_config,
1956
+ forward_context=fwd_gpu,
1957
+ copy_funcs=_COPY_FUNCS,
1958
+ block_table=block_table,
1959
+ req_ids=req_ids,
1960
+ num_accepted_tokens=num_accepted_tokens,
1961
+ mamba_state_idx=mamba_state_idx,
1962
+ num_scheduled_tokens=num_scheduled_tokens,
1963
+ num_computed_tokens=num_computed_tokens,
1964
+ num_draft_tokens=num_draft_tokens,
1965
+ device=device,
1966
+ )
1967
+
1968
+ # --- Assertions ---
1969
+ # With the bug (pre-240723d46), the kernel copies page_size_bytes
1970
+ # (768) for temporal state instead of 256 bytes, overwriting
1971
+ # conv_state[4] with conv_state[5]'s data.
1972
+ torch.testing.assert_close(
1973
+ conv_gpu,
1974
+ conv_py,
1975
+ msg=(
1976
+ "Conv state corrupted: temporal copy_size was likely "
1977
+ "state_block_stride instead of inner_size * elem_size"
1978
+ ),
1979
+ )
1980
+ torch.testing.assert_close(
1981
+ temp_gpu,
1982
+ temp_py,
1983
+ msg="Temporal state mismatch",
1984
+ )
1985
+
1986
+ expected_accepted = torch.tensor(
1987
+ batch_py.num_accepted_tokens_cpu[:num_reqs],
1988
+ dtype=torch.int32,
1989
+ device=device,
1990
+ )
1991
+ torch.testing.assert_close(
1992
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
1993
+ expected_accepted,
1994
+ msg="num_accepted_tokens mismatch",
1995
+ )
1996
+
1997
+ def test_temporal_copy_with_bias_ge_2(self, device, test_config):
1998
+ """
1999
+ Coverage test for the temporal-state block-table stride arithmetic
2000
+ when ``accept_token_bias >= 2``.
2001
+
2002
+ The kernel computes, for temporal (non-conv) states::
2003
+
2004
+ actual_src_block_idx = src_block_idx + accept_token_bias
2005
+ actual_src_block_id = block_table[req, actual_src_block_idx]
2006
+
2007
+ All prior regression tests exercise only ``bias == 1``, i.e. they
2008
+ only ever read one slot ahead of ``src_block_idx`` in the block
2009
+ table. An off-by-one (or missing scale) in the address computation
2010
+ on line 143 of ``mamba_utils.py`` would be invisible to every
2011
+ existing test but would silently read the wrong physical block on
2012
+ any speculative-decode cycle that accepts multiple tokens across a
2013
+ block boundary, feeding a stale hidden state forward one step.
2014
+
2015
+ Setup (block_size=16):
2016
+ - running = 28 + 2 - 0 = 30
2017
+ - new = 30 + 3 - 1 = 32
2018
+ - aligned = 32 >= 30 -> COPY needed
2019
+ - bias = 32 - 30 = 2 (key: >= 2)
2020
+ - dest_idx = 32 // 16 - 1 = 1
2021
+ - src_idx = 1 (same as dest -> exercises post-copy accepted=1 write)
2022
+ - temporal actual_src_block_idx = 1 + 2 = 3 (reads block_table[0, 3])
2023
+
2024
+ With identity block_ids = [0,1,2,3,...], an off-by-one that used
2025
+ bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3,
2026
+ producing a clear state-value mismatch against the Python
2027
+ reference.
2028
+ """
2029
+ cfg = test_config
2030
+ torch.manual_seed(7002)
2031
+
2032
+ req_ids = ["req_0"]
2033
+ num_computed_tokens = [28]
2034
+ num_scheduled_tokens = {"req_0": 2}
2035
+ num_draft_tokens: dict[str, int] = {}
2036
+ num_accepted_tokens = [3] # -> accept_token_bias = 2
2037
+ mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx
2038
+ block_ids_per_req = [list(range(8))]
2039
+
2040
+ layer_names = ["layer_0"]
2041
+ kv_cache_config = _make_kv_cache_config(cfg, layer_names)
2042
+
2043
+ (
2044
+ conv_state_py,
2045
+ temporal_state_py,
2046
+ conv_state_gpu,
2047
+ temporal_state_gpu,
2048
+ fwd_py,
2049
+ fwd_gpu,
2050
+ ) = _make_dual_layer_state(cfg, device)
2051
+ temporal_state_orig = temporal_state_py.clone()
2052
+
2053
+ # --- Python reference ---
2054
+ sched = _make_postprocess_scheduler_output(
2055
+ req_ids,
2056
+ num_scheduled_tokens,
2057
+ {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
2058
+ )
2059
+ batch_py = _make_input_batch(
2060
+ req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
2061
+ )
2062
+ requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
2063
+ copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)
2064
+
2065
+ postprocess_mamba(
2066
+ sched,
2067
+ kv_cache_config,
2068
+ batch_py,
2069
+ requests,
2070
+ fwd_py,
2071
+ _COPY_FUNCS,
2072
+ copy_bufs,
2073
+ )
2074
+ torch.accelerator.synchronize()
2075
+
2076
+ # --- GPU fused kernel ---
2077
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
2078
+ num_reqs = 1
2079
+ block_table = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
2080
+ block_table[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)
2081
+
2082
+ _run_gpu_postprocess(
2083
+ gpu_ctx,
2084
+ kv_cache_config=kv_cache_config,
2085
+ forward_context=fwd_gpu,
2086
+ copy_funcs=_COPY_FUNCS,
2087
+ block_table=block_table,
2088
+ req_ids=req_ids,
2089
+ num_accepted_tokens=num_accepted_tokens,
2090
+ mamba_state_idx=mamba_state_idx,
2091
+ num_scheduled_tokens=num_scheduled_tokens,
2092
+ num_computed_tokens=num_computed_tokens,
2093
+ num_draft_tokens=num_draft_tokens,
2094
+ device=device,
2095
+ )
2096
+
2097
+ # --- Ground truth: Python must have sourced temporal from block 3 ---
2098
+ actual_src_block_id = block_ids_per_req[0][3] # == 3
2099
+ dest_block_id = block_ids_per_req[0][1] # == 1
2100
+ torch.testing.assert_close(
2101
+ temporal_state_py[dest_block_id],
2102
+ temporal_state_orig[actual_src_block_id],
2103
+ msg=(
2104
+ "Python reference did not copy from block_ids[src+bias]=3; "
2105
+ "test preconditions are wrong"
2106
+ ),
2107
+ )
2108
+
2109
+ # --- GPU kernel must match Python byte-for-byte ---
2110
+ torch.testing.assert_close(
2111
+ conv_state_gpu,
2112
+ conv_state_py,
2113
+ msg="Conv state mismatch at accept_token_bias=2",
2114
+ )
2115
+ torch.testing.assert_close(
2116
+ temporal_state_gpu,
2117
+ temporal_state_py,
2118
+ msg=(
2119
+ "Temporal state mismatch at accept_token_bias=2: the kernel "
2120
+ "likely read the wrong slot of the block table "
2121
+ "(actual_src_block_idx stride arithmetic)"
2122
+ ),
2123
+ )
2124
+
2125
+ expected_accepted = torch.tensor(
2126
+ batch_py.num_accepted_tokens_cpu[:num_reqs],
2127
+ dtype=torch.int32,
2128
+ device=device,
2129
+ )
2130
+ torch.testing.assert_close(
2131
+ gpu_ctx.num_accepted_tokens_out[:num_reqs],
2132
+ expected_accepted,
2133
+ msg="num_accepted_tokens mismatch at accept_token_bias=2",
2134
+ )
vllm/v1/worker/gpu_model_runner.py
+65
-37
@@ -887,7 +887,7 @@ class GPUModelRunner(
887
self.execute_model_state: ExecuteModelState | None = None
888
self.kv_connector_output: KVConnectorOutput | None = None
889
self.mamba_state_idx: dict[str, int] = {}
890
- self._mamba_copy_bufs: mamba_utils.MambaCopyBuffers | None = None
890
+ self._mamba_bufs: mamba_utils.MambaBuffers | None = None
891
self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None
892
if self.cache_config.mamba_cache_mode == "all" and self.num_spec_tokens > 0:
893
self.mamba_prev_last_scheduled_idx = self._make_buffer(
@@ -992,15 +992,23 @@ class GPUModelRunner(
992
with_numpy=numpy,
993
)
994
995
- def _get_mamba_copy_bufs(self) -> mamba_utils.MambaCopyBuffers:
996
- if self._mamba_copy_bufs is None:
997
- self._mamba_copy_bufs = mamba_utils.MambaCopyBuffers.create(
998
- self.max_num_reqs,
999
- self.kv_cache_config,
1000
- self.model.get_mamba_state_copy_func(),
1001
- self._make_buffer,
995
+ def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers:
996
+ # Only reachable on the ``mamba_cache_mode == "align"`` path.
997
+ # The postprocess sub-object is additionally gated on spec
998
+ # decode + hybrid model.
999
+ assert self.cache_config.mamba_cache_mode == "align"
1000
+ if self._mamba_bufs is None:
1001
+ self._mamba_bufs = mamba_utils.MambaBuffers.create(
1002
+ max_num_reqs=self.max_num_reqs,
1003
+ kv_cache_config=self.kv_cache_config,
1004
+ copy_funcs=self.model.get_mamba_state_copy_func(),
1005
+ make_buffer=self._make_buffer,
1006
+ device=self.device,
1007
+ with_postprocess_align=(
1008
+ self.speculative_config is not None and self.model_config.is_hybrid
1009
+ ),
1010
)
1003
- return self._mamba_copy_bufs
1011
+ return self._mamba_bufs
1012
1013
def _init_model_kwargs(self):
1014
model_kwargs = dict[str, Any]()
@@ -1489,21 +1497,32 @@ class GPUModelRunner(
1497
if not self.speculative_config or not self.model_config.is_hybrid:
1498
return
1499
1492
- # TODO: Remove .cpu() sync to enable fully async for hybrid model;
1493
- # Use num_computed_tokens.gpu instead of req.num_computed_tokens to
1494
- # support aligned mamba cache mode.
1500
# Count the number of accepted tokens for each sequence.
1501
# Valid tokens are contiguous from position 0, so counting non-(-1)
1502
# tokens gives us the first -1 position (i.e., number of accepted).
1503
num_reqs = output_token_ids.size(0)
1504
self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1)
1505
1501
- is_align = self.cache_config.mamba_cache_mode == "align"
1502
- if is_align:
1503
- for i, num_tokens in enumerate(
1504
- self.num_accepted_tokens.gpu[:num_reqs].cpu().numpy()
1505
- ):
1506
- self.input_batch.num_accepted_tokens_cpu[i] = num_tokens
1506
+ if self.cache_config.mamba_cache_mode == "align":
1507
+ # Fused GPU postprocess: state copies + per-request accepted-token
1508
+ # update without CPU-GPU sync. The metadata
1509
+ # (num_scheduled_tokens, num_draft_tokens, num_computed_tokens) is
1510
+ # pre-staged to GPU buffers in _prepare_inputs.
1511
+ mamba_utils.postprocess_mamba_align_gpu(
1512
+ bufs=self._get_mamba_bufs(),
1513
+ num_reqs=num_reqs,
1514
+ num_accepted_tokens_gpu=self.num_accepted_tokens.gpu,
1515
+ num_accepted_tokens_cpu_tensor=(
1516
+ self.input_batch.num_accepted_tokens_cpu_tensor
1517
+ ),
1518
+ input_batch=self.input_batch,
1519
+ kv_cache_config=self.kv_cache_config,
1520
+ forward_context=self.compilation_config.static_forward_context,
1521
+ mamba_state_copy_funcs=self.model.get_mamba_state_copy_func(),
1522
+ )
1523
+
1524
+ assert self.num_accepted_tokens_event is not None
1525
+ self.num_accepted_tokens_event.record()
1526
else:
1527
self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_(
1528
self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True
@@ -1511,23 +1530,16 @@ class GPUModelRunner(
1530
assert self.num_accepted_tokens_event is not None
1531
self.num_accepted_tokens_event.record()
1532
1514
- mamba_utils.postprocess_mamba(
1515
- scheduler_output,
1516
- self.kv_cache_config,
1517
- self.cache_config,
1518
- self.input_batch,
1519
- self.requests,
1520
- self.mamba_state_idx,
1521
- self.num_spec_tokens,
1522
- num_reqs,
1523
- forward_context=(
1524
- self.compilation_config.static_forward_context if is_align else None
1525
- ),
1526
- mamba_state_copy_funcs=(
1527
- self.model.get_mamba_state_copy_func() if is_align else None
1528
- ),
1529
- copy_bufs=self._get_mamba_copy_bufs() if is_align else None,
1530
- )
1533
+ if self.cache_config.mamba_cache_mode == "all":
1534
+ mamba_utils.postprocess_mamba_all(
1535
+ scheduler_output,
1536
+ self.kv_cache_config,
1537
+ self.input_batch,
1538
+ self.requests,
1539
+ self.mamba_state_idx,
1540
+ self.num_spec_tokens,
1541
+ num_reqs,
1542
+ )
1543
1544
def _update_streaming_request(
1545
self, req_id: str, new_req_data: NewRequestData
@@ -4098,6 +4110,7 @@ class GPUModelRunner(
4110
if deferred_state_corrections_fn:
4111
deferred_state_corrections_fn()
4112
deferred_state_corrections_fn = None
4113
+ mamba_bufs = self._get_mamba_bufs()
4114
mamba_utils.preprocess_mamba(
4115
scheduler_output,
4116
self.kv_cache_config,
@@ -4107,7 +4120,7 @@ class GPUModelRunner(
4120
self.requests,
4121
self.compilation_config.static_forward_context,
4122
self.model.get_mamba_state_copy_func(),
4110
- self._get_mamba_copy_bufs(),
4123
+ mamba_bufs.preprocess,
4124
)
4125
# preprocess_mamba resets num_accepted_tokens_cpu to 1
4126
# for requests whose state was copied to a new block.
@@ -4118,6 +4131,21 @@ class GPUModelRunner(
4131
)
4132
self.num_accepted_tokens.copy_to_gpu(num_reqs)
4133
4134
+ # Stage per-request inputs for the fused postprocess kernel
4135
+ # only when that kernel will actually run. The kernel is
4136
+ # gated on spec-decode + hybrid (see MambaBuffers.create);
4137
+ # without it, ``mamba_bufs.postprocess_align`` is None and
4138
+ # the staging buffers don't exist.
4139
+ if mamba_bufs.postprocess_align is not None:
4140
+ mamba_utils.stage_postprocess_inputs_to_gpu(
4141
+ mamba_bufs.postprocess_align,
4142
+ scheduler_output,
4143
+ self.input_batch.req_ids,
4144
+ num_reqs,
4145
+ self.requests,
4146
+ self.mamba_state_idx,
4147
+ )
4148
+
4149
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
4150
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
4151
@@ -7114,7 +7142,7 @@ class GPUModelRunner(
7142
"""
7143
kv_cache_config = deepcopy(kv_cache_config)
7144
self.kv_cache_config = kv_cache_config
7117
- self._mamba_copy_bufs = None
7145
+ self._mamba_bufs = None
7146
self.may_add_encoder_only_layers_to_kv_cache_config()
7147
self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
7148
self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
vllm/v1/worker/mamba_utils.py
+646
-76
@@ -10,6 +10,8 @@ import torch
10
from vllm.config import CacheConfig
11
from vllm.model_executor.layers.mamba.mamba_utils import (
12
MambaStateCopyFunc,
13
+ get_conv_copy_spec,
14
+ get_temporal_copy_spec,
15
)
16
from vllm.triton_utils import tl, triton
17
from vllm.utils.math_utils import cdiv
@@ -20,6 +22,157 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState
22
from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
23
24
25
+@triton.jit
26
+def postprocess_mamba_fused_kernel(
27
+ # Decision inputs (per-request)
28
+ num_accepted_tokens_ptr,
29
+ mamba_state_idx_ptr,
30
+ num_scheduled_tokens_ptr,
31
+ num_computed_tokens_ptr,
32
+ num_draft_tokens_ptr,
33
+ # Per-group block table base addresses: int64[num_groups]. Each entry is
34
+ # the data_ptr of that group's persistent [max_reqs, max_blocks] int32
35
+ # block table.
36
+ block_table_ptrs_ptr,
37
+ block_table_stride_req: tl.int64, # stride between requests (in elements)
38
+ # Mamba state metadata (per-layer, per-state-type)
39
+ # These are 1D arrays indexed by (layer_idx * num_state_types + state_type_idx)
40
+ state_base_addrs_ptr, # base address of each state tensor
41
+ state_block_strides_ptr, # bytes per block for each state
42
+ state_elem_sizes_ptr, # element size for each state
43
+ state_inner_sizes_ptr, # number of elements in inner dimensions
44
+ state_conv_widths_ptr, # conv width for conv states (0 for temporal)
45
+ state_group_indices_ptr, # maps state_idx to group index in block table
46
+ # Output: num_accepted_tokens update (for src==dst case)
47
+ num_accepted_tokens_out_ptr,
48
+ # Runtime parameter (varies per batch - NOT constexpr to avoid recompilation)
49
+ num_reqs,
50
+ # Compile-time constants (fixed after model initialization)
51
+ # block_size: determined by model config, constant for all invocations
52
+ block_size: tl.constexpr,
53
+ # COPY_BLOCK_SIZE: fixed tuning parameter for memory copy loop
54
+ COPY_BLOCK_SIZE: tl.constexpr,
55
+):
56
+ """
57
+ Fused GPU kernel for postprocess_mamba that computes decisions AND performs
58
+ mamba state copies without any CPU-GPU synchronization.
59
+
60
+ Grid: (num_reqs, num_layers * num_state_types)
61
+ - program_id(0) = request index
62
+ - program_id(1) = state_idx (flattened index into layer/state_type metadata)
63
+
64
+ Note: num_layers and num_state_types are not passed as kernel parameters
65
+ because the kernel indexes directly into pre-flattened metadata arrays
66
+ using program_id(1). The grid dimensions encode the total state count.
67
+ """
68
+ req_idx = tl.program_id(0)
69
+ state_idx = tl.program_id(1)
70
+
71
+ # Bounds check
72
+ if req_idx >= num_reqs:
73
+ return
74
+
75
+ # Compute decision logic (mirrors postprocess_mamba Python reference)
76
+ num_accepted = tl.load(num_accepted_tokens_ptr + req_idx)
77
+ src_block_idx = tl.load(mamba_state_idx_ptr + req_idx)
78
+ num_scheduled = tl.load(num_scheduled_tokens_ptr + req_idx)
79
+ num_computed = tl.load(num_computed_tokens_ptr + req_idx)
80
+ num_draft = tl.load(num_draft_tokens_ptr + req_idx)
81
+
82
+ num_tokens_running_state = num_computed + num_scheduled - num_draft
83
+ new_num_computed = num_tokens_running_state + num_accepted - 1
84
+ aligned_new_computed = (new_num_computed // block_size) * block_size
85
+
86
+ needs_copy = aligned_new_computed >= num_tokens_running_state
87
+
88
+ if not needs_copy:
89
+ return
90
+
91
+ # Compute copy parameters
92
+ accept_token_bias = aligned_new_computed - num_tokens_running_state
93
+ dest_block_idx = aligned_new_computed // block_size - 1
94
+
95
+ # Load state metadata for this layer/state_type
96
+ state_base_addr = tl.load(state_base_addrs_ptr + state_idx)
97
+ state_block_stride = tl.load(state_block_strides_ptr + state_idx)
98
+ state_elem_size = tl.load(state_elem_sizes_ptr + state_idx)
99
+ state_inner_size = tl.load(state_inner_sizes_ptr + state_idx)
100
+ conv_width = tl.load(state_conv_widths_ptr + state_idx)
101
+
102
+ # Load the group index for this state, then index into the correct
103
+ # group's block table. Each mamba group has independently allocated
104
+ # physical blocks.
105
+ group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64)
106
+
107
+ # block_table_ptrs_ptr holds one pointer per group (each group owns its own
108
+ # block table). Reinterpret as int32* since block ids are int32.
109
+ group_base_addr = tl.load(block_table_ptrs_ptr + group_idx)
110
+ block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32))
111
+ block_table_base = block_table_typed + req_idx * block_table_stride_req
112
+
113
+ # Widen block ids to int64 before they reach `block_id * state_block_stride`
114
+ # below: state_block_stride can exceed 2**31 bytes for large mamba caches,
115
+ # and Triton would otherwise do the multiply in int32 and wrap.
116
+ src_block_id = tl.load(block_table_base + src_block_idx).to(tl.int64)
117
+ dest_block_id = tl.load(block_table_base + dest_block_idx).to(tl.int64)
118
+
119
+ # Compute source and destination addresses based on state type
120
+ # conv_width > 0 means this is a conv state (get_conv_copy_spec logic)
121
+ # conv_width == 0 means this is a temporal state (get_temporal_copy_spec logic)
122
+ is_conv_state = conv_width > 0
123
+
124
+ if is_conv_state:
125
+ # Conv state: copy
126
+ # state[block_table[req_idx, src_block_idx], accept_token_bias:]
127
+ # to
128
+ # state[block_table[req_idx, dest_block_idx], :conv_width - accept_token_bias]
129
+ src_offset = accept_token_bias.to(tl.int64) * state_inner_size * state_elem_size
130
+ src_addr = state_base_addr + src_block_id * state_block_stride + src_offset
131
+ dst_addr = state_base_addr + dest_block_id * state_block_stride
132
+ # Number of elements to copy:
133
+ # (conv_width - accept_token_bias) * inner_size
134
+ num_elems_to_copy = (conv_width - accept_token_bias).to(
135
+ tl.int64
136
+ ) * state_inner_size
137
+ copy_size = num_elems_to_copy * state_elem_size
138
+ else:
139
+ # Temporal state: copy
140
+ # state[block_table[req_idx, src_block_idx + accept_token_bias]]
141
+ # to
142
+ # state[block_table[req_idx, dest_block_idx]]
143
+ actual_src_block_idx = src_block_idx + accept_token_bias
144
+ actual_src_block_id = tl.load(block_table_base + actual_src_block_idx).to(
145
+ tl.int64
146
+ )
147
+ src_addr = state_base_addr + actual_src_block_id * state_block_stride
148
+ dst_addr = state_base_addr + dest_block_id * state_block_stride
149
+ # Use natural block data size (inner_size * elem_size), NOT
150
+ # state_block_stride which is the page stride and can exceed the
151
+ # actual data when the state tensor uses as_strided page padding.
152
+ copy_size = state_inner_size * state_elem_size
153
+
154
+ # Mirror postprocess_mamba's trailing
155
+ # if src_block_idx == dest_block_idx: num_accepted_tokens_cpu[i] = 1
156
+ # This runs whether or not the copy below is skipped (it's per-request, so
157
+ # only state_idx == 0 writes).
158
+ if src_block_idx == dest_block_idx and state_idx == 0:
159
+ tl.store(num_accepted_tokens_out_ptr + req_idx, 1)
160
+
161
+ # Mirror collect_mamba_copy_meta's early return: src==dst with no token
162
+ # bias means source and destination ranges coincide, so the copy is a
163
+ # no-op.
164
+ if src_block_idx == dest_block_idx and accept_token_bias == 0:
165
+ return
166
+
167
+ offsets = tl.arange(0, COPY_BLOCK_SIZE)
168
+ for i in range(0, copy_size, COPY_BLOCK_SIZE):
169
+ mask = (i + offsets) < copy_size
170
+ curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8))
171
+ curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8))
172
+ data = tl.load(curr_src, mask=mask)
173
+ tl.store(curr_dst, data, mask=mask)
174
+
175
+
176
@triton.jit
177
def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr):
178
pid = tl.program_id(0)
@@ -94,6 +247,328 @@ class MambaCopyBuffers:
247
)
248
249
250
+@dataclasses.dataclass
251
+class MambaSpecDecodeGPUContext:
252
+ """
253
+ Context for GPU-side Mamba state copy operations during the
254
+ fused postprocess path.
255
+
256
+ Only used when speculative decoding is enabled on a hybrid model
257
+ (and the mamba_cache_config is in align mode).
258
+
259
+ Precomputes memory layout metadata (base addresses, strides, element sizes)
260
+ so the GPU kernel can perform state copies without CPU-GPU sync.
261
+
262
+ State types are distinguished by conv_width: >0 for conv states (sliding
263
+ window with offset-based copies), 0 for temporal states (full block copies).
264
+ """
265
+
266
+ # Per-state metadata tensors (shape: [num_layers * num_state_types])
267
+ # These are populated from forward_context during the first forward pass
268
+ state_base_addrs: torch.Tensor # int64: base address of each state tensor
269
+ state_block_strides: torch.Tensor # int64: bytes per block
270
+ state_elem_sizes: torch.Tensor # int32: element size in bytes
271
+ state_inner_sizes: torch.Tensor # int64: elements in inner dimensions
272
+ state_conv_widths: torch.Tensor # int32: conv width (0 for temporal states)
273
+ state_group_indices: torch.Tensor # int32: maps state_idx to group index
274
+
275
+ # Configuration
276
+ block_size: int
277
+ num_layers: int
278
+ num_state_types: int
279
+ mamba_group_ids: list[int]
280
+ num_groups: int
281
+
282
+ # Output buffer for num_accepted_tokens updates
283
+ num_accepted_tokens_out: torch.Tensor
284
+
285
+ # Per-group block-table base addresses: int64[num_groups]. Populated in
286
+ # initialize_from_forward_context from the persistent per-group block
287
+ # table tensors (whose data_ptr is stable across steps).
288
+ block_table_ptrs: torch.Tensor
289
+ block_table_stride_req: int = 0
290
+
291
+ # Per-request staging buffers (CPU+GPU mirrors). The runner stages
292
+ # values into the CPU view in ``_prepare_inputs`` and the fused kernel
293
+ # reads the GPU side. These only exist when the postprocess kernel is
294
+ # enabled (spec decode + hybrid + align mode).
295
+ mamba_state_idx_buf: CpuGpuBuffer | None = None
296
+ num_scheduled_tokens_buf: CpuGpuBuffer | None = None
297
+ num_computed_tokens_buf: CpuGpuBuffer | None = None
298
+ num_draft_tokens_buf: CpuGpuBuffer | None = None
299
+
300
+ # Flag to track if metadata has been populated
301
+ is_initialized: bool = False
302
+
303
+ @classmethod
304
+ def create(
305
+ cls,
306
+ max_num_reqs: int,
307
+ kv_cache_config: KVCacheConfig,
308
+ num_state_types: int,
309
+ device: torch.device,
310
+ make_buffer: Callable[..., CpuGpuBuffer],
311
+ ) -> "MambaSpecDecodeGPUContext":
312
+ """Create context with allocated buffers (metadata populated later)."""
313
+ mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
314
+
315
+ # Count total layers across all mamba groups
316
+ num_layers = sum(
317
+ len(kv_cache_config.kv_cache_groups[gid].layer_names)
318
+ for gid in mamba_group_ids
319
+ )
320
+ total_states = num_layers * num_state_types
321
+
322
+ return cls(
323
+ state_base_addrs=torch.zeros(
324
+ total_states, dtype=torch.int64, device=device
325
+ ),
326
+ state_block_strides=torch.zeros(
327
+ total_states, dtype=torch.int64, device=device
328
+ ),
329
+ state_elem_sizes=torch.zeros(
330
+ total_states, dtype=torch.int32, device=device
331
+ ),
332
+ state_inner_sizes=torch.zeros(
333
+ total_states, dtype=torch.int64, device=device
334
+ ),
335
+ state_conv_widths=torch.zeros(
336
+ total_states, dtype=torch.int32, device=device
337
+ ),
338
+ state_group_indices=torch.zeros(
339
+ total_states, dtype=torch.int32, device=device
340
+ ),
341
+ block_size=mamba_spec.block_size,
342
+ num_layers=num_layers,
343
+ num_state_types=num_state_types,
344
+ mamba_group_ids=mamba_group_ids,
345
+ num_groups=len(mamba_group_ids),
346
+ num_accepted_tokens_out=torch.zeros(
347
+ max_num_reqs, dtype=torch.int32, device=device
348
+ ),
349
+ block_table_ptrs=torch.zeros(
350
+ len(mamba_group_ids), dtype=torch.int64, device=device
351
+ ),
352
+ mamba_state_idx_buf=make_buffer(max_num_reqs, dtype=torch.int32),
353
+ num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
354
+ num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
355
+ num_draft_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
356
+ is_initialized=False,
357
+ )
358
+
359
+ def initialize_from_forward_context(
360
+ self,
361
+ kv_cache_config: KVCacheConfig,
362
+ forward_context: dict[str, Any],
363
+ mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
364
+ block_tables: list[torch.Tensor],
365
+ ) -> None:
366
+ """
367
+ Extract and cache memory layout metadata from Mamba state tensors.
368
+
369
+ This method populates the pre-allocated metadata tensors with information
370
+ needed by `postprocess_mamba_fused_kernel` to perform state copies entirely
371
+ on the GPU without CPU-GPU synchronization.
372
+
373
+ For each Mamba layer and state type, the following metadata is extracted:
374
+ - state_base_addrs: GPU memory address (data_ptr) of the state tensor
375
+ - state_block_strides: Bytes between consecutive blocks (stride * elem_size)
376
+ - state_elem_sizes: Element size in bytes (e.g., 2 for float16)
377
+ - state_inner_sizes: For conv states, elements per conv position (stride(1)),
378
+ used to compute offset when slicing state[block, offset:]. For temporal
379
+ states, this field is unused (set to 1).
380
+ - state_conv_widths: Conv dimension size for conv states, 0 for temporal states
381
+
382
+ The conv vs temporal state type is detected by inspecting the copy function
383
+ name: functions containing "conv" are treated as conv states.
384
+
385
+ This method is idempotent - it only executes once (guarded by is_initialized
386
+ flag) since the metadata is static after model loading.
387
+
388
+ Args:
389
+ kv_cache_config: Configuration containing KV cache group info and
390
+ layer name mappings.
391
+ forward_context: Dictionary mapping layer names to attention objects,
392
+ populated after the model is loaded. Each attention object must
393
+ have a `kv_cache` attribute containing the list of state tensors.
394
+ mamba_state_copy_funcs: Tuple of copy functions (one per state type)
395
+ used to determine whether each state is a conv or temporal state.
396
+ block_tables: per-mamba-group persistent block-table tensors, in
397
+ the same order as `mamba_group_ids`. Their `data_ptr()` /
398
+ `stride(0)` are captured once for the kernel to index into.
399
+ """
400
+ if self.is_initialized:
401
+ return
402
+
403
+ idx = 0
404
+ for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids):
405
+ layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names
406
+ for layer_name in layer_names:
407
+ attention = forward_context[layer_name]
408
+ kv_caches: list[torch.Tensor] = attention.kv_cache
409
+
410
+ for state_type_idx, state in enumerate(kv_caches):
411
+ # Base address
412
+ self.state_base_addrs[idx] = state.data_ptr()
413
+
414
+ # Block stride (bytes between consecutive blocks)
415
+ # state shape: [num_blocks, ...], stride(0) = elements per block
416
+ if state.dim() > 1:
417
+ block_stride_elems = state.stride(0)
418
+ else:
419
+ block_stride_elems = state.numel()
420
+ self.state_block_strides[idx] = (
421
+ block_stride_elems * state.element_size()
422
+ )
423
+
424
+ # Element size
425
+ self.state_elem_sizes[idx] = state.element_size()
426
+
427
+ copy_func = mamba_state_copy_funcs[state_type_idx]
428
+ assert (
429
+ copy_func is get_conv_copy_spec
430
+ or copy_func is get_temporal_copy_spec
431
+ ), f"unexpected copy func: {copy_func}"
432
+ if copy_func is get_conv_copy_spec:
433
+ # Conv state: conv_width is state.size(1)
434
+ # inner_size is stride(1) = elements per conv position,
435
+ # used to compute byte offset for state[block, offset:]
436
+ conv_w = state.size(1) if state.dim() > 1 else 0
437
+ self.state_conv_widths[idx] = conv_w
438
+ if state.dim() > 2:
439
+ # stride(1) = product of dims[2:] for contiguous tensor
440
+ self.state_inner_sizes[idx] = state.stride(1)
441
+ else:
442
+ # 2D tensor: [num_blocks, conv_dim], no inner dims
443
+ self.state_inner_sizes[idx] = 1
444
+ else:
445
+ # Temporal state: inner_size = natural elements per
446
+ # block (prod of inner dims). The kernel uses this
447
+ # to compute copy_size = inner_size * elem_size,
448
+ # which gives the correct byte count even when the
449
+ # state tensor is as_strided with padded page strides
450
+ # (state_block_stride would be the page size, too big).
451
+ self.state_conv_widths[idx] = 0
452
+ self.state_inner_sizes[idx] = (
453
+ state[0].numel() if state.dim() > 1 else 1
454
+ )
455
+
456
+ self.state_group_indices[idx] = group_local_idx
457
+ idx += 1
458
+
459
+ # Cache per-group block-table base addresses and per-request stride.
460
+ # `block_tables[i]` is the persistent 2D int32 block-table tensor for
461
+ # `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the
462
+ # engine's lifetime, so we capture them once here.
463
+ assert len(block_tables) == self.num_groups, (
464
+ f"expected {self.num_groups} block tables, got {len(block_tables)}"
465
+ )
466
+ strides = {bt.stride(0) for bt in block_tables}
467
+ assert len(strides) == 1, (
468
+ f"all mamba block tables must share stride(0), got {strides}"
469
+ )
470
+ self.block_table_stride_req = int(next(iter(strides)))
471
+ for i, bt in enumerate(block_tables):
472
+ self.block_table_ptrs[i] = bt.data_ptr()
473
+
474
+ self.is_initialized = True
475
+
476
+ def run_fused_postprocess(
477
+ self,
478
+ num_reqs: int,
479
+ num_accepted_tokens_gpu: torch.Tensor,
480
+ mamba_state_idx_gpu: torch.Tensor,
481
+ num_scheduled_tokens_gpu: torch.Tensor,
482
+ num_computed_tokens_gpu: torch.Tensor,
483
+ num_draft_tokens_gpu: torch.Tensor,
484
+ ) -> None:
485
+ """
486
+ Run the fused postprocess_mamba kernel on GPU.
487
+
488
+ This computes decisions and performs mamba state copies entirely on GPU,
489
+ eliminating the CPU-GPU sync that was previously needed.
490
+
491
+ Args:
492
+ num_reqs: Number of active requests
493
+ num_accepted_tokens_gpu: [num_reqs] accepted token counts
494
+ mamba_state_idx_gpu: [num_reqs] source block indices
495
+ num_scheduled_tokens_gpu: [num_reqs] scheduled token counts
496
+ num_computed_tokens_gpu: [num_reqs] computed token counts
497
+ num_draft_tokens_gpu: [num_reqs] draft token counts
498
+ """
499
+ if num_reqs == 0 or not self.is_initialized:
500
+ return
501
+
502
+ # Initialize output to current values (unchanged unless src==dst)
503
+ self.num_accepted_tokens_out[:num_reqs].copy_(
504
+ num_accepted_tokens_gpu[:num_reqs]
505
+ )
506
+
507
+ total_states = self.num_layers * self.num_state_types
508
+ grid = (num_reqs, total_states)
509
+
510
+ postprocess_mamba_fused_kernel[grid](
511
+ num_accepted_tokens_gpu,
512
+ mamba_state_idx_gpu,
513
+ num_scheduled_tokens_gpu,
514
+ num_computed_tokens_gpu,
515
+ num_draft_tokens_gpu,
516
+ self.block_table_ptrs,
517
+ self.block_table_stride_req,
518
+ self.state_base_addrs,
519
+ self.state_block_strides,
520
+ self.state_elem_sizes,
521
+ self.state_inner_sizes,
522
+ self.state_conv_widths,
523
+ self.state_group_indices,
524
+ self.num_accepted_tokens_out,
525
+ num_reqs,
526
+ block_size=self.block_size,
527
+ COPY_BLOCK_SIZE=1024,
528
+ )
529
+
530
+
531
+@dataclasses.dataclass
532
+class MambaBuffers:
533
+ """Single owner for all mamba-specific runner buffers.
534
+
535
+ The two sub-objects have different gates:
536
+ ``preprocess`` is needed whenever ``mamba_cache_mode == "align"``;
537
+ ``postprocess_align`` is needed only when align is combined with
538
+ speculative decoding on a hybrid model, and is ``None`` otherwise.
539
+ """
540
+
541
+ preprocess: MambaCopyBuffers
542
+ postprocess_align: MambaSpecDecodeGPUContext | None
543
+
544
+ @classmethod
545
+ def create(
546
+ cls,
547
+ max_num_reqs: int,
548
+ kv_cache_config: KVCacheConfig,
549
+ copy_funcs: tuple[MambaStateCopyFunc, ...],
550
+ make_buffer: Callable[..., CpuGpuBuffer],
551
+ device: torch.device,
552
+ with_postprocess_align: bool,
553
+ ) -> "MambaBuffers":
554
+ return cls(
555
+ preprocess=MambaCopyBuffers.create(
556
+ max_num_reqs, kv_cache_config, copy_funcs, make_buffer
557
+ ),
558
+ postprocess_align=(
559
+ MambaSpecDecodeGPUContext.create(
560
+ max_num_reqs=max_num_reqs,
561
+ kv_cache_config=kv_cache_config,
562
+ num_state_types=len(copy_funcs),
563
+ device=device,
564
+ make_buffer=make_buffer,
565
+ )
566
+ if with_postprocess_align
567
+ else None
568
+ ),
569
+ )
570
+
571
+
572
def collect_mamba_copy_meta(
573
copy_bufs: MambaCopyBuffers,
574
kv_cache_config: KVCacheConfig,
@@ -228,91 +703,34 @@ def preprocess_mamba(
703
do_mamba_copy_block(copy_bufs)
704
705
231
-def postprocess_mamba(
706
+def postprocess_mamba_all(
707
scheduler_output: SchedulerOutput,
708
kv_cache_config: KVCacheConfig,
234
- cache_config: CacheConfig,
709
input_batch: GPUInputBatch,
710
requests: dict[str, CachedRequestState],
711
mamba_state_idx: dict[str, int],
712
num_spec_tokens: int,
713
num_reqs: int,
240
- *,
241
- forward_context: dict[str, Any] | None = None,
242
- mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...] | None = None,
243
- copy_bufs: MambaCopyBuffers | None = None,
714
):
715
+ """All-mode postprocess (only meaningful with num_spec_tokens > 0):
716
+ record per-request the block index of the last token scheduled this
717
+ step, so the next step can anchor its in-place writes when accepted
718
+ drafts leave the sequence at a non-block-aligned position.
719
"""
246
- Post-model-execute mamba prefix-caching bookkeeping. Dispatched by
247
- cache_config.mamba_cache_mode:
248
- - "align": if a block is converted from partial to full this step,
249
- copy the running state into the new full block.
250
- - "all" + num_spec_tokens > 0: record per-request the block index of
251
- the last token scheduled this step, so the next step can anchor
252
- its in-place writes when accepted drafts leave the sequence at a
253
- non-block-aligned position.
254
- """
255
- if cache_config.mamba_cache_mode == "align":
256
- assert forward_context is not None
257
- assert mamba_state_copy_funcs is not None
258
- assert copy_bufs is not None
259
- num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens
260
- scheduled_spec_decode_tokens_dict = (
261
- scheduler_output.scheduled_spec_decode_tokens
262
- )
263
- num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu
264
- mamba_group_ids = copy_bufs.mamba_group_ids
265
- mamba_spec = copy_bufs.mamba_spec
266
- copy_bufs.offset = 0
267
- for i, req_id in enumerate(input_batch.req_ids):
268
- req_state = requests[req_id]
269
- num_computed_tokens = req_state.num_computed_tokens
270
- num_draft_tokens = len(scheduled_spec_decode_tokens_dict.get(req_id, []))
271
- num_scheduled_tokens = num_scheduled_tokens_dict[req_id]
272
- num_accepted_tokens = num_accepted_tokens_cpu[i]
273
- num_tokens_running_state = (
274
- num_computed_tokens + num_scheduled_tokens - num_draft_tokens
275
- )
276
- new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens - 1
277
- aligned_new_computed_tokens = (
278
- new_num_computed_tokens // mamba_spec.block_size * mamba_spec.block_size
279
- )
280
- # TODO: how to ensure all blocks that cache_blocks called are cached here?
281
- if aligned_new_computed_tokens >= num_tokens_running_state:
282
- accept_token_bias = (
283
- aligned_new_computed_tokens - num_tokens_running_state
284
- )
285
- src_block_idx = mamba_state_idx[req_id]
286
- dest_block_idx = (
287
- aligned_new_computed_tokens // mamba_spec.block_size - 1
288
- )
289
- collect_mamba_copy_meta(
290
- copy_bufs,
291
- kv_cache_config,
292
- mamba_state_copy_funcs,
293
- mamba_group_ids,
294
- src_block_idx,
295
- dest_block_idx,
296
- accept_token_bias,
297
- req_state,
298
- forward_context,
299
- )
300
- if src_block_idx == dest_block_idx:
301
- num_accepted_tokens_cpu[i] = 1
302
- do_mamba_copy_block(copy_bufs)
303
- elif cache_config.mamba_cache_mode == "all" and num_spec_tokens > 0:
304
- _, mamba_spec = get_mamba_groups(kv_cache_config)
305
- block_size = mamba_spec.block_size
306
- full_decode_len = 1 + num_spec_tokens
307
- scheduled = scheduler_output.num_scheduled_tokens
308
- for req_id in input_batch.req_ids[:num_reqs]:
309
- num_query = scheduled.get(req_id, 0)
310
- if num_query == full_decode_len:
311
- req = requests[req_id]
312
- seq_len = req.num_computed_tokens + num_query
313
- mamba_state_idx[req_id] = max(0, (seq_len - 1) // block_size)
314
- else:
315
- mamba_state_idx.pop(req_id, None)
720
+ if num_spec_tokens <= 0:
721
+ return
722
+ _, mamba_spec = get_mamba_groups(kv_cache_config)
723
+ block_size = mamba_spec.block_size
724
+ full_decode_len = 1 + num_spec_tokens
725
+ scheduled = scheduler_output.num_scheduled_tokens
726
+ for req_id in input_batch.req_ids[:num_reqs]:
727
+ num_query = scheduled.get(req_id, 0)
728
+ if num_query == full_decode_len:
729
+ req = requests[req_id]
730
+ seq_len = req.num_computed_tokens + num_query
731
+ mamba_state_idx[req_id] = max(0, (seq_len - 1) // block_size)
732
+ else:
733
+ mamba_state_idx.pop(req_id, None)
734
735
736
def preprocess_mamba_all_specdec(
@@ -328,3 +746,155 @@ def preprocess_mamba_all_specdec(
746
np_view[i] = mamba_state_idx.get(req_id, -1)
747
np_view[num_reqs:].fill(-1)
748
prev_last_scheduled_idx_buf.copy_to_gpu()
749
+
750
+
751
+def postprocess_mamba_align_gpu(
752
+ *,
753
+ bufs: "MambaBuffers",
754
+ num_reqs: int,
755
+ num_accepted_tokens_gpu: torch.Tensor,
756
+ num_accepted_tokens_cpu_tensor: torch.Tensor,
757
+ input_batch: GPUInputBatch,
758
+ kv_cache_config: KVCacheConfig,
759
+ forward_context: dict[str, Any],
760
+ mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
761
+) -> None:
762
+ """GPU-side mamba postprocess for spec decode + hybrid + align mode.
763
+
764
+ Lazily binds the fused-kernel context to the persistent block tables and
765
+ forward-context state pointers on the first call, runs the fused kernel,
766
+ and async-copies the per-request accepted-token counts back to the input
767
+ batch's CPU tensor for the next iteration's preprocess.
768
+ """
769
+ ctx = bufs.postprocess_align
770
+ # Caller is responsible for gating on spec decode + hybrid; this assert is
771
+ # a tripwire if those gates ever drift apart.
772
+ assert ctx is not None
773
+ assert ctx.mamba_state_idx_buf is not None
774
+ assert ctx.num_scheduled_tokens_buf is not None
775
+ assert ctx.num_computed_tokens_buf is not None
776
+ assert ctx.num_draft_tokens_buf is not None
777
+
778
+ if not ctx.is_initialized:
779
+ ctx.initialize_from_forward_context(
780
+ kv_cache_config,
781
+ forward_context,
782
+ mamba_state_copy_funcs,
783
+ [
784
+ input_batch.block_table[gid].get_device_tensor(num_reqs)
785
+ for gid in ctx.mamba_group_ids
786
+ ],
787
+ )
788
+
789
+ ctx.run_fused_postprocess(
790
+ num_reqs=num_reqs,
791
+ num_accepted_tokens_gpu=num_accepted_tokens_gpu,
792
+ mamba_state_idx_gpu=ctx.mamba_state_idx_buf.gpu,
793
+ num_scheduled_tokens_gpu=ctx.num_scheduled_tokens_buf.gpu,
794
+ num_computed_tokens_gpu=ctx.num_computed_tokens_buf.gpu,
795
+ num_draft_tokens_gpu=ctx.num_draft_tokens_buf.gpu,
796
+ )
797
+
798
+ # ``num_accepted_tokens_out`` is pre-initialized from
799
+ # ``num_accepted_tokens_gpu``; the kernel only overwrites entries to 1
800
+ # when src_block_idx == dest_block_idx (copy within the same block), so
801
+ # the original count is preserved for everyone else.
802
+ num_accepted_tokens_cpu_tensor[:num_reqs].copy_(
803
+ ctx.num_accepted_tokens_out[:num_reqs], non_blocking=True
804
+ )
805
+
806
+
807
+def stage_postprocess_metadata_to_gpu(
808
+ scheduler_output: SchedulerOutput,
809
+ req_ids: list[str],
810
+ num_reqs: int,
811
+ requests: dict[str, CachedRequestState],
812
+ num_scheduled_tokens_buf: CpuGpuBuffer,
813
+ num_computed_tokens_buf: CpuGpuBuffer,
814
+ num_draft_tokens_buf: CpuGpuBuffer,
815
+) -> None:
816
+ """Stage per-request postprocess metadata into GPU buffers (non-blocking).
817
+
818
+ Walks ``req_ids[:num_reqs]`` in batch order and writes each request's
819
+ scheduled/computed/draft token counts into the matching pinned numpy
820
+ views, then issues three non-blocking H→D copies. These values don't
821
+ change between ``_prepare_inputs`` and ``_update_states_after_model_execute``.
822
+ The fused postprocess kernel indexes the resulting GPU tensors
823
+ by ``req_idx``.
824
+ """
825
+ scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
826
+ num_scheduled = scheduler_output.num_scheduled_tokens
827
+ scheduled_np = num_scheduled_tokens_buf.np
828
+ computed_np = num_computed_tokens_buf.np
829
+ draft_np = num_draft_tokens_buf.np
830
+ for i in range(num_reqs):
831
+ req_id = req_ids[i]
832
+ scheduled_np[i] = num_scheduled[req_id]
833
+ computed_np[i] = requests[req_id].num_computed_tokens
834
+ draft_np[i] = len(scheduled_spec_tokens.get(req_id, []))
835
+ num_scheduled_tokens_buf.copy_to_gpu(num_reqs)
836
+ num_computed_tokens_buf.copy_to_gpu(num_reqs)
837
+ num_draft_tokens_buf.copy_to_gpu(num_reqs)
838
+
839
+
840
+def stage_mamba_state_idx_to_gpu(
841
+ mamba_state_idx: dict[str, int],
842
+ req_ids: list[str],
843
+ num_reqs: int,
844
+ gpu_buf: CpuGpuBuffer,
845
+) -> None:
846
+ """Materialize ``mamba_state_idx`` into ``gpu_buf`` and copy to GPU.
847
+
848
+ Walks ``req_ids[:num_reqs]`` in batch order, writing each request's block
849
+ index into the buffer's pinned numpy view, then issues a non-blocking H→D
850
+ copy. The fused kernel indexes the resulting GPU tensor by ``req_idx``.
851
+
852
+ Invariant: ``preprocess_mamba`` must have run first for the same batch so
853
+ that every ``req_ids[i]`` has an entry in ``mamba_state_idx``.
854
+ """
855
+ np_view = gpu_buf.np
856
+ for i in range(num_reqs):
857
+ req_id = req_ids[i]
858
+ state_idx = mamba_state_idx.get(req_id)
859
+ assert state_idx is not None, (
860
+ f"mamba_state_idx missing entry for {req_id!r}; "
861
+ "preprocess_mamba must run before stage_mamba_state_idx_to_gpu"
862
+ )
863
+ np_view[i] = state_idx
864
+ gpu_buf.copy_to_gpu(num_reqs)
865
+
866
+
867
+def stage_postprocess_inputs_to_gpu(
868
+ ctx: MambaSpecDecodeGPUContext,
869
+ scheduler_output: SchedulerOutput,
870
+ req_ids: list[str],
871
+ num_reqs: int,
872
+ requests: dict[str, CachedRequestState],
873
+ mamba_state_idx: dict[str, int],
874
+) -> None:
875
+ """Stage all per-request inputs the fused mamba postprocess kernel reads.
876
+
877
+ Bundles ``stage_mamba_state_idx_to_gpu`` and
878
+ ``stage_postprocess_metadata_to_gpu`` into a single call so the runner
879
+ has one entry point for postprocess staging. Buffers live on ``ctx``
880
+ and only exist when the postprocess kernel is enabled.
881
+ """
882
+ assert ctx.mamba_state_idx_buf is not None
883
+ assert ctx.num_scheduled_tokens_buf is not None
884
+ assert ctx.num_computed_tokens_buf is not None
885
+ assert ctx.num_draft_tokens_buf is not None
886
+ stage_mamba_state_idx_to_gpu(
887
+ mamba_state_idx,
888
+ req_ids,
889
+ num_reqs,
890
+ ctx.mamba_state_idx_buf,
891
+ )
892
+ stage_postprocess_metadata_to_gpu(
893
+ scheduler_output,
894
+ req_ids,
895
+ num_reqs,
896
+ requests,
897
+ ctx.num_scheduled_tokens_buf,
898
+ ctx.num_computed_tokens_buf,
899
+ ctx.num_draft_tokens_buf,
900
+ )