Skip to content

feat: add per-tensor FP8 e4m3 quant/dequant ops for KV cache quantiza… - #1521

Open
BoBoDai wants to merge 2 commits into
InfiniTensor:mainfrom
BoBoDai:suppose-fp8
Open

feat: add per-tensor FP8 e4m3 quant/dequant ops for KV cache quantiza…#1521
BoBoDai wants to merge 2 commits into
InfiniTensor:mainfrom
BoBoDai:suppose-fp8

Conversation

@BoBoDai

@BoBoDai BoBoDai commented Aug 23, 2026

Copy link
Copy Markdown

feat: add per-tensor FP8 e4m3 quant/dequant ops for KV cache quantization

摘要

本 PR 新增两个算子:per_tensor_quant_fp8(量化)与 per_tensor_dequant_fp8(反量化),实现 FP8 e4m3 格式的对称 per-tensor 量化。二者为 InfiniLM 后续 FP8 KV cache 量化(--kv-cache-dtype fp8,尚未实现)提供前置算子基础:存储 K/V 时调用 quant 算子将状态转换为 FP8 e4m3,attention 读取时调用 dequant 算子还原,均为该路径每次 decode 所需的核心内核。本 PR 仅包含算子实现与配套测试,KV cache 侧的接入与基准测试属于后续工作。

动机:FP8 e4m3 是当前主流推理框架(vLLM、SGLang 等)KV cache 量化的常用格式,
在同等 1 字节占用下精度优于 INT8,契合低精度格式方向。

e4m3 格式

  • 布局:1 符号位 + 4 指数位(bias 7)+ 3 尾数位;取值范围 [-448, 448];最小正规数 2^-6,最小次正规数 2^-9
  • 仅支持对称量化(无 zero point),与主流 FP8 KV cache 实现保持一致。
  • float <-> e4m3 转换由内核手写实现(round-half-away-from-zero,饱和至 0x7E/0xFE),不依赖 CUDA fp8 转换头。
  • 解码侧保留次正规数(最低至 2^-9)。

API

算子接口完全参照现有 per_tensor_quant_int8per_tensor_dequant_int8 的实现约定:

| per_tensor_quant_fp8 | infiniopPerTensorQuantFp8(desc, ws, x_packed, x_scale, x, is_static, stream) |
| per_tensor_dequant_fp8 | infiniopPerTensorDequantFp8(desc, ws, x, x_packed, x_scale, stream) |

  • 输入输出:quant 为 x(fp16/bf16/fp32) -> x_packed(fp8/uint8) + x_scale(fp32, [1]);dequant 为 x_packed(fp8/uint8) + x_scale(fp32, [1]) -> x(fp16/bf16/fp32)。
  • 均支持动态(运行时计算 scale = max|x| / 448)与静态(调用方给定 scale)两种模式。
  • 提供标准的 infiniopGet*WorkspaceSize / infiniopDestroy*Descriptor 入口,并在 include/infiniop.hinclude/infinicore/ops.hpp 注册,含 infinicore C++ 封装(PerTensorQuantFp8 / PerTensorDequantFp8)。
  • 当前实现 CUDA(NVIDIA/QY)后端,grid-stride kernel,注册方式与 per_tensor_quant_int8 一致。

测试

test/infiniop/ 下新增三个测试文件(对照 int8 测试结构):

  • per_tensor_quant_fp8.py:动态/静态模式、strided 布局、F16/BF16/F32; 另含确定性边界用例(零、次正规数、进位舍入、精确 448、超范围饱和),以及全零输入保底用例(动态模式下 scale=0 不除零、输出全 0、scale 归一化为 1.0)。
  • per_tensor_dequant_fp8.py:形状/stride 覆盖;另含字节级边界用例(0x00/0x01/0x07/0x08/0x7E/0x7F/0x80/0xFF 等,覆盖零、次正规、最小正规、最大值与保留位模式)。
  • w8a8fp8.py:quant + dequant + matmul 端到端链路,两种模式——WMode.FP16(仅激活量化、权重保持 fp16,对应 KV cache 场景)与 WMode.FP8(完整对称 W8A8-FP8)。

验证

运行 python scripts/python_test.py --nvidia,所有算子测试全部通过。
image

运行 python scripts/python_test.py --cpu,所有算子测试全部通过。
image

后续工作(不在本 PR 范围内)

  • 将两个算子接入 InfiniLM KV cache 量化路径(KVQuantAlgo::FP8_E4M3--kv-cache-dtype fp8):存储 K/V 时以 per_tensor_quant_fp8 量化、读取时以 per_tensor_dequant_fp8 反量化,配合静态 per-cache scale。
  • 基准对比 fp16 / int8 / fp8 三种 KV cache 方案的精度与性能(PPL、内存、吞吐)。

@BoBoDai
BoBoDai requested a review from a team August 23, 2026 14:03
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