执行摘要

• 一句话：修复 Kimi-VL 编码器数据并行与 TP 逻辑
• 推荐动作：该 PR 值得精读，尤其是编码器 DP/TP 集成的模式、CUDA graph 捕获中条件逃逸设计，以及位置推断缓存的 LRU 实现。对于需要支持同类多模态模型的开发者有借鉴意义。

功能与动机

修复 Kimi-VL 在 TP=1 编码器-DP 模式下投影仪拿到列表而非张量，导致图像请求失败的问题。启用 DP 路径并保持 TP 路径正确。

实现拆解

1. MoonViT 层 TP 化：在 kimi_vl_moonvit.py 中为 MoonVitEncoderLayer 引入 ColumnParallelLinear、QKVParallelLinear、RowParallelLinear，添加 use_tensor_parallel 参数；在 multihead_attention 和 sdpa_attention 中新增 max_seqlen 参数，避免每层 GPU 同步；Learnable2DInterpPosEmb 添加推理 LRU 缓存（最大 256 条），避免重复 bicubic 插值。
2. 模型主干并行分发：在 kimi_vl.py 的 __init__ 中根据服务器参数 mm_enable_dp_encoder 设置 use_data_parallel，并据此配置 MoonVitPretrainedModel；get_image_feature 区分 DP 和 TP 路径，DP 时调用 run_dp_sharded_mrope_vision_model，否则直接 vision tower 调用并预计算 max_seqlen。
3. CUDA Graph 捕获条件：vit_cuda_graph_runner.py 中 _capture_context 感知 use_data_parallel，DP 模式下不进入 TP 通信捕获，避免 DP 分片图像工作时捕获 TP 集合通信。
4. DP 分片辅助函数升级：mm_utils.py 中 run_dp_sharded_mrope_vision_model 新增 rope_type 参数支持 Kimi-VL 的 rope_2d，并增加 keep_token_shape 配置。
5. 图像处理器 kwargs 过滤优化：hf_transformers_patches.py 重构 safe_call，缓存每个处理器类的接受 kwargs 集，避免每次调用反射和异常捕获，提升批图像预处理性能。
6. 配套测试：新增 CPU-only 测试覆盖编码器并行性、位置缓存淘汰、CUDA graph 上下文切换、kwargs 过滤缓存。

关键文件：

• python/sglang/srt/models/kimi_vl_moonvit.py（模块 模型层；类别 source；类型 core-logic；符号 init, multihead_attention, Learnable2DInterpPosEmb）: 核心变更，包括 TP 线性层引入、注意力同步优化、位置推断缓存
• python/sglang/srt/models/kimi_vl.py（模块 模型层；类别 source；类型 data-contract；符号 get_image_feature, init）: 修改了编码器并行路径分发和图像特征获取
• python/sglang/srt/multimodal/vit_cuda_graph_runner.py（模块 CUDA图；类别 source；类型 core-logic；符号 _capture_context）: CUDA graph 捕获上下文根据 DP 模式切换
• python/sglang/srt/utils/hf_transformers_patches.py（模块 HF补丁；类别 source；类型 core-logic；符号 safe_call）: 优化图像处理器 kwargs 过滤，缓存已接受的参数集合
• test/registered/unit/models/test_kimi_vl.py（模块 Kimi-VL 测试；类别 test；类型 test-coverage；符号 _VisionTower, init, call, _Projector）: 新增编码器并行性 CPU-only 测试

关键符号：multihead_attention, Learnable2DInterpPosEmb.forward, KimiVLForConditionalGeneration.get_image_feature, KimiVLForConditionalGeneration.init, MoonVitEncoderLayer, ViTCudaGraphRunner._capture_context, run_dp_sharded_mrope_vision_model, safe_call

关键源码片段

python/sglang/srt/models/kimi_vl_moonvit.py

核心变更，包括 TP 线性层引入、注意力同步优化、位置推断缓存

# multihead_attention: 避免每层 GPU 同步，通过 max_seqlen 参数预计算值传入
def multihead_attention(
    q, k, v,
    q_cu_seqlens=None, k_cu_seqlens=None,
    max_seqlen: Optional[int] = None, # 新参数，由调用方预计算
):
    # ... flash_attn 导入等 ...
    # 仅在 CPU 上验证形状，避免 GPU 同步（.item() 触发 CUDA 同步）
    if not q_cu_seqlens.is_cuda:
        assert q_cu_seqlens[-1] == q.shape[0], "q_cu_seqlens must sum to q.shape[0]"
        assert k_cu_seqlens[-1] == k.shape[0] == v.shape[0]
    if max_seqlen is None:
        max_seqlen = (q_cu_seqlens[1:] - q_cu_seqlens[:-1]).max().item()
    attn_out = flash_attn_varlen_func(
        q, k, v,
        q_cu_seqlens, k_cu_seqlens,
        max_seqlen, max_seqlen,
        causal=False,
    )
    return attn_out.flatten(start_dim=-2)
​
# Learnable2DInterpPosEmb.forward 中的推理缓存
if not self.training:
    cache_key = (shape, self.weight.dtype, self.weight.device)
    cached = self._interpolated_pos_emb_cache.get(cache_key)
    if cached is not None:
        pos_embs.append(cached)
        continue
    # ... 计算 interpolate ...
    self._interpolated_pos_emb_cache[cache_key] = interpolated
    # LRU 淘汰超出限制的条目
    if len(self._interpolated_pos_emb_cache) > _MAX_INFERENCE_POS_EMB_CACHE_ENTRIES:
        key_to_evict = next(iter(self._interpolated_pos_emb_cache))
        del self._interpolated_pos_emb_cache[key_to_evict]

评论区精华

审查者 gemini-code-assist[bot] 指出 _interpolated_pos_emb_cache 字典在推理中无界增长可能导致 GPU OOM，建议限制缓存大小并淘汰旧条目。开发者采纳此建议，在代码中添加了 _MAX_INFERENCE_POS_EMB_CACHE_ENTRIES = 256 以及基于 LRU 的淘汰策略。

• 推断缓存无界增长风险 (correctness): 建议限制缓存大小并实现 LRU 淘汰。开发者采纳并添加了 _MAX_INFERENCE_POS_EMB_CACHE_ENTRIES=256 和淘汰逻辑。

风险与影响

• 风险：新的 TP 线性层可能与其他量化设置（如 ModelSlim）冲突；DP 路径依赖 run_dp_sharded_mrope_vision_model 需要正确 rope_type；CUDA graph 捕获上下文改变可能影响其他 ViT 模型；位置缓存淘汰可能导致高频变化网格下性能下降；图像处理器 kwarg 缓存假定处理器类不变，运行时变更可能产生错误。
• 影响：对 Kimi-VL 用户：修复了编码器 DP 的基本功能，使模型可以正常完成图像请求；非 DP 下 MoonViT 层获得 TP 加速。性能优化包括减少 GPU 同步、避免异常路径、缓存插值结果，对批图像请求有明显提升。系统层面新增的测试用例在 CPU 上快速验证并行路径，降低回归风险。
• 风险标记：缓存淘汰可能影响性能, 并行配置错误风险, GPU 同步移除需验证, TP 线性层量化兼容性

关联脉络

• 暂无明显关联 PR