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

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

[Kernel][Mamba] Optimize Mamba2 SSD prefill Triton kernels (#35397)

Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com>

tomeras91 committed Mar 4, 2026 at 20:47 UTC 7faba503c403bc8c562888df3a841b6df104d042

6 files changed +155 -369

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

+3 -2

@@ -8,6 +8,7 @@ import torch

8 from packaging import version

10 from vllm import _custom_ops as ops

11 +from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp

12 from vllm.triton_utils import HAS_TRITON, tl, triton

13 from vllm.v1.attention.backends.utils import PAD_SLOT_ID

@@ -215,7 +216,7 @@ def _selective_scan_update_kernel(

216 mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate),

217 other=0.0,

218 ).to(tl.float32)

218 - dA = tl.exp(A * dt[:, None])

219 + dA = fast_exp(A * dt[:, None])

220 else:

221 dt = tl.load(dt_ptr).to(tl.float32)

222 if HAS_DT_BIAS:

@@ -223,7 +224,7 @@ def _selective_scan_update_kernel(

224 if DT_SOFTPLUS:

225 dt = softplus(dt)

226 A = tl.load(A_ptr).to(tl.float32)

226 - dA = tl.exp(A * dt) # scalar, not a matrix

227 + dA = fast_exp(A * dt) # scalar, not a matrix

228

229 B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)

230 C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)

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

+73 -2

@@ -8,6 +8,7 @@

9 from packaging import version

11 +from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp

12 from vllm.triton_utils import tl, triton

14 TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0")

@@ -15,6 +16,76 @@ TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0")

17 @triton.autotune(

18 configs=[

19 + # =================================================================

20 + # Higher warp count configs for better latency hiding

21 + # More warps = more instructions in flight = better memory latency hiding

22 + # =================================================================

23 + triton.Config(

24 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},

25 + num_stages=2,

26 + num_warps=8,

27 + ),

28 + triton.Config(

29 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

30 + num_stages=2,

31 + num_warps=8,

32 + ),

33 + triton.Config(

34 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},

35 + num_stages=2,

36 + num_warps=8,

37 + ),

38 + # Smaller tiles with more stages for software pipelining

39 + triton.Config(

40 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},

41 + num_stages=3,

42 + num_warps=4,

43 + ),

44 + triton.Config(

45 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64},

46 + num_stages=2,

47 + num_warps=4,

48 + ),

49 + # =================================================================

50 + # Low register pressure configs (num_stages=1) for large dstate

51 + # =================================================================

52 + triton.Config(

53 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},

54 + num_stages=1,

55 + num_warps=4,

56 + ),

57 + triton.Config(

58 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

59 + num_stages=1,

60 + num_warps=4,

61 + ),

62 + triton.Config(

63 + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

64 + num_stages=1,

65 + num_warps=4,

66 + ),

67 + triton.Config(

68 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

69 + num_stages=1,

70 + num_warps=4,

71 + ),

72 + # num_stages=2 configs - moderate register pressure

73 + triton.Config(

74 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},

75 + num_stages=2,

76 + num_warps=4,

77 + ),

78 + triton.Config(

79 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

80 + num_stages=2,

81 + num_warps=4,

82 + ),

83 + triton.Config(

84 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

85 + num_stages=2,

86 + num_warps=4,

87 + ),

88 + # Original configs for larger dstate values

89 triton.Config(

90 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64},

91 num_stages=3,

@@ -200,7 +271,7 @@ def _chunk_scan_fwd_kernel(

271 offs_m[:, None] * stride_C_seqlen + offs_k_dstate[None, :] * stride_C_dstate

272 )

273

203 - scale_m = tl.exp(dA_cs_m)

274 + scale_m = fast_exp(dA_cs_m)

275 if BLOCK_SIZE_DSTATE <= 128:

276 C = tl.load(

277 C_ptrs,

@@ -285,7 +356,7 @@ def _chunk_scan_fwd_kernel(

356 )

357 # If there's seq_idx, we already set cb[i, j] = 0 for seq_idx[i] != seq_idx[j].

358 # So we don't need masking wrt seq_idx here.

288 - cb *= tl.exp(dA_cs_m[:, None] - dA_cs_k[None, :])

359 + cb *= fast_exp(dA_cs_m[:, None] - dA_cs_k[None, :])

360 dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32)

361 cb *= dt_k

362 if IS_CAUSAL:

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

+30 -319

@@ -8,6 +8,7 @@

9 import torch

11 +from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp

12 from vllm.triton_utils import tl, triton

14 from .mamba_ssm import softplus

@@ -116,6 +117,34 @@ def _chunk_cumsum_fwd_kernel(

117

118 @triton.autotune(

119 configs=[

120 + # Small headdim/dstate configs (hdim<=64, dstate<=128) - increased parallelism

121 + triton.Config(

122 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},

123 + num_stages=3,

124 + num_warps=4,

125 + ),

126 + triton.Config(

127 + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

128 + num_stages=3,

129 + num_warps=4,

130 + ),

131 + triton.Config(

132 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},

133 + num_stages=3,

134 + num_warps=4,

135 + ),

136 + # Low register pressure configs for large dstate (dstate=128)

137 + triton.Config(

138 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},

139 + num_stages=2,

140 + num_warps=4,

141 + ),

142 + triton.Config(

143 + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64},

144 + num_stages=2,

145 + num_warps=4,

146 + ),

147 + # original configs for larger headdim/dstate values

148 triton.Config(

149 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64},

150 num_stages=3,

@@ -251,7 +280,7 @@ def _chunk_state_fwd_kernel(

280 dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0).to(

281 tl.float32

282 )

254 - scale = tl.exp(dA_cs_last - dA_cs_k) * dt_k

283 + scale = fast_exp(dA_cs_last - dA_cs_k) * dt_k

284 b *= scale[:, None]

285 b = b.to(x_ptr.dtype.element_ty)

286 acc += tl.dot(x, b)

@@ -273,238 +302,6 @@ def _chunk_state_fwd_kernel(

302 tl.store(states_ptrs, states, mask=c_mask)

303

304

276 -@triton.autotune(

277 - configs=[

278 - triton.Config(

279 - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64},

280 - num_stages=3,

281 - num_warps=8,

282 - ),

283 - triton.Config(

284 - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32},

285 - num_stages=4,

286 - num_warps=4,

287 - ),

288 - triton.Config(

289 - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},

290 - num_stages=4,

291 - num_warps=4,

292 - ),

293 - triton.Config(

294 - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

295 - num_stages=4,

296 - num_warps=4,

297 - ),

298 - triton.Config(

299 - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},

300 - num_stages=4,

301 - num_warps=4,

302 - ),

303 - triton.Config(

304 - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},

305 - num_stages=4,

306 - num_warps=4,

307 - ),

308 - triton.Config(

309 - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},

310 - num_stages=5,

311 - num_warps=2,

312 - ),

313 - triton.Config(

314 - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

315 - num_stages=5,

316 - num_warps=2,

317 - ),

318 - triton.Config(

319 - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},

320 - num_stages=4,

321 - num_warps=2,

322 - ),

323 - ],

324 - key=["hdim", "dstate", "chunk_size"],

325 -)

326 -@triton.jit

327 -def _chunk_state_varlen_kernel(

328 - # Pointers to matrices

329 - x_ptr,

330 - b_ptr,

331 - dt_ptr,

332 - dA_cumsum_ptr,

333 - chunk_states_ptr,

334 - cu_seqlens_ptr,

335 - states_ptr,

336 - initstates_ptr,

337 - # Matrix dimensions

338 - hdim: tl.constexpr,

339 - dstate: tl.constexpr,

340 - chunk_size: tl.constexpr,

341 - nheads_ngroups_ratio: tl.constexpr,

342 - # Strides

343 - stride_x_seqlen: tl.int64,

344 - stride_x_head: tl.int64,

345 - stride_x_hdim: tl.constexpr,

346 - stride_b_seqlen: tl.int64,

347 - stride_b_head: tl.int64,

348 - stride_b_dstate: tl.constexpr,

349 - stride_dt_head: tl.int64,

350 - stride_dt_chunk: tl.int64,

351 - stride_dt_csize: tl.constexpr,

352 - stride_dA_cs_head: tl.int64,

353 - stride_dA_cs_chunk: tl.int64,

354 - stride_dA_cs_csize: tl.constexpr,

355 - stride_chunk_states_chunk: tl.int64,

356 - stride_chunk_states_head: tl.int64,

357 - stride_chunk_states_hdim: tl.int64,

358 - stride_chunk_states_dstate: tl.constexpr,

359 - stride_states_batch: tl.int64,

360 - stride_states_head: tl.int64,

361 - stride_states_hdim: tl.int64,

362 - stride_states_dstate: tl.constexpr,

363 - stride_init_states_batch: tl.int64,

364 - stride_init_states_head: tl.int64,

365 - stride_init_states_hdim: tl.int64,

366 - stride_init_states_dstate: tl.constexpr,

367 - # Meta-parameters

368 - BLOCK_SIZE_M: tl.constexpr,

369 - BLOCK_SIZE_N: tl.constexpr,

370 - BLOCK_SIZE_K: tl.constexpr,

371 - HAS_INITSTATES: tl.constexpr,

372 -):

373 - pid_b = tl.program_id(axis=1)

374 - pid_h = tl.program_id(axis=2)

375 - num_pid_n = tl.cdiv(dstate, BLOCK_SIZE_N)

376 - pid_m = tl.program_id(axis=0) // num_pid_n

377 - pid_n = tl.program_id(axis=0) % num_pid_n

378 - end_idx = tl.load(cu_seqlens_ptr + pid_b + 1)

379 - pid_c = (end_idx - 1) // chunk_size

380 - b_ptr += (

381 - pid_c * chunk_size * stride_b_seqlen

382 - + (pid_h // nheads_ngroups_ratio) * stride_b_head

383 - )

384 - x_ptr += pid_c * chunk_size * stride_x_seqlen + pid_h * stride_x_head

385 - dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head

386 - dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head

387 - chunk_states_ptr += (

388 - pid_c * stride_chunk_states_chunk + pid_h * stride_chunk_states_head

389 - )

390 -

391 - if HAS_INITSTATES:

392 - # if there are init states provided, we differentiate between states (which

393 - # are boundary conditions at a chunk boundary) and initstates (which are boundary

394 - # conditions when a new example in a cont batch starts)

395 - initstates_ptr += pid_h * stride_init_states_head

396 -

397 - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)

398 - offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

399 - offs_k = tl.arange(0, BLOCK_SIZE_K)

400 - x_ptrs = x_ptr + (

401 - offs_m[:, None] * stride_x_hdim + offs_k[None, :] * stride_x_seqlen

402 - )

403 - b_ptrs = b_ptr + (

404 - offs_n[None, :] * stride_b_dstate + offs_k[:, None] * stride_b_seqlen

405 - )

406 - dt_ptrs = dt_ptr + offs_k * stride_dt_csize

407 - dA_cs_last = tl.load(

408 - dA_cumsum_ptr + (end_idx - pid_c * chunk_size - 1) * stride_dA_cs_csize

409 - ).to(tl.float32)

410 - dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize

411 -

412 - chunk_size_limit = end_idx - pid_c * chunk_size

413 - start_idx = tl.load(cu_seqlens_ptr + pid_b)

414 - start_idx_cur = tl.maximum(start_idx - pid_c * chunk_size, 0)

415 -

416 - acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

417 - for k in range(0, chunk_size_limit, BLOCK_SIZE_K):

418 - x = tl.load(

419 - x_ptrs,

420 - mask=(offs_m[:, None] < hdim)

421 - & (offs_k[None, :] < chunk_size_limit - k)

422 - & (offs_k[None, :] >= start_idx_cur - k),

423 - other=0.0,

424 - )

425 - b = tl.load(

426 - b_ptrs,

427 - mask=(offs_k[:, None] < chunk_size_limit - k)

428 - & (offs_n[None, :] < dstate)

429 - & (offs_k[:, None] >= start_idx_cur - k),

430 - other=0.0,

431 - ).to(tl.float32)

432 - dA_cs_k = tl.load(

433 - dA_cumsum_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0

434 - ).to(tl.float32)

435 - dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0).to(

436 - tl.float32

437 - )

438 - scale = tl.where(

439 - (offs_k >= start_idx_cur - k) & (offs_k < chunk_size_limit - k),

440 - tl.exp(dA_cs_last - dA_cs_k) * dt_k,

441 - 0.0,

442 - )

443 - b *= scale[:, None]

444 - b = b.to(x_ptr.dtype.element_ty)

445 - acc += tl.dot(x, b)

446 - x_ptrs += BLOCK_SIZE_K * stride_x_seqlen

447 - b_ptrs += BLOCK_SIZE_K * stride_b_seqlen

448 - dt_ptrs += BLOCK_SIZE_K * stride_dt_csize

449 - dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize

450 -

451 - # If the sequence starts after the last chunk idx, we don't need to add the contribution from the last chunk

452 - # If HAS_INITSTATES==True need to consider two possibilities

453 - # - if start_idx < pid_c * chunk_size, then we need to take the past_states_ptrs

454 - # - if state_idx >= pid * chunk_size, then we need to insert initstates

455 - if (

456 - (start_idx < pid_c * chunk_size) # first chunk

457 - or (HAS_INITSTATES)

458 - ):

459 - dA_cs_boundary = 0.0 # default

460 -

461 - if not HAS_INITSTATES:

462 - past_states_ptrs = chunk_states_ptr + (

463 - offs_m[:, None] * stride_chunk_states_hdim

464 - + offs_n[None, :] * stride_chunk_states_dstate

465 - )

466 - else:

467 - # - this seems repetitive, buts its to help the compiler

468 - if start_idx < pid_c * chunk_size:

469 - past_states_ptrs = chunk_states_ptr + (

470 - offs_m[:, None] * stride_chunk_states_hdim

471 - + offs_n[None, :] * stride_chunk_states_dstate

472 - )

473 - else:

474 - past_states_ptrs = initstates_ptr + (

475 - pid_b * stride_init_states_batch

476 - + offs_m[:, None] * stride_init_states_hdim

477 - + offs_n[None, :] * stride_init_states_dstate

478 - )

479 -

480 - # need to adjust the boundary

481 - if start_idx > pid_c * chunk_size:

482 - dA_cs_boundary = tl.load(

483 - dA_cumsum_ptr

484 - + (start_idx - pid_c * chunk_size - 1) * stride_dA_cs_csize

485 - ).to(tl.float32)

486 -

487 - past_states = tl.load(

488 - past_states_ptrs,

489 - mask=(offs_m[:, None] < hdim) & (offs_n[None, :] < dstate),

490 - other=0.0,

491 - ).to(tl.float32)

492 -

493 - scale = tl.exp(dA_cs_last - dA_cs_boundary)

494 - acc += past_states * scale

495 -

496 - states = acc.to(states_ptr.dtype.element_ty)

497 -

498 - states_ptr += pid_b * stride_states_batch + pid_h * stride_states_head

499 - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)

500 - offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

501 - states_ptrs = states_ptr + (

502 - offs_m[:, None] * stride_states_hdim + offs_n[None, :] * stride_states_dstate

503 - )

504 - c_mask = (offs_m[:, None] < hdim) & (offs_n[None, :] < dstate)

505 - tl.store(states_ptrs, states, mask=c_mask)

506 -

507 -

305 def _chunk_cumsum_fwd(

306 dt,

307 A,

@@ -612,89 +409,3 @@ def _chunk_state_fwd(

409 stride_dA_cs_csize=dA_cumsum.stride(2),

410 )

411 return states

615 -

616 -

617 -def chunk_state_varlen(

618 - B, x, dt, dA_cumsum, cu_seqlens, chunk_states, initial_states=None

619 -):

620 - total_seqlen, nheads, headdim = x.shape

621 - _, nchunks, chunk_size = dt.shape

622 - _, ngroups, dstate = B.shape

623 - batch = cu_seqlens.shape[0] - 1

624 - cu_seqlens = cu_seqlens.contiguous()

625 - assert nheads % ngroups == 0

626 - assert B.shape == (total_seqlen, ngroups, dstate)

627 - assert dt.shape == (nheads, nchunks, chunk_size)

628 - assert dA_cumsum.shape == dt.shape

629 - assert chunk_states.shape == (nchunks, nheads, headdim, dstate)

630 -

631 - if initial_states is not None:

632 - assert initial_states.shape == (batch, nheads, headdim, dstate)

633 -

634 - states = torch.empty(

635 - batch,

636 - nheads,

637 - headdim,

638 - dstate,

639 - dtype=chunk_states.dtype,

640 - device=chunk_states.device,

641 - )

642 -

643 - initial_states_strides = (

644 - (

645 - initial_states.stride(0),

646 - initial_states.stride(1),

647 - initial_states.stride(2),

648 - initial_states.stride(3),

649 - )

650 - if initial_states is not None

651 - else (0, 0, 0, 0)

652 - )

653 -

654 - grid = lambda META: (

655 - triton.cdiv(headdim, META["BLOCK_SIZE_M"])

656 - * triton.cdiv(dstate, META["BLOCK_SIZE_N"]),

657 - batch,

658 - nheads,

659 - )

660 - with torch.cuda.device(x.device.index):

661 - _chunk_state_varlen_kernel[grid](

662 - x_ptr=x,

663 - b_ptr=B,

664 - dt_ptr=dt,

665 - dA_cumsum_ptr=dA_cumsum,

666 - chunk_states_ptr=chunk_states,

667 - cu_seqlens_ptr=cu_seqlens,

668 - states_ptr=states,

669 - initstates_ptr=initial_states,

670 - hdim=headdim,

671 - dstate=dstate,

672 - chunk_size=chunk_size,

673 - nheads_ngroups_ratio=nheads // ngroups,

674 - stride_x_seqlen=x.stride(0),

675 - stride_x_head=x.stride(1),

676 - stride_x_hdim=x.stride(2),

677 - stride_b_seqlen=B.stride(0),

678 - stride_b_head=B.stride(1),

679 - stride_b_dstate=B.stride(2),

680 - stride_dt_head=dt.stride(0),

681 - stride_dt_chunk=dt.stride(1),

682 - stride_dt_csize=dt.stride(2),

683 - stride_dA_cs_head=dA_cumsum.stride(0),

684 - stride_dA_cs_chunk=dA_cumsum.stride(1),

685 - stride_dA_cs_csize=dA_cumsum.stride(2),

686 - stride_chunk_states_chunk=chunk_states.stride(0),

687 - stride_chunk_states_head=chunk_states.stride(1),

688 - stride_chunk_states_hdim=chunk_states.stride(2),

689 - stride_chunk_states_dstate=chunk_states.stride(3),

690 - stride_states_batch=states.stride(0),

691 - stride_states_head=states.stride(1),

692 - stride_states_hdim=states.stride(2),

693 - stride_states_dstate=states.stride(3),

694 - stride_init_states_batch=initial_states_strides[0],

695 - stride_init_states_head=initial_states_strides[1],

696 - stride_init_states_hdim=initial_states_strides[2],

697 - stride_init_states_dstate=initial_states_strides[3],

698 - HAS_INITSTATES=initial_states is not None,

699 - )

700 - return states

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

+3 -6

@@ -107,18 +107,15 @@ def _mamba_chunk_scan_combined_fwd(

107

108 # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries

109 # (middle term of factorization of off-diag blocks; A terms)

110 - # - for handling chunked prefill, this requires i) initial_states and

111 - # ii) seq_idx to be all specified.

112 - # - When a new seq_idx is detected, we will stop passing the prev_state

113 - # and switch accordingly to the init_state corresponding to the new seq_idx.

110 + # - parallelized across sequences using last_chunk_indices to derive

111 + # per-sequence chunk ranges. Each sequence's state passing runs independently.

112 states = _state_passing_fwd(

113 rearrange(states, "... p n -> ... (p n)"),

114 dA_cumsum, # (nheads, nchunks, chunk_size)

117 - cu_chunk_seqlens,

115 + last_chunk_indices,

116 initial_states=rearrange(initial_states, "... p n -> ... (p n)")

117 if initial_states is not None

118 else None, # (batch, nheads, headdim*dstate)

121 - seq_idx=seq_idx,

119 out_dtype=state_dtype if state_dtype is not None else C.dtype,

120 )

121 states = rearrange(states, "... (p n) -> ... p n", n=dstate)

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

+29 -40

@@ -8,6 +8,7 @@

9 import torch

11 +from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp

12 from vllm.triton_utils import tl, triton

@@ -29,12 +30,9 @@ def _state_passing_fwd_kernel(

30 out_ptr,

31 dA_cs_ptr,

32 initstates_ptr,

32 - seq_idx_ptr,

33 - cu_chunk_seqlens_ptr,

33 + last_chunk_indices_ptr,

34 # Matrix dimensions

35 dim: tl.constexpr,

36 - nchunks,

37 - seqlen,

36 chunk_size: tl.constexpr,

37 # Strides

38 stride_states_chunk: tl.int64,

@@ -49,55 +47,51 @@ def _state_passing_fwd_kernel(

47 stride_initstates_batch: tl.int64,

48 stride_initstates_head: tl.int64,

49 stride_initstates_dim: tl.constexpr,

52 - stride_seq_idx_chunk: tl.constexpr,

50 # Meta-parameters

51 HAS_INITSTATES: tl.constexpr,

52 BLOCK_SIZE: tl.constexpr,

53 ):

57 - pid_h = tl.program_id(axis=1)

54 pid_m = tl.program_id(axis=0)

55 + pid_b = tl.program_id(axis=1)

56 + pid_h = tl.program_id(axis=2)

60 - states_ptr += pid_h * stride_states_head

61 - dA_cs_ptr += pid_h * stride_dA_cs_head + (chunk_size - 1) * stride_dA_cs_csize

62 - out_ptr += pid_h * stride_out_head

58 + # Derive this sequence's chunk range from last_chunk_indices

59 + chunk_end = tl.load(last_chunk_indices_ptr + pid_b) + 1

60 + chunk_start = (

61 + tl.load(last_chunk_indices_ptr + pid_b - 1, mask=pid_b > 0, other=-1) + 1

62 + )

63 +

64 + # Offset pointers to this sequence's first chunk

65 + states_ptr += chunk_start * stride_states_chunk + pid_h * stride_states_head

66 + dA_cs_ptr += (

67 + pid_h * stride_dA_cs_head

68 + + chunk_start * stride_dA_cs_chunk

69 + + (chunk_size - 1) * stride_dA_cs_csize

70 + )

71 + out_ptr += chunk_start * stride_out_chunk + pid_h * stride_out_head

73 offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

74 states_ptrs = states_ptr + offs_m * stride_states_dim

75 out_ptrs = out_ptr + offs_m * stride_out_dim

77 + # Load initial state once — no per-chunk branching needed

78 if HAS_INITSTATES:

79 initstates_ptrs = (

80 initstates_ptr

81 + + pid_b * stride_initstates_batch

82 + pid_h * stride_initstates_head

83 + offs_m * stride_initstates_dim

84 )

74 -

85 states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

86 else:

87 states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)

79 - prev_seq_idx = 0

80 - for c in range(nchunks):

89 + # Loop over only this sequence's chunks — branchless

90 + nchunks_this_seq = chunk_end - chunk_start

91 + for _ in range(nchunks_this_seq):

92 new_states = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

93 dA_cs = tl.load(dA_cs_ptr).to(tl.float32)

83 - seq_idx = tl.load(seq_idx_ptr + c * stride_seq_idx_chunk)

84 - # we have started a new sequence

85 - if prev_seq_idx != seq_idx:

86 - if HAS_INITSTATES:

87 - initstates_ptrs = (

88 - initstates_ptr

89 - + seq_idx * stride_initstates_batch

90 - + pid_h * stride_initstates_head

91 - + offs_m * stride_initstates_dim

92 - )

93 - states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to(

94 - tl.float32

95 - )

96 - else:

97 - states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)

98 -

99 - prev_seq_idx = seq_idx

100 - states = tl.exp(dA_cs) * states + new_states

94 + states = fast_exp(dA_cs) * states + new_states

95 tl.store(out_ptrs, states, mask=offs_m < dim)

97 states_ptrs += stride_states_chunk

@@ -108,15 +102,14 @@ def _state_passing_fwd_kernel(

102 def _state_passing_fwd(

103 states,

104 dA_cumsum,

111 - cu_chunk_seqlens,

112 - seq_idx,

105 + last_chunk_indices,

106 initial_states=None,

107 out_dtype=None,

108 ):

109 nchunks, nheads, dim = states.shape

110 chunk_size = dA_cumsum.shape[-1]

111 + batch = last_chunk_indices.shape[0]

112 assert dA_cumsum.shape == (nheads, nchunks, chunk_size)

119 - seqlen = seq_idx.shape[-1]

113 out_dtype = states.dtype if out_dtype is None else out_dtype

114 out = torch.empty((nchunks, nheads, dim), device=states.device, dtype=out_dtype)

115

@@ -126,19 +119,16 @@ def _state_passing_fwd(

119 else (0, 0, 0)

120 )

121

129 - grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), nheads)

122 + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), batch, nheads)

123 with torch.cuda.device(states.device.index):

124 _state_passing_fwd_kernel[grid](

125 states_ptr=states,

126 out_ptr=out,

127 dA_cs_ptr=dA_cumsum,

128 initstates_ptr=initial_states,

136 - seq_idx_ptr=seq_idx,

137 - cu_chunk_seqlens_ptr=cu_chunk_seqlens,

129 + last_chunk_indices_ptr=last_chunk_indices,

130 dim=dim,

139 - nchunks=nchunks,

140 - seqlen=seqlen if seq_idx is not None else 0,

141 - chunk_size=chunk_size if seq_idx is not None else 0,

131 + chunk_size=chunk_size,

132 stride_states_chunk=states.stride(0),

133 stride_states_head=states.stride(1),

134 stride_states_dim=states.stride(2),

@@ -151,7 +141,6 @@ def _state_passing_fwd(

141 stride_initstates_batch=initial_states_strides[0],

142 stride_initstates_head=initial_states_strides[1],

143 stride_initstates_dim=initial_states_strides[2],

154 - stride_seq_idx_chunk=seq_idx.stride(0),

144 HAS_INITSTATES=initial_states is not None,

145 )

146 return out

vllm/model_executor/layers/mamba/ops/triton_helpers.py new

+17

@@ -0,0 +1,17 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +from vllm.triton_utils import tl, triton

5 +

6 +

7 +@triton.jit

8 +def fast_exp(x):

9 + """Faster alternative to tl.exp() using the hardware exp2 instruction.

10 +

11 + tl.math.exp2 maps directly to a single ex2.approx.f32 PTX instruction,

12 + while tl.exp goes through libdevice __nv_expf which adds function call

13 + overhead and extra range checking.

14 + """

15 + # exp(x) = exp2(x * log2(e)), where log2(e) = 1/ln(2) = 1.4426950408889634

16 + LOG2E = tl.constexpr(1.4426950408889634)

17 + return tl.math.exp2(LOG2E * x)