[Perf] DeepSeek-OCR-2 TTFT Optimize (#49531)
Signed-off-by: RED <outofthewoods@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
RED committed
Jul 26, 2026 at 13:53 UTC
30b0714031ae24f77832099cf120576dae28b17c
1 file changed
+32
-36
vllm/model_executor/models/deepencoder2.py
+32
-36
@@ -10,6 +10,8 @@
10
# This source code is licensed under the license found in the
11
# LICENSE file in the root directory of this source tree.
12
13
+from functools import lru_cache
14
+
15
import torch
16
import torch.nn as nn
17
import transformers
@@ -90,8 +92,6 @@ class CustomQwen2Decoder(PluggableLayer):
92
return_dict=None,
93
cache_position=None,
94
):
93
- # token_type_ids
94
- self._current_token_type_ids = token_type_ids
95
causal_mask_mapping = {
96
"full_attention": self._update_causal_mask(
97
attention_mask,
@@ -131,15 +131,12 @@ class CustomQwen2Decoder(PluggableLayer):
131
input_tensor.shape[1],
132
)
133
134
- token_type_ids = self._current_token_type_ids
135
-
134
# attention mask
135
causal_mask = self._create_custom_4d_mask(
136
sequence_length=sequence_length,
137
dtype=dtype,
138
device=device,
139
batch_size=batch_size,
142
- token_type_ids=token_type_ids,
140
)
141
142
# padding mask
@@ -150,44 +147,43 @@ class CustomQwen2Decoder(PluggableLayer):
147
148
return causal_mask
149
150
+ @classmethod
151
+ @lru_cache(maxsize=8)
152
+ def compute_mask_base(cls, sequence_length, dtype, device):
153
+ # token_type_ids is the fixed pattern [0]*n_query + [1]*n_query,
154
+ # identical across the batch, so the mask depends only on
155
+ # sequence_length: img tokens (first half) attend to
156
+ # everything, txt tokens (second half) attend causally among
157
+ # themselves. lru_cache keeps one batch-invariant [1, 1, S, S]
158
+ # mask per (S, dtype, device).
159
+ min_dtype = torch.finfo(dtype).min
160
+ n_query = sequence_length // 2
161
+ img = torch.arange(sequence_length, device=device) < n_query
162
+ txt = ~img
163
+ causal = torch.tril(
164
+ torch.ones(
165
+ sequence_length,
166
+ sequence_length,
167
+ dtype=torch.bool,
168
+ device=device,
169
+ )
170
+ )
171
+ allow = img[None, :] | (txt[:, None] & txt[None, :] & causal)
172
+ return torch.where(
173
+ allow,
174
+ torch.zeros((), dtype=dtype, device=device),
175
+ torch.full((), min_dtype, dtype=dtype, device=device),
176
+ )[None, None]
177
+
178
def _create_custom_4d_mask(
179
self,
180
sequence_length,
181
dtype,
182
device,
183
batch_size,
159
- token_type_ids,
184
):
161
- min_dtype = torch.finfo(dtype).min
162
-
163
- masks = []
164
- for b in range(batch_size):
165
- mask = torch.full(
166
- (sequence_length, sequence_length),
167
- fill_value=min_dtype,
168
- dtype=dtype,
169
- device=device,
170
- )
171
-
172
- type_ids = token_type_ids[b]
173
-
174
- image_positions = (type_ids == 0).nonzero(as_tuple=True)[0]
175
- text_positions = (type_ids == 1).nonzero(as_tuple=True)[0]
176
-
177
- # non-casual
178
- if len(image_positions) > 0:
179
- mask[image_positions[:, None], image_positions] = 0.0
180
-
181
- # causal
182
- for i, text_pos in enumerate(text_positions):
183
- if len(image_positions) > 0:
184
- mask[text_pos, image_positions] = 0.0
185
- mask[text_pos, text_positions[: i + 1]] = 0.0
186
-
187
- masks.append(mask)
188
-
189
- mask = torch.stack(masks, dim=0).unsqueeze(1)
190
- return mask
185
+ base = self.compute_mask_base(sequence_length, dtype, device)
186
+ return base.expand(batch_size, -1, -1, -1)
187
188
return CustomQwen2ModelInner(config)
189