Skip to content

[FEATURE] HY V4: fp8 KV cache dequant to BF16 for sparse attention (L… - #117

Open
alexanderbin123 wants to merge 2 commits into
v0.25.1from
feat/hyv4-fp8-kv-dequant-lightop
Open

alexanderbin123 wants to merge 2 commits into
v0.25.1from
feat/hyv4-fp8-kv-dequant-lightop

Conversation

@alexanderbin123

@alexanderbin123 alexanderbin123 commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

What

Adds an opt-in path for HY V4 sparse MLA, gated by VLLM_HCU_HYV4_FP8_KV_DEQUANT, that dequantizes the fp8 KV cache to BF16 and runs the geometry-agnostic BF16 sparse kernel (flash_mla_sparse_fwd), using LightOp's vendor-optimized HCU op for the gather + up-convert.

Why

HY V4's sparse MLA runs in mixed-batch mode (local num_heads=8 in TP8 < 32), so all tokens — prefill and decode — route through the single fp8 FlashMLA kernel. That kernel hardcodes DeepSeek's fp8_ds_mla geometry (pe_dim==64). When the env var is set, we intercept that call and run the BF16 sparse kernel on an upconverted cache instead, covering both prefill and decode.

This is the LightOp variant of the same feature: it swaps the standalone Triton gather/dequant kernel for LightOp's fused decode_gather_and_up_convert_with_indices, which does the gather + FP8→BF16 up-convert + compact-index remap in one vendor op.

How

  • fp8_kv_dequant.pygather_dequantize_fp8_ds_mla_cache delegates to LightOp's decode_gather_and_up_convert_with_indices. It reads only the topk-selected slots of the 656-byte fp8_ds_mla layout (512 fp8 NoPE + 4 fp32 tile scales + 64 bf16 RoPE) into a compact (num_tokens*topk, 576) BF16 buffer, and returns the sparse indices remapped to that buffer's rows (t*topk+k, -1 preserved). Gathering only the topk slots is ~60× less work than upconverting the whole pool; topk is constant and num_tokens is fixed per graph batch, so the output shape is static and the decode path stays CUDA-graph-capturable. flash_mla_sparse_fwd masks the -1 indices natively.
    • cache_seqlens is passed as a full-width cap (= topk), so validity depends solely on the -1 entries in topk_indices — matching the fp8 kernel (which is fed cache_lens = max_model_len, an upper bound) and the Triton path bit-for-bit.
    • Fails closed: if LightOp is missing or lacks the op, _resolve_lightop_gather raises rather than silently falling back to the fp8 kernel. The opt-in is deliberate.
  • hcu_sparse.py — intercepts _fp8_flash_mla_kernel when the env var is set and dispatches to the BF16 sparse kernel on the gathered/dequantized cache, passing the remapped indices. (Identical to the Triton variant; the swap is entirely behind gather_dequantize_fp8_ds_mla_cache.)
  • envs.py — declares VLLM_HCU_HYV4_FP8_KV_DEQUANT (default False).

Testing

  • Contract: LightOp output verified bit-exact against an independent Triton oracle on realistic fp8 data — valid rows byte-identical (max_abs_diff = 0 for both NoPE and RoPE), -1 rows zeroed, compact_indices == t*topk+k with -1 preserved.
  • Integration: TP8 smoke test on Hy4-preview-Channel-FP8-w8a8 with the env var enabled — warmup and both PIECEWISE and FULL CUDA-graph capture pass, coherent generation, LightOp op runs under graph capture.

Relationship to the Triton MR

This is an alternative backend for the same env-gated feature; the Triton variant is #116 (feat/hyv4-fp8-kv-dequant-decode). hcu_sparse.py and envs.py are identical between the two branches; only fp8_kv_dequant.py differs (LightOp vendor op vs Triton kernel). Both gather only the num_tokens * topk selected slots (not the whole pool). Pick one to merge.

…ightOp)

HY V4's sparse MLA runs in mixed-batch mode (num_heads=8 in TP8 < 32), so
all tokens (prefill and decode) route through the fp8 FlashMLA kernel. That
kernel hardcodes DeepSeek's fp8_ds_mla geometry. This adds an opt-in path,
gated by VLLM_HCU_HYV4_FP8_KV_DEQUANT, that dequantizes the fp8 KV cache to
BF16 and runs the geometry-agnostic BF16 sparse kernel (flash_mla_sparse_fwd)
instead, covering both prefill and decode.

This is the LightOp variant: the gather + FP8->BF16 up-convert + compact-index
remap is done by LightOp's fused vendor op
decode_gather_and_up_convert_with_indices instead of a standalone Triton
kernel.

- fp8_kv_dequant.py: gather_dequantize_fp8_ds_mla_cache delegates to LightOp's
  decode_gather_and_up_convert_with_indices over the 656-byte fp8_ds_mla
  layout (512 fp8 NoPE + 4 fp32 tile scales + 64 bf16 RoPE). It reads only the
  topk-selected slots and writes a compact (num_tokens*topk, 576) BF16 buffer,
  returning the sparse indices remapped to that buffer's rows (-1 preserved).
  Gathering only the topk slots is ~60x less work than upconverting the whole
  pool; topk is constant and num_tokens is fixed per graph batch, so the
  output shape is static and the decode path stays CUDA-graph-capturable.
  flash_mla_sparse_fwd masks the -1 indices natively. cache_seqlens is a
  full-width cap, so validity depends solely on -1 in topk_indices, matching
  the fp8 kernel (cache_lens = max_model_len) and the Triton path. Missing
  LightOp fails closed rather than silently falling back.
- hcu_sparse.py: intercept _fp8_flash_mla_kernel when the env var is set and
  dispatch to the BF16 sparse kernel on the gathered/dequantized cache,
  passing the remapped indices.
- envs.py: declare VLLM_HCU_HYV4_FP8_KV_DEQUANT (default False).

Verified with a TP8 smoke test on Hy4-preview-Channel-FP8-w8a8 with the env
var enabled: warmup and both PIECEWISE and FULL CUDA-graph capture pass,
coherent generation, and the "BF16 sparse prefill kernel" log line (emitted
during FULL capture) confirms the dequant path is taken. LightOp output also
checked bit-exact against an independent Triton reference on realistic fp8
data (valid rows byte-identical, -1 rows zeroed, compact indices t*topk+k).

@hygon-ai-ai-reviewer hygon-ai-ai-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Review

🟢 未发现有明确证据的问题

在已提供并完成审查的 diff 范围内未形成可确认问题。

变更概览

本次变更主要包含:

  • 新增 fp8_kv_dequant 模块,将 fp8_ds_mla KV 缓存中 topk 选中条目 gather 并反量化为紧凑 BF16 缓冲;并在设置 VLLM_HCU_HYV4_FP8_KV_DEQUANT 时,把 HY V4 的 fp8 稀疏注意力内核替换为基于反量化缓冲的 BF16 稀疏内核路径,返回值改为允许缺失 LSE。
  • 在环境变量配置中新增默认关闭的开关 VLLM_HCU_HYV4_FP8_KV_DEQUANT,用于选择在稀疏注意力前将 HY V4 的 fp8 KV cache 反量化为 BF16,而非调用硬编码 fp8_ds_mla 几何的 fp8 FlashMLA 内核。
文件审查摘要
文件 变更 审查结果
vllm_hcu/models/hy_v4/fp8_kv_dequant.py 新增 · +126/-0
vllm_hcu/models/hy_v4/hcu_sparse.py 修改 · +49/-1
vllm_hcu/platforms/envs.py 修改 · +9/-0
审查信息
  • 变更统计:3 个文件,+184/-1。
  • 覆盖情况:共 3 个文件,已完整审查 3 个。
  • 候选问题:0 项;证据复核过滤:0 项;发布前敏感信息保护:0 项。
  • 本服务只审查 GitHub 提供的 PR diff,未执行代码或重跑测试;结论仍需维护者核验。

@hygon-ai-ai-reviewer

Copy link
Copy Markdown

AI CI 失败分析

工作流:HCU PR CI
状态:失败

总结

  • Static gate and test selection:两个静态测试失败:test_lightop_api_boundary.py:544 扫描 vllm_hcu 发现 vllm_hcu/models/hy_v4/fp8_kv_dequant.py:50 存在 "moved top-level LightOp import 'op'" 违规(L332/L310);test_hcu_ci_selector.py:801 断言仓库文件含 '! -uid "$uid"' 未命中(L299/L331)。L342 汇总 2 failed,随后退出码 1。

  • ci-gate:日志直接证明:门禁步骤读到 STATIC_RESULT=failure(L36),命中并输出 'Static gate or selector failed.'(L45),随后以 exit 1 结束(L46),属确定性失败路径。上游 static/selector 任务自身为何失败不在本批日志内,无法确定。

与本次改动的关系

无法确定:不同失败原因与本次改动的关联判断不一致,需要人工核验。

建议处理

  1. 对照 test_lightop_api_boundary 的扫描规则核查该文件第50行导入方式是否允许,按规则调整导入位置或形式;对 hcu_ci_selector 失败,核对对应工作流文件是否缺少 '! -uid "$uid"' 片段。上述均需人工核验后再重跑测试确认。

  2. 最小处理方向:结合其它批次的 static 门禁/selector 任务日志定位其失败原因;本批可见门禁脚本按设计条件分支退出,无证据要求修改该汇总逻辑。

@hygon-ai-ai-reviewer hygon-ai-ai-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Review

🟢 未发现有明确证据的问题

在已提供并完成审查的 diff 范围内未形成可确认问题。

变更概览

本次变更主要包含:

  • 新增 HY V4 的 fp8 KV cache 反量化模块:通过环境变量选择性启用,调用 LightOp 融合算子把 fp8_ds_mla 布局缓存中 topk 选中的槽位 gather 并升转 BF16,输出紧凑缓冲与重映射索引,供 BF16 稀疏注意力 kernel 使用;LightOp 缺失或算子不存在时显式报错而非静默回退。
  • 在 HY V4 稀疏注意力的 fp8 kernel 入口增加环境变量开关:开启时将 fp8 KV cache 仅对 topk 选中的槽位反量化为 BF16 并改走 BF16 稀疏 kernel,同时函数返回类型的 lse 允许为 None。
  • 在平台环境开关注册中新增布尔开关 VLLM_HCU_HYV4_FP8_KV_DEQUANT(默认关闭),通过环境变量启用后,HY V4 fp8 KV cache 在稀疏注意力前反量化为 BF16,而非调用硬编码 DeepSeek fp8_ds_mla 几何的 fp8 FlashMLA 内核。
文件审查摘要
文件 变更 审查结果
vllm_hcu/models/hy_v4/fp8_kv_dequant.py 新增 · +126/-0
vllm_hcu/models/hy_v4/hcu_sparse.py 修改 · +49/-1
vllm_hcu/platforms/envs.py 修改 · +9/-0
审查信息
  • 变更统计:3 个文件,+184/-1。
  • 覆盖情况:共 3 个文件,已完整审查 3 个。
  • 候选问题:0 项;证据复核过滤:0 项;发布前敏感信息保护:0 项。
  • 本服务只审查 GitHub 提供的 PR diff,未执行代码或重跑测试;结论仍需维护者核验。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant