| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.distributed.parallel_state import ( |
| 8 | destroy_distributed_environment, |
| 9 | destroy_model_parallel, |
| 10 | init_distributed_environment, |
| 11 | initialize_model_parallel, |
| 12 | ) |
| 13 | from vllm.model_executor.models.conformer_encoder import RelPosMultiHeadAttention |
| 14 | |
| 15 | pytestmark = pytest.mark.cpu_test |
| 16 | |
| 17 | |
| 18 | class _ReferenceRelPosMultiHeadAttention(RelPosMultiHeadAttention): |
| 19 | """Pre-baddbmm relative attention implementation.""" |
| 20 | |
| 21 | def forward( |
| 22 | self, |
| 23 | q: torch.Tensor, |
| 24 | k: torch.Tensor, |
| 25 | v: torch.Tensor, |
| 26 | pos_emb: torch.Tensor, |
| 27 | mask: torch.Tensor | None = None, |
| 28 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 29 | batch_size, query_len = q.size(0), q.size(1) |
| 30 | residual = q |
| 31 | q, k, v = self.forward_qkv(q, k, v) |
| 32 | |
| 33 | q = q.transpose(1, 2) |
| 34 | position_batch = pos_emb.size(0) |
| 35 | p = self.linear_pos(pos_emb)[0].view(position_batch, -1, self.n_head, self.d_k) |
| 36 | p = p.transpose(1, 2) |
| 37 | |
| 38 | q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) |
| 39 | q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) |
| 40 | |
| 41 | matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) |
| 42 | matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) |
| 43 | matrix_bd = self._reference_rel_shift(matrix_bd) |
| 44 | |
| 45 | attention_scores = matrix_ac + matrix_bd |
| 46 | attention_scores.mul_(self.scale) |
| 47 | |
| 48 | output, attention = self.forward_attention(attention_scores, v, mask=mask) |
| 49 | output = self.forward_output(output, residual, batch_size, query_len) |
| 50 | return output, attention |
| 51 | |
| 52 | @staticmethod |
| 53 | def _reference_rel_shift(x: torch.Tensor) -> torch.Tensor: |
| 54 | batch, heads, query_len, position_len = x.shape |
| 55 | zero_pad = torch.zeros( |
| 56 | (batch, heads, query_len, 1), device=x.device, dtype=x.dtype |
| 57 | ) |
| 58 | x = torch.cat([zero_pad, x], dim=-1) |
| 59 | x = x.view(batch, heads, position_len + 1, query_len) |
| 60 | x = x[:, :, 1:].view(batch, heads, query_len, position_len) |
| 61 | return x[:, :, :, : position_len // 2 + 1] |
| 62 | |
| 63 | |
| 64 | @pytest.fixture |
| 65 | def tp1_cpu(tmp_path, default_vllm_config): |
| 66 | """Initialize a temporary single-rank CPU tensor-parallel group.""" |
| 67 | try: |
| 68 | init_distributed_environment( |
| 69 | world_size=1, |
| 70 | rank=0, |
| 71 | distributed_init_method=f"file://{tmp_path / 'distributed'}", |
| 72 | local_rank=0, |
| 73 | backend="gloo", |
| 74 | ) |
| 75 | initialize_model_parallel(tensor_model_parallel_size=1) |
| 76 | yield |
| 77 | finally: |
| 78 | destroy_model_parallel() |
| 79 | destroy_distributed_environment() |
| 80 | |
| 81 | |
| 82 | @pytest.mark.parametrize("sequence", [2, 8]) |
| 83 | def test_relative_attention_matches_reference(tp1_cpu, sequence: int) -> None: |
| 84 | torch.manual_seed(7) |
| 85 | batch, heads, model_dim = 2, 4, 32 |
| 86 | optimized = RelPosMultiHeadAttention(heads, model_dim) |
| 87 | reference = _ReferenceRelPosMultiHeadAttention(heads, model_dim) |
| 88 | |
| 89 | with torch.no_grad(): |
| 90 | for parameter in optimized.parameters(): |
| 91 | parameter.uniform_(-0.2, 0.2) |
| 92 | reference.load_state_dict(optimized.state_dict()) |
| 93 | |
| 94 | query = torch.randn(batch, sequence, model_dim) |
| 95 | key = torch.randn(batch, sequence, model_dim) |
| 96 | value = torch.randn(batch, sequence, model_dim) |
| 97 | pos_emb = torch.randn(1, 2 * sequence - 1, model_dim) |
| 98 | mask = torch.ones(batch, 1, sequence, dtype=torch.bool) |
| 99 | mask[:, :, -1] = False |
| 100 | |
| 101 | expected_output, expected_attention = reference( |
| 102 | query, key, value, pos_emb, mask=mask |
| 103 | ) |
| 104 | output, attention = optimized(query, key, value, pos_emb, mask=mask) |
| 105 | |
| 106 | torch.testing.assert_close(output, expected_output) |
| 107 | torch.testing.assert_close(attention, expected_attention) |
| 108 | assert output.shape == (batch, sequence, model_dim) |
| 109 | assert attention.shape == (batch, heads, sequence, sequence) |
| 110 | assert output.is_contiguous() |
| 111 | assert attention.is_contiguous() |