| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | 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_varlen, |
| 11 | ) |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.torch_utils import set_random_seed |
| 14 | from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata |
| 15 | |
| 16 | # All kernels exercised here are pure Triton, so they run on any backend |
| 17 | # that the vLLM platform layer treats as a CUDA-alike device or as XPU. |
| 18 | DEVICE = current_platform.device_type |
| 19 | |
| 20 | pytestmark = pytest.mark.skipif( |
| 21 | not (current_platform.is_cuda_alike() or current_platform.is_xpu()), |
| 22 | reason="Mamba2 SSD Triton kernels require a CUDA-alike or XPU device.", |
| 23 | ) |
| 24 | |
| 25 | # Added by the IBM Team, 2024 |
| 26 | |
| 27 | # Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/modules/ssd_minimal.py |
| 28 | |
| 29 | |
| 30 | # this is the segsum implementation taken from above |
| 31 | def segsum(x): |
| 32 | """Calculates segment sum.""" |
| 33 | T = x.size(-1) |
| 34 | x = repeat(x, "... d -> ... d e", e=T) |
| 35 | mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=-1) |
| 36 | x = x.masked_fill(~mask, 0) |
| 37 | x_segsum = torch.cumsum(x, dim=-2) |
| 38 | mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0) |
| 39 | x_segsum = x_segsum.masked_fill(~mask, -torch.inf) |
| 40 | return x_segsum |
| 41 | |
| 42 | |
| 43 | def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None): |
| 44 | """Arguments: |
| 45 | X: (batch, length, n_heads, d_head) |
| 46 | A: (batch, length, n_heads) |
| 47 | B: (batch, length, n_heads, d_state) |
| 48 | C: (batch, length, n_heads, d_state) |
| 49 | |
| 50 | Return: |
| 51 | Y: (batch, length, n_heads, d_head) |
| 52 | |
| 53 | """ |
| 54 | assert X.dtype == A.dtype == B.dtype == C.dtype |
| 55 | assert X.shape[1] % block_len == 0 |
| 56 | |
| 57 | # Rearrange into blocks/chunks |
| 58 | X, A, B, C = ( |
| 59 | rearrange(x, "b (c l) ... -> b c l ...", l=block_len) for x in (X, A, B, C) |
| 60 | ) |
| 61 | |
| 62 | A = rearrange(A, "b c l h -> b h c l") |
| 63 | A_cumsum = torch.cumsum(A, dim=-1) |
| 64 | |
| 65 | # 1. Compute the output for each intra-chunk (diagonal blocks) |
| 66 | L = torch.exp(segsum(A)) |
| 67 | Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X) |
| 68 | |
| 69 | # 2. Compute the state for each intra-chunk |
| 70 | # (right term of low-rank factorization of off-diagonal blocks; B terms) |
| 71 | decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) |
| 72 | states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X) |
| 73 | |
| 74 | # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at |
| 75 | # chunk boundaries |
| 76 | # (middle term of factorization of off-diag blocks; A terms) |
| 77 | if initial_states is None: |
| 78 | initial_states = torch.zeros_like(states[:, :1]) |
| 79 | states = torch.cat([initial_states, states], dim=1) |
| 80 | decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))) |
| 81 | new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states) |
| 82 | states, final_state = new_states[:, :-1], new_states[:, -1] |
| 83 | |
| 84 | # 4. Compute state -> output conversion per chunk |
| 85 | # (left term of low-rank factorization of off-diagonal blocks; C terms) |
| 86 | state_decay_out = torch.exp(A_cumsum) |
| 87 | Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, state_decay_out) |
| 88 | |
| 89 | # Add output of intra-chunk and inter-chunk terms |
| 90 | # (diagonal and off-diagonal blocks) |
| 91 | Y = rearrange(Y_diag + Y_off, "b c l h p -> b (c l) h p") |
| 92 | return Y, final_state |
| 93 | |
| 94 | |
| 95 | def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device=DEVICE): |
| 96 | set_random_seed(0) |
| 97 | A = -torch.exp(torch.rand(n_heads, dtype=itype, device=device)) |
| 98 | dt = F.softplus( |
| 99 | torch.randn(batch_size, seqlen, n_heads, dtype=itype, device=device) - 4 |
| 100 | ) |
| 101 | X = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device) |
| 102 | B = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device) |
| 103 | C = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device) |
| 104 | |
| 105 | return A, dt, X, B, C |
| 106 | |
| 107 | |
| 108 | def generate_continuous_batched_examples( |
| 109 | example_lens_by_batch, |
| 110 | num_examples, |
| 111 | full_length, |
| 112 | last_taken, |
| 113 | exhausted, |
| 114 | n_heads, |
| 115 | d_head, |
| 116 | itype, |
| 117 | device=DEVICE, |
| 118 | return_naive_ref=True, |
| 119 | ): |
| 120 | # this function generates a random examples of certain length |
| 121 | # and then cut according to "example_lens_by_batch" and feed |
| 122 | # them in continuous batches to the kernels. |
| 123 | # If if return_naive_ref=True, the naive torch implementation |
| 124 | # ssd_minimal_discrete will be used to compute and return |
| 125 | # reference output. |
| 126 | |
| 127 | # generate the full-length example |
| 128 | A, dt, X, B, C = generate_random_inputs( |
| 129 | num_examples, full_length, n_heads, d_head, itype |
| 130 | ) |
| 131 | |
| 132 | if return_naive_ref: |
| 133 | Y_min, final_state_min = ssd_minimal_discrete( |
| 134 | X * dt.unsqueeze(-1), A * dt, B, C, block_len=full_length // 4 |
| 135 | ) |
| 136 | |
| 137 | # internal function that outputs a cont batch of examples |
| 138 | # given a tuple of lengths for each example in the batch |
| 139 | # e.g., example_lens=(8, 4) means take 8 samples from first eg, |
| 140 | # 4 examples from second eg, etc |
| 141 | def get_continuous_batch(example_lens: tuple[int, ...]): |
| 142 | indices = [] |
| 143 | for i, x in enumerate(example_lens): |
| 144 | c = last_taken.get(i, 0) |
| 145 | indices.append((c, c + x)) |
| 146 | last_taken[i] = (c + x) % full_length |
| 147 | exhausted[i] = last_taken[i] == 0 |
| 148 | |
| 149 | return ( |
| 150 | torch.concat([x[i, s:e] for i, (s, e) in enumerate(indices)]).unsqueeze(0) |
| 151 | for x in (dt, X, B, C) |
| 152 | ) |
| 153 | |
| 154 | # internal function that maps "n" to the appropriate right boundary |
| 155 | # value when forming continuous batches from examples of length given |
| 156 | # by "full_length". |
| 157 | # - e.g., when n > full_length, returns n % full_length |
| 158 | # when n == full_length, returns full_length |
| 159 | def end_boundary(n: int): |
| 160 | return n - ((n - 1) // full_length) * full_length |
| 161 | |
| 162 | IND_E = None |
| 163 | for spec in example_lens_by_batch: |
| 164 | # get the (maybe partial) example seen in this cont batch |
| 165 | dt2, X2, B2, C2 = get_continuous_batch(spec) |
| 166 | |
| 167 | # get the metadata |
| 168 | cu_seqlens = torch.tensor((0,) + spec, device=device).cumsum(dim=0) |
| 169 | seq_idx = torch.zeros( |
| 170 | cu_seqlens[-1], dtype=torch.int32, device=cu_seqlens.device |
| 171 | ) |
| 172 | for i, (srt, end) in enumerate( |
| 173 | zip( |
| 174 | cu_seqlens, |
| 175 | cu_seqlens[1:], |
| 176 | ) |
| 177 | ): |
| 178 | seq_idx[srt:end] = i |
| 179 | |
| 180 | # for cont batch |
| 181 | if IND_E is None: |
| 182 | IND_S = [0 for _ in range(len(spec))] |
| 183 | else: |
| 184 | IND_S = [x % full_length for x in IND_E] |
| 185 | IND_E = [end_boundary(x + y) for x, y in zip(IND_S, spec)] |
| 186 | |
| 187 | # varlen has implicit batch=1 |
| 188 | dt2 = dt2.squeeze(0) |
| 189 | X2 = X2.squeeze(0) |
| 190 | B2 = B2.squeeze(0) |
| 191 | C2 = C2.squeeze(0) |
| 192 | yield ( |
| 193 | [Y_min[s, IND_S[s] : IND_E[s]] for s in range(num_examples)] |
| 194 | if return_naive_ref |
| 195 | else None, |
| 196 | cu_seqlens, |
| 197 | seq_idx, |
| 198 | (A, dt2, X2, B2, C2), |
| 199 | ) |
| 200 | |
| 201 | |
| 202 | @pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16]) |
| 203 | @pytest.mark.parametrize("n_heads", [4, 16, 32]) |
| 204 | @pytest.mark.parametrize("d_head", [5, 8, 32, 128]) |
| 205 | @pytest.mark.parametrize("seq_len_chunk_size", [(112, 16), (128, 32)]) |
| 206 | def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, itype): |
| 207 | # this tests the kernels on a single example (bs=1) |
| 208 | |
| 209 | # TODO: the bfloat16 case requires higher thresholds. To be investigated |
| 210 | |
| 211 | if itype == torch.bfloat16: |
| 212 | atol, rtol = 5e-2, 5e-2 |
| 213 | else: |
| 214 | atol, rtol = 8e-3, 5e-3 |
| 215 | |
| 216 | # set seed |
| 217 | batch_size = 1 # batch_size |
| 218 | # ssd_minimal_discrete requires chunk_size divide seqlen |
| 219 | # - this is only required for generating the reference seqs, |
| 220 | # it is not an operational limitation. |
| 221 | seqlen, chunk_size = seq_len_chunk_size |
| 222 | |
| 223 | A, dt, X, B, C = generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype) |
| 224 | |
| 225 | Y_min, final_state_min = ssd_minimal_discrete( |
| 226 | X * dt.unsqueeze(-1), A * dt, B, C, chunk_size |
| 227 | ) |
| 228 | |
| 229 | cu_seqlens = torch.tensor((0, seqlen), device=DEVICE).cumsum(dim=0) |
| 230 | cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = ( |
| 231 | compute_varlen_chunk_metadata(cu_seqlens, chunk_size) |
| 232 | ) |
| 233 | # varlen has implicit batch=1 |
| 234 | X = X.squeeze(0) |
| 235 | dt = dt.squeeze(0) |
| 236 | A = A.squeeze(0) |
| 237 | B = B.squeeze(0) |
| 238 | C = C.squeeze(0) |
| 239 | Y = torch.empty_like(X) |
| 240 | final_state = mamba_chunk_scan_combined_varlen( |
| 241 | X, |
| 242 | dt, |
| 243 | A, |
| 244 | B, |
| 245 | C, |
| 246 | chunk_size, |
| 247 | cu_seqlens=cu_seqlens.to(torch.int32), |
| 248 | cu_chunk_seqlens=cu_chunk_seqlens, |
| 249 | last_chunk_indices=last_chunk_indices, |
| 250 | seq_idx=seq_idx_chunks, |
| 251 | out=Y, |
| 252 | D=None, |
| 253 | ) |
| 254 | |
| 255 | # just test the last in sequence |
| 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 |
| 260 | torch.testing.assert_close( |
| 261 | final_state[:, -1].to(torch.float32), |
| 262 | final_state_min[:, -1].to(torch.float32), |
| 263 | atol=atol, |
| 264 | rtol=rtol, |
| 265 | ) |
| 266 | |
| 267 | |
| 268 | @pytest.mark.parametrize("itype", [torch.float32]) |
| 269 | @pytest.mark.parametrize("n_heads", [4, 8]) |
| 270 | @pytest.mark.parametrize("d_head", [5, 16, 32]) |
| 271 | @pytest.mark.parametrize( |
| 272 | "seq_len_chunk_size_cases", |
| 273 | [ |
| 274 | # small-ish chunk_size (8) |
| 275 | (64, 8, 2, [(64, 32), (64, 32)]), |
| 276 | (64, 8, 2, [(8, 8), (8, 8), (8, 8)]), # chunk size boundary |
| 277 | ( |
| 278 | 64, |
| 279 | 8, |
| 280 | 2, |
| 281 | [(4, 4), (4, 4), (4, 4), (4, 4)], |
| 282 | ), # chunk_size larger than cont batches |
| 283 | (64, 8, 5, [(64, 32, 16, 8, 8)]), |
| 284 | # large-ish chunk_size (256) |
| 285 | (64, 256, 1, [(5,), (1,), (1,), (1,)]), # irregular sizes with small sequences |
| 286 | ( |
| 287 | 64, |
| 288 | 256, |
| 289 | 2, |
| 290 | [(5, 30), (1, 2), (1, 2), (1, 2)], |
| 291 | ), # irregular sizes with small sequences |
| 292 | # we also need to test some large seqlen |
| 293 | # to catch errors with init states decay |
| 294 | (768, 128, 2, [(138, 225), (138, 225)]), |
| 295 | ], |
| 296 | ) |
| 297 | def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, itype): |
| 298 | # this test with multiple examples in a continuous batch |
| 299 | # (i.e. chunked prefill) |
| 300 | |
| 301 | seqlen, chunk_size, num_examples, cases = seq_len_chunk_size_cases |
| 302 | |
| 303 | # This test can have larger error for longer sequences |
| 304 | if seqlen > 256: |
| 305 | atol, rtol = 1e-2, 5e-3 |
| 306 | else: |
| 307 | atol, rtol = 5e-3, 5e-3 |
| 308 | |
| 309 | # hold state during the cutting process so we know if an |
| 310 | # example has been exhausted and needs to cycle |
| 311 | last_taken: dict = {} # map: eg -> pointer to last taken sample |
| 312 | exhausted: dict = {} # map: eg -> boolean indicating example is exhausted |
| 313 | |
| 314 | states = None |
| 315 | for Y_min, cu_seqlens, _token_seq_idx, ( |
| 316 | A, |
| 317 | dt, |
| 318 | X, |
| 319 | B, |
| 320 | C, |
| 321 | ) in generate_continuous_batched_examples( |
| 322 | cases, num_examples, seqlen, last_taken, exhausted, n_heads, d_head, itype |
| 323 | ): |
| 324 | cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = ( |
| 325 | compute_varlen_chunk_metadata(cu_seqlens, chunk_size) |
| 326 | ) |
| 327 | |
| 328 | Y = torch.empty_like(X) |
| 329 | new_states = mamba_chunk_scan_combined_varlen( |
| 330 | X, |
| 331 | dt, |
| 332 | A, |
| 333 | B, |
| 334 | C, |
| 335 | chunk_size, |
| 336 | cu_seqlens=cu_seqlens.to(torch.int32), |
| 337 | cu_chunk_seqlens=cu_chunk_seqlens, |
| 338 | last_chunk_indices=last_chunk_indices, |
| 339 | seq_idx=seq_idx_chunks, |
| 340 | out=Y, |
| 341 | D=None, |
| 342 | initial_states=states, |
| 343 | ) |
| 344 | |
| 345 | # just test the last in sequence |
| 346 | for i in range(num_examples): |
| 347 | # just test one dim and dstate |
| 348 | Y_eg = Y[cu_seqlens[i] : cu_seqlens[i + 1], 0, 0] |
| 349 | Y_min_eg = Y_min[i][:, 0, 0] |
| 350 | torch.testing.assert_close(Y_eg, Y_min_eg, atol=atol, rtol=rtol) |
| 351 | |
| 352 | # update states |
| 353 | states = new_states |
| 354 | for i, clear in exhausted.items(): |
| 355 | if clear: |
| 356 | states[i].fill_(0.0) |
| 357 | exhausted[i] = False |
| 358 | |
| 359 | |
| 360 | @pytest.mark.parametrize("chunk_size", [8, 256]) |
| 361 | @pytest.mark.parametrize( |
| 362 | "seqlens", |
| 363 | [(16, 20), (270, 88, 212, 203)], |
| 364 | ) |
| 365 | def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens): |
| 366 | # This test verifies the correctness of the chunked prefill implementation |
| 367 | # in the mamba2 ssd kernels, by comparing concatenation (in the sequence |
| 368 | # dimension) of chunked results with the full sequence result. |
| 369 | # It is different from test_mamba_chunk_scan_cont_batch by: |
| 370 | # 1. Not using the naive torch implementation (ssd_minimal_discrete) to get |
| 371 | # reference outputs. Instead, it compares chunked kernel outputs to full |
| 372 | # sequence kernel outputs. This is the most straightforward way to |
| 373 | # assert chunked prefill correctness. |
| 374 | # 2. It focuses on cases where sequences change in the middle of mamba |
| 375 | # chunks, and not necessarily on chunk boundaries. |
| 376 | |
| 377 | max_seqlen = max(seqlens) |
| 378 | # This test can have larger error for longer sequences |
| 379 | if max_seqlen > 256: |
| 380 | atol, rtol = 1e-2, 5e-3 |
| 381 | else: |
| 382 | atol, rtol = 5e-3, 5e-3 |
| 383 | |
| 384 | num_sequences = len(seqlens) |
| 385 | n_heads = 16 |
| 386 | d_head = 64 |
| 387 | itype = torch.float32 |
| 388 | |
| 389 | # hold state during the cutting process so we know if an |
| 390 | # example has been exhausted and needs to cycle |
| 391 | last_taken: dict = {} # map: eg -> pointer to last taken sample |
| 392 | exhausted: dict = {} # map: eg -> boolean indicating example is exhausted |
| 393 | _, cu_seqlens, seq_idx, (A, dt, X, B, C) = next( |
| 394 | generate_continuous_batched_examples( |
| 395 | [seqlens], |
| 396 | num_sequences, |
| 397 | max_seqlen, |
| 398 | last_taken, |
| 399 | exhausted, |
| 400 | n_heads, |
| 401 | d_head, |
| 402 | itype, |
| 403 | return_naive_ref=False, |
| 404 | ) |
| 405 | ) |
| 406 | seqlens = torch.tensor(seqlens, dtype=torch.int32, device=X.device) |
| 407 | device = X.device |
| 408 | |
| 409 | ## full seqlen computation |
| 410 | cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = ( |
| 411 | compute_varlen_chunk_metadata(cu_seqlens, chunk_size) |
| 412 | ) |
| 413 | Y_ref = torch.empty_like(X) |
| 414 | state_ref = mamba_chunk_scan_combined_varlen( |
| 415 | X, |
| 416 | dt, |
| 417 | A, |
| 418 | B, |
| 419 | C, |
| 420 | chunk_size, |
| 421 | cu_seqlens=cu_seqlens.to(torch.int32), |
| 422 | cu_chunk_seqlens=cu_chunk_seqlens, |
| 423 | last_chunk_indices=last_chunk_indices, |
| 424 | seq_idx=seq_idx_chunks, |
| 425 | out=Y_ref, |
| 426 | D=None, |
| 427 | initial_states=None, |
| 428 | ) |
| 429 | |
| 430 | ## chunked seqlen computation |
| 431 | # first chunk |
| 432 | chunked_seqlens = seqlens // 2 |
| 433 | chunked_cu_seqlens = torch.cat( |
| 434 | [torch.tensor([0], device=device), torch.cumsum(chunked_seqlens, dim=0)], dim=0 |
| 435 | ) |
| 436 | chunked_input_seq_len = chunked_cu_seqlens[-1] |
| 437 | X_chunked = torch.zeros_like(X)[:chunked_input_seq_len, ...] |
| 438 | dt_chunked = torch.zeros_like(dt)[:chunked_input_seq_len, ...] |
| 439 | B_chunked = torch.zeros_like(B)[:chunked_input_seq_len, ...] |
| 440 | C_chunked = torch.zeros_like(C)[:chunked_input_seq_len, ...] |
| 441 | for i in range(num_sequences): |
| 442 | chunk_f = lambda x, i: x[ |
| 443 | cu_seqlens[i] : cu_seqlens[i] + chunked_seqlens[i], ... |
| 444 | ] |
| 445 | |
| 446 | X_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f( |
| 447 | X, i |
| 448 | ) |
| 449 | dt_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f( |
| 450 | dt, i |
| 451 | ) |
| 452 | B_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f( |
| 453 | B, i |
| 454 | ) |
| 455 | C_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f( |
| 456 | C, i |
| 457 | ) |
| 458 | |
| 459 | cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = ( |
| 460 | compute_varlen_chunk_metadata(chunked_cu_seqlens, chunk_size) |
| 461 | ) |
| 462 | Y_partial = torch.empty_like(X_chunked) |
| 463 | partial_state = mamba_chunk_scan_combined_varlen( |
| 464 | X_chunked, |
| 465 | dt_chunked, |
| 466 | A, |
| 467 | B_chunked, |
| 468 | C_chunked, |
| 469 | chunk_size, |
| 470 | cu_seqlens=chunked_cu_seqlens.to(torch.int32), |
| 471 | cu_chunk_seqlens=cu_chunk_seqlens, |
| 472 | last_chunk_indices=last_chunk_indices, |
| 473 | seq_idx=seq_idx_chunks, |
| 474 | out=Y_partial, |
| 475 | D=None, |
| 476 | initial_states=None, |
| 477 | ) |
| 478 | |
| 479 | # remaining chunk |
| 480 | remaining_chunked_seqlens = seqlens - chunked_seqlens |
| 481 | remaining_chunked_cu_seqlens = torch.cat( |
| 482 | [ |
| 483 | torch.tensor([0], device=device), |
| 484 | torch.cumsum(remaining_chunked_seqlens, dim=0), |
| 485 | ], |
| 486 | dim=0, |
| 487 | ) |
| 488 | remaining_chunked_input_seq_len = remaining_chunked_cu_seqlens[-1] |
| 489 | remaining_X_chunked = torch.zeros_like(X)[:remaining_chunked_input_seq_len, ...] |
| 490 | remaining_dt_chunked = torch.zeros_like(dt)[:remaining_chunked_input_seq_len, ...] |
| 491 | remaining_B_chunked = torch.zeros_like(B)[:remaining_chunked_input_seq_len, ...] |
| 492 | remaining_C_chunked = torch.zeros_like(C)[:remaining_chunked_input_seq_len, ...] |
| 493 | for i in range(num_sequences): |
| 494 | remaining_chunk_f = lambda x, i: x[ |
| 495 | cu_seqlens[i] + chunked_seqlens[i] : cu_seqlens[i + 1], ... |
| 496 | ] |
| 497 | |
| 498 | remaining_X_chunked[ |
| 499 | remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ... |
| 500 | ] = remaining_chunk_f(X, i) |
| 501 | remaining_dt_chunked[ |
| 502 | remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ... |
| 503 | ] = remaining_chunk_f(dt, i) |
| 504 | remaining_B_chunked[ |
| 505 | remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ... |
| 506 | ] = remaining_chunk_f(B, i) |
| 507 | remaining_C_chunked[ |
| 508 | remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ... |
| 509 | ] = remaining_chunk_f(C, i) |
| 510 | |
| 511 | # assert input chunking is correct |
| 512 | concat_chunk_f = lambda pt1, pt2, i: torch.cat( |
| 513 | [ |
| 514 | pt1[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...], |
| 515 | pt2[ |
| 516 | remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], |
| 517 | ..., |
| 518 | ], |
| 519 | ], |
| 520 | dim=0, |
| 521 | ) |
| 522 | concat_batch_f = lambda pt1, pt2: torch.cat( |
| 523 | [concat_chunk_f(pt1, pt2, i) for i in range(num_sequences)], dim=0 |
| 524 | ) |
| 525 | |
| 526 | assert concat_batch_f(X_chunked, remaining_X_chunked).equal(X) |
| 527 | assert concat_batch_f(dt_chunked, remaining_dt_chunked).equal(dt) |
| 528 | assert concat_batch_f(B_chunked, remaining_B_chunked).equal(B) |
| 529 | assert concat_batch_f(C_chunked, remaining_C_chunked).equal(C) |
| 530 | |
| 531 | cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = ( |
| 532 | compute_varlen_chunk_metadata(remaining_chunked_cu_seqlens, chunk_size) |
| 533 | ) |
| 534 | |
| 535 | Y_chunked = torch.empty_like(remaining_X_chunked) |
| 536 | state_chunked = mamba_chunk_scan_combined_varlen( |
| 537 | remaining_X_chunked, |
| 538 | remaining_dt_chunked, |
| 539 | A, |
| 540 | remaining_B_chunked, |
| 541 | remaining_C_chunked, |
| 542 | chunk_size, |
| 543 | cu_seqlens=remaining_chunked_cu_seqlens.to(torch.int32), |
| 544 | cu_chunk_seqlens=cu_chunk_seqlens, |
| 545 | last_chunk_indices=last_chunk_indices, |
| 546 | seq_idx=seq_idx_chunks, |
| 547 | out=Y_chunked, |
| 548 | D=None, |
| 549 | initial_states=partial_state, |
| 550 | ) |
| 551 | Y = concat_batch_f(Y_partial, Y_chunked) |
| 552 | |
| 553 | # kernel chunked is same as kernel overall |
| 554 | for i in range(num_sequences): |
| 555 | Y_seq = Y[cu_seqlens[i] : cu_seqlens[i + 1], ...] |
| 556 | Y_ref_seq = Y_ref[cu_seqlens[i] : cu_seqlens[i + 1], ...] |
| 557 | torch.testing.assert_close( |
| 558 | Y_seq[: chunked_seqlens[i], ...], |
| 559 | Y_ref_seq[: chunked_seqlens[i], ...], |
| 560 | atol=atol, |
| 561 | rtol=rtol, |
| 562 | msg=lambda x, i=i: f"seq{i} output part1 " + x, |
| 563 | ) |
| 564 | torch.testing.assert_close( |
| 565 | Y_seq[chunked_seqlens[i] :, ...], |
| 566 | Y_ref_seq[chunked_seqlens[i] :, ...], |
| 567 | atol=atol, |
| 568 | rtol=rtol, |
| 569 | msg=lambda x, i=i: f"seq{i} output part2 " + x, |
| 570 | ) |
| 571 | |
| 572 | state_seq = state_chunked[i] |
| 573 | state_seq_ref = state_ref[i] |
| 574 | torch.testing.assert_close( |
| 575 | state_seq, |
| 576 | state_seq_ref, |
| 577 | atol=atol, |
| 578 | rtol=rtol, |
| 579 | msg=lambda x, i=i: f"seq{i} state " + x, |
| 580 | ) |