执行摘要

• 一句话：将 DSv4 attention 路径的序列/DCP JIT 内核迁移到共享 warmup 契约，消除运行时编译。
• 推荐动作：本 PR 是 DSv4 去 JIT 化系列中的关键里程碑，实现了高质量的内核迁移和契约建立。建议精读以下内容：
  1) VllmCuTeDSLJitKernel 和 VllmTritonJitKernel 基类的设计，理解 dispatch/get_warmup_keys/warmup_inputs 的协作模式；
  2) jit_warmup_triton_helper.py 中 kernel_launcher 装饰器的运行时/编译时适配逻辑；
  3) jit_warmup.py 中 _same_registration 去重优化。整体实现清晰、测试充分，值得作为后续迁移的参考模板。

功能与动机

消除 DSv4 推理路径的运行时 JIT 编译，是 Issue #49349（零 JIT 编译）追踪的关键步骤。PR 依赖于 #50175 等前置 PR，旨在将 --jit-monitor-mode error 下不再产生未追踪内核编译，提升首次推理延迟和可预测性。

实现拆解

1. 引入内核类契约：在 vllm/model_executor/warmup/ 下新增 VllmCuTeDSLJitKernel 基类与 cutedsl_kernel_launcher 装饰器（jit_warmup_cutedsl_helper.py），扩展 VllmTritonJitKernel 的 kernel_launcher 以支持 CPU 函数包装器和运行时启动器适配。
2. 迁移核心内核：将 dcp_indexer_cutedsl.py 中的 _pack_dcp_topk_candidates_triton_kernel 和 StableTopKFromGatheredCandidatesKernel 重构为 PackDCPTopkCandidatesKernel/StableTopKFromGatheredCandidatesKernel 类，新增 CompileKey、dispatch、get_warmup_keys、warmup_inputs 及 __call__ 方法。将 vllm/v1/attention/ops/common.py 中的 _pack_seq_kernel 和 _unpack_seq_triton_kernel 迁移为 PackSeqTritonKernel/UnpackSeqTritonKernel 类。将 vllm/v1/attention/ops/dcp.py 中的 _correct_attn_cp_out_kernel 迁移为 CorrectAttnCPOutKernel 类。
3. 迁移辅助内核：将 vllm/v1/worker/block_table.py 中的 ComputeSlotMappingKernel 从 VllmJitKernel 迁移到 VllmTritonJitKernel，新增 warmup_inputs 方法。
4. 优化注册去重：在 jit_warmup.py 中新增 _same_value 和 _same_registration 辅助函数，改进 JitWarmupRegistry._add 以支持基于值相等性的重复注册去重，避免深层模型（如 61 层 DSA）重复追踪。
5. 注册与测试：在 sparse_attn_indexer.py 中添加条件注册逻辑，为 DSv4 attention 内核调用 register_warmup()。新增 test_jit_warmup_triton_launcher.py 和 test_jit_warmup_cutedsl_launcher.py 测试文件，验证启动器支持编译/运行时适配、CPU 函数包装器及编译执行器复用。更新 .buildkite/test_areas/model_executor.yaml 添加新测试路径。

关键文件：

• vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py（模块 DSA 内核；类别 source；类型 core-logic；符号 PackDCPTopkCandidatesKernel, StableTopKFromGatheredCandidatesKernel, CompileKey, dispatch）: 核心内核迁移：将 DCP 候选打包 Triton 内核和稳定 Top-K CuTeDSL 内核重构为遵循 warmup 契约的类，是 PR 变更量最大、复杂度最高的文件。
• vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py（模块 Warmup 辅助；类别 source；类型 data-contract；符号 VllmCuTeDSLJitKernel, compile_cutedsl, cutedsl_kernel_launcher, cutedsl_fake_stream）: 新增基础设施：为 CuTeDSL 内核提供与 Triton 内核对称的 warmup 契约基类和启动器装饰器。
• vllm/v1/attention/ops/common.py（模块 注意力操作；类别 infra；类型 core-logic；符号 PackSeqTritonKernel, UnpackSeqTritonKernel, CompileKey, dispatch）: 序列打包/解包内核迁移：将核心序列操作 Triton 内核迁移到 warmup 契约，修复了 warmup 期间的类型转换崩溃。
• vllm/model_executor/warmup/jit_warmup.py（模块 Warmup 框架；类别 source；类型 data-contract；符号 _same_value, _same_registration, JitWarmupRegistry._add）: 基础架构优化：改进注册去重逻辑，通过值相等性比较减少深层模型的重复 warmup 追踪开销。
• tests/model_executor/test_jit_warmup_triton_launcher.py（模块 Triton 启动器测试；类别 test；类型 test-coverage；符号 _FakeTritonKernel, _TestTritonKernel, test_triton_launcher_supports_compile_and_runtime_adapters, test_triton_launcher_supports_cpu_function_wrappers）: 测试覆盖：新增 Triton 启动器单元测试，验证编译/运行时适配、CPU 函数包装器和命名输入传递。

关键符号：PackDCPTopkCandidatesKernel.dispatch, PackDCPTopkCandidatesKernel.get_warmup_keys, PackDCPTopkCandidatesKernel.warmup_inputs, StableTopKFromGatheredCandidatesKernel.compile, VllmCuTeDSLJitKernel.compile, cutedsl_kernel_launcher, PackSeqTritonKernel.dispatch, PackSeqTritonKernel.warmup_inputs, CorrectAttnCPOutKernel.dispatch, ComputeSlotMappingKernel.warmup_inputs, _same_value, _same_registration, JitWarmupRegistry._add

关键源码片段

vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py

核心内核迁移：将 DCP 候选打包 Triton 内核和稳定 Top-K CuTeDSL 内核重构为遵循 warmup 契约的类，是 PR 变更量最大、复杂度最高的文件。

# PackDCPTopkCandidatesKernel 类遵循 VllmTritonJitKernel 契约，
# 将 Triton 内核的所有 constexpr 特化参数封装为不可变 CompileKey。
# dispatch 方法负责将运行时参数映射为编译键。
class PackDCPTopkCandidatesKernel(
    VllmTritonJitKernel["PackDCPTopkCandidatesKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        dcp_rank: int
        dcp_world_size: int
        cp_interleave: int
        has_row_starts: bool
        topk: int
        block_size: int
​
    @staticmethod
    @triton.jit(
        do_not_specialize=[
            "logits_stride0", "logits_stride1",
            "topk_stride0", "topk_stride1",
            "packed_stride0", "packed_stride1", "packed_stride2",
            "num_cols",
        ]
    )
    def kernel(logits, topk_indices, packed, row_starts,
               logits_stride0, logits_stride1, ...
               DCP_RANK: tl.constexpr, DCP_WORLD_SIZE: tl.constexpr, ...):
        # 内核逻辑保持不变，仅参数布局重构为键控特化
        row = tl.program_id(0)
        tile = tl.program_id(1)
        cols = tile * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        mask = cols < TOPK
        ...
​
    def dispatch(self, *, has_row_starts: bool, **compile_key_fields: int) -> CompileKey:
        # 将运行时 has_row_starts 映射为 CompileKey 的一部分
        return self.CompileKey(**compile_key_fields, has_row_starts=has_row_starts)
​
    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
        # 根据 vllm_config 动态确定需要预热的编译键组合
        dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
        if dcp_world_size <= 1:
            return []
        ...
        return self._trace_dispatch(self.dispatch)(
            dcp_rank=dcp_rank, dcp_world_size=dcp_world_size,
            cp_interleave=cp_interleave, has_row_starts=(False, True),
            topk=topk, block_size=512,
        )
​
    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        # 返回假张量，用于编译时特化而不触发真实计算
        fp32_ptr = TritonWarmupTensor(torch.float32)
        int32_ptr = TritonWarmupTensor(torch.int32)
        return dict(
            logits=fp32_ptr,
            topk_indices=TritonWarmupTensor(torch.int32, shape=(1, compile_key.topk)),
            packed=fp32_ptr,
            row_starts_arg=int32_ptr,
            logits_stride0=1, logits_stride1=1, ...
            topk=compile_key.topk, block_size=compile_key.block_size,
        )
​
    @kernel_launcher
    def __call__(self, logits, topk_indices, packed, row_starts_arg, *,
                logits_stride0, ..., topk, block_size) -> LaunchSpec:
        grid = (topk_indices.shape[0], triton.cdiv(topk, block_size))
        return grid, dict(
            row_starts=row_starts_arg,
            DCP_RANK=dcp_rank, ...
            TOPK=topk, BLOCK_SIZE=block_size, num_warps=8,
        )

vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py

新增基础设施：为 CuTeDSL 内核提供与 Triton 内核对称的 warmup 契约基类和启动器装饰器。

# VllmCuTeDSLJitKernel 是 CuTeDSL 内核的 warmup 契约基类，
# 提供 compile 方法将 warmup 输入编译为执行器并缓存。
DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi"
​
CompileKeyT = TypeVar("CompileKeyT")
CuTeDSLLaunchSpec: TypeAlias = tuple[
    CompileKeyT, tuple[Any, ...], Mapping[str, Any] | None,
]
​
def compile_cutedsl(entry, *args, options=DEFAULT_CUTEDSL_COMPILE_OPTIONS, ...):
    # 使用假流编译 CuTeDSL 内核入口，避免需要真实 CUDA 流
    return cute.compile(entry, *args, cutedsl_fake_stream(...), options=options)
​
class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
    """CuTeDSL owner whose compiled executor is shared by warmup and runtime."""
    kernel: ClassVar[Any]
​
    @abstractmethod
    def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]:
        """Return fake arguments that compile one executor specialization."""
        raise NotImplementedError
​
    def compile(self, compile_key: CompileKeyT) -> None:
        if compile_key in self._compiled_cache:
            return
        # 编译一次，缓存执行器供后续 warmup 和 runtime 共享
        self._compiled_cache[compile_key] = compile_cutedsl(
            self.kernel(compile_key), *self.warmup_inputs(compile_key),
        )
​
def cutedsl_kernel_launcher(call_fn):
    """Invoke a cached CuTeDSL executor from a declarative __call__."""
    @wraps(call_fn)
    def wrapper(self, *args, **kwargs):
        compile_key, launch_args, runtime_context = call_fn(self, *args, **kwargs)
        executor = self._get_or_compile(compile_key, runtime_context=runtime_context)
        return executor(*launch_args)
    return wrapper

vllm/v1/attention/ops/common.py

序列打包/解包内核迁移：将核心序列操作 Triton 内核迁移到 warmup 契约，修复了 warmup 期间的类型转换崩溃。

# PackSeqTritonKernel 将长度转换移至包装器边界，
# 避免 warmup 路径访问 TritonWarmupTensor 无的方法。
class PackSeqTritonKernel(VllmTritonJitKernel["PackSeqTritonKernel.CompileKey"]):
    @dataclass(frozen=True)
    class CompileKey:
        dtype: torch.dtype
        pad_value: float | int
        pad_is_uint8: bool
        block_t: int
        block_d: int
​
    @staticmethod
    @triton.jit(do_not_specialize=["N", "D", "Lmax"])
    def kernel(x_ptr, out_ptr, lengths_ptr, N, D, Lmax,
               PAD_VALUE: tl.constexpr, PAD_IS_UINT8: tl.constexpr, ...):
        # 内核逻辑不变，N/D/Lmax 不参与特化
        pid_b = tl.program_id(0)
        pid_t = tl.program_id(1)
        pid_d = tl.program_id(2)
        off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
        ...
​
    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        # 返回假张量和标量，N=1, D=1, Lmax=1 满足编译要求
        return dict(
            x_reshaped=TritonWarmupTensor(compile_key.dtype),
            out=TritonWarmupTensor(compile_key.dtype),
            lengths=TritonWarmupTensor(torch.int32),
            N=1, D=1, Lmax=1,
            pad_value=compile_key.pad_value,
            pad_is_uint8=compile_key.pad_is_uint8,
            block_t=compile_key.block_t, block_d=compile_key.block_d,
        )
​
    @kernel_launcher
    def __call__(self, x_reshaped, out, lengths, *, N, D, Lmax,
                pad_value, pad_is_uint8, block_t, block_d) -> LaunchSpec:
        grid = (lengths.shape[0], triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d))
        return grid, dict(
            x_ptr=x_reshaped, lengths_ptr=lengths,
            PAD_VALUE=pad_value, PAD_IS_UINT8=pad_is_uint8,
            BLOCK_T=block_t, BLOCK_D=block_d, num_warps=4, num_stages=2,
        )

评论区精华

核心讨论点

1. 启动崩溃修复（正确性）：维护者 LucasWilkinson 指出迁移后的 PackSeqTritonKernel/UnpackSeqTritonKernel 在 compile() 期间因 TritonWarmupTensor 无 numel()/int() 方法导致启动崩溃（AttributeError）。作者回应已将类型转换移至包装器边界，避免在 warmup 路径访问这些属性。
2. 死参数清理（设计）：LucasWilkinson 建议删除 CorrectAttnCPOutKernel.dispatch 中未被 get_warmup_keys 调用的 6 个 runtime_*_stride 参数及其 is not None 分支。作者已移除这些死参数，直接内联连续布局默认值。
3. 硬编码参数计数（健壮性）：LucasWilkinson 指出 _runtime_launcher_arg_count=11 是硬编码的，若内核参数增减会导致缓存重放不对齐。作者已改为 len(self.kernel.arg_names) - len(self.kernel.constexprs) 动态计算。
4. 注册去重扩展（性能）：LucasWilkinson 指出每层 DSA 模型重复注册相同 (dtype, pad_value) 会触发 61 次 get_warmup_keys 追踪。作者在同一 PR 中通过 _same_registration 实现了基于值相等性的去重。
5. 测试覆盖建议：coderabbitai[bot] 建议扩展 test_jit_warmup.py 以覆盖值比较和异常捕获场景，但未被纳入本次合并。

• warmup 崩溃修复：TritonWarmupTensor 缺少 numel()/int() 方法 (correctness): 问题已修复：__call__ 使用 lengths.shape[0] 而非 .numel()，.int() 转换移至外层包装器，避免在 warmup 路径访问 TritonWarmupTensor 的非法方法。
• CorrectAttnCPOutKernel.dispatch 中未使用的 runtime_*_stride 参数 (design): 已修复：所有 6 个死参数及其分支已从 dispatch 中移除，替换为内联的连续布局默认值计算。
• 硬编码 _runtime_launcher_arg_count=11 (correctness): 已修复：改为动态计算 len(self.kernel.arg_names) - len(self.kernel.constexprs)。
• 深层模型重复注册去重优化 (performance): 已在本次 PR 中通过 _same_registration 实现基于值相等性的去重。
• AI 辅助 review 建议的测试扩展 (testing): 未在本次 PR 中采纳，可能留待后续。

风险与影响

• 风险：
  1. 核心路径变更风险：迁移涉及 DSv4 attention 的核心内核（打包、解包、DCP 校正），若 CompileKey 或 warmup_inputs 构造错误，可能导致运行时行为异常。但 PR 配套了全面的单元测试（test_jit_warmup_triton_launcher.py、test_jit_warmup_cutedsl_launcher.py）和 CI 覆盖。
  2. CuTeDSL 编译缓存：VllmCuTeDSLJitKernel.compile 使用 _compiled_cache 存储编译后的执行器，确保运行时复用。若 warmup_inputs 返回的假张量与真实输入类型不匹配，可能引发编译错误。测试通过 monkeypatch 验证了缓存行为。
  3. 注册去重逻辑：新引入的 _same_registration 依赖 _same_value 的 __eq__ 比较，若对象自定义了非幂等 __eq__，可能导致错误去重。但实际注册参数多为不可变标量或配置对象，风险较低。
  4. 条件注册依赖平台：sparse_attn_indexer.py 中的注册依赖 current_platform.is_cuda() and has_cutedsl()，确保仅在支持的平台上激活。
• 影响：
  1. 用户影响：DSv4 模型用户在启用 --enable-jit-warmup 时，首次推理延迟将显著降低，因为相关内核已在启动时预编译。不影响禁用该标志的用户。
  2. 系统影响：共享 warmup 契约（VllmTritonJitKernel/VllmCuTeDSLJitKernel）为后续其他模型（GLM-5.2/5.3、Gemma 4、gpt-oss、Kimi-K3）的去 JIT 化提供了标准化接口。
  3. 团队影响：建立了一套可复用的内核迁移模式（声明式 dispatch、编译键、warmup 输入），降低了未来迁移其他内核的认知成本。测试覆盖为框架扩展提供了安全网。
     - 风险标记：核心路径变更, warmup 契约接口变更, 注册去重逻辑引入

关联脉络

• PR #50175 父 PR：warmup 基础设施和 DSv4 前置迁移: 本 PR 依赖 #50175 作为前置 PR，提供了 warmup 基础设施和 DSv4 的部分内核迁移。
• PR #49627 父 (draft) PR：DSv4 完整 warmup 迁移草稿: 本 PR 是 #49627 草稿 PR 的拆分提交，负责序列和 DCP 内核的迁移部分。
• PR #53565 [3/N][warmup][DSv4] Migrate remaining DSA kernels: 紧随本 PR 的后续迁移，处理 DSv4 的其他 DSA 内核。
• PR #55348 [warmup][GLM-5.2/5.3] Migrate GLM models: 同系列去 JIT 化 PR，将 GLM-5.2/5.3 模型的内核迁移到相同契约。
• PR #55768 [warmup][Gemma 4] Migrate Gemma 4 kernels: 同系列去 JIT 化 PR，将 Gemma 4 模型的内核迁移到相同契约。