[Perf] Fuse CohereASR relative attention score accumulation (#55190)
Signed-off-by: levius <2114377220@qq.com> Co-authored-by: Codex <codex@openai.com>
Sun committed
Sep 18, 2026 at 14:18 UTC
42a33039eb43c2adc885329a1c26ee1a5d8cc0fb
2 files changed
+69
-5
tests/models/multimodal/test_cohere_asr.py
+57
-1
@@ -4,9 +4,65 @@
4
from types import SimpleNamespace
5
6
import pytest
7
+import torch
8
9
from vllm.config import SpeechToTextConfig
9
-from vllm.model_executor.models.cohere_asr import CohereAsrForConditionalGeneration
10
+from vllm.model_executor.models.cohere_asr import (
11
+ CohereAsrForConditionalGeneration,
12
+ RelPositionMultiHeadAttention,
13
+)
14
+
15
+
16
+class _ReferenceRelPositionMultiHeadAttention(RelPositionMultiHeadAttention):
17
+ """Relative attention implementation before score accumulation is fused."""
18
+
19
+ def forward(self, query, key, value, mask, pos_emb=None):
20
+ q, k, v = self.forward_qkv(query, key, value)
21
+ q = q.transpose(1, 2)
22
+
23
+ assert pos_emb is not None
24
+ n_batch_pos = pos_emb.size(0)
25
+ p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
26
+ p = p.transpose(1, 2)
27
+
28
+ q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
29
+ q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
30
+
31
+ matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
32
+ matrix_bd = self.rel_shift(matrix_bd)
33
+ matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
34
+ matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
35
+ scores = (matrix_ac + matrix_bd) / self.s_d_k
36
+ return self.forward_attention(v, scores, mask)
37
+
38
+
39
+@pytest.mark.parametrize(("query_len", "key_len"), [(2, 2), (3, 5)])
40
+def test_relative_attention_fuses_score_accumulation(monkeypatch, query_len, key_len):
41
+ torch.manual_seed(7)
42
+ batch, heads, model_dim = 2, 4, 32
43
+ optimized = RelPositionMultiHeadAttention(heads, model_dim, None, None)
44
+ reference = _ReferenceRelPositionMultiHeadAttention(heads, model_dim, None, None)
45
+ reference.load_state_dict(optimized.state_dict())
46
+
47
+ query = torch.randn(batch, query_len, model_dim)
48
+ key = torch.randn(batch, key_len, model_dim)
49
+ value = torch.randn(batch, key_len, model_dim)
50
+ pos_emb = torch.randn(1, query_len + key_len - 1, model_dim)
51
+ expected = reference(query, key, value, None, pos_emb)
52
+
53
+ original_baddbmm = torch.baddbmm
54
+ calls = 0
55
+
56
+ def tracked_baddbmm(*args, **kwargs):
57
+ nonlocal calls
58
+ calls += 1
59
+ return original_baddbmm(*args, **kwargs)
60
+
61
+ monkeypatch.setattr(torch, "baddbmm", tracked_baddbmm)
62
+ actual = optimized(query, key, value, None, pos_emb)
63
+
64
+ torch.testing.assert_close(actual, expected)
65
+ assert calls == 1
66
67
68
@pytest.mark.parametrize(
vllm/model_executor/models/cohere_asr.py
+12
-4
@@ -1206,10 +1206,18 @@ class RelPositionMultiHeadAttention(CohereASRMultiHeadAttention):
1206
matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
1207
matrix_bd = self.rel_shift(matrix_bd)
1208
1209
- # drops extra elements in the matrix_bd to match the matrix_ac's size
1210
- matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
1211
- matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
1212
- scores = (matrix_ac + matrix_bd) / self.s_d_k # (batch, head, time1, time2)
1209
+ # drops extra elements in matrix_bd to match the key sequence length
1210
+ matrix_bd = matrix_bd[:, :, :, : k.size(-2)]
1211
+ batch, heads, query_len, key_len = matrix_bd.shape
1212
+ scale = 1.0 / self.s_d_k
1213
+ scores = torch.baddbmm(
1214
+ matrix_bd.flatten(0, 1),
1215
+ q_with_bias_u.flatten(0, 1),
1216
+ k.transpose(-2, -1).flatten(0, 1),
1217
+ beta=scale,
1218
+ alpha=scale,
1219
+ ).view(batch, heads, query_len, key_len)
1220
+
1221
return self.forward_attention(v, scores, mask)
1222
1223