| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import dataclasses |
| 4 | import itertools |
| 5 | from collections.abc import Callable |
| 6 | from typing import Any, NamedTuple |
| 7 | |
| 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 | MambaStateCopyFuncsByType, |
| 14 | get_conv_copy_spec, |
| 15 | get_temporal_copy_spec, |
| 16 | is_conv_state_dim_first, |
| 17 | ) |
| 18 | from vllm.triton_utils import tl, triton |
| 19 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 20 | from vllm.utils.math_utils import cdiv |
| 21 | from vllm.v1.core.sched.output import SchedulerOutput |
| 22 | from vllm.v1.kv_cache_interface import ( |
| 23 | KVCacheConfig, |
| 24 | KVCacheGroupSpec, |
| 25 | MambaSpec, |
| 26 | UniformTypeKVCacheSpecs, |
| 27 | ) |
| 28 | from vllm.v1.utils import CpuGpuBuffer |
| 29 | from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | # 16 saturates HBM on H100/GB200 across the reqs=8..128 range in |
| 34 | # microbenchmarks |
| 35 | _TEMPORAL_TILES = 16 |
| 36 | |
| 37 | |
| 38 | @triton.jit(do_not_specialize=["num_requests"]) |
| 39 | def get_aligned_state_indices_multi_group_kernel( |
| 40 | block_table_ptrs_ptr, |
| 41 | seq_lens_ptr, |
| 42 | state_indices_ptr, |
| 43 | block_table_stride_req: tl.int64, |
| 44 | seq_lens_stride: tl.constexpr, |
| 45 | state_indices_stride_0: tl.constexpr, |
| 46 | state_indices_stride_1: tl.constexpr, |
| 47 | state_indices_stride_2: tl.constexpr, |
| 48 | num_requests, |
| 49 | CACHE_BLOCK_SIZE: tl.constexpr, |
| 50 | NUM_GROUPS: tl.constexpr, |
| 51 | BLOCK_GROUPS: tl.constexpr, |
| 52 | NUM_STATE_SLOTS: tl.constexpr, |
| 53 | BLOCK_STATE_SLOTS: tl.constexpr, |
| 54 | BLOCK_ROWS: tl.constexpr, |
| 55 | ): |
| 56 | rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) |
| 57 | valid_row = rows < num_requests |
| 58 | |
| 59 | seq_lens = tl.load( |
| 60 | seq_lens_ptr + rows * seq_lens_stride, |
| 61 | mask=valid_row, |
| 62 | other=1, |
| 63 | ) |
| 64 | first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0) |
| 65 | |
| 66 | # load multiple block table for each group |
| 67 | groups = tl.arange(0, BLOCK_GROUPS) |
| 68 | valid_group = groups < NUM_GROUPS |
| 69 | group_base_addrs = tl.load( |
| 70 | block_table_ptrs_ptr + groups, |
| 71 | mask=valid_group, |
| 72 | other=0, |
| 73 | ) |
| 74 | block_tables = group_base_addrs.to(tl.pointer_type(tl.int32)) |
| 75 | state_slots = tl.arange(0, BLOCK_STATE_SLOTS) |
| 76 | valid_state_slot = state_slots < NUM_STATE_SLOTS |
| 77 | state_indices = tl.load( |
| 78 | block_tables[:, None, None] |
| 79 | + rows[None, :, None] * block_table_stride_req |
| 80 | + first_state_slot[None, :, None] |
| 81 | + state_slots[None, None, :], |
| 82 | mask=( |
| 83 | valid_group[:, None, None] |
| 84 | & valid_row[None, :, None] |
| 85 | & valid_state_slot[None, None, :] |
| 86 | ), |
| 87 | ) |
| 88 | tl.store( |
| 89 | state_indices_ptr |
| 90 | + groups[:, None, None] * state_indices_stride_0 |
| 91 | + rows[None, :, None] * state_indices_stride_1 |
| 92 | + state_slots[None, None, :] * state_indices_stride_2, |
| 93 | state_indices, |
| 94 | mask=( |
| 95 | valid_group[:, None, None] |
| 96 | & valid_row[None, :, None] |
| 97 | & valid_state_slot[None, None, :] |
| 98 | ), |
| 99 | ) |
| 100 | |
| 101 | |
| 102 | @triton.jit |
| 103 | def _memcpy_u64_tiled( |
| 104 | src_addr, |
| 105 | dst_addr, |
| 106 | copy_size, |
| 107 | tile_idx, |
| 108 | COPY_BLOCK_SIZE: tl.constexpr, |
| 109 | NUM_TILES: tl.constexpr, |
| 110 | ): |
| 111 | """Head/body/tail memcpy with the u64 body split across ``NUM_TILES`` CTAs. |
| 112 | |
| 113 | Fast path (``src`` and ``dst`` share sub-8B alignment): tile 0 owns the |
| 114 | byte head that lifts dst to 8B and the 0-7 byte tail; body tiles vectorize |
| 115 | as u64 over the aligned interior. ``NUM_TILES=1`` collapses to a single- |
| 116 | CTA memcpy. Production callers derive both addresses from the same |
| 117 | ``state_base_addr + block_id * stride`` and always take this path. |
| 118 | |
| 119 | Slow path (mismatched sub-8B alignment): byte-wide tiled copy. Some |
| 120 | NVIDIA parts (e.g. GB200) reject misaligned u64 loads with |
| 121 | ``cudaErrorMisalignedAddress`` instead of accepting the misaligned-sector |
| 122 | cost, so we can't just let the fast path run with a misaligned src. |
| 123 | """ |
| 124 | src_addr_i = src_addr.to(tl.int64) |
| 125 | dst_addr_i = dst_addr.to(tl.int64) |
| 126 | |
| 127 | if ((src_addr_i ^ dst_addr_i) & 7) == 0: |
| 128 | head_bytes = tl.minimum(((-dst_addr_i) & 7).to(tl.int64), copy_size) |
| 129 | if tile_idx == 0: |
| 130 | head_off = tl.arange(0, 8) |
| 131 | head_mask = head_off < head_bytes |
| 132 | head_src = src_addr.to(tl.pointer_type(tl.uint8)) |
| 133 | head_dst = dst_addr.to(tl.pointer_type(tl.uint8)) |
| 134 | tl.store( |
| 135 | head_dst + head_off, |
| 136 | tl.load(head_src + head_off, mask=head_mask), |
| 137 | mask=head_mask, |
| 138 | ) |
| 139 | |
| 140 | # Body: u64 tiled. Rounding per_tile up to COPY_BLOCK_SIZE keeps every |
| 141 | # inner-loop iteration full-width vectorized; only the last non-empty |
| 142 | # tile can be masked. Late tiles fall off the end and iterate zero |
| 143 | # times. |
| 144 | body_bytes = copy_size - head_bytes |
| 145 | body_u64 = body_bytes // 8 |
| 146 | per_tile_u64_raw = tl.cdiv(body_u64, NUM_TILES) |
| 147 | per_tile_u64 = tl.cdiv(per_tile_u64_raw, COPY_BLOCK_SIZE) * COPY_BLOCK_SIZE |
| 148 | tile_start = tile_idx.to(tl.int64) * per_tile_u64 |
| 149 | tile_end = tl.minimum(tile_start + per_tile_u64, body_u64) |
| 150 | |
| 151 | src_body_u64 = (src_addr + head_bytes).to(tl.pointer_type(tl.uint64)) |
| 152 | dst_body_u64 = (dst_addr + head_bytes).to(tl.pointer_type(tl.uint64)) |
| 153 | offsets = tl.arange(0, COPY_BLOCK_SIZE) |
| 154 | for i in range(tile_start, tile_end, COPY_BLOCK_SIZE): |
| 155 | mask = (i + offsets) < tile_end |
| 156 | data = tl.load(src_body_u64 + i + offsets, mask=mask) |
| 157 | tl.store(dst_body_u64 + i + offsets, data, mask=mask) |
| 158 | |
| 159 | if tile_idx == 0: |
| 160 | tail_start = head_bytes + body_u64 * 8 |
| 161 | tail_bytes = copy_size - tail_start |
| 162 | tail_off = tl.arange(0, 8) |
| 163 | tail_src = (src_addr + tail_start).to(tl.pointer_type(tl.uint8)) |
| 164 | tail_dst = (dst_addr + tail_start).to(tl.pointer_type(tl.uint8)) |
| 165 | tail_mask = tail_off < tail_bytes |
| 166 | tl.store( |
| 167 | tail_dst + tail_off, |
| 168 | tl.load(tail_src + tail_off, mask=tail_mask), |
| 169 | mask=tail_mask, |
| 170 | ) |
| 171 | else: |
| 172 | src_u8 = src_addr.to(tl.pointer_type(tl.uint8)) |
| 173 | dst_u8 = dst_addr.to(tl.pointer_type(tl.uint8)) |
| 174 | per_tile_bytes_raw = tl.cdiv(copy_size, NUM_TILES) |
| 175 | per_tile_bytes = tl.cdiv(per_tile_bytes_raw, COPY_BLOCK_SIZE) * COPY_BLOCK_SIZE |
| 176 | tile_start = tile_idx.to(tl.int64) * per_tile_bytes |
| 177 | tile_end = tl.minimum(tile_start + per_tile_bytes, copy_size) |
| 178 | offsets = tl.arange(0, COPY_BLOCK_SIZE) |
| 179 | for i in range(tile_start, tile_end, COPY_BLOCK_SIZE): |
| 180 | mask = (i + offsets) < tile_end |
| 181 | data = tl.load(src_u8 + i + offsets, mask=mask) |
| 182 | tl.store(dst_u8 + i + offsets, data, mask=mask) |
| 183 | |
| 184 | |
| 185 | def _reinterpret_u64_as_i64(value: int) -> int: |
| 186 | """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" |
| 187 | return value if value < (1 << 63) else value - (1 << 64) |
| 188 | |
| 189 | |
| 190 | @triton.jit |
| 191 | def _copy_mamba_state_block( |
| 192 | state_idx, |
| 193 | bt_row_idx, |
| 194 | src_col, |
| 195 | dst_col, |
| 196 | token_bias, |
| 197 | block_table_ptrs_ptr, |
| 198 | block_table_stride_req, |
| 199 | state_base_addrs_ptr, |
| 200 | state_block_strides_ptr, |
| 201 | state_elem_sizes_ptr, |
| 202 | state_inner_sizes_ptr, |
| 203 | state_conv_widths_ptr, |
| 204 | state_group_indices_ptr, |
| 205 | # DS conv row metadata. Zero keeps the single-region copy path. |
| 206 | state_dim_row_count_ptr, |
| 207 | state_dim_row_stride_ptr, |
| 208 | tile_idx, |
| 209 | COPY_BLOCK_SIZE: tl.constexpr, |
| 210 | CONV_STATE_DIM_FIRST: tl.constexpr, |
| 211 | TEMPORAL_TILES: tl.constexpr, |
| 212 | ): |
| 213 | """Copy one (layer, state-type) mamba state block between block columns. |
| 214 | |
| 215 | Shared copy body of ``postprocess_mamba_fused_kernel`` and |
| 216 | ``precopy_mamba_align_fused_kernel``, mirroring the V1 copy specs |
| 217 | (``get_conv_copy_spec`` / ``get_temporal_copy_spec``): |
| 218 | - conv state (conv_width > 0): shift the window by ``token_bias`` tokens, |
| 219 | ``state[bt[src_col], token_bias:] -> |
| 220 | state[bt[dst_col], :conv_width - token_bias]`` |
| 221 | - temporal state: ``token_bias`` selects the accepted speculative column, |
| 222 | ``state[bt[src_col + token_bias]] -> state[bt[dst_col]]`` |
| 223 | |
| 224 | The caller owns the decision logic (which columns, whether to copy); this |
| 225 | device function only performs the byte copy for the given metadata slot. |
| 226 | |
| 227 | ``tile_idx`` in ``[0, TEMPORAL_TILES)`` partitions the temporal state's |
| 228 | u64 range into ``TEMPORAL_TILES`` contiguous, COPY_BLOCK_SIZE-aligned |
| 229 | slices, giving more CTAs to fill the SMs at small batch (multi-MiB |
| 230 | temporal copies otherwise leave the GPU under-filled). Conv states are |
| 231 | small; only ``tile_idx == 0`` copies them. ``TEMPORAL_TILES == 1`` and |
| 232 | ``tile_idx == 0`` reproduces the untiled behavior. |
| 233 | """ |
| 234 | state_base_addr = tl.load(state_base_addrs_ptr + state_idx) |
| 235 | state_block_stride = tl.load(state_block_strides_ptr + state_idx) |
| 236 | state_elem_size = tl.load(state_elem_sizes_ptr + state_idx) |
| 237 | state_inner_size = tl.load(state_inner_sizes_ptr + state_idx) |
| 238 | conv_width = tl.load(state_conv_widths_ptr + state_idx) |
| 239 | |
| 240 | # Load the group index for this state, then index into the correct |
| 241 | # group's block table. Each mamba group has independently allocated |
| 242 | # physical blocks. Reinterpret as int32* since block ids are int32. |
| 243 | group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) |
| 244 | group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) |
| 245 | block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32)) |
| 246 | block_table_base = block_table_typed + bt_row_idx * block_table_stride_req |
| 247 | |
| 248 | # Widen block ids to int64 before they reach `block_id * state_block_stride` |
| 249 | # below: state_block_stride can exceed 2**31 bytes for large mamba caches, |
| 250 | # and Triton would otherwise do the multiply in int32 and wrap. |
| 251 | dest_block_id = tl.load(block_table_base + dst_col).to(tl.int64) |
| 252 | dst_addr = state_base_addr + dest_block_id * state_block_stride |
| 253 | |
| 254 | is_conv_state = conv_width > 0 |
| 255 | |
| 256 | if CONV_STATE_DIM_FIRST and is_conv_state: |
| 257 | # Conv states are small; only tile 0 does the copy. Higher tiles |
| 258 | # early-return so they contribute nothing beyond a bounds check. |
| 259 | if tile_idx > 0: |
| 260 | return |
| 261 | # DS conv layout: state_len is the slide axis; copy per dim row. |
| 262 | src_block_id = tl.load(block_table_base + src_col).to(tl.int64) |
| 263 | dim_rows = tl.load(state_dim_row_count_ptr + state_idx) |
| 264 | row_stride = tl.load(state_dim_row_stride_ptr + state_idx) |
| 265 | src_block_addr = state_base_addr + src_block_id * state_block_stride |
| 266 | offsets = tl.arange(0, COPY_BLOCK_SIZE) |
| 267 | |
| 268 | # Stable row-to-lane ownership makes left shifts memmove-safe while |
| 269 | # exposing the dimension rows in parallel. All addresses retain |
| 270 | # state_elem_size alignment: tensor strides and token offsets are |
| 271 | # measured in whole elements before conversion to bytes. |
| 272 | num_dst_tokens = conv_width - token_bias |
| 273 | for token_idx in range(0, num_dst_tokens): |
| 274 | for row_base in range(0, dim_rows, COPY_BLOCK_SIZE): |
| 275 | rows = row_base + offsets |
| 276 | mask = rows < dim_rows |
| 277 | src_byte_addr = ( |
| 278 | src_block_addr |
| 279 | + rows * row_stride |
| 280 | + (token_idx + token_bias) * state_elem_size |
| 281 | ) |
| 282 | dst_byte_addr = ( |
| 283 | dst_addr + rows * row_stride + token_idx * state_elem_size |
| 284 | ) |
| 285 | if state_elem_size == 2: |
| 286 | src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16)) |
| 287 | dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16)) |
| 288 | data_u16 = tl.load(src_u16, mask=mask) |
| 289 | tl.store(dst_u16, data_u16, mask=mask) |
| 290 | elif state_elem_size == 4: |
| 291 | src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32)) |
| 292 | dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32)) |
| 293 | data_u32 = tl.load(src_u32, mask=mask) |
| 294 | tl.store(dst_u32, data_u32, mask=mask) |
| 295 | else: |
| 296 | for byte_idx in range(0, state_elem_size): |
| 297 | src_u8 = (src_byte_addr + byte_idx).to( |
| 298 | tl.pointer_type(tl.uint8) |
| 299 | ) |
| 300 | dst_u8 = (dst_byte_addr + byte_idx).to( |
| 301 | tl.pointer_type(tl.uint8) |
| 302 | ) |
| 303 | data_u8 = tl.load(src_u8, mask=mask) |
| 304 | tl.store(dst_u8, data_u8, mask=mask) |
| 305 | return |
| 306 | |
| 307 | if is_conv_state: |
| 308 | if tile_idx > 0: |
| 309 | return |
| 310 | # SD conv: copy |
| 311 | # state[bt[src_col], token_bias:] -> |
| 312 | # state[bt[dst_col], :conv_width - token_bias] |
| 313 | src_block_id = tl.load(block_table_base + src_col).to(tl.int64) |
| 314 | src_block_addr = state_base_addr + src_block_id * state_block_stride |
| 315 | token_bytes = state_inner_size * state_elem_size |
| 316 | num_dst_tokens = conv_width - token_bias |
| 317 | |
| 318 | # Distinct blocks and exact self-copies cannot have a destructive |
| 319 | # overlap, so retain the u64-vectorized single-CTA copy. |
| 320 | if src_block_id != dest_block_id or token_bias == 0: |
| 321 | src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes |
| 322 | copy_size = num_dst_tokens.to(tl.int64) * token_bytes |
| 323 | _memcpy_u64_tiled( |
| 324 | src_addr, |
| 325 | dst_addr, |
| 326 | copy_size, |
| 327 | tile_idx, |
| 328 | COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, |
| 329 | NUM_TILES=1, |
| 330 | ) |
| 331 | return |
| 332 | |
| 333 | # Copy tokens from low to high. Each token-sized source and destination |
| 334 | # region is disjoint, so same-block left shifts are memmove-safe |
| 335 | # without a barrier. |
| 336 | for token_idx in range(0, num_dst_tokens): |
| 337 | src_token = src_block_addr + (token_idx + token_bias) * token_bytes |
| 338 | dst_token = dst_addr + token_idx * token_bytes |
| 339 | _memcpy_u64_tiled( |
| 340 | src_token, |
| 341 | dst_token, |
| 342 | token_bytes, |
| 343 | tile_idx, |
| 344 | COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, |
| 345 | NUM_TILES=1, |
| 346 | ) |
| 347 | return |
| 348 | |
| 349 | # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] |
| 350 | # Body u64 range is partitioned across TEMPORAL_TILES CTAs to keep the |
| 351 | # SMs filled at small batch. |
| 352 | actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(tl.int64) |
| 353 | src_addr = state_base_addr + actual_src_block_id * state_block_stride |
| 354 | # Use natural block data size (inner_size * elem_size), NOT |
| 355 | # state_block_stride which is the page stride and can exceed the |
| 356 | # actual data when the state tensor uses as_strided page padding. |
| 357 | copy_size = state_inner_size * state_elem_size |
| 358 | _memcpy_u64_tiled( |
| 359 | src_addr, |
| 360 | dst_addr, |
| 361 | copy_size, |
| 362 | tile_idx, |
| 363 | COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, |
| 364 | NUM_TILES=TEMPORAL_TILES, |
| 365 | ) |
| 366 | |
| 367 | |
| 368 | @triton.jit(do_not_specialize=["num_reqs"]) |
| 369 | def postprocess_mamba_fused_kernel( |
| 370 | # Decision inputs (per-request) |
| 371 | num_accepted_tokens_ptr, |
| 372 | mamba_state_idx_ptr, |
| 373 | num_scheduled_tokens_ptr, |
| 374 | num_computed_tokens_ptr, |
| 375 | num_draft_tokens_ptr, |
| 376 | # Per-group block table base addresses: int64[num_groups]. Each entry is |
| 377 | # the data_ptr of that group's persistent [max_reqs, max_blocks] int32 |
| 378 | # block table. |
| 379 | block_table_ptrs_ptr, |
| 380 | block_table_stride_req: tl.int64, # stride between requests (in elements) |
| 381 | # Mamba state metadata, flattened across all layer/state tensors. Different |
| 382 | # Mamba groups may contribute different numbers of state tensors. |
| 383 | state_base_addrs_ptr, # base address of each state tensor |
| 384 | state_block_strides_ptr, # bytes per block for each state |
| 385 | state_elem_sizes_ptr, # element size for each state |
| 386 | state_inner_sizes_ptr, # number of elements in inner dimensions |
| 387 | state_conv_widths_ptr, # conv width for conv states (0 for temporal) |
| 388 | state_group_indices_ptr, # maps state_idx to group index in block table |
| 389 | # DS conv row metadata. Zero keeps the single-region copy path. |
| 390 | state_dim_row_count_ptr, # int32: per-block dim row count for DS conv |
| 391 | state_dim_row_stride_ptr, # int64: bytes between rows for DS conv |
| 392 | # Output: num_accepted_tokens update (for src==dst case) |
| 393 | num_accepted_tokens_out_ptr, |
| 394 | # Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The |
| 395 | # per-request decision arrays are in req-state-slot order; the block table |
| 396 | # is in batch order, so HAS_IDX_MAPPING splits the two indexings. |
| 397 | idx_mapping_ptr, |
| 398 | # Runtime parameter (varies per batch - NOT constexpr to avoid recompilation) |
| 399 | num_reqs, |
| 400 | # Compile-time constants (fixed after model initialization) |
| 401 | # block_size: determined by model config, constant for all invocations |
| 402 | block_size: tl.constexpr, |
| 403 | # COPY_BLOCK_SIZE: fixed tuning parameter for memory copy loop |
| 404 | COPY_BLOCK_SIZE: tl.constexpr, |
| 405 | CONV_STATE_DIM_FIRST: tl.constexpr, |
| 406 | # HAS_IDX_MAPPING: when True, program_id(0) is a batch index resolved to a |
| 407 | # req-state slot via idx_mapping_ptr (V2). When False, it is the req index. |
| 408 | HAS_IDX_MAPPING: tl.constexpr = False, |
| 409 | # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds |
| 410 | # the post-step new_num_computed value (V2 supplies the advanced count). |
| 411 | PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False, |
| 412 | # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across |
| 413 | # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a |
| 414 | # 3D grid (num_reqs, total_states, TEMPORAL_TILES). Default 1 preserves |
| 415 | # the existing 2D-grid contract. |
| 416 | TEMPORAL_TILES: tl.constexpr = 1, |
| 417 | ): |
| 418 | """Fused GPU kernel for postprocess_mamba that computes decisions AND performs |
| 419 | mamba state copies without any CPU-GPU synchronization. |
| 420 | |
| 421 | Grid: (num_reqs, num_states [, TEMPORAL_TILES]) |
| 422 | - program_id(0) = request/batch index |
| 423 | - program_id(1) = state_idx (flattened index into layer/state_type metadata) |
| 424 | - program_id(2) = temporal-copy tile index (0 when TEMPORAL_TILES == 1) |
| 425 | |
| 426 | The kernel indexes directly into pre-flattened metadata arrays using |
| 427 | program_id(1). The grid dimensions encode the total state count. |
| 428 | """ |
| 429 | batch_idx = tl.program_id(0) |
| 430 | state_idx = tl.program_id(1) |
| 431 | tile_idx = tl.program_id(2) |
| 432 | |
| 433 | # Bounds check |
| 434 | if batch_idx >= num_reqs: |
| 435 | return |
| 436 | |
| 437 | if HAS_IDX_MAPPING: |
| 438 | req_idx = tl.load(idx_mapping_ptr + batch_idx) |
| 439 | if req_idx < 0: |
| 440 | return |
| 441 | else: |
| 442 | req_idx = batch_idx |
| 443 | |
| 444 | # Compute decision logic (mirrors postprocess_mamba Python reference) |
| 445 | num_accepted = tl.load(num_accepted_tokens_ptr + req_idx) |
| 446 | src_block_idx = tl.load(mamba_state_idx_ptr + req_idx) |
| 447 | |
| 448 | if PRECOMPUTED_NEW_COMPUTED: |
| 449 | new_num_computed = tl.load(num_computed_tokens_ptr + req_idx) |
| 450 | num_tokens_running_state = new_num_computed - num_accepted + 1 |
| 451 | else: |
| 452 | num_scheduled = tl.load(num_scheduled_tokens_ptr + req_idx) |
| 453 | num_computed = tl.load(num_computed_tokens_ptr + req_idx) |
| 454 | num_draft = tl.load(num_draft_tokens_ptr + req_idx) |
| 455 | num_tokens_running_state = num_computed + num_scheduled - num_draft |
| 456 | new_num_computed = num_tokens_running_state + num_accepted - 1 |
| 457 | |
| 458 | aligned_new_computed = (new_num_computed // block_size) * block_size |
| 459 | |
| 460 | needs_copy = aligned_new_computed >= num_tokens_running_state |
| 461 | |
| 462 | if not needs_copy: |
| 463 | return |
| 464 | |
| 465 | # Compute copy parameters |
| 466 | accept_token_bias = aligned_new_computed - num_tokens_running_state |
| 467 | dest_block_idx = aligned_new_computed // block_size - 1 |
| 468 | |
| 469 | # Update accepted-token count before early exits (per-request, so only |
| 470 | # state_idx == 0 writes). Also guard on tile_idx == 0 so tiles > 0 |
| 471 | # (when TEMPORAL_TILES > 1) do not duplicate the store. |
| 472 | if src_block_idx == dest_block_idx and state_idx == 0 and tile_idx == 0: |
| 473 | tl.store(num_accepted_tokens_out_ptr + req_idx, 1) |
| 474 | |
| 475 | # Skip no-op self-copy. |
| 476 | if src_block_idx == dest_block_idx and accept_token_bias == 0: |
| 477 | return |
| 478 | |
| 479 | bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx |
| 480 | _copy_mamba_state_block( |
| 481 | state_idx, |
| 482 | bt_row_idx, |
| 483 | src_block_idx, |
| 484 | dest_block_idx, |
| 485 | accept_token_bias, |
| 486 | block_table_ptrs_ptr, |
| 487 | block_table_stride_req, |
| 488 | state_base_addrs_ptr, |
| 489 | state_block_strides_ptr, |
| 490 | state_elem_sizes_ptr, |
| 491 | state_inner_sizes_ptr, |
| 492 | state_conv_widths_ptr, |
| 493 | state_group_indices_ptr, |
| 494 | state_dim_row_count_ptr, |
| 495 | state_dim_row_stride_ptr, |
| 496 | tile_idx, |
| 497 | COPY_BLOCK_SIZE, |
| 498 | CONV_STATE_DIM_FIRST, |
| 499 | TEMPORAL_TILES, |
| 500 | ) |
| 501 | |
| 502 | |
| 503 | @triton.jit(do_not_specialize=["num_reqs"]) |
| 504 | def preprocess_mamba_align_fused_kernel( |
| 505 | idx_mapping_ptr, |
| 506 | state_idx_ptr, |
| 507 | num_computed_tokens_ptr, |
| 508 | query_start_loc_ptr, |
| 509 | num_accepted_tokens_ptr, |
| 510 | src_col_ptr, |
| 511 | src_off_ptr, |
| 512 | num_reqs, |
| 513 | BLOCK_SIZE: tl.constexpr, |
| 514 | MAMBA_BLOCK_SIZE: tl.constexpr, |
| 515 | ): |
| 516 | """Fused align preprocess: emit the pre-copy src column/offset AND advance |
| 517 | state_idx (with accepted-token reset) in a single launch (V2 align). |
| 518 | |
| 519 | Per batch_idx (0..num_reqs-1), resolving req slot via idx_mapping: |
| 520 | 1. Read pre-advance state_idx and num_accepted (last step's values). |
| 521 | 2. Store the pre-copy src columns for ``precopy_mamba_align_fused_kernel``: |
| 522 | - src_col = state_idx (the previous running block column) |
| 523 | - src_off = max(num_accepted - 1, 0) (the accepted-token bias) |
| 524 | 3. Advance state_idx to the new running block, and reset num_accepted to 1 |
| 525 | when a block boundary is crossed (so the migrated state, now at the |
| 526 | start of the new block, is read with the neutral bias). |
| 527 | """ |
| 528 | offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 529 | mask = offsets < num_reqs |
| 530 | req_indices = tl.load(idx_mapping_ptr + offsets, mask=mask, other=0) |
| 531 | |
| 532 | state_idx = tl.load(state_idx_ptr + req_indices, mask=mask, other=-1) |
| 533 | num_accepted = tl.load(num_accepted_tokens_ptr + req_indices, mask=mask, other=1) |
| 534 | |
| 535 | src_off = tl.maximum(num_accepted - 1, 0) |
| 536 | tl.store(src_col_ptr + req_indices, state_idx, mask=mask) |
| 537 | tl.store(src_off_ptr + req_indices, src_off, mask=mask) |
| 538 | |
| 539 | num_computed = tl.load(num_computed_tokens_ptr + req_indices, mask=mask, other=0) |
| 540 | query_start = tl.load(query_start_loc_ptr + offsets, mask=mask, other=0) |
| 541 | query_end = tl.load(query_start_loc_ptr + offsets + 1, mask=mask, other=0) |
| 542 | computed_after = num_computed + query_end - query_start |
| 543 | new_state_idx = (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1 |
| 544 | tl.store(state_idx_ptr + req_indices, new_state_idx, mask=mask) |
| 545 | should_reset = (state_idx >= 0) & (state_idx != new_state_idx) |
| 546 | tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset) |
| 547 | |
| 548 | |
| 549 | @triton.jit(do_not_specialize=["num_reqs"]) |
| 550 | def precopy_mamba_align_fused_kernel( |
| 551 | # Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by |
| 552 | # the V2 fused align preprocess kernel for the current step: |
| 553 | mamba_state_idx_ptr, # post-advance dst block column |
| 554 | src_col_ptr, # pre-advance src block column (-1 = fresh) |
| 555 | token_bias_ptr, # accepted-token bias = num_accepted - 1 (pre-reset) |
| 556 | # Same flattened state-layout metadata as postprocess_mamba_fused_kernel |
| 557 | block_table_ptrs_ptr, |
| 558 | block_table_stride_req: tl.int64, |
| 559 | state_base_addrs_ptr, |
| 560 | state_block_strides_ptr, |
| 561 | state_elem_sizes_ptr, |
| 562 | state_inner_sizes_ptr, |
| 563 | state_conv_widths_ptr, |
| 564 | state_group_indices_ptr, |
| 565 | state_dim_row_count_ptr, |
| 566 | state_dim_row_stride_ptr, |
| 567 | idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) |
| 568 | num_reqs, |
| 569 | COPY_BLOCK_SIZE: tl.constexpr, |
| 570 | CONV_STATE_DIM_FIRST: tl.constexpr, |
| 571 | HAS_IDX_MAPPING: tl.constexpr = True, |
| 572 | # TEMPORAL_TILES: see postprocess_mamba_fused_kernel. Default 1 preserves |
| 573 | # the 2D-grid contract; > 1 requires a 3D grid. |
| 574 | TEMPORAL_TILES: tl.constexpr = 1, |
| 575 | ): |
| 576 | """Pre-copy mamba "align" state across block boundaries. |
| 577 | |
| 578 | Before the forward pass, copy each request's last SSM/conv state from its |
| 579 | previous block column into the new window block column, so the kernels read |
| 580 | the initial state from the write-side block as usual (V1 align semantics). |
| 581 | Same per-(layer, state) copy semantics as ``postprocess_mamba_fused_kernel`` |
| 582 | (shared ``_copy_mamba_state_block`` body, i.e. the V1 ``preprocess_mamba`` |
| 583 | copy specs), but driven by the GPU-resident src columns so it needs no |
| 584 | CPU-GPU sync (async-scheduling safe). |
| 585 | |
| 586 | Grid: (num_reqs, num_states [, TEMPORAL_TILES]). V2 passes a |
| 587 | batch-to-state idx_mapping; V1 already stores the staged arrays in batch |
| 588 | order and uses HAS_IDX_MAPPING=False. |
| 589 | """ |
| 590 | batch_idx = tl.program_id(0) |
| 591 | state_idx = tl.program_id(1) |
| 592 | tile_idx = tl.program_id(2) |
| 593 | if batch_idx >= num_reqs: |
| 594 | return |
| 595 | if HAS_IDX_MAPPING: |
| 596 | req_idx = tl.load(idx_mapping_ptr + batch_idx) |
| 597 | if req_idx < 0: |
| 598 | return |
| 599 | else: |
| 600 | req_idx = batch_idx |
| 601 | |
| 602 | src_col = tl.load(src_col_ptr + req_idx) |
| 603 | dst_col = tl.load(mamba_state_idx_ptr + req_idx) |
| 604 | # Fresh state, or still writing the same block: kernels locate the initial |
| 605 | # state in-block via num_accepted (preserved when no boundary is crossed), |
| 606 | # so there is nothing to copy. |
| 607 | if src_col < 0 or src_col == dst_col: |
| 608 | return |
| 609 | |
| 610 | token_bias = tl.load(token_bias_ptr + req_idx) |
| 611 | _copy_mamba_state_block( |
| 612 | state_idx, |
| 613 | batch_idx, |
| 614 | src_col, |
| 615 | dst_col, |
| 616 | token_bias, |
| 617 | block_table_ptrs_ptr, |
| 618 | block_table_stride_req, |
| 619 | state_base_addrs_ptr, |
| 620 | state_block_strides_ptr, |
| 621 | state_elem_sizes_ptr, |
| 622 | state_inner_sizes_ptr, |
| 623 | state_conv_widths_ptr, |
| 624 | state_group_indices_ptr, |
| 625 | state_dim_row_count_ptr, |
| 626 | state_dim_row_stride_ptr, |
| 627 | tile_idx, |
| 628 | COPY_BLOCK_SIZE, |
| 629 | CONV_STATE_DIM_FIRST, |
| 630 | TEMPORAL_TILES, |
| 631 | ) |
| 632 | |
| 633 | |
| 634 | @triton.jit |
| 635 | def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): |
| 636 | pid = tl.program_id(0) |
| 637 | |
| 638 | src_ptr = tl.load(src_ptrs + pid) |
| 639 | dst_ptr = tl.load(dst_ptrs + pid) |
| 640 | size = tl.load(sizes + pid) |
| 641 | is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr |
| 642 | |
| 643 | offsets = tl.arange(0, BLOCK_SIZE) |
| 644 | for i in range(0, size, BLOCK_SIZE): |
| 645 | mask = (i + offsets) < size |
| 646 | |
| 647 | curr_src_ptr = (src_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) |
| 648 | curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) |
| 649 | |
| 650 | data = tl.load(curr_src_ptr, mask=mask) |
| 651 | if is_left_overlap: |
| 652 | # Preserve each lane's source before a lower-address lane stores |
| 653 | # over it. The condition is uniform within the program. |
| 654 | tl.debug_barrier() |
| 655 | tl.store(curr_dst_ptr, data, mask=mask) |
| 656 | |
| 657 | |
| 658 | def batch_memcpy(src_ptrs, dst_ptrs, sizes): |
| 659 | batch = src_ptrs.shape[0] |
| 660 | assert dst_ptrs.shape[0] == batch |
| 661 | assert sizes.shape[0] == batch |
| 662 | |
| 663 | grid = (batch,) |
| 664 | BLOCK_SIZE = 1024 |
| 665 | batch_memcpy_kernel[grid](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=BLOCK_SIZE) |
| 666 | |
| 667 | |
| 668 | def _get_mamba_spec_for_layer( |
| 669 | kv_cache_group: KVCacheGroupSpec, layer_name: str |
| 670 | ) -> MambaSpec: |
| 671 | kv_cache_spec = kv_cache_group.kv_cache_spec |
| 672 | if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): |
| 673 | kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name] |
| 674 | assert isinstance(kv_cache_spec, MambaSpec) |
| 675 | return kv_cache_spec |
| 676 | |
| 677 | |
| 678 | def get_mamba_group_ids(mamba_groups: dict[MambaSpec, list[int]]) -> list[int]: |
| 679 | """Return the sorted group ids of all mamba groups, without duplicates. |
| 680 | |
| 681 | A group whose layers use different ``MambaSpec``s is listed under one key |
| 682 | per spec, so the same group id can appear more than once. |
| 683 | """ |
| 684 | return sorted( |
| 685 | {group_id for group_ids in mamba_groups.values() for group_id in group_ids} |
| 686 | ) |
| 687 | |
| 688 | |
| 689 | def get_mamba_groups(kv_cache_config: KVCacheConfig) -> dict[MambaSpec, list[int]]: |
| 690 | """Return a mapping from each distinct MambaSpec to its sorted group_ids.""" |
| 691 | mamba_groups: dict[MambaSpec, set[int]] = {} |
| 692 | for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups): |
| 693 | kv_cache_spec = kv_cache_group.kv_cache_spec |
| 694 | if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): |
| 695 | kv_cache_spec = kv_cache_spec.first_spec |
| 696 | if not isinstance(kv_cache_spec, MambaSpec): |
| 697 | continue |
| 698 | for layer_name in kv_cache_group.layer_names: |
| 699 | mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name) |
| 700 | mamba_groups.setdefault(mamba_spec, set()).add(i) |
| 701 | assert len(mamba_groups) > 0, "no mamba layers in the model" |
| 702 | return {spec: sorted(ids) for spec, ids in mamba_groups.items()} |
| 703 | |
| 704 | |
| 705 | def validate_mamba_state_copy_funcs( |
| 706 | mamba_groups: dict[MambaSpec, list[int]], |
| 707 | copy_funcs: MambaStateCopyFuncsByType, |
| 708 | ) -> None: |
| 709 | for mamba_spec in mamba_groups: |
| 710 | assert mamba_spec.mamba_type in copy_funcs, ( |
| 711 | f"missing state copy funcs for {mamba_spec.mamba_type}" |
| 712 | ) |
| 713 | state_copy_funcs = copy_funcs[mamba_spec.mamba_type] |
| 714 | assert 0 < len(state_copy_funcs) <= len(mamba_spec.shapes), ( |
| 715 | f"{mamba_spec.mamba_type} declares {len(mamba_spec.shapes)} states, " |
| 716 | f"but provides {len(state_copy_funcs)} state copy funcs; expected " |
| 717 | "a non-empty copyable prefix" |
| 718 | ) |
| 719 | |
| 720 | |
| 721 | @dataclasses.dataclass |
| 722 | class MambaCopyBuffers: |
| 723 | src_ptrs: CpuGpuBuffer |
| 724 | dst_ptrs: CpuGpuBuffer |
| 725 | sizes: CpuGpuBuffer |
| 726 | mamba_group_ids: list[int] |
| 727 | mamba_spec: MambaSpec |
| 728 | offset: int = 0 |
| 729 | |
| 730 | @classmethod |
| 731 | def create( |
| 732 | cls, |
| 733 | max_num_reqs: int, |
| 734 | kv_cache_config: KVCacheConfig, |
| 735 | copy_funcs: MambaStateCopyFuncsByType, |
| 736 | make_buffer: Callable[..., CpuGpuBuffer], |
| 737 | ) -> "MambaCopyBuffers": |
| 738 | mamba_groups = get_mamba_groups(kv_cache_config) |
| 739 | mamba_spec = next(iter(mamba_groups)) |
| 740 | assert all( |
| 741 | spec.block_size == mamba_spec.block_size |
| 742 | and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks |
| 743 | and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode |
| 744 | for spec in mamba_groups |
| 745 | ), "all mamba groups must share cache scheduling parameters" |
| 746 | mamba_group_ids = get_mamba_group_ids(mamba_groups) |
| 747 | entries_per_req = sum( |
| 748 | len( |
| 749 | copy_funcs[ |
| 750 | _get_mamba_spec_for_layer( |
| 751 | kv_cache_config.kv_cache_groups[gid], layer_name |
| 752 | ).mamba_type |
| 753 | ] |
| 754 | ) |
| 755 | for gid in mamba_group_ids |
| 756 | for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names |
| 757 | ) |
| 758 | n = max_num_reqs * entries_per_req |
| 759 | |
| 760 | return cls( |
| 761 | src_ptrs=make_buffer(n, dtype=torch.uint64), |
| 762 | dst_ptrs=make_buffer(n, dtype=torch.uint64), |
| 763 | sizes=make_buffer(n, dtype=torch.int32), |
| 764 | mamba_group_ids=mamba_group_ids, |
| 765 | mamba_spec=mamba_spec, |
| 766 | ) |
| 767 | |
| 768 | |
| 769 | @dataclasses.dataclass |
| 770 | class MambaSpecDecodeGPUContext: |
| 771 | """Context for GPU-side Mamba state copy operations during the |
| 772 | fused postprocess path. |
| 773 | |
| 774 | Only used when speculative decoding is enabled on a hybrid model |
| 775 | (and the mamba_cache_config is in align mode). |
| 776 | |
| 777 | Precomputes memory layout metadata (base addresses, strides, element sizes) |
| 778 | so the GPU kernel can perform state copies without CPU-GPU sync. |
| 779 | |
| 780 | State types are distinguished by conv_width: >0 for conv states (sliding |
| 781 | window with offset-based copies), 0 for temporal states (full block copies). |
| 782 | """ |
| 783 | |
| 784 | # Per-state metadata tensors (shape: [num_states]). Mamba groups may have |
| 785 | # different numbers of state types (for example, GDN has two while PLE |
| 786 | # short-conv has one). |
| 787 | # These are populated from forward_context during the first forward pass |
| 788 | state_base_addrs: torch.Tensor # int64: base address of each state tensor |
| 789 | state_block_strides: torch.Tensor # int64: bytes per block |
| 790 | state_elem_sizes: torch.Tensor # int32: element size in bytes |
| 791 | state_inner_sizes: torch.Tensor # int64: elements in inner dimensions |
| 792 | state_conv_widths: torch.Tensor # int32: conv width (0 for temporal states) |
| 793 | state_group_indices: torch.Tensor # int32: maps state_idx to group index |
| 794 | # DS conv row metadata. Zero keeps the single-region copy path. |
| 795 | state_dim_row_count: torch.Tensor # int32: per-block dim row count |
| 796 | state_dim_row_stride: torch.Tensor # int64: bytes between rows |
| 797 | |
| 798 | # Configuration |
| 799 | block_size: int |
| 800 | num_states: int |
| 801 | mamba_group_ids: list[int] |
| 802 | num_groups: int |
| 803 | |
| 804 | # Output buffer for num_accepted_tokens updates |
| 805 | num_accepted_tokens_out: torch.Tensor |
| 806 | |
| 807 | # Per-group block-table base addresses: int64[num_groups]. Populated in |
| 808 | # initialize_from_forward_context from the persistent per-group block |
| 809 | # table tensors (whose data_ptr is stable across steps). |
| 810 | block_table_ptrs: torch.Tensor |
| 811 | block_table_stride_req: int = 0 |
| 812 | |
| 813 | # persistent output for the once-per-step, all-group aligned-index launch. |
| 814 | # shape: [num_groups, max_num_reqs, 1 + num_speculative_blocks]. |
| 815 | aligned_state_indices: torch.Tensor | None = None |
| 816 | |
| 817 | # Per-request staging buffers (CPU+GPU mirrors). The runner stages |
| 818 | # values into the CPU view in ``_prepare_inputs`` and the fused kernel |
| 819 | # reads the GPU side. These only exist when the postprocess kernel is |
| 820 | # enabled (spec decode + hybrid + align mode). |
| 821 | mamba_state_idx_buf: CpuGpuBuffer | None = None |
| 822 | num_scheduled_tokens_buf: CpuGpuBuffer | None = None |
| 823 | num_computed_tokens_buf: CpuGpuBuffer | None = None |
| 824 | num_draft_tokens_buf: CpuGpuBuffer | None = None |
| 825 | precopy_src_col_buf: CpuGpuBuffer | None = None |
| 826 | precopy_token_bias_buf: CpuGpuBuffer | None = None |
| 827 | |
| 828 | # Flag to track if metadata has been populated |
| 829 | is_initialized: bool = False |
| 830 | |
| 831 | @classmethod |
| 832 | def create( |
| 833 | cls, |
| 834 | max_num_reqs: int, |
| 835 | kv_cache_config: KVCacheConfig, |
| 836 | copy_funcs: MambaStateCopyFuncsByType, |
| 837 | device: torch.device, |
| 838 | make_buffer: Callable[..., CpuGpuBuffer], |
| 839 | ) -> "MambaSpecDecodeGPUContext": |
| 840 | """Create context with allocated buffers (metadata populated later).""" |
| 841 | mamba_groups = get_mamba_groups(kv_cache_config) |
| 842 | mamba_group_ids = get_mamba_group_ids(mamba_groups) |
| 843 | mamba_spec = next(iter(mamba_groups)) |
| 844 | assert all( |
| 845 | spec.block_size == mamba_spec.block_size |
| 846 | and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks |
| 847 | and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode |
| 848 | for spec in mamba_groups |
| 849 | ), "all mamba groups must share cache scheduling parameters" |
| 850 | copy_funcs_by_spec = { |
| 851 | spec: copy_funcs[spec.mamba_type] for spec in mamba_groups |
| 852 | } |
| 853 | |
| 854 | # Count physical state tensors across all Mamba groups. Different |
| 855 | # groups may expose different state specs. |
| 856 | total_states = sum( |
| 857 | len( |
| 858 | copy_funcs_by_spec[ |
| 859 | _get_mamba_spec_for_layer( |
| 860 | kv_cache_config.kv_cache_groups[gid], layer_name |
| 861 | ) |
| 862 | ] |
| 863 | ) |
| 864 | for gid in mamba_group_ids |
| 865 | for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names |
| 866 | ) |
| 867 | |
| 868 | return cls( |
| 869 | state_base_addrs=torch.zeros( |
| 870 | total_states, dtype=torch.int64, device=device |
| 871 | ), |
| 872 | state_block_strides=torch.zeros( |
| 873 | total_states, dtype=torch.int64, device=device |
| 874 | ), |
| 875 | state_elem_sizes=torch.zeros( |
| 876 | total_states, dtype=torch.int32, device=device |
| 877 | ), |
| 878 | state_inner_sizes=torch.zeros( |
| 879 | total_states, dtype=torch.int64, device=device |
| 880 | ), |
| 881 | state_conv_widths=torch.zeros( |
| 882 | total_states, dtype=torch.int32, device=device |
| 883 | ), |
| 884 | state_group_indices=torch.zeros( |
| 885 | total_states, dtype=torch.int32, device=device |
| 886 | ), |
| 887 | state_dim_row_count=torch.zeros( |
| 888 | total_states, dtype=torch.int32, device=device |
| 889 | ), |
| 890 | state_dim_row_stride=torch.zeros( |
| 891 | total_states, dtype=torch.int64, device=device |
| 892 | ), |
| 893 | block_size=mamba_spec.block_size, |
| 894 | num_states=total_states, |
| 895 | mamba_group_ids=mamba_group_ids, |
| 896 | num_groups=len(mamba_group_ids), |
| 897 | num_accepted_tokens_out=torch.zeros( |
| 898 | max_num_reqs, dtype=torch.int32, device=device |
| 899 | ), |
| 900 | block_table_ptrs=torch.zeros( |
| 901 | len(mamba_group_ids), dtype=torch.int64, device=device |
| 902 | ), |
| 903 | aligned_state_indices=torch.empty( |
| 904 | ( |
| 905 | len(mamba_group_ids), |
| 906 | max_num_reqs, |
| 907 | 1 + mamba_spec.num_speculative_blocks, |
| 908 | ), |
| 909 | dtype=torch.int32, |
| 910 | device=device, |
| 911 | ), |
| 912 | mamba_state_idx_buf=make_buffer(max_num_reqs, dtype=torch.int32), |
| 913 | num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32), |
| 914 | num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32), |
| 915 | num_draft_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32), |
| 916 | precopy_src_col_buf=make_buffer(max_num_reqs, dtype=torch.int32), |
| 917 | precopy_token_bias_buf=make_buffer(max_num_reqs, dtype=torch.int32), |
| 918 | is_initialized=False, |
| 919 | ) |
| 920 | |
| 921 | def initialize_from_forward_context( |
| 922 | self, |
| 923 | kv_cache_config: KVCacheConfig, |
| 924 | forward_context: dict[str, Any], |
| 925 | mamba_state_copy_funcs: MambaStateCopyFuncsByType, |
| 926 | block_tables: list[torch.Tensor], |
| 927 | ) -> None: |
| 928 | """Extract and cache memory layout metadata from Mamba state tensors. |
| 929 | |
| 930 | This method populates the pre-allocated metadata tensors with information |
| 931 | needed by `postprocess_mamba_fused_kernel` to perform state copies entirely |
| 932 | on the GPU without CPU-GPU synchronization. |
| 933 | |
| 934 | For each Mamba layer and state type, the following metadata is extracted: |
| 935 | - state_base_addrs: GPU memory address (data_ptr) of the state tensor |
| 936 | - state_block_strides: Bytes between consecutive blocks (stride * elem_size) |
| 937 | - state_elem_sizes: Element size in bytes (e.g., 2 for float16) |
| 938 | - state_inner_sizes: For conv states, elements per conv position (stride(1)), |
| 939 | used to compute offset when slicing state[block, offset:]. For temporal |
| 940 | states, this field is unused (set to 1). |
| 941 | - state_conv_widths: Conv dimension size for conv states, 0 for temporal states |
| 942 | |
| 943 | The conv vs temporal state type is detected by inspecting the copy function |
| 944 | name: functions containing "conv" are treated as conv states. |
| 945 | |
| 946 | This method is idempotent - it only executes once (guarded by is_initialized |
| 947 | flag) since the metadata is static after model loading. |
| 948 | |
| 949 | Args: |
| 950 | kv_cache_config: Configuration containing KV cache group info and |
| 951 | layer name mappings. |
| 952 | forward_context: Dictionary mapping layer names to attention objects, |
| 953 | populated after the model is loaded. Each attention object must |
| 954 | have a `kv_cache` attribute containing the list of state tensors. |
| 955 | mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to |
| 956 | copy functions. |
| 957 | block_tables: per-mamba-group persistent block-table tensors, in |
| 958 | the same order as `mamba_group_ids`. Their `data_ptr()` / |
| 959 | `stride(0)` are captured once for the kernel to index into. |
| 960 | |
| 961 | """ |
| 962 | if self.is_initialized: |
| 963 | return |
| 964 | # This only runs once per worker. |
| 965 | with gpu_sync_allowed(): |
| 966 | self._populate_metadata( |
| 967 | kv_cache_config, |
| 968 | forward_context, |
| 969 | mamba_state_copy_funcs, |
| 970 | block_tables, |
| 971 | ) |
| 972 | |
| 973 | def _populate_metadata( |
| 974 | self, |
| 975 | kv_cache_config: KVCacheConfig, |
| 976 | forward_context: dict[str, Any], |
| 977 | mamba_state_copy_funcs: MambaStateCopyFuncsByType, |
| 978 | block_tables: list[torch.Tensor], |
| 979 | ) -> None: |
| 980 | idx = 0 |
| 981 | for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): |
| 982 | kv_cache_group = kv_cache_config.kv_cache_groups[mamba_group_id] |
| 983 | layer_names = kv_cache_group.layer_names |
| 984 | for layer_name in layer_names: |
| 985 | mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name) |
| 986 | state_copy_funcs = mamba_state_copy_funcs[mamba_spec.mamba_type] |
| 987 | attention = forward_context[layer_name] |
| 988 | kv_caches: list[torch.Tensor] = attention.kv_cache |
| 989 | if len(kv_caches) < len(state_copy_funcs): |
| 990 | raise ValueError( |
| 991 | f"Expected at least {len(state_copy_funcs)} Mamba state " |
| 992 | f"tensors, got {len(kv_caches)}" |
| 993 | ) |
| 994 | for state_type_idx, copy_func in enumerate(state_copy_funcs): |
| 995 | state = kv_caches[state_type_idx] |
| 996 | # Base address |
| 997 | self.state_base_addrs[idx] = _reinterpret_u64_as_i64( |
| 998 | state.data_ptr() |
| 999 | ) |
| 1000 | |
| 1001 | # Block stride (bytes between consecutive blocks) |
| 1002 | # state shape: [num_blocks, ...], stride(0) = elements per block |
| 1003 | if state.dim() > 1: |
| 1004 | block_stride_elems = state.stride(0) |
| 1005 | else: |
| 1006 | block_stride_elems = state.numel() |
| 1007 | self.state_block_strides[idx] = ( |
| 1008 | block_stride_elems * state.element_size() |
| 1009 | ) |
| 1010 | |
| 1011 | # Element size |
| 1012 | self.state_elem_sizes[idx] = state.element_size() |
| 1013 | |
| 1014 | assert ( |
| 1015 | copy_func is get_conv_copy_spec |
| 1016 | or copy_func is get_temporal_copy_spec |
| 1017 | ), f"unexpected copy func: {copy_func}" |
| 1018 | if copy_func is get_conv_copy_spec: |
| 1019 | if state.dim() != 3: |
| 1020 | raise ValueError( |
| 1021 | "Expected 3D conv state cache, got " |
| 1022 | f"shape {tuple(state.shape)}" |
| 1023 | ) |
| 1024 | if is_conv_state_dim_first(): |
| 1025 | # DS layout: state_len is the slide axis. |
| 1026 | self.state_conv_widths[idx] = state.size(2) |
| 1027 | self.state_inner_sizes[idx] = 1 |
| 1028 | self.state_dim_row_count[idx] = state.size(1) |
| 1029 | self.state_dim_row_stride[idx] = ( |
| 1030 | state.stride(1) * state.element_size() |
| 1031 | ) |
| 1032 | else: |
| 1033 | # SD layout: dim is contiguous. |
| 1034 | self.state_conv_widths[idx] = state.size(1) |
| 1035 | self.state_inner_sizes[idx] = state.stride(1) |
| 1036 | else: |
| 1037 | # Temporal state: inner_size = natural elements per |
| 1038 | # block (prod of inner dims). The kernel uses this |
| 1039 | # to compute copy_size = inner_size * elem_size, |
| 1040 | # which gives the correct byte count even when the |
| 1041 | # state tensor is as_strided with padded page strides |
| 1042 | # (state_block_stride would be the page size, too big). |
| 1043 | self.state_conv_widths[idx] = 0 |
| 1044 | self.state_inner_sizes[idx] = ( |
| 1045 | state[0].numel() if state.dim() > 1 else 1 |
| 1046 | ) |
| 1047 | # Temporal copies vectorize with uint64 loads/stores. |
| 1048 | # The kernel's head/tail handles misalignment for |
| 1049 | # correctness, but unaligned base/stride costs in |
| 1050 | # throughput. |
| 1051 | base_addr = state.data_ptr() |
| 1052 | block_stride_bytes = block_stride_elems * state.element_size() |
| 1053 | if base_addr % 8 != 0: |
| 1054 | logger.warning_once( |
| 1055 | "layer %s: state.data_ptr() = %#x is not " |
| 1056 | "8B-aligned; _memcpy_u64_tiled uint64 " |
| 1057 | "vectorization will pay misaligned load cost", |
| 1058 | layer_name, |
| 1059 | base_addr, |
| 1060 | ) |
| 1061 | if block_stride_bytes % 8 != 0: |
| 1062 | logger.warning_once( |
| 1063 | "layer %s: block stride = %dB is not " |
| 1064 | "8B-aligned; _memcpy_u64_tiled uint64 " |
| 1065 | "vectorization will pay misaligned load cost", |
| 1066 | layer_name, |
| 1067 | block_stride_bytes, |
| 1068 | ) |
| 1069 | |
| 1070 | self.state_group_indices[idx] = group_local_idx |
| 1071 | idx += 1 |
| 1072 | |
| 1073 | assert idx == self.num_states |
| 1074 | |
| 1075 | # Cache per-group block-table base addresses and per-request stride. |
| 1076 | # `block_tables[i]` is the persistent 2D int32 block-table tensor for |
| 1077 | # `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the |
| 1078 | # engine's lifetime, so we capture them once here. |
| 1079 | assert len(block_tables) == self.num_groups, ( |
| 1080 | f"expected {self.num_groups} block tables, got {len(block_tables)}" |
| 1081 | ) |
| 1082 | strides = {bt.stride(0) for bt in block_tables} |
| 1083 | assert len(strides) == 1, ( |
| 1084 | f"all mamba block tables must share stride(0), got {strides}" |
| 1085 | ) |
| 1086 | self.block_table_stride_req = int(next(iter(strides))) |
| 1087 | for i, bt in enumerate(block_tables): |
| 1088 | self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr()) |
| 1089 | |
| 1090 | self.is_initialized = True |
| 1091 | |
| 1092 | def compute_aligned_state_indices( |
| 1093 | self, |
| 1094 | seq_lens: torch.Tensor, |
| 1095 | num_reqs: int, |
| 1096 | ) -> torch.Tensor: |
| 1097 | """Compute every Mamba group's aligned physical state IDs in one launch.""" |
| 1098 | assert self.is_initialized |
| 1099 | assert seq_lens.is_cuda |
| 1100 | assert 0 <= num_reqs <= seq_lens.shape[0] |
| 1101 | assert self.aligned_state_indices is not None |
| 1102 | assert num_reqs <= self.aligned_state_indices.shape[1] |
| 1103 | if num_reqs == 0: |
| 1104 | return self.aligned_state_indices[:, :0] |
| 1105 | |
| 1106 | num_state_slots = self.aligned_state_indices.shape[2] |
| 1107 | block_rows = 32 |
| 1108 | grid = (triton.cdiv(num_reqs, block_rows),) |
| 1109 | get_aligned_state_indices_multi_group_kernel[grid]( |
| 1110 | self.block_table_ptrs, |
| 1111 | seq_lens, |
| 1112 | self.aligned_state_indices, |
| 1113 | self.block_table_stride_req, |
| 1114 | seq_lens.stride(0), |
| 1115 | self.aligned_state_indices.stride(0), |
| 1116 | self.aligned_state_indices.stride(1), |
| 1117 | self.aligned_state_indices.stride(2), |
| 1118 | num_reqs, |
| 1119 | CACHE_BLOCK_SIZE=self.block_size, |
| 1120 | NUM_GROUPS=self.num_groups, |
| 1121 | BLOCK_GROUPS=triton.next_power_of_2(self.num_groups), |
| 1122 | NUM_STATE_SLOTS=num_state_slots, |
| 1123 | BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), |
| 1124 | BLOCK_ROWS=block_rows, |
| 1125 | num_warps=1, |
| 1126 | ) |
| 1127 | return self.aligned_state_indices[:, :num_reqs] |
| 1128 | |
| 1129 | def run_fused_postprocess( |
| 1130 | self, |
| 1131 | num_reqs: int, |
| 1132 | num_accepted_tokens_gpu: torch.Tensor, |
| 1133 | mamba_state_idx_gpu: torch.Tensor, |
| 1134 | num_scheduled_tokens_gpu: torch.Tensor, |
| 1135 | num_computed_tokens_gpu: torch.Tensor, |
| 1136 | num_draft_tokens_gpu: torch.Tensor, |
| 1137 | ) -> None: |
| 1138 | """Run the fused postprocess_mamba kernel on GPU. |
| 1139 | |
| 1140 | This computes decisions and performs mamba state copies entirely on GPU, |
| 1141 | eliminating the CPU-GPU sync that was previously needed. |
| 1142 | |
| 1143 | Args: |
| 1144 | num_reqs: Number of active requests |
| 1145 | num_accepted_tokens_gpu: [num_reqs] accepted token counts |
| 1146 | mamba_state_idx_gpu: [num_reqs] source block indices |
| 1147 | num_scheduled_tokens_gpu: [num_reqs] scheduled token counts |
| 1148 | num_computed_tokens_gpu: [num_reqs] computed token counts |
| 1149 | num_draft_tokens_gpu: [num_reqs] draft token counts |
| 1150 | |
| 1151 | """ |
| 1152 | if num_reqs == 0 or not self.is_initialized: |
| 1153 | return |
| 1154 | |
| 1155 | # Initialize output to current values (unchanged unless src==dst) |
| 1156 | self.num_accepted_tokens_out[:num_reqs].copy_( |
| 1157 | num_accepted_tokens_gpu[:num_reqs] |
| 1158 | ) |
| 1159 | |
| 1160 | total_states = self.num_states |
| 1161 | grid = (num_reqs, total_states, _TEMPORAL_TILES) |
| 1162 | |
| 1163 | postprocess_mamba_fused_kernel[grid]( |
| 1164 | num_accepted_tokens_gpu, |
| 1165 | mamba_state_idx_gpu, |
| 1166 | num_scheduled_tokens_gpu, |
| 1167 | num_computed_tokens_gpu, |
| 1168 | num_draft_tokens_gpu, |
| 1169 | self.block_table_ptrs, |
| 1170 | self.block_table_stride_req, |
| 1171 | self.state_base_addrs, |
| 1172 | self.state_block_strides, |
| 1173 | self.state_elem_sizes, |
| 1174 | self.state_inner_sizes, |
| 1175 | self.state_conv_widths, |
| 1176 | self.state_group_indices, |
| 1177 | self.state_dim_row_count, |
| 1178 | self.state_dim_row_stride, |
| 1179 | self.num_accepted_tokens_out, |
| 1180 | None, # idx_mapping: V1 decision arrays are already in req order |
| 1181 | num_reqs, |
| 1182 | block_size=self.block_size, |
| 1183 | COPY_BLOCK_SIZE=1024, |
| 1184 | CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), |
| 1185 | TEMPORAL_TILES=_TEMPORAL_TILES, |
| 1186 | ) |
| 1187 | |
| 1188 | def run_fused_precopy( |
| 1189 | self, |
| 1190 | num_reqs: int, |
| 1191 | state_idx_gpu: torch.Tensor, |
| 1192 | src_col_gpu: torch.Tensor, |
| 1193 | token_bias_gpu: torch.Tensor, |
| 1194 | idx_mapping: torch.Tensor | None, |
| 1195 | ) -> None: |
| 1196 | """Pre-copy each request's previous running block into its new window |
| 1197 | block before the forward pass (align boundary migration). |
| 1198 | |
| 1199 | Args: |
| 1200 | num_reqs: Number of active requests (batch order). |
| 1201 | state_idx_gpu: [max_reqs] post-advance dst block column per req slot. |
| 1202 | src_col_gpu: [max_reqs] pre-advance src block column (-1 = fresh). |
| 1203 | token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1). |
| 1204 | idx_mapping: optional [num_reqs] batch_idx -> req_state_idx. |
| 1205 | None means V1 batch order already equals request state order. |
| 1206 | |
| 1207 | """ |
| 1208 | if num_reqs == 0 or not self.is_initialized: |
| 1209 | return |
| 1210 | total_states = self.num_states |
| 1211 | grid = (num_reqs, total_states, _TEMPORAL_TILES) |
| 1212 | precopy_mamba_align_fused_kernel[grid]( |
| 1213 | state_idx_gpu, |
| 1214 | src_col_gpu, |
| 1215 | token_bias_gpu, |
| 1216 | self.block_table_ptrs, |
| 1217 | self.block_table_stride_req, |
| 1218 | self.state_base_addrs, |
| 1219 | self.state_block_strides, |
| 1220 | self.state_elem_sizes, |
| 1221 | self.state_inner_sizes, |
| 1222 | self.state_conv_widths, |
| 1223 | self.state_group_indices, |
| 1224 | self.state_dim_row_count, |
| 1225 | self.state_dim_row_stride, |
| 1226 | idx_mapping, |
| 1227 | num_reqs, |
| 1228 | COPY_BLOCK_SIZE=1024, |
| 1229 | CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), |
| 1230 | HAS_IDX_MAPPING=idx_mapping is not None, |
| 1231 | TEMPORAL_TILES=_TEMPORAL_TILES, |
| 1232 | ) |
| 1233 | |
| 1234 | def run_fused_postprocess_align( |
| 1235 | self, |
| 1236 | num_reqs: int, |
| 1237 | num_accepted_tokens_gpu: torch.Tensor, |
| 1238 | state_idx_gpu: torch.Tensor, |
| 1239 | new_num_computed_tokens_gpu: torch.Tensor, |
| 1240 | idx_mapping: torch.Tensor, |
| 1241 | ) -> None: |
| 1242 | """V2 align postprocess: save the running state to the block-aligned |
| 1243 | position after spec-decode acceptance leaves the sequence non-aligned. |
| 1244 | |
| 1245 | ``num_accepted_tokens_gpu`` is updated in place while the kernel reads |
| 1246 | from a snapshot to avoid cross-program races when the accepted position |
| 1247 | stays in the running block and the count is reset to 1. |
| 1248 | ``new_num_computed_tokens`` already holds the post-step computed count |
| 1249 | (PRECOMPUTED_NEW_COMPUTED). |
| 1250 | ``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING). |
| 1251 | """ |
| 1252 | if num_reqs == 0 or not self.is_initialized: |
| 1253 | return |
| 1254 | |
| 1255 | # V2 reads non-contiguous idx_mapping positions, so snapshot the whole |
| 1256 | # decision buffer rather than only [:num_reqs]. |
| 1257 | num_accepted_tokens_snapshot = self.num_accepted_tokens_out |
| 1258 | num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) |
| 1259 | |
| 1260 | total_states = self.num_states |
| 1261 | grid = (num_reqs, total_states, _TEMPORAL_TILES) |
| 1262 | postprocess_mamba_fused_kernel[grid]( |
| 1263 | num_accepted_tokens_snapshot, |
| 1264 | state_idx_gpu, |
| 1265 | None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED |
| 1266 | new_num_computed_tokens_gpu, |
| 1267 | None, # num_draft: unused under PRECOMPUTED_NEW_COMPUTED |
| 1268 | self.block_table_ptrs, |
| 1269 | self.block_table_stride_req, |
| 1270 | self.state_base_addrs, |
| 1271 | self.state_block_strides, |
| 1272 | self.state_elem_sizes, |
| 1273 | self.state_inner_sizes, |
| 1274 | self.state_conv_widths, |
| 1275 | self.state_group_indices, |
| 1276 | self.state_dim_row_count, |
| 1277 | self.state_dim_row_stride, |
| 1278 | num_accepted_tokens_gpu, |
| 1279 | idx_mapping, |
| 1280 | num_reqs, |
| 1281 | block_size=self.block_size, |
| 1282 | COPY_BLOCK_SIZE=1024, |
| 1283 | CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), |
| 1284 | HAS_IDX_MAPPING=True, |
| 1285 | PRECOMPUTED_NEW_COMPUTED=True, |
| 1286 | TEMPORAL_TILES=_TEMPORAL_TILES, |
| 1287 | ) |
| 1288 | |
| 1289 | |
| 1290 | @dataclasses.dataclass |
| 1291 | class MambaBuffers: |
| 1292 | """Single owner for all mamba-specific runner buffers. |
| 1293 | |
| 1294 | The two sub-objects have different gates: |
| 1295 | ``preprocess`` is needed whenever ``mamba_cache_mode == "align"``; |
| 1296 | ``postprocess_align`` is needed only when align is combined with |
| 1297 | speculative decoding on a hybrid model, and is ``None`` otherwise. |
| 1298 | """ |
| 1299 | |
| 1300 | preprocess: MambaCopyBuffers |
| 1301 | postprocess_align: MambaSpecDecodeGPUContext | None |
| 1302 | |
| 1303 | @classmethod |
| 1304 | def create( |
| 1305 | cls, |
| 1306 | max_num_reqs: int, |
| 1307 | kv_cache_config: KVCacheConfig, |
| 1308 | copy_funcs: MambaStateCopyFuncsByType, |
| 1309 | make_buffer: Callable[..., CpuGpuBuffer], |
| 1310 | device: torch.device, |
| 1311 | with_postprocess_align: bool, |
| 1312 | ) -> "MambaBuffers": |
| 1313 | return cls( |
| 1314 | preprocess=MambaCopyBuffers.create( |
| 1315 | max_num_reqs, kv_cache_config, copy_funcs, make_buffer |
| 1316 | ), |
| 1317 | postprocess_align=( |
| 1318 | MambaSpecDecodeGPUContext.create( |
| 1319 | max_num_reqs=max_num_reqs, |
| 1320 | kv_cache_config=kv_cache_config, |
| 1321 | copy_funcs=copy_funcs, |
| 1322 | device=device, |
| 1323 | make_buffer=make_buffer, |
| 1324 | ) |
| 1325 | if with_postprocess_align |
| 1326 | else None |
| 1327 | ), |
| 1328 | ) |
| 1329 | |
| 1330 | |
| 1331 | def collect_mamba_copy_meta( |
| 1332 | copy_bufs: MambaCopyBuffers, |
| 1333 | kv_cache_config: KVCacheConfig, |
| 1334 | mamba_state_copy_funcs: MambaStateCopyFuncsByType, |
| 1335 | mamba_group_ids: list[int], |
| 1336 | src_block_idx: int, |
| 1337 | dest_block_idx: int, |
| 1338 | accept_token_bias: int, |
| 1339 | req_state: CachedRequestState, |
| 1340 | forward_context: dict[str, Any], |
| 1341 | ) -> None: |
| 1342 | if src_block_idx == dest_block_idx and accept_token_bias == 0: |
| 1343 | return |
| 1344 | |
| 1345 | src_ptrs_np = copy_bufs.src_ptrs.np |
| 1346 | dst_ptrs_np = copy_bufs.dst_ptrs.np |
| 1347 | sizes_np = copy_bufs.sizes.np |
| 1348 | offset = copy_bufs.offset |
| 1349 | |
| 1350 | for mamba_group_id in mamba_group_ids: |
| 1351 | block_ids = req_state.block_ids[mamba_group_id] |
| 1352 | dest_block_id = block_ids[dest_block_idx] |
| 1353 | kv_cache_group = kv_cache_config.kv_cache_groups[mamba_group_id] |
| 1354 | layer_names = kv_cache_group.layer_names |
| 1355 | for layer_name in layer_names: |
| 1356 | mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name) |
| 1357 | state_copy_funcs = mamba_state_copy_funcs[mamba_spec.mamba_type] |
| 1358 | attention = forward_context[layer_name] |
| 1359 | kv_caches: list[torch.Tensor] = attention.kv_cache |
| 1360 | for state, state_copy_func in zip(kv_caches, state_copy_funcs): |
| 1361 | copy_spec = state_copy_func( |
| 1362 | state, block_ids, src_block_idx, accept_token_bias + 1 |
| 1363 | ) |
| 1364 | |
| 1365 | src_ptrs_np[offset] = copy_spec.start_addr |
| 1366 | dst_ptrs_np[offset] = state[dest_block_id].data_ptr() |
| 1367 | sizes_np[offset] = copy_spec.num_elements * state.element_size() |
| 1368 | offset += 1 |
| 1369 | |
| 1370 | copy_bufs.offset = offset |
| 1371 | |
| 1372 | |
| 1373 | def do_mamba_copy_block(copy_bufs: MambaCopyBuffers): |
| 1374 | n = copy_bufs.offset |
| 1375 | if n == 0: |
| 1376 | return |
| 1377 | batch_memcpy( |
| 1378 | copy_bufs.src_ptrs.copy_to_gpu(n), |
| 1379 | copy_bufs.dst_ptrs.copy_to_gpu(n), |
| 1380 | copy_bufs.sizes.copy_to_gpu(n), |
| 1381 | ) |
| 1382 | |
| 1383 | |
| 1384 | def cleanup_mamba_state_idx( |
| 1385 | scheduler_output: SchedulerOutput, |
| 1386 | mamba_state_idx: dict[str, int], |
| 1387 | ) -> None: |
| 1388 | """Pop stale `mamba_state_idx` entries for finished/preempted/resumed reqs. |
| 1389 | |
| 1390 | Force-preempted requests (e.g., during reset_prefix_cache / KV cache |
| 1391 | flush) appear in resumed_req_ids without a corresponding entry in |
| 1392 | preempted_req_ids, leaving stale entries that can point to block |
| 1393 | indices beyond the new (smaller) block allocation. |
| 1394 | """ |
| 1395 | finished_req_ids = scheduler_output.finished_req_ids |
| 1396 | preempted_req_ids = scheduler_output.preempted_req_ids or set() |
| 1397 | resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids |
| 1398 | for req_id in itertools.chain(finished_req_ids, preempted_req_ids, resumed_req_ids): |
| 1399 | mamba_state_idx.pop(req_id, None) |
| 1400 | |
| 1401 | |
| 1402 | class _FusedPrecopy(NamedTuple): |
| 1403 | """Resolved fused align pre-copy resources (all non-None once resolved).""" |
| 1404 | |
| 1405 | ctx: "MambaSpecDecodeGPUContext" |
| 1406 | state_idx: CpuGpuBuffer |
| 1407 | src_col: CpuGpuBuffer |
| 1408 | token_bias: CpuGpuBuffer |
| 1409 | |
| 1410 | |
| 1411 | def _resolve_fused_precopy( |
| 1412 | align_ctx: "MambaSpecDecodeGPUContext | None", |
| 1413 | ) -> _FusedPrecopy | None: |
| 1414 | """Bundle the fused-path buffers, or None for the scalar path. |
| 1415 | |
| 1416 | Returning one non-None bundle lets callers narrow all four members with a |
| 1417 | single ``is not None`` check instead of re-asserting each buffer per use. |
| 1418 | """ |
| 1419 | if align_ctx is None: |
| 1420 | return None |
| 1421 | assert align_ctx.mamba_state_idx_buf is not None |
| 1422 | assert align_ctx.precopy_src_col_buf is not None |
| 1423 | assert align_ctx.precopy_token_bias_buf is not None |
| 1424 | return _FusedPrecopy( |
| 1425 | align_ctx, |
| 1426 | align_ctx.mamba_state_idx_buf, |
| 1427 | align_ctx.precopy_src_col_buf, |
| 1428 | align_ctx.precopy_token_bias_buf, |
| 1429 | ) |
| 1430 | |
| 1431 | |
| 1432 | def preprocess_mamba( |
| 1433 | scheduler_output: SchedulerOutput, |
| 1434 | kv_cache_config: KVCacheConfig, |
| 1435 | cache_config: CacheConfig, |
| 1436 | mamba_state_idx: dict[str, int], |
| 1437 | input_batch: InputBatch, |
| 1438 | requests: dict[str, CachedRequestState], |
| 1439 | forward_context: dict[str, Any], |
| 1440 | mamba_state_copy_funcs: MambaStateCopyFuncsByType, |
| 1441 | copy_bufs: MambaCopyBuffers, |
| 1442 | align_ctx: MambaSpecDecodeGPUContext | None = None, |
| 1443 | ): |
| 1444 | """Copy the mamba state of previous step to the last |
| 1445 | (1 + num_speculative_blocks) block. |
| 1446 | """ |
| 1447 | fused = _resolve_fused_precopy(align_ctx) |
| 1448 | mamba_group_ids = copy_bufs.mamba_group_ids |
| 1449 | mamba_spec = copy_bufs.mamba_spec |
| 1450 | num_speculative_blocks = mamba_spec.num_speculative_blocks |
| 1451 | # TODO(Chen): we need to optimize this function a lot |
| 1452 | assert cache_config.enable_prefix_caching |
| 1453 | block_size = mamba_spec.block_size |
| 1454 | cleanup_mamba_state_idx(scheduler_output, mamba_state_idx) |
| 1455 | |
| 1456 | copy_bufs.offset = 0 |
| 1457 | num_reqs = len(input_batch.req_ids) |
| 1458 | |
| 1459 | if fused is not None: |
| 1460 | if num_reqs == 0: |
| 1461 | return |
| 1462 | if not fused.ctx.is_initialized: |
| 1463 | fused.ctx.initialize_from_forward_context( |
| 1464 | kv_cache_config, |
| 1465 | forward_context, |
| 1466 | mamba_state_copy_funcs, |
| 1467 | [ |
| 1468 | input_batch.block_table[gid].get_device_tensor(num_reqs) |
| 1469 | for gid in fused.ctx.mamba_group_ids |
| 1470 | ], |
| 1471 | ) |
| 1472 | |
| 1473 | fused.src_col.np[:num_reqs] = -1 |
| 1474 | fused.token_bias.np[:num_reqs] = 0 |
| 1475 | |
| 1476 | for i, req_id in enumerate(input_batch.req_ids): |
| 1477 | req_state = requests[req_id] |
| 1478 | prev_state_idx = mamba_state_idx.get(req_id) |
| 1479 | if prev_state_idx is None: |
| 1480 | # New / resumed request; num_computed_tokens == 0 gives -1. |
| 1481 | prev_state_idx = (req_state.num_computed_tokens - 1) // block_size |
| 1482 | |
| 1483 | num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] |
| 1484 | num_blocks = ( |
| 1485 | cdiv(req_state.num_computed_tokens + num_scheduled_tokens, block_size) |
| 1486 | + num_speculative_blocks |
| 1487 | ) |
| 1488 | # We always save the current running state at the last |
| 1489 | # (1 + num_speculative_blocks) block. |
| 1490 | # A corner case worth mention here: assume we have block_size = 4 and |
| 1491 | # num_speculative_tokens = 2. The request is [A, B, C] and contains 2 draft |
| 1492 | # tokens [draft 1, draft 2]. Then we will have: |
| 1493 | # Block 0: [A, B, C, draft 1] |
| 1494 | # Block 1: [draft 2, TOFILL, TOFILL, TOFILL] |
| 1495 | # Block 2: speculative block |
| 1496 | # Block 3: speculative block |
| 1497 | # And use block 1 to save the running state. |
| 1498 | curr_state_idx = num_blocks - 1 - num_speculative_blocks |
| 1499 | mamba_state_idx[req_id] = curr_state_idx |
| 1500 | if fused is not None: |
| 1501 | fused.state_idx.np[i] = curr_state_idx |
| 1502 | |
| 1503 | if prev_state_idx != -1 and prev_state_idx != curr_state_idx: |
| 1504 | accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 |
| 1505 | if fused is not None: |
| 1506 | assert accept_token_bias >= 0 |
| 1507 | fused.src_col.np[i] = prev_state_idx |
| 1508 | fused.token_bias.np[i] = accept_token_bias |
| 1509 | else: |
| 1510 | collect_mamba_copy_meta( |
| 1511 | copy_bufs, |
| 1512 | kv_cache_config, |
| 1513 | mamba_state_copy_funcs, |
| 1514 | mamba_group_ids, |
| 1515 | prev_state_idx, |
| 1516 | curr_state_idx, |
| 1517 | accept_token_bias, |
| 1518 | req_state, |
| 1519 | forward_context, |
| 1520 | ) |
| 1521 | input_batch.num_accepted_tokens_cpu[i] = 1 |
| 1522 | |
| 1523 | if fused is not None: |
| 1524 | fused.state_idx.copy_to_gpu(num_reqs) |
| 1525 | fused.src_col.copy_to_gpu(num_reqs) |
| 1526 | fused.token_bias.copy_to_gpu(num_reqs) |
| 1527 | fused.ctx.run_fused_precopy( |
| 1528 | num_reqs=num_reqs, |
| 1529 | state_idx_gpu=fused.state_idx.gpu, |
| 1530 | src_col_gpu=fused.src_col.gpu, |
| 1531 | token_bias_gpu=fused.token_bias.gpu, |
| 1532 | idx_mapping=None, |
| 1533 | ) |
| 1534 | else: |
| 1535 | do_mamba_copy_block(copy_bufs) |
| 1536 | |
| 1537 | |
| 1538 | def postprocess_mamba_align_gpu( |
| 1539 | *, |
| 1540 | bufs: "MambaBuffers", |
| 1541 | num_reqs: int, |
| 1542 | num_accepted_tokens_gpu: torch.Tensor, |
| 1543 | num_accepted_tokens_cpu_tensor: torch.Tensor, |
| 1544 | input_batch: InputBatch, |
| 1545 | kv_cache_config: KVCacheConfig, |
| 1546 | forward_context: dict[str, Any], |
| 1547 | mamba_state_copy_funcs: MambaStateCopyFuncsByType, |
| 1548 | ) -> None: |
| 1549 | """GPU-side mamba postprocess for spec decode + hybrid + align mode. |
| 1550 | |
| 1551 | Lazily binds the fused-kernel context to the persistent block tables and |
| 1552 | forward-context state pointers on the first call, runs the fused kernel, |
| 1553 | and async-copies the per-request accepted-token counts back to the input |
| 1554 | batch's CPU tensor for the next iteration's preprocess. |
| 1555 | """ |
| 1556 | ctx = bufs.postprocess_align |
| 1557 | # Caller is responsible for gating on spec decode + hybrid; this assert is |
| 1558 | # a tripwire if those gates ever drift apart. |
| 1559 | assert ctx is not None |
| 1560 | assert ctx.mamba_state_idx_buf is not None |
| 1561 | assert ctx.num_scheduled_tokens_buf is not None |
| 1562 | assert ctx.num_computed_tokens_buf is not None |
| 1563 | assert ctx.num_draft_tokens_buf is not None |
| 1564 | |
| 1565 | if not ctx.is_initialized: |
| 1566 | ctx.initialize_from_forward_context( |
| 1567 | kv_cache_config, |
| 1568 | forward_context, |
| 1569 | mamba_state_copy_funcs, |
| 1570 | [ |
| 1571 | input_batch.block_table[gid].get_device_tensor(num_reqs) |
| 1572 | for gid in ctx.mamba_group_ids |
| 1573 | ], |
| 1574 | ) |
| 1575 | |
| 1576 | ctx.run_fused_postprocess( |
| 1577 | num_reqs=num_reqs, |
| 1578 | num_accepted_tokens_gpu=num_accepted_tokens_gpu, |
| 1579 | mamba_state_idx_gpu=ctx.mamba_state_idx_buf.gpu, |
| 1580 | num_scheduled_tokens_gpu=ctx.num_scheduled_tokens_buf.gpu, |
| 1581 | num_computed_tokens_gpu=ctx.num_computed_tokens_buf.gpu, |
| 1582 | num_draft_tokens_gpu=ctx.num_draft_tokens_buf.gpu, |
| 1583 | ) |
| 1584 | |
| 1585 | # ``num_accepted_tokens_out`` is pre-initialized from |
| 1586 | # ``num_accepted_tokens_gpu``; the kernel only overwrites entries to 1 |
| 1587 | # when src_block_idx == dest_block_idx (copy within the same block), so |
| 1588 | # the original count is preserved for everyone else. |
| 1589 | num_accepted_tokens_cpu_tensor[:num_reqs].copy_( |
| 1590 | ctx.num_accepted_tokens_out[:num_reqs], non_blocking=True |
| 1591 | ) |
| 1592 | |
| 1593 | |
| 1594 | def stage_postprocess_inputs_to_gpu( |
| 1595 | ctx: MambaSpecDecodeGPUContext, |
| 1596 | scheduler_output: SchedulerOutput, |
| 1597 | req_ids: list[str], |
| 1598 | num_reqs: int, |
| 1599 | requests: dict[str, CachedRequestState], |
| 1600 | mamba_state_idx: dict[str, int], |
| 1601 | ) -> None: |
| 1602 | """Stage all per-request inputs the fused mamba postprocess kernel reads. |
| 1603 | |
| 1604 | Walks ``req_ids[:num_reqs]`` once, writing each request's mamba block |
| 1605 | index and scheduled/computed/draft token counts into the matching pinned |
| 1606 | numpy views, then issues four non-blocking H→D copies. The fused kernel |
| 1607 | indexes the resulting GPU tensors by ``req_idx``. Buffers live on ``ctx`` |
| 1608 | and only exist when the postprocess kernel is enabled. |
| 1609 | |
| 1610 | Invariant: ``preprocess_mamba`` must have run first for the same batch so |
| 1611 | that every ``req_ids[i]`` has an entry in ``mamba_state_idx``. |
| 1612 | """ |
| 1613 | assert ctx.mamba_state_idx_buf is not None |
| 1614 | assert ctx.num_scheduled_tokens_buf is not None |
| 1615 | assert ctx.num_computed_tokens_buf is not None |
| 1616 | assert ctx.num_draft_tokens_buf is not None |
| 1617 | |
| 1618 | scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens |
| 1619 | num_scheduled = scheduler_output.num_scheduled_tokens |
| 1620 | state_idx_np = ctx.mamba_state_idx_buf.np |
| 1621 | scheduled_np = ctx.num_scheduled_tokens_buf.np |
| 1622 | computed_np = ctx.num_computed_tokens_buf.np |
| 1623 | draft_np = ctx.num_draft_tokens_buf.np |
| 1624 | |
| 1625 | for i in range(num_reqs): |
| 1626 | req_id = req_ids[i] |
| 1627 | state_idx = mamba_state_idx.get(req_id) |
| 1628 | assert state_idx is not None, ( |
| 1629 | f"mamba_state_idx missing entry for {req_id!r}; " |
| 1630 | "preprocess_mamba must run before stage_postprocess_inputs_to_gpu" |
| 1631 | ) |
| 1632 | state_idx_np[i] = state_idx |
| 1633 | scheduled_np[i] = num_scheduled[req_id] |
| 1634 | computed_np[i] = requests[req_id].num_computed_tokens |
| 1635 | draft_np[i] = len(scheduled_spec_tokens.get(req_id, [])) |
| 1636 | |
| 1637 | ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) |
| 1638 | ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs) |
| 1639 | ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs) |
| 1640 | ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs) |