| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import functools |
| 5 | import types |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | |
| 11 | import vllm._custom_ops as ops |
| 12 | from vllm.model_executor.layers.mamba.ops.cpu import gdn_attention |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.torch_utils import set_random_seed |
| 15 | from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata |
| 16 | |
| 17 | if not current_platform.is_cpu(): |
| 18 | pytest.skip("skipping CPU-only tests", allow_module_level=True) |
| 19 | |
| 20 | set_random_seed(12345) |
| 21 | |
| 22 | NUM_HEADS = [ |
| 23 | (2, 4), |
| 24 | (4, 4), |
| 25 | ] |
| 26 | HEAD_DIMS = [ |
| 27 | (32, 32), |
| 28 | (64, 32), |
| 29 | ] |
| 30 | # chunk_gated_delta_rule_cpu (the chunked-prefill kernel) only supports |
| 31 | # head_dim == head_dim_v in {64, 128}; the decode-path update kernel above |
| 32 | # has no such restriction and keeps using the wider HEAD_DIMS list. |
| 33 | CHUNK_HEAD_DIMS = [ |
| 34 | (64, 64), |
| 35 | (128, 128), |
| 36 | ] |
| 37 | CHUNK_SIZE = 64 |
| 38 | CONV_DIM = 128 |
| 39 | CONV_KERNEL = 4 |
| 40 | DECODE_BATCH_SIZES = [1, 3, 5] |
| 41 | PREFILL_SEQ_LENS = [ |
| 42 | [1], |
| 43 | [1, 2, 3], |
| 44 | [CHUNK_SIZE - 1], |
| 45 | [CHUNK_SIZE], |
| 46 | [CHUNK_SIZE + 1], |
| 47 | [CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE + 1], |
| 48 | [2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1], |
| 49 | [4 * CHUNK_SIZE + 17], |
| 50 | ] |
| 51 | DECODE_DTYPE_CASES = [ |
| 52 | *[ |
| 53 | pytest.param( |
| 54 | batch_size, |
| 55 | num_heads, |
| 56 | head_dims, |
| 57 | torch.float32, |
| 58 | id=f"fp32-b{batch_size}-h{num_heads}-d{head_dims}", |
| 59 | ) |
| 60 | for batch_size in [1, 3, 5] |
| 61 | for num_heads in NUM_HEADS |
| 62 | for head_dims in HEAD_DIMS |
| 63 | ], |
| 64 | pytest.param(3, (2, 4), (32, 32), torch.bfloat16, id="bf16-representative"), |
| 65 | pytest.param(3, (2, 4), (32, 32), torch.float16, id="fp16-representative"), |
| 66 | ] |
| 67 | PREFILL_DTYPE_CASES = [ |
| 68 | *[ |
| 69 | pytest.param( |
| 70 | num_heads, |
| 71 | head_dims, |
| 72 | torch.float32, |
| 73 | id=f"fp32-h{num_heads}-d{head_dims}", |
| 74 | ) |
| 75 | for num_heads in NUM_HEADS |
| 76 | for head_dims in CHUNK_HEAD_DIMS |
| 77 | ], |
| 78 | pytest.param((2, 4), (64, 64), torch.bfloat16, id="bf16-representative"), |
| 79 | pytest.param((2, 4), (64, 64), torch.float16, id="fp16-representative"), |
| 80 | ] |
| 81 | |
| 82 | |
| 83 | @functools.lru_cache(maxsize=128, typed=False) |
| 84 | def tensor_cache( |
| 85 | elem_num: int, |
| 86 | dtype: torch.dtype, |
| 87 | ) -> torch.Tensor: |
| 88 | tensor = torch.rand(elem_num, dtype=dtype) |
| 89 | return tensor |
| 90 | |
| 91 | |
| 92 | def ref_l2norm( |
| 93 | x: torch.Tensor, |
| 94 | dim: int = -1, |
| 95 | eps: float = 1e-5, |
| 96 | ) -> torch.Tensor: |
| 97 | inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) |
| 98 | return x * inv_norm |
| 99 | |
| 100 | |
| 101 | def ref_gdn_gating( |
| 102 | A_log: torch.Tensor, |
| 103 | a: torch.Tensor, |
| 104 | b: torch.Tensor, |
| 105 | dt_bias: torch.Tensor, |
| 106 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 107 | softplus_x = F.softplus(a.float() + dt_bias.float(), beta=1.0, threshold=20.0) |
| 108 | g = -torch.exp(A_log.float()) * softplus_x |
| 109 | beta = torch.sigmoid(b.float()).to(dtype=b.dtype) |
| 110 | return g, beta |
| 111 | |
| 112 | |
| 113 | def ref_gated_delta_rule( |
| 114 | query: torch.Tensor, |
| 115 | key: torch.Tensor, |
| 116 | value: torch.Tensor, |
| 117 | a: torch.Tensor, |
| 118 | b: torch.Tensor, |
| 119 | A_log: torch.Tensor, |
| 120 | dt_bias: torch.Tensor, |
| 121 | initial_state: torch.Tensor, |
| 122 | cu_seqlens: torch.Tensor, |
| 123 | use_qk_l2norm_in_kernel: bool = False, |
| 124 | state_dtype: torch.dtype = torch.float32, |
| 125 | state_write_interval: int = 1, |
| 126 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 127 | g, beta = ref_gdn_gating(A_log, a, b, dt_bias) |
| 128 | out = torch.empty_like(value) |
| 129 | final_state = torch.empty_like(initial_state) |
| 130 | |
| 131 | for seq_idx in range(cu_seqlens.numel() - 1): |
| 132 | begin = int(cu_seqlens[seq_idx].item()) |
| 133 | end = int(cu_seqlens[seq_idx + 1].item()) |
| 134 | q_seq = query[:, begin:end] |
| 135 | k_seq = key[:, begin:end] |
| 136 | v_seq = value[:, begin:end] |
| 137 | g_seq = g[begin:end].unsqueeze(0) |
| 138 | beta_seq = beta[begin:end].unsqueeze(0) |
| 139 | initial_dtype = q_seq.dtype |
| 140 | |
| 141 | if use_qk_l2norm_in_kernel: |
| 142 | q_seq = ref_l2norm(q_seq, dim=-1) |
| 143 | k_seq = ref_l2norm(k_seq, dim=-1) |
| 144 | |
| 145 | if q_seq.shape[2] != v_seq.shape[2]: |
| 146 | repeat_factor = v_seq.shape[2] // q_seq.shape[2] |
| 147 | q_seq = q_seq.repeat_interleave(repeat_factor, dim=2) |
| 148 | k_seq = k_seq.repeat_interleave(repeat_factor, dim=2) |
| 149 | |
| 150 | q_seq, k_seq, v_seq, beta_seq, g_seq = [ |
| 151 | x.transpose(1, 2).contiguous().to(torch.float32) |
| 152 | for x in (q_seq, k_seq, v_seq, beta_seq, g_seq) |
| 153 | ] |
| 154 | |
| 155 | batch_size, num_heads, seq_len, head_dim = q_seq.shape |
| 156 | v_head_dim = v_seq.shape[-1] |
| 157 | q_seq = q_seq * (1 / (head_dim**0.5)) |
| 158 | out_seq = torch.empty( |
| 159 | batch_size, |
| 160 | num_heads, |
| 161 | seq_len, |
| 162 | v_head_dim, |
| 163 | dtype=v_seq.dtype, |
| 164 | ) |
| 165 | state = initial_state[seq_idx : seq_idx + 1].to(v_seq) |
| 166 | |
| 167 | for token_idx in range(seq_len): |
| 168 | q_t = q_seq[:, :, token_idx] |
| 169 | k_t = k_seq[:, :, token_idx] |
| 170 | v_t = v_seq[:, :, token_idx] |
| 171 | g_t = g_seq[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1) |
| 172 | beta_t = beta_seq[:, :, token_idx].unsqueeze(-1) |
| 173 | |
| 174 | state = state * g_t |
| 175 | kv_mem = (state * k_t.unsqueeze(-2)).sum(dim=-1) |
| 176 | delta = (v_t - kv_mem) * beta_t |
| 177 | state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2) |
| 178 | out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1) |
| 179 | if state_dtype != torch.float32 and ( |
| 180 | (token_idx + 1) % state_write_interval == 0 or token_idx + 1 == seq_len |
| 181 | ): |
| 182 | state = state.to(state_dtype).float() |
| 183 | |
| 184 | out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype) |
| 185 | final_state[seq_idx] = state.squeeze(0) |
| 186 | |
| 187 | return out, final_state |
| 188 | |
| 189 | |
| 190 | def gdn_inputs( |
| 191 | num_tokens: int, |
| 192 | num_heads: tuple[int, int], |
| 193 | head_dims: tuple[int, int], |
| 194 | ) -> tuple[torch.Tensor, ...]: |
| 195 | num_qk_heads, num_v_heads = num_heads |
| 196 | head_dim, v_head_dim = head_dims |
| 197 | q_shape = (1, num_tokens, num_qk_heads, head_dim) |
| 198 | q_numel = num_tokens * num_qk_heads * head_dim |
| 199 | q = tensor_cache(q_numel, torch.bfloat16).view(q_shape) |
| 200 | k = tensor_cache(q_numel, torch.bfloat16).view(q_shape) |
| 201 | |
| 202 | v_shape = (1, num_tokens, num_v_heads, v_head_dim) |
| 203 | v = tensor_cache(num_tokens * num_v_heads * v_head_dim, torch.bfloat16).view( |
| 204 | v_shape |
| 205 | ) |
| 206 | |
| 207 | gate_shape = (num_tokens, num_v_heads) |
| 208 | gate_numel = num_tokens * num_v_heads |
| 209 | a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) |
| 210 | b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) |
| 211 | A_log = tensor_cache(num_v_heads, torch.float32) |
| 212 | dt_bias = tensor_cache(num_v_heads, torch.bfloat16) |
| 213 | return q, k, v, a, b, A_log, dt_bias |
| 214 | |
| 215 | |
| 216 | @pytest.mark.parametrize("num_tokens", [1, 9]) |
| 217 | @pytest.mark.parametrize("num_v_heads", [4, 8]) |
| 218 | @torch.inference_mode() |
| 219 | def test_fused_gdn_gating_cpu( |
| 220 | num_tokens: int, |
| 221 | num_v_heads: int, |
| 222 | ) -> None: |
| 223 | gate_shape = (num_tokens, num_v_heads) |
| 224 | gate_numel = num_tokens * num_v_heads |
| 225 | a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) |
| 226 | b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) |
| 227 | A_log = tensor_cache(num_v_heads, torch.float32) |
| 228 | dt_bias = tensor_cache(num_v_heads, torch.bfloat16) |
| 229 | |
| 230 | g_ref, beta_ref = ref_gdn_gating(A_log, a, b, dt_bias) |
| 231 | g, beta = ops.fused_gdn_gating_cpu(A_log, a, b, dt_bias) |
| 232 | |
| 233 | torch.testing.assert_close(g, g_ref.unsqueeze(0), atol=1e-4, rtol=1e-4) |
| 234 | torch.testing.assert_close( |
| 235 | beta.float(), beta_ref.unsqueeze(0).float(), atol=5e-3, rtol=5e-3 |
| 236 | ) |
| 237 | |
| 238 | |
| 239 | # decode path |
| 240 | @pytest.mark.parametrize( |
| 241 | "batch_size,num_heads,head_dims,state_dtype", DECODE_DTYPE_CASES |
| 242 | ) |
| 243 | @torch.inference_mode() |
| 244 | def test_fused_sigmoid_gating_delta_rule_update_cpu( |
| 245 | batch_size: int, |
| 246 | num_heads: tuple[int, int], |
| 247 | head_dims: tuple[int, int], |
| 248 | state_dtype: torch.dtype, |
| 249 | ) -> None: |
| 250 | q, k, v, a, b, A_log, dt_bias = gdn_inputs( |
| 251 | num_tokens=batch_size, |
| 252 | num_heads=num_heads, |
| 253 | head_dims=head_dims, |
| 254 | ) |
| 255 | _, num_v_heads = num_heads |
| 256 | head_dim, v_head_dim = head_dims |
| 257 | state_indices = torch.arange(batch_size, dtype=torch.int32) |
| 258 | cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32) |
| 259 | state_shape = (batch_size, num_v_heads, head_dim, v_head_dim) |
| 260 | state = tensor_cache( |
| 261 | batch_size * num_v_heads * head_dim * v_head_dim, state_dtype |
| 262 | ).view(state_shape) |
| 263 | state_ref = state[state_indices].transpose(-1, -2).contiguous() |
| 264 | |
| 265 | out_ref, final_state_ref = ref_gated_delta_rule( |
| 266 | query=q, |
| 267 | key=k, |
| 268 | value=v, |
| 269 | a=a, |
| 270 | b=b, |
| 271 | A_log=A_log, |
| 272 | dt_bias=dt_bias, |
| 273 | initial_state=state_ref, |
| 274 | cu_seqlens=cu_seqlens, |
| 275 | use_qk_l2norm_in_kernel=True, |
| 276 | state_dtype=state_dtype, |
| 277 | ) |
| 278 | out_ref = out_ref.transpose(0, 1).contiguous() |
| 279 | |
| 280 | state_out = state.clone() |
| 281 | out = ops.fused_sigmoid_gating_delta_rule_update_cpu( |
| 282 | A_log=A_log, |
| 283 | dt_bias=dt_bias, |
| 284 | q=q, |
| 285 | k=k, |
| 286 | v=v, |
| 287 | a=a, |
| 288 | b=b, |
| 289 | initial_state_source=state_out, |
| 290 | initial_state_indices=state_indices, |
| 291 | cu_seqlens=cu_seqlens, |
| 292 | use_qk_l2norm_in_kernel=True, |
| 293 | ) |
| 294 | |
| 295 | torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) |
| 296 | torch.testing.assert_close( |
| 297 | state_out[state_indices].transpose(-1, -2).float(), |
| 298 | final_state_ref.float(), |
| 299 | atol=1e-2, |
| 300 | rtol=1e-2, |
| 301 | ) |
| 302 | |
| 303 | |
| 304 | # prefill path |
| 305 | @pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS) |
| 306 | @pytest.mark.parametrize("num_heads,head_dims,state_dtype", PREFILL_DTYPE_CASES) |
| 307 | @torch.inference_mode() |
| 308 | def test_chunk_gated_delta_rule_cpu( |
| 309 | seq_lens: list[int], |
| 310 | num_heads: tuple[int, int], |
| 311 | head_dims: tuple[int, int], |
| 312 | state_dtype: torch.dtype, |
| 313 | ) -> None: |
| 314 | total_tokens = sum(seq_lens) |
| 315 | q, k, v, a, b, A_log, dt_bias = gdn_inputs( |
| 316 | num_tokens=total_tokens, |
| 317 | num_heads=num_heads, |
| 318 | head_dims=head_dims, |
| 319 | ) |
| 320 | _, num_v_heads = num_heads |
| 321 | head_dim, v_head_dim = head_dims |
| 322 | cu_seqlens = torch.tensor( |
| 323 | [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 |
| 324 | ) |
| 325 | initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim) |
| 326 | initial_state = tensor_cache( |
| 327 | len(seq_lens) * num_v_heads * head_dim * v_head_dim, state_dtype |
| 328 | ).view(initial_state_shape) |
| 329 | initial_state_ref = initial_state.transpose(-1, -2).contiguous() |
| 330 | |
| 331 | out_ref, final_state_ref = ref_gated_delta_rule( |
| 332 | query=q, |
| 333 | key=k, |
| 334 | value=v, |
| 335 | a=a, |
| 336 | b=b, |
| 337 | A_log=A_log, |
| 338 | dt_bias=dt_bias, |
| 339 | initial_state=initial_state_ref, |
| 340 | cu_seqlens=cu_seqlens, |
| 341 | use_qk_l2norm_in_kernel=True, |
| 342 | state_dtype=state_dtype, |
| 343 | state_write_interval=CHUNK_SIZE, |
| 344 | ) |
| 345 | |
| 346 | g, beta = ref_gdn_gating(A_log, a, b, dt_bias) |
| 347 | out, final_state = ops.chunk_gated_delta_rule_cpu( |
| 348 | query=q, |
| 349 | key=k, |
| 350 | value=v, |
| 351 | g=g.unsqueeze(0), |
| 352 | beta=beta.unsqueeze(0), |
| 353 | initial_state=initial_state, |
| 354 | output_final_state=True, |
| 355 | cu_seqlens=cu_seqlens, |
| 356 | head_first=False, |
| 357 | use_qk_l2norm_in_kernel=True, |
| 358 | initial_state_indices=torch.arange(len(seq_lens), dtype=torch.int32), |
| 359 | ) |
| 360 | |
| 361 | torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) |
| 362 | torch.testing.assert_close( |
| 363 | final_state.transpose(-1, -2).float(), |
| 364 | final_state_ref.float(), |
| 365 | atol=1e-2, |
| 366 | rtol=1e-2, |
| 367 | ) |
| 368 | |
| 369 | |
| 370 | # (total_tokens, split) pairs mimicking where chunked prefill breaks a sequence |
| 371 | # across two scheduler steps: chunk-aligned and non-aligned splits. |
| 372 | TWO_CALL_SPLITS = [ |
| 373 | (2 * CHUNK_SIZE, CHUNK_SIZE), |
| 374 | (2 * CHUNK_SIZE + 17, CHUNK_SIZE), |
| 375 | (2 * CHUNK_SIZE + 17, CHUNK_SIZE + 9), |
| 376 | (4 * CHUNK_SIZE + 17, 2 * CHUNK_SIZE), |
| 377 | (3 * CHUNK_SIZE, CHUNK_SIZE + 1), |
| 378 | ] |
| 379 | |
| 380 | |
| 381 | @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) |
| 382 | @pytest.mark.parametrize("num_heads", [(2, 4)]) |
| 383 | @pytest.mark.parametrize("head_dims", [(64, 64)]) |
| 384 | @pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16, torch.float16]) |
| 385 | @torch.inference_mode() |
| 386 | def test_chunk_gated_delta_rule_cpu_two_call_split( |
| 387 | total_tokens: int, |
| 388 | split: int, |
| 389 | num_heads: tuple[int, int], |
| 390 | head_dims: tuple[int, int], |
| 391 | state_dtype: torch.dtype, |
| 392 | ) -> None: |
| 393 | """A prefill split into two calls (the second seeded with the first's |
| 394 | ``final_state`` and a rebased ``cu_seqlens``) must match the single-call |
| 395 | result, mimicking the cross-scheduler-step handoff in |
| 396 | ``cpu_gdn_attention_core``. |
| 397 | """ |
| 398 | q, k, v, a, b, A_log, dt_bias = gdn_inputs( |
| 399 | num_tokens=total_tokens, |
| 400 | num_heads=num_heads, |
| 401 | head_dims=head_dims, |
| 402 | ) |
| 403 | _, num_v_heads = num_heads |
| 404 | head_dim, v_head_dim = head_dims |
| 405 | |
| 406 | g, beta = ref_gdn_gating(A_log, a, b, dt_bias) |
| 407 | g = g.unsqueeze(0) # [1, T, HV] |
| 408 | beta = beta.unsqueeze(0) |
| 409 | |
| 410 | zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=state_dtype) |
| 411 | |
| 412 | # Reference: whole sequence in one call, no initial state. |
| 413 | out_full, final_full = ops.chunk_gated_delta_rule_cpu( |
| 414 | query=q, |
| 415 | key=k, |
| 416 | value=v, |
| 417 | g=g, |
| 418 | beta=beta, |
| 419 | initial_state=zero_state, |
| 420 | output_final_state=True, |
| 421 | cu_seqlens=torch.tensor([0, total_tokens], dtype=torch.int32), |
| 422 | head_first=False, |
| 423 | use_qk_l2norm_in_kernel=True, |
| 424 | initial_state_indices=torch.zeros(1, dtype=torch.int32), |
| 425 | ) |
| 426 | |
| 427 | # Call 1: tokens [0:split], no initial state, capture final state. |
| 428 | out1, state1 = ops.chunk_gated_delta_rule_cpu( |
| 429 | query=q[:, :split], |
| 430 | key=k[:, :split], |
| 431 | value=v[:, :split], |
| 432 | g=g[:, :split], |
| 433 | beta=beta[:, :split], |
| 434 | initial_state=zero_state, |
| 435 | output_final_state=True, |
| 436 | cu_seqlens=torch.tensor([0, split], dtype=torch.int32), |
| 437 | head_first=False, |
| 438 | use_qk_l2norm_in_kernel=True, |
| 439 | initial_state_indices=torch.zeros(1, dtype=torch.int32), |
| 440 | ) |
| 441 | # Call 2: tokens [split:T] seeded with call 1's final state and a cu_seqlens |
| 442 | # rebased to start at 0, as cpu_gdn_attention_core continues a prefill chunk. |
| 443 | tail = total_tokens - split |
| 444 | out2, state2 = ops.chunk_gated_delta_rule_cpu( |
| 445 | query=q[:, split:], |
| 446 | key=k[:, split:], |
| 447 | value=v[:, split:], |
| 448 | g=g[:, split:], |
| 449 | beta=beta[:, split:], |
| 450 | initial_state=state1, |
| 451 | output_final_state=True, |
| 452 | cu_seqlens=torch.tensor([0, tail], dtype=torch.int32), |
| 453 | head_first=False, |
| 454 | use_qk_l2norm_in_kernel=True, |
| 455 | initial_state_indices=torch.zeros(1, dtype=torch.int32), |
| 456 | ) |
| 457 | |
| 458 | out_split = torch.cat([out1, out2], dim=1) |
| 459 | |
| 460 | # State must be near-exact; output allows a looser bound for the bf16 round-trip. |
| 461 | torch.testing.assert_close(state2.float(), final_full.float(), atol=1e-3, rtol=1e-3) |
| 462 | torch.testing.assert_close(out_split, out_full, atol=2e-2, rtol=2e-2) |
| 463 | |
| 464 | |
| 465 | def _conv_inputs(total_tokens: int): |
| 466 | x = tensor_cache(total_tokens * CONV_DIM, torch.bfloat16).view( |
| 467 | total_tokens, CONV_DIM |
| 468 | ) |
| 469 | weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view( |
| 470 | CONV_DIM, CONV_KERNEL |
| 471 | ) |
| 472 | bias = tensor_cache(CONV_DIM, torch.bfloat16) |
| 473 | return x, weight, bias |
| 474 | |
| 475 | |
| 476 | def _conv_states( |
| 477 | num_slots: int, |
| 478 | state_len: int, |
| 479 | dim: int = CONV_DIM, |
| 480 | layout: str = "SD", |
| 481 | ) -> torch.Tensor: |
| 482 | if layout == "SD": |
| 483 | return torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16).transpose( |
| 484 | 1, 2 |
| 485 | ) |
| 486 | assert layout == "DS" |
| 487 | return torch.zeros(num_slots, dim, state_len, dtype=torch.bfloat16) |
| 488 | |
| 489 | |
| 490 | def _sd_conv_states( |
| 491 | num_slots: int, state_len: int, dim: int = CONV_DIM |
| 492 | ) -> torch.Tensor: |
| 493 | storage = torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16) |
| 494 | return storage.transpose(1, 2) |
| 495 | |
| 496 | |
| 497 | def _maybe_pack_conv_weight( |
| 498 | weight: torch.Tensor, is_weight_packed: bool |
| 499 | ) -> torch.Tensor: |
| 500 | return ops.causal_conv1d_weight_pack(weight) if is_weight_packed else weight |
| 501 | |
| 502 | |
| 503 | @torch.inference_mode() |
| 504 | def test_causal_conv1d_update_cpu_preserves_sd_state( |
| 505 | monkeypatch: pytest.MonkeyPatch, |
| 506 | ) -> None: |
| 507 | """SD state and indices reach the native operator unchanged.""" |
| 508 | num_slots = 2 |
| 509 | state_len = CONV_KERNEL - 1 |
| 510 | conv_state = _sd_conv_states(num_slots, state_len) |
| 511 | state_indices = torch.tensor([1], dtype=torch.int32) |
| 512 | x = torch.zeros(1, CONV_DIM, dtype=torch.bfloat16) |
| 513 | recorded = {} |
| 514 | |
| 515 | def record_update(**kwargs): |
| 516 | recorded.update(kwargs) |
| 517 | return x |
| 518 | |
| 519 | monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) |
| 520 | monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", record_update) |
| 521 | out = gdn_attention._causal_conv1d_update_cpu( |
| 522 | x=x, |
| 523 | conv_states=conv_state, |
| 524 | weight=torch.zeros(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16), |
| 525 | bias=None, |
| 526 | silu_activation=False, |
| 527 | conv_state_indices=state_indices, |
| 528 | is_weight_packed=False, |
| 529 | ) |
| 530 | |
| 531 | assert out is x |
| 532 | assert recorded["conv_states"] is conv_state |
| 533 | assert recorded["conv_state_indices"] is state_indices |
| 534 | |
| 535 | |
| 536 | @torch.inference_mode() |
| 537 | def test_spec_aware_mixed_routing_preserves_token_order( |
| 538 | monkeypatch: pytest.MonkeyPatch, |
| 539 | ) -> None: |
| 540 | num_tokens = 4 |
| 541 | projection = torch.arange(num_tokens * 16, dtype=torch.float32).view(num_tokens, 16) |
| 542 | mixed_qkv, b, a = projection[:, :4], projection[:, 4:6], projection[:, 6:8] |
| 543 | assert all(not tensor.is_contiguous() for tensor in (mixed_qkv, b, a)) |
| 544 | |
| 545 | spec_indices = torch.tensor([0, 2]) |
| 546 | nonspec_indices = torch.tensor([1, 3]) |
| 547 | metadata = GDNAttentionMetadata( |
| 548 | num_prefills=1, |
| 549 | num_prefill_tokens=0, |
| 550 | num_decodes=0, |
| 551 | num_decode_tokens=0, |
| 552 | num_spec_decodes=0, |
| 553 | num_spec_decode_tokens=0, |
| 554 | num_actual_tokens=0, |
| 555 | spec_sequence_masks=torch.ones(1, dtype=torch.bool), |
| 556 | spec_token_indx=spec_indices, |
| 557 | non_spec_token_indx=nonspec_indices, |
| 558 | ) |
| 559 | routed = [] |
| 560 | |
| 561 | def record(*args): |
| 562 | routed.append(args[2:5]) |
| 563 | return args[2] |
| 564 | |
| 565 | monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True) |
| 566 | monkeypatch.setattr(gdn_attention, "_spec_forward", record) |
| 567 | monkeypatch.setattr(gdn_attention, "_spec_aware_nonspec_subset", record) |
| 568 | |
| 569 | layer = types.SimpleNamespace( |
| 570 | kv_cache=[torch.empty(1, 4, 6), torch.empty(1, 1, 1, 1)] |
| 571 | ) |
| 572 | core_attn_out = torch.empty_like(mixed_qkv) |
| 573 | gdn_attention._cpu_gdn_attention_spec_aware( |
| 574 | layer=layer, |
| 575 | attn_metadata_i=metadata, |
| 576 | mixed_qkv=mixed_qkv, |
| 577 | b=b, |
| 578 | a=a, |
| 579 | core_attn_out=core_attn_out, |
| 580 | width=CONV_KERNEL, |
| 581 | state_len=6, |
| 582 | ) |
| 583 | |
| 584 | expected_inputs = (mixed_qkv, b, a) |
| 585 | assert len(routed) == 2 |
| 586 | for actual_inputs, indices in zip(routed, (spec_indices, nonspec_indices)): |
| 587 | for actual, expected in zip(actual_inputs, expected_inputs): |
| 588 | assert actual.is_contiguous() |
| 589 | torch.testing.assert_close(actual, expected.index_select(0, indices)) |
| 590 | torch.testing.assert_close(core_attn_out, mixed_qkv) |
| 591 | |
| 592 | |
| 593 | @torch.inference_mode() |
| 594 | def test_spec_aware_nonspec_materializes_state_indices( |
| 595 | monkeypatch: pytest.MonkeyPatch, |
| 596 | ) -> None: |
| 597 | block_table = torch.arange(8, dtype=torch.int32).view(2, 4) |
| 598 | state_indices = block_table[:, 0] |
| 599 | assert not state_indices.is_contiguous() |
| 600 | |
| 601 | metadata = GDNAttentionMetadata( |
| 602 | num_prefills=2, |
| 603 | num_prefill_tokens=4, |
| 604 | num_decodes=0, |
| 605 | num_decode_tokens=0, |
| 606 | num_spec_decodes=0, |
| 607 | num_spec_decode_tokens=0, |
| 608 | num_actual_tokens=4, |
| 609 | non_spec_state_indices_tensor=state_indices, |
| 610 | non_spec_query_start_loc=torch.tensor([0, 2, 4], dtype=torch.int32), |
| 611 | has_initial_state=torch.tensor([False, False]), |
| 612 | ) |
| 613 | |
| 614 | recorded_indices = None |
| 615 | |
| 616 | def causal_conv1d_fwd_cpu(**kwargs): |
| 617 | nonlocal recorded_indices |
| 618 | recorded_indices = kwargs["cache_indices"] |
| 619 | return kwargs["x"] |
| 620 | |
| 621 | def fused_gdn_gating_cpu(**kwargs): |
| 622 | return kwargs["a"], kwargs["b"] |
| 623 | |
| 624 | def chunk_gated_delta_rule_cpu(**kwargs): |
| 625 | out = torch.zeros(1, 4, 1, 1) |
| 626 | return out, kwargs["initial_state"] |
| 627 | |
| 628 | monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True) |
| 629 | monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) |
| 630 | monkeypatch.setattr( |
| 631 | gdn_attention.ops, "causal_conv1d_fwd_cpu", causal_conv1d_fwd_cpu |
| 632 | ) |
| 633 | monkeypatch.setattr(gdn_attention.ops, "fused_gdn_gating_cpu", fused_gdn_gating_cpu) |
| 634 | monkeypatch.setattr( |
| 635 | gdn_attention.ops, |
| 636 | "chunk_gated_delta_rule_cpu", |
| 637 | chunk_gated_delta_rule_cpu, |
| 638 | ) |
| 639 | |
| 640 | layer = types.SimpleNamespace( |
| 641 | activation="silu", |
| 642 | conv1d=types.SimpleNamespace(weight=torch.empty(0), bias=None), |
| 643 | A_log=torch.empty(0), |
| 644 | dt_bias=torch.empty(0), |
| 645 | rearrange_mixed_qkv=lambda x: ( |
| 646 | x[:, :1].view(1, 4, 1, 1), |
| 647 | x[:, :1].view(1, 4, 1, 1), |
| 648 | x[:, :1].view(1, 4, 1, 1), |
| 649 | ), |
| 650 | ) |
| 651 | gdn_attention._spec_aware_nonspec( |
| 652 | layer=layer, |
| 653 | attn_metadata_i=metadata, |
| 654 | mixed_qkv=torch.zeros(4, 4), |
| 655 | b=torch.zeros(4, 1), |
| 656 | a=torch.zeros(4, 1), |
| 657 | core_attn_out=torch.zeros(4, 1, 1), |
| 658 | conv_buf=torch.zeros(8, 1, 6), |
| 659 | ssm_state=torch.zeros(8, 1, 1, 1), |
| 660 | width=4, |
| 661 | ) |
| 662 | |
| 663 | assert recorded_indices is not None |
| 664 | assert recorded_indices.is_contiguous() |
| 665 | torch.testing.assert_close( |
| 666 | recorded_indices, torch.tensor([0, 4], dtype=torch.int32) |
| 667 | ) |
| 668 | |
| 669 | |
| 670 | @torch.inference_mode() |
| 671 | @pytest.mark.parametrize(("has_amx", "has_arm_bf16"), [(True, False), (False, True)]) |
| 672 | def test_spec_forward_prepares_native_conv_metadata( |
| 673 | has_amx: bool, |
| 674 | has_arm_bf16: bool, |
| 675 | monkeypatch: pytest.MonkeyPatch, |
| 676 | ) -> None: |
| 677 | block_table = torch.tensor([[0, 1], [4, 5]], dtype=torch.int32) |
| 678 | state_indices = block_table[:, 0] |
| 679 | accepted_counts = torch.tensor([1, 4], dtype=torch.int32) |
| 680 | assert not state_indices.is_contiguous() |
| 681 | metadata = GDNAttentionMetadata( |
| 682 | num_prefills=0, |
| 683 | num_prefill_tokens=0, |
| 684 | num_decodes=0, |
| 685 | num_decode_tokens=0, |
| 686 | num_spec_decodes=2, |
| 687 | num_spec_decode_tokens=8, |
| 688 | num_actual_tokens=8, |
| 689 | spec_state_indices_tensor=block_table, |
| 690 | spec_query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32), |
| 691 | num_accepted_tokens=accepted_counts, |
| 692 | ) |
| 693 | forwarded_indices = None |
| 694 | forwarded_counts = None |
| 695 | forwarded_is_weight_packed = None |
| 696 | |
| 697 | def causal_conv1d_update_cpu(**kwargs): |
| 698 | nonlocal forwarded_counts, forwarded_indices, forwarded_is_weight_packed |
| 699 | forwarded_indices = kwargs["conv_state_indices"] |
| 700 | forwarded_counts = kwargs["num_accepted_tokens"] |
| 701 | forwarded_is_weight_packed = kwargs["is_weight_packed"] |
| 702 | return kwargs["x"] |
| 703 | |
| 704 | monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: has_amx) |
| 705 | monkeypatch.setattr(gdn_attention, "is_arm_bf16", lambda: has_arm_bf16) |
| 706 | monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) |
| 707 | monkeypatch.setattr( |
| 708 | gdn_attention.ops, "causal_conv1d_update_cpu", causal_conv1d_update_cpu |
| 709 | ) |
| 710 | monkeypatch.setattr( |
| 711 | gdn_attention.ops, |
| 712 | "fused_sigmoid_gating_delta_rule_update_spec_cpu", |
| 713 | lambda **kwargs: kwargs["q"], |
| 714 | ) |
| 715 | |
| 716 | layer = types.SimpleNamespace( |
| 717 | activation="silu", |
| 718 | conv1d=types.SimpleNamespace(weight=torch.empty(1, CONV_KERNEL), bias=None), |
| 719 | A_log=None, |
| 720 | dt_bias=None, |
| 721 | rearrange_mixed_qkv=lambda x: (x.unsqueeze(0),) * 3, |
| 722 | ) |
| 723 | gdn_attention._spec_forward( |
| 724 | layer=layer, |
| 725 | attn_metadata_i=metadata, |
| 726 | mixed_qkv_spec=torch.zeros(8, 1, dtype=torch.bfloat16), |
| 727 | b_spec=torch.empty(0), |
| 728 | a_spec=torch.empty(0), |
| 729 | conv_buf=torch.empty(0), |
| 730 | ssm_state=torch.empty(0), |
| 731 | width=CONV_KERNEL, |
| 732 | state_len=0, |
| 733 | ) |
| 734 | |
| 735 | expected = ( |
| 736 | (forwarded_indices, torch.tensor([0, 4], dtype=torch.int32)), |
| 737 | (forwarded_counts, torch.tensor([1, 4], dtype=torch.int32)), |
| 738 | ) |
| 739 | for actual, reference in expected: |
| 740 | assert actual is not None |
| 741 | assert actual.is_contiguous() |
| 742 | assert actual.dtype == torch.int32 |
| 743 | assert forwarded_is_weight_packed is True |
| 744 | torch.testing.assert_close(actual, reference) |
| 745 | |
| 746 | |
| 747 | @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) |
| 748 | @torch.inference_mode() |
| 749 | def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> None: |
| 750 | """Non-AMX conv-state handoff: a two-call split (the second seeded via |
| 751 | ``has_initial_state=True`` from the conv_states the first wrote back) must |
| 752 | match the single-call result. |
| 753 | """ |
| 754 | from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( |
| 755 | causal_conv1d_fn_cpu as causal_conv1d_torch, |
| 756 | ) |
| 757 | |
| 758 | x, weight, bias = _conv_inputs(total_tokens) |
| 759 | state_len = CONV_KERNEL - 1 |
| 760 | # [num_slots, conv_dim, state_len]; slot 0 used here. |
| 761 | conv_states_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) |
| 762 | conv_states_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) |
| 763 | |
| 764 | # x is [conv_dim, T] for causal_conv1d_torch. |
| 765 | xt = x.transpose(0, 1).contiguous() |
| 766 | |
| 767 | out_full = causal_conv1d_torch( |
| 768 | x=xt, |
| 769 | weight=weight, |
| 770 | bias=bias, |
| 771 | conv_states=conv_states_full, |
| 772 | query_start_loc=torch.tensor([0, total_tokens], dtype=torch.int32), |
| 773 | cache_indices=torch.tensor([0], dtype=torch.int32), |
| 774 | has_initial_state=torch.tensor([False]), |
| 775 | activation="silu", |
| 776 | ) |
| 777 | |
| 778 | out1 = causal_conv1d_torch( |
| 779 | x=xt[:, :split], |
| 780 | weight=weight, |
| 781 | bias=bias, |
| 782 | conv_states=conv_states_split, |
| 783 | query_start_loc=torch.tensor([0, split], dtype=torch.int32), |
| 784 | cache_indices=torch.tensor([0], dtype=torch.int32), |
| 785 | has_initial_state=torch.tensor([False]), |
| 786 | activation="silu", |
| 787 | ) |
| 788 | out2 = causal_conv1d_torch( |
| 789 | x=xt[:, split:], |
| 790 | weight=weight, |
| 791 | bias=bias, |
| 792 | conv_states=conv_states_split, |
| 793 | query_start_loc=torch.tensor([0, total_tokens - split], dtype=torch.int32), |
| 794 | cache_indices=torch.tensor([0], dtype=torch.int32), |
| 795 | has_initial_state=torch.tensor([True]), |
| 796 | activation="silu", |
| 797 | ) |
| 798 | out_split = torch.cat([out1, out2], dim=1) |
| 799 | |
| 800 | torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2) |
| 801 | |
| 802 | |
| 803 | @pytest.mark.skipif( |
| 804 | not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()), |
| 805 | reason="native causal conv1d requires AMX or Arm BF16 support", |
| 806 | ) |
| 807 | @torch.inference_mode() |
| 808 | def test_causal_conv1d_update_cpu_accepts_wide_state() -> None: |
| 809 | state_len = CONV_KERNEL - 1 |
| 810 | wide_state_len = state_len + 5 |
| 811 | batch_size = 3 |
| 812 | is_weight_packed = True |
| 813 | x, weight, bias = _conv_inputs(batch_size) |
| 814 | conv_state_indices = torch.tensor([2, 0, 1], dtype=torch.int32) |
| 815 | |
| 816 | narrow_state = _sd_conv_states(batch_size, state_len) |
| 817 | narrow_state.copy_( |
| 818 | tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state) |
| 819 | ) |
| 820 | wide_state = _sd_conv_states(batch_size, wide_state_len) |
| 821 | wide_state[:, :, :state_len].copy_(narrow_state) |
| 822 | wide_state[:, :, state_len:].fill_(7) |
| 823 | wide_tail = wide_state[:, :, state_len:].clone() |
| 824 | |
| 825 | conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed) |
| 826 | out_narrow = ops.causal_conv1d_update_cpu( |
| 827 | x=x, |
| 828 | conv_states=narrow_state, |
| 829 | weight=conv_weight, |
| 830 | bias=bias, |
| 831 | silu_activation=True, |
| 832 | conv_state_indices=conv_state_indices, |
| 833 | is_weight_packed=is_weight_packed, |
| 834 | ) |
| 835 | out_wide = ops.causal_conv1d_update_cpu( |
| 836 | x=x, |
| 837 | conv_states=wide_state, |
| 838 | weight=conv_weight, |
| 839 | bias=bias, |
| 840 | silu_activation=True, |
| 841 | conv_state_indices=conv_state_indices, |
| 842 | is_weight_packed=is_weight_packed, |
| 843 | ) |
| 844 | |
| 845 | torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2) |
| 846 | torch.testing.assert_close( |
| 847 | wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0 |
| 848 | ) |
| 849 | torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0) |
| 850 | |
| 851 | |
| 852 | def _ref_causal_conv1d_update_cpu_multi( |
| 853 | x: torch.Tensor, |
| 854 | conv_states: torch.Tensor, |
| 855 | weight: torch.Tensor, |
| 856 | bias: torch.Tensor | None, |
| 857 | silu_activation: bool, |
| 858 | conv_state_indices: torch.Tensor, |
| 859 | num_accepted_tokens: torch.Tensor, |
| 860 | ) -> torch.Tensor: |
| 861 | batch_size, seq_len, dim = x.shape |
| 862 | state_len = conv_states.size(2) |
| 863 | conv_out = torch.empty_like(x) |
| 864 | conv_weight = weight.unsqueeze(1) |
| 865 | |
| 866 | for i in range(batch_size): |
| 867 | slot = int(conv_state_indices[i].item()) |
| 868 | offset = int(num_accepted_tokens[i].item()) - 1 |
| 869 | state = conv_states[slot] |
| 870 | x_seq = x[i].transpose(0, 1).to(state.dtype) |
| 871 | prior = state[:, offset : offset + CONV_KERNEL - 1] |
| 872 | conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) |
| 873 | out = F.conv1d(conv_in, conv_weight, bias, groups=dim)[0] |
| 874 | if silu_activation: |
| 875 | out = F.silu(out) |
| 876 | conv_out[i] = out.transpose(0, 1).to(conv_out.dtype) |
| 877 | keep = state[:, offset + 1 : offset + 1 + (state_len - seq_len)] |
| 878 | state.copy_(torch.cat([keep, x_seq], dim=-1)) |
| 879 | |
| 880 | return conv_out |
| 881 | |
| 882 | |
| 883 | @pytest.mark.skipif( |
| 884 | not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()), |
| 885 | reason="native causal conv1d requires AMX or Arm BF16 support", |
| 886 | ) |
| 887 | @pytest.mark.parametrize( |
| 888 | ( |
| 889 | "batch_size, " |
| 890 | "seq_len, " |
| 891 | "accepted_counts, " |
| 892 | "has_bias, " |
| 893 | "silu_activation, " |
| 894 | "is_weight_packed" |
| 895 | ), |
| 896 | [ |
| 897 | (1, 1, [1], False, False, False), |
| 898 | (1, 1, [1], True, True, True), |
| 899 | (4, 4, [1, 2, 3, 4], False, True, False), |
| 900 | (4, 4, [4, 3, 2, 1], True, False, True), |
| 901 | (4, 16, [1, 5, 10, 16], False, False, True), |
| 902 | (4, 16, [16, 10, 5, 1], True, True, False), |
| 903 | ], |
| 904 | ) |
| 905 | @torch.inference_mode() |
| 906 | def test_causal_conv1d_update_cpu_multi_token_matches_python( |
| 907 | batch_size: int, |
| 908 | seq_len: int, |
| 909 | accepted_counts: list[int], |
| 910 | has_bias: bool, |
| 911 | silu_activation: bool, |
| 912 | is_weight_packed: bool, |
| 913 | ) -> None: |
| 914 | dim = 96 |
| 915 | state_len = seq_len + 2 |
| 916 | x = tensor_cache(batch_size * seq_len * dim, torch.bfloat16).view( |
| 917 | batch_size, seq_len, dim |
| 918 | ) |
| 919 | weight = tensor_cache(dim * CONV_KERNEL, torch.bfloat16).view(dim, CONV_KERNEL) |
| 920 | bias = tensor_cache(dim, torch.bfloat16) if has_bias else None |
| 921 | conv_state_indices = torch.arange(batch_size - 1, -1, -1, dtype=torch.int32) |
| 922 | num_accepted_tokens = torch.tensor(accepted_counts, dtype=torch.int32) |
| 923 | |
| 924 | conv_states_ref = _sd_conv_states(batch_size, state_len, dim) |
| 925 | conv_states_ref.copy_( |
| 926 | tensor_cache(conv_states_ref.numel(), torch.bfloat16).view_as(conv_states_ref) |
| 927 | ) |
| 928 | conv_states = conv_states_ref.clone() |
| 929 | |
| 930 | conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed) |
| 931 | out = ops.causal_conv1d_update_cpu( |
| 932 | x=x, |
| 933 | conv_states=conv_states, |
| 934 | weight=conv_weight, |
| 935 | bias=bias, |
| 936 | silu_activation=silu_activation, |
| 937 | conv_state_indices=conv_state_indices, |
| 938 | is_weight_packed=is_weight_packed, |
| 939 | num_accepted_tokens=num_accepted_tokens, |
| 940 | ) |
| 941 | ref_out = _ref_causal_conv1d_update_cpu_multi( |
| 942 | x=x, |
| 943 | conv_states=conv_states_ref, |
| 944 | weight=weight, |
| 945 | bias=bias, |
| 946 | silu_activation=silu_activation, |
| 947 | conv_state_indices=conv_state_indices, |
| 948 | num_accepted_tokens=num_accepted_tokens, |
| 949 | ) |
| 950 | |
| 951 | torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2) |
| 952 | torch.testing.assert_close(conv_states, conv_states_ref, atol=0, rtol=0) |
| 953 | |
| 954 | |
| 955 | @pytest.mark.skipif( |
| 956 | not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()), |
| 957 | reason="native causal conv1d requires AMX or Arm BF16 support", |
| 958 | ) |
| 959 | @pytest.mark.parametrize("num_accepted", [0, 17]) |
| 960 | @torch.inference_mode() |
| 961 | def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count( |
| 962 | num_accepted: int, |
| 963 | ) -> None: |
| 964 | batch_size = 1 |
| 965 | seq_len = 16 |
| 966 | dim = 96 |
| 967 | state_len = seq_len + 2 |
| 968 | x = torch.zeros(batch_size, seq_len, dim, dtype=torch.bfloat16) |
| 969 | weight = torch.zeros(dim, CONV_KERNEL, dtype=torch.bfloat16) |
| 970 | conv_states = _sd_conv_states(batch_size, state_len, dim) |
| 971 | |
| 972 | with pytest.raises(RuntimeError, match="num_accepted_tokens must be in.*seqlen"): |
| 973 | ops.causal_conv1d_update_cpu( |
| 974 | x=x, |
| 975 | conv_states=conv_states, |
| 976 | weight=weight, |
| 977 | bias=None, |
| 978 | silu_activation=True, |
| 979 | conv_state_indices=torch.tensor([0], dtype=torch.int32), |
| 980 | is_weight_packed=False, |
| 981 | num_accepted_tokens=torch.tensor([num_accepted], dtype=torch.int32), |
| 982 | ) |
| 983 | |
| 984 | |
| 985 | @pytest.mark.skipif( |
| 986 | not (torch.cpu._is_avx512_bf16_supported() or gdn_attention.is_arm_bf16()), |
| 987 | reason="causal_conv1d_fwd_cpu requires AVX-512BF16" |
| 988 | " (Intel Xeon or AMD EPYC) or Arm BF16 support", |
| 989 | ) |
| 990 | @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) |
| 991 | @torch.inference_mode() |
| 992 | def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> None: |
| 993 | """C++ prefill conv op must honor ``has_initial_state`` so a two-call split |
| 994 | matches the single-call result. |
| 995 | |
| 996 | Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``) |
| 997 | ignoring the carried conv state on continued chunks. Runs on aarch64 cpus with |
| 998 | bf16 support or any AVX-512BF16 CPU since conv.cpp uses VDPBF16PS, not AMX tiles. |
| 999 | """ |
| 1000 | state_len = CONV_KERNEL - 1 |
| 1001 | x, weight, bias = _conv_inputs(total_tokens) |
| 1002 | |
| 1003 | def native(x_seg, conv_states, has_init): |
| 1004 | seq = x_seg.shape[0] |
| 1005 | return ops.causal_conv1d_fwd_cpu( |
| 1006 | x=x_seg.transpose(0, 1), # [dim, seq]; stride(-2)==1 (view of [seq,dim]) |
| 1007 | weight=weight, |
| 1008 | bias=bias, |
| 1009 | conv_states=conv_states, |
| 1010 | query_start_loc=torch.tensor([0, seq], dtype=torch.int32), |
| 1011 | cache_indices=torch.tensor([0], dtype=torch.int32), |
| 1012 | has_initial_state=torch.tensor([has_init]), |
| 1013 | silu_activation=True, |
| 1014 | is_weight_packed=False, |
| 1015 | ).contiguous() |
| 1016 | |
| 1017 | # conv_state layout passed by the AMX branch: [num_slots, dim, state_len]. |
| 1018 | cs_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) |
| 1019 | out_full = native(x, cs_full, False) |
| 1020 | |
| 1021 | cs_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) |
| 1022 | out1 = native(x[:split], cs_split, False) |
| 1023 | out2 = native(x[split:], cs_split, True) |
| 1024 | out_split = torch.cat([out1, out2], dim=1) |
| 1025 | |
| 1026 | torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2) |
| 1027 | |
| 1028 | |
| 1029 | @pytest.mark.skipif( |
| 1030 | not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()), |
| 1031 | reason="native causal conv1d requires AMX or Arm BF16 support", |
| 1032 | ) |
| 1033 | @torch.inference_mode() |
| 1034 | def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None: |
| 1035 | state_len = CONV_KERNEL - 1 |
| 1036 | wide_state_len = state_len + 5 |
| 1037 | is_weight_packed = True |
| 1038 | seq_lens = [CHUNK_SIZE - 1, CHUNK_SIZE + 5] |
| 1039 | total_tokens = sum(seq_lens) |
| 1040 | x, weight, bias = _conv_inputs(total_tokens) |
| 1041 | query_start_loc = torch.tensor([0, seq_lens[0], total_tokens], dtype=torch.int32) |
| 1042 | cache_indices = torch.tensor([2, 0], dtype=torch.int32) |
| 1043 | has_initial_state = torch.tensor([True, False]) |
| 1044 | |
| 1045 | narrow_state = _sd_conv_states(3, state_len) |
| 1046 | narrow_state.copy_( |
| 1047 | tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state) |
| 1048 | ) |
| 1049 | wide_state = _sd_conv_states(3, wide_state_len) |
| 1050 | wide_state[:, :, :state_len].copy_(narrow_state) |
| 1051 | wide_state[:, :, state_len:].fill_(7) |
| 1052 | wide_tail = wide_state[:, :, state_len:].clone() |
| 1053 | |
| 1054 | conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed) |
| 1055 | out_narrow = ops.causal_conv1d_fwd_cpu( |
| 1056 | x=x.transpose(0, 1), |
| 1057 | weight=conv_weight, |
| 1058 | bias=bias, |
| 1059 | conv_states=narrow_state, |
| 1060 | query_start_loc=query_start_loc, |
| 1061 | cache_indices=cache_indices, |
| 1062 | has_initial_state=has_initial_state, |
| 1063 | silu_activation=True, |
| 1064 | is_weight_packed=is_weight_packed, |
| 1065 | ) |
| 1066 | out_wide = ops.causal_conv1d_fwd_cpu( |
| 1067 | x=x.transpose(0, 1), |
| 1068 | weight=conv_weight, |
| 1069 | bias=bias, |
| 1070 | conv_states=wide_state, |
| 1071 | query_start_loc=query_start_loc, |
| 1072 | cache_indices=cache_indices, |
| 1073 | has_initial_state=has_initial_state, |
| 1074 | silu_activation=True, |
| 1075 | is_weight_packed=is_weight_packed, |
| 1076 | ) |
| 1077 | |
| 1078 | torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2) |
| 1079 | torch.testing.assert_close( |
| 1080 | wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0 |
| 1081 | ) |
| 1082 | torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0) |
| 1083 | |
| 1084 | |
| 1085 | @torch.inference_mode() |
| 1086 | def test_batch_memcpy_cpu_fallback() -> None: |
| 1087 | """The ctypes batch_memcpy fallback (used when triton-cpu is absent) must |
| 1088 | copy each src into its dst, validating the (src_ptrs, dst_ptrs, sizes) |
| 1089 | argument order against ctypes.memmove(dst, src, size). |
| 1090 | """ |
| 1091 | from vllm.triton_utils.dispatcher import ( |
| 1092 | KernelOverride, |
| 1093 | register_kernels, |
| 1094 | ) |
| 1095 | from vllm.utils.cpu_triton_utils import _batch_memcpy_impl |
| 1096 | |
| 1097 | register_kernels( |
| 1098 | {"vllm.v1.worker.mamba_utils.batch_memcpy_kernel": _batch_memcpy_impl} |
| 1099 | ) |
| 1100 | from vllm.v1.worker.mamba_utils import batch_memcpy_kernel |
| 1101 | |
| 1102 | assert isinstance(batch_memcpy_kernel, KernelOverride) |
| 1103 | |
| 1104 | # Varied byte sizes, including a non-power-of-two run. |
| 1105 | sizes_bytes = [256, 1024, 17 * 4, 4096] |
| 1106 | srcs = [torch.rand(n // 4, dtype=torch.float32) for n in sizes_bytes] |
| 1107 | dsts = [torch.zeros_like(s) for s in srcs] |
| 1108 | |
| 1109 | src_ptrs = torch.tensor([s.data_ptr() for s in srcs], dtype=torch.uint64) |
| 1110 | dst_ptrs = torch.tensor([d.data_ptr() for d in dsts], dtype=torch.uint64) |
| 1111 | sizes = torch.tensor(sizes_bytes, dtype=torch.int32) |
| 1112 | |
| 1113 | batch_memcpy_kernel[(len(srcs),)](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=1024) |
| 1114 | |
| 1115 | for src, dst in zip(srcs, dsts): |
| 1116 | torch.testing.assert_close(dst, src) |
| 1117 | |
| 1118 | |
| 1119 | # --------------------------------------------------------------------------- |
| 1120 | # C++ conv (conv.cpp) uses AArch64BF16 or VDPBF16PS, not AMX tiles, so it can |
| 1121 | # run on any AVX-512BF16 CPU; weight is packed on this same predicate at load time. |
| 1122 | # --------------------------------------------------------------------------- |
| 1123 | |
| 1124 | _HAS_AVX512_BF16 = torch.cpu._is_avx512_bf16_supported() |
| 1125 | |
| 1126 | _STATE_LEN = CONV_KERNEL - 1 |
| 1127 | |
| 1128 | CONV_EQUIV_SEQ_LENS = [ |
| 1129 | [1], |
| 1130 | [7], |
| 1131 | [64], |
| 1132 | [65], |
| 1133 | [1, 2, 3], |
| 1134 | [63, 64, 65], |
| 1135 | [128, 129], |
| 1136 | ] |
| 1137 | |
| 1138 | |
| 1139 | def _conv_fp32_oracle(x, weight, bias, seq_lens, activation="silu"): |
| 1140 | """High-precision conv reference: everything in float32, no bf16 rounding. |
| 1141 | x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim].""" |
| 1142 | xf = x.float() |
| 1143 | wf = weight.float().unsqueeze(1) |
| 1144 | bf = bias.float() |
| 1145 | out = torch.empty_like(xf) |
| 1146 | start = 0 |
| 1147 | for n in seq_lens: |
| 1148 | seg = xf[start : start + n].transpose(0, 1).unsqueeze(0) # [1, dim, n] |
| 1149 | conv_in = F.pad(seg, (_STATE_LEN, 0)) |
| 1150 | seg_out = F.conv1d(conv_in, wf, bf, padding=0, groups=CONV_DIM)[..., -n:] |
| 1151 | if activation in ("silu", "swish"): |
| 1152 | seg_out = F.silu(seg_out) |
| 1153 | out[start : start + n] = seg_out.squeeze(0).transpose(0, 1) |
| 1154 | start += n |
| 1155 | return out |
| 1156 | |
| 1157 | |
| 1158 | def _run_prefill_torch(x, weight, bias, seq_lens): |
| 1159 | from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( |
| 1160 | causal_conv1d_fn_cpu as causal_conv1d_torch, |
| 1161 | ) |
| 1162 | |
| 1163 | num_seqs = len(seq_lens) |
| 1164 | conv_states = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype) |
| 1165 | qsl = torch.tensor( |
| 1166 | [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 |
| 1167 | ) |
| 1168 | out = causal_conv1d_torch( |
| 1169 | x=x.transpose(0, 1).contiguous(), |
| 1170 | weight=weight, |
| 1171 | bias=bias, |
| 1172 | conv_states=conv_states, |
| 1173 | query_start_loc=qsl, |
| 1174 | cache_indices=torch.arange(num_seqs, dtype=torch.int32), |
| 1175 | has_initial_state=torch.zeros(num_seqs, dtype=torch.bool), |
| 1176 | activation="silu", |
| 1177 | ) |
| 1178 | return out.transpose(0, 1).contiguous(), conv_states |
| 1179 | |
| 1180 | |
| 1181 | def _run_prefill_cpp(x, weight, bias, seq_lens, is_weight_packed=False): |
| 1182 | num_seqs = len(seq_lens) |
| 1183 | packed_w = ops.causal_conv1d_weight_pack(weight) if is_weight_packed else weight |
| 1184 | if is_weight_packed: |
| 1185 | # C++-branch layout: kv-cache "SD" [slots, state_len, dim] transposed to |
| 1186 | # [slots, dim, state_len] (a non-contiguous view). |
| 1187 | conv_state = torch.zeros( |
| 1188 | num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype |
| 1189 | ).transpose(1, 2) |
| 1190 | else: |
| 1191 | conv_state = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype) |
| 1192 | qsl = torch.tensor( |
| 1193 | [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 |
| 1194 | ) |
| 1195 | out = ops.causal_conv1d_fwd_cpu( |
| 1196 | x=x.transpose(0, 1), |
| 1197 | weight=packed_w, |
| 1198 | bias=bias, |
| 1199 | conv_states=conv_state, |
| 1200 | query_start_loc=qsl, |
| 1201 | cache_indices=torch.arange(num_seqs, dtype=torch.int32), |
| 1202 | has_initial_state=torch.zeros(num_seqs, dtype=torch.bool), |
| 1203 | silu_activation=True, |
| 1204 | is_weight_packed=is_weight_packed, |
| 1205 | ) |
| 1206 | return out.transpose(0, 1).contiguous(), conv_state |
| 1207 | |
| 1208 | |
| 1209 | def _python_native_fwd(**kwargs): |
| 1210 | from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( |
| 1211 | causal_conv1d_fn_cpu, |
| 1212 | ) |
| 1213 | |
| 1214 | return causal_conv1d_fn_cpu( |
| 1215 | x=kwargs["x"], |
| 1216 | weight=kwargs["weight"], |
| 1217 | bias=kwargs["bias"], |
| 1218 | conv_states=kwargs["conv_states"], |
| 1219 | query_start_loc=kwargs["query_start_loc"], |
| 1220 | cache_indices=kwargs["cache_indices"], |
| 1221 | has_initial_state=kwargs["has_initial_state"], |
| 1222 | activation=kwargs["silu_activation"], |
| 1223 | ) |
| 1224 | |
| 1225 | |
| 1226 | @torch.inference_mode() |
| 1227 | def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state( |
| 1228 | monkeypatch: pytest.MonkeyPatch, |
| 1229 | ) -> None: |
| 1230 | dim = 8 |
| 1231 | seq_lens = [2, 3] |
| 1232 | state_len = _STATE_LEN |
| 1233 | state_indices = torch.tensor([2, 0], dtype=torch.int32) |
| 1234 | query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32) |
| 1235 | has_initial_state = torch.tensor([False, True]) |
| 1236 | x_tokens, weight, bias = _conv_inputs(sum(seq_lens)) |
| 1237 | x_tokens = x_tokens[:, :dim] |
| 1238 | weight = weight[:dim] |
| 1239 | bias = bias[:dim] |
| 1240 | x = x_tokens.transpose(0, 1).contiguous() |
| 1241 | logical_state = ( |
| 1242 | torch.arange(3 * dim * state_len, dtype=torch.float32) |
| 1243 | .view(3, dim, state_len) |
| 1244 | .to(torch.bfloat16) |
| 1245 | ) |
| 1246 | state_sd = _conv_states(3, state_len, dim, "SD") |
| 1247 | state_ds = _conv_states(3, state_len, dim, "DS") |
| 1248 | state_sd.copy_(logical_state) |
| 1249 | state_ds.copy_(logical_state) |
| 1250 | recorded = [] |
| 1251 | |
| 1252 | def record_fwd(**kwargs): |
| 1253 | recorded.append( |
| 1254 | ( |
| 1255 | kwargs["conv_states"].shape, |
| 1256 | kwargs["conv_states"].stride(), |
| 1257 | kwargs["cache_indices"].clone(), |
| 1258 | ) |
| 1259 | ) |
| 1260 | return _python_native_fwd(**kwargs) |
| 1261 | |
| 1262 | monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_fwd_cpu", record_fwd) |
| 1263 | monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) |
| 1264 | out_sd = gdn_attention._causal_conv1d_fwd_cpu( |
| 1265 | x=x, |
| 1266 | weight=weight, |
| 1267 | bias=bias, |
| 1268 | conv_states=state_sd, |
| 1269 | query_start_loc=query_start_loc, |
| 1270 | cache_indices=state_indices, |
| 1271 | has_initial_state=has_initial_state, |
| 1272 | silu_activation=True, |
| 1273 | is_weight_packed=False, |
| 1274 | ) |
| 1275 | monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True) |
| 1276 | out_ds = gdn_attention._causal_conv1d_fwd_cpu( |
| 1277 | x=x, |
| 1278 | weight=weight, |
| 1279 | bias=bias, |
| 1280 | conv_states=state_ds, |
| 1281 | query_start_loc=query_start_loc, |
| 1282 | cache_indices=state_indices, |
| 1283 | has_initial_state=has_initial_state, |
| 1284 | silu_activation=True, |
| 1285 | is_weight_packed=False, |
| 1286 | ) |
| 1287 | |
| 1288 | torch.testing.assert_close(out_ds, out_sd) |
| 1289 | torch.testing.assert_close(state_ds, state_sd) |
| 1290 | assert recorded[0][2].equal(state_indices) |
| 1291 | assert recorded[1][0] == (2, dim, state_len) |
| 1292 | assert recorded[1][1] == (dim * state_len, 1, dim) |
| 1293 | assert recorded[1][2].equal(torch.arange(2, dtype=torch.int32)) |
| 1294 | |
| 1295 | |
| 1296 | @torch.inference_mode() |
| 1297 | def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback( |
| 1298 | monkeypatch: pytest.MonkeyPatch, |
| 1299 | ) -> None: |
| 1300 | dim = 8 |
| 1301 | seq_len = 4 |
| 1302 | state_len = seq_len + 5 |
| 1303 | state = _conv_states(4, state_len, dim, "DS") |
| 1304 | state.copy_( |
| 1305 | torch.arange(state.numel(), dtype=torch.float32).view_as(state).to(state.dtype) |
| 1306 | ) |
| 1307 | original_state = state.clone() |
| 1308 | state_indices = torch.tensor([3, 1, 0], dtype=torch.int32) |
| 1309 | num_accepted_tokens = torch.tensor([1, 4, 2], dtype=torch.int32) |
| 1310 | mutated_scratch = torch.arange( |
| 1311 | state_indices.numel() * dim * state_len, dtype=torch.float32 |
| 1312 | ).view(state_indices.numel(), dim, state_len) |
| 1313 | mutated_scratch = (mutated_scratch + 200).to(state.dtype) |
| 1314 | recorded = {} |
| 1315 | |
| 1316 | def mutate_update(**kwargs): |
| 1317 | recorded.update(kwargs) |
| 1318 | kwargs["conv_states"].copy_(mutated_scratch) |
| 1319 | return torch.zeros_like(kwargs["x"]) |
| 1320 | |
| 1321 | monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True) |
| 1322 | monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", mutate_update) |
| 1323 | gdn_attention._causal_conv1d_update_cpu( |
| 1324 | x=torch.zeros(state_indices.numel(), seq_len, dim, dtype=state.dtype), |
| 1325 | conv_states=state, |
| 1326 | weight=torch.zeros(dim, CONV_KERNEL, dtype=state.dtype), |
| 1327 | bias=None, |
| 1328 | silu_activation=False, |
| 1329 | conv_state_indices=state_indices, |
| 1330 | is_weight_packed=False, |
| 1331 | num_accepted_tokens=num_accepted_tokens, |
| 1332 | ) |
| 1333 | |
| 1334 | expected_state = original_state.clone() |
| 1335 | expected_state.index_copy_(0, state_indices.to(torch.long), mutated_scratch) |
| 1336 | torch.testing.assert_close(state, expected_state) |
| 1337 | torch.testing.assert_close( |
| 1338 | state.index_select(0, state_indices)[:, :, _STATE_LEN:], |
| 1339 | mutated_scratch[:, :, _STATE_LEN:], |
| 1340 | ) |
| 1341 | torch.testing.assert_close( |
| 1342 | recorded["conv_state_indices"], |
| 1343 | torch.arange(state_indices.numel(), dtype=torch.int32), |
| 1344 | ) |
| 1345 | torch.testing.assert_close(recorded["num_accepted_tokens"], num_accepted_tokens) |
| 1346 | |
| 1347 | |
| 1348 | @pytest.mark.skipif( |
| 1349 | not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()), |
| 1350 | reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support", |
| 1351 | ) |
| 1352 | @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS) |
| 1353 | @torch.inference_mode() |
| 1354 | def test_conv_cpp_matches_torch(seq_lens): |
| 1355 | """C++ causal_conv1d_fwd_cpu matches the torch fallback within bf16 tol.""" |
| 1356 | x, weight, bias = _conv_inputs(sum(seq_lens)) |
| 1357 | out_torch, state_torch = _run_prefill_torch(x, weight, bias, seq_lens) |
| 1358 | out_cpp, state_cpp = _run_prefill_cpp(x, weight, bias, seq_lens) |
| 1359 | torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2) |
| 1360 | torch.testing.assert_close(state_cpp, state_torch, atol=1e-2, rtol=1e-2) |
| 1361 | |
| 1362 | |
| 1363 | @pytest.mark.skipif( |
| 1364 | not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()), |
| 1365 | reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support", |
| 1366 | ) |
| 1367 | @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS) |
| 1368 | @torch.inference_mode() |
| 1369 | def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens): |
| 1370 | """Swapping torch -> C++ conv must not increase error vs an fp32 oracle.""" |
| 1371 | x, weight, bias = _conv_inputs(sum(seq_lens)) |
| 1372 | oracle = _conv_fp32_oracle(x, weight, bias, seq_lens) |
| 1373 | out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens) |
| 1374 | out_cpp, _ = _run_prefill_cpp(x, weight, bias, seq_lens) |
| 1375 | err_torch = (out_torch.float() - oracle).abs().mean().item() |
| 1376 | err_cpp = (out_cpp.float() - oracle).abs().mean().item() |
| 1377 | assert err_cpp <= err_torch + 1e-3, ( |
| 1378 | f"C++ conv less accurate than torch: " |
| 1379 | f"err_cpp={err_cpp:.2e} err_torch={err_torch:.2e}" |
| 1380 | ) |
| 1381 | |
| 1382 | |
| 1383 | @pytest.mark.skipif( |
| 1384 | not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()), |
| 1385 | reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support", |
| 1386 | ) |
| 1387 | @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS) |
| 1388 | @torch.inference_mode() |
| 1389 | def test_conv_cpp_weight_packed_matches_torch(seq_lens): |
| 1390 | """The exact runtime prefill sequence (packed weight + SD-layout |
| 1391 | conv_state view + is_weight_packed=True) must match the torch fallback. Validates |
| 1392 | the native packing + layout handoff.""" |
| 1393 | x, weight, bias = _conv_inputs(sum(seq_lens)) |
| 1394 | out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens) |
| 1395 | out_weight_packed, _ = _run_prefill_cpp( |
| 1396 | x, weight, bias, seq_lens, is_weight_packed=True |
| 1397 | ) |
| 1398 | torch.testing.assert_close(out_weight_packed, out_torch, atol=1e-2, rtol=1e-2) |
| 1399 | |
| 1400 | |
| 1401 | @pytest.mark.skipif( |
| 1402 | not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()), |
| 1403 | reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support", |
| 1404 | ) |
| 1405 | @pytest.mark.parametrize("batch", DECODE_BATCH_SIZES) |
| 1406 | @pytest.mark.parametrize( |
| 1407 | ("dim", "has_bias", "silu_activation", "is_weight_packed"), |
| 1408 | [ |
| 1409 | (32, False, False, False), |
| 1410 | (32, True, True, True), |
| 1411 | (96, False, True, True), |
| 1412 | (96, True, False, False), |
| 1413 | (8192, False, True, False), |
| 1414 | (8192, True, True, True), |
| 1415 | ], |
| 1416 | ) |
| 1417 | @torch.inference_mode() |
| 1418 | def test_conv_update_cpp_matches_torch( |
| 1419 | batch, dim, has_bias, silu_activation, is_weight_packed |
| 1420 | ): |
| 1421 | """Decode preserves indexed, padded state across successive signed inputs.""" |
| 1422 | from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( |
| 1423 | causal_conv1d_update_torch, |
| 1424 | ) |
| 1425 | |
| 1426 | torch.manual_seed(0) |
| 1427 | weight = torch.randn(dim, CONV_KERNEL, dtype=torch.bfloat16) |
| 1428 | bias = torch.randn(dim, dtype=torch.bfloat16) if has_bias else None |
| 1429 | conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed) |
| 1430 | indices = torch.arange(batch, 0, -1, dtype=torch.int32) |
| 1431 | storage_ref = torch.randn(batch + 2, _STATE_LEN + 2, dim, dtype=torch.bfloat16) |
| 1432 | storage_cpp = storage_ref.clone() |
| 1433 | cs_cpp = storage_cpp[:, :_STATE_LEN].transpose(1, 2) |
| 1434 | |
| 1435 | for _ in range(3): |
| 1436 | x = torch.randn(batch, dim, dtype=torch.bfloat16) |
| 1437 | cs_torch = storage_ref[indices, :_STATE_LEN].transpose(1, 2).contiguous() |
| 1438 | out_torch = causal_conv1d_update_torch( |
| 1439 | x=x.unsqueeze(-1), |
| 1440 | conv_state=cs_torch, |
| 1441 | weight=weight, |
| 1442 | bias=bias, |
| 1443 | activation="silu" if silu_activation else None, |
| 1444 | ).squeeze(-1) |
| 1445 | storage_ref[indices, :_STATE_LEN] = cs_torch.transpose(1, 2) |
| 1446 | out_cpp = ops.causal_conv1d_update_cpu( |
| 1447 | x=x, |
| 1448 | conv_states=cs_cpp, |
| 1449 | weight=conv_weight, |
| 1450 | bias=bias, |
| 1451 | silu_activation=silu_activation, |
| 1452 | conv_state_indices=indices, |
| 1453 | is_weight_packed=is_weight_packed, |
| 1454 | ) |
| 1455 | torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2) |
| 1456 | torch.testing.assert_close(storage_cpp, storage_ref, atol=0, rtol=0) |
| 1457 | |
| 1458 | |
| 1459 | @pytest.mark.skipif( |
| 1460 | not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()), |
| 1461 | reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support", |
| 1462 | ) |
| 1463 | @torch.inference_mode() |
| 1464 | def test_conv_weight_pack_roundtrip_unpacked_matches(): |
| 1465 | """`causal_conv1d_weight_pack` repacks the (dim, width) weight, so the C++ conv op |
| 1466 | with packed weight + is_weight_packed=True must equal unpacked weight |
| 1467 | + is_weight_packed=False. Guards the spec-decode contract: the runtime packs |
| 1468 | `layer.conv1d.weight` in place and stashes the original as |
| 1469 | `_cpu_unpacked_conv_weight` for the torch spec-decode path, so both must produce |
| 1470 | identical math. |
| 1471 | """ |
| 1472 | seq_lens = [7, 64, 65] |
| 1473 | x, weight, bias = _conv_inputs(sum(seq_lens)) |
| 1474 | out_unpacked, state_unpacked = _run_prefill_cpp( |
| 1475 | x, weight, bias, seq_lens, is_weight_packed=False |
| 1476 | ) |
| 1477 | out_packed, state_packed = _run_prefill_cpp( |
| 1478 | x, weight, bias, seq_lens, is_weight_packed=True |
| 1479 | ) |
| 1480 | torch.testing.assert_close(out_packed, out_unpacked, atol=0, rtol=0) |
| 1481 | torch.testing.assert_close(state_packed, state_unpacked, atol=0, rtol=0) |
| 1482 | |
| 1483 | |
| 1484 | @pytest.mark.skipif( |
| 1485 | not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()), |
| 1486 | reason="C++ causal_conv1d requires AVX-512BF16 or AArch64 BF16", |
| 1487 | ) |
| 1488 | @torch.inference_mode() |
| 1489 | def test_spec_decode_unpacked_conv_weight_stash(): |
| 1490 | """Spec-decode correctness: AVX-512BF16 and AArch64 BF16 CPUs pack ``conv1d.weight`` |
| 1491 | in place at load time and stash the original ``(dim, width)`` tensor as |
| 1492 | ``_cpu_unpacked_conv_weight``; the spec-decode path must use that stash since |
| 1493 | reading the packed weight directly produces garbage. Verifies the recovered |
| 1494 | weight equals the original and that torch F.conv1d agrees with the C++ conv |
| 1495 | fed the packed weight. |
| 1496 | """ |
| 1497 | from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm |
| 1498 | |
| 1499 | torch.manual_seed(0) |
| 1500 | # conv1d weight is stored [dim, 1, width]; bias [dim]. |
| 1501 | orig_2d = torch.rand(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16) |
| 1502 | bias = torch.rand(CONV_DIM, dtype=torch.bfloat16) |
| 1503 | |
| 1504 | conv = torch.nn.Module() |
| 1505 | conv.weight = torch.nn.Parameter( |
| 1506 | orig_2d.view(CONV_DIM, 1, CONV_KERNEL).clone(), requires_grad=False |
| 1507 | ) |
| 1508 | |
| 1509 | # Load-time dispatch: packs weight in place + stashes the unpacked copy. |
| 1510 | dispatch_cpu_unquantized_gemm(conv, remove_weight=False) |
| 1511 | |
| 1512 | # 1. The stash must exist and equal the original (dim, width) weight. |
| 1513 | assert hasattr(conv, "_cpu_unpacked_conv_weight"), ( |
| 1514 | "dispatch_cpu_unquantized_gemm did not stash _cpu_unpacked_conv_weight " |
| 1515 | "on an AArch64-BF16 or AVX-512BF16 CPU" |
| 1516 | ) |
| 1517 | stash = conv._cpu_unpacked_conv_weight |
| 1518 | torch.testing.assert_close(stash, orig_2d, atol=0, rtol=0) |
| 1519 | |
| 1520 | # 2. Mirror _unpacked_conv_weight()'s lookup: stash wins over conv.weight. |
| 1521 | recovered = getattr(conv, "_cpu_unpacked_conv_weight", None) |
| 1522 | assert recovered is not None |
| 1523 | # conv.weight is now packed; using it directly (the bug) would differ. |
| 1524 | packed_weight = conv.weight # [dim, 1, width], packed contents |
| 1525 | |
| 1526 | # 3. Spec-path torch conv with the recovered unpacked weight must match the |
| 1527 | # nonspec C++ conv with the packed weight (is_weight_packed=True). |
| 1528 | seq_lens = [7, 65] |
| 1529 | total = sum(seq_lens) |
| 1530 | x = tensor_cache(total * CONV_DIM, torch.bfloat16).view(total, CONV_DIM) |
| 1531 | |
| 1532 | out_torch, _ = _run_prefill_torch(x, recovered, bias, seq_lens) |
| 1533 | |
| 1534 | num_seqs = len(seq_lens) |
| 1535 | cs = torch.zeros(num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype).transpose(1, 2) |
| 1536 | qsl = torch.tensor( |
| 1537 | [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 |
| 1538 | ) |
| 1539 | out_cpp = ( |
| 1540 | ops.causal_conv1d_fwd_cpu( |
| 1541 | x=x.transpose(0, 1), |
| 1542 | weight=packed_weight.view(CONV_DIM, CONV_KERNEL), |
| 1543 | bias=bias, |
| 1544 | conv_states=cs, |
| 1545 | query_start_loc=qsl, |
| 1546 | cache_indices=torch.arange(num_seqs, dtype=torch.int32), |
| 1547 | has_initial_state=torch.zeros(num_seqs, dtype=torch.bool), |
| 1548 | silu_activation=True, |
| 1549 | is_weight_packed=True, |
| 1550 | ) |
| 1551 | .transpose(0, 1) |
| 1552 | .contiguous() |
| 1553 | ) |
| 1554 | |
| 1555 | torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2) |