执行摘要

• 一句话：修复 KVBlockZeroer 对非均匀 KV 缓存页面大小的断言崩溃
• 推荐动作：强烈建议使用 GLM-5.2 或类似含多结构 KV 缓存组模型的用户升级到此提交。本 PR 展示了一个清晰的从假设到一般化的重构模式：将编译时常量改为运行时 per-segment 参数，避免核函数多实例化，值得在类似场景中复用。

功能与动机

关联 Issue #49696 报告 GLM-5.2 FP8 模型启动时 KVBlockZeroer 抛出断言错误：Non-uniform page sizes: 2112 vs 10496。Root cause 是 #47574 启用了 mixed-precision KV cache 的零初始化，但 KVBlockZeroer 假设所有 KV 缓存组页面大小相同。GLM-5.2 使用 MLA（kv_lora_rank: 512, v_head_dim: 256）加 DSA indexer（index_n_heads: 32, index_head_dim: 128），产生两个不同页面大小的缓存组，导致崩溃。

实现拆解

1. 修改 Triton 内核签名 (vllm/v1/worker/utils.py)
   移除编译时常量 PAGE_SIZE_EL，新增 seg_page_sizes_ptr 运行时张量指针。内核按段加载实际页面大小，若 chunk 索引超出该段页面大小则提前 return，确保不同大小的段可在同一网格内安全运行。

2. 重构 KVBlockZeroer.init 与 _meta 结构 (vllm/v1/worker/utils.py)
   将 page_size_el: int | None 替换为 seg_page_sizes: list[int] 收集所有段的页面大小。_meta 从 4 元组 (seg_addrs, page_size_el, block_size, n_segs) 扩展为 5 元组 (seg_addrs, seg_page_sizes, max_chunks, block_size, n_segs)，其中 max_chunks = max(seg_page_sizes) // block_size 用于计算网格尺寸并作为内核的 MAX_CHUNKS 参数。

3. 更新 Qwen Triton 预热模块 (vllm/model_executor/warmup/qwen_triton_warmup.py)
   修改 _ZeroKvWarmupConfig 字段及 _zero_kv_warmup_config 的解包逻辑，_warm_zero_kv_blocks_kernel 中按 seg_page_sizes.max() 分配临时缓冲区，并传递新参数。

4. 添加单元测试 (tests/v1/worker/test_kv_block_zeroer.py)
   新增 test_non_uniform_page_sizes，使用两个尺寸不同的存储段（10496 与 2112 个 int32 元素），验证零初始化后各段块正确归零，覆盖非均匀场景。同时更新已有测试适配新的 _meta 结构。

5. 运行级验证 由维护者 mgoin 在 B300 上对 DeepSeek-V4-Flash-DSpark 和 GLM-5.2-NVFP4 进行了端到端启动测试，确认修复有效。

关键文件：

• vllm/v1/worker/utils.py（模块 Worker；类别 source；类型 core-logic；符号 _zero_kv_blocks_kernel, KVBlockZeroer）: 核心改造文件：修改了 Triton kernel 以支持 per-segment 页面大小，重构了 KVBlockZeroer 的初始化逻辑与内部元数据结构。
• tests/v1/worker/test_kv_block_zeroer.py（模块 测试；类别 test；类型 test-coverage；符号 test_non_uniform_page_sizes, largest_power_of_2_divisor）: 新增非均匀页面大小的单元测试，确保零初始化内核正确性；同时更新现有测试适配 _meta 新格式。
• vllm/model_executor/warmup/qwen_triton_warmup.py（模块 预热；类别 source；类型 data-contract；符号 _ZeroKvWarmupConfig, _zero_kv_warmup_config, _warm_zero_kv_blocks_kernel）: 需要适配 _meta 新格式的预热模块，变更 ZeroKvWarmupConfig 结构及预热函数参数。

关键符号：_zero_kv_blocks_kernel, KVBlockZeroer.init, KVBlockZeroer.zero_block_ids, _zero_kv_warmup_config, _warm_zero_kv_blocks_kernel, test_non_uniform_page_sizes

关键源码片段

vllm/v1/worker/utils.py

核心改造文件：修改了 Triton kernel 以支持 per-segment 页面大小，重构了 KVBlockZeroer 的初始化逻辑与内部元数据结构。

# --- Triton kernel 改造：从编译时常量 PAGE_SIZE_EL 改为运行时 per-segment 加载 ---
@triton.jit
def _zero_kv_blocks_kernel(
    seg_addrs_ptr,
    seg_page_sizes_ptr, # 新增：每个段的页面大小（int32 elements）
    block_ids_ptr,
    n_blocks,
    N_SEGS: tl.constexpr,
    MAX_CHUNKS: tl.constexpr, # 替换 PAGE_SIZE_EL：所有段中最大 chunks 数
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    work_per_block = N_SEGS * MAX_CHUNKS
    block_index = pid // work_per_block
    if block_index >= n_blocks:
        return
    remainder = pid % work_per_block
    seg_index = remainder // MAX_CHUNKS
    chunk_index = remainder % MAX_CHUNKS
    # 运行时加载当前段的页面大小，支持不同段尺寸
    page_size_el = tl.load(seg_page_sizes_ptr + seg_index)
    if chunk_index >= page_size_el // BLOCK_SIZE:
        return # chunk 超出该段范围，提前退出
    block_id = tl.load(block_ids_ptr + block_index)
    seg_addr = tl.load(seg_addrs_ptr + seg_index)
    ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))
    offset = (
        block_id.to(tl.int64) * page_size_el.to(tl.int64)
        + chunk_index.to(tl.int64) * BLOCK_SIZE
    )
    cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)
    tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32))
​
​
# --- KVBlockZeroer 初始化：收集 per-segment 页面大小 ---
class KVBlockZeroer:
    def __init__(self, device, attn_groups_iter, kernel_block_sizes, cache_dtype, static_forward_context, runner_only_attn_layers=None):
        self.device = device
        self._meta: tuple[torch.Tensor, torch.Tensor, int, int, int] | None = None # 5 元组
        if runner_only_attn_layers is None:
            runner_only_attn_layers = set()
        seen_ptrs: set[int] = set()
        seg_addrs: list[int] = []
        seg_page_sizes: list[int] = [] # 替代原先的 page_size_el: int | None
        for group in attn_groups_iter:
            spec = group.kv_cache_spec
            if not isinstance(spec, FullAttentionSpec):
                continue
            # ... 计算 kernel_block_el, cur_page_el ...
            # 不再断言统一页面大小，改为收集
            seg_page_sizes.append(cur_page_el)
        # 后续计算 max_chunks = max(seg_page_sizes) // block_size

tests/v1/worker/test_kv_block_zeroer.py

新增非均匀页面大小的单元测试，确保零初始化内核正确性；同时更新现有测试适配 _meta 新格式。

@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_non_uniform_page_sizes():
    """两个段具有不同页面大小 (如 MLA + DSA indexer)。"""
    device = torch.device("cuda")
    num_blocks = 4
    page_size_a = 10496 # int32 elements
    page_size_b = 2112
​
    storage_a = torch.ones((num_blocks, page_size_a), dtype=torch.int32, device=device)
    storage_b = torch.ones((num_blocks, page_size_b), dtype=torch.int32, device=device)
​
    zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
    zeroer.device = device
​
    seg_page_sizes = [page_size_a, page_size_b]
    max_ps = max(seg_page_sizes)
    # 选取最保守的 BLOCK_SIZE：所有页面大小的 2 的幂因子的最小值，不超过 1024
    def largest_power_of_2_divisor(n):
        return n & -n
    blk_size = min(min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), 1024)
​
    zeroer._meta = (
        torch.tensor(
            [storage_a.data_ptr(), storage_b.data_ptr()],
            dtype=torch.uint64, device=device,
        ),
        torch.tensor(seg_page_sizes, dtype=torch.int64, device=device), # seg_page_sizes
        max_ps // blk_size, # max_chunks
        blk_size, # block_size
        2, # n_segs
    )
​
    stream = torch.cuda.Stream()
    with torch.cuda.stream(stream):
        zeroer.zero_block_ids([1, 2]) # 将第 1、2 块清零
    stream.synchronize()
​
    for storage in (storage_a, storage_b):
        assert torch.all(storage[0] == 1) # 第 0 块未被清零
        assert torch.all(storage[1] == 0) # 第 1 块清零
        assert torch.all(storage[2] == 0) # 第 2 块清零
        assert torch.all(storage[3] == 1) # 第 3 块未被清零

评论区精华

• PR 作者在 body 中坦言“vibe coded this to fix #49696, but it's unclear if GLM needs this as it is neither Mamba nor non-uniform KV dtype model so it doesn't have the NaN/Inf reinterpret_cast hazard”，表明对 GLM 实际是否需要零初始化的不确定性，但修复本身是安全且必要的。

• 维护者 mgoin 批准并附上了详细的验证命令，在 B300 上成功启动两个含非均匀 KV 缓存组的模型，确认修复正确。

• 非均匀页面大小支持合理性 (design): 采纳运行时 per-segment 页面大小方案，无需统一断言。

风险与影响

• 风险：
  1. 回归风险：Triton 内核参数从 PAGE_SIZE_EL: tl.constexpr 改为运行时 seg_page_sizes_ptr，可能引入额外的全局内存加载，但每个段仅加载一次，开销可忽略。现有统一页面大小的路径（如 Mamba 模型）仍可正常工作，因为所有段页面大小相同，MAX_CHUNKS 退化为统一值。
  2. 数据契约风险：_meta 元组结构变更，所有消费方（zero_block_ids、qwen_triton_warmup.py中的解包）均已同步更新。若存在外部扩展直接访问 _meta，可能造成兼容性问题，但该属性为内部实现。
  3. 测试覆盖：新增了非均匀页面大小的显式测试，并与现有测试一同运行，覆盖了基本回归场景。
• 影响：
  • 用户影响：修复 GLM-5.2 FP8 模型启动崩溃，用户无需修改配置即可正常使用。对于其他含多个 KV 缓存组但页面大小不同的模型（如未来架构），也自动受益。
  • 系统影响：零初始化内核仍为单次网格启动，未引入额外 kernel launch，性能影响极小。
  • 团队影响：代码量 +92/-31，改动集中在 3 个文件，评审成本低；数据契约变更已一并处理，无遗留适配项。
  • 风险标记：核心路径变更, 数据契约调整, 新增测试覆盖

关联脉络

• PR #49696 KVBlockZeroer crashes with non-uniform page sizes on GLM-5.2 FP8: 关联 issue，直接触发本 PR 的修复需求。
• PR #47574 Enable KV cache zeroing for mixed-precision KV cache: 根因 PR：改变了 needs_kv_cache_zeroing 逻辑，使 GLM-5.2 的零初始化被启用，暴露了 KVBlockZeroer 的统一页面大小假设缺陷。
• PR #48597 [Perf][GLM-5.2] Blackwell decode optimizations: 同系列模型优化，涉及 MLA + DSA indexer 的 KV 缓存组，本 PR 的修复使这些优化可以正常结合 KV 缓存零初始化。