[Perf] Accumulate Conformer attention scores with baddbmm (#55062)
Signed-off-by: levius <2114377220@qq.com> Co-authored-by: Codex <codex@openai.com>
Sun committed
Sep 3, 2026 at 19:17 UTC
848ab131bcdb5264bcff0d802f47f7d4adb0f548
2 files changed
+124
-3
tests/models/multimodal/test_conformer_encoder.py
new
+111
@@ -0,0 +1,111 @@
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()
vllm/model_executor/models/conformer_encoder.py
+13
-3
@@ -206,12 +206,22 @@ class RelPosMultiHeadAttention(EncoderMultiHeadAttention):
206
q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
207
q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
208
209
- matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
209
matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
210
matrix_bd = self._rel_shift(matrix_bd)
211
213
- attn_scores = matrix_ac + matrix_bd
214
- attn_scores.mul_(self.scale)
212
+ batch_size, num_heads, query_len, key_len = matrix_bd.shape
213
+ attn_scores = torch.baddbmm(
214
+ matrix_bd.reshape(-1, query_len, key_len),
215
+ q_with_bias_u.reshape(-1, query_len, self.d_k),
216
+ k.transpose(-2, -1).reshape(-1, self.d_k, key_len),
217
+ beta=self.scale,
218
+ alpha=self.scale,
219
+ ).view(
220
+ batch_size,
221
+ num_heads,
222
+ query_len,
223
+ key_len,
224
+ )
225
226
output, attn = self.forward_attention(attn_scores, v, mask=mask)
227
output = self.forward_output(output, residual, sz_b, len_q)