执行摘要

• 一句话：优化DeepSeek-OCR-2注意力掩码计算，TTFT降46%
• 推荐动作：值得精读。该PR清晰展示了如何通过分析数据模式（批次不变性）消除冗余CPU→GPU操作，并利用@lru_cache以极小成本实现加速。设计决策中关于缓存key范围的取舍（保留dtype和device）体现了对实际部署场景的考量。

功能与动机

create_custom_4d_mask函数在CPU上循环构造掩码，是TTFT的瓶颈。PR body中指明该函数是CPU-bound，torch profiler显示aten::_index_put_impl_调用达229,296次，消耗52.2%的CPU时间。优化目标是将这一计算移至GPU并消除冗余索引操作。

实现拆解

1. 移除实例状态：删除CustomQwen2ModelInner.forward中的_current_token_type_ids实例变量，不再传递token_type_ids，因为mask生成不再依赖每个样本的token_type_ids。
2. 新增类方法compute_mask_base：在CustomQwen2ModelInner中声明类方法，使用@classmethod和@lru_cache(maxsize=8)缓存。该方法根据固定模式（前一半image token、后一半text token）在GPU上一次性计算[1,1,S,S]掩码基座。
3. 简化_create_custom_4d_mask：原方法遍历batch、逐位置填充mask。改为调用compute_mask_base获得基座，然后通过expand(batch_size,-1,-1,-1)广播到batch维度。
4. 调整调用链：_update_causal_mask不再传递token_type_ids，直接调用_create_custom_4d_mask并处理padding mask。
5. 性能验证：未添加单元测试，但PR中附带了详细的TTFT基准脚本和分析，证明优化效果。

关键文件：

• vllm/model_executor/models/deepencoder2.py（模块 视觉编码；类别 source；类型 core-logic；符号 CustomQwen2Decoder, compute_mask_base, _create_custom_4d_mask, _update_causal_mask）: 核心性能优化文件，重构了掩码计算和缓存逻辑

关键符号：compute_mask_base, _create_custom_4d_mask, _update_causal_mask

关键源码片段

vllm/model_executor/models/deepencoder2.py

核心性能优化文件，重构了掩码计算和缓存逻辑

from functools import lru_cache
import torch
​
# ... 在 CustomQwen2ModelInner 类中 ...
​
@classmethod
@lru_cache(maxsize=8)
def compute_mask_base(cls, sequence_length: int, dtype: torch.dtype, device: torch.device):
    """
    Compute the base 4D attention mask for DeepSeek-OCR-2.
    The mask is batch-invariant because token_type_ids follows a fixed pattern:
    first half of tokens are image tokens (non-causal), second half are text tokens (causal).
    This method computes a single [1, 1, S, S] mask and relies on expand()
    in the caller to broadcast over batch dimension.
    """
    min_dtype = torch.finfo(dtype).min
    n_query = sequence_length // 2 # image token count (also text token count)
    # Indices of image tokens (first half)
    img = torch.arange(sequence_length, device=device) < n_query
    txt = ~img
    # Standard causal mask: lower triangular
    causal = torch.tril(torch.ones(sequence_length, sequence_length,
                                    dtype=torch.bool, device=device))
    # Allow: image token attends to all; text token attends to all image tokens
    # plus causal among text tokens.
    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] # add head and batch dims
​
def _create_custom_4d_mask(self, sequence_length, dtype, device, batch_size):
    """Compute the full mask by expanding the cached base to batch size."""
    base = self.compute_mask_base(sequence_length, dtype, device)
    return base.expand(batch_size, -1, -1, -1)

评论区精华

Review中，Isotr0py对缓存有效性提出质疑，担心sequence_length变化导致缓存低效；LiuLi1998回应实际只有288和512两个值，缓存有界。随后Isotr0py建议改用@classmethod+@lru_cache，并提交commit实现。最终方案保留了dtype和device作为缓存key，以兼容不同配置。

• 缓存key范围和实现方式 (design): 保留(sequence_length, dtype, device)作为缓存key，采用@lru_cache，LiuLi1998测试通过。
• 批次不变性假设验证 (correctness): 通过基准测试和profile确认正确性。

风险与影响

• 风险：风险较低。缓存key包括(sequence_length, dtype, device)，典型部署下组合有限；lru_cache的maxsize=8限制了条目数。无动态控制流变化，正确性通过基准测试验证。未修改batch维度的计算，扩展不会引入新回归。若未来模型变种导致更多sequence_length值，需注意缓存命中率，但当前场景无虞。
• 影响：直接影响使用DeepSeek-OCR-2模型的用户，TTFT显著降低（p50降46%），吞吐量提升83%。对其他模型无影响。团队维护成本极低——仅一个文件改动，无新增依赖，部署无需额外配置。
• 风险标记：缓存假设依赖固定token_type_ids模式, 仅针对OCR-2模型优化

关联脉络

• 暂无明显关联 PR