opt: optimatize fp8kv performance

Author: niushengxiao

lightllm/common/basemodel/triton_kernel/q_per_head_fp8_quant.py

@@ -8,12 +8,10 @@
 def _per_head_max_reduce_kernel(
     Q,
     Scales,
-    BatchIds,
     StartLoc,
     stride_q_t,
     stride_q_h,
     stride_scales_b,
-    SET_BATCH_IDS: tl.constexpr,
     FP8_MAX: tl.constexpr,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
@@ -32,8 +30,6 @@ def _per_head_max_reduce_kernel(
         mask = (t_idx[:, None] < end_loc) & (q_range[None, :] < stride_q_h)
         q_vals = tl.load(q_ptrs, mask=mask, other=0.0)
         max_val = tl.maximum(tl.max(q_vals.abs()), max_val)
-        if SET_BATCH_IDS:
-            tl.store(BatchIds + t_idx, b_id, mask=t_idx < end_loc)
 
     scale = tl.where(max_val > 0, max_val / FP8_MAX, 1.0)
     scale_ptr = Scales + b_id * stride_scales_b + h_id
@@ -73,29 +69,29 @@ def _apply_quantization_kernel(
 
 
 @torch.no_grad()
-def q_per_head_fp8_quant(q, seq_lens, b1_start_loc):
+def q_per_head_fp8_quant(q, seq_lens, b1_start_loc, scale_out=None, batch_ids=None):
     T, H, D = q.shape
     B = seq_lens.shape[0]
-    device = q.device
-
-    q_out = torch.empty_like(q, dtype=torch.float8_e4m3fn)
-    scales = torch.empty((B, H), dtype=torch.float32, device=device)
-    batch_ids = torch.zeros((T,), dtype=torch.int32, device=device)
 
     BLOCK_D = triton.next_power_of_2(D)
     BLOCK_T = 256
     num_warps = 4
     num_stages = 2
+
+    q_out = torch.empty_like(q, dtype=torch.float8_e4m3fn)
+    if scale_out is None:
+        scale_out = torch.empty((B, H), dtype=torch.float32, device=q.device)
+    if batch_ids is None:
+        batch_ids = torch.repeat_interleave(torch.arange(B, device=q.device), seq_lens)
+
     _per_head_max_reduce_kernel[(B, H)](
         q,
-        scales,
-        batch_ids,
+        scale_out,
         b1_start_loc,
         q.stride(0),
         q.stride(1),
-        scales.stride(0),
+        scale_out.stride(0),
         FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
-        SET_BATCH_IDS=B > 1,
         BLOCK_T=BLOCK_T,
         BLOCK_D=BLOCK_D,
         num_warps=num_warps,
@@ -106,19 +102,19 @@ def q_per_head_fp8_quant(q, seq_lens, b1_start_loc):
         q,
         q_out,
         batch_ids,
-        scales,
+        scale_out,
         q.stride(0),
         q.stride(1),
         q_out.stride(0),
         q_out.stride(1),
-        scales.stride(0),
+        scale_out.stride(0),
         FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
         FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
         BLOCK_D=BLOCK_D,
         num_warps=num_warps,
         num_stages=num_stages,
     )
-    return q_out, scales
+    return q_out, scale_out
 
 
 def ref_q_per_head_fp8_quant(q, seq_lens):

lightllm/models/llama/flashattention_infer_struct.py

@@ -28,8 +28,16 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
         if self.is_prefill:
             self.cu_seqlens_q = self.b1_cu_q_seq_len.int()
             self.cu_seqlens_k = self.b1_cu_kv_seq_len.int()
-            self.page_table = torch.empty((self.batch_size, self.max_seq_len), dtype=torch.int32).to(input_ids.device)
+            self.page_table = torch.empty(
+                (self.batch_size, self.max_seq_len), dtype=torch.int32, device=input_ids.device
+            )
             self.page_table.copy_(model.req_manager.req_to_token_indexs[self.b_req_idx, : self.max_seq_len])
+            if "calibration_fp8kv" in model.mode:
+                device = input_ids.device
+                self.q_scale = torch.empty(
+                    (self.batch_size, self.mem_manager.head_num), dtype=torch.float32, device=device
+                )
+                self.batch_ids = torch.repeat_interleave(torch.arange(self.batch_size, device=device), self.b_q_seq_len)
         else:
             # Meta information of flashattention for decoding
             self.cu_seqlens_q = self.b1_cu_q_seq_len.int()
@@ -43,12 +51,38 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                     : self.batch_size * model.graph_max_len_in_batch
                 ].reshape(self.batch_size, model.graph_max_len_in_batch)
             else:
-                self.page_table = torch.empty((self.batch_size, self.max_len_in_batch), dtype=torch.int32).to(
-                    input_ids.device
+                self.page_table = torch.empty(
+                    (self.batch_size, self.max_len_in_batch), dtype=torch.int32, device=input_ids.device
                 )
 
             self.page_table[:, :max_seq_len_k].copy_(
                 model.req_manager.req_to_token_indexs[self.b_req_idx, :max_seq_len_k]
             )
             self.page_table[:, max_seq_len_k:].fill_(0)
+
+        if "calibration_fp8kv" in model.mode:
+            offline_scales = self.mem_manager.offline_fp8_quant_manager.scales
+            head_num = self.mem_manager.head_num
+            self.k_descale = (
+                offline_scales[:, :head_num]
+                .view(-1, 1, head_num)
+                .expand(offline_scales.shape[0], self.batch_size, head_num)
+                if offline_scales is not None
+                else torch.ones(
+                    (self.mem_manager.layer_num, self.batch_size, head_num),
+                    dtype=torch.float32,
+                    device=input_ids.device,
+                )
+            )
+            self.v_descale = (
+                offline_scales[:, head_num:]
+                .view(-1, 1, head_num)
+                .expand(offline_scales.shape[0], self.batch_size, head_num)
+                if offline_scales is not None
+                else torch.ones(
+                    (self.mem_manager.layer_num, self.batch_size, head_num),
+                    dtype=torch.float32,
+                    device=input_ids.device,
+                )
+            )
         return

lightllm/models/llama/layer_infer/transformer_layer_infer.py

@@ -132,7 +132,7 @@ def _bind_attention(self):
                     LlamaTransformerLayerInfer._token_decode_attention_flashinfer_fp8, self
                 )
             else:
-                raise Exception("fp8 kvcache only support fa3 and flashinfer backend")
+                raise Exception("calibration fp8 kvcache only support fa3 and flashinfer backend")
         elif "triton_flashdecoding" in self.mode:
             self._token_attention_kernel = partial(
                 LlamaTransformerLayerInfer._token_decode_attention_flashdecoding, self
@@ -333,6 +333,13 @@ def _context_attention_flashattention(self, q, kv, infer_state: FlashAttentionSt
     def _context_attention_flashattention_fp8(
         self, q, kv, infer_state: FlashAttentionStateInfo, layer_weight, out=None
     ):
+        q, q_scale = q_per_head_fp8_quant(
+            q.view(q.shape[0], self.tp_k_head_num_, -1),
+            infer_state.b_seq_len,
+            infer_state.cu_seqlens_q,
+            infer_state.q_scale,
+            infer_state.batch_ids,
+        )
         cache_k = (
             (infer_state.mem_manager.kv_buffer[self.layer_num_][:, : self.tp_k_head_num_, :])
             .reshape(-1, 1, self.tp_k_head_num_, self.head_dim_)
@@ -347,43 +354,21 @@ def _context_attention_flashattention_fp8(
             .reshape(-1, 1, self.tp_v_head_num_, self.head_dim_)
             .view(torch.float8_e4m3fn)
         )
-        q, q_scale = q_per_head_fp8_quant(
-            q.view(q.shape[0], self.tp_k_head_num_, -1),
-            infer_state.b_seq_len,
-            infer_state.cu_seqlens_q,
-        )
-        q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
-        q_descale = q_scale
-        ones_scales = torch.ones((infer_state.batch_size, self.tp_k_head_num_), device=q.device, dtype=torch.float32)
-        offline_scales = infer_state.mem_manager.offline_fp8_quant_manager.scales
-        k_descale = (
-            offline_scales[self.layer_num_][: self.tp_k_head_num_].expand(infer_state.batch_size, self.tp_k_head_num_)
-            if offline_scales is not None
-            else ones_scales
-        )
-        v_descale = (
-            offline_scales[self.layer_num_][self.tp_k_head_num_ :].expand(infer_state.batch_size, self.tp_k_head_num_)
-            if offline_scales is not None
-            else ones_scales
-        )
-        Lq = q.shape[-1]
-        sm_scale = 1.0 / (Lq ** 0.5)
         o = flash_attn_with_kvcache(
-            q=q,
+            q=q.view(-1, self.tp_q_head_num_, self.head_dim_),
             k_cache=cache_k,
             v_cache=cache_v,
             page_table=infer_state.page_table,
             cache_seqlens=infer_state.b_seq_len,
             cu_seqlens_q=infer_state.cu_seqlens_q,
             cu_seqlens_k_new=infer_state.cu_seqlens_k,
             max_seqlen_q=infer_state.q_max_seq_len,
-            softmax_scale=sm_scale,
             causal=True,
             window_size=(-1, -1),
             softcap=0.0,
-            q_descale=q_descale,
-            k_descale=k_descale,
-            v_descale=v_descale,
+            q_descale=q_scale,
+            k_descale=infer_state.k_descale[self.layer_num_],
+            v_descale=infer_state.v_descale[self.layer_num_],
             return_softmax_lse=False,
         )
         return o
@@ -867,38 +852,21 @@ def _token_decode_attention_flashattention_fp8(
             .view(torch.float8_e4m3fn)
         )
         q, q_scale = scaled_fp8_quant(q.view(q.shape[0] * self.tp_k_head_num_, -1), use_per_token_if_dynamic=True)
-        q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
-        q_descale = q_scale.view(q.shape[0], self.tp_k_head_num_)
-        ones_scales = torch.ones((infer_state.batch_size, self.tp_k_head_num_), device=q.device, dtype=torch.float32)
-        offline_scales = infer_state.mem_manager.offline_fp8_quant_manager.scales
-        k_descale = (
-            offline_scales[self.layer_num_][: self.tp_k_head_num_].expand(infer_state.batch_size, self.tp_k_head_num_)
-            if offline_scales is not None
-            else ones_scales
-        )
-        v_descale = (
-            offline_scales[self.layer_num_][self.tp_k_head_num_ :].expand(infer_state.batch_size, self.tp_k_head_num_)
-            if offline_scales is not None
-            else ones_scales
-        )
-        Lq = q.shape[-1]
-        sm_scale = 1.0 / (Lq ** 0.5)
         o = flash_attn_with_kvcache(
-            q=q,
+            q=q.view(-1, self.tp_q_head_num_, self.head_dim_),
             k_cache=cache_k,
             v_cache=cache_v,
             page_table=infer_state.page_table,
             cache_seqlens=infer_state.b_seq_len,
             cu_seqlens_q=infer_state.cu_seqlens_q,
             cu_seqlens_k_new=infer_state.cu_seqlens_k,
             max_seqlen_q=1,
-            softmax_scale=sm_scale,
             causal=False,
             window_size=(-1, -1),
             softcap=0.0,
-            q_descale=q_descale,
-            k_descale=k_descale,
-            v_descale=v_descale,
+            q_descale=q_scale.view(infer_state.batch_size, self.tp_k_head_num_),
+            k_descale=infer_state.k_descale[self.layer_num_],
+            v_descale=infer_state.v_descale[self.layer_num_],
             return_softmax_lse=False,
         )
         return o