[Perf] Fuse DeepEncoder relative bias in Triton attention (#55629)
Signed-off-by: levius <2114377220@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Sun committed
Sep 8, 2026 at 21:22 UTC
472d8c912566b539182c7154f582eaf1ba84bab1
1 file changed
+191
-7
vllm/model_executor/models/deepencoder.py
+191
-7
@@ -22,11 +22,177 @@ from vllm.model_executor.custom_op import PluggableLayer
22
from vllm.model_executor.layers.attention import MMEncoderAttention
23
from vllm.model_executor.layers.conv import Conv2dLayer
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
+from vllm.platforms import current_platform
26
+from vllm.triton_utils import LOG2E, tl, triton
27
28
from .clip import CLIPEncoder, CLIPVisionEmbeddings
29
from .utils import AutoWeightsLoader
30
31
32
+@triton.jit
33
+def _deepencoder_rel_pos_attention_kernel(
34
+ q_ptr,
35
+ k_ptr,
36
+ v_ptr,
37
+ rel_h_ptr,
38
+ rel_w_ptr,
39
+ out_ptr,
40
+ scale,
41
+ stride_qb,
42
+ stride_qh,
43
+ stride_qm,
44
+ stride_qd,
45
+ stride_kb,
46
+ stride_kh,
47
+ stride_kn,
48
+ stride_kd,
49
+ stride_vb,
50
+ stride_vh,
51
+ stride_vn,
52
+ stride_vd,
53
+ stride_rhb,
54
+ stride_rhh,
55
+ stride_rhm,
56
+ stride_rhk,
57
+ stride_rwb,
58
+ stride_rwh,
59
+ stride_rwm,
60
+ stride_rwk,
61
+ stride_ob,
62
+ stride_oh,
63
+ stride_om,
64
+ stride_od,
65
+ num_tokens: tl.constexpr,
66
+ num_heads: tl.constexpr,
67
+ key_width: tl.constexpr,
68
+ BLOCK_M: tl.constexpr,
69
+ BLOCK_N: tl.constexpr,
70
+ BLOCK_D: tl.constexpr,
71
+):
72
+ query_block = tl.program_id(0)
73
+ batch_head = tl.program_id(1)
74
+ batch = batch_head // num_heads
75
+ head = batch_head % num_heads
76
+
77
+ offs_m = query_block * BLOCK_M + tl.arange(0, BLOCK_M)
78
+ offs_n = tl.arange(0, BLOCK_N)
79
+ offs_d = tl.arange(0, BLOCK_D)
80
+ mask_m = offs_m < num_tokens
81
+
82
+ q_offsets = (
83
+ batch * stride_qb
84
+ + head * stride_qh
85
+ + offs_m[:, None] * stride_qm
86
+ + offs_d[None, :] * stride_qd
87
+ )
88
+ q = tl.load(q_ptr + q_offsets, mask=mask_m[:, None], other=0.0)
89
+
90
+ row_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
91
+ row_sum = tl.zeros([BLOCK_M], tl.float32)
92
+ accumulator = tl.zeros([BLOCK_M, BLOCK_D], tl.float32)
93
+ qk_scale = scale * LOG2E
94
+
95
+ for start_n in range(0, num_tokens, BLOCK_N):
96
+ key_indices = start_n + offs_n
97
+ mask_n = key_indices < num_tokens
98
+ k_offsets = (
99
+ batch * stride_kb
100
+ + head * stride_kh
101
+ + key_indices[None, :] * stride_kn
102
+ + offs_d[:, None] * stride_kd
103
+ )
104
+ k = tl.load(k_ptr + k_offsets, mask=mask_n[None, :], other=0.0)
105
+ scores = tl.dot(q, k) * qk_scale
106
+
107
+ key_h = key_indices // key_width
108
+ key_w = key_indices % key_width
109
+ rel_h_offsets = (
110
+ batch * stride_rhb
111
+ + head * stride_rhh
112
+ + offs_m[:, None] * stride_rhm
113
+ + key_h[None, :] * stride_rhk
114
+ )
115
+ rel_w_offsets = (
116
+ batch * stride_rwb
117
+ + head * stride_rwh
118
+ + offs_m[:, None] * stride_rwm
119
+ + key_w[None, :] * stride_rwk
120
+ )
121
+ score_mask = mask_m[:, None] & mask_n[None, :]
122
+ rel_h = tl.load(rel_h_ptr + rel_h_offsets, mask=score_mask, other=0.0)
123
+ rel_w = tl.load(rel_w_ptr + rel_w_offsets, mask=score_mask, other=0.0)
124
+ scores += (rel_h + rel_w) * LOG2E
125
+ scores = tl.where(score_mask, scores, -float("inf"))
126
+
127
+ next_max = tl.maximum(row_max, tl.max(scores, axis=1))
128
+ probabilities = tl.math.exp2(scores - next_max[:, None])
129
+ alpha = tl.math.exp2(row_max - next_max)
130
+ next_sum = row_sum * alpha + tl.sum(probabilities, axis=1)
131
+ accumulator *= alpha[:, None]
132
+
133
+ v_offsets = (
134
+ batch * stride_vb
135
+ + head * stride_vh
136
+ + key_indices[:, None] * stride_vn
137
+ + offs_d[None, :] * stride_vd
138
+ )
139
+ value = tl.load(v_ptr + v_offsets, mask=mask_n[:, None], other=0.0)
140
+ accumulator = tl.dot(probabilities.to(value.dtype), value, accumulator)
141
+ row_max = next_max
142
+ row_sum = next_sum
143
+
144
+ output = accumulator / row_sum[:, None]
145
+ out_offsets = (
146
+ batch * stride_ob
147
+ + head * stride_oh
148
+ + offs_m[:, None] * stride_om
149
+ + offs_d[None, :] * stride_od
150
+ )
151
+ tl.store(out_ptr + out_offsets, output, mask=mask_m[:, None])
152
+
153
+
154
+def deepencoder_rel_pos_attention(
155
+ q: torch.Tensor,
156
+ k: torch.Tensor,
157
+ v: torch.Tensor,
158
+ rel_h: torch.Tensor,
159
+ rel_w: torch.Tensor,
160
+ width: int,
161
+ scale: float,
162
+) -> torch.Tensor:
163
+ """Run DeepEncoder global attention with fused relative-position bias."""
164
+ batch, heads, num_tokens, head_dim = q.shape
165
+
166
+ output = torch.empty_like(q)
167
+ block_m = 64
168
+ block_n = 64
169
+ grid = (triton.cdiv(num_tokens, block_m), batch * heads)
170
+ _deepencoder_rel_pos_attention_kernel[grid](
171
+ q,
172
+ k,
173
+ v,
174
+ rel_h,
175
+ rel_w,
176
+ output,
177
+ scale,
178
+ *q.stride(),
179
+ *k.stride(),
180
+ *v.stride(),
181
+ *rel_h.stride(),
182
+ *rel_w.stride(),
183
+ *output.stride(),
184
+ num_tokens=num_tokens,
185
+ num_heads=heads,
186
+ key_width=width,
187
+ BLOCK_M=block_m,
188
+ BLOCK_N=block_n,
189
+ BLOCK_D=head_dim,
190
+ num_warps=4,
191
+ num_stages=2,
192
+ )
193
+ return output
194
+
195
+
196
class MLPBlock(nn.Module):
197
def __init__(
198
self,
@@ -234,6 +400,7 @@ class Block(nn.Module):
400
num_heads=num_heads,
401
qkv_bias=qkv_bias,
402
use_rel_pos=use_rel_pos,
403
+ use_triton_attention=window_size == 0,
404
rel_pos_zero_init=rel_pos_zero_init,
405
input_size=input_size if window_size == 0 else (window_size, window_size),
406
)
@@ -277,6 +444,7 @@ class RelPosAttention(PluggableLayer):
444
num_heads: int = 8,
445
qkv_bias: bool = True,
446
use_rel_pos: bool = False,
447
+ use_triton_attention: bool = False,
448
rel_pos_zero_init: bool = True,
449
input_size: tuple[int, int] | None = None,
450
) -> None:
@@ -285,6 +453,8 @@ class RelPosAttention(PluggableLayer):
453
dim (int): Number of input channels.
454
num_heads (int): Number of attention heads.
455
qkv_bias (bool): If True, add a learnable bias to query, key, value.
456
+ use_triton_attention (bool): If True, fuse relative position bias in
457
+ a Triton attention kernel on CUDA-alike platforms.
458
rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
459
input_size (tuple(int, int) or None): Input resolution for calculating the relative
460
positional parameter size.
@@ -298,6 +468,7 @@ class RelPosAttention(PluggableLayer):
468
self.proj = nn.Linear(dim, dim)
469
470
self.use_rel_pos = use_rel_pos
471
+ self.use_triton_attention = use_triton_attention
472
if self.use_rel_pos:
473
assert input_size is not None, (
474
"Input size must be provided if using relative positional encoding."
@@ -333,14 +504,27 @@ class RelPosAttention(PluggableLayer):
504
rel_w = rel_w.view(
505
B, self.num_heads, rel_w.size(1), rel_w.size(2), rel_w.size(3)
506
)
336
- attn_bias = (rel_h + rel_w).view(
337
- B, self.num_heads, rel_h.size(2), rel_h.size(3) * rel_w.size(4)
338
- )
339
- x = torch.nn.functional.scaled_dot_product_attention(
340
- q, k, v, attn_mask=attn_bias
341
- )
507
+ use_triton = self.use_triton_attention and current_platform.is_cuda_alike()
508
+ if use_triton:
509
+ x = deepencoder_rel_pos_attention(
510
+ q,
511
+ k,
512
+ v,
513
+ rel_h.squeeze(-1),
514
+ rel_w.squeeze(-2),
515
+ W,
516
+ self.scale,
517
+ )
518
+ else:
519
+ attn_bias = (rel_h + rel_w).view(
520
+ B,
521
+ self.num_heads,
522
+ rel_h.size(2),
523
+ rel_h.size(3) * rel_w.size(4),
524
+ )
525
+ x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias)
526
else:
343
- x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
527
+ x = F.scaled_dot_product_attention(q, k, v)
528
529
x = (
530
x.view(B, self.num_heads, H, W, -1)