| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | import torch.nn.functional as F |
| 8 | from einops import rearrange |
| 9 | |
| 10 | from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( |
| 11 | causal_conv1d_fn, |
| 12 | causal_conv1d_update, |
| 13 | ) |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.torch_utils import set_random_seed |
| 16 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID |
| 17 | |
| 18 | DEVICE = current_platform.device_type |
| 19 | |
| 20 | pytestmark = pytest.mark.skipif( |
| 21 | not ( |
| 22 | current_platform.is_cuda_alike() |
| 23 | or current_platform.is_xpu() |
| 24 | or current_platform.is_cpu() |
| 25 | ), |
| 26 | reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU", |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | def causal_conv1d_ref( |
| 31 | x: torch.Tensor, |
| 32 | weight: torch.Tensor, |
| 33 | bias: torch.Tensor | None = None, |
| 34 | initial_states: torch.Tensor | None = None, |
| 35 | return_final_states: bool = False, |
| 36 | final_states_out: torch.Tensor | None = None, |
| 37 | activation: str | None = "silu", |
| 38 | ): |
| 39 | """x: (batch, dim, seqlen) |
| 40 | weight: (dim, width) |
| 41 | bias: (dim,) |
| 42 | initial_states: (batch, dim, width - 1) |
| 43 | final_states_out: (batch, dim, width - 1) |
| 44 | |
| 45 | out: (batch, dim, seqlen) |
| 46 | """ |
| 47 | if activation not in [None, "silu", "swish"]: |
| 48 | raise NotImplementedError("activation must be None, silu, or swish") |
| 49 | dtype_in = x.dtype |
| 50 | x = x.to(weight.dtype) |
| 51 | seqlen = x.shape[-1] |
| 52 | dim, width = weight.shape |
| 53 | if initial_states is None: |
| 54 | out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) |
| 55 | else: |
| 56 | x = torch.cat([initial_states, x], dim=-1) |
| 57 | out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) |
| 58 | out = out[..., :seqlen] |
| 59 | if return_final_states: |
| 60 | final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( |
| 61 | dtype_in |
| 62 | ) # (batch, dim, width - 1) |
| 63 | if final_states_out is not None: |
| 64 | final_states_out.copy_(final_states) |
| 65 | else: |
| 66 | final_states_out = final_states |
| 67 | out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) |
| 68 | return (out, None) if not return_final_states else (out, final_states_out) |
| 69 | |
| 70 | |
| 71 | def causal_conv1d_update_ref( |
| 72 | x, conv_state, weight, bias=None, activation=None, cache_seqlens=None |
| 73 | ): |
| 74 | """x: (batch, dim) or (batch, dim, seqlen) |
| 75 | conv_state: (batch, dim, state_len), where state_len >= width - 1 |
| 76 | weight: (dim, width) |
| 77 | bias: (dim,) |
| 78 | cache_seqlens: (batch,), dtype int32. |
| 79 | If not None, the conv_state is treated as a circular buffer. |
| 80 | The conv_state will be updated by copying x to the |
| 81 | conv_state starting at the index |
| 82 | @cache_seqlens % state_len before performing the convolution. |
| 83 | |
| 84 | out: (batch, dim) or (batch, dim, seqlen) |
| 85 | """ |
| 86 | if activation not in [None, "silu", "swish"]: |
| 87 | raise NotImplementedError("activation must be None, silu, or swish") |
| 88 | dtype_in = x.dtype |
| 89 | unsqueeze = x.dim() == 2 |
| 90 | if unsqueeze: |
| 91 | x = x.unsqueeze(-1) |
| 92 | batch, dim, seqlen = x.shape |
| 93 | width = weight.shape[1] |
| 94 | state_len = conv_state.shape[-1] |
| 95 | assert conv_state.shape == (batch, dim, state_len) |
| 96 | assert weight.shape == (dim, width) |
| 97 | if cache_seqlens is None: |
| 98 | x_new = torch.cat([conv_state, x], dim=-1).to( |
| 99 | weight.dtype |
| 100 | ) # (batch, dim, state_len + seqlen) |
| 101 | conv_state.copy_(x_new[:, :, -state_len:]) |
| 102 | else: |
| 103 | width_idx = torch.arange( |
| 104 | -(width - 1), 0, dtype=torch.long, device=x.device |
| 105 | ).unsqueeze(0) + cache_seqlens.unsqueeze(1) |
| 106 | width_idx = ( |
| 107 | torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) |
| 108 | ) |
| 109 | x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) |
| 110 | copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze( |
| 111 | 0 |
| 112 | ) + cache_seqlens.unsqueeze(1) |
| 113 | copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) |
| 114 | conv_state.scatter_(2, copy_idx, x) |
| 115 | out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[ |
| 116 | :, :, -seqlen: |
| 117 | ] |
| 118 | if unsqueeze: |
| 119 | out = out.squeeze(-1) |
| 120 | return (out if activation is None else F.silu(out)).to(dtype=dtype_in) |
| 121 | |
| 122 | |
| 123 | @pytest.mark.parametrize("itype", [torch.bfloat16]) |
| 124 | @pytest.mark.parametrize("silu_activation", [False, True]) |
| 125 | @pytest.mark.parametrize("has_bias", [False, True]) |
| 126 | @pytest.mark.parametrize("seqlen", [1]) |
| 127 | @pytest.mark.parametrize("width", [4]) |
| 128 | @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) |
| 129 | def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, itype): |
| 130 | device = DEVICE |
| 131 | rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) |
| 132 | if itype == torch.bfloat16: |
| 133 | rtol, atol = 1e-2, 5e-2 |
| 134 | # set seed |
| 135 | set_random_seed(0) |
| 136 | batch = 2 |
| 137 | x = torch.randn(batch, dim, seqlen, device=device, dtype=itype) |
| 138 | x_ref = x.clone() |
| 139 | # +1 entry to reserve index 0 as null block |
| 140 | conv_state = torch.randn(batch + 1, dim, width - 1, device=device, dtype=itype) |
| 141 | |
| 142 | weight = torch.randn(dim, width, device=device, dtype=itype) |
| 143 | bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None |
| 144 | # Start indices from 1, skipping null block at index 0 |
| 145 | conv_state_indices = torch.arange(1, batch + 1, dtype=torch.int32, device=device) |
| 146 | conv_state_ref = conv_state[conv_state_indices].detach().clone() |
| 147 | activation = None if not silu_activation else "silu" |
| 148 | |
| 149 | out = causal_conv1d_update( |
| 150 | x, |
| 151 | conv_state, |
| 152 | weight, |
| 153 | bias, |
| 154 | activation=activation, |
| 155 | conv_state_indices=conv_state_indices, |
| 156 | ) |
| 157 | out_ref = causal_conv1d_update_ref( |
| 158 | x_ref, conv_state_ref, weight, bias, activation=activation |
| 159 | ) |
| 160 | |
| 161 | assert torch.equal(conv_state[conv_state_indices], conv_state_ref) |
| 162 | assert torch.allclose(out, out_ref, rtol=rtol, atol=atol) |
| 163 | |
| 164 | |
| 165 | @pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16]) |
| 166 | @pytest.mark.parametrize("silu_activation", [False, True]) |
| 167 | @pytest.mark.parametrize("has_bias", [False, True]) |
| 168 | @pytest.mark.parametrize("seqlen", [1, 3]) |
| 169 | @pytest.mark.parametrize("width", [3, 4]) |
| 170 | @pytest.mark.parametrize("dim", [2048 + 16, 4096]) |
| 171 | # tests correctness in case subset of the sequences are padded |
| 172 | @pytest.mark.parametrize("with_padding", [True, False]) |
| 173 | @pytest.mark.parametrize("batch_size", [3]) |
| 174 | def test_causal_conv1d_update_with_batch_gather( |
| 175 | batch_size, with_padding, dim, width, seqlen, has_bias, silu_activation, itype |
| 176 | ): |
| 177 | device = DEVICE |
| 178 | rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) |
| 179 | if itype == torch.bfloat16: |
| 180 | rtol, atol = 1e-2, 5e-2 |
| 181 | |
| 182 | # set seed |
| 183 | set_random_seed(0) |
| 184 | |
| 185 | padding = 5 if with_padding else 0 |
| 186 | padded_batch_size = batch_size + padding |
| 187 | # total_entries = number of cache line |
| 188 | total_entries = 10 * batch_size |
| 189 | |
| 190 | # x will be (batch, dim, seqlen) with contiguous along dim-axis |
| 191 | x = torch.randn( |
| 192 | padded_batch_size, seqlen, dim, device=device, dtype=itype |
| 193 | ).transpose(1, 2) |
| 194 | |
| 195 | x_ref = x.clone() |
| 196 | |
| 197 | # +1 to exclude index 0 (null block) |
| 198 | conv_state_indices = (torch.randperm(total_entries - 1)[:batch_size] + 1).to( |
| 199 | dtype=torch.int32, device=device |
| 200 | ) |
| 201 | unused_states_bool = torch.ones(total_entries, dtype=torch.bool, device=device) |
| 202 | unused_states_bool[conv_state_indices] = False |
| 203 | padded_state_indices = torch.concat( |
| 204 | [ |
| 205 | conv_state_indices, |
| 206 | torch.as_tensor( |
| 207 | [NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device |
| 208 | ), |
| 209 | ], |
| 210 | dim=0, |
| 211 | ) |
| 212 | |
| 213 | # conv_state will be (cache_lines, dim, state_len) |
| 214 | # with contiguous along dim-axis |
| 215 | conv_state = torch.randn( |
| 216 | total_entries, width - 1, dim, device=device, dtype=itype |
| 217 | ).transpose(1, 2) |
| 218 | |
| 219 | conv_state_for_padding_test = conv_state.clone() |
| 220 | |
| 221 | weight = torch.randn(dim, width, device=device, dtype=itype) |
| 222 | bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None |
| 223 | conv_state_ref = conv_state[conv_state_indices, :].detach().clone() |
| 224 | activation = None if not silu_activation else "silu" |
| 225 | |
| 226 | out = causal_conv1d_update( |
| 227 | x, |
| 228 | conv_state, |
| 229 | weight, |
| 230 | bias, |
| 231 | activation=activation, |
| 232 | conv_state_indices=padded_state_indices, |
| 233 | ) |
| 234 | out_ref = causal_conv1d_update_ref( |
| 235 | x_ref[:batch_size], conv_state_ref, weight, bias, activation=activation |
| 236 | ) |
| 237 | |
| 238 | assert torch.equal(conv_state[conv_state_indices, :], conv_state_ref) |
| 239 | assert torch.equal( |
| 240 | conv_state[unused_states_bool], conv_state_for_padding_test[unused_states_bool] |
| 241 | ) |
| 242 | assert torch.allclose(out[:batch_size], out_ref, rtol=rtol, atol=atol) |
| 243 | |
| 244 | |
| 245 | @pytest.mark.parametrize("itype", [torch.bfloat16]) |
| 246 | @pytest.mark.parametrize("silu_activation", [True]) |
| 247 | @pytest.mark.parametrize("has_bias", [True]) |
| 248 | @pytest.mark.parametrize("width", [4]) |
| 249 | @pytest.mark.parametrize("seqlen", [8, 249, 4096]) |
| 250 | @pytest.mark.parametrize("dim", [64, 4096]) |
| 251 | @pytest.mark.parametrize("with_padding", [True, False]) |
| 252 | @pytest.mark.parametrize("batch", [4, 10]) |
| 253 | def test_causal_conv1d_varlen( |
| 254 | batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype |
| 255 | ): |
| 256 | device = DEVICE |
| 257 | if not current_platform.is_cpu(): |
| 258 | torch.accelerator.empty_cache() |
| 259 | rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) |
| 260 | if itype == torch.bfloat16: |
| 261 | rtol, atol = 1e-2, 5e-2 |
| 262 | # set seed |
| 263 | set_random_seed(0) |
| 264 | seqlens = [] |
| 265 | batch_size = batch |
| 266 | padding = 3 if with_padding else 0 |
| 267 | padded_batch_size = batch_size + padding |
| 268 | nsplits = padded_batch_size - 1 |
| 269 | |
| 270 | eos_pos = torch.randperm(seqlen - 1)[:nsplits].sort().values |
| 271 | |
| 272 | seqlens.append( |
| 273 | torch.diff( |
| 274 | torch.cat([torch.tensor([-1]), eos_pos, torch.tensor([seqlen - 1])]) |
| 275 | ).tolist() |
| 276 | ) |
| 277 | assert sum(seqlens[-1]) == seqlen |
| 278 | assert all(s > 0 for s in seqlens[-1]) |
| 279 | |
| 280 | total_entries = batch_size * 10 |
| 281 | cumsum = torch.cumsum(torch.tensor(seqlens[0]), dim=0).to(torch.int32) |
| 282 | cumsum = torch.concat([torch.tensor([0], dtype=torch.int32), cumsum], dim=0) |
| 283 | x = rearrange( |
| 284 | torch.randn(1, seqlen, 4096 + dim + 64, device=device, dtype=itype), |
| 285 | "b s d -> b d s", |
| 286 | )[:, 4096 : 4096 + dim, :] |
| 287 | |
| 288 | weight = torch.randn(dim, width, device=device, dtype=itype) |
| 289 | |
| 290 | bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None |
| 291 | x_ref = x.clone() |
| 292 | weight_ref = weight.clone() |
| 293 | bias_ref = bias.clone() if bias is not None else None |
| 294 | activation = None if not silu_activation else "silu" |
| 295 | final_states = torch.randn( |
| 296 | total_entries, width - 1, dim, device=x.device, dtype=x.dtype |
| 297 | ).transpose(1, 2) |
| 298 | final_states_ref = final_states.clone() |
| 299 | has_initial_states = torch.randint( |
| 300 | 0, 2, (cumsum.shape[0] - 1,), dtype=torch.bool, device=x.device |
| 301 | ) |
| 302 | # +1 to exclude index 0 (null block) |
| 303 | state_indices = ( |
| 304 | torch.randperm(total_entries - 1, dtype=torch.int32, device=x.device)[ |
| 305 | :batch_size |
| 306 | ] |
| 307 | + 1 |
| 308 | ) |
| 309 | padded_state_indices = torch.concat( |
| 310 | [ |
| 311 | state_indices, |
| 312 | torch.as_tensor( |
| 313 | [NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device |
| 314 | ), |
| 315 | ], |
| 316 | dim=-1, |
| 317 | ) |
| 318 | out = causal_conv1d_fn( |
| 319 | x.squeeze(0), |
| 320 | weight, |
| 321 | bias=bias, |
| 322 | conv_states=final_states, |
| 323 | query_start_loc=cumsum.to(device), |
| 324 | cache_indices=padded_state_indices, |
| 325 | has_initial_state=has_initial_states, |
| 326 | activation=activation, |
| 327 | ) |
| 328 | |
| 329 | out_ref = [] |
| 330 | out_ref_b = [] |
| 331 | |
| 332 | splits = [torch.split(var, seqlens[0], dim=-1) for var in (x_ref)] |
| 333 | for i in range(len(seqlens[0])): |
| 334 | x_s = [v[i].unsqueeze(0) for v in splits][0] |
| 335 | if padded_state_indices[i] == NULL_BLOCK_ID: |
| 336 | continue |
| 337 | out_ref_b.append( |
| 338 | causal_conv1d_ref( |
| 339 | x_s, |
| 340 | weight_ref, |
| 341 | bias_ref, |
| 342 | activation=activation, |
| 343 | return_final_states=True, |
| 344 | final_states_out=final_states_ref[padded_state_indices[i]].unsqueeze(0), |
| 345 | initial_states=final_states_ref[padded_state_indices[i]].unsqueeze(0) |
| 346 | if has_initial_states[i] |
| 347 | else None, |
| 348 | ) |
| 349 | ) |
| 350 | out_ref.append(torch.cat([t[0] for t in out_ref_b], dim=2)) |
| 351 | out_ref_tensor = torch.cat(out_ref, dim=0) |
| 352 | |
| 353 | assert torch.allclose( |
| 354 | final_states[state_indices], |
| 355 | final_states_ref[state_indices], |
| 356 | rtol=rtol, |
| 357 | atol=atol, |
| 358 | ) |
| 359 | unpadded_out = out[:, : out_ref_tensor.shape[-1]] |
| 360 | assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol) |