#54774 [Model] Add Cohere Compass model

原始 PR 作者 am-cohere 合并时间 2026-09-07 22:36 文件变更 4 提交数 6 评论 29 代码增减 +2366 / -1

执行摘要

为 vLLM 添加 Cohere Compass (North Micro Vision Instruct) 多模态视觉语言模型支持。

随着 Cohere 发布 North-Micro-Vision-Instruct 模型并在 HuggingFace transformers 中添加支持（transformers#47878），vLLM 需要跟进支持该模型的推理。PR body 中明确说明：'Following the release of https://huggingface.co/CohereLabs/North-Micro-Vision-Instruct and addition to transformers in https://github.com/huggingface/transformers/pull/47878 we add a new model "Cohere Compass" aka North Micro Vision Instruct.' 作者提到此前已有 PR#52772 的尝试，本次 PR 为新的迭代版本。

本 PR 是一个标准的新增模型支持 PR，实现了完整的推理链路。值得关注的设计决策包括：

1) Triton 内核加速位置嵌入插值的选择；
2) 对现有 CohereForCausalLM 基类的复用；
3) M-ROPE 检测逻辑的泛化修改。建议精读 cohere_compass.py 中的视觉编码器实现和 Triton 内核，评估是否可进一步抽象为通用组件。CodeRabbit 提出的 CUDA 图参数问题值得后续跟进。

讨论亮点

代码重复问题：DarkLight1337 指出 Triton 位置嵌入插值内核可能复制自 Qwen3-VL，建议直接导入而非重复实现。作者后续删除了相关方法（'delete unused method' 提交），但最终版本中仍保留了 Triton 内核的独立实现。未使用方法清理：DarkLight1337 发现并指出存在未使用的方法，作者随后删除了该方法。CI 测试兼容性：CI 失败是因为 CohereCompass 模型在 vLLM 依赖的 transformers 版本中尚不存在。DarkLight1337 建议在测试注册中设置 min_transformers_version 来跳过测试，作者按此建议修改后测试通过。CodeRabbit 建议：CodeRabbit 发现 prepare_encoder_metadata 在 CUDA 图 replay 分支中可能传递了错误的 batch-size 参数（应使用 max_frames_per_batch 而非 max_batch_size），但此建议未被作者在本次 PR 中修复。

实现拆解

新增模型实现文件 (vllm/model_executor/models/cohere_compass.py)：新增 2354 行代码，实现完整的 CohereCompassForConditionalGeneration 模型，包括视觉编码器 (CohereVisionEncoder)、语言模型头部 (CohereCompassForCausalLM)、多模态处理器、图像输入处理和权重映射。关键类包括 CohereVisionAttention（视觉注意力）、CohereVisionBlock（视觉块）、CohereCompassVisual（视觉模型主类）和 CohereCompassForConditionalGeneration（主模型类）。
Triton 双线性位置嵌入插值内核：在同文件中实现 _bilinear_pos_embed_kernel Triton 内核及 triton_pos_embed_interpolate 启动函数，将视觉编码器的位置嵌入插值从多次 eager CUDA kernel 合并为单次 Triton 启动，同时提供 pos_embed_interpolate_native 作为后备方案。
模型注册：在 vllm/model_executor/models/registry.py 的 _MULTIMODAL_MODELS 字典中添加 CohereCompassForConditionalGeneration 条目，映射到 cohere_compass 模块。
M-ROPE 检测逻辑扩展：修改 vllm/transformers_utils/config.py 中的 _uses_mrope 函数，使其支持 rope_parameters 为字典值（而非仅字符串键）的格式，以兼容 CohereCompass 的配置。
测试覆盖：在 tests/models/registry.py 的 _MULTIMODAL_EXAMPLE_MODELS 中添加 CohereCompass 的测试配置，设置 min_transformers_version="5.16.0" 以确保测试环境兼容。

文件	模块	状态	重要度
`vllm/model_executor/models/cohere_compass.py`	模型执行	added	9.17
`vllm/model_executor/models/registry.py`	模型注册	modified	5.68
`vllm/transformers_utils/config.py`	配置工具	modified	4.96
`tests/models/registry.py`	模型测试	modified	4.2

关键符号

_bilinear_pos_embed_kernel triton_pos_embed_interpolate pos_embed_interpolate_native CohereVisionEncoder.__init__ CohereVisionEncoder.forward CohereCompassVisual.__init__ CohereCompassVisual.forward CohereCompassForConditionalGeneration.__init__ CohereCompassForConditionalGeneration.forward _uses_mrope

关键源码片段

vllm/model_executor/models/cohere_compass.py core-logic

核心模型实现文件，包含完整的 CohereCompassForConditionalGeneration 模型定义、视觉编码器、多模态处理器和 Triton 加速内核。

# Triton 双线性位置嵌入插值内核：将视觉编码器的位置嵌入插值从多次 eager CUDA kernel 合并为单次 Triton 启动
# 该内核同时处理空间合并重排序，输出可直接与 patch embeddings 相加
if HAS_TRITON:

    @triton.jit
    def _bilinear_pos_embed_kernel(
        embed_ptr,
        output_ptr,
        H,
        W,
        h_scale,
        w_scale,
        NUM_GRID: tl.constexpr,
        M_SIZE: tl.constexpr,
        HIDDEN_DIM: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        """Fused bilinear pos-embed interpolation with spatial-merge reorder."""
        pid = tl.program_id(0)
        total_spatial = H * W
        spatial_idx = pid % total_spatial

        # 计算空间合并后的块索引和局部索引
        num_blocks_w = W // M_SIZE
        block_idx = spatial_idx // (M_SIZE * M_SIZE)
        local_idx = spatial_idx % (M_SIZE * M_SIZE)
        br = block_idx // num_blocks_w
        bc = block_idx % num_blocks_w
        lr = local_idx // M_SIZE
        lc = local_idx % M_SIZE
        row = br * M_SIZE + lr
        col = bc * M_SIZE + lc

        # 计算双线性插值坐标
        h_frac = row.to(tl.float32) * h_scale
        w_frac = col.to(tl.float32) * w_scale

        # 双线性插值权重计算
        hf = tl.math.floor(h_frac).to(tl.int32)
        wf = tl.math.floor(w_frac).to(tl.int32)
        hc = tl.minimum(hf + 1, NUM_GRID - 1)
        wc = tl.minimum(wf + 1, NUM_GRID - 1)

        dh = h_frac - hf.to(tl.float32)
        dw = w_frac - wf.to(tl.float32)
        w11 = dh * dw
        w10 = dh - w11
        w01 = dw - w11
        w00 = 1.0 - dh - w01

        # 按隐藏维度块进行向量化计算
        out_dtype = output_ptr.dtype.element_ty
        w00_c = w00.to(out_dtype)
        w01_c = w01.to(out_dtype)
        w10_c = w10.to(out_dtype)
        w11_c = w11.to(out_dtype)

        for d in tl.range(0, HIDDEN_DIM, BLOCK_D):
            cols = d + tl.arange(0, BLOCK_D)
            mask = cols < HIDDEN_DIM

            e00 = tl.load(embed_ptr + off00 + cols, mask=mask)
            e01 = tl.load(embed_ptr + off01 + cols, mask=mask)
            e10 = tl.load(embed_ptr + off10 + cols, mask=mask)
            e11 = tl.load(embed_ptr + off11 + cols, mask=mask)

            val = w00_c * e00 + w01_c * e01 + w10_c * e10 + w11_c * e11

            tl.store(output_ptr + out_off + cols, val, mask=mask)

vllm/transformers_utils/config.py core-logic

配置检测逻辑修改，扩展 M-ROPE 检测以支持 CohereCompass 的 rope_parameters 字典格式。

# 扩展 M-ROPE 检测逻辑，支持 rope_parameters 为字典值的格式
# CohereCompass 模型的 rope_parameters 是一个字典而非字符串键
# 修改前仅检查 "mrope_section" in rope_parameters
# 修改后同时检查字典值中是否包含 "mrope_section"
def _uses_mrope(config: PretrainedConfig) -> bool:
    rope_parameters = getattr(config, "rope_parameters", None)
    if rope_parameters is None:
        return False

    return "mrope_section" in rope_parameters or any(
        isinstance(params, dict) and "mrope_section" in params
        for params in rope_parameters.values()
    )

评论区精华

Triton 内核是否应从 Qwen3-VL 导入而非重复实现 设计

DarkLight1337 指出 Triton 双线性位置嵌入插值内核可能复制自 Qwen3-VL，建议直接导入以避免代码重复。

结论：作者后续通过删除相关方法进行了调整，但最终版本中仍保留了独立的 Triton 内核实现。讨论标记为 'Resolved offline'。 · 已解决

CI 测试因 transformers 版本不兼容而失败 测试

CI 失败是因为 CohereCompass 模型在 vLLM 依赖的 transformers 版本中尚不存在。DarkLight1337 建议设置 min_transformers_version 来跳过测试。

结论：作者按建议在测试注册中设置 min_transformers_version="5.16.0"，测试通过。 · 已解决

未使用方法的清理 style

DarkLight1337 指出存在未使用的方法，建议删除。

结论：作者在后续提交中删除了该方法。 · 已解决

CUDA 图 replay 路径的 batch-size 参数传递 正确性

CodeRabbit 通过静态分析发现 prepare_encoder_metadata 在 replay 分支中可能传递了错误的 batch-size 参数（应使用 max_frames_per_batch 而非 max_batch_size）。

结论：该建议未被作者在本次 PR 中修复，属于未解决的潜在问题。 · 待处理

风险与影响

兼容性风险：模型依赖 transformers >= 5.16.0 中的 CohereCompassConfig 等类，若用户环境 transformers 版本较低将无法使用。代码重复风险：Triton 双线性位置嵌入插值内核与 Qwen3-VL 的实现可能存在重复，维护时需同步更新。CUDA 图元数据参数风险：CodeRabbit 指出的 prepare_encoder_metadata 参数问题（在 replay 路径传递 max_batch_size 而非 max_frames_per_batch）可能在特定帧数配置下导致 CUDA 图重放缓冲区与捕获时不一致。新模型首次引入的未知风险：作为全新模型实现，缺乏长期使用验证，可能存在边缘情况的正确性问题。

用户影响：支持 Cohere Compass 模型的 vLLM 推理，使用户可以通过 vLLM API 使用该多模态视觉语言模型。系统影响：扩展了 vLLM 的多模态模型支持范围，_uses_mrope 的修改影响所有使用 M-ROPE 的模型的配置检测逻辑。团队影响：为 Cohere 系列模型在 vLLM 中的后续支持奠定基础（已有 Cohere2 Vision），但需关注代码重复维护成本。

新模型首次引入代码重复风险 CUDA 图参数潜在错误依赖版本要求

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。

完整报告

执行摘要

一句话：为 vLLM 添加 Cohere Compass (North Micro Vision Instruct) 多模态视觉语言模型支持。
推荐动作：本 PR 是一个标准的新增模型支持 PR，实现了完整的推理链路。值得关注的设计决策包括：
1) Triton 内核加速位置嵌入插值的选择；
2) 对现有 CohereForCausalLM 基类的复用；
3) M-ROPE 检测逻辑的泛化修改。建议精读 cohere_compass.py 中的视觉编码器实现和 Triton 内核，评估是否可进一步抽象为通用组件。CodeRabbit 提出的 CUDA 图参数问题值得后续跟进。

功能与动机

实现拆解

新增模型实现文件 (vllm/model_executor/models/cohere_compass.py)：新增 2354 行代码，实现完整的 CohereCompassForConditionalGeneration 模型，包括视觉编码器 (CohereVisionEncoder)、语言模型头部 (CohereCompassForCausalLM)、多模态处理器、图像输入处理和权重映射。关键类包括 CohereVisionAttention（视觉注意力）、CohereVisionBlock（视觉块）、CohereCompassVisual（视觉模型主类）和 CohereCompassForConditionalGeneration（主模型类）。
Triton 双线性位置嵌入插值内核：在同文件中实现 _bilinear_pos_embed_kernel Triton 内核及 triton_pos_embed_interpolate 启动函数，将视觉编码器的位置嵌入插值从多次 eager CUDA kernel 合并为单次 Triton 启动，同时提供 pos_embed_interpolate_native 作为后备方案。
模型注册：在 vllm/model_executor/models/registry.py 的 _MULTIMODAL_MODELS 字典中添加 CohereCompassForConditionalGeneration 条目，映射到 cohere_compass 模块。
M-ROPE 检测逻辑扩展：修改 vllm/transformers_utils/config.py 中的 _uses_mrope 函数，使其支持 rope_parameters 为字典值（而非仅字符串键）的格式，以兼容 CohereCompass 的配置。
测试覆盖：在 tests/models/registry.py 的 _MULTIMODAL_EXAMPLE_MODELS 中添加 CohereCompass 的测试配置，设置 min_transformers_version="5.16.0" 以确保测试环境兼容。

关键文件：

vllm/model_executor/models/cohere_compass.py（模块模型执行；类别 source；类型 core-logic；符号 _bilinear_pos_embed_kernel, triton_pos_embed_interpolate, pos_embed_interpolate_native, Cohere_VLImagePixelInputs）: 核心模型实现文件，包含完整的 CohereCompassForConditionalGeneration 模型定义、视觉编码器、多模态处理器和 Triton 加速内核。
vllm/model_executor/models/registry.py（模块模型注册；类别 source；类型 data-contract）: 模型注册表，将 CohereCompassForConditionalGeneration 映射到正确的模块和类，使 vLLM 能够自动发现和加载该模型。
vllm/transformers_utils/config.py（模块配置工具；类别 source；类型 core-logic；符号 _uses_mrope）: 配置检测逻辑修改，扩展 M-ROPE 检测以支持 CohereCompass 的 rope_parameters 字典格式。
tests/models/registry.py（模块模型测试；类别 test；类型 test-coverage）: 测试注册表，添加 CohereCompass 的测试配置，确保模型在 CI 中得到验证。

关键符号：_bilinear_pos_embed_kernel, triton_pos_embed_interpolate, pos_embed_interpolate_native, CohereVisionEncoder.init, CohereVisionEncoder.forward, CohereCompassVisual.init, CohereCompassVisual.forward, CohereCompassForConditionalGeneration.init, CohereCompassForConditionalGeneration.forward, _uses_mrope

关键源码片段

`vllm/model_executor/models/cohere_compass.py`

核心模型实现文件，包含完整的 CohereCompassForConditionalGeneration 模型定义、视觉编码器、多模态处理器和 Triton 加速内核。

# Triton 双线性位置嵌入插值内核：将视觉编码器的位置嵌入插值从多次 eager CUDA kernel 合并为单次 Triton 启动
# 该内核同时处理空间合并重排序，输出可直接与 patch embeddings 相加
if HAS_TRITON:

    @triton.jit
    def _bilinear_pos_embed_kernel(
        embed_ptr,
        output_ptr,
        H,
        W,
        h_scale,
        w_scale,
        NUM_GRID: tl.constexpr,
        M_SIZE: tl.constexpr,
        HIDDEN_DIM: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        """Fused bilinear pos-embed interpolation with spatial-merge reorder."""
        pid = tl.program_id(0)
        total_spatial = H * W
        spatial_idx = pid % total_spatial

        # 计算空间合并后的块索引和局部索引
        num_blocks_w = W // M_SIZE
        block_idx = spatial_idx // (M_SIZE * M_SIZE)
        local_idx = spatial_idx % (M_SIZE * M_SIZE)
        br = block_idx // num_blocks_w
        bc = block_idx % num_blocks_w
        lr = local_idx // M_SIZE
        lc = local_idx % M_SIZE
        row = br * M_SIZE + lr
        col = bc * M_SIZE + lc

        # 计算双线性插值坐标
        h_frac = row.to(tl.float32) * h_scale
        w_frac = col.to(tl.float32) * w_scale

        # 双线性插值权重计算
        hf = tl.math.floor(h_frac).to(tl.int32)
        wf = tl.math.floor(w_frac).to(tl.int32)
        hc = tl.minimum(hf + 1, NUM_GRID - 1)
        wc = tl.minimum(wf + 1, NUM_GRID - 1)

        dh = h_frac - hf.to(tl.float32)
        dw = w_frac - wf.to(tl.float32)
        w11 = dh * dw
        w10 = dh - w11
        w01 = dw - w11
        w00 = 1.0 - dh - w01

        # 按隐藏维度块进行向量化计算
        out_dtype = output_ptr.dtype.element_ty
        w00_c = w00.to(out_dtype)
        w01_c = w01.to(out_dtype)
        w10_c = w10.to(out_dtype)
        w11_c = w11.to(out_dtype)

        for d in tl.range(0, HIDDEN_DIM, BLOCK_D):
            cols = d + tl.arange(0, BLOCK_D)
            mask = cols < HIDDEN_DIM

            e00 = tl.load(embed_ptr + off00 + cols, mask=mask)
            e01 = tl.load(embed_ptr + off01 + cols, mask=mask)
            e10 = tl.load(embed_ptr + off10 + cols, mask=mask)
            e11 = tl.load(embed_ptr + off11 + cols, mask=mask)

            val = w00_c * e00 + w01_c * e01 + w10_c * e10 + w11_c * e11

            tl.store(output_ptr + out_off + cols, val, mask=mask)

`vllm/transformers_utils/config.py`

配置检测逻辑修改，扩展 M-ROPE 检测以支持 CohereCompass 的 rope_parameters 字典格式。

# 扩展 M-ROPE 检测逻辑，支持 rope_parameters 为字典值的格式
# CohereCompass 模型的 rope_parameters 是一个字典而非字符串键
# 修改前仅检查 "mrope_section" in rope_parameters
# 修改后同时检查字典值中是否包含 "mrope_section"
def _uses_mrope(config: PretrainedConfig) -> bool:
    rope_parameters = getattr(config, "rope_parameters", None)
    if rope_parameters is None:
        return False

    return "mrope_section" in rope_parameters or any(
        isinstance(params, dict) and "mrope_section" in params
        for params in rope_parameters.values()
    )

评论区精华

Triton 内核是否应从 Qwen3-VL 导入而非重复实现 (design): 作者后续通过删除相关方法进行了调整，但最终版本中仍保留了独立的 Triton 内核实现。讨论标记为'Resolved offline'。
CI 测试因 transformers 版本不兼容而失败 (testing): 作者按建议在测试注册中设置 min_transformers_version="5.16.0"，测试通过。
未使用方法的清理 (style): 作者在后续提交中删除了该方法。
CUDA 图 replay 路径的 batch-size 参数传递 (correctness): 该建议未被作者在本次 PR 中修复，属于未解决的潜在问题。

风险与影响

风险：兼容性风险：模型依赖 transformers >= 5.16.0 中的 CohereCompassConfig 等类，若用户环境 transformers 版本较低将无法使用。代码重复风险：Triton 双线性位置嵌入插值内核与 Qwen3-VL 的实现可能存在重复，维护时需同步更新。CUDA 图元数据参数风险：CodeRabbit 指出的 prepare_encoder_metadata 参数问题（在 replay 路径传递 max_batch_size 而非 max_frames_per_batch）可能在特定帧数配置下导致 CUDA 图重放缓冲区与捕获时不一致。新模型首次引入的未知风险：作为全新模型实现，缺乏长期使用验证，可能存在边缘情况的正确性问题。
影响：用户影响：支持 Cohere Compass 模型的 vLLM 推理，使用户可以通过 vLLM API 使用该多模态视觉语言模型。系统影响：扩展了 vLLM 的多模态模型支持范围，_uses_mrope 的修改影响所有使用 M-ROPE 的模型的配置检测逻辑。团队影响：为 Cohere 系列模型在 vLLM 中的后续支持奠定基础（已有 Cohere2 Vision），但需关注代码重复维护成本。
风险标记：新模型首次引入, 代码重复风险, CUDA 图参数潜在错误, 依赖版本要求

关联脉络

PR #52772 [Model] Add Cohere Compass model: 本 PR 是 PR#52772 的替代版本，作者在评论中明确说明 'I have opened this PR in lieu of https://github.com/vllm-project/vllm/pull/52772'
PR #55735 [Frontend] Migrate input-validation errors to VLLMValidationError: 同仓库近期 PR，涉及前端错误处理迁移，与本 PR 无直接技术关联，但同属于模型支持和前端改进的演进脉络

#54774 [Model] Add Cohere Compass model

执行摘要

为 vLLM 添加 Cohere Compass (North Micro Vision Instruct) 多模态视觉语言模型支持。

实现拆解

评论区精华

风险与影响

关联 Issue

未识别关联 Issue

完整报告

参与讨论