test_mamba_ssm_ssd.py 19.5 KB
Newer Older
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
3
4
5
6
7
8
9

import pytest
import torch
import torch.nn.functional as F
from einops import rearrange, repeat

from vllm.model_executor.layers.mamba.ops.ssd_combined import (
10
11
    mamba_chunk_scan_combined_varlen,
)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
12
from vllm.platforms import current_platform
13
from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
14
15
16
17
18
19
20
21
22
23
24

# Added by the IBM Team, 2024

# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/modules/ssd_minimal.py


# this is the segsum implementation taken from above
def segsum(x):
    """Calculates segment sum."""
    T = x.size(-1)
    x = repeat(x, "... d -> ... d e", e=T)
25
    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=-1)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
26
27
    x = x.masked_fill(~mask, 0)
    x_segsum = torch.cumsum(x, dim=-2)
28
    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
    x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
    return x_segsum


def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None):
    """
    Arguments:
        X: (batch, length, n_heads, d_head)
        A: (batch, length, n_heads)
        B: (batch, length, n_heads, d_state)
        C: (batch, length, n_heads, d_state)
    Return:
        Y: (batch, length, n_heads, d_head)
    """
    assert X.dtype == A.dtype == B.dtype == C.dtype
    assert X.shape[1] % block_len == 0

    # Rearrange into blocks/chunks
47
48
49
    X, A, B, C = (
        rearrange(x, "b (c l) ... -> b c l ...", l=block_len) for x in (X, A, B, C)
    )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

    A = rearrange(A, "b c l h -> b h c l")
    A_cumsum = torch.cumsum(A, dim=-1)

    # 1. Compute the output for each intra-chunk (diagonal blocks)
    L = torch.exp(segsum(A))
    Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X)

    # 2. Compute the state for each intra-chunk
    # (right term of low-rank factorization of off-diagonal blocks; B terms)
    decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)

    # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at
    #    chunk boundaries
    # (middle term of factorization of off-diag blocks; A terms)
    if initial_states is None:
        initial_states = torch.zeros_like(states[:, :1])
    states = torch.cat([initial_states, states], dim=1)
    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
    states, final_state = new_states[:, :-1], new_states[:, -1]

    # 4. Compute state -> output conversion per chunk
    # (left term of low-rank factorization of off-diagonal blocks; C terms)
    state_decay_out = torch.exp(A_cumsum)
76
    Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, state_decay_out)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
77
78
79
80
81
82
83

    # Add output of intra-chunk and inter-chunk terms
    # (diagonal and off-diagonal blocks)
    Y = rearrange(Y_diag + Y_off, "b c l h p -> b (c l) h p")
    return Y, final_state


84
def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device="cuda"):
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
85
    current_platform.seed_everything(0)
86
    A = -torch.exp(torch.rand(n_heads, dtype=itype, device=device))
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
87
    dt = F.softplus(
88
89
90
91
92
        torch.randn(batch_size, seqlen, n_heads, dtype=itype, device=device) - 4
    )
    X = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device)
    B = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device)
    C = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
93
94
95
96

    return A, dt, X, B, C


97
98
99
100
101
102
103
104
105
106
107
108
def generate_continuous_batched_examples(
    example_lens_by_batch,
    num_examples,
    full_length,
    last_taken,
    exhausted,
    n_heads,
    d_head,
    itype,
    device="cuda",
    return_naive_ref=True,
):
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
109
110
    # this function generates a random examples of certain length
    # and then cut according to "example_lens_by_batch" and feed
111
112
113
114
    # them in continuous batches to the kernels.
    # If if return_naive_ref=True, the naive torch implementation
    # ssd_minimal_discrete will be used to compute and return
    # reference output.
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
115
116

    # generate the full-length example
117
118
119
    A, dt, X, B, C = generate_random_inputs(
        num_examples, full_length, n_heads, d_head, itype
    )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
120

121
    if return_naive_ref:
122
123
124
        Y_min, final_state_min = ssd_minimal_discrete(
            X * dt.unsqueeze(-1), A * dt, B, C, block_len=full_length // 4
        )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
125
126
127
128
129

    # internal function that outputs a cont batch of examples
    # given a tuple of lengths for each example in the batch
    # e.g., example_lens=(8, 4) means take 8 samples from first eg,
    #       4 examples from second eg, etc
130
    def get_continuous_batch(example_lens: tuple[int, ...]):
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
131
132
133
134
135
136
137
        indices = []
        for i, x in enumerate(example_lens):
            c = last_taken.get(i, 0)
            indices.append((c, c + x))
            last_taken[i] = (c + x) % full_length
            exhausted[i] = last_taken[i] == 0

138
139
140
141
        return (
            torch.concat([x[i, s:e] for i, (s, e) in enumerate(indices)]).unsqueeze(0)
            for x in (dt, X, B, C)
        )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

    # internal function that maps "n" to the appropriate right boundary
    # value when forming continuous batches from examples of length given
    # by "full_length".
    # - e.g., when n > full_length, returns n % full_length
    #         when n == full_length, returns full_length
    def end_boundary(n: int):
        return n - ((n - 1) // full_length) * full_length

    IND_E = None
    for spec in example_lens_by_batch:
        # get the (maybe partial) example seen in this cont batch
        dt2, X2, B2, C2 = get_continuous_batch(spec)

        # get the metadata
157
158
159
160
161
162
        cu_seqlens = torch.tensor((0,) + spec, device=device).cumsum(dim=0)
        seq_idx = torch.zeros(
            cu_seqlens[-1], dtype=torch.int32, device=cu_seqlens.device
        )
        for i, (srt, end) in enumerate(
            zip(
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
163
164
                cu_seqlens,
                cu_seqlens[1:],
165
166
            )
        ):
167
            seq_idx[srt:end] = i
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
168
169
170
171
172
173
174
175

        # for cont batch
        if IND_E is None:
            IND_S = [0 for _ in range(len(spec))]
        else:
            IND_S = [x % full_length for x in IND_E]
        IND_E = [end_boundary(x + y) for x, y in zip(IND_S, spec)]

176
177
178
179
180
        # varlen has implicit batch=1
        dt2 = dt2.squeeze(0)
        X2 = X2.squeeze(0)
        B2 = B2.squeeze(0)
        C2 = C2.squeeze(0)
181
182
183
184
185
186
187
188
        yield (
            [Y_min[s, IND_S[s] : IND_E[s]] for s in range(num_examples)]
            if return_naive_ref
            else None,
            cu_seqlens,
            seq_idx,
            (A, dt2, X2, B2, C2),
        )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
189
190


191
192
193
@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("n_heads", [4, 16, 32])
@pytest.mark.parametrize("d_head", [5, 8, 32, 128])
194
@pytest.mark.parametrize("seq_len_chunk_size", [(112, 16), (128, 32)])
195
def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, itype):
196
    # this tests the kernels on a single example (bs=1)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
197

198
199
200
201
202
203
204
    # TODO: the bfloat16 case requires higher thresholds. To be investigated

    if itype == torch.bfloat16:
        atol, rtol = 5e-2, 5e-2
    else:
        atol, rtol = 8e-3, 5e-3

Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
205
206
207
208
209
210
211
    # set seed
    batch_size = 1  # batch_size
    # ssd_minimal_discrete requires chunk_size divide seqlen
    # - this is only required for generating the reference seqs,
    #   it is not an operational limitation.
    seqlen, chunk_size = seq_len_chunk_size

212
    A, dt, X, B, C = generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
213

214
215
216
    Y_min, final_state_min = ssd_minimal_discrete(
        X * dt.unsqueeze(-1), A * dt, B, C, chunk_size
    )
217

218
219
    cu_seqlens = torch.tensor((0, seqlen), device="cuda").cumsum(dim=0)
    cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
220
221
        compute_varlen_chunk_metadata(cu_seqlens, chunk_size)
    )
222
223
224
225
226
227
    # varlen has implicit batch=1
    X = X.squeeze(0)
    dt = dt.squeeze(0)
    A = A.squeeze(0)
    B = B.squeeze(0)
    C = C.squeeze(0)
228
    Y = torch.empty_like(X)
229
230
231
232
233
234
235
236
237
238
239
240
241
242
    final_state = mamba_chunk_scan_combined_varlen(
        X,
        dt,
        A,
        B,
        C,
        chunk_size,
        cu_seqlens=cu_seqlens.to(torch.int32),
        cu_chunk_seqlens=cu_chunk_seqlens,
        last_chunk_indices=last_chunk_indices,
        seq_idx=seq_idx_chunks,
        out=Y,
        D=None,
    )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
243
244

    # just test the last in sequence
245
    torch.testing.assert_close(Y[-1], Y_min[0, -1], atol=atol, rtol=rtol)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
246
247
248

    # just test the last head
    # NOTE, in the kernel we always cast states to fp32
249
250
251
252
253
254
    torch.testing.assert_close(
        final_state[:, -1].to(torch.float32),
        final_state_min[:, -1].to(torch.float32),
        atol=atol,
        rtol=rtol,
    )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
255
256


257
258
259
@pytest.mark.parametrize("itype", [torch.float32])
@pytest.mark.parametrize("n_heads", [4, 8])
@pytest.mark.parametrize("d_head", [5, 16, 32])
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
260
261
262
263
264
265
@pytest.mark.parametrize(
    "seq_len_chunk_size_cases",
    [
        # small-ish chunk_size (8)
        (64, 8, 2, [(64, 32), (64, 32)]),
        (64, 8, 2, [(8, 8), (8, 8), (8, 8)]),  # chunk size boundary
266
267
268
269
270
271
        (
            64,
            8,
            2,
            [(4, 4), (4, 4), (4, 4), (4, 4)],
        ),  # chunk_size larger than cont batches
272
        (64, 8, 5, [(64, 32, 16, 8, 8)]),
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
273
        # large-ish chunk_size (256)
274
275
276
277
278
279
280
        (64, 256, 1, [(5,), (1,), (1,), (1,)]),  # irregular sizes with small sequences
        (
            64,
            256,
            2,
            [(5, 30), (1, 2), (1, 2), (1, 2)],
        ),  # irregular sizes with small sequences
281
282
283
        # we also need to test some large seqlen
        # to catch errors with init states decay
        (768, 128, 2, [(138, 225), (138, 225)]),
284
285
286
    ],
)
def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, itype):
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
287
288
289
290
291
    # this test with multiple examples in a continuous batch
    # (i.e. chunked prefill)

    seqlen, chunk_size, num_examples, cases = seq_len_chunk_size_cases

292
293
294
    # This test can have larger error for longer sequences
    if seqlen > 256:
        atol, rtol = 1e-2, 5e-3
295
296
297
    else:
        atol, rtol = 5e-3, 5e-3

Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
298
299
    # hold state during the cutting process so we know if an
    # example has been exhausted and needs to cycle
300
301
    last_taken: dict = {}  # map: eg -> pointer to last taken sample
    exhausted: dict = {}  # map: eg -> boolean indicating example is exhausted
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
302
303

    states = None
304
    for Y_min, cu_seqlens, _token_seq_idx, (
305
306
307
308
309
310
311
312
        A,
        dt,
        X,
        B,
        C,
    ) in generate_continuous_batched_examples(
        cases, num_examples, seqlen, last_taken, exhausted, n_heads, d_head, itype
    ):
313
        cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
314
315
            compute_varlen_chunk_metadata(cu_seqlens, chunk_size)
        )
316

317
        Y = torch.empty_like(X)
318
        new_states = mamba_chunk_scan_combined_varlen(
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
319
320
321
322
323
324
            X,
            dt,
            A,
            B,
            C,
            chunk_size,
325
326
327
328
329
            cu_seqlens=cu_seqlens.to(torch.int32),
            cu_chunk_seqlens=cu_chunk_seqlens,
            last_chunk_indices=last_chunk_indices,
            seq_idx=seq_idx_chunks,
            out=Y,
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
330
331
332
333
334
335
336
            D=None,
            initial_states=states,
        )

        # just test the last in sequence
        for i in range(num_examples):
            # just test one dim and dstate
337
            Y_eg = Y[cu_seqlens[i] : cu_seqlens[i + 1], 0, 0]
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
338
            Y_min_eg = Y_min[i][:, 0, 0]
339
            torch.testing.assert_close(Y_eg, Y_min_eg, atol=atol, rtol=rtol)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
340
341
342
343
344

        # update states
        states = new_states
        for i, clear in exhausted.items():
            if clear:
345
                states[i].fill_(0.0)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
346
                exhausted[i] = False
347
348
349


@pytest.mark.parametrize("chunk_size", [8, 256])
350
351
@pytest.mark.parametrize(
    "seqlens",
352
    [(16, 20), (270, 88, 212, 203)],
353
)
354
355
356
357
358
def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
    # This test verifies the correctness of the chunked prefill implementation
    # in the mamba2 ssd kernels, by comparing concatenation (in the sequence
    # dimension) of chunked results with the full sequence result.
    # It is different from test_mamba_chunk_scan_cont_batch by:
co63oc's avatar
co63oc committed
359
    # 1. Not using the naive torch implementation (ssd_minimal_discrete) to get
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
    #    reference outputs. Instead, it compares chunked kernel outputs to full
    #    sequence kernel outputs. This is the most straightforward way to
    #    assert chunked prefill correctness.
    # 2. It focuses on cases where sequences change in the middle of mamba
    #    chunks, and not necessarily on chunk boundaries.

    max_seqlen = max(seqlens)
    # This test can have larger error for longer sequences
    if max_seqlen > 256:
        atol, rtol = 1e-2, 5e-3
    else:
        atol, rtol = 5e-3, 5e-3

    num_sequences = len(seqlens)
    n_heads = 16
    d_head = 64
    itype = torch.float32

    # hold state during the cutting process so we know if an
    # example has been exhausted and needs to cycle
    last_taken: dict = {}  # map: eg -> pointer to last taken sample
    exhausted: dict = {}  # map: eg -> boolean indicating example is exhausted
    _, cu_seqlens, seq_idx, (A, dt, X, B, C) = next(
383
384
385
386
387
388
389
390
391
392
393
394
        generate_continuous_batched_examples(
            [seqlens],
            num_sequences,
            max_seqlen,
            last_taken,
            exhausted,
            n_heads,
            d_head,
            itype,
            return_naive_ref=False,
        )
    )
395
396
397
398
    seqlens = torch.tensor(seqlens, dtype=torch.int32, device=X.device)
    device = X.device

    ## full seqlen computation
399
    cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
400
401
        compute_varlen_chunk_metadata(cu_seqlens, chunk_size)
    )
402
    Y_ref = torch.empty_like(X)
403
    state_ref = mamba_chunk_scan_combined_varlen(
404
405
406
407
408
409
        X,
        dt,
        A,
        B,
        C,
        chunk_size,
410
411
412
413
414
        cu_seqlens=cu_seqlens.to(torch.int32),
        cu_chunk_seqlens=cu_chunk_seqlens,
        last_chunk_indices=last_chunk_indices,
        seq_idx=seq_idx_chunks,
        out=Y_ref,
415
416
417
418
419
420
421
        D=None,
        initial_states=None,
    )

    ## chunked seqlen computation
    # first chunk
    chunked_seqlens = seqlens // 2
422
423
424
    chunked_cu_seqlens = torch.cat(
        [torch.tensor([0], device=device), torch.cumsum(chunked_seqlens, dim=0)], dim=0
    )
425
    chunked_input_seq_len = chunked_cu_seqlens[-1]
426
427
428
429
    X_chunked = torch.zeros_like(X)[:chunked_input_seq_len, ...]
    dt_chunked = torch.zeros_like(dt)[:chunked_input_seq_len, ...]
    B_chunked = torch.zeros_like(B)[:chunked_input_seq_len, ...]
    C_chunked = torch.zeros_like(C)[:chunked_input_seq_len, ...]
430
    for i in range(num_sequences):
431
432
433
        chunk_f = lambda x, i: x[
            cu_seqlens[i] : cu_seqlens[i] + chunked_seqlens[i], ...
        ]
434

435
436
437
438
439
440
441
442
443
444
445
446
        X_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
            X, i
        )
        dt_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
            dt, i
        )
        B_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
            B, i
        )
        C_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
            C, i
        )
447

448
    cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
449
450
        compute_varlen_chunk_metadata(chunked_cu_seqlens, chunk_size)
    )
451
    Y_partial = torch.empty_like(X_chunked)
452
    partial_state = mamba_chunk_scan_combined_varlen(
453
454
455
456
457
458
        X_chunked,
        dt_chunked,
        A,
        B_chunked,
        C_chunked,
        chunk_size,
459
460
461
462
463
        cu_seqlens=chunked_cu_seqlens.to(torch.int32),
        cu_chunk_seqlens=cu_chunk_seqlens,
        last_chunk_indices=last_chunk_indices,
        seq_idx=seq_idx_chunks,
        out=Y_partial,
464
465
466
467
468
469
        D=None,
        initial_states=None,
    )

    # remaining chunk
    remaining_chunked_seqlens = seqlens - chunked_seqlens
470
471
472
473
474
475
476
    remaining_chunked_cu_seqlens = torch.cat(
        [
            torch.tensor([0], device=device),
            torch.cumsum(remaining_chunked_seqlens, dim=0),
        ],
        dim=0,
    )
477
    remaining_chunked_input_seq_len = remaining_chunked_cu_seqlens[-1]
478
479
480
481
    remaining_X_chunked = torch.zeros_like(X)[:remaining_chunked_input_seq_len, ...]
    remaining_dt_chunked = torch.zeros_like(dt)[:remaining_chunked_input_seq_len, ...]
    remaining_B_chunked = torch.zeros_like(B)[:remaining_chunked_input_seq_len, ...]
    remaining_C_chunked = torch.zeros_like(C)[:remaining_chunked_input_seq_len, ...]
482
    for i in range(num_sequences):
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
        remaining_chunk_f = lambda x, i: x[
            cu_seqlens[i] + chunked_seqlens[i] : cu_seqlens[i + 1], ...
        ]

        remaining_X_chunked[
            remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
        ] = remaining_chunk_f(X, i)
        remaining_dt_chunked[
            remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
        ] = remaining_chunk_f(dt, i)
        remaining_B_chunked[
            remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
        ] = remaining_chunk_f(B, i)
        remaining_C_chunked[
            remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
        ] = remaining_chunk_f(C, i)
499
500

    # assert input chunking is correct
501
502
503
504
505
506
507
    concat_chunk_f = lambda pt1, pt2, i: torch.cat(
        [
            pt1[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...],
            pt2[
                remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1],
                ...,
            ],
508
        ],
509
510
511
512
513
        dim=0,
    )
    concat_batch_f = lambda pt1, pt2: torch.cat(
        [concat_chunk_f(pt1, pt2, i) for i in range(num_sequences)], dim=0
    )
514
515
516
517
518
519

    assert concat_batch_f(X_chunked, remaining_X_chunked).equal(X)
    assert concat_batch_f(dt_chunked, remaining_dt_chunked).equal(dt)
    assert concat_batch_f(B_chunked, remaining_B_chunked).equal(B)
    assert concat_batch_f(C_chunked, remaining_C_chunked).equal(C)

520
    cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
521
522
        compute_varlen_chunk_metadata(remaining_chunked_cu_seqlens, chunk_size)
    )
523
524

    Y_chunked = torch.empty_like(remaining_X_chunked)
525
    state_chunked = mamba_chunk_scan_combined_varlen(
526
527
528
529
530
531
        remaining_X_chunked,
        remaining_dt_chunked,
        A,
        remaining_B_chunked,
        remaining_C_chunked,
        chunk_size,
532
533
534
535
536
        cu_seqlens=remaining_chunked_cu_seqlens.to(torch.int32),
        cu_chunk_seqlens=cu_chunk_seqlens,
        last_chunk_indices=last_chunk_indices,
        seq_idx=seq_idx_chunks,
        out=Y_chunked,
537
538
539
540
541
542
543
        D=None,
        initial_states=partial_state,
    )
    Y = concat_batch_f(Y_partial, Y_chunked)

    # kernel chunked is same as kernel overall
    for i in range(num_sequences):
544
545
        Y_seq = Y[cu_seqlens[i] : cu_seqlens[i + 1], ...]
        Y_ref_seq = Y_ref[cu_seqlens[i] : cu_seqlens[i + 1], ...]
546
        torch.testing.assert_close(
547
548
            Y_seq[: chunked_seqlens[i], ...],
            Y_ref_seq[: chunked_seqlens[i], ...],
549
550
            atol=atol,
            rtol=rtol,
551
552
            msg=lambda x, i=i: f"seq{i} output part1 " + x,
        )
553
        torch.testing.assert_close(
554
555
            Y_seq[chunked_seqlens[i] :, ...],
            Y_ref_seq[chunked_seqlens[i] :, ...],
556
557
            atol=atol,
            rtol=rtol,
558
559
            msg=lambda x, i=i: f"seq{i} output part2 " + x,
        )
560
561
562
563
564
565
566
567

        state_seq = state_chunked[i]
        state_seq_ref = state_ref[i]
        torch.testing.assert_close(
            state_seq,
            state_seq_ref,
            atol=atol,
            rtol=rtol,
568
569
            msg=lambda x, i=i: f"seq{i} state " + x,
        )