| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright (c) 2024, Tri Dao, Albert Gu. |
| 5 | # Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/ssd_combined.py |
| 6 | |
| 7 | # ruff: noqa: E501 |
| 8 | |
| 9 | import torch |
| 10 | from einops import rearrange |
| 11 | from packaging import version |
| 12 | |
| 13 | from vllm.triton_utils import triton |
| 14 | |
| 15 | from .ssd_bmm import _bmm_chunk_fwd |
| 16 | from .ssd_chunk_scan import _chunk_scan_fwd |
| 17 | from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd |
| 18 | from .ssd_state_passing import _state_passing_fwd |
| 19 | |
| 20 | TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0") |
| 21 | |
| 22 | |
| 23 | def is_int_pow_2(n): |
| 24 | return isinstance(n, int) and n > 0 and (n & (n - 1)) == 0 |
| 25 | |
| 26 | |
| 27 | def _mamba_chunk_scan_combined_fwd( |
| 28 | x, |
| 29 | dt, |
| 30 | A, |
| 31 | B, |
| 32 | C, |
| 33 | chunk_size, |
| 34 | out, |
| 35 | D=None, |
| 36 | z=None, |
| 37 | dt_bias=None, |
| 38 | initial_states=None, |
| 39 | return_intermediate_states=False, |
| 40 | seq_idx=None, |
| 41 | cu_seqlens=None, |
| 42 | cu_chunk_seqlens=None, |
| 43 | last_chunk_indices=None, |
| 44 | dt_softplus=False, |
| 45 | dt_limit=(0.0, float("inf")), |
| 46 | state_dtype=None, |
| 47 | ): |
| 48 | assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2" |
| 49 | seqlen, nheads, headdim = x.shape |
| 50 | _, ngroups, dstate = B.shape |
| 51 | assert nheads % ngroups == 0 |
| 52 | assert B.shape == (seqlen, ngroups, dstate) |
| 53 | assert dt.shape == (seqlen, nheads) |
| 54 | assert A.shape == (nheads,) |
| 55 | assert C.shape == B.shape |
| 56 | if z is not None: |
| 57 | assert z.shape == x.shape |
| 58 | if D is not None: |
| 59 | assert D.shape == (nheads, headdim) or D.shape == (nheads,) |
| 60 | if seq_idx is not None: |
| 61 | assert seq_idx.shape == (cu_chunk_seqlens.shape[0] - 1,) |
| 62 | if B.stride(-1) != 1: |
| 63 | B = B.contiguous() |
| 64 | if C.stride(-1) != 1: |
| 65 | C = C.contiguous() |
| 66 | if ( |
| 67 | x.stride(-1) != 1 and x.stride(0) != 1 |
| 68 | ): # Either M or K dimension should be contiguous |
| 69 | x = x.contiguous() |
| 70 | if ( |
| 71 | z is not None and z.stride(-1) != 1 and z.stride(0) != 1 |
| 72 | ): # Either M or K dimension should be contiguous |
| 73 | z = z.contiguous() |
| 74 | if D is not None and D.stride(-1) != 1: |
| 75 | D = D.contiguous() |
| 76 | assert cu_seqlens is not None, "Assuming varlen input - must supply cu_seqlens" |
| 77 | |
| 78 | if initial_states is not None: |
| 79 | assert initial_states.shape == (len(cu_seqlens) - 1, nheads, headdim, dstate) |
| 80 | |
| 81 | # This function executes 5 sub-functions for computing mamba |
| 82 | # - a good resource is the blog https://goombalab.github.io/blog/2024/mamba2-part3-algorithm/ |
| 83 | # which has a minimal implementation to understand the below operations |
| 84 | # - as explained by the blog, mamba is a special case of causal attention |
| 85 | # - the idea is to chunk the attention matrix and compute each |
| 86 | # submatrix separately using different optimizations. |
| 87 | # - see the blog and paper for a visualization of the submatrices |
| 88 | # which we refer to in the comments below |
| 89 | |
| 90 | # 1. Compute chunked cumsum of A * dt |
| 91 | # - here dt may go through a softplus activation |
| 92 | dA_cumsum, dt = _chunk_cumsum_fwd( |
| 93 | dt, |
| 94 | A, |
| 95 | chunk_size, |
| 96 | cu_chunk_seqlens, |
| 97 | dt_bias=dt_bias, |
| 98 | dt_softplus=dt_softplus, |
| 99 | dt_limit=dt_limit, |
| 100 | ) |
| 101 | |
| 102 | # 2. Compute the state for each intra-chunk |
| 103 | # (right term of low-rank factorization of off-diagonal blocks; B terms) |
| 104 | states = _chunk_state_fwd( |
| 105 | B, x, dt, dA_cumsum, cu_chunk_seqlens, states_in_fp32=True |
| 106 | ) |
| 107 | |
| 108 | # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries |
| 109 | # (middle term of factorization of off-diag blocks; A terms) |
| 110 | # - parallelized across sequences using last_chunk_indices to derive |
| 111 | # per-sequence chunk ranges. Each sequence's state passing runs independently. |
| 112 | states = _state_passing_fwd( |
| 113 | rearrange(states, "... p n -> ... (p n)"), |
| 114 | dA_cumsum, # (nheads, nchunks, chunk_size) |
| 115 | last_chunk_indices, |
| 116 | initial_states=rearrange(initial_states, "... p n -> ... (p n)") |
| 117 | if initial_states is not None |
| 118 | else None, # (batch, nheads, headdim*dstate) |
| 119 | out_dtype=state_dtype if state_dtype is not None else C.dtype, |
| 120 | ) |
| 121 | states = rearrange(states, "... (p n) -> ... p n", n=dstate) |
| 122 | |
| 123 | # 4. Compute batched matrix multiply for C_j^T B_i terms |
| 124 | CB = _bmm_chunk_fwd(C, B, chunk_size, cu_chunk_seqlens, output_dtype=torch.float32) |
| 125 | |
| 126 | # 5. Scan and compute the diagonal blocks, taking into |
| 127 | # account past causal states. |
| 128 | # - if initial states are provided, then states information will be |
| 129 | # augmented with initial_states. |
| 130 | # - to do this properly, we need to account for example changes in |
| 131 | # the continuous batch, therefore we introduce pseudo chunks, which is |
| 132 | # a chunk that is split up each time an example changes. |
| 133 | # - in each (pseudo) chunk, we detect if the previous (pseudo) chunk had |
| 134 | # a seq_idx change, in which case we take states information from |
| 135 | # init_states. |
| 136 | _chunk_scan_fwd( |
| 137 | CB, |
| 138 | x, |
| 139 | dt, |
| 140 | dA_cumsum, |
| 141 | C, |
| 142 | states, |
| 143 | cu_chunk_seqlens, |
| 144 | out, # in-place update |
| 145 | seq_idx, |
| 146 | D=D, |
| 147 | z=z, |
| 148 | initial_states=initial_states, |
| 149 | ) |
| 150 | |
| 151 | if return_intermediate_states: |
| 152 | return states |
| 153 | else: |
| 154 | return states[last_chunk_indices] |
| 155 | |
| 156 | |
| 157 | def mamba_chunk_scan_combined_varlen( |
| 158 | x, |
| 159 | dt, |
| 160 | A, |
| 161 | B, |
| 162 | C, |
| 163 | chunk_size, |
| 164 | cu_seqlens, |
| 165 | cu_chunk_seqlens, |
| 166 | last_chunk_indices, |
| 167 | seq_idx, |
| 168 | out, |
| 169 | D=None, |
| 170 | z=None, |
| 171 | dt_bias=None, |
| 172 | initial_states=None, |
| 173 | dt_softplus=False, |
| 174 | dt_limit=(0.0, float("inf")), |
| 175 | return_intermediate_states=False, |
| 176 | state_dtype=None, |
| 177 | ): |
| 178 | """Argument: |
| 179 | x: (seqlen, nheads, headdim) |
| 180 | dt: (seqlen, nheads) |
| 181 | A: (nheads) |
| 182 | B: (seqlen, ngroups, dstate) |
| 183 | C: (seqlen, ngroups, dstate) |
| 184 | chunk_size: int |
| 185 | cu_seqlens: (batch + 1,) |
| 186 | cu_chunk_seqlens: (nchunks + 1,) |
| 187 | last_chunk_indices: (batch,) |
| 188 | seq_idx: (nchunks,) |
| 189 | out: (seqlen, nheads, headdim) preallocated output tensor |
| 190 | D: (nheads, headdim) or (nheads,) |
| 191 | z: (seqlen, nheads, headdim) |
| 192 | dt_bias: (nheads,) |
| 193 | initial_states: (batch, nheads, headdim, dstate) |
| 194 | dt_softplus: Whether to apply softplus to dt |
| 195 | out: (seqlen, nheads, headdim) preallocated output tensor |
| 196 | state_dtype: The data type of the ssm state |
| 197 | Return: |
| 198 | varlen_states: (batch, nheads, headdim, dstate) |
| 199 | """ |
| 200 | assert cu_seqlens is not None, "cu_seqlens must be provided assuming varlen input" |
| 201 | assert seq_idx is not None |
| 202 | |
| 203 | varlen_states = _mamba_chunk_scan_combined_fwd( |
| 204 | x, |
| 205 | dt, |
| 206 | A, |
| 207 | B, |
| 208 | C, |
| 209 | chunk_size, |
| 210 | out, |
| 211 | D=D, |
| 212 | z=z, |
| 213 | dt_bias=dt_bias, |
| 214 | initial_states=initial_states, |
| 215 | return_intermediate_states=return_intermediate_states, |
| 216 | seq_idx=seq_idx, |
| 217 | cu_seqlens=cu_seqlens, |
| 218 | cu_chunk_seqlens=cu_chunk_seqlens, |
| 219 | last_chunk_indices=last_chunk_indices, |
| 220 | dt_softplus=dt_softplus, |
| 221 | dt_limit=dt_limit, |
| 222 | state_dtype=state_dtype, |
| 223 | ) |
| 224 | |
| 225 | return varlen_states |
| 226 | |
| 227 | |
| 228 | from vllm.platforms import current_platform # noqa: E402 |
| 229 | |
| 230 | if current_platform.is_cpu(): |
| 231 | import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_mamba_ssm |
| 232 | |
| 233 | _mamba_chunk_scan_combined_fwd = cpu_mamba_ssm._mamba_chunk_scan_combined_fwd_cpu # type: ignore |