feat: add triton_fp8kv mode for deepseek2

Author: niushengxiao

lightllm/common/deepseek2_fp8kv_mem_manager.py

@@ -0,0 +1,8 @@
+import torch
+from .deepseek2_mem_manager import Deepseek2MemoryManager
+
+
+class Deepseek2FP8KVMemoryManager(Deepseek2MemoryManager):
+    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
+        # scale被追加到kv_buffer末尾, 因此加2, dtype统一改成uint8
+        super().__init__(size, torch.uint8, head_num, head_dim + 2, layer_num, always_copy, mem_fraction)

lightllm/common/mem_utils.py

@@ -18,6 +18,8 @@ def select_mem_manager_class(mode):
     elif "triton_int8kv" in mode:
         memory_manager_class = INT8KVMemoryManager
         logger.info("Model kv cache using mode triton int8kv")
+    elif "triton_fp8kv" in mode:
+        raise Exception("currently only for deepseek")
     else:
         memory_manager_class = MemoryManager
         logger.info("Model kv cache using mode normal")

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -8,12 +8,16 @@
 from lightllm.models.deepseek2.triton_kernel.destindex_copy_kv import destindex_copy_kv
 from lightllm.models.deepseek2.triton_kernel.context_flashattention_nopad import (
     context_attention_fwd,
+    context_attention_fwd_fp8,
     context_attention_fwd_no_prompt_cache,
 )
 from lightllm.models.deepseek2.triton_kernel.context_flashattention_nopad_with_v import context_attention_fwd_with_v
 from lightllm.models.deepseek2.triton_kernel.sample_kv import sample_kv
 
-from lightllm.models.deepseek2.triton_kernel.gqa_flash_decoding import gqa_token_decode_attention_flash_decoding
+from lightllm.models.deepseek2.triton_kernel.gqa_flash_decoding import (
+    gqa_token_decode_attention_flash_decoding,
+    gqa_token_decode_attention_flash_decoding_fp8,
+)
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward
 from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd
@@ -22,6 +26,7 @@
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 import os
+from lightllm.common.quantization import vLLMFP8w8a8QuantizationMethod
 
 
 class Deepseek2TransformerLayerInfer(LlamaTransformerLayerInfer):
@@ -67,19 +72,12 @@ def __init__(self, layer_num, tp_rank, world_size, network_config, mode=[]):
         self.enable_opt_decoding_mha = os.getenv("ENABLE_OPT_DECODE_MHA", "False").upper() in ["ON", "TRUE", "1"]
         return
 
-    def _bind_attention(self):
-        if self.enable_cc_method:
-            self._context_attention_kernel = partial(
-                Deepseek2TransformerLayerInfer._context_attention_kernel_with_CC, self
-            )
-        else:
-            self._context_attention_kernel = partial(
-                Deepseek2TransformerLayerInfer._context_attention_kernel_origin, self
-            )
-        self._token_attention_kernel = partial(
-            Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashdecoding, self
-        )
-        self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
+    def _bind_func(self):
+        super()._bind_func()
+        self._bind_ffn()
+        return
+
+    def _bind_ffn(self):
         if self.is_moe:
             if self.enable_dp:
                 if os.environ.get("MOE_MODE", "TP") == "TP":
@@ -92,6 +90,36 @@ def _bind_attention(self):
         else:
             self._ffn = partial(LlamaTransformerLayerInfer._ffn, self)
 
+    def _bind_attention(self):
+        if "triton_fp8kv" in self.mode:
+            self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_fp8, self)
+            self._token_attention_kernel = partial(
+                Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashdecoding_fp8, self
+            )
+        else:
+            self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
+            self._token_attention_kernel = partial(
+                Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashdecoding, self
+            )
+        if self.enable_cc_method:
+            if "triton_fp8kv" in self.mode:
+                self._context_attention_kernel = partial(
+                    Deepseek2TransformerLayerInfer._context_attention_kernel_with_CC_fp8, self
+                )
+            else:
+                self._context_attention_kernel = partial(
+                    Deepseek2TransformerLayerInfer._context_attention_kernel_with_CC, self
+                )
+        else:
+            if "triton_fp8kv" in self.mode:
+                self._context_attention_kernel = partial(
+                    Deepseek2TransformerLayerInfer._context_attention_kernel_origin_fp8, self
+                )
+            else:
+                self._context_attention_kernel = partial(
+                    Deepseek2TransformerLayerInfer._context_attention_kernel_origin, self
+                )
+
     def _get_qkv(
         self,
         input: torch.Tensor,
@@ -133,9 +161,19 @@ def _get_o(
         o_tensor = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.qk_nope_head_dim))
         return o_tensor
 
-    def _decompress_kv(self, kv, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight):
+    def _decompress_kv(
+        self, kv, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, is_fp8
+    ):
         if infer_state.use_dynamic_prompt_cache:
-            kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+            if is_fp8:
+                kv = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, :-2].view(torch.float8_e4m3fn)
+                kv_scale = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, -2:].view(torch.bfloat16)
+                k_scale = self.alloc_tensor([infer_state.total_token_num, 1], dtype=kv_scale.dtype)
+            else:
+                kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+                kv_scale = None
+                k_scale = None
+
             compressed_kv = self.alloc_tensor(
                 [infer_state.total_token_num, 1, layer_weight.kv_lora_rank], dtype=kv.dtype
             )
@@ -147,7 +185,12 @@ def _decompress_kv(self, kv, infer_state: Deepseek2InferStateInfo, layer_weight:
                 infer_state.b_req_idx,
                 infer_state.b_seq_len,
                 infer_state.req_manager.req_to_token_indexs,
+                kv_scale,
+                k_scale,
             )
+            if k_scale is not None:
+                compressed_kv = compressed_kv.to(k_scale.dtype) * k_scale.unsqueeze(-1)
+                k_rope = k_rope.to(k_scale.dtype) * k_scale.unsqueeze(-1)
         else:
             compressed_kv, k_rope = torch.split(  # (b*s, 1, kv_lora + qk_r)
                 kv, [layer_weight.kv_lora_rank, layer_weight.qk_rope_head_dim], dim=-1
@@ -177,7 +220,33 @@ def _context_attention_kernel_with_CC(
         layer_weight: Deepseek2TransformerLayerWeight,
         out=None,
     ) -> torch.Tensor:
-        k_nope, k_rope, v = self._decompress_kv(kv, infer_state, layer_weight)
+        k_nope, k_rope, v = self._decompress_kv(kv, infer_state, layer_weight, False)
+        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
+        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
+        context_attention_fwd_with_v(
+            q_nope,
+            q_rope,
+            k_nope,
+            k_rope,
+            v,
+            o_tensor.view(-1, self.tp_q_head_num_, q_nope.shape[-1]),
+            infer_state.b_start_loc,
+            infer_state.b_seq_len,
+            infer_state.b_ready_cache_len,
+            infer_state.max_len_in_batch,
+            self.softmax_scale,
+        )
+        return o_tensor
+
+    def _context_attention_kernel_with_CC_fp8(
+        self,
+        q: torch.Tensor,
+        kv,
+        infer_state: Deepseek2InferStateInfo,
+        layer_weight: Deepseek2TransformerLayerWeight,
+        out=None,
+    ) -> torch.Tensor:
+        k_nope, k_rope, v = self._decompress_kv(kv, infer_state, layer_weight, True)
         q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
         o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
         context_attention_fwd_with_v(
@@ -237,6 +306,50 @@ def _context_attention_kernel_origin(
 
         return o_tensor
 
+    def _context_attention_kernel_origin_fp8(
+        self,
+        q: torch.Tensor,
+        kv,
+        infer_state: Deepseek2InferStateInfo,
+        layer_weight: Deepseek2TransformerLayerWeight,
+        out=None,
+    ) -> torch.Tensor:
+        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
+        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
+        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
+        if infer_state.use_dynamic_prompt_cache:
+            kv = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, :-2].view(torch.float8_e4m3fn)
+            kv_scale = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, -2:].view(torch.bfloat16)
+            context_attention_fwd_fp8(
+                q_nope,
+                q_rope,
+                kv[:, :, : -self.qk_rope_head_dim],
+                kv[:, :, -self.qk_rope_head_dim :],
+                kv_scale,
+                o_tensor.view(-1, self.tp_q_head_num_, self.kv_lora_rank),
+                infer_state.b_req_idx,
+                infer_state.b_start_loc,
+                infer_state.b_seq_len,
+                infer_state.b_ready_cache_len,
+                infer_state.max_len_in_batch,
+                infer_state.req_manager.req_to_token_indexs,
+                self.softmax_scale,
+            )
+        else:
+            context_attention_fwd_no_prompt_cache(
+                q_nope,
+                q_rope,
+                kv[:, :, : -self.qk_rope_head_dim],
+                kv[:, :, -self.qk_rope_head_dim :],
+                o_tensor.view(-1, self.tp_q_head_num_, self.kv_lora_rank),
+                infer_state.b_start_loc,
+                infer_state.b_seq_len,
+                infer_state.max_len_in_batch,
+                self.softmax_scale,
+            )
+
+        return o_tensor
+
     def _token_gqa_decode_attention_flashdecoding(
         self, q, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):
@@ -279,6 +392,29 @@ def _token_gqa_decode_attention_flashdecoding(
                 alloc_tensor_func=self.alloc_tensor,
             )
 
+    def _token_gqa_decode_attention_flashdecoding_fp8(
+        self, q, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
+    ):
+        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
+        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
+
+        kv = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, :-2].view(torch.float8_e4m3fn)
+        kv_scale = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, -2:].view(torch.bfloat16)
+        return gqa_token_decode_attention_flash_decoding_fp8(
+            q_nope,
+            q_rope,
+            kv[:, :, : -self.qk_rope_head_dim],
+            kv[:, :, -self.qk_rope_head_dim :],
+            kv_scale,
+            infer_state,
+            self.tp_q_head_num_,
+            self.kv_lora_rank,
+            self.qk_rope_head_dim,
+            self.qk_nope_head_dim,
+            self.softmax_scale,
+            alloc_tensor_func=self.alloc_tensor,
+        )
+
     def _splitfuse_attention_kernel(
         self, q, infer_state: SplitFuseInferStateInfo, layer_weight, out=None
     ) -> torch.Tensor:
@@ -321,6 +457,20 @@ def _copy_kv_to_mem_cache_normal(self, buffer, mem_index, mem_manager):
         )
         return
 
+    def _copy_kv_to_mem_cache_fp8(self, buffer, mem_index, mem_manager):
+        quant_method = vLLMFP8w8a8QuantizationMethod()
+        quant, scale = quant_method.quantize_scaled_mm_fp8(buffer.reshape(-1, buffer.shape[-1]))
+        destindex_copy_kv(
+            quant.T.unsqueeze(1)[:, :, : self.kv_lora_rank].view(torch.uint8),
+            quant.T.unsqueeze(1)[:, :, self.kv_lora_rank :].view(torch.uint8),
+            mem_index,
+            mem_manager.kv_buffer[self.layer_num_][:, :, : self.kv_lora_rank],
+            mem_manager.kv_buffer[self.layer_num_][:, :, self.kv_lora_rank : -2],
+            mem_manager.kv_buffer[self.layer_num_][:, :, -2:],
+            scale.to(buffer.dtype).view(torch.uint8),
+        )
+        return
+
     def _ffn_dp(
         self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
     ) -> torch.Tensor:

lightllm/models/deepseek2/model.py

@@ -7,6 +7,7 @@
 
 from lightllm.models.llama.model import LlamaTpPartModel
 from lightllm.common.deepseek2_mem_manager import Deepseek2MemoryManager
+from lightllm.common.deepseek2_fp8kv_mem_manager import Deepseek2FP8KVMemoryManager
 from lightllm.utils.log_utils import init_logger
 
 
@@ -48,7 +49,10 @@ def _verify_params(self):
         return super()._verify_params()
 
     def _init_mem_manager(self):
-        self.mem_manager = Deepseek2MemoryManager(
+        manager_class = Deepseek2MemoryManager
+        if "triton_fp8kv" in self.mode:
+            manager_class = Deepseek2FP8KVMemoryManager
+        self.mem_manager = manager_class(
             self.max_total_token_num,
             dtype=self.data_type,
             head_num=1,