Prhub

#51288 [Test] Add packed DeepSeek-V4 KV zeroer geometry regression

原始 PR 作者 coltonottley 合并时间 2026-08-08 03:11 文件变更 1 提交数 1 评论 3 代码增减 +195 / -0

执行摘要

新增打包 DeepSeek-V4 KV 清零几何回归测试

50276 指出 KVBlockZeroer 在打包 KV 视图下将页面大小同时用作块地址步幅与清零元素数,会清除相邻打包数据并对最终逻辑块越界写。本 PR 为该内存安全修复提供回归保护:如 PR body 所述,测试实例化真实生产 seam 并验证构造函数元数据将完整打包行块步幅与每层有效清零跨度分开,防止该缺陷回归。

值得精读。该测试展示了如何为内部数据布局编写 constructor-crossing 回归测试:不重新实现任何布局算术,而是实例化真实生产代码链并以元数据语义作为断言目标;同时用新旧格式兼容分支让失败信息直接指向缺陷本质。对 KV 缓存布局、MLA 和 KVBlockZeroer 相关改动者是很好的防护样例。

讨论亮点

PR 的 review 无技术交锋:维护者 njhill 通过 /ci run 触发 CI 后批准合并(“Thanks @coltonottley”);claude[bot] 提示 fork 仓库自动审核被禁用。PR body 中的“回归判别器”说明最值得注意:测试兼容旧五字段与新六字段 _meta,使失败成为语义性判别(zero span 应为 37376 B 的有效页,而非整个打包行 149760 B),而非 schema 形状错误。

实现拆解

  1. 新增测试文件 tests/v1/worker/test_dsv4_packed_zeroer_geometry.py(本 PR 唯一变更,+195 行),以 pytest.mark.cpu_test 标记,保证纯 CPU 执行、不启动任何 CUDA/Triton 内核。
  2. 构建规格:为 4 个 self_attn 层构造 MLAAttentionSpecblock_size=256head_size=512cache_dtype_str='fp8_ds_mla'alignment=576compress_ratio=4),经 UniformTypeKVCacheSpecs.from_specs 得到统一缓存规格。
  3. 调用真实布局规划器 _get_packed_kv_cache_layout,得到块步幅(4 * page_bytes)与各层偏移;再调用 _reshape_attention_kv_cache 生成每层打包视图,校验 data_ptr 相对偏移与行步幅。
  4. 实例化 KVBlockZeroer,传入 DeepseekV4FlashMLABackendkernel_block_sizes=[256] 与打包视图上下文;通过 _extract_segments 读取 zeroer._meta 段表,兼容修复前五字段与修复后六字段两种格式。
  5. 断言语义:4 个段对应 4 个打包层、每段块步幅等于完整打包行、清零跨度等于有效页 real_page_size_bytes、block 99 的写区间 [off + 99*bs, off + 99*bs + zs) 不越界、最高偏移段末块精确结束于 total_size - alignment_gap,且清零跨度严格小于块步幅。
  6. 验证配套:作者在 #50276 head 上跑通 pytestruffmypy,在 pre-fix parent 上确认语义性失败(zero span 为整个打包行 149760 B 而非有效页 37376 B)。
文件 模块 状态 重要度
tests/v1/worker/test_dsv4_packed_zeroer_geometry.py KV 清零 added 6.84

关键符号

_extract_segments test_dsv4_packed_zeroer_geometry

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

CI 启动与运行 other

维护者 njhill 评论 `/ci run`,github-actions 确认 CI 构建已在 Buildkite 运行(build 82721)。

结论:CI 正常执行,未发现失败记录。 · 已解决

PR 审核 other

njhill 批准 PR 并致谢作者;claude[bot] 说明 fork 仓库自动审核被禁用。

结论:已合并。 · 已解决

风险与影响

纯测试变更,无生产路径修改,运行时风险极低。风险集中在维护耦合:测试依赖 KVBlockZeroer._meta 六字段内部结构以及 _get_packed_kv_cache_layout_reshape_attention_kv_cache 等内部 API,后续重构需同步维护;_extract_segments 的五/六字段分支提供一定兼容缓冲。测试仅在 CPU 上验证构造函数元数据,不覆盖真实 Triton/CUDA 内核清零行为,不能替代端到端验证。另外本 PR 明确堆叠在 #50276 之上,若 #50276 未先合入,测试会失败。

对用户与系统无运行时影响,不改变任何推理行为。对工程团队,该测试为 V1 KV 缓存清零这一内存安全关键路径提供回归防护,能在 CI 中快速执行(CPU 测试,无 GPU 占用),保护 DeepSeek-V4 打包 KV 布局的后续修改。

纯测试变更 依赖 #50276 先合入 耦合 KVBlockZeroer 内部元数据

关联 Issue

#50276 [Bugfix] Fix packed KV block zeroing stride

完整报告

参与讨论