feat: fp8kv support

Author: niushengxiao

docs/CN/source/index.rst

@@ -49,6 +49,7 @@ Lightllm 整合了众多的开源方案的优点，包括但不限于 FasterTran
    :caption: 部署教程
 
    DeepSeek R1 部署 <tutorial/deepseek_deployment>
+   FP8 KV 量化与校准 <tutorial/fp8_kv_quantization>
    多级缓存部署 <tutorial/multi_level_cache_deployment>
    多模态部署 <tutorial/multimodal>
    奖励模型部署 <tutorial/reward_model>

docs/CN/source/tutorial/fp8_kv_quantization.rst

@@ -0,0 +1,98 @@
+.. _tutorial/fp8_kv_quantization_cn:
+
+FP8 KV 量化与校准指南
+======================
+
+本章节介绍 LightLLM 中 FP8 KV 推理的使用方式，包括：
+
+- 使用校准文件进行推理（``fp8kv``）
+- FA3 与 FlashInfer 后端下的量化粒度差异
+- 常见报错与排查建议
+
+功能概览
+--------
+
+LightLLM 的 FP8 KV 推理需要准备好的校准文件（``kv_cache_calib.json``），
+并通过 ``--kv_quant_calibration_config_path`` 加载。
+你可以直接使用 ``test/advanced_config/`` 目录下已有的校准文件，
+也可以使用 `LightCompress <https://github.com/ModelTC/LightCompress>`_ 工具导出，或使用自有兼容文件。
+
+后端与量化粒度
+--------------
+
+当前行为如下：
+
+- ``fa3``: 使用 ``per_head``（每个 head 独立 scale）
+- ``flashinfer``: 使用 ``per_tensor``（K/V 各一个标量 scale）
+
+因此，校准文件与后端强相关：
+
+- ``fa3`` 对应 ``per_head`` 校准文件，应配合 ``fa3`` 推理。
+- ``flashinfer`` 对应 ``per_tensor`` 校准文件，应配合 ``flashinfer`` 推理。
+
+不建议混用不同后端的校准文件。
+
+使用校准文件启动 FP8 推理
+-------------------------
+
+推理模式示例（FA3）：
+
+.. code-block:: console
+
+    $ python -m lightllm.server.api_server \
+        --model_dir /path/to/model \
+        --llm_kv_type fp8kv \
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3 \
+        --kv_quant_calibration_config_path /path/to/kv_cache_calib.json
+
+推理模式示例（FlashInfer）：
+
+.. code-block:: console
+
+    $ python -m lightllm.server.api_server \
+        --model_dir /path/to/model \
+        --llm_kv_type fp8kv \
+        --llm_prefill_att_backend flashinfer \
+        --llm_decode_att_backend flashinfer \
+        --kv_quant_calibration_config_path /path/to/kv_cache_calib.json
+
+说明：
+
+- ``fp8kv`` 模式必须提供 ``--kv_quant_calibration_config_path``。
+- 建议推理时的 attention backend 与校准文件要求保持一致。
+
+校准文件格式
+------------
+
+``kv_cache_calib.json`` 主要字段包括：
+
+- ``quant_type``: ``per_head`` 或 ``per_tensor``
+- ``num_layers``: 层数
+- ``num_head``: 总 head 数
+- ``scales_shape``: scale 张量形状
+- ``scales``: 实际 scale 数值
+- ``qmin`` / ``qmax``: FP8 范围参数
+
+加载校准文件时，会校验模型架构、层数、head 数及量化类型是否匹配。
+
+多卡说明
+--------
+
+在多卡（TP）场景下，系统会根据当前 rank 自动切分本地需要的 head 对应 scale。
+你仍然只需要提供一份全量 ``kv_cache_calib.json``。
+
+常见问题
+--------
+
+1. 启动时报错需要 ``--kv_quant_calibration_config_path``
+
+   说明你使用了 ``--llm_kv_type fp8kv`` 但未传入校准文件路径。
+
+2. 报错 ``quant_type not match``
+
+   通常是后端与校准文件类型不一致。例如拿 ``per_head`` 文件去跑 ``flashinfer``。
+
+3. 切换后端后效果异常
+
+   建议使用与目标后端匹配的校准文件，不要跨后端复用不兼容文件。

docs/EN/source/index.rst

@@ -48,6 +48,7 @@ Documentation List
    :caption: Deployment Tutorials
 
    DeepSeek R1 Deployment <tutorial/deepseek_deployment>
+   FP8 KV Quantization and Calibration <tutorial/fp8_kv_quantization>
    Multi-Level Cache Deployment <tutorial/multi_level_cache_deployment>
    Multimodal Deployment <tutorial/multimodal>
    Reward Model Deployment <tutorial/reward_model>

docs/EN/source/tutorial/fp8_kv_quantization.rst

@@ -0,0 +1,98 @@
+.. _tutorial/fp8_kv_quantization_en:
+
+FP8 KV Quantization and Calibration Guide
+=========================================
+
+This chapter describes FP8 KV inference in LightLLM, including:
+
+- Running inference with calibration data (``fp8kv``)
+- Quantization granularity differences between FA3 and FlashInfer
+- Common errors and troubleshooting
+
+Overview
+--------
+
+LightLLM FP8 KV inference requires a prepared calibration file (``kv_cache_calib.json``),
+which is loaded by ``--kv_quant_calibration_config_path``.
+You can use calibration files provided in ``test/advanced_config/``,
+export one with `LightCompress <https://github.com/ModelTC/LightCompress>`_, or use your own compatible file.
+
+Backend and Quantization Granularity
+------------------------------------
+
+Current behavior:
+
+- ``fa3``: ``per_head`` scales (independent scale per head)
+- ``flashinfer``: ``per_tensor`` scales (one scalar for K and one scalar for V)
+
+Calibration files are backend-dependent:
+
+- ``per_head`` files for ``fa3`` should be used with ``fa3`` inference.
+- ``per_tensor`` files for ``flashinfer`` should be used with ``flashinfer`` inference.
+
+Avoid mixing calibration files across different backends.
+
+Start FP8 Inference with Calibration
+------------------------------------
+
+Inference mode example (FA3):
+
+.. code-block:: console
+
+    $ python -m lightllm.server.api_server \
+        --model_dir /path/to/model \
+        --llm_kv_type fp8kv \
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3 \
+        --kv_quant_calibration_config_path /path/to/kv_cache_calib.json
+
+Inference mode example (FlashInfer):
+
+.. code-block:: console
+
+    $ python -m lightllm.server.api_server \
+        --model_dir /path/to/model \
+        --llm_kv_type fp8kv \
+        --llm_prefill_att_backend flashinfer \
+        --llm_decode_att_backend flashinfer \
+        --kv_quant_calibration_config_path /path/to/kv_cache_calib.json
+
+Notes:
+
+- ``fp8kv`` requires ``--kv_quant_calibration_config_path``.
+- Keep the inference backend consistent with the backend expected by the calibration file.
+
+Calibration File Schema
+-----------------------
+
+Key fields in ``kv_cache_calib.json``:
+
+- ``quant_type``: ``per_head`` or ``per_tensor``
+- ``num_layers``: number of layers
+- ``num_head``: total number of heads
+- ``scales_shape``: shape of the scale tensor
+- ``scales``: actual scale values
+- ``qmin`` / ``qmax``: FP8 numeric range parameters
+
+At load time, LightLLM validates architecture, layer count, head count, and quantization type.
+
+Multi-GPU Note
+--------------
+
+In multi-GPU (TP) setups, LightLLM slices the global scales to local rank heads automatically.
+You only need to provide one full ``kv_cache_calib.json`` file.
+
+Common Issues
+-------------
+
+1. Error says ``--kv_quant_calibration_config_path`` is required
+
+   You are using ``--llm_kv_type fp8kv`` without a calibration file path.
+
+2. ``quant_type not match`` error
+
+   Usually caused by backend/file mismatch (for example, using a ``per_head`` file with ``flashinfer``).
+
+3. Abnormal quality after backend switch
+
+   Use a calibration file that matches the target backend instead of reusing an incompatible file.

lightllm/common/basemodel/attention/create_utils.py

@@ -35,6 +35,10 @@
         # "fa3": Fp8Fa3AttBackend,
         # "flashinfer": Fp8FlashInferAttBackend,
     },
+    "fp8kv": {
+        "fa3": Fp8Fa3AttBackend,
+        "flashinfer": Fp8FlashInferAttBackend,
+    },
 }
 
 mla_data_type_to_backend = {

lightllm/common/basemodel/attention/fa3/fp8.py

@@ -89,19 +89,21 @@ def _fp8_prefill_att(
     ) -> torch.Tensor:
         self.backend: Fp8Fa3AttBackend = self.backend  # for typing
 
+        q_head_num = q.shape[1]
+        q_head_dim = q.shape[2]
+        k_head_num = k.shape[1]
         q, q_scale = q_per_head_fp8_quant(
-            q,
+            q.reshape(q.shape[0], k_head_num, -1),
             self.infer_state.b_seq_len,
             self.cu_seqlens_q,
-            self.mid_token_batch_ids,
+            token_batch_ids=self.mid_token_batch_ids,
         )
-        k_head_num = k.shape[1]
         k_head_dim = k.shape[2]
         cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn)
         cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn)
         layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self)
         o = flash_attn_with_kvcache(
-            q=q,
+            q=q.reshape(-1, q_head_num, q_head_dim),
             k_cache=cache_k,
             v_cache=cache_v,
             page_table=self.page_table,
@@ -200,9 +202,11 @@ def _fp8_decode_att(
         layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self)
 
         q_head_num = q.shape[1]
-        q, q_scale = scaled_fp8_quant(q.view(q.shape[0] * k_head_num, -1), use_per_token_if_dynamic=True)
+        if scaled_fp8_quant is None:
+            raise ImportError("scaled_fp8_quant is unavailable. Please install vllm to enable FP8 decode attention.")
+        q, q_scale = scaled_fp8_quant(q.reshape(q.shape[0] * k_head_num, -1), use_per_token_if_dynamic=True)
         o = flash_attn_with_kvcache(
-            q=q.view(-1, q_head_num, k_head_dim),
+            q=q.reshape(-1, q_head_num, k_head_dim),
             k_cache=cache_k,
             v_cache=cache_v,
             page_table=self.page_table,

lightllm/common/basemodel/attention/flashinfer/fp8.py

@@ -20,7 +20,7 @@ def create_att_decode_state(self, infer_state) -> "Fp8FlashInferDecodeAttState":
 
 @dataclasses.dataclass
 class Fp8FlashInferPrefillAttState(FlashInferPrefillAttState):
-    offline_scales: torch.Tensor = None
+    offline_scales: list = None
 
     def init_state(self):
         super().init_state()
@@ -68,7 +68,7 @@ def _fp8_prefill_att(
 
 @dataclasses.dataclass
 class Fp8FlashInferDecodeAttState(FlashInferDecodeAttState):
-    offline_scales: torch.Tensor = None
+    offline_scales: list = None
 
     def init_state(self):
         super().init_state()

lightllm/common/kv_cache_mem_manager/__init__.py

@@ -1,6 +1,5 @@
 from .mem_manager import MemoryManager, ReadOnlyStaticsMemoryManager
 from .calibration_fp8kv_mem_manager import CalibrationFP8KVMemoryManager
-from .export_calibration_mem_manager import ExportCalibrationMemoryManager
 from .ppl_int8kv_mem_manager import PPLINT8KVMemoryManager
 from .ppl_int4kv_mem_manager import PPLINT4KVMemoryManager
 from .deepseek2_mem_manager import Deepseek2MemoryManager
@@ -10,7 +9,6 @@
     "MemoryManager",
     "ReadOnlyStaticsMemoryManager",
     "CalibrationFP8KVMemoryManager",
-    "ExportCalibrationMemoryManager",
     "PPLINT4KVMemoryManager",
     "PPLINT8KVMemoryManager",
     "Deepseek2MemoryManager",

lightllm/common/kv_cache_mem_manager/calibration_fp8kv_mem_manager.py

@@ -1,6 +1,28 @@
+import torch
+from typing import Tuple, Any
 from .offline_fp8_quant_mem_manager import OfflineFP8QuantMemManager
 
 
 class CalibrationFP8KVMemoryManager(OfflineFP8QuantMemManager):
     def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
-        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction, is_export_mode=False)
+        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction)
+
+    def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+        """
+        推理模式：使用预计算的FP8 scale将kv量化后拷贝到kv_buffer中
+        """
+        from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8
+
+        scales = self.scales
+        destindex_copy_kv_fp8(
+            kv,
+            mem_index,
+            scales[layer_index] if scales is not None else None,
+            self.kv_buffer[layer_index].view(torch.float8_e4m3fn),
+        )
+        return
+
+    def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]:
+        k = self.kv_buffer[layer_index][:, : self.head_num, :]
+        v = self.kv_buffer[layer_index][:, self.head_num :, :]
+        return k, v