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

Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 22:35 UTC

[Model] Mamba2 preallocate SSM output tensor to avoid d2d copy overhead (#21075)

Signed-off-by: Chih-Chieh Yang <7364402+cyang49@users.noreply.github.com> Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>

Chih-Chieh Yang committed Aug 2, 2025 at 04:59 UTC b690e34824fd5a5c4054a0c0468ebfb6aa1dd215

9 files changed +144 -118

tests/kernels/mamba/test_mamba_ssm.py

+40 -34

@@ -365,6 +365,7 @@ def test_selective_state_update(dim, dstate, has_z, itype):

365 batch_size = 1

366 state = torch.randn(batch_size, dim, dstate, dtype=itype, device=device)

367 x = torch.randn(batch_size, dim, device=device, dtype=itype)

368 + out = torch.empty_like(x)

369 dt = torch.randn(batch_size, dim, device=device, dtype=itype)

370 dt_bias = torch.rand(dim, device=device) - 4.0

371 A = -torch.rand(dim, dstate, device=device) - 1.0

@@ -373,16 +374,17 @@ def test_selective_state_update(dim, dstate, has_z, itype):

374 D = torch.randn(dim, device=device)

375 z = torch.randn_like(x) if has_z else None

376 state_ref = state.detach().clone()

376 - out = selective_state_update(state,

377 - x,

378 - dt,

379 - A,

380 - B,

381 - C,

382 - D=D,

383 - z=z,

384 - dt_bias=dt_bias,

385 - dt_softplus=True)

377 + selective_state_update(state,

378 + x,

379 + dt,

380 + A,

381 + B,

382 + C,

383 + D=D,

384 + z=z,

385 + dt_bias=dt_bias,

386 + dt_softplus=True,

387 + out=out)

388 out_ref = selective_state_update_ref(state_ref,

389 x,

390 dt,

@@ -581,6 +583,7 @@ def test_selective_state_update_with_batch_indices(with_padding, dim, dstate,

583 ],

584 dim=0)

585 x = torch.randn(padded_batch_size, dim, device=device, dtype=itype)

586 + out = torch.empty_like(x)

587 dt = torch.randn(padded_batch_size, dim, device=device, dtype=itype)

588 dt_bias = torch.rand(dim, device=device) - 4.0

589 A = -torch.rand(dim, dstate, device=device) - 1.0

@@ -590,18 +593,19 @@ def test_selective_state_update_with_batch_indices(with_padding, dim, dstate,

593 z = torch.randn_like(x) if has_z else None

594 state_ref = state[state_indices, :].clone()

595 state_before = state.clone()

593 - out = selective_state_update(state,

594 - x,

595 - dt,

596 - A,

597 - B,

598 - C,

599 - D=D,

600 - z=z,

601 - dt_bias=dt_bias,

602 - dt_softplus=True,

603 - state_batch_indices=padded_state_indices,

604 - pad_slot_id=PAD_SLOT_ID)

596 + selective_state_update(state,

597 + x,

598 + dt,

599 + A,

600 + B,

601 + C,

602 + D=D,

603 + z=z,

604 + dt_bias=dt_bias,

605 + dt_softplus=True,

606 + state_batch_indices=padded_state_indices,

607 + pad_slot_id=PAD_SLOT_ID,

608 + out=out)

609 out_ref = selective_state_update_ref(state_ref,

610 x[:batch_size],

611 dt[:batch_size],

@@ -665,6 +669,7 @@ def test_selective_state_update_with_heads_with_batch_indices(

669 dtype=torch.int32, device=device)

670

671 x = torch.randn(batch_size, nheads, headdim, device=device, dtype=itype)

672 + out = torch.empty_like(x)

673 if not tie_hdim:

674 dt = torch.randn(batch_size,

675 nheads,

@@ -691,18 +696,19 @@ def test_selective_state_update_with_heads_with_batch_indices(

696 C = torch.randn(batch_size, ngroups, dstate, device=device)

697 z = torch.randn_like(x) if has_z else None

698 state_ref = state[state_indices, :].detach().clone()

694 - out = selective_state_update(state,

695 - x,

696 - dt,

697 - A,

698 - B,

699 - C,

700 - D=D,

701 - z=z,

702 - dt_bias=dt_bias,

703 - dt_softplus=True,

704 - state_batch_indices=state_indices,

705 - pad_slot_id=PAD_SLOT_ID)

699 + selective_state_update(state,

700 + x,

701 + dt,

702 + A,

703 + B,

704 + C,

705 + D=D,

706 + z=z,

707 + dt_bias=dt_bias,

708 + dt_softplus=True,

709 + state_batch_indices=state_indices,

710 + pad_slot_id=PAD_SLOT_ID,

711 + out=out)

712 out_ref = selective_state_update_ref(state_ref,

713 x,

714 dt,

tests/kernels/mamba/test_mamba_ssm_ssd.py

+13 -10

@@ -212,15 +212,16 @@ def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size,

212

213 Y_min, final_state_min = ssd_minimal_discrete(X * dt.unsqueeze(-1), A * dt,

214 B, C, chunk_size)

215 -

216 - Y, final_state = mamba_chunk_scan_combined(X,

217 - dt,

218 - A,

219 - B,

220 - C,

221 - chunk_size,

222 - D=None,

223 - return_final_states=True)

215 + Y = torch.empty_like(X)

216 + final_state = mamba_chunk_scan_combined(X,

217 + dt,

218 + A,

219 + B,

220 + C,

221 + chunk_size,

222 + D=None,

223 + return_final_states=True,

224 + out=Y)

225

226 # just test the last in sequence

227 torch.testing.assert_close(Y[:, -1], Y_min[:, -1], atol=atol, rtol=rtol)

@@ -292,7 +293,8 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,

293 _query_start_loc_to_chunk_indices_offsets(

294 cu_seqlens, chunk_size, cu_seqlens[-1])

295

295 - Y, new_states = mamba_chunk_scan_combined(

296 + Y = torch.empty_like(X)

297 + new_states = mamba_chunk_scan_combined(

298 X,

299 dt,

300 A,

@@ -306,6 +308,7 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,

308 chunk_offsets=chunk_offsets,

309 return_varlen_states=True,

310 initial_states=states,

311 + out=Y,

312 )

313

314 # just test the last in sequence

vllm/model_executor/layers/mamba/mamba_mixer.py

+4 -2

@@ -220,7 +220,8 @@ class MambaMixer(CustomOp):

220 has_initial_state=attn_metadata.context_lens_tensor > 0,

221 query_start_loc=attn_metadata.query_start_loc)

222 else:

223 - scan_outputs = selective_state_update(

223 + scan_outputs = torch.empty_like(hidden_states.transpose(0, 1))

224 + selective_state_update(

225 mamba_cache_params.ssm_state,

226 hidden_states.transpose(0, 1),

227 discrete_time_step.transpose(0, 1),

@@ -231,7 +232,8 @@ class MambaMixer(CustomOp):

232 gate.transpose(0, 1),

233 time_proj_bias,

234 dt_softplus=True,

234 - state_batch_indices=mamba_cache_params.state_indices_tensor)

235 + state_batch_indices=mamba_cache_params.state_indices_tensor,

236 + out=scan_outputs)

237 scan_outputs = scan_outputs.transpose(0, 1)

238

239 # 4. Final linear projection

vllm/model_executor/layers/mamba/mamba_mixer2.py

+32 -22

@@ -541,7 +541,6 @@ class MambaMixer2(MambaBase, CustomOp):

541 # NOTE: V0 put prefill before decode, v1 puts decode before prefill

542 # Separate prefill and decode by splitting varlen input

543 # Split along token dimension

544 - # NOTE: V0 put prefill before decode, v1 puts decode before prefill

544 if envs.VLLM_USE_V1:

545 hidden_states_B_C_d, hidden_states_B_C_p = torch.split(

546 hidden_states_B_C[:num_actual_tokens],

@@ -583,7 +582,28 @@ class MambaMixer2(MambaBase, CustomOp):

582 1]

583 if has_prefill else None)

584

586 - ssd_output_list = []

585 + # Preallocate output tensor to avoid memcpy cost for merging prefill

586 + # and decode outputs

587 + preallocated_ssm_out = torch.empty(

588 + [

589 + num_prefill_tokens + num_decodes,

590 + (self.num_heads // self.tp_size) * self.head_dim

591 + ],

592 + dtype=hidden_states.dtype,

593 + device=hidden_states.device,

594 + )

595 + if envs.VLLM_USE_V1:

596 + preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(

597 + preallocated_ssm_out,

598 + [num_decodes, num_prefill_tokens],

599 + dim=0,

600 + )

601 + else:

602 + preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split(

603 + preallocated_ssm_out,

604 + [num_prefill_tokens, num_decodes],

605 + dim=0,

606 + )

607

608 # Process prefill requests

609 if has_prefill:

@@ -623,7 +643,8 @@ class MambaMixer2(MambaBase, CustomOp):

643 has_initial_states_p[:num_prefills, None, None, None],

644 ssm_state[state_indices_tensor_p], 0)

645

626 - scan_output, varlen_state = mamba_chunk_scan_combined(

646 + # NOTE: final output is an in-place update of out tensor

647 + varlen_state = mamba_chunk_scan_combined(

648 hidden_states_p.view(1, num_prefill_tokens,

649 self.num_heads // self.tp_size,

650 self.head_dim),

@@ -646,15 +667,14 @@ class MambaMixer2(MambaBase, CustomOp):

667 return_final_states=False,

668 dt_softplus=True,

669 dt_limit=(0.0, float("inf")),

670 + out=preallocated_ssm_out_p.view(1, num_prefill_tokens, -1,

671 + self.head_dim),

672 )

673

674 # update ssm states

675 # - varlen state is a (num_prefills, nheads, headdim, dstate) tensor

676 ssm_state[state_indices_tensor_p] = varlen_state

677

655 - # - reshape

656 - ssd_output_list.append(scan_output.view(num_prefill_tokens, -1))

657 -

678 # Process decode requests

679 if has_decode:

680 # 2. Convolution sequence transformation

@@ -684,8 +704,8 @@ class MambaMixer2(MambaBase, CustomOp):

704 # - the hidden is reshaped into (bs, num_heads, head_dim)

705 # - mamba_cache_params.ssm_state's slots will be selected

706 # using state_indices_tensor_d

687 -

688 - hidden_states_d = selective_state_update(

707 + # NOTE: final output is an in-place update of out tensor

708 + selective_state_update(

709 ssm_state,

710 hidden_states_d,

711 dt_d,

@@ -697,26 +717,16 @@ class MambaMixer2(MambaBase, CustomOp):

717 dt_bias=dt_bias,

718 dt_softplus=True,

719 state_batch_indices=state_indices_tensor_d,

720 + out=preallocated_ssm_out_d.view(num_decodes, -1,

721 + self.head_dim),

722 )

723

702 - if envs.VLLM_USE_V1:

703 - ssd_output_list.insert(

704 - 0,

705 - hidden_states_d.view(-1, (self.num_heads // self.tp_size) *

706 - self.head_dim))

707 - else:

708 - ssd_output_list.append(

709 - hidden_states_d.view(-1, (self.num_heads // self.tp_size) *

710 - self.head_dim))

711 -

712 - # Merge prefill and decode outputs before passing to gated MLP

713 - hidden_states = torch.vstack(ssd_output_list)

714 -

724 # 4. gated MLP

725 # GatedRMSNorm internally applying SiLU to the gate

726 # SiLU is applied internally before normalization, unlike standard

727 # norm usage

719 - hidden_states = self.norm(hidden_states, gate[:num_actual_tokens])

728 + hidden_states = self.norm(preallocated_ssm_out,

729 + gate[:num_actual_tokens])

730

731 # 5. Final linear projection

732 output[:num_actual_tokens], _ = self.out_proj(hidden_states)

vllm/model_executor/layers/mamba/ops/mamba_ssm.py

+8 -8

@@ -205,7 +205,8 @@ def selective_state_update(state,

205 dt_bias=None,

206 dt_softplus=False,

207 state_batch_indices=None,

208 - pad_slot_id=PAD_SLOT_ID):

208 + pad_slot_id=PAD_SLOT_ID,

209 + out=None):

210 """

211 Argument:

212 state: (batch, dim, dstate) or (batch, nheads, dim, dstate)

@@ -223,10 +224,9 @@ def selective_state_update(state,

224 for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]

225 in this case, the kernel will not process entries at

226 indices 0 and 3

226 - Return:

227 - out: (batch, dim) or (batch, nheads, dim)

227 + out: Preallocated ssm output tensor. Assume same shape as x.

228 + In-place updated.

229 """

229 - has_heads = state.dim() > 3

230 if state.dim() == 3:

231 state = state.unsqueeze(1)

232 if x.dim() == 2:

@@ -245,6 +245,8 @@ def selective_state_update(state,

245 z = z.unsqueeze(1)

246 if dt_bias is not None and dt_bias.dim() == 1:

247 dt_bias = dt_bias.unsqueeze(0)

248 + if out.dim() == 2:

249 + out = out.unsqueeze(1)

250

251 _, nheads, dim, dstate = state.shape

252 batch = x.shape[0]

@@ -264,7 +266,8 @@ def selective_state_update(state,

266 assert dt_bias.shape == (nheads, dim)

267 if state_batch_indices is not None:

268 assert state_batch_indices.shape == (batch, )

267 - out = torch.empty_like(x)

269 + assert out.shape == x.shape

270 +

271 grid = lambda META: (triton.cdiv(dim, META['BLOCK_SIZE_M']), batch, nheads)

272 z_strides = ((z.stride(0), z.stride(1), z.stride(2)) if z is not None else

273 (0, 0, 0))

@@ -328,9 +331,6 @@ def selective_state_update(state,

331 BLOCK_SIZE_M,

332 num_warps=num_warps,

333 )

331 - if not has_heads:

332 - out = out.squeeze(1)

333 - return out

334

335

336 def selective_scan_fn(u,

vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py

+5 -14

@@ -454,6 +454,7 @@ def _chunk_scan_fwd(

454 chunk_indices=None,

455 chunk_offsets=None,

456 initial_states=None,

457 + out=None,

458 ):

459 batch, seqlen, nheads, headdim = x.shape

460 _, _, nchunks, chunk_size = dt.shape

@@ -483,20 +484,10 @@ def _chunk_scan_fwd(

484 else:

485 chunk_indices, chunk_offsets = None, None

486

486 - # Allocates output.

487 - out = torch.empty(batch,

488 - seqlen,

489 - nheads,

490 - headdim,

491 - device=x.device,

492 - dtype=x.dtype)

487 + assert out.shape == x.shape

488 +

489 if z is not None:

494 - out_x = torch.empty(batch,

495 - seqlen,

496 - nheads,

497 - headdim,

498 - device=x.device,

499 - dtype=x.dtype)

490 + out_x = torch.empty_like(x)

491 assert out_x.stride() == out.stride()

492 else:

493 out_x = None

@@ -579,4 +570,4 @@ def _chunk_scan_fwd(

570 IS_TRITON_22=TRITON_22,

571 HAS_INITSTATES=initial_states is not None,

572 )

582 - return out, out_x

573 + return out_x

vllm/model_executor/layers/mamba/ops/ssd_combined.py

+16 -12

@@ -36,7 +36,8 @@ def _mamba_chunk_scan_combined_fwd(x,

36 chunk_offsets=None,

37 cu_seqlens=None,

38 dt_softplus=False,

39 - dt_limit=(0.0, float("inf"))):

39 + dt_limit=(0.0, float("inf")),

40 + out=None):

41 batch, seqlen, nheads, headdim = x.shape

42 _, _, ngroups, dstate = B.shape

43 assert nheads % ngroups == 0

@@ -134,7 +135,7 @@ def _mamba_chunk_scan_combined_fwd(x,

135 # - in each (pseudo) chunk, we detect if the previous (pseudo) chunk had

136 # a seq_idx change, in which case we take states information from

137 # init_states.

137 - out, out_x = _chunk_scan_fwd(

138 + out_x = _chunk_scan_fwd(

139 CB,

140 x,

141 dt,

@@ -147,9 +148,10 @@ def _mamba_chunk_scan_combined_fwd(x,

148 chunk_indices=chunk_indices,

149 chunk_offsets=chunk_offsets,

150 initial_states=initial_states,

151 + out=out,

152 )

153 if cu_seqlens is None:

152 - return out, out_x, dt, dA_cumsum, states, final_states

154 + return out_x, dt, dA_cumsum, states, final_states

155 else:

156 assert batch == 1, "passing cu_seqlens to get the varlen states is only supported if batch dimension is 1"

157 varlen_states = chunk_state_varlen(

@@ -161,7 +163,7 @@ def _mamba_chunk_scan_combined_fwd(x,

163 states.squeeze(0),

164 initial_states=initial_states,

165 )

164 - return out, out_x, dt, dA_cumsum, states, final_states, varlen_states

166 + return out_x, dt, dA_cumsum, states, final_states, varlen_states

167

168

169 def mamba_chunk_scan_combined(x,

@@ -180,6 +182,7 @@ def mamba_chunk_scan_combined(x,

182 cu_seqlens=None,

183 dt_softplus=False,

184 dt_limit=(0.0, float("inf")),

185 + out=None,

186 return_final_states=False,

187 return_varlen_states=False):

188 """

@@ -197,15 +200,14 @@ def mamba_chunk_scan_combined(x,

200 seq_idx: (batch, seqlen)

201 cu_seqlens: (num_sequences + 1) or None, only used if return_varlen_states is True

202 dt_softplus: Whether to apply softplus to dt

200 - Return:

201 - out: (batch, seqlen, nheads, headdim)

203 + out: Preallocated output tensor

204 """

205

206 if not return_varlen_states:

207 cu_seqlens = None

208 else:

209 assert cu_seqlens is not None, "cu_seqlens must be provided if return_varlen_states is True"

208 - out, out_x, dt_out, dA_cumsum, states, final_states, *rest = _mamba_chunk_scan_combined_fwd(

210 + out_x, dt_out, dA_cumsum, states, final_states, *rest = _mamba_chunk_scan_combined_fwd(

211 x,

212 dt,

213 A,

@@ -221,12 +223,14 @@ def mamba_chunk_scan_combined(x,

223 chunk_offsets=chunk_offsets,

224 cu_seqlens=cu_seqlens,

225 dt_softplus=dt_softplus,

224 - dt_limit=dt_limit)

226 + dt_limit=dt_limit,

227 + out=out)

228 if not return_varlen_states:

226 - return out if not return_final_states else (out, final_states)

229 + if not return_final_states:

230 + return

231 + else:

232 + return final_states

233 else:

234 varlen_states = rest[0]

229 - return (out,

230 - varlen_states) if not return_final_states else (out,

231 - final_states,

235 + return (varlen_states) if not return_final_states else (final_states,

236 varlen_states)

vllm/model_executor/models/phi4flash.py

+4 -2

@@ -387,7 +387,8 @@ class Phi4Mamba(nn.Module):

387 has_initial_state=attn_metadata.context_lens_tensor > 0,

388 query_start_loc=attn_metadata.query_start_loc)

389 else:

390 - scan_outputs = selective_state_update(

390 + scan_outputs = torch.empty_like(hidden_states.transpose(0, 1))

391 + selective_state_update(

392 mamba_cache_params.ssm_state,

393 hidden_states.transpose(0, 1),

394 discrete_time_step.transpose(0, 1),

@@ -400,7 +401,8 @@ class Phi4Mamba(nn.Module):

401 None if self.yoco_kv else gate.transpose(0, 1),

402 time_proj_bias,

403 dt_softplus=True,

403 - state_batch_indices=mamba_cache_params.state_indices_tensor)

404 + state_batch_indices=mamba_cache_params.state_indices_tensor,

405 + out=scan_outputs)

406 scan_outputs = scan_outputs.transpose(0, 1)

407

408 # 4. Final linear projection

vllm/model_executor/models/plamo2.py

+22 -14

@@ -257,7 +257,21 @@ class Plamo2MambaMixer(nn.Module):

257 query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills + 1]

258 if has_prefill else None)

259

260 - ssd_output_list = []

260 + # Preallocate output tensor to avoid memcpy cost for merging prefill

261 + # and decode outputs

262 + preallocated_ssm_out = torch.empty(

263 + [

264 + num_prefill_tokens + num_decodes,

265 + (self.num_heads // self.tp_size) * self.head_dim

266 + ],

267 + dtype=hidden_states.dtype,

268 + device=hidden_states.device,

269 + )

270 + preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split(

271 + preallocated_ssm_out,

272 + [num_prefill_tokens, num_decodes],

273 + dim=0,

274 + )

275

276 # Process prefill requests

277 if has_prefill:

@@ -290,7 +304,7 @@ class Plamo2MambaMixer(nn.Module):

304 initial_states = torch.where(

305 mamba2_metadata.has_initial_states[:, None, None, None],

306 mamba_cache_params.ssm_state[state_indices_tensor_p], 0)

293 - scan_output, varlen_state = mamba_chunk_scan_combined(

307 + varlen_state = mamba_chunk_scan_combined(

308 hidden_states_p.view(1, num_prefill_tokens,

309 self.num_heads // self.tp_size,

310 self.head_dim),

@@ -312,15 +326,14 @@ class Plamo2MambaMixer(nn.Module):

326 return_final_states=False,

327 dt_softplus=True,

328 dt_limit=(0.0, float("inf")),

329 + out=preallocated_ssm_out_p.view(1, num_prefill_tokens, -1,

330 + self.head_dim),

331 )

332

333 # update ssm states

334 # - varlen state is a (batch, nheads, headdim, dstate) tensor

335 mamba_cache_params.ssm_state[state_indices_tensor_p] = varlen_state

336

321 - # - reshape

322 - ssd_output_list.append(scan_output.view(num_prefill_tokens, -1))

323 -

337 # Process decode requests

338 if has_decode:

339 # 2. Convolution sequence transformation

@@ -349,8 +362,7 @@ class Plamo2MambaMixer(nn.Module):

362 # - the hidden is reshaped into (bs, num_heads, head_dim)

363 # - mamba_cache_params.ssm_state's slots will be selected

364 # using state_indices_tensor_d

352 -

353 - hidden_states_d = selective_state_update(

365 + selective_state_update(

366 mamba_cache_params.ssm_state,

367 hidden_states_d,

368 dt,

@@ -362,17 +374,13 @@ class Plamo2MambaMixer(nn.Module):

374 dt_bias=dt_bias,

375 dt_softplus=True,

376 state_batch_indices=state_indices_tensor_d,

377 + out=preallocated_ssm_out_d.view(num_decodes, -1,

378 + self.head_dim),

379 )

380 assert self.num_heads % self.tp_size == 0

367 - ssd_output_list.append(

368 - hidden_states_d.view(-1, (self.num_heads // self.tp_size) *

369 - self.head_dim))

370 -

371 - # Merge prefill and decode outputs before passing to MLP

372 - hidden_states = torch.vstack(ssd_output_list)

381

382 # 4. Final linear projection

375 - out = self.out_proj(hidden_states)

383 + out = self.out_proj(preallocated_ssm_out)

384 return out

385

386