[FEATURE] HY V4: fp8 KV cache dequant to BF16 for sparse attention (L… - #117
alexanderbin123 wants to merge 2 commits into
Conversation
…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).
There was a problem hiding this comment.
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,未执行代码或重跑测试;结论仍需维护者核验。
AI CI 失败分析工作流:HCU PR CI 总结
与本次改动的关系无法确定:不同失败原因与本次改动的关联判断不一致,需要人工核验。 建议处理
|
There was a problem hiding this comment.
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,未执行代码或重跑测试;结论仍需维护者核验。
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=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.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.py—gather_dequantize_fp8_ds_mla_cachedelegates to LightOp'sdecode_gather_and_up_convert_with_indices. It reads only thetopk-selected slots of the 656-bytefp8_ds_mlalayout (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,-1preserved). Gathering only thetopkslots is ~60× less work than upconverting the whole pool;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.cache_seqlensis passed as a full-width cap (= topk), so validity depends solely on the-1entries intopk_indices— matching the fp8 kernel (which is fedcache_lens = max_model_len, an upper bound) and the Triton path bit-for-bit._resolve_lightop_gatherraises rather than silently falling back to the fp8 kernel. The opt-in is deliberate.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. (Identical to the Triton variant; the swap is entirely behindgather_dequantize_fp8_ds_mla_cache.)envs.py— declaresVLLM_HCU_HYV4_FP8_KV_DEQUANT(defaultFalse).Testing
-1rows zeroed,compact_indices == t*topk+kwith-1preserved.Hy4-preview-Channel-FP8-w8a8with 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.pyandenvs.pyare identical between the two branches; onlyfp8_kv_dequant.pydiffers (LightOp vendor op vs Triton kernel). Both gather only thenum_tokens * topkselected slots (not the whole pool). Pick one to merge.