| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | import torch.nn.functional as F |
| 6 | from torch import Tensor |
| 7 | |
| 8 | import vllm._custom_ops as ops |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.utils.math_utils import cdiv |
| 11 | |
| 12 | |
| 13 | def ref_mla( |
| 14 | out: Tensor, # (bs, num_heads, v_head_dim) |
| 15 | query: Tensor, # (bs, num_heads, head_dim) |
| 16 | kv_cache: Tensor, # (num_blocks, block_size, head_dim) |
| 17 | scale: float, |
| 18 | block_tables: Tensor, # (bs, max_num_blocks) |
| 19 | seq_lens: Tensor, # (bs,) |
| 20 | ): |
| 21 | bs, num_heads, v_head_dim = out.shape |
| 22 | head_dim = query.shape[2] |
| 23 | |
| 24 | for i in range(bs): |
| 25 | # gather and flatten KV-cache |
| 26 | kv = kv_cache[block_tables[i]] # (max_num_blocks, block_size, head_dim) |
| 27 | kv = kv.view(1, -1, head_dim)[:, : seq_lens[i]] # (1, seq_len, head_dim) |
| 28 | v = kv[:, :, :v_head_dim] |
| 29 | |
| 30 | q = query[i].view(num_heads, 1, head_dim) |
| 31 | o = F.scaled_dot_product_attention(q, kv, v, scale=scale, enable_gqa=True) |
| 32 | out[i] = o.view(num_heads, v_head_dim) |
| 33 | |
| 34 | return out |
| 35 | |
| 36 | |
| 37 | @pytest.mark.parametrize("bs", [4]) |
| 38 | @pytest.mark.parametrize("mean_seq_len", [256]) |
| 39 | @pytest.mark.parametrize("h_q", [16]) |
| 40 | @pytest.mark.parametrize("d", [576]) |
| 41 | @pytest.mark.parametrize("dv", [512]) |
| 42 | @pytest.mark.parametrize("block_size", [16]) |
| 43 | @pytest.mark.parametrize("dtype", [torch.float, torch.half, torch.bfloat16]) |
| 44 | @pytest.mark.parametrize("varlen", [False, True]) |
| 45 | @pytest.mark.cpu_model |
| 46 | @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only") |
| 47 | def test_mla_decode_cpu( |
| 48 | bs: int, |
| 49 | mean_seq_len: int, |
| 50 | h_q: int, |
| 51 | d: int, |
| 52 | dv: int, |
| 53 | block_size: int, |
| 54 | dtype: torch.dtype, |
| 55 | varlen: bool, |
| 56 | ): |
| 57 | torch.set_default_dtype(dtype) |
| 58 | torch.manual_seed(0) |
| 59 | |
| 60 | scale = d ** (-0.5) |
| 61 | if varlen: |
| 62 | seq_lens = torch.empty(bs).normal_(mean_seq_len, mean_seq_len / 2) |
| 63 | seq_lens = seq_lens.clip(2).to(torch.int32) |
| 64 | else: |
| 65 | seq_lens = torch.full((bs,), mean_seq_len, dtype=torch.int32) |
| 66 | max_seq_len = seq_lens.max().item() |
| 67 | seqlen_pad = cdiv(max_seq_len, 256) * 256 # is this necessary? |
| 68 | |
| 69 | q = torch.randn(bs, h_q, d) |
| 70 | block_table = torch.arange(bs * seqlen_pad // block_size, dtype=torch.int32) |
| 71 | block_table = block_table.view(bs, seqlen_pad // block_size) |
| 72 | |
| 73 | kv_cache = torch.randn(block_table.numel(), block_size, d) |
| 74 | for i, seq_len in enumerate(seq_lens.tolist()): |
| 75 | kv_cache.view(bs, seqlen_pad, d)[i, seq_len:] = float("nan") |
| 76 | |
| 77 | out_mla = q.new_zeros(bs, h_q, dv) |
| 78 | ops.mla_decode_kvcache_cpu(out_mla, q, kv_cache, scale, block_table, seq_lens) |
| 79 | |
| 80 | out_ref = q.new_zeros(bs, h_q, dv) |
| 81 | ref_mla(out_ref, q, kv_cache, scale, block_table, seq_lens) |
| 82 | |
| 83 | assert not out_mla.isnan().any(), "Likely read out of bounds" |
| 84 | torch.testing.assert_close(out_mla, out_ref) |