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

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, 05:30 UTC

[V1] - Split Prefill and Decode for Mamba1 models (#22653)

Signed-off-by: amirk <amirk@ai21.com> Signed-off-by: asafg <asafg@ai21.com> Co-authored-by: asafg <asafg@ai21.com> Co-authored-by: Asaf Joseph Gardin <39553475+Josephasafg@users.noreply.github.com>

amirai21 committed Aug 15, 2025 at 11:59 UTC fe91ce9591760e3b58fc24845fa80364fbcdd07f

3 files changed +253 -95

tests/models/language/generation/test_hybrid.py

+13

@@ -57,6 +57,13 @@ V1_SUPPORTED_MODELS = [

57 # Avoid OOM

58 MAX_NUM_SEQS = 4

60 +# Once we add support for FCG in Mamba1, this list will be removed and tests

61 +# all test cases will use enforce_eager=False

62 +ENFORCE_EAGER_MODELS_V1 = [

63 + "state-spaces/mamba-130m-hf",

64 + "ai21labs/Jamba-tiny-dev",

65 +]

66 +

68 @pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)

69 @pytest.mark.parametrize("max_tokens", [64])

@@ -94,13 +101,19 @@ def test_models(

101 example_prompts, max_tokens, num_logprobs)

102

103 if model in V1_SUPPORTED_MODELS:

104 + enforce_eager = False

105 with monkeypatch.context() as m:

106 m.setenv("VLLM_USE_V1", "1")

107 if model in HYBRID_MODELS:

108 # required due to reorder_batch behaviour

109 m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")

110 +

111 + if model in ENFORCE_EAGER_MODELS_V1:

112 + enforce_eager = True

113 +

114 with vllm_runner(model,

115 max_num_seqs=MAX_NUM_SEQS,

116 + enforce_eager=enforce_eager,

117 enable_prefix_caching=False) as vllm_model:

118 vllm_v1_outputs = vllm_model.generate_greedy_logprobs(

119 example_prompts, max_tokens, num_logprobs)

vllm/model_executor/layers/mamba/mamba_mixer.py

+219 -90

@@ -1,13 +1,14 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from typing import Optional

4 +from typing import NamedTuple, Optional

6 import torch

7 from torch import nn

8 from torch.nn.parameter import Parameter

10 from vllm import envs

11 +from vllm.attention.backends.abstract import AttentionMetadata

12 from vllm.config import get_current_vllm_config

13 from vllm.distributed.parallel_state import (

14 get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)

@@ -154,13 +155,38 @@ class MambaMixer(MambaBase, CustomOp):

155

156 self.prefix = prefix

157

158 + def _ssm_transform(

159 + self, x: torch.Tensor

160 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

161 + if self.is_lora_enabled:

162 + # Lora kernel requires contiguous tensor.

163 + ssm_params = self.x_proj(x.contiguous())[0]

164 + else:

165 + ssm_params = self.x_proj(x)[0]

166 + time_step, B, C = torch.split(

167 + ssm_params,

168 + [self.time_step_rank, self.ssm_state_size, self.ssm_state_size],

169 + dim=-1)

170 + if self.use_rms_norm:

171 + assert self.dt_layernorm is not None

172 + assert self.b_layernorm is not None

173 + assert self.c_layernorm is not None

174 + time_step = self.dt_layernorm(time_step.contiguous())

175 + B = self.b_layernorm(B.contiguous())

176 + C = self.c_layernorm(C.contiguous())

177 + discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)

178 + return discrete_time_step, B, C

179 +

180 def forward(self,

181 hidden_states: torch.Tensor,

182 mamba_cache_params: Optional[MambaCacheParams] = None):

183 if not envs.VLLM_USE_V1:

184 return CustomOp.forward(self, hidden_states, mamba_cache_params)

185 else:

163 - return self.forward_cuda(hidden_states, mamba_cache_params)

186 + return self.forward_cuda(

187 + hidden_states,

188 + mamba_cache_params,

189 + )

190

191 def forward_native(self,

192 hidden_states: torch.Tensor,

@@ -170,6 +196,27 @@ class MambaMixer(MambaBase, CustomOp):

196 def forward_cuda(self,

197 hidden_states: torch.Tensor,

198 mamba_cache_params: Optional[MambaCacheParams] = None):

199 + """

200 + Run the Mamba-1 SSM pipeline.

201 +

202 + Steps

203 + -----

204 + 1. Apply the gated-MLP linear projection to the raw input.

205 + 2. Pass the projected sequence through the convolutional mixing layer.

206 + 3. Feed the result into the State-Space Model (SSM) blocks.

207 + 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x)

208 + to produce contextual representations.

209 + 5. Project the contextualised sequence back

210 + to the output embedding dimension.

211 +

212 + Batch handling

213 + --------------

214 + Prefill and decode tokens are processed by dedicated CUDA

215 + kernels for both the convolutional (conv1d) and SSM stages.

216 + In the case of a mixed batch (containing both prefill and

217 + decode tokens), both sets of kernels are executed independently

218 + and their outputs are concatenated before the final output projection.

219 + """

220

221 forward_context: ForwardContext = get_forward_context()

222 attn_metadata = forward_context.attn_metadata

@@ -185,126 +232,142 @@ class MambaMixer(MambaBase, CustomOp):

232 self_kv_cache = self.kv_cache[forward_context.virtual_engine]

233 conv_state = self_kv_cache[0].transpose(-1, -2)

234 ssm_state = self_kv_cache[1]

188 - has_initial_state = mamba1_metadata.has_initial_states

189 - context_lens_tensor = mamba1_metadata.context_lens_tensor

235 + has_initial_states = mamba1_metadata.has_initial_states

236 else:

237 + assert isinstance(attn_metadata, AttentionMetadata)

238 assert mamba_cache_params is not None

239 conv_state = mamba_cache_params.conv_state

240 ssm_state = mamba_cache_params.ssm_state

241 state_indices_tensor = mamba_cache_params.state_indices_tensor

242 query_start_loc = attn_metadata.query_start_loc

243 context_lens_tensor = attn_metadata.context_lens_tensor

197 -

244 + has_initial_states = None

245 if context_lens_tensor is not None:

199 - has_initial_state = context_lens_tensor > 0

246 + has_initial_states = context_lens_tensor > 0

247

248 # 1. Gated MLP's linear projection

249 projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)

203 - hidden_states, gate = projected_states.chunk(2, dim=-2)

250 + hidden_states_BC, gate = projected_states.chunk(2, dim=-2)

251

205 - # 2. Convolution sequence transformation

252 conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),

253 self.conv1d.weight.size(2))

254

255 if envs.VLLM_USE_V1 and attn_metadata is None:

256 # V1 profile run

211 - hidden_states = hidden_states.contiguous()

212 - return self.out_proj(hidden_states.transpose(-2, -1))[0]

213 -

214 - if query_start_loc is not None and context_lens_tensor is not None:

215 - # |---------- N-1 iteration --------|

216 - # |---------------- N iteration ---------------------|

217 - # |- tokenA -|......................|-- newTokens ---|

218 - # |---------- context_len ----------|

219 - # |-------------------- seq_len ---------------------|

220 - # |-- query_len ---|

221 - hidden_states = causal_conv1d_fn(

222 - hidden_states,

257 + hidden_states_BC = hidden_states_BC.contiguous()

258 + return self.out_proj(hidden_states_BC.transpose(-2, -1))[0]

259 +

260 + num_prefill_tokens = attn_metadata.num_prefill_tokens # token count

261 + num_decode_tokens = attn_metadata.num_decode_tokens

262 + num_prefills = attn_metadata.num_prefills # request count

263 + num_decodes = attn_metadata.num_decode_tokens # token count (=request)

264 + has_prefill = num_prefill_tokens > 0

265 + has_decode = num_decode_tokens > 0

266 +

267 + prefill_decode_split = split_batch_to_prefill_and_decode(

268 + hidden_states_BC,

269 + gate,

270 + state_indices_tensor,

271 + query_start_loc,

272 + has_initial_states,

273 + num_prefill_tokens,

274 + num_decode_tokens,

275 + num_prefills,

276 + num_decodes,

277 + )

278 + hidden_states_BC_p = prefill_decode_split.hidden_states_BC_p

279 + hidden_states_BC_d = prefill_decode_split.hidden_states_BC_d

280 + gate_p = prefill_decode_split.gate_p

281 + gate_d = prefill_decode_split.gate_d

282 + state_indices_tensor_p = prefill_decode_split.state_indices_tensor_p

283 + state_indices_tensor_d = prefill_decode_split.state_indices_tensor_d

284 + query_start_loc_p = prefill_decode_split.query_start_loc_p

285 + has_initial_states_p = prefill_decode_split.has_initial_states_p

286 +

287 + ssm_outputs = []

288 +

289 + if has_prefill:

290 + # 2. Convolution sequence transformation

291 + conv_out_p = causal_conv1d_fn(

292 + hidden_states_BC_p,

293 conv_weights,

224 - bias=self.conv1d.bias,

294 + self.conv1d.bias,

295 activation=self.activation,

296 conv_states=conv_state,

227 - has_initial_state=has_initial_state,

228 - cache_indices=state_indices_tensor,

229 - query_start_loc=query_start_loc)

230 - else:

231 - hidden_states = causal_conv1d_update(

232 - hidden_states.transpose(0, 1),

297 + has_initial_state=has_initial_states_p,

298 + cache_indices=state_indices_tensor_p,

299 + query_start_loc=query_start_loc_p)

300 + # 3. State Space Model sequence transformations.

301 + discrete_time_step_p, B_p, C_p = self._ssm_transform(

302 + conv_out_p.transpose(-2, -1))

303 + time_proj_bias = self._time_proj_bias()

304 +

305 + # 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x)

306 + scan_out_p = selective_scan_fn(

307 + conv_out_p,

308 + ssm_state,

309 + discrete_time_step_p,

310 + self.A,

311 + B_p.transpose(-2, -1),

312 + C_p.transpose(-2, -1),

313 + self.D.float(),

314 + gate_p,

315 + time_proj_bias,

316 + delta_softplus=True,

317 + cache_indices=state_indices_tensor_p,

318 + has_initial_state=has_initial_states_p,

319 + query_start_loc=query_start_loc_p)

320 + ssm_outputs.append(scan_out_p)

321 +

322 + if has_decode:

323 + # 2. Convolution sequence transformation

324 + conv_out_d = causal_conv1d_update(

325 + hidden_states_BC_d.transpose(0, 1),

326 conv_state,

327 conv_weights,

328 self.conv1d.bias,

329 self.activation,

237 - conv_state_indices=state_indices_tensor)

238 - hidden_states = hidden_states.transpose(0, 1)

330 + conv_state_indices=state_indices_tensor_d).transpose(0, 1)

331

240 - # 3. State Space Model sequence transformation

241 - # 3.a. input varying initialization of time_step, B and C

332 + # 3. State Space Model sequence transformation.

333 + discrete_time_step_d, B_d, C_d = self._ssm_transform(

334 + conv_out_d.transpose(-2, -1))

335 + time_proj_bias = self._time_proj_bias()

336

243 - if self.is_lora_enabled:

244 - # lora kernel requires contiguous tensor

245 - ssm_parameters = self.x_proj(

246 - hidden_states.transpose(-2, -1).contiguous())[0]

247 - else:

248 - ssm_parameters = self.x_proj(hidden_states.transpose(-2, -1))[0]

249 -

250 - time_step, B, C = torch.split(

251 - ssm_parameters,

252 - [self.time_step_rank, self.ssm_state_size, self.ssm_state_size],

253 - dim=-1,

254 - )

255 - if self.use_rms_norm:

256 - assert self.dt_layernorm is not None

257 - assert self.b_layernorm is not None

258 - assert self.c_layernorm is not None

259 - time_step = self.dt_layernorm(time_step.contiguous())

260 - B = self.b_layernorm(B.contiguous())

261 - C = self.c_layernorm(C.contiguous())

262 -

263 - discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)

264 - # 3.c perform the recurrence y ← SSM(A, B, C)(x)

265 - time_proj_bias = (self.dt_proj.bias.float() if hasattr(

266 - self.dt_proj, "bias") else None)

267 -

268 - if query_start_loc is not None and context_lens_tensor is not None:

269 - scan_outputs = selective_scan_fn(

270 - hidden_states,

271 - ssm_state,

272 - discrete_time_step,

273 - self.A,

274 - B.transpose(-2, -1),

275 - C.transpose(-2, -1),

276 - self.D.float(),

277 - gate,

278 - time_proj_bias,

279 - delta_softplus=True,

280 - cache_indices=state_indices_tensor,

281 - has_initial_state=has_initial_state,

282 - query_start_loc=query_start_loc)

283 - else:

284 - scan_outputs = torch.empty_like(hidden_states.transpose(0, 1))

337 + # 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x)

338 + scan_outputs_d = torch.empty_like(

339 + hidden_states_BC_d.transpose(0, 1))

340 selective_state_update(ssm_state,

286 - hidden_states.transpose(0, 1),

287 - discrete_time_step.transpose(0, 1),

341 + conv_out_d.transpose(0, 1),

342 + discrete_time_step_d.transpose(0, 1),

343 self.A,

289 - B,

290 - C,

344 + B_d,

345 + C_d,

346 self.D,

292 - gate.transpose(0, 1),

347 + gate_d.transpose(0, 1),

348 time_proj_bias,

349 dt_softplus=True,

295 - state_batch_indices=state_indices_tensor,

296 - out=scan_outputs)

297 - scan_outputs = scan_outputs.transpose(0, 1)

298 -

299 - # 4. Final linear projection

300 - if self.is_lora_enabled:

301 - # lora kernel requires contiguous tensor

302 - contextualized_states = self.out_proj(

303 - scan_outputs.transpose(-2, -1).contiguous())[0]

350 + state_batch_indices=state_indices_tensor_d,

351 + out=scan_outputs_d)

352 + scan_outputs_d = scan_outputs_d.transpose(0, 1)

353 +

354 + if envs.VLLM_USE_V1:

355 + ssm_outputs.insert(0, scan_outputs_d)

356 + else:

357 + ssm_outputs.append(scan_outputs_d)

358 +

359 + scan_outputs_combined = ssm_outputs[0] if len(

360 + ssm_outputs) == 1 else torch.cat(ssm_outputs, dim=-1)

361 +

362 + # 5. Final output projection

363 + if self.is_lora_enabled: # Lora kernel requires contiguous tensor.

364 + scan_outputs_combined = scan_outputs_combined.transpose(

365 + -2, -1).contiguous()

366 + out = self.out_proj(scan_outputs_combined)[0]

367 else:

305 - contextualized_states = self.out_proj(

306 - scan_outputs.transpose(-2, -1))[0]

307 - return contextualized_states

368 + out = self.out_proj(scan_outputs_combined.transpose(-2, -1))[0]

369 +

370 + return out

371

372 def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:

373 return MambaStateShapeCalculator.mamba1_state_shape(

@@ -317,3 +380,69 @@ class MambaMixer(MambaBase, CustomOp):

380 @property

381 def mamba_type(self) -> str:

382 return "mamba1"

383 +

384 + def _time_proj_bias(self) -> Optional[torch.Tensor]:

385 + if hasattr(self.dt_proj, "bias") and self.dt_proj.bias is not None:

386 + return self.dt_proj.bias.float()

387 + return None

388 +

389 +

390 +class PrefillDecodeSplit(NamedTuple):

391 + hidden_states_BC_p: torch.Tensor

392 + hidden_states_BC_d: torch.Tensor

393 + gate_p: torch.Tensor

394 + gate_d: torch.Tensor

395 + state_indices_tensor_p: torch.Tensor

396 + state_indices_tensor_d: torch.Tensor

397 + query_start_loc_p: Optional[torch.Tensor]

398 + has_initial_states_p: Optional[torch.Tensor]

399 +

400 +

401 +def split_batch_to_prefill_and_decode(

402 + hidden_states_BC: torch.Tensor,

403 + gate: torch.Tensor,

404 + state_indices_tensor: torch.Tensor,

405 + query_start_loc: torch.Tensor,

406 + has_initial_states: Optional[torch.Tensor],

407 + num_prefill_tokens: int,

408 + num_decode_tokens: int,

409 + num_prefills: int,

410 + num_decodes: int,

411 +) -> PrefillDecodeSplit:

412 + if envs.VLLM_USE_V1:

413 + # In v1, decode tokens come first, then prefill tokens.

414 + hidden_states_BC_d, hidden_states_BC_p = torch.split(

415 + hidden_states_BC, [num_decode_tokens, num_prefill_tokens], dim=-1)

416 + gate_d, gate_p = torch.split(gate,

417 + [num_decode_tokens, num_prefill_tokens],

418 + dim=-1)

419 + state_indices_tensor_d, state_indices_tensor_p = torch.split(

420 + state_indices_tensor, [num_decodes, num_prefills], dim=0)

421 + query_start_loc_p = (query_start_loc[-num_prefills - 1:] -

422 + num_decodes if num_prefills > 0 else None)

423 + has_initial_states_p = has_initial_states[-num_prefills:] if (

424 + has_initial_states is not None and num_prefills > 0) else None

425 + else:

426 + # In v0, prefill tokens come first, then decode tokens.

427 + hidden_states_BC_p, hidden_states_BC_d = torch.split(

428 + hidden_states_BC, [num_prefill_tokens, num_decode_tokens], dim=-1)

429 + gate_p, gate_d = torch.split(gate,

430 + [num_prefill_tokens, num_decode_tokens],

431 + dim=-1)

432 + state_indices_tensor_p, state_indices_tensor_d = torch.split(

433 + state_indices_tensor, [num_prefills, num_decodes], dim=0)

434 + query_start_loc_p = (query_start_loc[:num_prefills +

435 + 1] if num_prefills > 0 else None)

436 + has_initial_states_p = has_initial_states[:num_prefills] if (

437 + has_initial_states is not None and num_prefills > 0) else None

438 +

439 + return PrefillDecodeSplit(

440 + hidden_states_BC_p=hidden_states_BC_p,

441 + hidden_states_BC_d=hidden_states_BC_d,

442 + gate_p=gate_p,

443 + gate_d=gate_d,

444 + state_indices_tensor_p=state_indices_tensor_p,

445 + state_indices_tensor_d=state_indices_tensor_d,

446 + query_start_loc_p=query_start_loc_p,

447 + has_initial_states_p=has_initial_states_p,

448 + )

vllm/v1/attention/backends/mamba1_attn.py

+21 -5

@@ -2,14 +2,15 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from dataclasses import dataclass

5 -from typing import ClassVar

5 +from typing import ClassVar, Optional

7 import torch

9 from vllm.attention.backends.abstract import AttentionBackend

10 from vllm.config import VllmConfig

11 from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,

12 - CommonAttentionMetadata)

12 + CommonAttentionMetadata,

13 + split_decodes_and_prefills)

14 from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec

@@ -25,12 +26,15 @@ class Mamba1AttentionMetadata:

26 query_start_loc: torch.Tensor

27 context_lens_tensor: torch.Tensor

28 state_indices_tensor: torch.Tensor

28 - has_initial_states: torch.Tensor

29 + has_initial_states: Optional[torch.Tensor]

30 + num_prefills: int

31 + num_prefill_tokens: int

32 + num_decodes: int

33 + num_decode_tokens: int

36 class Mamba1AttentionMetadataBuilder(

37 AttentionMetadataBuilder[Mamba1AttentionMetadata]):

33 -

38 reorder_batch_threshold: ClassVar[int] = 1

40 def __init__(

@@ -57,11 +61,23 @@ class Mamba1AttentionMetadataBuilder(

61 state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]

62 context_lens_tensor = common_attn_metadata.num_computed_tokens_cpu.to(

63 query_start_loc.device)

60 - has_initial_states = (context_lens_tensor > 0)

64 +

65 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

66 + split_decodes_and_prefills(common_attn_metadata,

67 + decode_threshold=1))

68 +

69 + has_initial_states = None

70 +

71 + if num_prefills > 0:

72 + has_initial_states = context_lens_tensor > 0

74 return Mamba1AttentionMetadata(

75 query_start_loc=query_start_loc,

76 context_lens_tensor=context_lens_tensor,

77 has_initial_states=has_initial_states,

78 state_indices_tensor=state_indices_tensor,

79 + num_prefills=num_prefills,

80 + num_prefill_tokens=num_prefill_tokens,

81 + num_decodes=num_decodes,

82 + num_decode_tokens=num_decode_tokens,

83 )