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
[Bugfix] Fix mamba2 prefill chunking (#23279)
Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Signed-off-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
tomeras91 committed
Sep 8, 2025 at 14:42 UTC
e041314184e3b1c4d797b2b025a5a9f54ade5a00
5 files changed
+348
-34
tests/kernels/mamba/test_mamba_ssm_ssd.py
+225
-8
@@ -115,21 +115,27 @@ def generate_continuous_batched_examples(example_lens_by_batch,
115
n_heads,
116
d_head,
117
itype,
118
- device='cuda'):
118
+ device='cuda',
119
+ return_naive_ref=True):
120
121
# this function generates a random examples of certain length
122
# and then cut according to "example_lens_by_batch" and feed
122
- # them in continuous batches to the kernels
123
+ # them in continuous batches to the kernels.
124
+ # If if return_naive_ref=True, the naive torch implementation
125
+ # ssd_minimal_discrete will be used to compute and return
126
+ # reference output.
127
128
# generate the full-length example
129
A, dt, X, B, C = generate_random_inputs(num_examples, full_length, n_heads,
130
d_head, itype)
131
128
- Y_min, final_state_min = ssd_minimal_discrete(X * dt.unsqueeze(-1),
129
- A * dt,
130
- B,
131
- C,
132
- block_len=full_length // 4)
132
+ if return_naive_ref:
133
+ Y_min, final_state_min = ssd_minimal_discrete(X * dt.unsqueeze(-1),
134
+ A * dt,
135
+ B,
136
+ C,
137
+ block_len=full_length //
138
+ 4)
139
140
# internal function that outputs a cont batch of examples
141
# given a tuple of lengths for each example in the batch
@@ -179,7 +185,8 @@ 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
182
- yield ([Y_min[s, IND_S[s]:IND_E[s]] for s in range(num_examples)],
188
+ yield ([Y_min[s, IND_S[s]:IND_E[s]]
189
+ for s in range(num_examples)] if return_naive_ref else None,
190
cu_seqlens, seq_idx.unsqueeze(0), (A, dt2, X2, B2, C2))
191
192
@@ -324,3 +331,213 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
331
if clear:
332
states[i].fill_(0.)
333
exhausted[i] = False
334
+
335
+
336
+@pytest.mark.parametrize("chunk_size", [8, 256])
337
+@pytest.mark.parametrize("seqlens", [
338
+ (16, 2, 8, 13),
339
+ (270, 88, 212, 203),
340
+ (16, 20),
341
+])
342
+def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
343
+
344
+ # This test verifies the correctness of the chunked prefill implementation
345
+ # in the mamba2 ssd kernels, by comparing concatenation (in the sequence
346
+ # dimension) of chunked results with the full sequence result.
347
+ # It is different from test_mamba_chunk_scan_cont_batch by:
348
+ # 1. Not using the naive torch implementaion (ssd_minimal_discrete) to get
349
+ # reference outputs. Instead, it compares chunked kernel outputs to full
350
+ # sequence kernel outputs. This is the most straightforward way to
351
+ # assert chunked prefill correctness.
352
+ # 2. It focuses on cases where sequences change in the middle of mamba
353
+ # chunks, and not necessarily on chunk boundaries.
354
+
355
+ max_seqlen = max(seqlens)
356
+ # This test can have larger error for longer sequences
357
+ if max_seqlen > 256:
358
+ atol, rtol = 1e-2, 5e-3
359
+ else:
360
+ atol, rtol = 5e-3, 5e-3
361
+
362
+ num_sequences = len(seqlens)
363
+ n_heads = 16
364
+ d_head = 64
365
+ itype = torch.float32
366
+
367
+ # hold state during the cutting process so we know if an
368
+ # example has been exhausted and needs to cycle
369
+ last_taken: dict = {} # map: eg -> pointer to last taken sample
370
+ exhausted: dict = {} # map: eg -> boolean indicating example is exhausted
371
+ _, cu_seqlens, seq_idx, (A, dt, X, B, C) = next(
372
+ generate_continuous_batched_examples([seqlens],
373
+ num_sequences,
374
+ max_seqlen,
375
+ last_taken,
376
+ exhausted,
377
+ n_heads,
378
+ d_head,
379
+ itype,
380
+ return_naive_ref=False))
381
+ seqlens = torch.tensor(seqlens, dtype=torch.int32, device=X.device)
382
+ device = X.device
383
+
384
+ ## full seqlen computation
385
+ chunk_indices, chunk_offsets = \
386
+ _query_start_loc_to_chunk_indices_offsets(
387
+ cu_seqlens, chunk_size, cu_seqlens[-1])
388
+ Y_ref = torch.empty_like(X)
389
+ state_ref = mamba_chunk_scan_combined(
390
+ X,
391
+ dt,
392
+ A,
393
+ B,
394
+ C,
395
+ chunk_size,
396
+ D=None,
397
+ cu_seqlens=cu_seqlens,
398
+ seq_idx=seq_idx,
399
+ chunk_indices=chunk_indices,
400
+ chunk_offsets=chunk_offsets,
401
+ return_varlen_states=True,
402
+ initial_states=None,
403
+ out=Y_ref,
404
+ )
405
+
406
+ ## chunked seqlen computation
407
+ # first chunk
408
+ chunked_seqlens = seqlens // 2
409
+ chunked_cu_seqlens = torch.cat([
410
+ torch.tensor([0], device=device),
411
+ torch.cumsum(chunked_seqlens, dim=0)
412
+ ],
413
+ dim=0)
414
+ chunked_seq_idx = torch.repeat_interleave(
415
+ torch.arange(len(chunked_seqlens), device=device),
416
+ chunked_seqlens,
417
+ output_size=chunked_cu_seqlens[-1]).unsqueeze(0).to(torch.int32)
418
+ 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, ...]
423
+ for i in range(num_sequences):
424
+ # fmt: off
425
+ chunk_f = lambda x, i: x[:, cu_seqlens[i]:cu_seqlens[i] + chunked_seqlens[i], ...] # noqa: E501
426
+
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
431
+ # fmt: on
432
+
433
+ chunk_indices, chunk_offsets = \
434
+ _query_start_loc_to_chunk_indices_offsets(
435
+ chunked_cu_seqlens, chunk_size, chunked_cu_seqlens[-1])
436
+ Y_partial = torch.empty_like(X_chunked)
437
+ partial_state = mamba_chunk_scan_combined(
438
+ X_chunked,
439
+ dt_chunked,
440
+ A,
441
+ B_chunked,
442
+ C_chunked,
443
+ chunk_size,
444
+ D=None,
445
+ cu_seqlens=chunked_cu_seqlens,
446
+ seq_idx=chunked_seq_idx,
447
+ chunk_indices=chunk_indices,
448
+ chunk_offsets=chunk_offsets,
449
+ return_varlen_states=True,
450
+ initial_states=None,
451
+ out=Y_partial,
452
+ )
453
+
454
+ # remaining chunk
455
+ remaining_chunked_seqlens = seqlens - chunked_seqlens
456
+ remaining_chunked_cu_seqlens = torch.cat([
457
+ torch.tensor([0], device=device),
458
+ torch.cumsum(remaining_chunked_seqlens, dim=0)
459
+ ],
460
+ dim=0)
461
+ remaining_chunked_seq_idx = torch.repeat_interleave(
462
+ torch.arange(len(remaining_chunked_seqlens), device=device),
463
+ remaining_chunked_seqlens,
464
+ output_size=remaining_chunked_cu_seqlens[-1]).unsqueeze(0).to(
465
+ torch.int32)
466
+ remaining_chunked_input_seq_len = remaining_chunked_cu_seqlens[-1]
467
+ # 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
472
+ 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
474
+
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
479
+
480
+ # assert input chunking is correct
481
+ 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],...],
484
+ ],
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
487
+ # fmt: on
488
+
489
+ assert concat_batch_f(X_chunked, remaining_X_chunked).equal(X)
490
+ assert concat_batch_f(dt_chunked, remaining_dt_chunked).equal(dt)
491
+ assert concat_batch_f(B_chunked, remaining_B_chunked).equal(B)
492
+ assert concat_batch_f(C_chunked, remaining_C_chunked).equal(C)
493
+
494
+ chunk_indices, chunk_offsets = \
495
+ _query_start_loc_to_chunk_indices_offsets(
496
+ remaining_chunked_cu_seqlens,
497
+ chunk_size,
498
+ remaining_chunked_cu_seqlens[-1])
499
+
500
+ Y_chunked = torch.empty_like(remaining_X_chunked)
501
+ state_chunked = mamba_chunk_scan_combined(
502
+ remaining_X_chunked,
503
+ remaining_dt_chunked,
504
+ A,
505
+ remaining_B_chunked,
506
+ remaining_C_chunked,
507
+ chunk_size,
508
+ D=None,
509
+ cu_seqlens=remaining_chunked_cu_seqlens,
510
+ seq_idx=remaining_chunked_seq_idx,
511
+ chunk_indices=chunk_indices,
512
+ chunk_offsets=chunk_offsets,
513
+ return_varlen_states=True,
514
+ initial_states=partial_state,
515
+ out=Y_chunked,
516
+ )
517
+ Y = concat_batch_f(Y_partial, Y_chunked)
518
+
519
+ # kernel chunked is same as kernel overall
520
+ 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], ...]
523
+ torch.testing.assert_close(
524
+ Y_seq[:, :chunked_seqlens[i], ...],
525
+ Y_ref_seq[:, :chunked_seqlens[i], ...],
526
+ atol=atol,
527
+ rtol=rtol,
528
+ msg=lambda x: f"seq{i} output part1 " + x) # noqa: B023
529
+ torch.testing.assert_close(
530
+ Y_seq[:, chunked_seqlens[i]:, ...],
531
+ Y_ref_seq[:, chunked_seqlens[i]:, ...],
532
+ atol=atol,
533
+ rtol=rtol,
534
+ msg=lambda x: f"seq{i} output part2 " + x) # noqa: B023
535
+
536
+ state_seq = state_chunked[i]
537
+ state_seq_ref = state_ref[i]
538
+ torch.testing.assert_close(
539
+ state_seq,
540
+ state_seq_ref,
541
+ atol=atol,
542
+ rtol=rtol,
543
+ msg=lambda x: f"seq{i} state " + x) # noqa: B023
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+3
@@ -289,6 +289,9 @@ def _chunk_scan_fwd_kernel(
289
290
# get the cs at the offset boundary
291
# - c_off == 0 is a passthrough
292
+ # - We need dA_cs at the boundary, defined by c_off - no need
293
+ # to increase pointer by pid_m (it is a constant offset,
294
+ # i.e. the same for all blocks)
295
dA_cs_m_boundary = tl.load(
296
dA_cumsum_ptr + (c_off - 1) * stride_dA_cs_csize,
297
mask=(((c_off - 1) > -1) and ((c_off) < chunk_size)),
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+6
-3
@@ -106,21 +106,24 @@ def _mamba_chunk_scan_combined_fwd(x,
106
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
107
# (middle term of factorization of off-diag blocks; A terms)
108
# - for handling chunked prefill, this requires i) initial_states
109
- # ii) seq_idx and iii) is_cont_batched to be all specified.
109
+ # ii) seq_idx iii) is_cont_batched and (iv) chunk_offsets to be all specified.
110
# - When a new seq_idx is detected, we will stop passing the prev_state
111
# and switch accordingly to the init_state corresponding to the new seq_idx.
112
+ # - We will also make sure that the dA_cumsum is taken only from the start of the
113
+ # sequence (hence we need the full dA_cumsum tensor and not just the values at chunk boundaries)
114
# - this will ensure that states will be updated with the rightmost flushed seq_idx
115
# of the previous chunk. This implies that the first chunk of states is either 0
116
# or equal to init_states of the first example.
117
states, final_states = _state_passing_fwd(
118
rearrange(states, "... p n -> ... (p n)"),
117
- dA_cumsum[:, :, :, -1],
119
+ dA_cumsum,
120
initial_states=rearrange(initial_states, "... p n -> ... (p n)")
121
if initial_states is not None else None,
122
seq_idx=seq_idx,
123
chunk_size=chunk_size,
124
out_dtype=state_dtype if state_dtype is not None else C.dtype,
123
- is_cont_batched=cu_seqlens is not None)
125
+ is_cont_batched=cu_seqlens is not None,
126
+ chunk_offsets=chunk_offsets)
127
states, final_states = (rearrange(t, "... (p n) -> ... p n", n=dstate)
128
for t in [states, final_states])
129
vllm/model_executor/layers/mamba/ops/ssd_state_passing.py
+62
-20
@@ -31,6 +31,8 @@ def _state_passing_fwd_kernel(
31
dA_cs_ptr,
32
initstates_ptr,
33
seq_idx_ptr,
34
+ chunk_offsets_ptr,
35
+ chunk_meta_num,
36
# Matrix dimensions
37
dim,
38
nchunks,
@@ -51,6 +53,7 @@ def _state_passing_fwd_kernel(
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,
@@ -66,7 +69,8 @@ def _state_passing_fwd_kernel(
69
pid_h = tl.program_id(axis=2)
70
pid_m = tl.program_id(axis=0)
71
states_ptr += pid_b * stride_states_batch + pid_h * stride_states_head
69
- dA_cs_ptr += pid_b * stride_dA_cs_batch + pid_h * stride_dA_cs_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
76
if HAS_INITSTATES:
@@ -95,35 +99,62 @@ def _state_passing_fwd_kernel(
99
100
tl.store(out_ptrs, states, mask=offs_m < dim)
101
out_ptrs += stride_out_chunk
98
- seq_idx = 0
102
+ prev_seq_idx_chunk_end = 0
103
+ logical_chunk_idx = 0
104
for c in range(nchunks):
105
new_states = tl.load(states_ptrs, mask=offs_m < dim,
106
other=0.0).to(tl.float32)
107
dA_cs = tl.load(dA_cs_ptr).to(tl.float32)
103
- scale = tl.exp(dA_cs)
108
+ 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.
107
- # - that is given by seq_idx_new below.
108
- seq_idx_new = tl.load(seq_idx_ptr +
109
- (min((c + 1) * chunk_size, seqlen) - 1) *
110
- stride_seq_idx_seqlen)
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:
112
- if IS_CONT_BATCHED and seq_idx != seq_idx_new:
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
117
- initstates_ptrs = initstates_ptr + seq_idx_new * stride_initstates_batch
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:
124
- scale = tl.where(seq_idx_new == seq_idx, scale, 0.0)
154
+ scale_mask = seq_idx_chunk_end == prev_seq_idx_chunk_end
155
+ prev_seq_idx_chunk_end = seq_idx_chunk_end
156
126
- seq_idx = seq_idx_new
157
+ scale = tl.where(scale_mask, tl.exp(dA_cs), 0.0)
158
states = scale * states + new_states
159
if c < nchunks - 1:
160
tl.store(out_ptrs, states, mask=offs_m < dim)
@@ -136,28 +167,36 @@ def _state_passing_fwd_kernel(
167
168
def _state_passing_fwd(
169
states,
139
- dA_chunk_cumsum,
170
+ dA_cumsum,
171
initial_states=None,
172
seq_idx=None,
173
chunk_size=None,
174
out_dtype=None,
175
is_cont_batched=False,
176
+ chunk_offsets=None,
177
):
178
batch, nchunks, nheads, dim = states.shape
147
- assert dA_chunk_cumsum.shape == (batch, nheads, nchunks)
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
153
- assert seq_idx is not None, ""
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:
160
- assert chunk_size is not None
200
seqlen = seq_idx.shape[-1]
201
assert seq_idx.shape == (batch, seqlen)
202
out_dtype = states.dtype if out_dtype is None else out_dtype
@@ -173,13 +212,15 @@ def _state_passing_fwd(
212
states,
213
out,
214
final_states,
176
- dA_chunk_cumsum,
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,
182
- chunk_size if seq_idx is not None else 0,
223
+ chunk_size,
224
states.stride(0),
225
states.stride(1),
226
states.stride(2),
@@ -191,9 +232,10 @@ def _state_passing_fwd(
232
final_states.stride(0),
233
final_states.stride(1),
234
final_states.stride(2),
194
- dA_chunk_cumsum.stride(0),
195
- dA_chunk_cumsum.stride(2),
196
- dA_chunk_cumsum.stride(1),
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)),
vllm/v1/attention/backends/mamba2_attn.py
+52
-3
@@ -16,9 +16,58 @@ from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
16
from vllm.v1.kv_cache_interface import AttentionSpec
17
18
19
-def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,
20
- chunk_size: int,
21
- total_seqlens: int):
19
+def _query_start_loc_to_chunk_indices_offsets(
20
+ query_start_loc: torch.Tensor, chunk_size: int,
21
+ total_seqlens: int) -> tuple[torch.Tensor, torch.Tensor]:
22
+ """
23
+ Args:
24
+ query_start_loc (torch.Tensor): 1D tensor of cumulative sequence
25
+ lengths, shape (num_seqs + 1,).
26
+ The first element should be 0. Each entry represents the starting
27
+ index of a sequence in the flattened token array.
28
+ chunk_size (int): The size of each physical mamba chunk
29
+ (number of tokens per chunk).
30
+ total_seqlens (int): The total number of tokens in the batch.
31
+
32
+ Returns:
33
+ Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
34
+ - chunk_indices (torch.Tensor): 1D tensor of indices
35
+ indicating the physical chunk for each logical chunk.
36
+ - chunk_offsets (torch.Tensor): 1D tensor of offsets
37
+ indicating the starting index of each logical chunk within
38
+ its physical chunk.
39
+
40
+ This function computes the chunk indices and offsets for the given
41
+ query_start_loc and chunk_size. Both are tensors of integers with length N,
42
+ where N is the number of logical (pseudo) chunks.
43
+ A logical chunk is a sequence of tokens that are all part of the same
44
+ sequence and are all in the same physical mamba chunk.
45
+ In other words, a logical chunk changes every time we cross a sequence
46
+ boundary or a physical mamba chunk boundary.
47
+ Logical chunks are needed to handle batched requests with initial states
48
+ (see _state_passing_fwd and _chunk_scan_fwd).
49
+ The chunk_indices tensor contains the index of the physical chunk for each
50
+ logical chunk.
51
+ The chunk_offsets tensor contains the offset (AKA starting index) of the
52
+ logical chunk in the physical chunk.
53
+
54
+ Example:
55
+ query_start_loc = [0, 5, 10]
56
+ chunk_size = 8
57
+ total_seqlens = 10
58
+ -> chunk_indices = [0, 0, 1]
59
+ -> chunk_offsets = [0, 5, 0]
60
+
61
+ In this example, we have 2 sequences, each with 5 tokens. The physical
62
+ chunk size is 8 tokens.
63
+ We have three logical chunks:
64
+ - the first logical chunk starts at token 0 in the first physical chunk
65
+ and contains all 5 tokens from the first sequence
66
+ - the second logical chunk starts at token 5 in the first physical chunk
67
+ and contains first 3 tokens from the second sequence
68
+ - the third logical chunk starts at token 0 in the second physical chunk
69
+ and contains the remaining 2 tokens from the second sequence
70
+ """
71
72
cu_seqlens = query_start_loc[1:] # remove prepended 0
73