add Fa3 (#858)

Co-authored-by: baishihao <[email protected]>

Author: shihaobai

lightllm/common/basemodel/basemodel.py

@@ -76,6 +76,7 @@ def __init__(self, kvargs):
         self._verify_must()
         self._verify_params()
         self._init_quant()
+        self._init_inferstate_cls()
 
         # 更连续的显存分配可以有更好的性能
         if self.max_total_token_num is None:
@@ -107,6 +108,9 @@ def _init_config(self):
             self.config["vocab_size"] = self.finetune_config.vocab_size
         return
 
+    def _init_inferstate_cls(self):
+        pass
+
     @final
     def _verify_must(self):
         assert self.config["num_attention_heads"] % self.tp_world_size_ == 0

lightllm/common/basemodel/infer_struct.py

@@ -49,6 +49,6 @@ def copy_for_cuda_graph(self, new_infer_state):
         for attr_name, attr_value in vars(new_infer_state).items():
             if isinstance(attr_value, torch.Tensor):
                 attr_ = getattr(self, attr_name, None)
-                if attr_ is not None:
+                if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr():
                     attr_.copy_(attr_value, non_blocking=True)
         return

lightllm/models/deepseek2/flashattention_infer_struct.py

@@ -0,0 +1,55 @@
+import os
+import torch
+import numpy as np
+import torch.distributed as dist
+from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
+from lightllm.utils.dist_utils import get_current_device_id
+
+
+class Deepseek2FlashAttentionStateInfo(Deepseek2InferStateInfo):
+    _shared_page_table_buffer = None
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
+        if cls._shared_page_table_buffer is None:
+            cls._shared_page_table_buffer = [
+                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
+                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
+            ]
+        return cls._shared_page_table_buffer
+
+    def init_some_extra_state(self, model, input_ids: torch.Tensor):
+        super().init_some_extra_state(model, input_ids)
+        if self.is_prefill:
+            self.cu_seqlens_q = torch.nn.functional.pad(
+                torch.cumsum(self.b_seq_len - self.b_ready_cache_len, dim=0, dtype=torch.int32), (1, 0)
+            )
+            self.cu_seqlens_k = torch.cat([self.b_start_loc, self.b_start_loc[-1:] + self.b_seq_len[-1:]], dim=0)
+            self.page_table = torch.empty((self.batch_size, self.max_seq_len), dtype=torch.int32).to(input_ids.device)
+            self.page_table.copy_(model.req_manager.req_to_token_indexs[self.b_req_idx, : self.max_seq_len])
+        else:
+            # Meta information of flashattention for decoding
+            self.cu_seqlens_q = torch.arange(0, self.batch_size + 1, dtype=torch.int32, device=input_ids.device)
+            self.cu_seqlens_k = torch.cat([self.b_start_loc, self.b_start_loc[-1:] + self.b_seq_len[-1:]], dim=0)
+            b_seq_len_numpy = self.b_seq_len.cpu().numpy()
+            max_seq_len_k = b_seq_len_numpy.max()
+            if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
+                page_buffer = Deepseek2FlashAttentionStateInfo.get_page_table_buffer(
+                    model.graph_max_batch_size, model.graph_max_len_in_batch
+                )
+                self.page_table = page_buffer[self.microbatch_index][
+                    : 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[:, :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)
+        return

lightllm/models/deepseek2/infer_struct.py

@@ -12,13 +12,11 @@ def __init__(self):
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
         super().init_some_extra_state(model, input_ids)
-        # 只有 decode 阶段使用 ppl 的优化算子才会有这个管理变量
         if not self.is_prefill:
             self.kv_starts = torch.cat([self.b_start_loc, self.b_start_loc[-1:] + self.b_seq_len[-1:]], dim=0)
             self.total_token_num_tensor = torch.sum(self.b_seq_len)
 
         if self.is_prefill:
             self.b_kv_start_loc = self.b_seq_len.cumsum(dim=0) - self.b_seq_len
             self.max_value_in_b_seq_len = self.b_seq_len.max().item()
-
         return

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -29,6 +29,14 @@
 from lightllm.distributed.communication_op import all_gather, all_gather_into_tensor, all_reduce, reduce_scatter_tensor
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.dist_utils import get_global_world_size
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+try:
+    from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
+except:
+    logger.warning("sgl_kernel is not installed, or the installed version does not support fa3!")
 
 
 class Deepseek2TransformerLayerInfer(LlamaTransformerLayerInfer):
@@ -93,7 +101,11 @@ def _bind_attention(self):
             )
         else:
             self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-            if get_env_start_args().enable_flashinfer_decode:
+            if get_env_start_args().enable_fa3:
+                self._token_attention_kernel = partial(
+                    Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashattention, self
+                )
+            elif get_env_start_args().enable_flashinfer_decode:
                 self._token_attention_kernel = partial(
                     Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashinfer, self
                 )
@@ -112,7 +124,11 @@ def _bind_attention(self):
                         Deepseek2TransformerLayerInfer._context_attention_kernel_with_CC_fp8, self
                     )
             else:
-                if get_env_start_args().enable_flashinfer_prefill:
+                if get_env_start_args().enable_fa3:
+                    self._context_attention_kernel = partial(
+                        Deepseek2TransformerLayerInfer._context_attention_flashattention_kernel_with_CC, self
+                    )
+                elif get_env_start_args().enable_flashinfer_prefill:
                     self._context_attention_kernel = partial(
                         Deepseek2TransformerLayerInfer._context_attention_flashinfer_kernel_with_CC, self
                     )
@@ -278,6 +294,30 @@ def _decompress_kv(
         k_nope, v = torch.split(kv_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
         return k_nope, k_rope, v
 
+    def _context_attention_flashattention_kernel_with_CC(
+        self,
+        q: torch.Tensor,
+        kv,
+        infer_state: Deepseek2FlashInferStateInfo,
+        layer_weight: Deepseek2TransformerLayerWeight,
+        out=None,
+    ) -> torch.Tensor:
+        k_nope, k_rope, v = self._decompress_kv(kv, infer_state, layer_weight, False)
+        k = torch.cat([k_nope, torch.repeat_interleave(k_rope, self.tp_q_head_num_, dim=-2)], dim=-1)
+        o_tensor = flash_attn_varlen_func(
+            q=q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim),
+            k=k.view(-1, self.tp_k_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim),
+            v=v.view(-1, self.tp_v_head_num_, self.v_head_dim),
+            cu_seqlens_q=infer_state.cu_seqlens_q,
+            cu_seqlens_k=infer_state.cu_seqlens_k,
+            max_seqlen_q=infer_state.q_max_seq_len,
+            max_seqlen_k=infer_state.max_seq_len,
+            softmax_scale=self.softmax_scale,
+            causal=True,
+            return_softmax_lse=False,
+        )
+        return o_tensor
+
     def _context_attention_flashinfer_kernel_with_CC(
         self,
         q: torch.Tensor,
@@ -450,6 +490,35 @@ def _context_attention_kernel_origin_fp8(
 
         return o_tensor
 
+    def _token_gqa_decode_attention_flashattention(
+        self, q, infer_state: Deepseek2FlashInferStateInfo, 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_]
+        k_rope = kv[:, :, -self.qk_rope_head_dim :].reshape(-1, 1, 1, self.qk_rope_head_dim)
+        kv_nope = kv[:, :, : -self.qk_rope_head_dim].reshape(-1, 1, 1, self.kv_lora_rank)
+        k_descale, v_descale = None, None
+        o_tensor = flash_attn_with_kvcache(
+            q=q_rope,
+            k_cache=k_rope,
+            v_cache=kv_nope,
+            qv=q_nope,
+            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=self.softmax_scale,
+            causal=True,
+            window_size=(-1, -1),
+            softcap=0.0,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            return_softmax_lse=False,
+        )
+        return o_tensor
+
     def _token_gqa_decode_attention_flashinfer(
         self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):

lightllm/models/deepseek2/model.py

@@ -4,6 +4,7 @@
 from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
 from lightllm.models.deepseek2.flashinfer_struct import Deepseek2FlashInferStateInfo
+from lightllm.models.deepseek2.flashattention_infer_struct import Deepseek2FlashAttentionStateInfo
 from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights
 
 from lightllm.models.llama.model import LlamaTpPartModel
@@ -62,11 +63,17 @@ def __init__(self, kvargs):
         self.enable_flashinfer = (
             get_env_start_args().enable_flashinfer_prefill or get_env_start_args().enable_flashinfer_decode
         )
-        if self.enable_flashinfer:
-            self.infer_state_class = Deepseek2FlashInferStateInfo
         super().__init__(kvargs)
         return
 
+    def _init_inferstate_cls(self):
+        if get_env_start_args().enable_fa3:
+            self.infer_state_class = Deepseek2FlashAttentionStateInfo
+        elif self.enable_flashinfer:
+            self.infer_state_class = Deepseek2FlashInferStateInfo
+        if self.enable_flashinfer:
+            self.flashinfer_extra_state = FlashInferStateExtraInfo(self)
+
     def _init_some_value(self):
         super()._init_some_value()
         self.tp_k_head_num_ = 1
@@ -77,8 +84,6 @@ def _init_some_value(self):
         self.q_lora_rank = self.config["q_lora_rank"]
         self.kv_lora_rank = self.config["kv_lora_rank"]
         self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim
-        if self.enable_flashinfer:
-            self.flashinfer_extra_state = FlashInferStateExtraInfo(self)
 
     def _init_custom(self):
         self._init_to_get_yarn_rotary()

lightllm/models/llama/flashattention_infer_struct.py

@@ -0,0 +1,57 @@
+import os
+import torch
+import numpy as np
+import torch.distributed as dist
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
+
+
+class FlashAttentionStateInfo(LlamaInferStateInfo):
+    _shared_page_table_buffer = None
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
+        if cls._shared_page_table_buffer is None:
+            cls._shared_page_table_buffer = [
+                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
+                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
+            ]
+        return cls._shared_page_table_buffer
+
+    def init_some_extra_state(self, model, input_ids: torch.Tensor):
+        super().init_some_extra_state(model, input_ids)
+        if self.is_prefill:
+            self.cu_seqlens_q = torch.nn.functional.pad(
+                torch.cumsum(self.b_seq_len - self.b_ready_cache_len, dim=0, dtype=torch.int32), (1, 0)
+            )
+            self.cu_seqlens_k = torch.cat([self.b_start_loc, self.b_start_loc[-1:] + self.b_seq_len[-1:]], dim=0)
+            self.page_table = torch.empty((self.batch_size, self.max_seq_len), dtype=torch.int32).to(input_ids.device)
+            self.page_table.copy_(model.req_manager.req_to_token_indexs[self.b_req_idx, : self.max_seq_len])
+        else:
+            # Meta information of flashattention for decoding
+            self.cu_seqlens_q = torch.arange(0, self.batch_size + 1, dtype=torch.int32, device=input_ids.device)
+            self.cu_seqlens_k = torch.cat([self.b_start_loc, self.b_start_loc[-1:] + self.b_seq_len[-1:]], dim=0)
+            b_seq_len_numpy = self.b_seq_len.cpu().numpy()
+            max_seq_len_k = b_seq_len_numpy.max()
+            if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
+                page_buffer = FlashAttentionStateInfo.get_page_table_buffer(
+                    model.graph_max_batch_size, model.graph_max_len_in_batch
+                )
+                self.page_table = page_buffer[self.microbatch_index][
+                    : 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[:, :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)
+        return

lightllm/models/llama/infer_struct.py

@@ -15,6 +15,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
             b_seq_len_numpy = self.b_seq_len.cpu().numpy()
             self.max_seq_len = b_seq_len_numpy.max()
             b_ready_cache_len_numpy = self.b_ready_cache_len.cpu().numpy()
+            self.q_max_seq_len = (b_seq_len_numpy - b_ready_cache_len_numpy).max()
             position_ids = torch.from_numpy(
                 np.concatenate(
                     [np.arange(b_ready_cache_len_numpy[i], b_seq_len_numpy[i]) for i in range(len(b_seq_len_numpy))],