[FEATURE] HY V4: optional fp8 KV cache dequant to BF16 for sparse att… - #116
alexanderbin123 wants to merge 1 commit into
Conversation
…ention 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. - fp8_kv_dequant.py: a fused Triton gather+dequant kernel 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, and returns 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. - 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. Kernel numerics also checked bit-exact against an independent torch reference on realistic fp8 data.
8b24446 to
c222098
Compare
AI CI 失败分析工作流：HCU PR CI 总结
与本次改动的关系无法确定：本批提供的 diff 信息显示 0/3 文件有可用 patch，无法核对 PR 是否修改了断言涉及的 workflow 文件或该测试本身；Run SHA 与 PR SHA 不同也不能单独证明关联或无关。 建议处理
|
![hygon-ai-ai-reviewer[bot]](https://avatars.githubusercontent.com/u/280652917?s=60&v=4){kind=small}
There was a problem hiding this comment.
AI Review
🟢 未发现有明确证据的问题
在已提供并完成审查的 diff 范围内未形成可确认问题。
变更概览
本次变更主要包含：
- 新增可选路径：环境变量开启时，将 fp8_ds_mla KV 缓存中仅 sparse kernel 实际读取的 topk 槽位反量化为 BF16 紧凑缓冲，并把注意力计算改走几何无关的 BF16 稀疏核，从而绕过 fp8 核硬编码的布局约束，同时保持 CUDA-graph 可捕获的固定输出形状。
- 在环境变量定义中新增布尔开关 VLLM_HCU_HYV4_FP8_KV_DEQUANT（默认关闭），用于控制 HY V4 的 fp8 KV cache 在稀疏注意力前反量化为 BF16 而非调用写死 pe_dim==64 几何的 fp8 FlashMLA 内核，并同时覆盖 prefill 与 decode 路径。
文件审查摘要
| 文件 | 变更 | 审查结果 |
|---|---|---|
vllm_hcu/models/hy_v4/fp8_kv_dequant.py |
新增 · +146/-0 | — |
vllm_hcu/models/hy_v4/hcu_sparse.py |
修改 · +49/-1 | — |
vllm_hcu/platforms/envs.py |
修改 · +9/-0 | — |
审查信息
- 变更统计：3 个文件，+204/-1。
- 覆盖情况：共 3 个文件，已完整审查 3 个。
- 候选问题：0 项；证据复核过滤：0 项；发布前敏感信息保护：0 项。
- 本服务只审查 GitHub 提供的 PR diff，未执行代码或重跑测试；结论仍需维护者核验。
AI CI 失败分析工作流：HCU PR CI 总结
与本次改动的关系无法确定：已定位的直接异常对本次改动的关联判断一致。 建议处理
|
1 participant
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 a fused Triton gather+dequant kernel.Why
HY V4's sparse MLA runs in mixed-batch mode (local
num_heads=8in TP8< 32), so all tokens — prefill and decode — route through the single fp8 FlashMLA kernel. That kernel hardcodes DeepSeek'sfp8_ds_mlageometry (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.How
fp8_kv_dequant.py— a fused Triton gather+dequant kernel over the 656-bytefp8_ds_mlalayout (512 fp8 NoPE + 4 fp32 tile scales + 64 bf16 RoPE). It does not upconvert the whole pool: the kernel is launched over exactlynum_tokens * topkprograms, gathering and dequantizing only the topk-selected slots into a compact(num_tokens*topk, 576)BF16 buffer, and returns the sparse indices remapped to that buffer's rows (t*topk+k,-1preserved). Gathering only the topk slots is ~60× less work than upconverting the whole pool (≈16K slots vs ≈1M slots/layer).topkis constant andnum_tokensis fixed per graph batch, so the output shape is static and the decode path stays CUDA-graph-capturable.flash_mla_sparse_fwdmasks the-1indices natively.hcu_sparse.py— intercepts_fp8_flash_mla_kernelwhen the env var is set and dispatches to the BF16 sparse kernel on the gathered/dequantized cache, passing the remapped indices.envs.py— declaresVLLM_HCU_HYV4_FP8_KV_DEQUANT(defaultFalse).Testing
Hy4-preview-Channel-FP8-w8a8with 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.Related
A LightOp-backed variant of the same env-gated feature is #117 (
feat/hyv4-fp8-kv-dequant-lightop) —hcu_sparse.pyandenvs.pyare identical; onlyfp8_kv_dequant.pydiffers (LightOp vendor op vs this Triton kernel). Pick one to merge.