feat: add deepseekv2_bf16kv and deepseekv2_fp8kv modes

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

@@ -3,6 +3,9 @@
 from lightllm.common.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager
 from lightllm.common.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager
 from lightllm.utils.log_utils import init_logger
+from lightllm.common.deepseek2_mem_manager import Deepseek2MemoryManager
+from lightllm.common.deepseek2_fp8kv_mem_manager import Deepseek2FP8KVMemoryManager
+
 
 logger = init_logger(__name__)
 
@@ -18,6 +21,12 @@ 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 "deepseek2_bf16kv" in mode:
+        memory_manager_class = Deepseek2MemoryManager
+        logger.info("Model kv cache using mode deepseek2 bf16kv")
+    elif "deepseek2_fp8kv" in mode:
+        memory_manager_class = Deepseek2FP8KVMemoryManager
+        logger.info("Model kv cache using mode deepseek2 fp8kv")
     else:
         memory_manager_class = MemoryManager
         logger.info("Model kv cache using mode normal")

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -22,6 +22,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):
@@ -68,6 +69,10 @@ def __init__(self, layer_num, tp_rank, world_size, network_config, mode=[]):
         return
 
     def _bind_attention(self):
+        if "deepseek2_bf16kv" in self.mode:
+            self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
+        else:
+            self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_fp8, self)
         if self.enable_cc_method:
             self._context_attention_kernel = partial(
                 Deepseek2TransformerLayerInfer._context_attention_kernel_with_CC, self
@@ -79,7 +84,6 @@ def _bind_attention(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)
         if self.is_moe:
             if self.enable_dp:
                 if os.environ.get("MOE_MODE", "TP") == "TP":
@@ -135,7 +139,15 @@ def _get_o(
 
     def _decompress_kv(self, kv, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight):
         if infer_state.use_dynamic_prompt_cache:
-            kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+            if "deepseek2_bf16kv" in self.mode:
+                kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+                kv_scale = None
+                k_scale = None
+            else:
+                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)
+
             compressed_kv = self.alloc_tensor(
                 [infer_state.total_token_num, 1, layer_weight.kv_lora_rank], dtype=kv.dtype
             )
@@ -147,7 +159,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
@@ -264,12 +281,18 @@ def _token_gqa_decode_attention_flashdecoding(
             )
             return o_tensor
         else:
-            kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+            if "deepseek2_bf16kv" in self.mode:
+                kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+                kv_scale = None
+            else:
+                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(
                 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,
@@ -321,6 +344,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,7 +7,9 @@
 
 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
+from lightllm.common.mem_utils import select_mem_manager_class
 
 
 logger = init_logger(__name__)
@@ -48,14 +50,17 @@ def _verify_params(self):
         return super()._verify_params()
 
     def _init_mem_manager(self):
-        self.mem_manager = Deepseek2MemoryManager(
+        self.mem_manager = select_mem_manager_class(self.mode)(
             self.max_total_token_num,
             dtype=self.data_type,
             head_num=1,
             head_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"],
             layer_num=self.config["num_hidden_layers"],
             mem_fraction=self.mem_fraction,
         )
+        assert isinstance(self.mem_manager, Deepseek2MemoryManager) or isinstance(
+            self.mem_manager, Deepseek2FP8KVMemoryManager
+        )
         return
 
     def _init_weights(self):

lightllm/models/deepseek2/triton_kernel/destindex_copy_kv.py

@@ -12,9 +12,11 @@ def _is_power_of_two(n):
 def _fwd_kernel_destindex_copy_kv(
     KV_nope,
     KV_rope,
+    KV_scale,
     Dest_loc,
     O_nope,
     O_rope,
+    O_scale,
     stride_kv_nope_bs,
     stride_kv_nope_h,
     stride_kv_nope_d,
@@ -29,6 +31,7 @@ def _fwd_kernel_destindex_copy_kv(
     stride_o_rope_d,
     kv_nope_head_num,
     kv_rope_head_num,
+    HAS_SCALE: tl.constexpr,
     BLOCK_DMODEL_NOPE: tl.constexpr,
     BLOCK_DMODEL_ROPE: tl.constexpr,
 ):
@@ -47,13 +50,20 @@ def _fwd_kernel_destindex_copy_kv(
     kv_nope = tl.load(kv_nope_ptrs)
     kv_rope = tl.load(kv_rope_ptrs)
 
+    if HAS_SCALE:
+        offs_d_scale = tl.arange(0, 2)
+        o_scale_ptrs = O_scale + dest_index * stride_o_rope_bs + stride_o_rope_d * offs_d_scale[None, :]
+        kv_scale_ptrs = KV_scale + cur_index * 2 + offs_d_scale[None, :]
+        kv_scale = tl.load(kv_scale_ptrs)
+        tl.store(o_scale_ptrs, kv_scale)
+
     tl.store(o_nope_ptrs, kv_nope)
     tl.store(o_rope_ptrs, kv_rope)
     return
 
 
 @torch.no_grad()
-def destindex_copy_kv(KV_nope, KV_rope, DestLoc, O_nope, O_rope):
+def destindex_copy_kv(KV_nope, KV_rope, DestLoc, O_nope, O_rope, O_scale=None, KV_scale=None):
     seq_len = DestLoc.shape[0]
     kv_nope_head_num = KV_nope.shape[1]
     kv_rope_head_num = KV_rope.shape[1]
@@ -71,9 +81,11 @@ def destindex_copy_kv(KV_nope, KV_rope, DestLoc, O_nope, O_rope):
     _fwd_kernel_destindex_copy_kv[grid](
         KV_nope,
         KV_rope,
+        KV_scale,
         DestLoc,
         O_nope,
         O_rope,
+        O_scale,
         KV_nope.stride(0),
         KV_nope.stride(1),
         KV_nope.stride(2),
@@ -88,6 +100,7 @@ def destindex_copy_kv(KV_nope, KV_rope, DestLoc, O_nope, O_rope):
         O_rope.stride(2),
         kv_nope_head_num,
         kv_rope_head_num,
+        HAS_SCALE=KV_scale is not None,
         BLOCK_DMODEL_NOPE=kv_nope_head_dim,
         BLOCK_DMODEL_ROPE=kv_rope_head_dim,
         num_warps=num_warps,

lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py

@@ -16,6 +16,7 @@ def gqa_token_decode_attention_flash_decoding(
     q_rope,
     kv_nope,
     kv_rope,
+    kv_scale,
     infer_state,
     q_head_num,
     kv_lora_rank,
@@ -63,6 +64,7 @@ def gqa_token_decode_attention_flash_decoding(
         q_rope.view(calcu_shape2),
         kv_nope,
         kv_rope,
+        kv_scale,
         infer_state.req_manager.req_to_token_indexs,
         infer_state.b_req_idx,
         infer_state.b_seq_len,
@@ -111,6 +113,7 @@ def gqa_token_decode_attention_flash_decoding(
         q_rope.view(calcu_shape2),
         kv_nope,
         kv_rope,
+        kv_scale,
         infer_state.req_manager.req_to_token_indexs,
         infer_state.b_req_idx,
         infer_state.b_seq_len,

lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py

@@ -10,6 +10,7 @@ def _fwd_kernel_flash_decode_stage1_padding(
     Q_rope,
     KV_nope,
     KV_rope,
+    KV_scale,
     sm_scale,
     Req_to_tokens,
     B_req_idx,
@@ -35,11 +36,13 @@ def _fwd_kernel_flash_decode_stage1_padding(
     stride_mid_od,
     stride_mid_o_eh,
     stride_mid_o_es,
+    stride_kv_scaled_bs,
     block_size_ptr,
     num_sm,
     head_group_num,
     head_num,
     batch_size,
+    HAS_SCALE: tl.constexpr,
     Q_HEAD_NUM: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_ROPE_DMODEL: tl.constexpr,
@@ -108,9 +111,14 @@ def _fwd_kernel_flash_decode_stage1_padding(
                 )
                 off_kv = kv_loc[None, :] * stride_kv_bs + offs_d[:, None]
                 kv = tl.load(KV_nope + off_kv, mask=seq_n_mask[None, :], other=0.0)
-                att_value = tl.dot(q, kv)
                 off_rope_kv = kv_loc[None, :] * stride_kv_rope_bs + offs_rope_d[:, None]
                 rope_kv = tl.load(KV_rope + off_rope_kv, mask=seq_n_mask[None, :], other=0.0)
+                if HAS_SCALE:
+                    off_kv_scale = kv_loc[None, :] * stride_kv_scaled_bs
+                    kv_scale = tl.load(KV_scale + off_kv_scale, mask=seq_n_mask[None, :], other=0.0)
+                    kv = (kv * kv_scale).to(kv_scale.dtype)
+                    rope_kv = (rope_kv * kv_scale).to(kv_scale.dtype)
+                att_value = tl.dot(q, kv)
                 att_value += tl.dot(q_rope, rope_kv)
 
                 att_value *= sm_scale
@@ -167,6 +175,7 @@ def flash_decode_stage1(
     q_rope,
     kv_nope,
     kv_rope,
+    kv_scale,
     Req_to_tokens,
     B_req_idx,
     B_Seqlen,
@@ -201,6 +210,7 @@ def flash_decode_stage1(
         q_rope,
         kv_nope,
         kv_rope,
+        kv_scale,
         softmax_scale,
         Req_to_tokens,
         B_req_idx,
@@ -214,11 +224,13 @@ def flash_decode_stage1(
         *kv_rope.stride(),
         *mid_out.stride(),
         *mid_out_logsumexp.stride(),
+        kv_scale.stride(0) if kv_scale is not None else 0,
         in_block_seq,
         num_sm=1,
         head_group_num=head_group_num,
         head_num=q_head_num,
         batch_size=batch_size,
+        HAS_SCALE=1 if kv_scale is not None else 0,
         Q_HEAD_NUM=Q_HEAD_NUM,
         BLOCK_DMODEL=q_nope_dim,
         BLOCK_ROPE_DMODEL=q_rope_dim,
@@ -243,6 +255,7 @@ def flash_decode_stage1(
         q_rope,
         kv_nope,
         kv_rope,
+        kv_scale,
         softmax_scale,
         Req_to_tokens,
         B_req_idx,
@@ -256,11 +269,13 @@ def flash_decode_stage1(
         *kv_rope.stride(),
         *mid_out.stride(),
         *mid_out_logsumexp.stride(),
+        kv_scale.stride(0) if kv_scale is not None else 0,
         in_block_seq,
         num_sm=num_sm,
         head_group_num=head_group_num,
         head_num=q_head_num,
         batch_size=batch_size,
+        HAS_SCALE=1 if kv_scale is not None else 0,
         Q_HEAD_NUM=Q_HEAD_NUM,
         BLOCK_DMODEL=q_nope_dim,
         BLOCK_ROPE_DMODEL=q_rope_dim,

lightllm/models/deepseek2/triton_kernel/sample_kv.py

@@ -10,16 +10,20 @@
 @triton.jit
 def _sample_kv_kernel(
     KV_input,
+    KV_scale,
     KV_nope,
     KV_rope,
+    K_scale,
     B_start_loc,
     B_Seqlen,
     Req_to_tokens,
     B_req_idx,
     stride_input_dim,
+    stride_scale_dim,
     stride_nope_dim,
     stride_rope_dim,
     stride_req_to_tokens_b,
+    HAS_SCALE: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_ROPE_DMODEL: tl.constexpr,
@@ -52,6 +56,11 @@ def _sample_kv_kernel(
     rope_ptrs = KV_rope + off_rope
     tl.store(nope_ptrs, kv_nope, mask=offs_m[:, None] < block_end_loc)
     tl.store(rope_ptrs, kv_rope, mask=offs_m[:, None] < block_end_loc)
+    if HAS_SCALE:
+        kv_scale = tl.load(KV_scale + kv_loc * stride_scale_dim, mask=offs_m < block_end_loc)
+        off_k_scale = cur_batch_start_loc + offs_m
+        k_scale_ptrs = K_scale + off_k_scale
+        tl.store(k_scale_ptrs, kv_scale, mask=offs_m < block_end_loc)
     return
 
 
@@ -63,6 +72,8 @@ def sample_kv(
     b_req_idx,
     b_seq_len,
     req_to_token_indexs,
+    kv_scale=None,
+    k_scale=None,
 ):
     BLOCK = 128 if not TESLA else 64
 
@@ -85,16 +96,20 @@ def sample_kv(
     b_start_loc = torch.cat([torch.zeros([1], device=b_seq_len.device, dtype=b_seq_len.dtype), b_seq_len[1:].cumsum(0)])
     _sample_kv_kernel[grid](
         kv_input,
+        kv_scale,
         kv_nope,
         kv_rope,
+        k_scale,
         b_start_loc,
         b_seq_len,
         req_to_token_indexs,
         b_req_idx,
         kv_input.stride(0),
+        kv_scale.stride(0) if kv_scale is not None else 0,
         kv_nope.stride(0),
         kv_rope.stride(0),
         req_to_token_indexs.stride(0),
+        HAS_SCALE=kv_scale is not None,
         BLOCK_M=BLOCK,
         BLOCK_DMODEL=nope_dim,
         BLOCK_ROPE_DMODEL=rope_dim,