@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4bf699d310

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Kernel] Support DS Mamba tail copy for MTP align mode (#45473)

Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Co-authored-by: Thomas Parnell <tom.parnell@gmail.com>

sungsoo ha committed Jun 16, 2026 at 15:50 UTC 4bf699d310300edf1df2df2713e3d6516438a798

3 files changed +182 -34

tests/v1/worker/test_mamba_utils.py

+115

@@ -2132,3 +2132,118 @@ class TestPostprocessMambaFusedKernel:

2132 expected_accepted,

2133 msg="num_accepted_tokens mismatch at accept_token_bias=2",

2134 )

2135 +

2136 + def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd(

2137 + self, device, test_config, monkeypatch

2138 + ):

2139 + """DS conv postprocess should match SD when accept_token_bias > 0."""

2140 + from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils

2141 +

2142 + cfg = test_config

2143 + torch.manual_seed(38898)

2144 +

2145 + req_ids = ["req_0"]

2146 + num_computed_tokens = [30]

2147 + num_scheduled_tokens = {"req_0": 1}

2148 + num_draft_tokens: dict[str, int] = {}

2149 + num_accepted_tokens = [2] # Results in accept_token_bias = 1

2150 + mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx

2151 + block_ids_per_req = [list(range(8))]

2152 +

2153 + layer_names = ["layer_0"]

2154 + kv_cache_config = _make_kv_cache_config(cfg, layer_names)

2155 +

2156 + num_reqs = len(req_ids)

2157 + block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)

2158 + block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

2159 +

2160 + # Same logical conv state in SD and DS layouts.

2161 + sd_source_conv = torch.randn(

2162 + cfg.num_blocks,

2163 + cfg.conv_width,

2164 + cfg.conv_inner_dim,

2165 + dtype=cfg.dtype,

2166 + device=device,

2167 + )

2168 + ds_source_conv = sd_source_conv.permute(0, 2, 1).contiguous()

2169 + sd_source_temporal = torch.randn(

2170 + cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device

2171 + )

2172 +

2173 + # SD GPU path. Default layout is SD.

2174 + model_mamba_utils.get_conv_state_layout.cache_clear()

2175 + sd_conv = sd_source_conv.clone()

2176 + sd_temporal = sd_source_temporal.clone()

2177 + forward_context_sd = {

2178 + "layer_0": _make_mock_attention(sd_conv, sd_temporal),

2179 + }

2180 + gpu_ctx_sd = _make_gpu_ctx(cfg, kv_cache_config, device)

2181 + _run_gpu_postprocess(

2182 + gpu_ctx_sd,

2183 + kv_cache_config=kv_cache_config,

2184 + forward_context=forward_context_sd,

2185 + copy_funcs=_COPY_FUNCS,

2186 + block_table=block_table_gpu,

2187 + req_ids=req_ids,

2188 + num_accepted_tokens=num_accepted_tokens,

2189 + mamba_state_idx=mamba_state_idx,

2190 + num_scheduled_tokens=num_scheduled_tokens,

2191 + num_computed_tokens=num_computed_tokens,

2192 + num_draft_tokens=num_draft_tokens,

2193 + device=device,

2194 + )

2195 + torch.accelerator.synchronize()

2196 +

2197 + # Sanity: SD path actually modified the state (copy was performed).

2198 + assert not torch.equal(sd_conv, sd_source_conv), (

2199 + "SD baseline did not modify conv state; test setup is wrong"

2200 + )

2201 +

2202 + # DS GPU path on the DS twin.

2203 + monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "DS")

2204 + model_mamba_utils.get_conv_state_layout.cache_clear()

2205 + try:

2206 + ds_conv = ds_source_conv.clone()

2207 + ds_temporal = sd_source_temporal.clone()

2208 + forward_context_ds = {

2209 + "layer_0": _make_mock_attention(ds_conv, ds_temporal),

2210 + }

2211 + gpu_ctx_ds = _make_gpu_ctx(cfg, kv_cache_config, device)

2212 + _run_gpu_postprocess(

2213 + gpu_ctx_ds,

2214 + kv_cache_config=kv_cache_config,

2215 + forward_context=forward_context_ds,

2216 + copy_funcs=_COPY_FUNCS,

2217 + block_table=block_table_gpu,

2218 + req_ids=req_ids,

2219 + num_accepted_tokens=num_accepted_tokens,

2220 + mamba_state_idx=mamba_state_idx,

2221 + num_scheduled_tokens=num_scheduled_tokens,

2222 + num_computed_tokens=num_computed_tokens,

2223 + num_draft_tokens=num_draft_tokens,

2224 + device=device,

2225 + )

2226 + torch.accelerator.synchronize()

2227 + finally:

2228 + # Reset the lru cache so other tests see the default layout again.

2229 + model_mamba_utils.get_conv_state_layout.cache_clear()

2230 +

2231 + # DS bytes, un-permuted, should match the SD result.

2232 + torch.testing.assert_close(

2233 + ds_conv.permute(0, 2, 1).contiguous(),

2234 + sd_conv,

2235 + msg=(

2236 + "DS conv post-kernel does not match SD baseline; the DS "

2237 + "row-loop in postprocess_mamba_fused_kernel is wrong."

2238 + ),

2239 + )

2240 + torch.testing.assert_close(

2241 + ds_temporal,

2242 + sd_temporal,

2243 + msg="DS temporal state diverged from SD",

2244 + )

2245 + torch.testing.assert_close(

2246 + gpu_ctx_ds.num_accepted_tokens_out[:num_reqs],

2247 + gpu_ctx_sd.num_accepted_tokens_out[:num_reqs],

2248 + msg="DS num_accepted_tokens diverged from SD",

2249 + )

vllm/model_executor/layers/mamba/mamba_utils.py

+6 -10

@@ -303,18 +303,14 @@ def get_conv_copy_spec(

303 src_block_id = block_ids[cur_block_idx]

304 offset = num_accepted_tokens - 1

305 if is_conv_state_dim_first():

306 - # DS layout: (num_blocks, dim, state_len) — state_len is last.

307 - if offset > 0:

308 - # Slicing along the last dim yields a non-contiguous view

309 - # because features (dim) are strided by state_len.

310 - raise NotImplementedError(

311 - "DS conv state layout does not yet support speculative "

312 - "decoding with mamba_cache_mode='align' "

313 - "(num_accepted_tokens > 1)."

314 - )

306 + # DS offset > 0 is handled by the fused postprocess kernel.

307 + assert offset == 0, (

308 + "DS conv state with num_accepted_tokens > 1 must be handled by "

309 + "the fused postprocess kernel, not get_conv_copy_spec"

310 + )

311 src_state = state[src_block_id]

312 else:

317 - # SD layout: (num_blocks, state_len, dim) — dim contiguous.

313 + # SD layout: (num_blocks, state_len, dim), with dim contiguous.

314 src_state = state[src_block_id, offset:]

315 return MambaCopySpec(

316 start_addr=src_state.data_ptr(), num_elements=src_state.numel()

vllm/v1/worker/mamba_utils.py

+61 -24

@@ -12,6 +12,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (

12 MambaStateCopyFunc,

13 get_conv_copy_spec,

14 get_temporal_copy_spec,

15 + is_conv_state_dim_first,

16 )

17 from vllm.triton_utils import tl, triton

18 from vllm.utils.math_utils import cdiv

@@ -43,6 +44,9 @@ def postprocess_mamba_fused_kernel(

44 state_inner_sizes_ptr, # number of elements in inner dimensions

45 state_conv_widths_ptr, # conv width for conv states (0 for temporal)

46 state_group_indices_ptr, # maps state_idx to group index in block table

47 + # DS conv row metadata. Zero keeps the single-region copy path.

48 + state_dim_row_count_ptr, # int32: per-block dim row count for DS conv

49 + state_dim_row_stride_ptr, # int64: bytes between rows for DS conv

50 # Output: num_accepted_tokens update (for src==dst case)

51 num_accepted_tokens_out_ptr,

52 # Runtime parameter (varies per batch - NOT constexpr to avoid recompilation)

@@ -52,6 +56,7 @@ def postprocess_mamba_fused_kernel(

56 block_size: tl.constexpr,

57 # COPY_BLOCK_SIZE: fixed tuning parameter for memory copy loop

58 COPY_BLOCK_SIZE: tl.constexpr,

59 + CONV_STATE_DIM_FIRST: tl.constexpr,

60 ):

61 """

62 Fused GPU kernel for postprocess_mamba that computes decisions AND performs

@@ -121,8 +126,35 @@ def postprocess_mamba_fused_kernel(

126 # conv_width == 0 means this is a temporal state (get_temporal_copy_spec logic)

127 is_conv_state = conv_width > 0

128

129 + # Update accepted-token count before early exits.

130 + if src_block_idx == dest_block_idx and state_idx == 0:

131 + tl.store(num_accepted_tokens_out_ptr + req_idx, 1)

132 +

133 + # Skip no-op self-copy.

134 + if src_block_idx == dest_block_idx and accept_token_bias == 0:

135 + return

136 +

137 + if CONV_STATE_DIM_FIRST and is_conv_state:

138 + dim_rows = tl.load(state_dim_row_count_ptr + state_idx)

139 + row_stride = tl.load(state_dim_row_stride_ptr + state_idx)

140 + per_row_bytes = (conv_width - accept_token_bias).to(tl.int64) * state_elem_size

141 + bias_bytes = accept_token_bias.to(tl.int64) * state_elem_size

142 + src_block_addr = state_base_addr + src_block_id * state_block_stride

143 + dst_block_addr = state_base_addr + dest_block_id * state_block_stride

144 + offsets = tl.arange(0, COPY_BLOCK_SIZE)

145 + for d in range(0, dim_rows):

146 + row_src = src_block_addr + d * row_stride + bias_bytes

147 + row_dst = dst_block_addr + d * row_stride

148 + for i in range(0, per_row_bytes, COPY_BLOCK_SIZE):

149 + mask = (i + offsets) < per_row_bytes

150 + curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8))

151 + curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8))

152 + data = tl.load(curr_src, mask=mask)

153 + tl.store(curr_dst, data, mask=mask)

154 + return

155 +

156 if is_conv_state:

125 - # Conv state: copy

157 + # SD conv: copy

158 # state[block_table[req_idx, src_block_idx], accept_token_bias:]

159 # to

160 # state[block_table[req_idx, dest_block_idx], :conv_width - accept_token_bias]

@@ -151,19 +183,6 @@ def postprocess_mamba_fused_kernel(

183 # actual data when the state tensor uses as_strided page padding.

184 copy_size = state_inner_size * state_elem_size

185

154 - # Mirror postprocess_mamba's trailing

155 - # if src_block_idx == dest_block_idx: num_accepted_tokens_cpu[i] = 1

156 - # This runs whether or not the copy below is skipped (it's per-request, so

157 - # only state_idx == 0 writes).

158 - if src_block_idx == dest_block_idx and state_idx == 0:

159 - tl.store(num_accepted_tokens_out_ptr + req_idx, 1)

160 -

161 - # Mirror collect_mamba_copy_meta's early return: src==dst with no token

162 - # bias means source and destination ranges coincide, so the copy is a

163 - # no-op.

164 - if src_block_idx == dest_block_idx and accept_token_bias == 0:

165 - return

166 -

186 offsets = tl.arange(0, COPY_BLOCK_SIZE)

187 for i in range(0, copy_size, COPY_BLOCK_SIZE):

188 mask = (i + offsets) < copy_size

@@ -271,6 +290,9 @@ class MambaSpecDecodeGPUContext:

290 state_inner_sizes: torch.Tensor # int64: elements in inner dimensions

291 state_conv_widths: torch.Tensor # int32: conv width (0 for temporal states)

292 state_group_indices: torch.Tensor # int32: maps state_idx to group index

293 + # DS conv row metadata. Zero keeps the single-region copy path.

294 + state_dim_row_count: torch.Tensor # int32: per-block dim row count

295 + state_dim_row_stride: torch.Tensor # int64: bytes between rows

296

297 # Configuration

298 block_size: int

@@ -338,6 +360,12 @@ class MambaSpecDecodeGPUContext:

360 state_group_indices=torch.zeros(

361 total_states, dtype=torch.int32, device=device

362 ),

363 + state_dim_row_count=torch.zeros(

364 + total_states, dtype=torch.int32, device=device

365 + ),

366 + state_dim_row_stride=torch.zeros(

367 + total_states, dtype=torch.int64, device=device

368 + ),

369 block_size=mamba_spec.block_size,

370 num_layers=num_layers,

371 num_state_types=num_state_types,

@@ -430,17 +458,23 @@ class MambaSpecDecodeGPUContext:

458 or copy_func is get_temporal_copy_spec

459 ), f"unexpected copy func: {copy_func}"

460 if copy_func is get_conv_copy_spec:

433 - # Conv state: conv_width is state.size(1)

434 - # inner_size is stride(1) = elements per conv position,

435 - # used to compute byte offset for state[block, offset:]

436 - conv_w = state.size(1) if state.dim() > 1 else 0

437 - self.state_conv_widths[idx] = conv_w

438 - if state.dim() > 2:

439 - # stride(1) = product of dims[2:] for contiguous tensor

440 - self.state_inner_sizes[idx] = state.stride(1)

441 - else:

442 - # 2D tensor: [num_blocks, conv_dim], no inner dims

461 + if state.dim() != 3:

462 + raise ValueError(

463 + "Expected 3D conv state cache, got "

464 + f"shape {tuple(state.shape)}"

465 + )

466 + if is_conv_state_dim_first():

467 + # DS layout: state_len is the slide axis.

468 + self.state_conv_widths[idx] = state.size(2)

469 self.state_inner_sizes[idx] = 1

470 + self.state_dim_row_count[idx] = state.size(1)

471 + self.state_dim_row_stride[idx] = (

472 + state.stride(1) * state.element_size()

473 + )

474 + else:

475 + # SD layout: dim is contiguous.

476 + self.state_conv_widths[idx] = state.size(1)

477 + self.state_inner_sizes[idx] = state.stride(1)

478 else:

479 # Temporal state: inner_size = natural elements per

480 # block (prod of inner dims). The kernel uses this

@@ -521,10 +555,13 @@ class MambaSpecDecodeGPUContext:

555 self.state_inner_sizes,

556 self.state_conv_widths,

557 self.state_group_indices,

558 + self.state_dim_row_count,

559 + self.state_dim_row_stride,

560 self.num_accepted_tokens_out,

561 num_reqs,

562 block_size=self.block_size,

563 COPY_BLOCK_SIZE=1024,

564 + CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),

565 )

566

567