[Perf][Hybrid] 3D-grid tiling of the state-copy Triton kernels (#49436)
Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com> Co-authored-by: Claude <noreply@anthropic.com>
Francesco Fusco committed
Aug 10, 2026 at 17:02 UTC
fac808b36f502d0d992509a187dca94c68b0360a
3 files changed
+347
-68
tests/kernels/mamba/test_memcpy_u64_tiled.py
new
+147
@@ -0,0 +1,147 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit test for ``_memcpy_u64_tiled``.
4
+
5
+Guards the head/body/tail alignment handling and the u64 body's tile
6
+partitioning used by ``postprocess_mamba_fused_kernel`` and
7
+``precopy_mamba_align_fused_kernel``. The device function must be
8
+byte-identical to a ``dst[dst_off:dst_off+copy_size] = src[src_off:src_off+
9
+copy_size]`` slice copy for every combination of:
10
+
11
+* ``copy_size``: the degenerate 0-byte case, sub-8B (head-only), 8B-aligned
12
+ bodies, and a multi-tile body that spans several ``COPY_BLOCK_SIZE``
13
+ iterations.
14
+* ``(src_off, dst_off)`` pairs: aligned, shared sub-8B misalignment (fast
15
+ path exercising head/tail masking), and mismatched src/dst alignment
16
+ (byte-fallback path). The kernel branches on ``(src ^ dst) & 7``, so a
17
+ full 5x5 product would re-exercise the same two codepaths.
18
+* ``NUM_TILES``: 1 (SD conv callsite, single-CTA memcpy) and
19
+ ``_TEMPORAL_TILES`` (temporal callsite, u64 range partitioned across
20
+ CTAs).
21
+* ``COPY_BLOCK_SIZE``: the production value 1024 (single-tile at these
22
+ sizes) and a small value 8 so the existing ``copy_size`` cases actually
23
+ cross tile boundaries under ``NUM_TILES=_TEMPORAL_TILES``. The
24
+ partitioning math is COPY_BLOCK_SIZE-agnostic, so testing at a small
25
+ value covers boundary arithmetic for all values.
26
+"""
27
+
28
+from __future__ import annotations
29
+
30
+import torch
31
+
32
+from vllm.platforms import current_platform
33
+from vllm.triton_utils import tl, triton
34
+from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, _memcpy_u64_tiled
35
+
36
+try:
37
+ import pytest
38
+
39
+ pytestmark = pytest.mark.skipif(
40
+ not current_platform.is_cuda(),
41
+ reason="_memcpy_u64_tiled needs CUDA/Triton",
42
+ )
43
+ _parametrize = pytest.mark.parametrize
44
+except ModuleNotFoundError: # allow running directly as ``python <thisfile>``
45
+ pytest = None
46
+
47
+ def _parametrize(_name, _values):
48
+ def _deco(fn):
49
+ return fn
50
+
51
+ return _deco
52
+
53
+
54
+@triton.jit
55
+def _memcpy_wrapper_kernel(
56
+ src_ptr,
57
+ dst_ptr,
58
+ copy_size,
59
+ src_off,
60
+ dst_off,
61
+ COPY_BLOCK_SIZE: tl.constexpr,
62
+ NUM_TILES: tl.constexpr,
63
+):
64
+ """Thin harness: one CTA per tile, exercising the device function."""
65
+ tile_idx = tl.program_id(0)
66
+ src_addr = src_ptr.to(tl.int64) + src_off
67
+ dst_addr = dst_ptr.to(tl.int64) + dst_off
68
+ _memcpy_u64_tiled(
69
+ src_addr,
70
+ dst_addr,
71
+ copy_size,
72
+ tile_idx,
73
+ COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
74
+ NUM_TILES=NUM_TILES,
75
+ )
76
+
77
+
78
+# Copy sizes: 0/1/7 (all-head), 8 (all-body when dst-aligned), 15 (body+tail,
79
+# tail-mask path), 16 (small tiled body), 1024 and 4 KiB (spans all 16 tiles
80
+# at COPY_BLOCK_SIZE=8, with 4 KiB adding multi-iteration-per-tile coverage;
81
+# both single-tile at COPY_BLOCK_SIZE=1024).
82
+_COPY_SIZES = [0, 1, 7, 8, 15, 16, 1024, 4 * 1024]
83
+# (src_off, dst_off) pairs. Torch tensors are 256B-aligned at data_ptr, so
84
+# slicing by these bytes yields a controlled sub-8B alignment. The kernel
85
+# branches on ``(src ^ dst) & 7``, so we cover both sides plus the aligned
86
+# baseline; the full 5x5 product added no coverage.
87
+_ALIGN_PAIRS = [
88
+ (0, 0), # fully aligned; head/tail masked out
89
+ (3, 3), # shared misalignment; fast path with head_bytes=5
90
+ (7, 7), # shared, extreme; head_bytes=1 with non-empty body
91
+ (1, 3), # mismatched; byte load/store fallback
92
+ (3, 1), # mismatched, opposite direction
93
+]
94
+_MAX_ALIGN_OFF = max(o for pair in _ALIGN_PAIRS for o in pair)
95
+# (NUM_TILES, COPY_BLOCK_SIZE) configs. NUM_TILES=1 is the SD conv callsite
96
+# (single-CTA memcpy) where tile partitioning collapses, so COPY_BLOCK_SIZE
97
+# is not observable — one value suffices. NUM_TILES=_TEMPORAL_TILES is the
98
+# temporal callsite: 1024 matches the production launch; 8 shrinks per-tile
99
+# rounding so ``_COPY_SIZES``' 1024/4096 cases cross tile boundaries
100
+# end-to-end (partitioning math is COPY_BLOCK_SIZE-agnostic, so a small
101
+# value proves it for all values).
102
+_TILE_BLOCK_CONFIGS = [
103
+ (1, 1024),
104
+ (_TEMPORAL_TILES, 1024),
105
+ (_TEMPORAL_TILES, 8),
106
+]
107
+
108
+
109
+@_parametrize("copy_size", _COPY_SIZES)
110
+@_parametrize("src_off,dst_off", _ALIGN_PAIRS)
111
+@_parametrize("num_tiles,copy_block_size", _TILE_BLOCK_CONFIGS)
112
+def test_memcpy_u64_tiled_matches_slice_copy(
113
+ copy_size, src_off, dst_off, num_tiles, copy_block_size
114
+):
115
+ device = torch.device("cuda")
116
+ torch.manual_seed(0)
117
+
118
+ # Pad by the max possible offset so the reference slice is in-bounds.
119
+ slack = _MAX_ALIGN_OFF + 8
120
+ src = torch.randint(0, 256, (copy_size + slack,), dtype=torch.uint8, device=device)
121
+ # Start dst from a distinct random pattern so the "unchanged region"
122
+ # check catches any accidental out-of-range writes.
123
+ dst = torch.randint(0, 256, (copy_size + slack,), dtype=torch.uint8, device=device)
124
+ ref = dst.clone()
125
+ ref[dst_off : dst_off + copy_size] = src[src_off : src_off + copy_size]
126
+
127
+ _memcpy_wrapper_kernel[(num_tiles,)](
128
+ src,
129
+ dst,
130
+ copy_size,
131
+ src_off,
132
+ dst_off,
133
+ COPY_BLOCK_SIZE=copy_block_size,
134
+ NUM_TILES=num_tiles,
135
+ )
136
+ torch.accelerator.synchronize()
137
+
138
+ # Byte-identical: the memcpy is dtype-agnostic; any diff is a bug.
139
+ torch.testing.assert_close(dst, ref, rtol=0, atol=0)
140
+
141
+
142
+if __name__ == "__main__":
143
+ for cs in _COPY_SIZES:
144
+ for so, do in _ALIGN_PAIRS:
145
+ for nt, cbs in _TILE_BLOCK_CONFIGS:
146
+ test_memcpy_u64_tiled_matches_slice_copy(cs, so, do, nt, cbs)
147
+ print("ok")
tests/kernels/mamba/test_precopy_mamba_align.py
+24
-12
@@ -31,7 +31,7 @@ import torch
31
from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils
32
from vllm.platforms import current_platform
33
from vllm.v1.worker import mamba_utils as worker_mamba_utils
34
-from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel
34
+from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, precopy_mamba_align_fused_kernel
35
36
_parametrize: Callable[..., Callable[[Any], Any]]
37
@@ -150,9 +150,14 @@ def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs, conv_dim_first
150
@_parametrize("num_reqs", [1, 4, 16])
151
@_parametrize("token_bias", [0, 1, 2])
152
@_parametrize("has_idx_mapping", [True, False])
153
+@_parametrize("temporal_tiles", [1, _TEMPORAL_TILES])
154
@_cuda_required
155
def test_precopy_matches_v1_copy_specs(
155
- num_reqs, token_bias, has_idx_mapping, conv_state_dim_first
156
+ num_reqs,
157
+ token_bias,
158
+ has_idx_mapping,
159
+ conv_state_dim_first,
160
+ temporal_tiles,
161
):
162
device = torch.device("cuda")
163
torch.manual_seed(0)
@@ -191,7 +196,7 @@ def test_precopy_matches_v1_copy_specs(
196
)
197
bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
198
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
194
- grid = (num_reqs, NUM_LAYERS * 2)
199
+ grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles)
200
precopy_mamba_align_fused_kernel[grid](
201
dst_col,
202
src_col,
@@ -211,6 +216,7 @@ def test_precopy_matches_v1_copy_specs(
216
COPY_BLOCK_SIZE=1024,
217
CONV_STATE_DIM_FIRST=conv_state_dim_first,
218
HAS_IDX_MAPPING=has_idx_mapping,
219
+ TEMPORAL_TILES=temporal_tiles,
220
)
221
torch.accelerator.synchronize()
222
@@ -413,12 +419,18 @@ def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bi
419
420
421
if __name__ == "__main__":
416
- for nr in (1, 4, 16):
417
- for tb in (0, 1, 2):
418
- for mapping in (True, False):
419
- for dim_first in (False, True):
420
- test_precopy_matches_v1_copy_specs(nr, tb, mapping, dim_first)
421
- print(
422
- f"OK num_reqs={nr} token_bias={tb} "
423
- f"has_idx_mapping={mapping} conv_dim_first={dim_first}"
424
- )
422
+ from itertools import product
423
+
424
+ _CASES = product(
425
+ (1, 4, 16), # num_reqs
426
+ (0, 1, 2), # token_bias
427
+ (True, False), # has_idx_mapping
428
+ (False, True), # conv_state_dim_first
429
+ (1, _TEMPORAL_TILES), # temporal_tiles
430
+ )
431
+ for nr, tb, mapping, dim_first, tt in _CASES:
432
+ test_precopy_matches_v1_copy_specs(nr, tb, mapping, dim_first, tt)
433
+ print(
434
+ f"OK num_reqs={nr} token_bias={tb} has_idx_mapping={mapping} "
435
+ f"conv_dim_first={dim_first} temporal_tiles={tt}"
436
+ )
vllm/v1/worker/mamba_utils.py
+176
-56
@@ -8,6 +8,7 @@ from typing import Any, NamedTuple
8
import torch
9
10
from vllm.config import CacheConfig
11
+from vllm.logger import init_logger
12
from vllm.model_executor.layers.mamba.mamba_utils import (
13
MambaStateCopyFunc,
14
get_conv_copy_spec,
@@ -22,6 +23,95 @@ from vllm.v1.utils import CpuGpuBuffer
23
from vllm.v1.worker.gpu_input_batch import CachedRequestState
24
from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
25
26
+logger = init_logger(__name__)
27
+
28
+# 16 saturates HBM on H100/GB200 across the reqs=8..128 range in
29
+# microbenchmarks
30
+_TEMPORAL_TILES = 16
31
+
32
+
33
+@triton.jit
34
+def _memcpy_u64_tiled(
35
+ src_addr,
36
+ dst_addr,
37
+ copy_size,
38
+ tile_idx,
39
+ COPY_BLOCK_SIZE: tl.constexpr,
40
+ NUM_TILES: tl.constexpr,
41
+):
42
+ """Head/body/tail memcpy with the u64 body split across ``NUM_TILES`` CTAs.
43
+
44
+ Fast path (``src`` and ``dst`` share sub-8B alignment): tile 0 owns the
45
+ byte head that lifts dst to 8B and the 0-7 byte tail; body tiles vectorize
46
+ as u64 over the aligned interior. ``NUM_TILES=1`` collapses to a single-
47
+ CTA memcpy. Production callers derive both addresses from the same
48
+ ``state_base_addr + block_id * stride`` and always take this path.
49
+
50
+ Slow path (mismatched sub-8B alignment): byte-wide tiled copy. Some
51
+ NVIDIA parts (e.g. GB200) reject misaligned u64 loads with
52
+ ``cudaErrorMisalignedAddress`` instead of accepting the misaligned-sector
53
+ cost, so we can't just let the fast path run with a misaligned src.
54
+ """
55
+ src_addr_i = src_addr.to(tl.int64)
56
+ dst_addr_i = dst_addr.to(tl.int64)
57
+
58
+ if ((src_addr_i ^ dst_addr_i) & 7) == 0:
59
+ head_bytes = tl.minimum(((-dst_addr_i) & 7).to(tl.int64), copy_size)
60
+ if tile_idx == 0:
61
+ head_off = tl.arange(0, 8)
62
+ head_mask = head_off < head_bytes
63
+ head_src = src_addr.to(tl.pointer_type(tl.uint8))
64
+ head_dst = dst_addr.to(tl.pointer_type(tl.uint8))
65
+ tl.store(
66
+ head_dst + head_off,
67
+ tl.load(head_src + head_off, mask=head_mask),
68
+ mask=head_mask,
69
+ )
70
+
71
+ # Body: u64 tiled. Rounding per_tile up to COPY_BLOCK_SIZE keeps every
72
+ # inner-loop iteration full-width vectorized; only the last non-empty
73
+ # tile can be masked. Late tiles fall off the end and iterate zero
74
+ # times.
75
+ body_bytes = copy_size - head_bytes
76
+ body_u64 = body_bytes // 8
77
+ per_tile_u64_raw = tl.cdiv(body_u64, NUM_TILES)
78
+ per_tile_u64 = tl.cdiv(per_tile_u64_raw, COPY_BLOCK_SIZE) * COPY_BLOCK_SIZE
79
+ tile_start = tile_idx.to(tl.int64) * per_tile_u64
80
+ tile_end = tl.minimum(tile_start + per_tile_u64, body_u64)
81
+
82
+ src_body_u64 = (src_addr + head_bytes).to(tl.pointer_type(tl.uint64))
83
+ dst_body_u64 = (dst_addr + head_bytes).to(tl.pointer_type(tl.uint64))
84
+ offsets = tl.arange(0, COPY_BLOCK_SIZE)
85
+ for i in range(tile_start, tile_end, COPY_BLOCK_SIZE):
86
+ mask = (i + offsets) < tile_end
87
+ data = tl.load(src_body_u64 + i + offsets, mask=mask)
88
+ tl.store(dst_body_u64 + i + offsets, data, mask=mask)
89
+
90
+ if tile_idx == 0:
91
+ tail_start = head_bytes + body_u64 * 8
92
+ tail_bytes = copy_size - tail_start
93
+ tail_off = tl.arange(0, 8)
94
+ tail_src = (src_addr + tail_start).to(tl.pointer_type(tl.uint8))
95
+ tail_dst = (dst_addr + tail_start).to(tl.pointer_type(tl.uint8))
96
+ tail_mask = tail_off < tail_bytes
97
+ tl.store(
98
+ tail_dst + tail_off,
99
+ tl.load(tail_src + tail_off, mask=tail_mask),
100
+ mask=tail_mask,
101
+ )
102
+ else:
103
+ src_u8 = src_addr.to(tl.pointer_type(tl.uint8))
104
+ dst_u8 = dst_addr.to(tl.pointer_type(tl.uint8))
105
+ per_tile_bytes_raw = tl.cdiv(copy_size, NUM_TILES)
106
+ per_tile_bytes = tl.cdiv(per_tile_bytes_raw, COPY_BLOCK_SIZE) * COPY_BLOCK_SIZE
107
+ tile_start = tile_idx.to(tl.int64) * per_tile_bytes
108
+ tile_end = tl.minimum(tile_start + per_tile_bytes, copy_size)
109
+ offsets = tl.arange(0, COPY_BLOCK_SIZE)
110
+ for i in range(tile_start, tile_end, COPY_BLOCK_SIZE):
111
+ mask = (i + offsets) < tile_end
112
+ data = tl.load(src_u8 + i + offsets, mask=mask)
113
+ tl.store(dst_u8 + i + offsets, data, mask=mask)
114
+
115
116
@triton.jit
117
def _copy_mamba_state_block(
@@ -41,8 +131,10 @@ def _copy_mamba_state_block(
131
# DS conv row metadata. Zero keeps the single-region copy path.
132
state_dim_row_count_ptr,
133
state_dim_row_stride_ptr,
134
+ tile_idx,
135
COPY_BLOCK_SIZE: tl.constexpr,
136
CONV_STATE_DIM_FIRST: tl.constexpr,
137
+ TEMPORAL_TILES: tl.constexpr,
138
):
139
"""Copy one (layer, state-type) mamba state block between block columns.
140
@@ -57,6 +149,13 @@ def _copy_mamba_state_block(
149
150
The caller owns the decision logic (which columns, whether to copy); this
151
device function only performs the byte copy for the given metadata slot.
152
+
153
+ ``tile_idx`` in ``[0, TEMPORAL_TILES)`` partitions the temporal state's
154
+ u64 range into ``TEMPORAL_TILES`` contiguous, COPY_BLOCK_SIZE-aligned
155
+ slices, giving more CTAs to fill the SMs at small batch (multi-MiB
156
+ temporal copies otherwise leave the GPU under-filled). Conv states are
157
+ small; only ``tile_idx == 0`` copies them. ``TEMPORAL_TILES == 1`` and
158
+ ``tile_idx == 0`` reproduces the untiled behavior.
159
"""
160
state_base_addr = tl.load(state_base_addrs_ptr + state_idx)
161
state_block_stride = tl.load(state_block_strides_ptr + state_idx)
@@ -81,6 +180,10 @@ def _copy_mamba_state_block(
180
is_conv_state = conv_width > 0
181
182
if CONV_STATE_DIM_FIRST and is_conv_state:
183
+ # Conv states are small; only tile 0 does the copy. Higher tiles
184
+ # early-return so they contribute nothing beyond a bounds check.
185
+ if tile_idx > 0:
186
+ return
187
# DS conv layout: state_len is the slide axis; copy per dim row.
188
src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
189
dim_rows = tl.load(state_dim_row_count_ptr + state_idx)
@@ -101,53 +204,46 @@ def _copy_mamba_state_block(
204
return
205
206
if is_conv_state:
207
+ if tile_idx > 0:
208
+ return
209
# SD conv: copy
210
# state[bt[src_col], token_bias:] ->
211
# state[bt[dst_col], :conv_width - token_bias]
212
+ # Small per-block bytes (~60-80 KiB) make tiling degenerate, so
213
+ # conv runs as a single-CTA memcpy (NUM_TILES=1).
214
src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
215
src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size
216
src_addr = state_base_addr + src_block_id * state_block_stride + src_offset
110
- num_elems_to_copy = (conv_width - token_bias).to(tl.int64) * state_inner_size
111
- copy_size = num_elems_to_copy * state_elem_size
112
- offsets = tl.arange(0, COPY_BLOCK_SIZE)
113
- for i in range(0, copy_size, COPY_BLOCK_SIZE):
114
- mask = (i + offsets) < copy_size
115
- curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8))
116
- curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8))
117
- data = tl.load(curr_src, mask=mask)
118
- tl.store(curr_dst, data, mask=mask)
217
+ copy_size = (
218
+ (conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size
219
+ )
220
+ _memcpy_u64_tiled(
221
+ src_addr,
222
+ dst_addr,
223
+ copy_size,
224
+ tile_idx,
225
+ COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
226
+ NUM_TILES=1,
227
+ )
228
return
229
230
# Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]]
231
+ # Body u64 range is partitioned across TEMPORAL_TILES CTAs to keep the
232
+ # SMs filled at small batch.
233
actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(tl.int64)
234
src_addr = state_base_addr + actual_src_block_id * state_block_stride
235
# Use natural block data size (inner_size * elem_size), NOT
236
# state_block_stride which is the page stride and can exceed the
237
# actual data when the state tensor uses as_strided page padding.
238
copy_size = state_inner_size * state_elem_size
128
-
129
- # Vectorize via uint64 (8B per thread → LDG.64/STG.64): both temporal
130
- # and SD conv produce src/dst addresses aligned to a full token slice
131
- # (inner_size * elem_size) and a copy_size that's a multiple of it,
132
- # which is 8B-aligned for all state dtypes in use. A masked byte tail
133
- # covers any remaining 0-7 bytes (only reachable for sub-8B slices).
134
- copy_size_u64 = copy_size // 8
135
- src_u64 = src_addr.to(tl.pointer_type(tl.uint64))
136
- dst_u64 = dst_addr.to(tl.pointer_type(tl.uint64))
137
- offsets = tl.arange(0, COPY_BLOCK_SIZE)
138
- for i in range(0, copy_size_u64, COPY_BLOCK_SIZE):
139
- mask = (i + offsets) < copy_size_u64
140
- data = tl.load(src_u64 + i + offsets, mask=mask)
141
- tl.store(dst_u64 + i + offsets, data, mask=mask)
142
-
143
- tail_start = copy_size_u64 * 8
144
- tail_bytes = copy_size - tail_start
145
- tail_off = tl.arange(0, 8)
146
- tail_src = (src_addr + tail_start).to(tl.pointer_type(tl.uint8))
147
- tail_dst = (dst_addr + tail_start).to(tl.pointer_type(tl.uint8))
148
- tail_mask = tail_off < tail_bytes
149
- tail_data = tl.load(tail_src + tail_off, mask=tail_mask)
150
- tl.store(tail_dst + tail_off, tail_data, mask=tail_mask)
239
+ _memcpy_u64_tiled(
240
+ src_addr,
241
+ dst_addr,
242
+ copy_size,
243
+ tile_idx,
244
+ COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
245
+ NUM_TILES=TEMPORAL_TILES,
246
+ )
247
248
249
@triton.jit(do_not_specialize=["num_reqs"])
@@ -194,14 +290,20 @@ def postprocess_mamba_fused_kernel(
290
# PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds
291
# the post-step new_num_computed value (V2 supplies the advanced count).
292
PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False,
293
+ # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across
294
+ # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a
295
+ # 3D grid (num_reqs, total_states, TEMPORAL_TILES). Default 1 preserves
296
+ # the existing 2D-grid contract.
297
+ TEMPORAL_TILES: tl.constexpr = 1,
298
):
299
"""
300
Fused GPU kernel for postprocess_mamba that computes decisions AND performs
301
mamba state copies without any CPU-GPU synchronization.
302
202
- Grid: (num_reqs, num_layers * num_state_types)
303
+ Grid: (num_reqs, num_layers * num_state_types [, TEMPORAL_TILES])
304
- program_id(0) = request/batch index
305
- program_id(1) = state_idx (flattened index into layer/state_type metadata)
306
+ - program_id(2) = temporal-copy tile index (0 when TEMPORAL_TILES == 1)
307
308
Note: num_layers and num_state_types are not passed as kernel parameters
309
because the kernel indexes directly into pre-flattened metadata arrays
@@ -209,6 +311,7 @@ def postprocess_mamba_fused_kernel(
311
"""
312
batch_idx = tl.program_id(0)
313
state_idx = tl.program_id(1)
314
+ tile_idx = tl.program_id(2)
315
316
# Bounds check
317
if batch_idx >= num_reqs:
@@ -247,8 +350,9 @@ def postprocess_mamba_fused_kernel(
350
dest_block_idx = aligned_new_computed // block_size - 1
351
352
# Update accepted-token count before early exits (per-request, so only
250
- # state_idx == 0 writes).
251
- if src_block_idx == dest_block_idx and state_idx == 0:
353
+ # state_idx == 0 writes). Also guard on tile_idx == 0 so tiles > 0
354
+ # (when TEMPORAL_TILES > 1) do not duplicate the store.
355
+ if src_block_idx == dest_block_idx and state_idx == 0 and tile_idx == 0:
356
tl.store(num_accepted_tokens_out_ptr + req_idx, 1)
357
358
# Skip no-op self-copy.
@@ -272,8 +376,10 @@ def postprocess_mamba_fused_kernel(
376
state_group_indices_ptr,
377
state_dim_row_count_ptr,
378
state_dim_row_stride_ptr,
379
+ tile_idx,
380
COPY_BLOCK_SIZE,
381
CONV_STATE_DIM_FIRST,
382
+ TEMPORAL_TILES,
383
)
384
385
@@ -346,6 +452,9 @@ def precopy_mamba_align_fused_kernel(
452
COPY_BLOCK_SIZE: tl.constexpr,
453
CONV_STATE_DIM_FIRST: tl.constexpr,
454
HAS_IDX_MAPPING: tl.constexpr = True,
455
+ # TEMPORAL_TILES: see postprocess_mamba_fused_kernel. Default 1 preserves
456
+ # the 2D-grid contract; > 1 requires a 3D grid.
457
+ TEMPORAL_TILES: tl.constexpr = 1,
458
):
459
"""Pre-copy mamba "align" state across block boundaries.
460
@@ -357,12 +466,13 @@ def precopy_mamba_align_fused_kernel(
466
copy specs), but driven by the GPU-resident src columns so it needs no
467
CPU-GPU sync (async-scheduling safe).
468
360
- Grid: (num_reqs, num_layers * num_state_types). V2 passes a batch-to-state
361
- idx_mapping; V1 already stores the staged arrays in batch order and uses
362
- HAS_IDX_MAPPING=False.
469
+ Grid: (num_reqs, num_layers * num_state_types [, TEMPORAL_TILES]). V2 passes
470
+ a batch-to-state idx_mapping; V1 already stores the staged arrays in batch
471
+ order and uses HAS_IDX_MAPPING=False.
472
"""
473
batch_idx = tl.program_id(0)
474
state_idx = tl.program_id(1)
475
+ tile_idx = tl.program_id(2)
476
if batch_idx >= num_reqs:
477
return
478
if HAS_IDX_MAPPING:
@@ -397,8 +507,10 @@ def precopy_mamba_align_fused_kernel(
507
state_group_indices_ptr,
508
state_dim_row_count_ptr,
509
state_dim_row_stride_ptr,
510
+ tile_idx,
511
COPY_BLOCK_SIZE,
512
CONV_STATE_DIM_FIRST,
513
+ TEMPORAL_TILES,
514
)
515
516
@@ -701,23 +813,28 @@ class MambaSpecDecodeGPUContext:
813
self.state_inner_sizes[idx] = (
814
state[0].numel() if state.dim() > 1 else 1
815
)
704
- # Temporal copies are vectorized with uint64
705
- # loads/stores; base pointer and block stride must
706
- # be 8B-aligned (tail loop handles copy_size % 8).
816
+ # Temporal copies vectorize with uint64 loads/stores.
817
+ # The kernel's head/tail handles misalignment for
818
+ # correctness, but unaligned base/stride costs in
819
+ # throughput.
820
base_addr = state.data_ptr()
821
block_stride_bytes = block_stride_elems * state.element_size()
709
- assert base_addr % 8 == 0, (
710
- f"layer {layer_name}: state.data_ptr() = "
711
- f"{base_addr:#x} is not 8B-aligned; "
712
- f"_copy_mamba_state_block uint64 "
713
- f"vectorization requires it"
714
- )
715
- assert block_stride_bytes % 8 == 0, (
716
- f"layer {layer_name}: block stride = "
717
- f"{block_stride_bytes}B is not 8B-aligned; "
718
- f"_copy_mamba_state_block uint64 "
719
- f"vectorization requires it"
720
- )
822
+ if base_addr % 8 != 0:
823
+ logger.warning_once(
824
+ "layer %s: state.data_ptr() = %#x is not "
825
+ "8B-aligned; _memcpy_u64_tiled uint64 "
826
+ "vectorization will pay misaligned load cost",
827
+ layer_name,
828
+ base_addr,
829
+ )
830
+ if block_stride_bytes % 8 != 0:
831
+ logger.warning_once(
832
+ "layer %s: block stride = %dB is not "
833
+ "8B-aligned; _memcpy_u64_tiled uint64 "
834
+ "vectorization will pay misaligned load cost",
835
+ layer_name,
836
+ block_stride_bytes,
837
+ )
838
839
self.state_group_indices[idx] = group_local_idx
840
idx += 1
@@ -771,7 +888,7 @@ class MambaSpecDecodeGPUContext:
888
)
889
890
total_states = self.num_layers * self.num_state_types
774
- grid = (num_reqs, total_states)
891
+ grid = (num_reqs, total_states, _TEMPORAL_TILES)
892
893
postprocess_mamba_fused_kernel[grid](
894
num_accepted_tokens_gpu,
@@ -795,6 +912,7 @@ class MambaSpecDecodeGPUContext:
912
block_size=self.block_size,
913
COPY_BLOCK_SIZE=1024,
914
CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
915
+ TEMPORAL_TILES=_TEMPORAL_TILES,
916
)
917
918
def run_fused_precopy(
@@ -819,7 +937,7 @@ class MambaSpecDecodeGPUContext:
937
if num_reqs == 0 or not self.is_initialized:
938
return
939
total_states = self.num_layers * self.num_state_types
822
- grid = (num_reqs, total_states)
940
+ grid = (num_reqs, total_states, _TEMPORAL_TILES)
941
precopy_mamba_align_fused_kernel[grid](
942
state_idx_gpu,
943
src_col_gpu,
@@ -839,6 +957,7 @@ class MambaSpecDecodeGPUContext:
957
COPY_BLOCK_SIZE=1024,
958
CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
959
HAS_IDX_MAPPING=idx_mapping is not None,
960
+ TEMPORAL_TILES=_TEMPORAL_TILES,
961
)
962
963
def run_fused_postprocess_align(
@@ -868,7 +987,7 @@ class MambaSpecDecodeGPUContext:
987
num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu)
988
989
total_states = self.num_layers * self.num_state_types
871
- grid = (num_reqs, total_states)
990
+ grid = (num_reqs, total_states, _TEMPORAL_TILES)
991
postprocess_mamba_fused_kernel[grid](
992
num_accepted_tokens_snapshot,
993
state_idx_gpu,
@@ -893,6 +1012,7 @@ class MambaSpecDecodeGPUContext:
1012
CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
1013
HAS_IDX_MAPPING=True,
1014
PRECOMPUTED_NEW_COMPUTED=True,
1015
+ TEMPORAL_TILES=_TEMPORAL_TILES,
1016
)
1017
1018