iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] DeepSeek-OCR-2 TTFT Optimize (#49531)

Browse Source 
Signed-off-by: RED <outofthewoods@qq.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
RED
2026-07-26 05:53:06 +00:00
committed by GitHub
co-authored by Isotr0py
parent 2e860de498
commit 30b0714031
1 changed files with 32 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/models/deepencoder2.py
+32 -36
View File
@@ -10,6 +10,8 @@
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from functools import lru_cache
import torch
import torch.nn as nn
import transformers
@@ -90,8 +92,6 @@ class CustomQwen2Decoder(PluggableLayer):
return_dict=None,
cache_position=None,
):
# token_type_ids
self._current_token_type_ids = token_type_ids
causal_mask_mapping = {
"full_attention": self._update_causal_mask(
attention_mask,
@@ -131,15 +131,12 @@ class CustomQwen2Decoder(PluggableLayer):
input_tensor.shape[1],
)
token_type_ids = self._current_token_type_ids
# attention mask
causal_mask = self._create_custom_4d_mask(
sequence_length=sequence_length,
dtype=dtype,
device=device,
batch_size=batch_size,
token_type_ids=token_type_ids,
)
# padding mask
@@ -150,44 +147,43 @@ class CustomQwen2Decoder(PluggableLayer):
return causal_mask
@classmethod
@lru_cache(maxsize=8)
def compute_mask_base(cls, sequence_length, dtype, device):
# token_type_ids is the fixed pattern [0]*n_query + [1]*n_query,
# identical across the batch, so the mask depends only on
# sequence_length: img tokens (first half) attend to
# everything, txt tokens (second half) attend causally among
# themselves. lru_cache keeps one batch-invariant [1, 1, S, S]
# mask per (S, dtype, device).
min_dtype = torch.finfo(dtype).min
n_query = sequence_length // 2
img = torch.arange(sequence_length, device=device) < n_query
txt = ~img
causal = torch.tril(
torch.ones(
sequence_length,
sequence_length,
dtype=torch.bool,
device=device,
)
)
allow = img[None, :] | (txt[:, None] & txt[None, :] & causal)
return torch.where(
allow,
torch.zeros((), dtype=dtype, device=device),
torch.full((), min_dtype, dtype=dtype, device=device),
)[None, None]
def _create_custom_4d_mask(
self,
sequence_length,
dtype,
device,
batch_size,
token_type_ids,
):
min_dtype = torch.finfo(dtype).min
masks = []
for b in range(batch_size):
mask = torch.full(
(sequence_length, sequence_length),
fill_value=min_dtype,
dtype=dtype,
device=device,
)
type_ids = token_type_ids[b]
image_positions = (type_ids == 0).nonzero(as_tuple=True)[0]
text_positions = (type_ids == 1).nonzero(as_tuple=True)[0]
# non-casual
if len(image_positions) > 0:
mask[image_positions[:, None], image_positions] = 0.0
# causal
for i, text_pos in enumerate(text_positions):
if len(image_positions) > 0:
mask[text_pos, image_positions] = 0.0
mask[text_pos, text_positions[: i + 1]] = 0.0
masks.append(mask)
mask = torch.stack(masks, dim=0).unsqueeze(1)
return mask
base = self.compute_mask_base(sequence_length, dtype, device)
return base.expand(batch_size, -1, -1, -1)
return CustomQwen2ModelInner(config)