[Kernel][Mamba] Fused-kernel support for align-mode DS-conv state migration with num_accepted_tokens > 1 (#49291)
Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
sungsoo ha committed
Jul 29, 2026 at 03:54 UTC
f51193b9aefe665904d793ca782002320224d27a
5 files changed
+409
-64
tests/kernels/mamba/test_precopy_mamba_align.py
+257
-20
@@ -2,9 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Equivalence test for ``precopy_mamba_align_fused_kernel``.
4
5
-The V2 "align" pre-copy must migrate mamba state across block boundaries with
6
-byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
7
-``get_temporal_copy_spec``):
5
+The fused "align" pre-copy must migrate mamba state across block boundaries
6
+with byte-identical semantics to the scalar V1 copy specs
7
+(``get_conv_copy_spec`` / ``get_temporal_copy_spec``):
8
9
* conv state (SD layout, conv_width > 0): shift the sliding window by
10
``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` ->
@@ -14,20 +14,27 @@ byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
14
``state[bt[dst_col]]``.
15
16
The kernel must also no-op when ``src_col < 0`` (fresh request) or
17
-``src_col == dst_col`` (no boundary crossed).
17
+``src_col == dst_col`` (no boundary crossed). V2 callers pass an explicit
18
+``idx_mapping``; V1 align preprocessing launches in batch order with
19
+``idx_mapping=None``.
20
"""
21
22
from __future__ import annotations
23
24
+from types import SimpleNamespace
25
+
26
+import numpy as np
27
import torch
28
29
+from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils
30
from vllm.platforms import current_platform
31
+from vllm.v1.worker import mamba_utils as worker_mamba_utils
32
from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel
33
34
try:
35
import pytest
36
30
- pytestmark = pytest.mark.skipif(
37
+ _cuda_required = pytest.mark.skipif(
38
not current_platform.is_cuda(),
39
reason="precopy_mamba_align_fused_kernel needs CUDA/Triton",
40
)
@@ -35,6 +42,9 @@ try:
42
except ModuleNotFoundError: # allow running directly as ``python <thisfile>``
43
pytest = None
44
45
+ def _cuda_required(fn):
46
+ return fn
47
+
48
def _parametrize(_name, _values):
49
def _deco(fn):
50
return fn
@@ -49,13 +59,20 @@ SSM_SHAPE = (4, 16, 16)
59
MAX_COLS = 8
60
61
52
-def _build_state(num_blocks, device):
62
+def _build_state(num_blocks, device, conv_state_dim_first):
63
"""Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools."""
64
convs, ssms = [], []
65
for _ in range(NUM_LAYERS):
66
+ conv_shape = (
67
+ (num_blocks, CONV_DIM, CONV_WIDTH)
68
+ if conv_state_dim_first
69
+ else (num_blocks, CONV_WIDTH, CONV_DIM)
70
+ )
71
convs.append(
72
torch.randn(
58
- num_blocks, CONV_WIDTH, CONV_DIM, dtype=torch.bfloat16, device=device
73
+ *conv_shape,
74
+ dtype=torch.bfloat16,
75
+ device=device,
76
)
77
)
78
ssms.append(
@@ -64,7 +81,7 @@ def _build_state(num_blocks, device):
81
return convs, ssms
82
83
67
-def _build_meta(convs, ssms, device):
84
+def _build_meta(convs, ssms, device, conv_state_dim_first):
85
"""Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer."""
86
n = NUM_LAYERS * 2
87
base = torch.zeros(n, dtype=torch.int64, device=device)
@@ -82,8 +99,14 @@ def _build_meta(convs, ssms, device):
99
base[i] = conv.data_ptr()
100
blk_stride[i] = conv.stride(0) * conv.element_size()
101
elem[i] = conv.element_size()
85
- width[i] = conv.size(1)
86
- inner[i] = conv.stride(1)
102
+ if conv_state_dim_first:
103
+ width[i] = conv.size(2)
104
+ inner[i] = 1
105
+ drc[i] = conv.size(1)
106
+ drs[i] = conv.stride(1) * conv.element_size()
107
+ else:
108
+ width[i] = conv.size(1)
109
+ inner[i] = conv.stride(1)
110
i += 1
111
# ssm (temporal): width = 0, inner = elems per block
112
base[i] = ssm.data_ptr()
@@ -95,7 +118,7 @@ def _build_meta(convs, ssms, device):
118
return base, blk_stride, elem, inner, width, group, drc, drs
119
120
98
-def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
121
+def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs, conv_dim_first):
122
"""Apply the V1 copy semantics on clones, reading from the pre-copy state."""
123
conv_pre = [c.clone() for c in convs]
124
ssm_pre = [s.clone() for s in ssms]
@@ -108,14 +131,24 @@ def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
131
sblk, dblk = int(bt[r, sc]), int(bt[r, dc])
132
tblk = int(bt[r, sc + tb]) # temporal src column shifted by bias
133
for layer in range(NUM_LAYERS):
111
- conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
134
+ if conv_dim_first:
135
+ conv_ref[layer][dblk, :, : CONV_WIDTH - tb] = conv_pre[layer][
136
+ sblk, :, tb:
137
+ ]
138
+ else:
139
+ conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
140
ssm_ref[layer][dblk] = ssm_pre[layer][tblk]
141
return conv_ref, ssm_ref
142
143
144
+@_parametrize("conv_state_dim_first", [False, True])
145
@_parametrize("num_reqs", [1, 4, 16])
146
@_parametrize("token_bias", [0, 1, 2])
118
-def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
147
+@_parametrize("has_idx_mapping", [True, False])
148
+@_cuda_required
149
+def test_precopy_matches_v1_copy_specs(
150
+ num_reqs, token_bias, has_idx_mapping, conv_state_dim_first
151
+):
152
device = torch.device("cuda")
153
torch.manual_seed(0)
154
# Distinct physical block per (req, col) so copies never alias.
@@ -136,13 +169,20 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
169
if num_reqs >= 2:
170
dst_col[1] = 1 # src_col == dst_col -> no copy
171
139
- convs, ssms = _build_state(num_blocks, device)
172
+ convs, ssms = _build_state(num_blocks, device, conv_state_dim_first)
173
conv_ref, ssm_ref = _reference(
141
- convs, ssms, bt.cpu(), src_col.cpu(), dst_col.cpu(), bias.cpu(), num_reqs
174
+ convs,
175
+ ssms,
176
+ bt.cpu(),
177
+ src_col.cpu(),
178
+ dst_col.cpu(),
179
+ bias.cpu(),
180
+ num_reqs,
181
+ conv_state_dim_first,
182
)
183
184
base, blk_stride, elem, inner, width, group, drc, drs = _build_meta(
145
- convs, ssms, device
185
+ convs, ssms, device, conv_state_dim_first
186
)
187
bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
188
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
@@ -161,10 +201,11 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
201
group,
202
drc,
203
drs,
164
- idx_mapping,
204
+ idx_mapping if has_idx_mapping else None,
205
num_reqs,
206
COPY_BLOCK_SIZE=1024,
167
- CONV_STATE_DIM_FIRST=False,
207
+ CONV_STATE_DIM_FIRST=conv_state_dim_first,
208
+ HAS_IDX_MAPPING=has_idx_mapping,
209
)
210
torch.accelerator.synchronize()
211
@@ -173,8 +214,204 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
214
torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0)
215
216
217
+def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch):
218
+ monkeypatch.setattr(
219
+ layer_mamba_utils,
220
+ "is_conv_state_dim_first",
221
+ lambda: True,
222
+ )
223
+ state = torch.empty((2, CONV_DIM, CONV_WIDTH), dtype=torch.bfloat16)
224
+
225
+ with pytest.raises(AssertionError, match="num_accepted_tokens > 1"):
226
+ layer_mamba_utils.get_conv_copy_spec(
227
+ state=state,
228
+ block_ids=[0, 1],
229
+ cur_block_idx=0,
230
+ num_accepted_tokens=3,
231
+ )
232
+
233
+
234
+class _FakeCpuGpuBuffer:
235
+ def __init__(self, n):
236
+ self.np = np.zeros(n, dtype=np.int32)
237
+ self.gpu = object()
238
+ self.copy_sizes = []
239
+
240
+ def copy_to_gpu(self, n=None):
241
+ self.copy_sizes.append(n)
242
+ return self.gpu
243
+
244
+
245
+class _FakePrecopyContext:
246
+ def __init__(self, n):
247
+ self.is_initialized = True
248
+ self.mamba_group_ids = [0]
249
+ self.mamba_state_idx_buf = _FakeCpuGpuBuffer(n)
250
+ self.precopy_src_col_buf = _FakeCpuGpuBuffer(n)
251
+ self.precopy_token_bias_buf = _FakeCpuGpuBuffer(n)
252
+ self.calls = []
253
+
254
+ def initialize_from_forward_context(self, *args, **kwargs):
255
+ raise AssertionError("test context is pre-initialized")
256
+
257
+ def run_fused_precopy(
258
+ self,
259
+ *,
260
+ num_reqs,
261
+ state_idx_gpu,
262
+ src_col_gpu,
263
+ token_bias_gpu,
264
+ idx_mapping,
265
+ ):
266
+ self.calls.append(
267
+ {
268
+ "num_reqs": num_reqs,
269
+ "state_idx": self.mamba_state_idx_buf.np[:num_reqs].copy(),
270
+ "src_col": self.precopy_src_col_buf.np[:num_reqs].copy(),
271
+ "token_bias": self.precopy_token_bias_buf.np[:num_reqs].copy(),
272
+ "idx_mapping": idx_mapping,
273
+ }
274
+ )
275
+
276
+
277
+def _make_preprocess_case(token_bias):
278
+ req_ids = ["fresh", "same", "cross_a", "cross_b"]
279
+ scheduler_output = SimpleNamespace(
280
+ finished_req_ids=set(),
281
+ preempted_req_ids=set(),
282
+ scheduled_cached_reqs=SimpleNamespace(resumed_req_ids=set()),
283
+ num_scheduled_tokens={
284
+ "fresh": 1,
285
+ "same": 1,
286
+ "cross_a": 1,
287
+ "cross_b": 2,
288
+ },
289
+ )
290
+ input_batch = SimpleNamespace(
291
+ req_ids=req_ids,
292
+ num_accepted_tokens_cpu=np.array(
293
+ [token_bias + 1, token_bias + 1, token_bias + 1, 2],
294
+ dtype=np.int32,
295
+ ),
296
+ )
297
+ requests = {
298
+ "fresh": SimpleNamespace(req_id="fresh", num_computed_tokens=0),
299
+ "same": SimpleNamespace(req_id="same", num_computed_tokens=5),
300
+ "cross_a": SimpleNamespace(req_id="cross_a", num_computed_tokens=8),
301
+ "cross_b": SimpleNamespace(req_id="cross_b", num_computed_tokens=7),
302
+ }
303
+ mamba_state_idx = {"same": 1, "cross_a": 0, "cross_b": 1}
304
+ return scheduler_output, input_batch, requests, mamba_state_idx
305
+
306
+
307
+@_parametrize("token_bias", [1, 2])
308
+def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bias):
309
+ block_size = 4
310
+ mamba_spec = SimpleNamespace(block_size=block_size, num_speculative_blocks=1)
311
+ cache_config = SimpleNamespace(enable_prefix_caching=True)
312
+ kv_cache_config = SimpleNamespace()
313
+ scalar_copy_calls = []
314
+
315
+ def fake_collect(
316
+ copy_bufs,
317
+ kv_cache_config,
318
+ mamba_state_copy_funcs,
319
+ mamba_group_ids,
320
+ src_block_idx,
321
+ dest_block_idx,
322
+ accept_token_bias,
323
+ req_state,
324
+ forward_context,
325
+ ):
326
+ scalar_copy_calls.append(
327
+ (
328
+ req_state.req_id,
329
+ src_block_idx,
330
+ dest_block_idx,
331
+ accept_token_bias,
332
+ )
333
+ )
334
+
335
+ monkeypatch.setattr(worker_mamba_utils, "collect_mamba_copy_meta", fake_collect)
336
+ monkeypatch.setattr(
337
+ worker_mamba_utils, "do_mamba_copy_block", lambda copy_bufs: None
338
+ )
339
+
340
+ scalar_case = _make_preprocess_case(token_bias)
341
+ fused_case = _make_preprocess_case(token_bias)
342
+
343
+ scalar_copy_bufs = SimpleNamespace(
344
+ mamba_group_ids=[0],
345
+ mamba_spec=mamba_spec,
346
+ offset=0,
347
+ )
348
+ worker_mamba_utils.preprocess_mamba(
349
+ scheduler_output=scalar_case[0],
350
+ kv_cache_config=kv_cache_config,
351
+ cache_config=cache_config,
352
+ mamba_state_idx=scalar_case[3],
353
+ input_batch=scalar_case[1],
354
+ requests=scalar_case[2],
355
+ forward_context={},
356
+ mamba_state_copy_funcs=(),
357
+ copy_bufs=scalar_copy_bufs,
358
+ )
359
+
360
+ ctx = _FakePrecopyContext(len(fused_case[1].req_ids))
361
+ fused_copy_bufs = SimpleNamespace(
362
+ mamba_group_ids=[0],
363
+ mamba_spec=mamba_spec,
364
+ offset=0,
365
+ )
366
+ worker_mamba_utils.preprocess_mamba(
367
+ scheduler_output=fused_case[0],
368
+ kv_cache_config=kv_cache_config,
369
+ cache_config=cache_config,
370
+ mamba_state_idx=fused_case[3],
371
+ input_batch=fused_case[1],
372
+ requests=fused_case[2],
373
+ forward_context={},
374
+ mamba_state_copy_funcs=(),
375
+ copy_bufs=fused_copy_bufs,
376
+ align_ctx=ctx,
377
+ )
378
+
379
+ assert fused_case[3] == scalar_case[3]
380
+ np.testing.assert_array_equal(
381
+ fused_case[1].num_accepted_tokens_cpu,
382
+ scalar_case[1].num_accepted_tokens_cpu,
383
+ )
384
+ assert scalar_copy_calls == [
385
+ ("cross_a", 0, 2, token_bias),
386
+ ("cross_b", 1, 2, 1),
387
+ ]
388
+ assert len(ctx.calls) == 1
389
+ call = ctx.calls[0]
390
+ assert call["num_reqs"] == len(fused_case[1].req_ids)
391
+ assert call["idx_mapping"] is None
392
+ np.testing.assert_array_equal(call["state_idx"], np.array([0, 1, 2, 2]))
393
+ np.testing.assert_array_equal(call["src_col"], np.array([-1, -1, 0, 1]))
394
+ np.testing.assert_array_equal(call["token_bias"], np.array([0, 0, token_bias, 1]))
395
+ fused_copy_calls = [
396
+ (req_id, int(src), int(dst), int(bias))
397
+ for req_id, src, dst, bias in zip(
398
+ fused_case[1].req_ids,
399
+ call["src_col"],
400
+ call["state_idx"],
401
+ call["token_bias"],
402
+ )
403
+ if int(src) != -1 and int(src) != int(dst)
404
+ ]
405
+ assert fused_copy_calls == scalar_copy_calls
406
+
407
+
408
if __name__ == "__main__":
409
for nr in (1, 4, 16):
410
for tb in (0, 1, 2):
179
- test_precopy_matches_v1_copy_specs(nr, tb)
180
- print(f"OK num_reqs={nr} token_bias={tb}")
411
+ for mapping in (True, False):
412
+ for dim_first in (False, True):
413
+ test_precopy_matches_v1_copy_specs(nr, tb, mapping, dim_first)
414
+ print(
415
+ f"OK num_reqs={nr} token_bias={tb} "
416
+ f"has_idx_mapping={mapping} conv_dim_first={dim_first}"
417
+ )
tests/v1/e2e/general/test_mamba_prefix_cache.py
+38
-6
@@ -338,6 +338,30 @@ def get_fake_process_mamba_fn(
338
assert copy_info[0][-1] == expected_temporal_src
339
assert copy_info[1][-1] == expected_temporal_dest
340
341
+ def check_fused_copy_info(
342
+ action: tuple[int, int],
343
+ align_ctx: mamba_utils.MambaSpecDecodeGPUContext,
344
+ ):
345
+ # Align + spec-decode on a hybrid model routes the pre-copy through the
346
+ # fused kernel (preprocess_mamba -> run_fused_precopy) instead of
347
+ # do_mamba_copy_block, so copy_info is never populated. Verify from the
348
+ # fused buffers (req-0 scope, mirroring check_copy_info):
349
+ # - the copy DECISION: src_col is -1 iff no pre-copy is scheduled;
350
+ # - the DESTINATION column: state_idx == action[1] (curr_state_idx;
351
+ # maps directly through block_ids, exactly as check_copy_info's dst).
352
+ # The source column is NOT asserted here: on the scalar path the source
353
+ # address is produced by the per-state copy func with an accept-token
354
+ # bias offset (collect_mamba_copy_meta), so prev_state_idx does not map
355
+ # to action[0] by plain equality. Source block-level exactness (incl.
356
+ # the accept-bias) is covered by test_precopy_mamba_align.py.
357
+ src_col = int(align_ctx.precopy_src_col_buf.np[0])
358
+ state_idx = int(align_ctx.mamba_state_idx_buf.np[0])
359
+ if action == (-1, -1):
360
+ assert src_col == -1
361
+ else:
362
+ assert src_col != -1
363
+ assert state_idx == action[1]
364
+
365
def fake_preprocess_mamba_fn(
366
scheduler_output: SchedulerOutput,
367
kv_cache_config: KVCacheConfig,
@@ -348,6 +372,7 @@ def get_fake_process_mamba_fn(
372
forward_context: dict[str, Any],
373
mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
374
copy_bufs: mamba_utils.MambaCopyBuffers,
375
+ align_ctx: mamba_utils.MambaSpecDecodeGPUContext | None = None,
376
):
377
nonlocal copy_info
378
copy_info = None
@@ -361,14 +386,21 @@ def get_fake_process_mamba_fn(
386
forward_context,
387
mamba_state_copy_funcs,
388
copy_bufs,
389
+ align_ctx,
390
)
391
if cur_step_action is not None:
366
- check_copy_info(
367
- cur_step_action.preprocess_copy_idx,
368
- kv_cache_config,
369
- forward_context,
370
- input_batch,
371
- )
392
+ if align_ctx is not None:
393
+ check_fused_copy_info(
394
+ cur_step_action.preprocess_copy_idx,
395
+ align_ctx,
396
+ )
397
+ else:
398
+ check_copy_info(
399
+ cur_step_action.preprocess_copy_idx,
400
+ kv_cache_config,
401
+ forward_context,
402
+ input_batch,
403
+ )
404
return ret
405
406
def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+5
-12
@@ -9,10 +9,6 @@ import torch.nn as nn
9
10
from vllm.config import VllmConfig
11
from vllm.config.compilation import CUDAGraphMode
12
-from vllm.model_executor.layers.mamba.mamba_utils import (
13
- get_conv_copy_spec,
14
- is_conv_state_dim_first,
15
-)
12
from vllm.triton_utils import tl, triton
13
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
14
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
@@ -133,14 +129,11 @@ class MambaHybridModelState(DefaultModelState):
129
) -> MambaSpecDecodeGPUContext:
130
if self._mamba_ctx is None:
131
copy_funcs = self.model.get_mamba_state_copy_func()
136
- # The fused copy kernels shift conv windows assuming the SD layout;
137
- # the DS layout cannot express a >0 spec-decode shift as a single
138
- # contiguous copy (mirrors get_conv_copy_spec's NotImplementedError).
139
- if get_conv_copy_spec in copy_funcs and is_conv_state_dim_first():
140
- assert self.vllm_config.speculative_config is None, (
141
- "DS conv state layout does not support mamba align state "
142
- "copies with speculative decoding"
143
- )
132
+ # Both SD and DS conv layouts support a >0 spec-decode shift: the
133
+ # fused pre-copy kernel (``_copy_mamba_state_block``) applies the
134
+ # ``token_bias = num_accepted - 1`` window shift per conv layout
135
+ # (SD: contiguous slice; DS: per-dim-row strided slice), matching
136
+ # the V1 ``get_conv_copy_spec`` semantics.
137
self._mamba_ctx = MambaSpecDecodeGPUContext.create(
138
max_num_reqs=self.max_num_reqs,
139
kv_cache_config=kv_cache_config,
vllm/v1/worker/gpu_model_runner.py
+1
@@ -4335,6 +4335,7 @@ class GPUModelRunner(
4335
self.compilation_config.static_forward_context,
4336
self.model.get_mamba_state_copy_func(),
4337
mamba_bufs.preprocess,
4338
+ align_ctx=mamba_bufs.postprocess_align,
4339
)
4340
# preprocess_mamba resets num_accepted_tokens_cpu to 1
4341
# for requests whose state was copied to a new block.
vllm/v1/worker/mamba_utils.py
+108
-26
@@ -3,7 +3,7 @@
3
import dataclasses
4
import itertools
5
from collections.abc import Callable
6
-from typing import Any
6
+from typing import Any, NamedTuple
7
8
import torch
9
@@ -348,8 +348,9 @@ def precopy_mamba_align_fused_kernel(
348
num_reqs,
349
COPY_BLOCK_SIZE: tl.constexpr,
350
CONV_STATE_DIM_FIRST: tl.constexpr,
351
+ HAS_IDX_MAPPING: tl.constexpr = True,
352
):
352
- """Pre-copy mamba "align" state across block boundaries on the V2 runner.
353
+ """Pre-copy mamba "align" state across block boundaries.
354
355
Before the forward pass, copy each request's last SSM/conv state from its
356
previous block column into the new window block column, so the kernels read
@@ -359,16 +360,20 @@ def precopy_mamba_align_fused_kernel(
360
copy specs), but driven by the GPU-resident src columns so it needs no
361
CPU-GPU sync (async-scheduling safe).
362
362
- Grid: (num_reqs, num_layers * num_state_types); block tables are indexed by
363
- batch row, per-request state by req_idx via idx_mapping (V2 layout).
363
+ Grid: (num_reqs, num_layers * num_state_types). V2 passes a batch-to-state
364
+ idx_mapping; V1 already stores the staged arrays in batch order and uses
365
+ HAS_IDX_MAPPING=False.
366
"""
367
batch_idx = tl.program_id(0)
368
state_idx = tl.program_id(1)
369
if batch_idx >= num_reqs:
370
return
369
- req_idx = tl.load(idx_mapping_ptr + batch_idx)
370
- if req_idx < 0:
371
- return
371
+ if HAS_IDX_MAPPING:
372
+ req_idx = tl.load(idx_mapping_ptr + batch_idx)
373
+ if req_idx < 0:
374
+ return
375
+ else:
376
+ req_idx = batch_idx
377
378
src_col = tl.load(src_col_ptr + req_idx)
379
dst_col = tl.load(mamba_state_idx_ptr + req_idx)
@@ -527,6 +532,8 @@ class MambaSpecDecodeGPUContext:
532
num_scheduled_tokens_buf: CpuGpuBuffer | None = None
533
num_computed_tokens_buf: CpuGpuBuffer | None = None
534
num_draft_tokens_buf: CpuGpuBuffer | None = None
535
+ precopy_src_col_buf: CpuGpuBuffer | None = None
536
+ precopy_token_bias_buf: CpuGpuBuffer | None = None
537
538
# Flag to track if metadata has been populated
539
is_initialized: bool = False
@@ -590,6 +597,8 @@ class MambaSpecDecodeGPUContext:
597
num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
598
num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
599
num_draft_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
600
+ precopy_src_col_buf=make_buffer(max_num_reqs, dtype=torch.int32),
601
+ precopy_token_bias_buf=make_buffer(max_num_reqs, dtype=torch.int32),
602
is_initialized=False,
603
)
604
@@ -797,17 +806,18 @@ class MambaSpecDecodeGPUContext:
806
state_idx_gpu: torch.Tensor,
807
src_col_gpu: torch.Tensor,
808
token_bias_gpu: torch.Tensor,
800
- idx_mapping: torch.Tensor,
809
+ idx_mapping: torch.Tensor | None,
810
) -> None:
811
"""Pre-copy each request's previous running block into its new window
803
- block before the forward pass (V2 align boundary migration).
812
+ block before the forward pass (align boundary migration).
813
814
Args:
815
num_reqs: Number of active requests (batch order).
816
state_idx_gpu: [max_reqs] post-advance dst block column per req slot.
817
src_col_gpu: [max_reqs] pre-advance src block column (-1 = fresh).
818
token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1).
810
- idx_mapping: [num_reqs] batch_idx -> req_state_idx (-1 to skip).
819
+ idx_mapping: optional [num_reqs] batch_idx -> req_state_idx.
820
+ None means V1 batch order already equals request state order.
821
"""
822
if num_reqs == 0 or not self.is_initialized:
823
return
@@ -831,6 +841,7 @@ class MambaSpecDecodeGPUContext:
841
num_reqs,
842
COPY_BLOCK_SIZE=1024,
843
CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
844
+ HAS_IDX_MAPPING=idx_mapping is not None,
845
)
846
847
def run_fused_postprocess_align(
@@ -989,6 +1000,36 @@ def cleanup_mamba_state_idx(
1000
mamba_state_idx.pop(req_id, None)
1001
1002
1003
+class _FusedPrecopy(NamedTuple):
1004
+ """Resolved fused align pre-copy resources (all non-None once resolved)."""
1005
+
1006
+ ctx: "MambaSpecDecodeGPUContext"
1007
+ state_idx: CpuGpuBuffer
1008
+ src_col: CpuGpuBuffer
1009
+ token_bias: CpuGpuBuffer
1010
+
1011
+
1012
+def _resolve_fused_precopy(
1013
+ align_ctx: "MambaSpecDecodeGPUContext | None",
1014
+) -> _FusedPrecopy | None:
1015
+ """Bundle the fused-path buffers, or None for the scalar path.
1016
+
1017
+ Returning one non-None bundle lets callers narrow all four members with a
1018
+ single ``is not None`` check instead of re-asserting each buffer per use.
1019
+ """
1020
+ if align_ctx is None:
1021
+ return None
1022
+ assert align_ctx.mamba_state_idx_buf is not None
1023
+ assert align_ctx.precopy_src_col_buf is not None
1024
+ assert align_ctx.precopy_token_bias_buf is not None
1025
+ return _FusedPrecopy(
1026
+ align_ctx,
1027
+ align_ctx.mamba_state_idx_buf,
1028
+ align_ctx.precopy_src_col_buf,
1029
+ align_ctx.precopy_token_bias_buf,
1030
+ )
1031
+
1032
+
1033
def preprocess_mamba(
1034
scheduler_output: SchedulerOutput,
1035
kv_cache_config: KVCacheConfig,
@@ -999,11 +1040,13 @@ def preprocess_mamba(
1040
forward_context: dict[str, Any],
1041
mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
1042
copy_bufs: MambaCopyBuffers,
1043
+ align_ctx: MambaSpecDecodeGPUContext | None = None,
1044
):
1045
"""
1046
Copy the mamba state of previous step to the last
1047
(1 + num_speculative_blocks) block.
1048
"""
1049
+ fused = _resolve_fused_precopy(align_ctx)
1050
mamba_group_ids = copy_bufs.mamba_group_ids
1051
mamba_spec = copy_bufs.mamba_spec
1052
num_speculative_blocks = mamba_spec.num_speculative_blocks
@@ -1013,20 +1056,37 @@ def preprocess_mamba(
1056
cleanup_mamba_state_idx(scheduler_output, mamba_state_idx)
1057
1058
copy_bufs.offset = 0
1059
+ num_reqs = len(input_batch.req_ids)
1060
+
1061
+ if fused is not None:
1062
+ if num_reqs == 0:
1063
+ return
1064
+ if not fused.ctx.is_initialized:
1065
+ fused.ctx.initialize_from_forward_context(
1066
+ kv_cache_config,
1067
+ forward_context,
1068
+ mamba_state_copy_funcs,
1069
+ [
1070
+ input_batch.block_table[gid].get_device_tensor(num_reqs)
1071
+ for gid in fused.ctx.mamba_group_ids
1072
+ ],
1073
+ )
1074
+
1075
+ fused.src_col.np[:num_reqs] = -1
1076
+ fused.token_bias.np[:num_reqs] = 0
1077
+
1078
for i, req_id in enumerate(input_batch.req_ids):
1079
req_state = requests[req_id]
1080
prev_state_idx = mamba_state_idx.get(req_id)
1081
if prev_state_idx is None:
1020
- # new / resumed request, no previous state
1021
- # if num_computed_tokens is 0, prev_state_idx will be -1
1082
+ # New / resumed request; num_computed_tokens == 0 gives -1.
1083
prev_state_idx = (req_state.num_computed_tokens - 1) // block_size
1084
1085
num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
1025
- num_blocks: int = (
1086
+ num_blocks = (
1087
cdiv(req_state.num_computed_tokens + num_scheduled_tokens, block_size)
1088
+ num_speculative_blocks
1089
)
1029
-
1090
# We always save the current running state at the last
1091
# (1 + num_speculative_blocks) block.
1092
# A corner case worth mention here: assume we have block_size = 4 and
@@ -1039,20 +1099,42 @@ def preprocess_mamba(
1099
# And use block 1 to save the running state.
1100
curr_state_idx = num_blocks - 1 - num_speculative_blocks
1101
mamba_state_idx[req_id] = curr_state_idx
1102
+ if fused is not None:
1103
+ fused.state_idx.np[i] = curr_state_idx
1104
+
1105
if prev_state_idx != -1 and prev_state_idx != curr_state_idx:
1043
- collect_mamba_copy_meta(
1044
- copy_bufs,
1045
- kv_cache_config,
1046
- mamba_state_copy_funcs,
1047
- mamba_group_ids,
1048
- prev_state_idx,
1049
- curr_state_idx,
1050
- input_batch.num_accepted_tokens_cpu[i] - 1,
1051
- req_state,
1052
- forward_context,
1053
- )
1106
+ accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1
1107
+ if fused is not None:
1108
+ assert accept_token_bias >= 0
1109
+ fused.src_col.np[i] = prev_state_idx
1110
+ fused.token_bias.np[i] = accept_token_bias
1111
+ else:
1112
+ collect_mamba_copy_meta(
1113
+ copy_bufs,
1114
+ kv_cache_config,
1115
+ mamba_state_copy_funcs,
1116
+ mamba_group_ids,
1117
+ prev_state_idx,
1118
+ curr_state_idx,
1119
+ accept_token_bias,
1120
+ req_state,
1121
+ forward_context,
1122
+ )
1123
input_batch.num_accepted_tokens_cpu[i] = 1
1055
- do_mamba_copy_block(copy_bufs)
1124
+
1125
+ if fused is not None:
1126
+ fused.state_idx.copy_to_gpu(num_reqs)
1127
+ fused.src_col.copy_to_gpu(num_reqs)
1128
+ fused.token_bias.copy_to_gpu(num_reqs)
1129
+ fused.ctx.run_fused_precopy(
1130
+ num_reqs=num_reqs,
1131
+ state_idx_gpu=fused.state_idx.gpu,
1132
+ src_col_gpu=fused.src_col.gpu,
1133
+ token_bias_gpu=fused.token_bias.gpu,
1134
+ idx_mapping=None,
1135
+ )
1136
+ else:
1137
+ do_mamba_copy_block(copy_bufs)
1138
1139
1140
def postprocess_mamba_all(