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
[Model] Mamba2 varlen refactor (#21467)
Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com> Co-authored-by: RishiAstra <40644327+RishiAstra@users.noreply.github.com>
Chih-Chieh Yang committed
Sep 26, 2025 at 07:31 UTC
2b6b1d78092826c789514c5a901d1e872d4602dd
10 files changed
+723
-865
tests/kernels/mamba/test_mamba_ssm_ssd.py
+71
-54
@@ -7,7 +7,7 @@ import torch.nn.functional as F
7
from einops import rearrange, repeat
8
9
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
10
- mamba_chunk_scan_combined)
10
+ mamba_chunk_scan_combined_varlen)
11
from vllm.platforms import current_platform
12
from vllm.v1.attention.backends.mamba2_attn import (
13
_query_start_loc_to_chunk_indices_offsets)
@@ -185,9 +185,14 @@ def generate_continuous_batched_examples(example_lens_by_batch,
185
IND_S = [x % full_length for x in IND_E]
186
IND_E = [end_boundary(x + y) for x, y in zip(IND_S, spec)]
187
188
+ # varlen has implicit batch=1
189
+ dt2 = dt2.squeeze(0)
190
+ X2 = X2.squeeze(0)
191
+ B2 = B2.squeeze(0)
192
+ C2 = C2.squeeze(0)
193
yield ([Y_min[s, IND_S[s]:IND_E[s]]
194
for s in range(num_examples)] if return_naive_ref else None,
190
- cu_seqlens, seq_idx.unsqueeze(0), (A, dt2, X2, B2, C2))
195
+ cu_seqlens, seq_idx, (A, dt2, X2, B2, C2))
196
197
198
@pytest.mark.parametrize("itype",
@@ -198,7 +203,7 @@ def generate_continuous_batched_examples(example_lens_by_batch,
203
def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size,
204
itype):
205
201
- # this tests the kernels on a single example (no batching)
206
+ # this tests the kernels on a single example (bs=1)
207
208
# TODO: the bfloat16 case requires higher thresholds. To be investigated
209
@@ -219,23 +224,40 @@ def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size,
224
225
Y_min, final_state_min = ssd_minimal_discrete(X * dt.unsqueeze(-1), A * dt,
226
B, C, chunk_size)
227
+
228
+ cu_seqlens = torch.tensor((0, seqlen), device='cuda').cumsum(dim=0)
229
+ seq_idx = torch.zeros(seqlen, dtype=torch.int32, device=cu_seqlens.device)
230
+
231
+ chunk_indices, chunk_offsets = \
232
+ _query_start_loc_to_chunk_indices_offsets(
233
+ cu_seqlens, chunk_size, cu_seqlens[-1])
234
+
235
+ # varlen has implicit batch=1
236
+ X = X.squeeze(0)
237
+ dt = dt.squeeze(0)
238
+ A = A.squeeze(0)
239
+ B = B.squeeze(0)
240
+ C = C.squeeze(0)
241
Y = torch.empty_like(X)
223
- final_state = mamba_chunk_scan_combined(X,
224
- dt,
225
- A,
226
- B,
227
- C,
228
- chunk_size,
229
- D=None,
230
- return_final_states=True,
231
- out=Y)
242
+ final_state = mamba_chunk_scan_combined_varlen(X,
243
+ dt,
244
+ A,
245
+ B,
246
+ C,
247
+ chunk_size,
248
+ D=None,
249
+ cu_seqlens=cu_seqlens,
250
+ seq_idx=seq_idx,
251
+ chunk_indices=chunk_indices,
252
+ chunk_offsets=chunk_offsets,
253
+ out=Y)
254
255
# just test the last in sequence
234
- torch.testing.assert_close(Y[:, -1], Y_min[:, -1], atol=atol, rtol=rtol)
256
+ torch.testing.assert_close(Y[-1], Y_min[0, -1], atol=atol, rtol=rtol)
257
258
# just test the last head
259
# NOTE, in the kernel we always cast states to fp32
238
- torch.testing.assert_close(final_state[:, -1],
260
+ torch.testing.assert_close(final_state[:, -1].to(torch.float32),
261
final_state_min[:, -1].to(torch.float32),
262
atol=atol,
263
rtol=rtol)
@@ -300,7 +322,7 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
322
cu_seqlens, chunk_size, cu_seqlens[-1])
323
324
Y = torch.empty_like(X)
303
- new_states = mamba_chunk_scan_combined(
325
+ new_states = mamba_chunk_scan_combined_varlen(
326
X,
327
dt,
328
A,
@@ -312,7 +334,6 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
334
seq_idx=seq_idx,
335
chunk_indices=chunk_indices,
336
chunk_offsets=chunk_offsets,
315
- return_varlen_states=True,
337
initial_states=states,
338
out=Y,
339
)
@@ -321,7 +342,7 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
342
for i in range(num_examples):
343
344
# just test one dim and dstate
324
- Y_eg = Y[0, cu_seqlens[i]:cu_seqlens[i + 1], 0, 0]
345
+ Y_eg = Y[cu_seqlens[i]:cu_seqlens[i + 1], 0, 0]
346
Y_min_eg = Y_min[i][:, 0, 0]
347
torch.testing.assert_close(Y_eg, Y_min_eg, atol=atol, rtol=rtol)
348
@@ -386,7 +407,7 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
407
_query_start_loc_to_chunk_indices_offsets(
408
cu_seqlens, chunk_size, cu_seqlens[-1])
409
Y_ref = torch.empty_like(X)
389
- state_ref = mamba_chunk_scan_combined(
410
+ state_ref = mamba_chunk_scan_combined_varlen(
411
X,
412
dt,
413
A,
@@ -398,7 +419,6 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
419
seq_idx=seq_idx,
420
chunk_indices=chunk_indices,
421
chunk_offsets=chunk_offsets,
401
- return_varlen_states=True,
422
initial_states=None,
423
out=Y_ref,
424
)
@@ -414,27 +434,27 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
434
chunked_seq_idx = torch.repeat_interleave(
435
torch.arange(len(chunked_seqlens), device=device),
436
chunked_seqlens,
417
- output_size=chunked_cu_seqlens[-1]).unsqueeze(0).to(torch.int32)
437
+ output_size=chunked_cu_seqlens[-1]).to(torch.int32)
438
chunked_input_seq_len = chunked_cu_seqlens[-1]
419
- X_chunked = torch.zeros_like(X)[:, :chunked_input_seq_len, ...]
420
- dt_chunked = torch.zeros_like(dt)[:, :chunked_input_seq_len, ...]
421
- B_chunked = torch.zeros_like(B)[:, :chunked_input_seq_len, ...]
422
- C_chunked = torch.zeros_like(C)[:, :chunked_input_seq_len, ...]
439
+ X_chunked = torch.zeros_like(X)[:chunked_input_seq_len, ...]
440
+ dt_chunked = torch.zeros_like(dt)[:chunked_input_seq_len, ...]
441
+ B_chunked = torch.zeros_like(B)[:chunked_input_seq_len, ...]
442
+ C_chunked = torch.zeros_like(C)[:chunked_input_seq_len, ...]
443
for i in range(num_sequences):
444
# fmt: off
425
- chunk_f = lambda x, i: x[:, cu_seqlens[i]:cu_seqlens[i] + chunked_seqlens[i], ...] # noqa: E501
445
+ chunk_f = lambda x, i: x[cu_seqlens[i]:cu_seqlens[i] + chunked_seqlens[i], ...] # noqa: E501
446
427
- X_chunked[:, chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(X, i) # noqa: E501
428
- dt_chunked[:, chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(dt, i) # noqa: E501
429
- B_chunked[:, chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(B, i) # noqa: E501
430
- C_chunked[:, chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(C, i) # noqa: E501
447
+ X_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(X, i) # noqa: E501
448
+ dt_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(dt, i) # noqa: E501
449
+ B_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(B, i) # noqa: E501
450
+ C_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(C, i) # noqa: E501
451
# fmt: on
452
453
chunk_indices, chunk_offsets = \
454
_query_start_loc_to_chunk_indices_offsets(
455
chunked_cu_seqlens, chunk_size, chunked_cu_seqlens[-1])
456
Y_partial = torch.empty_like(X_chunked)
437
- partial_state = mamba_chunk_scan_combined(
457
+ partial_state = mamba_chunk_scan_combined_varlen(
458
X_chunked,
459
dt_chunked,
460
A,
@@ -446,7 +466,6 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
466
seq_idx=chunked_seq_idx,
467
chunk_indices=chunk_indices,
468
chunk_offsets=chunk_offsets,
449
- return_varlen_states=True,
469
initial_states=None,
470
out=Y_partial,
471
)
@@ -461,29 +480,28 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
480
remaining_chunked_seq_idx = torch.repeat_interleave(
481
torch.arange(len(remaining_chunked_seqlens), device=device),
482
remaining_chunked_seqlens,
464
- output_size=remaining_chunked_cu_seqlens[-1]).unsqueeze(0).to(
465
- torch.int32)
483
+ output_size=remaining_chunked_cu_seqlens[-1]).to(torch.int32)
484
remaining_chunked_input_seq_len = remaining_chunked_cu_seqlens[-1]
485
# fmt: off
468
- remaining_X_chunked = torch.zeros_like(X)[:, :remaining_chunked_input_seq_len, ...] # noqa: E501
469
- remaining_dt_chunked = torch.zeros_like(dt)[:, :remaining_chunked_input_seq_len, ...] # noqa: E501
470
- remaining_B_chunked = torch.zeros_like(B)[:, :remaining_chunked_input_seq_len, ...] # noqa: E501
471
- remaining_C_chunked = torch.zeros_like(C)[:, :remaining_chunked_input_seq_len, ...] # noqa: E501
486
+ remaining_X_chunked = torch.zeros_like(X)[:remaining_chunked_input_seq_len, ...] # noqa: E501
487
+ remaining_dt_chunked = torch.zeros_like(dt)[:remaining_chunked_input_seq_len, ...] # noqa: E501
488
+ remaining_B_chunked = torch.zeros_like(B)[:remaining_chunked_input_seq_len, ...] # noqa: E501
489
+ remaining_C_chunked = torch.zeros_like(C)[:remaining_chunked_input_seq_len, ...] # noqa: E501
490
for i in range(num_sequences):
473
- remaining_chunk_f = lambda x, i: x[:, cu_seqlens[i] + chunked_seqlens[i]:cu_seqlens[i+1], ...] # noqa: E501
491
+ remaining_chunk_f = lambda x, i: x[cu_seqlens[i] + chunked_seqlens[i]:cu_seqlens[i+1], ...] # noqa: E501
492
475
- remaining_X_chunked[:, remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(X, i) # noqa: E501
476
- remaining_dt_chunked[:, remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(dt, i) # noqa: E501
477
- remaining_B_chunked[:, remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(B, i) # noqa: E501
478
- remaining_C_chunked[:, remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(C, i) # noqa: E501
493
+ remaining_X_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(X, i) # noqa: E501
494
+ remaining_dt_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(dt, i) # noqa: E501
495
+ remaining_B_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(B, i) # noqa: E501
496
+ remaining_C_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(C, i) # noqa: E501
497
498
# assert input chunking is correct
499
concat_chunk_f = lambda pt1, pt2, i: torch.cat([
482
- pt1[:,chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1],...],
483
- pt2[:,remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1],...],
500
+ pt1[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1],...],
501
+ pt2[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1],...],
502
],
485
- dim=1)
486
- concat_batch_f = lambda pt1, pt2: torch.cat([concat_chunk_f(pt1, pt2, i) for i in range(num_sequences)], dim=1) # noqa: E501
503
+ dim=0)
504
+ concat_batch_f = lambda pt1, pt2: torch.cat([concat_chunk_f(pt1, pt2, i) for i in range(num_sequences)], dim=0) # noqa: E501
505
# fmt: on
506
507
assert concat_batch_f(X_chunked, remaining_X_chunked).equal(X)
@@ -498,7 +516,7 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
516
remaining_chunked_cu_seqlens[-1])
517
518
Y_chunked = torch.empty_like(remaining_X_chunked)
501
- state_chunked = mamba_chunk_scan_combined(
519
+ state_chunked = mamba_chunk_scan_combined_varlen(
520
remaining_X_chunked,
521
remaining_dt_chunked,
522
A,
@@ -510,7 +528,6 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
528
seq_idx=remaining_chunked_seq_idx,
529
chunk_indices=chunk_indices,
530
chunk_offsets=chunk_offsets,
513
- return_varlen_states=True,
531
initial_states=partial_state,
532
out=Y_chunked,
533
)
@@ -518,17 +535,17 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
535
536
# kernel chunked is same as kernel overall
537
for i in range(num_sequences):
521
- Y_seq = Y[:, cu_seqlens[i]:cu_seqlens[i + 1], ...]
522
- Y_ref_seq = Y_ref[:, cu_seqlens[i]:cu_seqlens[i + 1], ...]
538
+ Y_seq = Y[cu_seqlens[i]:cu_seqlens[i + 1], ...]
539
+ Y_ref_seq = Y_ref[cu_seqlens[i]:cu_seqlens[i + 1], ...]
540
torch.testing.assert_close(
524
- Y_seq[:, :chunked_seqlens[i], ...],
525
- Y_ref_seq[:, :chunked_seqlens[i], ...],
541
+ Y_seq[:chunked_seqlens[i], ...],
542
+ Y_ref_seq[:chunked_seqlens[i], ...],
543
atol=atol,
544
rtol=rtol,
545
msg=lambda x: f"seq{i} output part1 " + x) # noqa: B023
546
torch.testing.assert_close(
530
- Y_seq[:, chunked_seqlens[i]:, ...],
531
- Y_ref_seq[:, chunked_seqlens[i]:, ...],
547
+ Y_seq[chunked_seqlens[i]:, ...],
548
+ Y_ref_seq[chunked_seqlens[i]:, ...],
549
atol=atol,
550
rtol=rtol,
551
msg=lambda x: f"seq{i} output part2 " + x) # noqa: B023
vllm/model_executor/layers/mamba/mamba_mixer2.py
+10
-13
@@ -29,7 +29,7 @@ from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated
29
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
30
selective_state_update)
31
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
32
- mamba_chunk_scan_combined)
32
+ mamba_chunk_scan_combined_varlen)
33
from vllm.model_executor.layers.quantization import QuantizationConfig
34
from vllm.model_executor.model_loader.weight_utils import (
35
LoaderFunction, composed_weight_loader, sharded_weight_loader)
@@ -504,6 +504,7 @@ class MambaMixer2(MambaBase, CustomOp):
504
seq_idx_p = attn_metadata.seq_idx_p
505
chunk_indices_p = attn_metadata.chunk_indices_p
506
chunk_offsets_p = attn_metadata.chunk_offsets_p
507
+ query_start_loc_p = attn_metadata.query_start_loc_p
508
509
# 1. Gated MLP's linear projection
510
projected_states, _ = self.in_proj(hidden_states)
@@ -545,6 +546,7 @@ class MambaMixer2(MambaBase, CustomOp):
546
out, _ = self.out_proj(hidden_states)
547
return out
548
549
+ # NOTE: V0 put prefill before decode, v1 puts decode before prefill
550
num_prefills = attn_metadata.num_prefills # request count
551
num_decodes = attn_metadata.num_decode_tokens # token count (=request)
552
num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
@@ -570,9 +572,6 @@ class MambaMixer2(MambaBase, CustomOp):
572
[num_decodes, num_prefills],
573
dim=0,
574
)
573
- query_start_loc_p = (
574
- attn_metadata.query_start_loc[-num_prefills - 1:] -
575
- num_decodes if has_prefill else None)
575
576
# Preallocate output tensor to avoid memcpy cost for merging prefill
577
# and decode outputs
@@ -620,15 +619,15 @@ class MambaMixer2(MambaBase, CustomOp):
619
ssm_state[state_indices_tensor_p], 0)
620
621
# NOTE: final output is an in-place update of out tensor
623
- varlen_state = mamba_chunk_scan_combined(
624
- hidden_states_p.view(1, num_prefill_tokens,
622
+ varlen_states = mamba_chunk_scan_combined_varlen(
623
+ hidden_states_p.view(num_prefill_tokens,
624
self.num_heads // self.tp_size,
625
self.head_dim),
627
- dt_p.unsqueeze(0),
626
+ dt_p,
627
self.A,
629
- B_p.view(1, num_prefill_tokens, self.n_groups // self.tp_size,
628
+ B_p.view(num_prefill_tokens, self.n_groups // self.tp_size,
629
-1),
631
- C_p.view(1, num_prefill_tokens, self.n_groups // self.tp_size,
630
+ C_p.view(num_prefill_tokens, self.n_groups // self.tp_size,
631
-1),
632
chunk_size=chunk_size,
633
D=self.D,
@@ -639,17 +638,15 @@ class MambaMixer2(MambaBase, CustomOp):
638
chunk_offsets=chunk_offsets_p,
639
cu_seqlens=query_start_loc_p,
640
initial_states=initial_states,
642
- return_varlen_states=True,
643
- return_final_states=False,
641
dt_softplus=True,
642
dt_limit=(0.0, float("inf")),
646
- out=preallocated_ssm_out_p.view(1, num_prefill_tokens, -1,
643
+ out=preallocated_ssm_out_p.view(num_prefill_tokens, -1,
644
self.head_dim),
645
state_dtype=ssm_state.dtype)
646
647
# update ssm states
648
# - varlen state is a (num_prefills, nheads, headdim, dstate) tensor
652
- ssm_state[state_indices_tensor_p] = varlen_state
649
+ ssm_state[state_indices_tensor_p] = varlen_states
650
651
# Process decode requests
652
if has_decode:
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+1
-1
@@ -427,7 +427,7 @@ def causal_conv1d_fn(
427
batch_ptr = metadata.batch_ptr
428
token_chunk_offset_ptr = metadata.token_chunk_offset_ptr
429
else:
430
- seqlens = np.diff(query_start_loc.to('cpu'))
430
+ seqlens = query_start_loc.diff().to('cpu')
431
args = seqlens
432
MAX_NUM_PROGRAMS = 1024
433
vllm/model_executor/layers/mamba/ops/ssd_bmm.py
+69
-89
@@ -99,34 +99,28 @@ def _bmm_chunk_fwd_kernel(
99
seq_idx_ptr,
100
# Matrix dimensions
101
seqlen,
102
- chunk_size,
103
- K,
104
- ngroups,
105
- stride_a_batch,
106
- stride_a_seqlen,
107
- stride_a_head,
108
- stride_ak,
109
- stride_b_batch,
110
- stride_b_seqlen,
111
- stride_b_head,
112
- stride_bk,
113
- stride_out_batch,
114
- stride_out_chunk,
115
- stride_out_head,
116
- stride_outm,
117
- stride_outn,
118
- stride_seq_idx_batch,
119
- stride_seq_idx_seqlen,
102
+ chunk_size: tl.constexpr,
103
+ K: tl.constexpr,
104
+ ngroups: tl.constexpr,
105
+ stride_a_seqlen: tl.int64,
106
+ stride_a_head: tl.int64,
107
+ stride_ak: tl.constexpr,
108
+ stride_b_seqlen: tl.int64,
109
+ stride_b_head: tl.int64,
110
+ stride_bk: tl.constexpr,
111
+ stride_out_chunk: tl.int64,
112
+ stride_out_head: tl.int64,
113
+ stride_outm: tl.int64,
114
+ stride_outn: tl.constexpr,
115
+ stride_seq_idx_seqlen: tl.constexpr,
116
# Meta-parameters
117
IS_CAUSAL: tl.constexpr,
118
dot_dtype: tl.constexpr,
123
- HAS_SEQ_IDX: tl.constexpr,
119
BLOCK_SIZE_M: tl.constexpr,
120
BLOCK_SIZE_N: tl.constexpr,
121
BLOCK_SIZE_K: tl.constexpr,
122
):
128
- pid_b = tl.program_id(axis=1)
129
- pid_ch = tl.program_id(axis=2).to(tl.int64)
123
+ pid_ch = tl.program_id(axis=1).to(tl.int64)
124
pid_c = pid_ch // ngroups
125
pid_h = pid_ch - pid_c * ngroups
126
num_pid_n = tl.cdiv(chunk_size, BLOCK_SIZE_N)
@@ -135,10 +129,10 @@ def _bmm_chunk_fwd_kernel(
129
if IS_CAUSAL:
130
if pid_n * BLOCK_SIZE_N >= (pid_m + 1) * BLOCK_SIZE_M:
131
return
138
- a_ptr += pid_b * stride_a_batch + pid_c * chunk_size * stride_a_seqlen + pid_h * stride_a_head
139
- b_ptr += pid_b * stride_b_batch + pid_c * chunk_size * stride_b_seqlen + pid_h * stride_b_head
140
- if HAS_SEQ_IDX:
141
- seq_idx_ptr += pid_b * stride_seq_idx_batch + pid_c * chunk_size * stride_seq_idx_seqlen
132
+ a_ptr += pid_c * chunk_size * stride_a_seqlen + pid_h * stride_a_head
133
+ b_ptr += pid_c * chunk_size * stride_b_seqlen + pid_h * stride_b_head
134
+
135
+ seq_idx_ptr += pid_c * chunk_size * stride_seq_idx_seqlen
136
137
offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
138
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
@@ -150,6 +144,8 @@ def _bmm_chunk_fwd_kernel(
144
chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size)
145
146
acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
147
+
148
+ # compute a * b.T
149
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
150
a = tl.load(a_ptrs,
151
mask=(offs_m[:, None] < chunk_size_limit) &
@@ -165,18 +161,19 @@ def _bmm_chunk_fwd_kernel(
161
162
offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
163
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
168
- if HAS_SEQ_IDX:
169
- chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size)
170
- seq_idx_m = tl.load(seq_idx_ptr + offs_m * stride_seq_idx_seqlen,
171
- mask=offs_m < chunk_size_limit,
172
- other=-1)
173
- seq_idx_n = tl.load(seq_idx_ptr + offs_n * stride_seq_idx_seqlen,
174
- mask=offs_n < chunk_size_limit,
175
- other=-2)
176
- acc = tl.where(seq_idx_m[:, None] == seq_idx_n[None, :], acc, 0.0)
177
- out = acc.to(out_ptr.dtype.element_ty)
164
179
- out_ptr += pid_b * stride_out_batch + pid_c * stride_out_chunk + pid_h * stride_out_head
165
+ # Zero out the results that are not from the same request
166
+ # in the varlen batch
167
+ seq_idx_m = tl.load(seq_idx_ptr + offs_m * stride_seq_idx_seqlen,
168
+ mask=offs_m < chunk_size_limit,
169
+ other=-1)
170
+ seq_idx_n = tl.load(seq_idx_ptr + offs_n * stride_seq_idx_seqlen,
171
+ mask=offs_n < chunk_size_limit,
172
+ other=-2)
173
+ acc = tl.where(seq_idx_m[:, None] == seq_idx_n[None, :], acc, 0.0)
174
+
175
+ out = acc.to(out_ptr.dtype.element_ty)
176
+ out_ptr += pid_c * stride_out_chunk + pid_h * stride_out_head
177
out_ptrs = out_ptr + (stride_outm * offs_m[:, None] +
178
offs_n[None, :] * stride_outn)
179
tl.store(out_ptrs,
@@ -185,78 +182,61 @@ def _bmm_chunk_fwd_kernel(
182
(offs_n[None, :] < chunk_size))
183
184
188
-def _bmm_chunk_fwd(a,
189
- b,
190
- chunk_size,
191
- seq_idx=None,
192
- causal=False,
193
- output_dtype=None):
185
+def _bmm_chunk_fwd(a, b, chunk_size, seq_idx, causal=False, output_dtype=None):
186
"""
187
Argument:
196
- a: (batch, seqlen, k) or (batch, seqlen, ngroups, k)
197
- b: (batch, seqlen, k) or (batch, seqlen, ngroups, k)
198
- seq_idx: (batch, seqlen) or None. out[i, j] for seq_idx[i] != seq_idx[j] will be zeroed out.
188
+ a: (seqlen, ngroups, k)
189
+ b: (seqlen, ngroups, k)
190
+ seq_idx: (seqlen,). out[i, j] for seq_idx[i] != seq_idx[j] will be zeroed out.
191
causal: if True, then out[i, j] for i > j will be arbitrary, only out[i, j] for i <= j are
192
guaranteed to be correct.
193
Return:
202
- out: (batch, nchunks, chunk_size, chunk_size) or (batch, nchunks, ngroups, chunk_size, chunk_size)
194
+ out: (nchunks, ngroups, chunk_size, chunk_size)
195
"""
204
- # Check constraints.
205
- has_groups = a.dim() == 4
206
- if not has_groups:
207
- batch, seqlen, k = a.shape
208
- else:
209
- batch, seqlen, ngroups, k = a.shape
196
+ seqlen, ngroups, k = a.shape
197
assert b.shape == a.shape
211
- if seq_idx is not None:
212
- assert seq_idx.shape == (batch, seqlen)
213
- if a.stride(-1) != 1 and a.stride(1) != 1:
198
+ assert seq_idx is not None
199
+ assert seq_idx.shape == (seqlen, )
200
+ if a.stride(-1) != 1 and a.stride(0) != 1:
201
a = a.contiguous()
215
- if b.stride(-1) != 1 and b.stride(1) != 1:
202
+ if b.stride(-1) != 1 and b.stride(0) != 1:
203
b = b.contiguous()
204
+
205
nchunks = math.ceil(seqlen / chunk_size)
206
# Allocates output.
207
out_dtype = a.dtype if output_dtype is None else output_dtype
220
- out = torch.empty(
221
- (batch, nchunks, chunk_size, chunk_size) if not has_groups else
222
- (batch, nchunks, ngroups, chunk_size, chunk_size),
223
- device=a.device,
224
- dtype=out_dtype)
208
+ out = torch.empty((nchunks, ngroups, chunk_size, chunk_size),
209
+ device=a.device,
210
+ dtype=out_dtype)
211
dot_dtype = (tl.bfloat16
212
if a.dtype == torch.bfloat16 or b.dtype == torch.bfloat16 else
213
(tl.float16 if a.dtype == torch.float16
214
or b.dtype == torch.float16 else tl.float32))
215
grid = lambda META: (triton.cdiv(
216
chunk_size, META['BLOCK_SIZE_M']) * triton.cdiv(
231
- chunk_size, META['BLOCK_SIZE_N']), batch, nchunks
232
- if not has_groups else nchunks * ngroups)
217
+ chunk_size, META['BLOCK_SIZE_N']), nchunks * ngroups)
218
with torch.cuda.device(a.device.index):
219
_bmm_chunk_fwd_kernel[grid](
235
- a,
236
- b,
237
- out,
238
- seq_idx,
239
- seqlen,
240
- chunk_size,
241
- k,
242
- ngroups if has_groups else 1,
243
- a.stride(0),
244
- a.stride(1),
245
- 0 if not has_groups else a.stride(2),
246
- a.stride(-1),
247
- b.stride(0),
248
- b.stride(1),
249
- 0 if not has_groups else b.stride(2),
250
- b.stride(-1),
251
- out.stride(0),
252
- out.stride(1),
253
- 0 if not has_groups else out.stride(2),
254
- out.stride(-2),
255
- out.stride(-1),
256
- *((seq_idx.stride(0),
257
- seq_idx.stride(1)) if seq_idx is not None else (0, 0)),
258
- causal,
259
- dot_dtype,
260
- HAS_SEQ_IDX=seq_idx is not None,
220
+ a_ptr=a,
221
+ b_ptr=b,
222
+ out_ptr=out,
223
+ seq_idx_ptr=seq_idx,
224
+ seqlen=seqlen,
225
+ chunk_size=chunk_size,
226
+ K=k,
227
+ ngroups=ngroups,
228
+ stride_a_seqlen=a.stride(0),
229
+ stride_a_head=a.stride(1),
230
+ stride_ak=a.stride(2),
231
+ stride_b_seqlen=b.stride(0),
232
+ stride_b_head=b.stride(1),
233
+ stride_bk=b.stride(2),
234
+ stride_out_chunk=out.stride(0),
235
+ stride_out_head=out.stride(1),
236
+ stride_outm=out.stride(-2),
237
+ stride_outn=out.stride(-1),
238
+ stride_seq_idx_seqlen=seq_idx.stride(0),
239
+ IS_CAUSAL=causal,
240
+ dot_dtype=dot_dtype,
241
)
242
return out
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+174
-221
@@ -6,7 +6,6 @@
6
7
# ruff: noqa: E501,SIM102
8
9
-import torch
9
from packaging import version
10
11
from vllm.triton_utils import tl, triton
@@ -114,7 +113,6 @@ def _chunk_scan_fwd_kernel(
113
x_ptr,
114
z_ptr,
115
out_ptr,
117
- out_x_ptr,
116
dt_ptr,
117
dA_cumsum_ptr,
118
seq_idx_ptr,
@@ -126,60 +124,49 @@ def _chunk_scan_fwd_kernel(
124
chunk_offsets_ptr,
125
chunk_meta_num,
126
# Matrix dimensions
129
- chunk_size,
130
- hdim,
131
- dstate,
132
- batch,
127
+ chunk_size: tl.constexpr,
128
+ hdim: tl.constexpr,
129
+ dstate: tl.constexpr,
130
seqlen,
134
- nheads_ngroups_ratio,
131
+ nheads_ngroups_ratio: tl.constexpr,
132
# Strides
136
- stride_cb_batch,
137
- stride_cb_chunk,
138
- stride_cb_head,
139
- stride_cb_csize_m,
140
- stride_cb_csize_k,
141
- stride_x_batch,
142
- stride_x_seqlen,
143
- stride_x_head,
144
- stride_x_hdim,
145
- stride_z_batch,
146
- stride_z_seqlen,
147
- stride_z_head,
148
- stride_z_hdim,
149
- stride_out_batch,
150
- stride_out_seqlen,
151
- stride_out_head,
152
- stride_out_hdim,
153
- stride_dt_batch,
154
- stride_dt_chunk,
155
- stride_dt_head,
156
- stride_dt_csize,
157
- stride_dA_cs_batch,
158
- stride_dA_cs_chunk,
159
- stride_dA_cs_head,
160
- stride_dA_cs_csize,
161
- stride_seq_idx_batch,
162
- stride_seq_idx_seqlen,
163
- stride_C_batch,
164
- stride_C_seqlen,
165
- stride_C_head,
166
- stride_C_dstate,
167
- stride_states_batch,
168
- stride_states_chunk,
169
- stride_states_head,
170
- stride_states_hdim,
171
- stride_states_dstate,
172
- stride_init_states_batch,
173
- stride_init_states_head,
174
- stride_init_states_hdim,
175
- stride_init_states_dstate,
176
- stride_D_head,
133
+ stride_cb_chunk: tl.int64,
134
+ stride_cb_head: tl.int64,
135
+ stride_cb_csize_m: tl.int64,
136
+ stride_cb_csize_k: tl.constexpr,
137
+ stride_x_seqlen: tl.int64,
138
+ stride_x_head: tl.int64,
139
+ stride_x_hdim: tl.constexpr,
140
+ stride_z_seqlen: tl.int64,
141
+ stride_z_head: tl.int64,
142
+ stride_z_hdim: tl.constexpr,
143
+ stride_out_seqlen: tl.int64,
144
+ stride_out_head: tl.int64,
145
+ stride_out_hdim: tl.constexpr,
146
+ stride_dt_chunk: tl.int64,
147
+ stride_dt_head: tl.int64,
148
+ stride_dt_csize: tl.constexpr,
149
+ stride_dA_cs_chunk: tl.int64,
150
+ stride_dA_cs_head: tl.int64,
151
+ stride_dA_cs_csize: tl.constexpr,
152
+ stride_seq_idx_seqlen: tl.constexpr,
153
+ stride_C_seqlen: tl.int64,
154
+ stride_C_head: tl.int64,
155
+ stride_C_dstate: tl.constexpr,
156
+ stride_states_chunk: tl.int64,
157
+ stride_states_head: tl.int64,
158
+ stride_states_hdim: tl.int64,
159
+ stride_states_dstate: tl.constexpr,
160
+ stride_init_states_batch: tl.int64,
161
+ stride_init_states_head: tl.int64,
162
+ stride_init_states_hdim: tl.int64,
163
+ stride_init_states_dstate: tl.constexpr,
164
+ stride_D_head: tl.constexpr,
165
# Meta-parameters
166
IS_CAUSAL: tl.constexpr,
167
HAS_D: tl.constexpr,
168
D_HAS_HDIM: tl.constexpr,
169
HAS_Z: tl.constexpr,
182
- HAS_SEQ_IDX: tl.constexpr,
170
BLOCK_SIZE_M: tl.constexpr,
171
BLOCK_SIZE_N: tl.constexpr,
172
BLOCK_SIZE_K: tl.constexpr,
@@ -187,9 +174,7 @@ def _chunk_scan_fwd_kernel(
174
IS_TRITON_22: tl.constexpr,
175
HAS_INITSTATES: tl.constexpr,
176
):
190
- pid_bc = tl.program_id(axis=1).to(tl.int64)
191
- pid_c = pid_bc // batch
192
- pid_b = pid_bc - pid_c * batch
177
+ pid_c = tl.program_id(axis=1).to(tl.int64)
178
if not HAS_INITSTATES:
179
c_idx = pid_c
180
c_off = 0
@@ -201,53 +186,51 @@ def _chunk_scan_fwd_kernel(
186
num_pid_n = tl.cdiv(hdim, BLOCK_SIZE_N)
187
pid_m = tl.program_id(axis=0) // num_pid_n
188
pid_n = tl.program_id(axis=0) % num_pid_n
204
- cb_ptr += pid_b * stride_cb_batch + c_idx * stride_cb_chunk + (
205
- pid_h // nheads_ngroups_ratio) * stride_cb_head
206
- x_ptr += pid_b * stride_x_batch + c_idx * chunk_size * stride_x_seqlen + pid_h * stride_x_head
207
- dt_ptr += pid_b * stride_dt_batch + c_idx * stride_dt_chunk + pid_h * stride_dt_head
208
- dA_cumsum_ptr += pid_b * stride_dA_cs_batch + c_idx * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
209
- C_ptr += pid_b * stride_C_batch + c_idx * chunk_size * stride_C_seqlen + (
189
+ cb_ptr += c_idx * stride_cb_chunk + (pid_h //
190
+ nheads_ngroups_ratio) * stride_cb_head
191
+ x_ptr += c_idx * chunk_size * stride_x_seqlen + pid_h * stride_x_head
192
+ dt_ptr += c_idx * stride_dt_chunk + pid_h * stride_dt_head
193
+ dA_cumsum_ptr += c_idx * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
194
+ C_ptr += c_idx * chunk_size * stride_C_seqlen + (
195
pid_h // nheads_ngroups_ratio) * stride_C_head
196
197
# M-block offsets and prev states
198
# - logic in next block may override these if there is an active offset
199
offs_m = pid_m * BLOCK_SIZE_M + c_off + tl.arange(0, BLOCK_SIZE_M)
215
- prev_states_ptr = states_ptr + pid_b * stride_states_batch + c_idx * stride_states_chunk + pid_h * stride_states_head
200
+ prev_states_ptr = states_ptr + c_idx * stride_states_chunk + pid_h * stride_states_head
201
prev_states_hdim = stride_states_hdim
202
prev_states_dstate = stride_states_dstate
203
204
chunk_size_limit = min(chunk_size, seqlen - c_idx * chunk_size)
220
- if HAS_SEQ_IDX:
221
- seq_idx_ptr += pid_b * stride_seq_idx_batch + c_idx * chunk_size * stride_seq_idx_seqlen
222
-
223
- # - we only need seq_idx_prev to be aligned to chunk boundary
224
- seq_idx_prev = tl.load(seq_idx_ptr - stride_seq_idx_seqlen,
225
- mask=c_idx >= 1,
226
- other=0)
227
-
228
- if HAS_INITSTATES:
229
- # if there are init states, we only need seq_idx_m to point
230
- # what is the current seq_idx
231
-
232
- # get current seq idx
233
- if (pid_m * BLOCK_SIZE_M + c_off) < chunk_size_limit:
234
- seq_idx_m = tl.load(
235
- seq_idx_ptr +
236
- (pid_m * BLOCK_SIZE_M + c_off) * stride_seq_idx_seqlen, )
237
-
238
- # - recall that in ssd_state_passing, for the case c_off == 0
239
- # i.e., the very first sequence, we made states_ptr hold its initial state
240
- # so this edge case is taken care of
241
- if ((c_off == 0) and
242
- (seq_idx_prev != seq_idx_m
243
- ) # if a seq is changed exactly on boundary
244
- or (c_off > 0) # implies a new example (pseudo chunk)
245
- ):
246
-
247
- # - replace prev_states_ptr with init_states
248
- prev_states_ptr = initstates_ptr + seq_idx_m * stride_init_states_batch + pid_h * stride_init_states_head
249
- prev_states_hdim = stride_init_states_hdim # override strides
250
- prev_states_dstate = stride_init_states_dstate
205
+
206
+ seq_idx_ptr += c_idx * chunk_size * stride_seq_idx_seqlen
207
+ # - we only need seq_idx_prev to be aligned to chunk boundary
208
+ seq_idx_prev = tl.load(seq_idx_ptr - stride_seq_idx_seqlen,
209
+ mask=c_idx >= 1,
210
+ other=0)
211
+
212
+ if HAS_INITSTATES:
213
+ # if there are init states, we only need seq_idx_m to point
214
+ # what is the current seq_idx
215
+
216
+ # get current seq idx
217
+ if (pid_m * BLOCK_SIZE_M + c_off) < chunk_size_limit:
218
+ seq_idx_m = tl.load(
219
+ seq_idx_ptr +
220
+ (pid_m * BLOCK_SIZE_M + c_off) * stride_seq_idx_seqlen, )
221
+
222
+ # - recall that in ssd_state_passing, for the case c_off == 0
223
+ # i.e., the very first sequence, we made states_ptr hold its initial state
224
+ # so this edge case is taken care of
225
+ if ((c_off == 0) and (seq_idx_prev != seq_idx_m
226
+ ) # if a seq is changed exactly on boundary
227
+ or (c_off > 0) # implies a new example (pseudo chunk)
228
+ ):
229
+
230
+ # - replace prev_states_ptr with init_states
231
+ prev_states_ptr = initstates_ptr + seq_idx_m * stride_init_states_batch + pid_h * stride_init_states_head
232
+ prev_states_hdim = stride_init_states_hdim # override strides
233
+ prev_states_dstate = stride_init_states_dstate
234
235
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
236
dA_cs_m = tl.load(dA_cumsum_ptr + offs_m * stride_dA_cs_csize,
@@ -256,7 +239,6 @@ def _chunk_scan_fwd_kernel(
239
240
# - handle chunk state limit
241
if HAS_INITSTATES:
259
-
242
# have to split this if otherwise compilation will have problems
243
dA_cs_m_boundary = 0.0
244
@@ -296,13 +278,11 @@ def _chunk_scan_fwd_kernel(
278
dA_cumsum_ptr + (c_off - 1) * stride_dA_cs_csize,
279
mask=(((c_off - 1) > -1) and ((c_off) < chunk_size)),
280
other=0.0).to(tl.float32)
299
-
300
- if HAS_SEQ_IDX:
281
+ else:
282
# - handle seq idx when HAS_INITSTATES==False
302
- if not HAS_INITSTATES:
303
- seq_idx_m = tl.load(seq_idx_ptr + offs_m * stride_seq_idx_seqlen,
304
- mask=offs_m < chunk_size_limit,
305
- other=-1)
283
+ seq_idx_m = tl.load(seq_idx_ptr + offs_m * stride_seq_idx_seqlen,
284
+ mask=offs_m < chunk_size_limit,
285
+ other=-1)
286
287
acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
288
@@ -319,18 +299,15 @@ def _chunk_scan_fwd_kernel(
299
prev_states_ptrs = prev_states_ptr + (
300
offs_n[None, :] * prev_states_hdim +
301
offs_k_dstate[:, None] * prev_states_dstate)
322
- if HAS_SEQ_IDX:
323
-
324
- if not HAS_INITSTATES:
325
- # - this is for continuous batching where there is no init states
326
- scale_m = tl.where(seq_idx_m == seq_idx_prev, tl.exp(dA_cs_m),
327
- 0.0)
328
- else:
329
- # - if there is initstates, we will rely on prev_states, no zeroing
330
- # required.
331
- scale_m = tl.exp(dA_cs_m - dA_cs_m_boundary)
302
+
303
+ if not HAS_INITSTATES:
304
+ # - this is for continuous batching where there is no init states
305
+ scale_m = tl.where(seq_idx_m == seq_idx_prev, tl.exp(dA_cs_m), 0.0)
306
else:
333
- scale_m = tl.exp(dA_cs_m)
307
+ # - if there is initstates, we will rely on prev_states, no zeroing
308
+ # required.
309
+ scale_m = tl.exp(dA_cs_m - dA_cs_m_boundary)
310
+
311
if BLOCK_SIZE_DSTATE <= 128:
312
C = tl.load(C_ptrs,
313
mask=(offs_m[:, None] < chunk_size_limit) &
@@ -416,15 +393,7 @@ def _chunk_scan_fwd_kernel(
393
acc += x_residual * D
394
395
if HAS_Z:
419
- out_x_ptr += pid_b * stride_out_batch + c_idx * chunk_size * stride_out_seqlen + pid_h * stride_out_head
420
- out_x_ptrs = out_x_ptr + (stride_out_seqlen * offs_out_m[:, None] +
421
- offs_out_n[None, :])
422
- tl.store(out_x_ptrs,
423
- acc,
424
- mask=(offs_out_m[:, None] < chunk_size_limit) &
425
- (offs_out_n[None, :] < hdim))
426
-
427
- z_ptr += pid_b * stride_z_batch + c_idx * chunk_size * stride_z_seqlen + pid_h * stride_z_head
396
+ z_ptr += c_idx * chunk_size * stride_z_seqlen + pid_h * stride_z_head
397
z_ptrs = z_ptr + (stride_z_seqlen * offs_out_m[:, None] +
398
stride_z_hdim * offs_out_n[None, :])
399
z = tl.load(z_ptrs,
@@ -433,7 +402,7 @@ def _chunk_scan_fwd_kernel(
402
other=0.0).to(tl.float32)
403
acc *= z * tl.sigmoid(z)
404
436
- out_ptr += pid_b * stride_out_batch + c_idx * chunk_size * stride_out_seqlen + pid_h * stride_out_head
405
+ out_ptr += c_idx * chunk_size * stride_out_seqlen + pid_h * stride_out_head
406
out_ptrs = out_ptr + (stride_out_seqlen * offs_out_m[:, None] +
407
offs_out_n[None, :] * stride_out_hdim)
408
tl.store(out_ptrs,
@@ -449,126 +418,110 @@ def _chunk_scan_fwd(
418
dA_cumsum,
419
C,
420
states,
421
+ out,
422
+ seq_idx,
423
D=None,
424
z=None,
454
- seq_idx=None,
425
chunk_indices=None,
426
chunk_offsets=None,
427
initial_states=None,
458
- out=None,
428
):
460
- batch, seqlen, nheads, headdim = x.shape
461
- _, _, nchunks, chunk_size = dt.shape
462
- _, _, ngroups, dstate = C.shape
429
+ assert seq_idx is not None, "this implementation requires seq_idx"
430
+
431
+ seqlen, nheads, headdim = x.shape
432
+ _, nchunks, chunk_size = dt.shape
433
+ _, ngroups, dstate = C.shape
434
assert nheads % ngroups == 0
464
- assert C.shape == (batch, seqlen, ngroups, dstate)
465
- assert cb.shape == (batch, nchunks, ngroups, chunk_size, chunk_size)
466
- if z is not None:
467
- assert z.shape == x.shape
435
+ assert C.shape == (seqlen, ngroups, dstate)
436
+ assert cb.shape == (nchunks, ngroups, chunk_size, chunk_size)
437
if D is not None:
438
assert D.shape == (nheads, headdim) or D.shape == (nheads, )
470
- assert dt.shape == (batch, nheads, nchunks, chunk_size)
471
- assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
472
- assert states.shape == (batch, nchunks, nheads, headdim, dstate)
473
-
474
- if seq_idx is not None:
475
- assert seq_idx.shape == (batch, seqlen)
476
-
477
- if initial_states is not None:
478
- # with initial states, we need to take care of how
479
- # seq_idx crosses the boundaries
480
- assert batch == 1, "chunk scan only supports initial states with batch 1"
481
- assert chunk_indices is not None and chunk_offsets is not None, \
482
- "chunk_indices and chunk_offsets should have been set"
483
- else:
484
- chunk_indices, chunk_offsets = None, None
485
- else:
486
- chunk_indices, chunk_offsets = None, None
487
-
488
- assert out.shape == x.shape
489
-
439
if z is not None:
491
- out_x = torch.empty_like(x)
492
- assert out_x.stride() == out.stride()
440
+ assert z.shape == x.shape
441
+ assert dt.shape == (nheads, nchunks, chunk_size)
442
+ assert dA_cumsum.shape == (nheads, nchunks, chunk_size)
443
+ assert states.shape == (nchunks, nheads, headdim, dstate)
444
+ assert seq_idx.shape == (seqlen, )
445
+
446
+ if initial_states is not None:
447
+ # with initial states, we need to take care of how
448
+ # seq_idx crosses the boundaries
449
+ assert chunk_indices is not None and chunk_offsets is not None, \
450
+ "chunk_indices and chunk_offsets should have been set"
451
else:
494
- out_x = None
452
+ chunk_indices, chunk_offsets = None, None
453
454
grid = lambda META: (
455
triton.cdiv(chunk_size, META['BLOCK_SIZE_M']) * triton.cdiv(
498
- headdim, META['BLOCK_SIZE_N']), batch * nchunks
456
+ headdim, META['BLOCK_SIZE_N']), nchunks
457
if chunk_offsets is None else len(chunk_offsets), nheads)
500
- z_strides = ((z.stride(0), z.stride(1), z.stride(2),
501
- z.stride(3)) if z is not None else (0, 0, 0, 0))
458
+
459
+ z_strides = ((z.stride(0), z.stride(1), z.stride(2)) if z is not None else
460
+ (0, 0, 0))
461
+ initial_states_strides = ((initial_states.stride(0),
462
+ initial_states.stride(1),
463
+ initial_states.stride(2),
464
+ initial_states.stride(3))
465
+ if initial_states is not None else (0, 0, 0, 0))
466
+
467
_chunk_scan_fwd_kernel[grid](
503
- cb,
504
- x,
505
- z,
506
- out,
507
- out_x,
508
- dt,
509
- dA_cumsum,
510
- seq_idx,
511
- C,
512
- states,
513
- D,
514
- initial_states,
515
- chunk_indices,
516
- chunk_offsets,
517
- len(chunk_indices) if chunk_indices is not None else 0,
518
- chunk_size,
519
- headdim,
520
- dstate,
521
- batch,
522
- seqlen,
523
- nheads // ngroups,
524
- cb.stride(0),
525
- cb.stride(1),
526
- cb.stride(2),
527
- cb.stride(3),
528
- cb.stride(4),
529
- x.stride(0),
530
- x.stride(1),
531
- x.stride(2),
532
- x.stride(3),
533
- z_strides[0],
534
- z_strides[1],
535
- z_strides[2],
536
- z_strides[3],
537
- out.stride(0),
538
- out.stride(1),
539
- out.stride(2),
540
- out.stride(3),
541
- dt.stride(0),
542
- dt.stride(2),
543
- dt.stride(1),
544
- dt.stride(3),
545
- dA_cumsum.stride(0),
546
- dA_cumsum.stride(2),
547
- dA_cumsum.stride(1),
548
- dA_cumsum.stride(3),
549
- *((seq_idx.stride(0), seq_idx.stride(1)) if seq_idx is not None else
550
- (0, 0)),
551
- C.stride(0),
552
- C.stride(1),
553
- C.stride(2),
554
- C.stride(3),
555
- states.stride(0),
556
- states.stride(1),
557
- states.stride(2),
558
- states.stride(3),
559
- states.stride(4),
560
- *((initial_states.stride(0), initial_states.stride(1),
561
- initial_states.stride(2),
562
- initial_states.stride(3)) if initial_states is not None else
563
- (0, 0, 0, 0)),
564
- D.stride(0) if D is not None else 0,
565
- True,
566
- D is not None,
567
- D.dim() == 2 if D is not None else True,
568
- BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
468
+ cb_ptr=cb,
469
+ x_ptr=x,
470
+ z_ptr=z,
471
+ out_ptr=out,
472
+ dt_ptr=dt,
473
+ dA_cumsum_ptr=dA_cumsum,
474
+ seq_idx_ptr=seq_idx,
475
+ C_ptr=C,
476
+ states_ptr=states,
477
+ D_ptr=D,
478
+ initstates_ptr=initial_states,
479
+ chunk_indices_ptr=chunk_indices,
480
+ chunk_offsets_ptr=chunk_offsets,
481
+ chunk_meta_num=len(chunk_indices) if chunk_indices is not None else 0,
482
+ chunk_size=chunk_size,
483
+ hdim=headdim,
484
+ dstate=dstate,
485
+ seqlen=seqlen,
486
+ nheads_ngroups_ratio=nheads // ngroups,
487
+ stride_cb_chunk=cb.stride(0),
488
+ stride_cb_head=cb.stride(1),
489
+ stride_cb_csize_m=cb.stride(2),
490
+ stride_cb_csize_k=cb.stride(3),
491
+ stride_x_seqlen=x.stride(0),
492
+ stride_x_head=x.stride(1),
493
+ stride_x_hdim=x.stride(2),
494
+ stride_z_seqlen=z_strides[0],
495
+ stride_z_head=z_strides[1],
496
+ stride_z_hdim=z_strides[2],
497
+ stride_out_seqlen=out.stride(0),
498
+ stride_out_head=out.stride(1),
499
+ stride_out_hdim=out.stride(2),
500
+ stride_dt_chunk=dt.stride(1),
501
+ stride_dt_head=dt.stride(0),
502
+ stride_dt_csize=dt.stride(2),
503
+ stride_dA_cs_chunk=dA_cumsum.stride(1),
504
+ stride_dA_cs_head=dA_cumsum.stride(0),
505
+ stride_dA_cs_csize=dA_cumsum.stride(2),
506
+ stride_seq_idx_seqlen=seq_idx.stride(0),
507
+ stride_C_seqlen=C.stride(0),
508
+ stride_C_head=C.stride(1),
509
+ stride_C_dstate=C.stride(2),
510
+ stride_states_chunk=states.stride(0),
511
+ stride_states_head=states.stride(1),
512
+ stride_states_hdim=states.stride(2),
513
+ stride_states_dstate=states.stride(3),
514
+ stride_init_states_batch=initial_states_strides[0],
515
+ stride_init_states_head=initial_states_strides[1],
516
+ stride_init_states_hdim=initial_states_strides[2],
517
+ stride_init_states_dstate=initial_states_strides[3],
518
+ stride_D_head=D.stride(0) if D is not None else 0,
519
+ IS_CAUSAL=True,
520
+ HAS_D=D is not None,
521
+ D_HAS_HDIM=D.dim() == 2 if D is not None else True,
522
HAS_Z=z is not None,
570
- HAS_SEQ_IDX=seq_idx is not None,
523
+ BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
524
IS_TRITON_22=TRITON_22,
525
HAS_INITSTATES=initial_states is not None,
526
)
574
- return out_x
527
+ return
vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py
+202
-228
@@ -35,41 +35,35 @@ def _chunk_cumsum_fwd_kernel(
35
dt_out_ptr,
36
dA_cumsum_ptr,
37
# Matrix dimension
38
- batch,
38
seqlen,
40
- nheads,
41
- chunk_size,
42
- dt_min,
43
- dt_max,
39
+ nheads: tl.constexpr,
40
+ chunk_size: tl.constexpr,
41
+ dt_min: tl.constexpr,
42
+ dt_max: tl.constexpr,
43
# Strides
45
- stride_dt_batch,
46
- stride_dt_seqlen,
47
- stride_dt_head,
48
- stride_A_head,
49
- stride_dt_bias_head,
50
- stride_dt_out_batch,
51
- stride_dt_out_chunk,
52
- stride_dt_out_head,
53
- stride_dt_out_csize,
54
- stride_dA_cs_batch,
55
- stride_dA_cs_chunk,
56
- stride_dA_cs_head,
57
- stride_dA_cs_csize,
44
+ stride_dt_seqlen: tl.int64,
45
+ stride_dt_head: tl.constexpr,
46
+ stride_A_head: tl.constexpr,
47
+ stride_dt_bias_head: tl.constexpr,
48
+ stride_dt_out_head: tl.int64,
49
+ stride_dt_out_chunk: tl.int64,
50
+ stride_dt_out_csize: tl.constexpr,
51
+ stride_dA_cs_head: tl.int64,
52
+ stride_dA_cs_chunk: tl.int64,
53
+ stride_dA_cs_csize: tl.constexpr,
54
# Meta-parameters
55
DT_SOFTPLUS: tl.constexpr,
56
HAS_DT_BIAS: tl.constexpr,
57
BLOCK_SIZE_H: tl.constexpr,
58
BLOCK_SIZE_CHUNK: tl.constexpr,
59
):
64
- pid_b = tl.program_id(axis=0)
65
-
60
# if dt is long, may cause problems, so use 64 bit
61
# https://github.com/triton-lang/triton/issues/1058
68
- pid_c = tl.program_id(axis=1).to(tl.int64)
69
- pid_h = tl.program_id(axis=2)
70
- dt_ptr += pid_b * stride_dt_batch + pid_c * chunk_size * stride_dt_seqlen
71
- dt_out_ptr += pid_b * stride_dt_out_batch + pid_c * stride_dt_out_chunk
72
- dA_cumsum_ptr += pid_b * stride_dA_cs_batch + pid_c * stride_dA_cs_chunk
62
+ pid_c = tl.program_id(axis=0).to(tl.int64)
63
+ pid_h = tl.program_id(axis=1)
64
+ dt_ptr += pid_c * chunk_size * stride_dt_seqlen
65
+ dt_out_ptr += pid_c * stride_dt_out_chunk
66
+ dA_cumsum_ptr += pid_c * stride_dA_cs_chunk
67
68
offs_h = pid_h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H)
69
offs_c = tl.arange(0, BLOCK_SIZE_CHUNK)
@@ -93,9 +87,8 @@ def _chunk_cumsum_fwd_kernel(
87
dt += dt_bias[:, None]
88
if DT_SOFTPLUS:
89
dt = tl.where(dt <= 20.0, softplus(dt), dt)
96
- # As of Triton 2.2.0, tl.clamp is not available yet
97
- # dt = tl.clamp(dt, dt_min, dt_max)
98
- dt = tl.minimum(tl.maximum(dt, dt_min), dt_max)
90
+
91
+ dt = tl.clamp(dt, dt_min, dt_max)
92
dt = tl.where(
93
(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size_limit), dt,
94
0.0)
@@ -197,56 +190,46 @@ def _chunk_state_fwd_kernel(
190
dA_cumsum_ptr,
191
seq_idx_ptr,
192
# Matrix dimensions
200
- hdim,
201
- dstate,
202
- chunk_size,
203
- batch,
193
+ hdim: tl.constexpr,
194
+ dstate: tl.constexpr,
195
+ chunk_size: tl.constexpr,
196
seqlen,
205
- nheads_ngroups_ratio,
197
+ nheads_ngroups_ratio: tl.constexpr,
198
# Strides
207
- stride_x_batch,
208
- stride_x_seqlen,
209
- stride_x_head,
210
- stride_x_hdim,
211
- stride_b_batch,
212
- stride_b_seqlen,
213
- stride_b_head,
214
- stride_b_dstate,
215
- stride_states_batch,
216
- stride_states_chunk,
217
- stride_states_head,
218
- stride_states_hdim,
219
- stride_states_dstate,
220
- stride_dt_batch,
221
- stride_dt_chunk,
222
- stride_dt_head,
223
- stride_dt_csize,
224
- stride_dA_cs_batch,
225
- stride_dA_cs_chunk,
226
- stride_dA_cs_head,
227
- stride_dA_cs_csize,
228
- stride_seq_idx_batch,
229
- stride_seq_idx_seqlen,
199
+ stride_x_seqlen: tl.int64,
200
+ stride_x_head: tl.int64,
201
+ stride_x_hdim: tl.constexpr,
202
+ stride_b_seqlen: tl.int64,
203
+ stride_b_head: tl.int64,
204
+ stride_b_dstate: tl.constexpr,
205
+ stride_states_chunk: tl.int64,
206
+ stride_states_head: tl.int64,
207
+ stride_states_hdim: tl.int64,
208
+ stride_states_dstate: tl.constexpr,
209
+ stride_dt_head: tl.int64,
210
+ stride_dt_chunk: tl.int64,
211
+ stride_dt_csize: tl.constexpr,
212
+ stride_dA_cs_head: tl.int64,
213
+ stride_dA_cs_chunk: tl.int64,
214
+ stride_dA_cs_csize: tl.constexpr,
215
+ stride_seq_idx_seqlen: tl.constexpr,
216
# Meta-parameters
231
- HAS_SEQ_IDX: tl.constexpr,
217
BLOCK_SIZE_M: tl.constexpr,
218
BLOCK_SIZE_N: tl.constexpr,
219
BLOCK_SIZE_K: tl.constexpr,
220
):
236
- pid_bc = tl.program_id(axis=1).to(tl.int64)
237
- pid_c = pid_bc // batch
238
- pid_b = pid_bc - pid_c * batch
221
+ pid_c = tl.program_id(axis=1).to(tl.int64)
222
pid_h = tl.program_id(axis=2)
223
num_pid_n = tl.cdiv(dstate, BLOCK_SIZE_N)
224
pid_m = tl.program_id(axis=0) // num_pid_n
225
pid_n = tl.program_id(axis=0) % num_pid_n
243
- b_ptr += pid_b * stride_b_batch + pid_c * chunk_size * stride_b_seqlen + (
226
+ b_ptr += pid_c * chunk_size * stride_b_seqlen + (
227
pid_h // nheads_ngroups_ratio) * stride_b_head
245
- x_ptr += pid_b * stride_x_batch + pid_c * chunk_size * stride_x_seqlen + pid_h * stride_x_head
246
- dt_ptr += pid_b * stride_dt_batch + pid_c * stride_dt_chunk + pid_h * stride_dt_head
247
- dA_cumsum_ptr += pid_b * stride_dA_cs_batch + pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
248
- if HAS_SEQ_IDX:
249
- seq_idx_ptr += pid_b * stride_seq_idx_batch + pid_c * chunk_size * stride_seq_idx_seqlen
228
+ x_ptr += pid_c * chunk_size * stride_x_seqlen + pid_h * stride_x_head
229
+ dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head
230
+ dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
231
+
232
+ seq_idx_ptr += pid_c * chunk_size * stride_seq_idx_seqlen
233
234
offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
235
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
@@ -259,13 +242,11 @@ def _chunk_state_fwd_kernel(
242
dA_cs_last = tl.load(dA_cumsum_ptr +
243
(chunk_size - 1) * stride_dA_cs_csize).to(tl.float32)
244
dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize
262
- if HAS_SEQ_IDX:
263
- seq_idx_ptrs = seq_idx_ptr + offs_k * stride_seq_idx_seqlen
245
246
+ seq_idx_ptrs = seq_idx_ptr + offs_k * stride_seq_idx_seqlen
247
chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size)
266
- if HAS_SEQ_IDX:
267
- seq_idx_last = tl.load(seq_idx_ptr +
268
- (chunk_size_limit - 1) * stride_seq_idx_seqlen)
248
+ seq_idx_last = tl.load(seq_idx_ptr +
249
+ (chunk_size_limit - 1) * stride_seq_idx_seqlen)
250
251
acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
252
for k in range(0, chunk_size_limit, BLOCK_SIZE_K):
@@ -280,29 +261,28 @@ def _chunk_state_fwd_kernel(
261
dA_cs_k = tl.load(dA_cumsum_ptrs,
262
mask=offs_k < chunk_size_limit - k,
263
other=0.0).to(tl.float32)
283
- if HAS_SEQ_IDX:
284
- seq_idx_k = tl.load(seq_idx_ptrs,
285
- mask=offs_k < chunk_size_limit - k,
286
- other=-1)
264
+
265
+ seq_idx_k = tl.load(seq_idx_ptrs,
266
+ mask=offs_k < chunk_size_limit - k,
267
+ other=-1)
268
dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k,
269
other=0.0).to(tl.float32)
289
- if not HAS_SEQ_IDX:
290
- scale = tl.exp(dA_cs_last - dA_cs_k) * dt_k
291
- else:
292
- scale = tl.where(seq_idx_k == seq_idx_last,
293
- tl.exp(dA_cs_last - dA_cs_k) * dt_k, 0.0)
270
+
271
+ scale = tl.where(seq_idx_k == seq_idx_last,
272
+ tl.exp(dA_cs_last - dA_cs_k) * dt_k, 0.0)
273
b *= scale[:, None]
274
b = b.to(x_ptr.dtype.element_ty)
275
acc += tl.dot(x, b)
276
+
277
x_ptrs += BLOCK_SIZE_K * stride_x_seqlen
278
b_ptrs += BLOCK_SIZE_K * stride_b_seqlen
279
dt_ptrs += BLOCK_SIZE_K * stride_dt_csize
280
dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize
301
- if HAS_SEQ_IDX:
302
- seq_idx_ptrs += BLOCK_SIZE_K * stride_seq_idx_seqlen
281
+ seq_idx_ptrs += BLOCK_SIZE_K * stride_seq_idx_seqlen
282
+
283
states = acc.to(states_ptr.dtype.element_ty)
284
305
- states_ptr += pid_b * stride_states_batch + pid_c * stride_states_chunk + pid_h * stride_states_head
285
+ states_ptr += pid_c * stride_states_chunk + pid_h * stride_states_head
286
offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
287
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
288
states_ptrs = states_ptr + (offs_m[:, None] * stride_states_hdim +
@@ -400,36 +380,35 @@ def _chunk_state_varlen_kernel(
380
states_ptr,
381
initstates_ptr,
382
# Matrix dimensions
403
- hdim,
404
- dstate,
405
- chunk_size,
406
- seqlen,
407
- nheads_ngroups_ratio,
383
+ hdim: tl.constexpr,
384
+ dstate: tl.constexpr,
385
+ chunk_size: tl.constexpr,
386
+ nheads_ngroups_ratio: tl.constexpr,
387
# Strides
409
- stride_x_seqlen,
410
- stride_x_head,
411
- stride_x_hdim,
412
- stride_b_seqlen,
413
- stride_b_head,
414
- stride_b_dstate,
415
- stride_dt_chunk,
416
- stride_dt_head,
417
- stride_dt_csize,
418
- stride_dA_cs_chunk,
419
- stride_dA_cs_head,
420
- stride_dA_cs_csize,
421
- stride_chunk_states_chunk,
422
- stride_chunk_states_head,
423
- stride_chunk_states_hdim,
424
- stride_chunk_states_dstate,
425
- stride_states_batch,
426
- stride_states_head,
427
- stride_states_hdim,
428
- stride_states_dstate,
429
- stride_init_states_batch,
430
- stride_init_states_head,
431
- stride_init_states_hdim,
432
- stride_init_states_dstate,
388
+ stride_x_seqlen: tl.int64,
389
+ stride_x_head: tl.int64,
390
+ stride_x_hdim: tl.constexpr,
391
+ stride_b_seqlen: tl.int64,
392
+ stride_b_head: tl.int64,
393
+ stride_b_dstate: tl.constexpr,
394
+ stride_dt_head: tl.int64,
395
+ stride_dt_chunk: tl.int64,
396
+ stride_dt_csize: tl.constexpr,
397
+ stride_dA_cs_head: tl.int64,
398
+ stride_dA_cs_chunk: tl.int64,
399
+ stride_dA_cs_csize: tl.constexpr,
400
+ stride_chunk_states_chunk: tl.int64,
401
+ stride_chunk_states_head: tl.int64,
402
+ stride_chunk_states_hdim: tl.int64,
403
+ stride_chunk_states_dstate: tl.constexpr,
404
+ stride_states_batch: tl.int64,
405
+ stride_states_head: tl.int64,
406
+ stride_states_hdim: tl.int64,
407
+ stride_states_dstate: tl.constexpr,
408
+ stride_init_states_batch: tl.int64,
409
+ stride_init_states_head: tl.int64,
410
+ stride_init_states_hdim: tl.int64,
411
+ stride_init_states_dstate: tl.constexpr,
412
# Meta-parameters
413
BLOCK_SIZE_M: tl.constexpr,
414
BLOCK_SIZE_N: tl.constexpr,
@@ -558,52 +537,47 @@ def _chunk_cumsum_fwd(dt,
537
dt_bias=None,
538
dt_softplus=False,
539
dt_limit=(0.0, float("inf"))):
561
- batch, seqlen, nheads = dt.shape
540
+ seqlen, nheads = dt.shape
541
assert A.shape == (nheads, )
542
if dt_bias is not None:
543
assert dt_bias.shape == (nheads, )
544
nchunks = math.ceil(seqlen / chunk_size)
566
- dt_out = torch.empty(batch,
567
- nheads,
545
+ dt_out = torch.empty(nheads,
546
nchunks,
547
chunk_size,
548
device=dt.device,
549
dtype=torch.float32)
572
- dA_cumsum = torch.empty(batch,
573
- nheads,
550
+ dA_cumsum = torch.empty(nheads,
551
nchunks,
552
chunk_size,
553
device=dt.device,
554
dtype=torch.float32)
578
- grid_chunk_cs = lambda META: (batch, nchunks,
555
+ grid_chunk_cs = lambda META: (nchunks,
556
triton.cdiv(nheads, META['BLOCK_SIZE_H']))
557
with torch.cuda.device(dt.device.index):
558
_chunk_cumsum_fwd_kernel[grid_chunk_cs](
582
- dt,
583
- A,
584
- dt_bias,
585
- dt_out,
586
- dA_cumsum,
587
- batch,
588
- seqlen,
589
- nheads,
590
- chunk_size,
591
- dt_limit[0],
592
- dt_limit[1],
593
- dt.stride(0),
594
- dt.stride(1),
595
- dt.stride(2),
596
- A.stride(0),
597
- dt_bias.stride(0) if dt_bias is not None else 0,
598
- dt_out.stride(0),
599
- dt_out.stride(2),
600
- dt_out.stride(1),
601
- dt_out.stride(3),
602
- dA_cumsum.stride(0),
603
- dA_cumsum.stride(2),
604
- dA_cumsum.stride(1),
605
- dA_cumsum.stride(3),
606
- dt_softplus,
559
+ dt_ptr=dt,
560
+ A_ptr=A,
561
+ dt_bias_ptr=dt_bias,
562
+ dt_out_ptr=dt_out,
563
+ dA_cumsum_ptr=dA_cumsum,
564
+ seqlen=seqlen,
565
+ nheads=nheads,
566
+ chunk_size=chunk_size,
567
+ dt_min=dt_limit[0],
568
+ dt_max=dt_limit[1],
569
+ stride_dt_seqlen=dt.stride(0),
570
+ stride_dt_head=dt.stride(1),
571
+ stride_A_head=A.stride(0),
572
+ stride_dt_bias_head=dt_bias.stride(0)
573
+ if dt_bias is not None else 0,
574
+ stride_dt_out_head=dt_out.stride(0),
575
+ stride_dt_out_chunk=dt_out.stride(1),
576
+ stride_dt_out_csize=dt_out.stride(2),
577
+ stride_dA_cs_head=dA_cumsum.stride(0),
578
+ stride_dA_cs_chunk=dA_cumsum.stride(1),
579
+ stride_dA_cs_csize=dA_cumsum.stride(2),
580
+ DT_SOFTPLUS=dt_softplus,
581
HAS_DT_BIAS=dt_bias is not None,
582
BLOCK_SIZE_CHUNK=triton.next_power_of_2(chunk_size),
583
)
@@ -617,63 +591,57 @@ def _chunk_state_fwd(B,
591
seq_idx=None,
592
states=None,
593
states_in_fp32=True):
620
- batch, seqlen, nheads, headdim = x.shape
621
- _, _, nchunks, chunk_size = dt.shape
622
- _, _, ngroups, dstate = B.shape
594
+ seqlen, nheads, headdim = x.shape
595
+ _, nchunks, chunk_size = dt.shape
596
+ _, ngroups, dstate = B.shape
597
assert nheads % ngroups == 0
624
- assert B.shape == (batch, seqlen, ngroups, dstate)
625
- assert dt.shape == (batch, nheads, nchunks, chunk_size)
598
+ assert B.shape == (seqlen, ngroups, dstate)
599
+ assert dt.shape == (nheads, nchunks, chunk_size)
600
assert dA_cumsum.shape == dt.shape
627
- if seq_idx is not None:
628
- assert seq_idx.shape == (batch, seqlen)
601
+
602
+ assert seq_idx is not None
603
+ assert seq_idx.shape == (seqlen, )
604
+
605
if states is not None:
630
- assert states.shape == (batch, nchunks, nheads, headdim, dstate)
606
+ assert states.shape == (nchunks, nheads, headdim, dstate)
607
else:
608
states_dtype = torch.float32 if states_in_fp32 else B.dtype
633
- states = torch.empty((batch, nchunks, nheads, headdim, dstate),
609
+ states = torch.empty((nchunks, nheads, headdim, dstate),
610
device=x.device,
611
dtype=states_dtype)
636
- grid = lambda META: (
637
- triton.cdiv(headdim, META['BLOCK_SIZE_M']) * triton.cdiv(
638
- dstate, META['BLOCK_SIZE_N']), batch * nchunks, nheads)
612
+
613
+ grid = lambda META: (triton.cdiv(headdim, META['BLOCK_SIZE_M']) * triton.
614
+ cdiv(dstate, META['BLOCK_SIZE_N']), nchunks, nheads)
615
with torch.cuda.device(x.device.index):
616
_chunk_state_fwd_kernel[grid](
641
- x,
642
- B,
643
- states,
644
- dt,
645
- dA_cumsum,
646
- seq_idx,
647
- headdim,
648
- dstate,
649
- chunk_size,
650
- batch,
651
- seqlen,
652
- nheads // ngroups,
653
- x.stride(0),
654
- x.stride(1),
655
- x.stride(2),
656
- x.stride(3),
657
- B.stride(0),
658
- B.stride(1),
659
- B.stride(2),
660
- B.stride(-1),
661
- states.stride(0),
662
- states.stride(1),
663
- states.stride(2),
664
- states.stride(3),
665
- states.stride(4),
666
- dt.stride(0),
667
- dt.stride(2),
668
- dt.stride(1),
669
- dt.stride(3),
670
- dA_cumsum.stride(0),
671
- dA_cumsum.stride(2),
672
- dA_cumsum.stride(1),
673
- dA_cumsum.stride(3),
674
- *((seq_idx.stride(0),
675
- seq_idx.stride(1)) if seq_idx is not None else (0, 0)),
676
- HAS_SEQ_IDX=seq_idx is not None,
617
+ x_ptr=x,
618
+ b_ptr=B,
619
+ states_ptr=states,
620
+ dt_ptr=dt,
621
+ dA_cumsum_ptr=dA_cumsum,
622
+ seq_idx_ptr=seq_idx,
623
+ hdim=headdim,
624
+ dstate=dstate,
625
+ chunk_size=chunk_size,
626
+ seqlen=seqlen,
627
+ nheads_ngroups_ratio=nheads // ngroups,
628
+ stride_x_seqlen=x.stride(0),
629
+ stride_x_head=x.stride(1),
630
+ stride_x_hdim=x.stride(2),
631
+ stride_b_seqlen=B.stride(0),
632
+ stride_b_head=B.stride(1),
633
+ stride_b_dstate=B.stride(2),
634
+ stride_states_chunk=states.stride(0),
635
+ stride_states_head=states.stride(1),
636
+ stride_states_hdim=states.stride(2),
637
+ stride_states_dstate=states.stride(3),
638
+ stride_dt_head=dt.stride(0),
639
+ stride_dt_chunk=dt.stride(1),
640
+ stride_dt_csize=dt.stride(2),
641
+ stride_dA_cs_head=dA_cumsum.stride(0),
642
+ stride_dA_cs_chunk=dA_cumsum.stride(1),
643
+ stride_dA_cs_csize=dA_cumsum.stride(2),
644
+ stride_seq_idx_seqlen=seq_idx.stride(0),
645
)
646
return states
647
@@ -705,46 +673,52 @@ def chunk_state_varlen(B,
673
dstate,
674
dtype=chunk_states.dtype,
675
device=chunk_states.device)
676
+
677
+ initial_states_strides = ((initial_states.stride(0),
678
+ initial_states.stride(1),
679
+ initial_states.stride(2),
680
+ initial_states.stride(3))
681
+ if initial_states is not None else (0, 0, 0, 0))
682
+
683
grid = lambda META: (triton.cdiv(headdim, META['BLOCK_SIZE_M']) * triton.
684
cdiv(dstate, META['BLOCK_SIZE_N']), batch, nheads)
685
with torch.cuda.device(x.device.index):
686
_chunk_state_varlen_kernel[grid](
712
- x,
713
- B,
714
- dt,
715
- dA_cumsum,
716
- chunk_states,
717
- cu_seqlens,
718
- states,
719
- initial_states,
720
- headdim,
721
- dstate,
722
- chunk_size,
723
- total_seqlen,
724
- nheads // ngroups,
725
- x.stride(0),
726
- x.stride(1),
727
- x.stride(2),
728
- B.stride(0),
729
- B.stride(1),
730
- B.stride(2),
731
- dt.stride(1),
732
- dt.stride(0),
733
- dt.stride(2),
734
- dA_cumsum.stride(1),
735
- dA_cumsum.stride(0),
736
- dA_cumsum.stride(2),
737
- chunk_states.stride(0),
738
- chunk_states.stride(1),
739
- chunk_states.stride(2),
740
- chunk_states.stride(3),
741
- states.stride(0),
742
- states.stride(1),
743
- states.stride(2),
744
- states.stride(3),
745
- *((initial_states.stride(0), initial_states.stride(1),
746
- initial_states.stride(2),
747
- initial_states.stride(3)) if initial_states is not None else
748
- (0, 0, 0, 0)),
687
+ x_ptr=x,
688
+ b_ptr=B,
689
+ dt_ptr=dt,
690
+ dA_cumsum_ptr=dA_cumsum,
691
+ chunk_states_ptr=chunk_states,
692
+ cu_seqlens_ptr=cu_seqlens,
693
+ states_ptr=states,
694
+ initstates_ptr=initial_states,
695
+ hdim=headdim,
696
+ dstate=dstate,
697
+ chunk_size=chunk_size,
698
+ nheads_ngroups_ratio=nheads // ngroups,
699
+ stride_x_seqlen=x.stride(0),
700
+ stride_x_head=x.stride(1),
701
+ stride_x_hdim=x.stride(2),
702
+ stride_b_seqlen=B.stride(0),
703
+ stride_b_head=B.stride(1),
704
+ stride_b_dstate=B.stride(2),
705
+ stride_dt_head=dt.stride(0),
706
+ stride_dt_chunk=dt.stride(1),
707
+ stride_dt_csize=dt.stride(2),
708
+ stride_dA_cs_head=dA_cumsum.stride(0),
709
+ stride_dA_cs_chunk=dA_cumsum.stride(1),
710
+ stride_dA_cs_csize=dA_cumsum.stride(2),
711
+ stride_chunk_states_chunk=chunk_states.stride(0),
712
+ stride_chunk_states_head=chunk_states.stride(1),
713
+ stride_chunk_states_hdim=chunk_states.stride(2),
714
+ stride_chunk_states_dstate=chunk_states.stride(3),
715
+ stride_states_batch=states.stride(0),
716
+ stride_states_head=states.stride(1),
717
+ stride_states_hdim=states.stride(2),
718
+ stride_states_dstate=states.stride(3),
719
+ stride_init_states_batch=initial_states_strides[0],
720
+ stride_init_states_head=initial_states_strides[1],
721
+ stride_init_states_hdim=initial_states_strides[2],
722
+ stride_init_states_dstate=initial_states_strides[3],
723
HAS_INITSTATES=initial_states is not None)
724
return states
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+73
-83
@@ -31,6 +31,7 @@ def _mamba_chunk_scan_combined_fwd(x,
31
B,
32
C,
33
chunk_size,
34
+ out,
35
D=None,
36
z=None,
37
dt_bias=None,
@@ -41,14 +42,13 @@ def _mamba_chunk_scan_combined_fwd(x,
42
cu_seqlens=None,
43
dt_softplus=False,
44
dt_limit=(0.0, float("inf")),
44
- state_dtype=None,
45
- out=None):
45
+ state_dtype=None):
46
assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2"
47
- batch, seqlen, nheads, headdim = x.shape
48
- _, _, ngroups, dstate = B.shape
47
+ seqlen, nheads, headdim = x.shape
48
+ _, ngroups, dstate = B.shape
49
assert nheads % ngroups == 0
50
- assert B.shape == (batch, seqlen, ngroups, dstate)
51
- assert dt.shape == (batch, seqlen, nheads)
50
+ assert B.shape == (seqlen, ngroups, dstate)
51
+ assert dt.shape == (seqlen, nheads)
52
assert A.shape == (nheads, )
53
assert C.shape == B.shape
54
if z is not None:
@@ -56,25 +56,24 @@ def _mamba_chunk_scan_combined_fwd(x,
56
if D is not None:
57
assert D.shape == (nheads, headdim) or D.shape == (nheads, )
58
if seq_idx is not None:
59
- assert seq_idx.shape == (batch, seqlen)
59
+ assert seq_idx.shape == (seqlen, )
60
if B.stride(-1) != 1:
61
B = B.contiguous()
62
if C.stride(-1) != 1:
63
C = C.contiguous()
64
if x.stride(-1) != 1 and x.stride(
65
- 1) != 1: # Either M or K dimension should be contiguous
65
+ 0) != 1: # Either M or K dimension should be contiguous
66
x = x.contiguous()
67
if z is not None and z.stride(-1) != 1 and z.stride(
68
- 1) != 1: # Either M or K dimension should be contiguous
68
+ 0) != 1: # Either M or K dimension should be contiguous
69
z = z.contiguous()
70
if D is not None and D.stride(-1) != 1:
71
D = D.contiguous()
72
+ assert cu_seqlens is not None, "Assuming varlen input - must supply cu_seqlens"
73
+
74
if initial_states is not None:
73
- if cu_seqlens is None:
74
- assert initial_states.shape == (batch, nheads, headdim, dstate)
75
- else:
76
- assert initial_states.shape == (len(cu_seqlens) - 1, nheads,
77
- headdim, dstate)
75
+ assert initial_states.shape == (len(cu_seqlens) - 1, nheads, headdim,
76
+ dstate)
77
78
# This function executes 5 sub-functions for computing mamba
79
# - a good resource is the blog https://goombalab.github.io/blog/2024/mamba2-part3-algorithm/
@@ -114,18 +113,16 @@ def _mamba_chunk_scan_combined_fwd(x,
113
# - this will ensure that states will be updated with the rightmost flushed seq_idx
114
# of the previous chunk. This implies that the first chunk of states is either 0
115
# or equal to init_states of the first example.
117
- states, final_states = _state_passing_fwd(
116
+ states = _state_passing_fwd(
117
rearrange(states, "... p n -> ... (p n)"),
119
- dA_cumsum,
118
+ dA_cumsum, # (nheads, nchunks, chunk_size)
119
initial_states=rearrange(initial_states, "... p n -> ... (p n)")
121
- if initial_states is not None else None,
120
+ if initial_states is not None else
121
+ None, # (batch, nheads, headdim*dstate)
122
seq_idx=seq_idx,
123
- chunk_size=chunk_size,
123
out_dtype=state_dtype if state_dtype is not None else C.dtype,
125
- is_cont_batched=cu_seqlens is not None,
124
chunk_offsets=chunk_offsets)
127
- states, final_states = (rearrange(t, "... (p n) -> ... p n", n=dstate)
128
- for t in [states, final_states])
125
+ states = rearrange(states, "... (p n) -> ... p n", n=dstate)
126
127
# 4. Compute batched matrix multiply for C_j^T B_i terms
128
CB = _bmm_chunk_fwd(C,
@@ -144,87 +141,88 @@ def _mamba_chunk_scan_combined_fwd(x,
141
# - in each (pseudo) chunk, we detect if the previous (pseudo) chunk had
142
# a seq_idx change, in which case we take states information from
143
# init_states.
147
- out_x = _chunk_scan_fwd(
144
+ _chunk_scan_fwd(
145
CB,
146
x,
147
dt,
148
dA_cumsum,
149
C,
150
states,
151
+ out, # in-place update
152
+ seq_idx,
153
D=D,
154
z=z,
156
- seq_idx=seq_idx,
155
chunk_indices=chunk_indices,
156
chunk_offsets=chunk_offsets,
157
initial_states=initial_states,
160
- out=out,
158
)
162
- if cu_seqlens is None:
163
- return out_x, dt, dA_cumsum, states, final_states
164
- else:
165
- assert batch == 1, "passing cu_seqlens to get the varlen states is only supported if batch dimension is 1"
166
- varlen_states = chunk_state_varlen(
167
- B.squeeze(0),
168
- x.squeeze(0),
169
- dt.squeeze(0),
170
- dA_cumsum.squeeze(0),
171
- cu_seqlens,
172
- states.squeeze(0),
173
- initial_states=initial_states,
174
- )
175
- return out_x, dt, dA_cumsum, states, final_states, varlen_states
176
-
177
-
178
-def mamba_chunk_scan_combined(x,
179
- dt,
180
- A,
181
- B,
182
- C,
183
- chunk_size,
184
- D=None,
185
- z=None,
186
- dt_bias=None,
187
- initial_states=None,
188
- seq_idx=None,
189
- chunk_indices=None,
190
- chunk_offsets=None,
191
- cu_seqlens=None,
192
- dt_softplus=False,
193
- dt_limit=(0.0, float("inf")),
194
- out=None,
195
- return_final_states=False,
196
- return_varlen_states=False,
197
- state_dtype=None):
159
+
160
+ varlen_states = chunk_state_varlen(
161
+ B,
162
+ x,
163
+ dt,
164
+ dA_cumsum,
165
+ cu_seqlens,
166
+ states,
167
+ initial_states=initial_states,
168
+ )
169
+
170
+ return varlen_states
171
+
172
+
173
+def mamba_chunk_scan_combined_varlen(
174
+ x,
175
+ dt,
176
+ A,
177
+ B,
178
+ C,
179
+ chunk_size,
180
+ cu_seqlens,
181
+ seq_idx,
182
+ out,
183
+ D=None,
184
+ z=None,
185
+ dt_bias=None,
186
+ initial_states=None,
187
+ chunk_indices=None,
188
+ chunk_offsets=None,
189
+ dt_softplus=False,
190
+ dt_limit=(0.0, float("inf")),
191
+ state_dtype=None,
192
+):
193
"""
194
Argument:
200
- x: (batch, seqlen, nheads, headdim)
201
- dt: (batch, seqlen, nheads)
195
+ x: (seqlen, nheads, headdim)
196
+ dt: (seqlen, nheads)
197
A: (nheads)
203
- B: (batch, seqlen, ngroups, dstate)
204
- C: (batch, seqlen, ngroups, dstate)
198
+ B: (seqlen, ngroups, dstate)
199
+ C: (seqlen, ngroups, dstate)
200
chunk_size: int
201
+ seq_idx: (seqlen)
202
+ cu_seqlens: (batch + 1)
203
+ out: (seqlen, nheads, headdim) preallocated output tensor
204
D: (nheads, headdim) or (nheads,)
207
- z: (batch, seqlen, nheads, headdim)
205
+ z: (seqlen, nheads, headdim)
206
dt_bias: (nheads,)
207
initial_states: (batch, nheads, headdim, dstate)
210
- seq_idx: (batch, seqlen)
211
- cu_seqlens: (num_sequences + 1) or None, only used if return_varlen_states is True
208
dt_softplus: Whether to apply softplus to dt
213
- out: Preallocated output tensor
209
+ out: (seqlen, nheads, headdim) preallocated output tensor
210
state_dtype: The data type of the ssm state
211
+ Return:
212
+ varlen_states: (batch, nheads, headdim, dstate)
213
"""
214
217
- if not return_varlen_states:
218
- cu_seqlens = None
219
- else:
220
- assert cu_seqlens is not None, "cu_seqlens must be provided if return_varlen_states is True"
221
- out_x, dt_out, dA_cumsum, states, final_states, *rest = _mamba_chunk_scan_combined_fwd(
215
+ assert cu_seqlens is not None, "cu_seqlens must be provided assuming varlen input"
216
+ assert seq_idx is not None
217
+
218
+ varlen_states = _mamba_chunk_scan_combined_fwd(
219
x,
220
dt,
221
A,
222
B,
223
C,
224
chunk_size,
225
+ out,
226
D=D,
227
z=z,
228
dt_bias=dt_bias,
@@ -235,14 +233,6 @@ def mamba_chunk_scan_combined(x,
233
cu_seqlens=cu_seqlens,
234
dt_softplus=dt_softplus,
235
dt_limit=dt_limit,
238
- out=out,
236
state_dtype=state_dtype)
240
- if not return_varlen_states:
241
- if not return_final_states:
242
- return
243
- else:
244
- return final_states
245
- else:
246
- varlen_states = rest[0]
247
- return (varlen_states) if not return_final_states else (final_states,
248
- varlen_states)
237
+
238
+ return varlen_states
vllm/model_executor/layers/mamba/ops/ssd_state_passing.py
+110
-158
@@ -27,64 +27,46 @@ def _state_passing_fwd_kernel(
27
# Pointers to matrices
28
states_ptr,
29
out_ptr,
30
- final_states_ptr,
30
dA_cs_ptr,
31
initstates_ptr,
32
seq_idx_ptr,
33
chunk_offsets_ptr,
34
chunk_meta_num,
35
# Matrix dimensions
37
- dim,
36
+ dim: tl.constexpr,
37
nchunks,
38
seqlen,
40
- chunk_size,
39
+ chunk_size: tl.constexpr,
40
# Strides
42
- stride_states_batch,
43
- stride_states_chunk,
44
- stride_states_head,
45
- stride_states_dim,
46
- stride_out_batch,
47
- stride_out_chunk,
48
- stride_out_head,
49
- stride_out_dim,
50
- stride_final_states_batch,
51
- stride_final_states_head,
52
- stride_final_states_dim,
53
- stride_dA_cs_batch,
54
- stride_dA_cs_chunk,
55
- stride_dA_cs_head,
56
- stride_dA_cs_csize,
57
- stride_initstates_batch,
58
- stride_initstates_head,
59
- stride_initstates_dim,
60
- stride_seq_idx_batch,
61
- stride_seq_idx_seqlen,
41
+ stride_states_chunk: tl.int64,
42
+ stride_states_head: tl.int64,
43
+ stride_states_dim: tl.constexpr,
44
+ stride_out_chunk: tl.int64,
45
+ stride_out_head: tl.int64,
46
+ stride_out_dim: tl.constexpr,
47
+ stride_dA_cs_head: tl.int64,
48
+ stride_dA_cs_chunk: tl.int64,
49
+ stride_dA_cs_csize: tl.constexpr,
50
+ stride_initstates_batch: tl.int64,
51
+ stride_initstates_head: tl.int64,
52
+ stride_initstates_dim: tl.constexpr,
53
+ stride_seq_idx_seqlen: tl.constexpr,
54
# Meta-parameters
55
HAS_INITSTATES: tl.constexpr,
64
- HAS_SEQ_IDX: tl.constexpr,
65
- IS_CONT_BATCHED: tl.constexpr,
56
BLOCK_SIZE: tl.constexpr,
57
):
68
- pid_b = tl.program_id(axis=1)
69
- pid_h = tl.program_id(axis=2)
58
+ pid_h = tl.program_id(axis=1)
59
pid_m = tl.program_id(axis=0)
71
- states_ptr += pid_b * stride_states_batch + pid_h * stride_states_head
72
- dA_cs_ptr += pid_b * stride_dA_cs_batch + pid_h * stride_dA_cs_head + (
73
- chunk_size - 1) * stride_dA_cs_csize
74
- out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head
75
- final_states_ptr += pid_b * stride_final_states_batch + pid_h * stride_final_states_head
60
+ states_ptr += pid_h * stride_states_head
61
+ dA_cs_ptr += pid_h * stride_dA_cs_head + (chunk_size -
62
+ 1) * stride_dA_cs_csize
63
+ out_ptr += pid_h * stride_out_head
64
if HAS_INITSTATES:
65
initstates_ptr += pid_h * stride_initstates_head
78
- if not IS_CONT_BATCHED:
79
- initstates_ptr += pid_b * stride_initstates_batch
80
-
81
- if HAS_SEQ_IDX:
82
- seq_idx_ptr += pid_b * stride_seq_idx_batch
66
67
offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
68
states_ptrs = states_ptr + offs_m * stride_states_dim
69
out_ptrs = out_ptr + offs_m * stride_out_dim
87
- final_states_ptrs = final_states_ptr + offs_m * stride_final_states_dim
70
71
# - states will be the past state of the sequence that continues on the current check
72
if not HAS_INITSTATES:
@@ -101,65 +83,63 @@ def _state_passing_fwd_kernel(
83
out_ptrs += stride_out_chunk
84
prev_seq_idx_chunk_end = 0
85
logical_chunk_idx = 0
104
- for c in range(nchunks):
86
+ for c in range(nchunks - 1):
87
new_states = tl.load(states_ptrs, mask=offs_m < dim,
88
other=0.0).to(tl.float32)
89
dA_cs = tl.load(dA_cs_ptr).to(tl.float32)
90
scale_mask = True
109
- if HAS_SEQ_IDX:
110
- # - the seq to pass forward is the one that is flushed to the right
111
- # boundary.
112
- # - that is given by seq_idx_chunk_end below: the sequence index at the end of the chunk.
113
- seq_idx_chunk_end = tl.load(seq_idx_ptr + (min(
114
- (c + 1) * chunk_size, seqlen) - 1) * stride_seq_idx_seqlen)
115
- if HAS_INITSTATES:
116
- if IS_CONT_BATCHED and prev_seq_idx_chunk_end != seq_idx_chunk_end:
117
- # this means in the current chunk the rightmost flushed seq
118
- # has changed.
119
- # - so we do not propagate the state from previous chunk
120
- # - but rather we load that sequence's init state
121
- initstates_ptrs = initstates_ptr + seq_idx_chunk_end * stride_initstates_batch
122
-
123
- # - update state with seq_idx_new's init state
124
- states = tl.load(initstates_ptrs,
125
- mask=offs_m < dim,
126
- other=0.0).to(tl.float32)
127
-
128
- # - we need to consider the cumsum only of the last sequence in the chunk
129
- # - find its starting position (given by c_off of the logical chunk index)
130
- # - and subtract the cumsum just before that position from the total cumsum
131
- # - first, update the logical chunk index (add the number of sequences in the current physical chunk):
132
- # sequence index at the start of the current chunk
133
- seq_idx_chunk_start = tl.load(seq_idx_ptr +
134
- min(c * chunk_size, seqlen) *
135
- stride_seq_idx_seqlen)
136
- logical_chunk_idx += seq_idx_chunk_end - seq_idx_chunk_start
137
- # - load the chunk offset:
138
- c_off = tl.load(chunk_offsets_ptr + logical_chunk_idx,
139
- mask=logical_chunk_idx < chunk_meta_num,
140
- other=0)
141
- # - if offset is 0, then the sequence starts at the beginning of the chunk, and we don't need to subtract anything
142
- if c_off > 0:
143
- # - dA_cs_ptr currently points to the cumsum at the end of the chunk - subtract the chunk size and add the offset
144
- dA_cs_boundary = tl.load(
145
- dA_cs_ptr - (chunk_size - 1) * stride_dA_cs_csize +
146
- (c_off - 1) * stride_dA_cs_csize,
147
- mask=(c_off - 1) > -1 and c_off < chunk_size,
148
- other=0.0)
149
- dA_cs -= dA_cs_boundary
150
-
151
- # - increment logical chunk index for every physical chunk
152
- logical_chunk_idx += 1
153
- else:
154
- scale_mask = seq_idx_chunk_end == prev_seq_idx_chunk_end
155
- prev_seq_idx_chunk_end = seq_idx_chunk_end
91
+ # - the seq to pass forward is the one that is flushed to the right
92
+ # boundary.
93
+ # - that is given by seq_idx_chunk_end below: the sequence index at the end of the chunk.
94
+ seq_idx_chunk_end = tl.load(seq_idx_ptr +
95
+ (min((c + 1) * chunk_size, seqlen) - 1) *
96
+ stride_seq_idx_seqlen)
97
+
98
+ if HAS_INITSTATES:
99
+ if prev_seq_idx_chunk_end != seq_idx_chunk_end:
100
+ # this means in the current chunk the rightmost flushed seq
101
+ # has changed.
102
+ # - so we do not propagate the state from previous chunk
103
+ # - but rather we load that sequence's init state
104
+ initstates_ptrs = initstates_ptr + seq_idx_chunk_end * stride_initstates_batch
105
+
106
+ # - update state with seq_idx_new's init state
107
+ states = tl.load(initstates_ptrs, mask=offs_m < dim,
108
+ other=0.0).to(tl.float32)
109
+
110
+ # - we need to consider the cumsum only of the last sequence in the chunk
111
+ # - find its starting position (given by c_off of the logical chunk index)
112
+ # - and subtract the cumsum just before that position from the total cumsum
113
+ # - first, update the logical chunk index (add the number of sequences in the current physical chunk):
114
+ # sequence index at the start of the current chunk
115
+ seq_idx_chunk_start = tl.load(seq_idx_ptr +
116
+ min(c * chunk_size, seqlen) *
117
+ stride_seq_idx_seqlen)
118
+ logical_chunk_idx += seq_idx_chunk_end - seq_idx_chunk_start
119
+ # - load the chunk offset:
120
+ c_off = tl.load(chunk_offsets_ptr + logical_chunk_idx,
121
+ mask=logical_chunk_idx < chunk_meta_num,
122
+ other=0)
123
+ # - if offset is 0, then the sequence starts at the beginning of the chunk, and we don't need to subtract anything
124
+ if c_off > 0:
125
+ # - dA_cs_ptr currently points to the cumsum at the end of the chunk - subtract the chunk size and add the offset
126
+ dA_cs_boundary = tl.load(
127
+ dA_cs_ptr - (chunk_size - 1) * stride_dA_cs_csize +
128
+ (c_off - 1) * stride_dA_cs_csize,
129
+ mask=(c_off - 1) > -1 and c_off < chunk_size,
130
+ other=0.0)
131
+ dA_cs -= dA_cs_boundary
132
+
133
+ # - increment logical chunk index for every physical chunk
134
+ logical_chunk_idx += 1
135
+ else:
136
+ scale_mask = seq_idx_chunk_end == prev_seq_idx_chunk_end
137
+ prev_seq_idx_chunk_end = seq_idx_chunk_end
138
139
scale = tl.where(scale_mask, tl.exp(dA_cs), 0.0)
140
states = scale * states + new_states
159
- if c < nchunks - 1:
160
- tl.store(out_ptrs, states, mask=offs_m < dim)
161
- else:
162
- tl.store(final_states_ptrs, states, mask=offs_m < dim)
141
+ tl.store(out_ptrs, states, mask=offs_m < dim)
142
+
143
states_ptrs += stride_states_chunk
144
dA_cs_ptr += stride_dA_cs_chunk
145
out_ptrs += stride_out_chunk
@@ -168,81 +148,53 @@ def _state_passing_fwd_kernel(
148
def _state_passing_fwd(
149
states,
150
dA_cumsum,
151
+ seq_idx,
152
+ chunk_offsets,
153
initial_states=None,
172
- seq_idx=None,
173
- chunk_size=None,
154
out_dtype=None,
175
- is_cont_batched=False,
176
- chunk_offsets=None,
155
):
178
- batch, nchunks, nheads, dim = states.shape
179
- if chunk_size is None:
180
- chunk_size = dA_cumsum.shape[-1]
181
- else:
182
- assert chunk_size == dA_cumsum.shape[-1]
183
- assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
184
- if initial_states is not None:
185
- if is_cont_batched:
186
- # - if cu_seqlens is provided, then the initial states
187
- # are used for continuous batching. In which case we
188
- # require seq_idx to be provided
189
- assert seq_idx is not None, "seq_idx must be provided for continuous batching"
190
- # - we also need chunk_offsets to be provided, to account
191
- # for computation of dA_cumsum from the start of the
192
- # sequence
193
- assert chunk_offsets is not None, "chunk_offsets must be provided for continuous batching"
194
- else:
195
- # - this is the regular batching case, where initial
196
- # states are used are for each example of the batch.
197
- assert initial_states.shape == (batch, nheads, dim)
198
-
199
- if seq_idx is not None:
200
- seqlen = seq_idx.shape[-1]
201
- assert seq_idx.shape == (batch, seqlen)
156
+ nchunks, nheads, dim = states.shape
157
+ chunk_size = dA_cumsum.shape[-1]
158
+ assert dA_cumsum.shape == (nheads, nchunks, chunk_size)
159
+ seqlen = seq_idx.shape[-1]
160
out_dtype = states.dtype if out_dtype is None else out_dtype
203
- out = torch.empty((batch, nchunks, nheads, dim),
161
+ out = torch.empty((nchunks, nheads, dim),
162
device=states.device,
163
dtype=out_dtype)
206
- final_states = torch.empty((batch, nheads, dim),
207
- device=states.device,
208
- dtype=torch.float32)
209
- grid = lambda META: (triton.cdiv(dim, META['BLOCK_SIZE']), batch, nheads)
164
+
165
+ initial_states_strides = ((initial_states.stride(0),
166
+ initial_states.stride(1),
167
+ initial_states.stride(2))
168
+ if initial_states is not None else (0, 0, 0))
169
+
170
+ grid = lambda META: (triton.cdiv(dim, META['BLOCK_SIZE']), nheads)
171
with torch.cuda.device(states.device.index):
172
_state_passing_fwd_kernel[grid](
212
- states,
213
- out,
214
- final_states,
215
- dA_cumsum,
216
- initial_states,
217
- seq_idx,
218
- chunk_offsets,
219
- len(chunk_offsets) if chunk_offsets is not None else 0,
220
- dim,
221
- nchunks,
222
- seqlen if seq_idx is not None else 0,
223
- chunk_size,
224
- states.stride(0),
225
- states.stride(1),
226
- states.stride(2),
227
- states.stride(3),
228
- out.stride(0),
229
- out.stride(1),
230
- out.stride(2),
231
- out.stride(3),
232
- final_states.stride(0),
233
- final_states.stride(1),
234
- final_states.stride(2),
235
- dA_cumsum.stride(0),
236
- dA_cumsum.stride(2),
237
- dA_cumsum.stride(1),
238
- dA_cumsum.stride(3),
239
- *((initial_states.stride(0), initial_states.stride(1),
240
- initial_states.stride(2)) if initial_states is not None else
241
- (0, 0, 0)),
242
- *((seq_idx.stride(0),
243
- seq_idx.stride(1)) if seq_idx is not None else (0, 0)),
173
+ states_ptr=states,
174
+ out_ptr=out,
175
+ dA_cs_ptr=dA_cumsum,
176
+ initstates_ptr=initial_states,
177
+ seq_idx_ptr=seq_idx,
178
+ chunk_offsets_ptr=chunk_offsets,
179
+ chunk_meta_num=len(chunk_offsets)
180
+ if chunk_offsets is not None else 0,
181
+ dim=dim,
182
+ nchunks=nchunks,
183
+ seqlen=seqlen if seq_idx is not None else 0,
184
+ chunk_size=chunk_size if seq_idx is not None else 0,
185
+ stride_states_chunk=states.stride(0),
186
+ stride_states_head=states.stride(1),
187
+ stride_states_dim=states.stride(2),
188
+ stride_out_chunk=out.stride(0),
189
+ stride_out_head=out.stride(1),
190
+ stride_out_dim=out.stride(2),
191
+ stride_dA_cs_head=dA_cumsum.stride(0),
192
+ stride_dA_cs_chunk=dA_cumsum.stride(1),
193
+ stride_dA_cs_csize=dA_cumsum.stride(2),
194
+ stride_initstates_batch=initial_states_strides[0],
195
+ stride_initstates_head=initial_states_strides[1],
196
+ stride_initstates_dim=initial_states_strides[2],
197
+ stride_seq_idx_seqlen=seq_idx.stride(0),
198
HAS_INITSTATES=initial_states is not None,
245
- HAS_SEQ_IDX=seq_idx is not None,
246
- IS_CONT_BATCHED=is_cont_batched,
199
)
248
- return out, final_states
200
+ return out
vllm/model_executor/models/plamo2.py
+9
-13
@@ -35,7 +35,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
35
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
36
selective_state_update)
37
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
38
- mamba_chunk_scan_combined)
38
+ mamba_chunk_scan_combined_varlen)
39
from vllm.model_executor.layers.quantization import QuantizationConfig
40
from vllm.model_executor.layers.rotary_embedding import get_rope
41
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -262,6 +262,7 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
262
seq_idx_p = attn_metadata.seq_idx_p
263
chunk_indices_p = attn_metadata.chunk_indices_p
264
chunk_offsets_p = attn_metadata.chunk_offsets_p
265
+ query_start_loc_p = attn_metadata.query_start_loc_p
266
267
# 1. Gated MLP's linear projection
268
projected_states = self.in_proj(hidden_states)
@@ -302,9 +303,6 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
303
[num_decodes, num_prefills],
304
dim=0,
305
)
305
- query_start_loc_p = (
306
- attn_metadata.query_start_loc[-num_prefills - 1:] -
307
- num_decodes if has_prefill else None)
306
307
# Preallocate output tensor to avoid memcpy cost for merging prefill
308
# and decode outputs
@@ -356,17 +354,17 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
354
has_initial_states_p[:, None, None, None],
355
ssm_state[state_indices_tensor_p], 0)
356
359
- varlen_state = mamba_chunk_scan_combined(
360
- hidden_states_p.view(1, num_prefill_tokens,
357
+ varlen_state = mamba_chunk_scan_combined_varlen(
358
+ hidden_states_p.view(num_prefill_tokens,
359
self.num_heads // self.tp_size,
360
self.head_dim),
363
- dt.unsqueeze(0),
361
+ dt,
362
self.A,
365
- B.view(1, num_prefill_tokens, 1, -1),
366
- C.view(1, num_prefill_tokens, 1, -1),
363
+ B.view(num_prefill_tokens, 1, -1),
364
+ C.view(num_prefill_tokens, 1, -1),
365
chunk_size=chunk_size,
366
D=self.D,
369
- z=gate_p.view(1, num_prefill_tokens,
367
+ z=gate_p.view(num_prefill_tokens,
368
self.num_heads // self.tp_size, self.head_dim),
369
dt_bias=self.dt_bias,
370
seq_idx=seq_idx_p,
@@ -374,11 +372,9 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
372
chunk_offsets=chunk_offsets_p,
373
cu_seqlens=query_start_loc_p,
374
initial_states=initial_states,
377
- return_varlen_states=True,
378
- return_final_states=False,
375
dt_softplus=True,
376
dt_limit=(0.0, float("inf")),
381
- out=preallocated_ssm_out_p.view(1, num_prefill_tokens, -1,
377
+ out=preallocated_ssm_out_p.view(num_prefill_tokens, -1,
378
self.head_dim),
379
state_dtype=ssm_state.dtype,
380
)
vllm/v1/attention/backends/mamba2_attn.py
+4
-5
@@ -115,7 +115,7 @@ class Mamba2AttentionMetadata:
115
num_prefill_tokens: int
116
num_decodes: int
117
num_decode_tokens: int
118
- query_start_loc: torch.Tensor
118
+ query_start_loc_p: torch.Tensor
119
seq_lens: torch.Tensor
120
121
prep_initial_states: bool
@@ -151,7 +151,7 @@ class Mamba2AttentionMetadataBuilder(
151
common_attn_metadata: CommonAttentionMetadata,
152
fast_build: bool = False) -> Mamba2AttentionMetadata:
153
num_reqs = common_attn_metadata.num_reqs
154
- query_start_loc = common_attn_metadata.query_start_loc
154
+ query_start_loc_p = None
155
seq_lens = common_attn_metadata.seq_lens
156
157
seq_idx_p = None
@@ -179,7 +179,7 @@ class Mamba2AttentionMetadataBuilder(
179
num_computed_tokens_cpu[num_reqs - num_prefills:num_reqs] > 0)
180
prep_initial_states = torch.any(has_initial_states_cpu).item()
181
has_initial_states_p = has_initial_states_cpu.to(
182
- query_start_loc.device)
182
+ common_attn_metadata.query_start_loc.device)
183
184
query_start_loc_p = common_attn_metadata.query_start_loc[
185
-num_prefills - 1:] - num_decode_tokens
@@ -190,7 +190,6 @@ class Mamba2AttentionMetadataBuilder(
190
device=query_start_loc_p.device),
191
query_start_loc_p.diff(),
192
output_size=num_prefill_tokens)
193
- seq_idx_p.unsqueeze_(0)
193
194
# We compute metadata for chunked prefill once at the top level
195
# model forward and reuse them in mamba layers. If not needed,
@@ -217,7 +216,7 @@ class Mamba2AttentionMetadataBuilder(
216
num_prefill_tokens=num_prefill_tokens,
217
num_decodes=num_decodes,
218
num_decode_tokens=num_decode_tokens,
220
- query_start_loc=query_start_loc,
219
+ query_start_loc_p=query_start_loc_p,
220
seq_lens=seq_lens,
221
prep_initial_states=prep_initial_states,
222
chunk_size=self.chunk_size,