| 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_chunk_state.py |
| 6 | |
| 7 | # ruff: noqa: E501 |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp |
| 12 | from vllm.triton_utils import tl, triton |
| 13 | |
| 14 | from .mamba_ssm import softplus |
| 15 | |
| 16 | |
| 17 | @triton.autotune( |
| 18 | configs=[ |
| 19 | triton.Config({"BLOCK_SIZE_H": 2}), |
| 20 | triton.Config({"BLOCK_SIZE_H": 4}), |
| 21 | triton.Config({"BLOCK_SIZE_H": 8}), |
| 22 | triton.Config({"BLOCK_SIZE_H": 16}), |
| 23 | triton.Config({"BLOCK_SIZE_H": 32}), |
| 24 | triton.Config({"BLOCK_SIZE_H": 64}), |
| 25 | ], |
| 26 | key=["chunk_size", "nheads"], |
| 27 | ) |
| 28 | @triton.jit |
| 29 | def _chunk_cumsum_fwd_kernel( |
| 30 | # Pointers to matrices |
| 31 | dt_ptr, |
| 32 | A_ptr, |
| 33 | dt_bias_ptr, |
| 34 | dt_out_ptr, |
| 35 | dA_cumsum_ptr, |
| 36 | cu_chunk_seqlens_ptr, |
| 37 | # Matrix dimension |
| 38 | nheads: tl.constexpr, |
| 39 | chunk_size: tl.constexpr, |
| 40 | dt_min: tl.constexpr, |
| 41 | dt_max: tl.constexpr, |
| 42 | # Strides |
| 43 | stride_dt_seqlen: tl.int64, |
| 44 | stride_dt_head: tl.constexpr, |
| 45 | stride_A_head: tl.constexpr, |
| 46 | stride_dt_bias_head: tl.constexpr, |
| 47 | stride_dt_out_head: tl.int64, |
| 48 | stride_dt_out_chunk: tl.int64, |
| 49 | stride_dt_out_csize: tl.constexpr, |
| 50 | stride_dA_cs_head: tl.int64, |
| 51 | stride_dA_cs_chunk: tl.int64, |
| 52 | stride_dA_cs_csize: tl.constexpr, |
| 53 | # Meta-parameters |
| 54 | DT_SOFTPLUS: tl.constexpr, |
| 55 | HAS_DT_BIAS: tl.constexpr, |
| 56 | BLOCK_SIZE_H: tl.constexpr, |
| 57 | BLOCK_SIZE_CHUNK: tl.constexpr, |
| 58 | ): |
| 59 | # if dt is long, may cause problems, so use 64 bit |
| 60 | # https://github.com/triton-lang/triton/issues/1058 |
| 61 | pid_c = tl.program_id(axis=0).to(tl.int64) |
| 62 | pid_h = tl.program_id(axis=1) |
| 63 | |
| 64 | chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) |
| 65 | chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) |
| 66 | |
| 67 | dt_ptr += chunk_seqlen_start * stride_dt_seqlen |
| 68 | dt_out_ptr += pid_c * stride_dt_out_chunk |
| 69 | dA_cumsum_ptr += pid_c * stride_dA_cs_chunk |
| 70 | |
| 71 | offs_h = pid_h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H) |
| 72 | offs_c = tl.arange(0, BLOCK_SIZE_CHUNK) |
| 73 | dt_ptrs = dt_ptr + ( |
| 74 | offs_h[:, None] * stride_dt_head + offs_c[None, :] * stride_dt_seqlen |
| 75 | ) |
| 76 | A_ptrs = A_ptr + offs_h * stride_A_head |
| 77 | dt_out_ptrs = dt_out_ptr + ( |
| 78 | offs_h[:, None] * stride_dt_out_head + offs_c[None, :] * stride_dt_out_csize |
| 79 | ) |
| 80 | dA_cs_ptrs = dA_cumsum_ptr + ( |
| 81 | offs_h[:, None] * stride_dA_cs_head + offs_c[None, :] * stride_dA_cs_csize |
| 82 | ) |
| 83 | chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start |
| 84 | |
| 85 | dt = tl.load( |
| 86 | dt_ptrs, |
| 87 | mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size_limit), |
| 88 | other=0.0, |
| 89 | ).to(tl.float32) |
| 90 | if HAS_DT_BIAS: |
| 91 | dt_bias = tl.load( |
| 92 | dt_bias_ptr + offs_h * stride_dt_bias_head, mask=offs_h < nheads, other=0.0 |
| 93 | ).to(tl.float32) |
| 94 | dt += dt_bias[:, None] |
| 95 | if DT_SOFTPLUS: |
| 96 | dt = tl.where(dt <= 20.0, softplus(dt), dt) |
| 97 | |
| 98 | dt = tl.clamp(dt, dt_min, dt_max) |
| 99 | dt = tl.where( |
| 100 | (offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size_limit), dt, 0.0 |
| 101 | ) |
| 102 | tl.store( |
| 103 | dt_out_ptrs, |
| 104 | dt, |
| 105 | mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size), |
| 106 | ) |
| 107 | A = tl.load(A_ptrs, mask=offs_h < nheads, other=0.0).to(tl.float32) |
| 108 | dA = dt * A[:, None] |
| 109 | dA_cs = tl.cumsum(dA, axis=1) |
| 110 | tl.store( |
| 111 | dA_cs_ptrs, |
| 112 | dA_cs, |
| 113 | mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size), |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | @triton.autotune( |
| 118 | configs=[ |
| 119 | # Small headdim/dstate configs (hdim<=64, dstate<=128) - increased parallelism |
| 120 | triton.Config( |
| 121 | {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, |
| 122 | num_stages=3, |
| 123 | num_warps=4, |
| 124 | ), |
| 125 | triton.Config( |
| 126 | {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, |
| 127 | num_stages=3, |
| 128 | num_warps=4, |
| 129 | ), |
| 130 | triton.Config( |
| 131 | {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, |
| 132 | num_stages=3, |
| 133 | num_warps=4, |
| 134 | ), |
| 135 | # Low register pressure configs for large dstate (dstate=128) |
| 136 | triton.Config( |
| 137 | {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, |
| 138 | num_stages=2, |
| 139 | num_warps=4, |
| 140 | ), |
| 141 | triton.Config( |
| 142 | {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64}, |
| 143 | num_stages=2, |
| 144 | num_warps=4, |
| 145 | ), |
| 146 | # original configs for larger headdim/dstate values |
| 147 | triton.Config( |
| 148 | {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64}, |
| 149 | num_stages=3, |
| 150 | num_warps=8, |
| 151 | ), |
| 152 | triton.Config( |
| 153 | {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32}, |
| 154 | num_stages=4, |
| 155 | num_warps=4, |
| 156 | ), |
| 157 | triton.Config( |
| 158 | {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, |
| 159 | num_stages=4, |
| 160 | num_warps=4, |
| 161 | ), |
| 162 | triton.Config( |
| 163 | {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, |
| 164 | num_stages=4, |
| 165 | num_warps=4, |
| 166 | ), |
| 167 | triton.Config( |
| 168 | {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, |
| 169 | num_stages=4, |
| 170 | num_warps=4, |
| 171 | ), |
| 172 | triton.Config( |
| 173 | {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, |
| 174 | num_stages=4, |
| 175 | num_warps=4, |
| 176 | ), |
| 177 | triton.Config( |
| 178 | {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, |
| 179 | num_stages=5, |
| 180 | num_warps=2, |
| 181 | ), |
| 182 | triton.Config( |
| 183 | {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, |
| 184 | num_stages=5, |
| 185 | num_warps=2, |
| 186 | ), |
| 187 | triton.Config( |
| 188 | {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, |
| 189 | num_stages=4, |
| 190 | num_warps=2, |
| 191 | ), |
| 192 | ], |
| 193 | key=["hdim", "dstate", "chunk_size"], |
| 194 | ) |
| 195 | @triton.jit |
| 196 | def _chunk_state_fwd_kernel( |
| 197 | # Pointers to matrices |
| 198 | x_ptr, |
| 199 | b_ptr, |
| 200 | states_ptr, |
| 201 | dt_ptr, |
| 202 | dA_cumsum_ptr, |
| 203 | cu_chunk_seqlens_ptr, |
| 204 | # Matrix dimensions |
| 205 | hdim: tl.constexpr, |
| 206 | dstate: tl.constexpr, |
| 207 | chunk_size: tl.constexpr, |
| 208 | nheads_ngroups_ratio: tl.constexpr, |
| 209 | # Strides |
| 210 | stride_x_seqlen: tl.int64, |
| 211 | stride_x_head: tl.int64, |
| 212 | stride_x_hdim: tl.constexpr, |
| 213 | stride_b_seqlen: tl.int64, |
| 214 | stride_b_head: tl.int64, |
| 215 | stride_b_dstate: tl.constexpr, |
| 216 | stride_states_chunk: tl.int64, |
| 217 | stride_states_head: tl.int64, |
| 218 | stride_states_hdim: tl.int64, |
| 219 | stride_states_dstate: tl.constexpr, |
| 220 | stride_dt_head: tl.int64, |
| 221 | stride_dt_chunk: tl.int64, |
| 222 | stride_dt_csize: tl.constexpr, |
| 223 | stride_dA_cs_head: tl.int64, |
| 224 | stride_dA_cs_chunk: tl.int64, |
| 225 | stride_dA_cs_csize: tl.constexpr, |
| 226 | # Meta-parameters |
| 227 | BLOCK_SIZE_M: tl.constexpr, |
| 228 | BLOCK_SIZE_N: tl.constexpr, |
| 229 | BLOCK_SIZE_K: tl.constexpr, |
| 230 | ): |
| 231 | pid_c = tl.program_id(axis=1).to(tl.int64) |
| 232 | pid_h = tl.program_id(axis=2) |
| 233 | num_pid_n = tl.cdiv(dstate, BLOCK_SIZE_N) |
| 234 | pid_m = tl.program_id(axis=0) // num_pid_n |
| 235 | pid_n = tl.program_id(axis=0) % num_pid_n |
| 236 | chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) |
| 237 | chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) |
| 238 | b_ptr += ( |
| 239 | chunk_seqlen_start * stride_b_seqlen |
| 240 | + (pid_h // nheads_ngroups_ratio) * stride_b_head |
| 241 | ) |
| 242 | x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head |
| 243 | dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head |
| 244 | dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head |
| 245 | |
| 246 | offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) |
| 247 | offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) |
| 248 | offs_k = tl.arange(0, BLOCK_SIZE_K) |
| 249 | x_ptrs = x_ptr + ( |
| 250 | offs_m[:, None] * stride_x_hdim + offs_k[None, :] * stride_x_seqlen |
| 251 | ) |
| 252 | b_ptrs = b_ptr + ( |
| 253 | offs_n[None, :] * stride_b_dstate + offs_k[:, None] * stride_b_seqlen |
| 254 | ) |
| 255 | dt_ptrs = dt_ptr + offs_k * stride_dt_csize |
| 256 | dA_cs_last = tl.load(dA_cumsum_ptr + (chunk_size - 1) * stride_dA_cs_csize).to( |
| 257 | tl.float32 |
| 258 | ) |
| 259 | dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize |
| 260 | |
| 261 | chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start |
| 262 | |
| 263 | acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) |
| 264 | for k in range(0, chunk_size_limit, BLOCK_SIZE_K): |
| 265 | x = tl.load( |
| 266 | x_ptrs, |
| 267 | mask=(offs_m[:, None] < hdim) & (offs_k[None, :] < chunk_size_limit - k), |
| 268 | other=0.0, |
| 269 | ) |
| 270 | b = tl.load( |
| 271 | b_ptrs, |
| 272 | mask=(offs_k[:, None] < chunk_size_limit - k) & (offs_n[None, :] < dstate), |
| 273 | other=0.0, |
| 274 | ).to(tl.float32) |
| 275 | dA_cs_k = tl.load( |
| 276 | dA_cumsum_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0 |
| 277 | ).to(tl.float32) |
| 278 | dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0).to( |
| 279 | tl.float32 |
| 280 | ) |
| 281 | scale = fast_exp(tl.minimum(dA_cs_last - dA_cs_k, 0.0)) * dt_k |
| 282 | b *= scale[:, None] |
| 283 | b = b.to(x_ptr.dtype.element_ty) |
| 284 | acc += tl.dot(x, b) |
| 285 | |
| 286 | x_ptrs += BLOCK_SIZE_K * stride_x_seqlen |
| 287 | b_ptrs += BLOCK_SIZE_K * stride_b_seqlen |
| 288 | dt_ptrs += BLOCK_SIZE_K * stride_dt_csize |
| 289 | dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize |
| 290 | |
| 291 | states = acc.to(states_ptr.dtype.element_ty) |
| 292 | |
| 293 | states_ptr += pid_c * stride_states_chunk + pid_h * stride_states_head |
| 294 | offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) |
| 295 | offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) |
| 296 | states_ptrs = states_ptr + ( |
| 297 | offs_m[:, None] * stride_states_hdim + offs_n[None, :] * stride_states_dstate |
| 298 | ) |
| 299 | c_mask = (offs_m[:, None] < hdim) & (offs_n[None, :] < dstate) |
| 300 | tl.store(states_ptrs, states, mask=c_mask) |
| 301 | |
| 302 | |
| 303 | def _chunk_cumsum_fwd( |
| 304 | dt, |
| 305 | A, |
| 306 | chunk_size, |
| 307 | cu_chunk_seqlens, |
| 308 | dt_bias=None, |
| 309 | dt_softplus=False, |
| 310 | dt_limit=(0.0, float("inf")), |
| 311 | ): |
| 312 | seqlen, nheads = dt.shape |
| 313 | assert A.shape == (nheads,) |
| 314 | if dt_bias is not None: |
| 315 | assert dt_bias.shape == (nheads,) |
| 316 | nchunks = cu_chunk_seqlens.shape[0] - 1 |
| 317 | dt_out = torch.empty( |
| 318 | nheads, nchunks, chunk_size, device=dt.device, dtype=torch.float32 |
| 319 | ) |
| 320 | dA_cumsum = torch.empty( |
| 321 | nheads, nchunks, chunk_size, device=dt.device, dtype=torch.float32 |
| 322 | ) |
| 323 | grid_chunk_cs = lambda META: (nchunks, triton.cdiv(nheads, META["BLOCK_SIZE_H"])) |
| 324 | with torch.accelerator.device_index(dt.device.index): |
| 325 | _chunk_cumsum_fwd_kernel[grid_chunk_cs]( |
| 326 | dt_ptr=dt, |
| 327 | A_ptr=A, |
| 328 | dt_bias_ptr=dt_bias, |
| 329 | dt_out_ptr=dt_out, |
| 330 | dA_cumsum_ptr=dA_cumsum, |
| 331 | cu_chunk_seqlens_ptr=cu_chunk_seqlens, |
| 332 | nheads=nheads, |
| 333 | chunk_size=chunk_size, |
| 334 | dt_min=dt_limit[0], |
| 335 | dt_max=dt_limit[1], |
| 336 | stride_dt_seqlen=dt.stride(0), |
| 337 | stride_dt_head=dt.stride(1), |
| 338 | stride_A_head=A.stride(0), |
| 339 | stride_dt_bias_head=dt_bias.stride(0) if dt_bias is not None else 0, |
| 340 | stride_dt_out_head=dt_out.stride(0), |
| 341 | stride_dt_out_chunk=dt_out.stride(1), |
| 342 | stride_dt_out_csize=dt_out.stride(2), |
| 343 | stride_dA_cs_head=dA_cumsum.stride(0), |
| 344 | stride_dA_cs_chunk=dA_cumsum.stride(1), |
| 345 | stride_dA_cs_csize=dA_cumsum.stride(2), |
| 346 | DT_SOFTPLUS=dt_softplus, |
| 347 | HAS_DT_BIAS=dt_bias is not None, |
| 348 | BLOCK_SIZE_CHUNK=triton.next_power_of_2(chunk_size), |
| 349 | ) |
| 350 | return dA_cumsum, dt_out |
| 351 | |
| 352 | |
| 353 | def _chunk_state_fwd( |
| 354 | B, x, dt, dA_cumsum, cu_chunk_seqlens, states=None, states_in_fp32=True |
| 355 | ): |
| 356 | seqlen, nheads, headdim = x.shape |
| 357 | _, nchunks, chunk_size = dt.shape |
| 358 | _, ngroups, dstate = B.shape |
| 359 | assert nheads % ngroups == 0 |
| 360 | assert B.shape == (seqlen, ngroups, dstate) |
| 361 | assert dt.shape == (nheads, nchunks, chunk_size) |
| 362 | assert dA_cumsum.shape == dt.shape |
| 363 | |
| 364 | if states is not None: |
| 365 | assert states.shape == (nchunks, nheads, headdim, dstate) |
| 366 | else: |
| 367 | states_dtype = torch.float32 if states_in_fp32 else B.dtype |
| 368 | states = torch.empty( |
| 369 | (nchunks, nheads, headdim, dstate), device=x.device, dtype=states_dtype |
| 370 | ) |
| 371 | |
| 372 | grid = lambda META: ( |
| 373 | triton.cdiv(headdim, META["BLOCK_SIZE_M"]) |
| 374 | * triton.cdiv(dstate, META["BLOCK_SIZE_N"]), |
| 375 | nchunks, |
| 376 | nheads, |
| 377 | ) |
| 378 | with torch.accelerator.device_index(x.device.index): |
| 379 | _chunk_state_fwd_kernel[grid]( |
| 380 | x_ptr=x, |
| 381 | b_ptr=B, |
| 382 | states_ptr=states, |
| 383 | dt_ptr=dt, |
| 384 | dA_cumsum_ptr=dA_cumsum, |
| 385 | cu_chunk_seqlens_ptr=cu_chunk_seqlens, |
| 386 | hdim=headdim, |
| 387 | dstate=dstate, |
| 388 | chunk_size=chunk_size, |
| 389 | nheads_ngroups_ratio=nheads // ngroups, |
| 390 | stride_x_seqlen=x.stride(0), |
| 391 | stride_x_head=x.stride(1), |
| 392 | stride_x_hdim=x.stride(2), |
| 393 | stride_b_seqlen=B.stride(0), |
| 394 | stride_b_head=B.stride(1), |
| 395 | stride_b_dstate=B.stride(2), |
| 396 | stride_states_chunk=states.stride(0), |
| 397 | stride_states_head=states.stride(1), |
| 398 | stride_states_hdim=states.stride(2), |
| 399 | stride_states_dstate=states.stride(3), |
| 400 | stride_dt_head=dt.stride(0), |
| 401 | stride_dt_chunk=dt.stride(1), |
| 402 | stride_dt_csize=dt.stride(2), |
| 403 | stride_dA_cs_head=dA_cumsum.stride(0), |
| 404 | stride_dA_cs_chunk=dA_cumsum.stride(1), |
| 405 | stride_dA_cs_csize=dA_cumsum.stride(2), |
| 406 | ) |
| 407 | return states |