feat: add flashinfer prefilled operator in the attention module

niushengxiao · niushengxiao · commit edae8c1d6ca8 · 2025-02-25T15:57:54.000+08:00
diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py
@@ -36,6 +36,7 @@ class TpPartBaseModel:
     infer_state_class = InferStateInfo
 
     def __init__(self, kvargs):
+        self.infer_state = self.infer_state_class()
         self.run_mode = kvargs["run_mode"]
         self.tp_rank_ = kvargs["tp_rank"]
         self.world_size_ = kvargs["world_size"]
@@ -330,7 +331,9 @@ def _decode(
         b_seq_len,
         multimodal_params,
     ):
-        infer_state = self.infer_state_class()
+        infer_state = self.infer_state
+        if self.graph is None or self.graph.need_capture(batch_size):
+            infer_state = self.infer_state_class()
         infer_state.is_prefill = False
         infer_state.batch_size = batch_size
         infer_state.total_token_num = total_token_num
diff --git a/lightllm/models/deepseek2/infer_struct.py b/lightllm/models/deepseek2/infer_struct.py
@@ -3,15 +3,20 @@
 import numpy as np
 import torch.distributed as dist
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
-import flashinfer
 
 
 class Deepseek2InferStateInfo(LlamaInferStateInfo):
     def __init__(self):
         super().__init__()
         self.kv_starts = None
+        self.prefill_wrapper = None
+        self.decode_wrapper = None
         self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
+        self.enable_flashinfer_prefilled = os.getenv("ENABLE_FLASHINFER_PREFILLED", "False").upper() in [
+            "ON",
+            "TRUE",
+            "1",
+        ]
         self.enable_flashinfer_decode_mla = os.getenv("ENABLE_FLASHINFER_DECODE_MLA", "False").upper() in [
             "ON",
             "TRUE",
@@ -20,12 +25,24 @@ 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.enable_flashinfer_decode_mla:
-                self.workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(input_ids.device)
+                import flashinfer
+                from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
+
+                self.tp_q_head_num = (
+                    model.tp_q_head_num_ * model.world_size_ if self.enable_dp else model.tp_q_head_num_
+                )
+                self.kv_lora_rank = model.kv_lora_rank
+                self.qk_rope_head_dim = model.qk_rope_head_dim
+                self.qk_nope_head_dim = model.qk_nope_head_dim
+                self.softmax_scale = model.softmax_scale
+                self.q_data_type = model.data_type
+                self.kv_data_type = model.data_type
+
                 self.q_indptr = torch.arange(self.batch_size + 1, dtype=torch.int32).to(input_ids.device)
                 self.kv_indices = torch.empty(self.batch_size * model.max_seq_length, dtype=torch.int32).to(
                     input_ids.device
@@ -38,38 +55,63 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                     self.max_len_in_batch,
                     self.kv_indices,
                 )
-                self.wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
-                    self.workspace_buffer,
-                    backend="fa2",
-                    use_cuda_graph=True,
-                    qo_indptr=self.q_indptr,
-                    kv_indices=self.kv_indices,
-                    kv_indptr=self.kv_starts,
-                    kv_len_arr=self.b_seq_len,
+                if self.decode_wrapper is None:
+                    self.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
+                        model.workspace_buffer,
+                        use_cuda_graph=True,
+                        qo_indptr=self.q_indptr,
+                        kv_indices=self.kv_indices,
+                        kv_indptr=self.kv_starts,
+                        kv_len_arr=self.b_seq_len,
+                    )
+                    self.decode_wrapper.plan(
+                        self.q_indptr,
+                        self.kv_starts,
+                        self.kv_indices,
+                        self.b_seq_len,
+                        self.tp_q_head_num,
+                        self.kv_lora_rank,
+                        self.qk_rope_head_dim,
+                        1,
+                        False,  # causal
+                        self.softmax_scale,
+                        self.q_data_type,
+                        self.kv_data_type,
+                    )
+        else:
+            self.b_kv_start_loc = self.b_seq_len.cumsum(dim=0) - self.b_seq_len
+            if self.enable_flashinfer_prefilled:
+                import flashinfer
+
+                self.tp_q_head_num = (
+                    model.tp_q_head_num_ * model.world_size_ if self.enable_dp else model.tp_q_head_num_
                 )
-                self.head_num = model.tp_q_head_num_ * model.world_size_ if self.enable_dp else model.tp_q_head_num_
-                self.kv_lora_rank = model.kv_lora_rank
                 self.qk_rope_head_dim = model.qk_rope_head_dim
+                self.qk_nope_head_dim = model.qk_nope_head_dim
                 self.softmax_scale = model.softmax_scale
                 self.q_data_type = model.data_type
-                self.kv_data_type = model.data_type
-                self.wrapper.plan(
-                    self.q_indptr,
-                    self.kv_starts,
-                    self.kv_indices,
-                    self.b_seq_len,
-                    self.head_num,
-                    self.kv_lora_rank,
-                    self.qk_rope_head_dim,
-                    1,
-                    False,  # causal
-                    self.softmax_scale,
-                    self.q_data_type,
-                    self.kv_data_type,
-                )
 
-        if self.is_prefill:
-            self.b_kv_start_loc = self.b_seq_len.cumsum(dim=0) - self.b_seq_len
+                q_starts = torch.cat(
+                    [self.b_start_loc, self.b_start_loc[-1:] + (self.b_seq_len - self.b_ready_cache_len)[-1:]], dim=0
+                ).int()
+                kv_starts = torch.cat(
+                    [self.b_kv_start_loc, self.b_kv_start_loc[-1:] + self.b_seq_len[-1:]], dim=0
+                ).int()
+                if self.prefill_wrapper is None:
+                    self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper(
+                        model.workspace_buffer, "NHD"
+                    )
+                self.prefill_wrapper.plan(
+                    qo_indptr=q_starts,
+                    kv_indptr=kv_starts,
+                    num_qo_heads=self.tp_q_head_num,
+                    num_kv_heads=self.tp_q_head_num,
+                    head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim,
+                    head_dim_vo=self.qk_nope_head_dim,
+                    q_data_type=self.q_data_type,
+                    causal=True,
+                    sm_scale=self.softmax_scale,
+                )
 
         if self.enable_dp:
             rank = dist.get_rank()
@@ -89,19 +131,19 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
 
     def copy_for_cuda_graph(self, new_infer_state):
         super().copy_for_cuda_graph(new_infer_state)
-        if self.enable_flashinfer_decode_mla:
-            self.wrapper.plan(
-                self.q_indptr,
-                self.kv_starts,
-                self.kv_indices,
-                self.b_seq_len,
-                self.head_num,
-                self.kv_lora_rank,
-                self.qk_rope_head_dim,
+        if self.enable_flashinfer_decode_mla and not self.is_prefill:
+            self.decode_wrapper.plan(
+                new_infer_state.q_indptr,
+                new_infer_state.kv_starts,
+                new_infer_state.kv_indices,
+                new_infer_state.b_seq_len,
+                new_infer_state.tp_q_head_num,
+                new_infer_state.kv_lora_rank,
+                new_infer_state.qk_rope_head_dim,
                 1,
                 False,  # causal
-                self.softmax_scale,
-                self.q_data_type,
-                self.kv_data_type,
+                new_infer_state.softmax_scale,
+                new_infer_state.q_data_type,
+                new_infer_state.kv_data_type,
             )
         return
diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py
@@ -69,6 +69,11 @@ def __init__(self, layer_num, tp_rank, world_size, network_config, mode=[]):
         self.num_heads = network_config["num_attention_heads"]
         self.num_kv_heads = network_config["num_key_value_heads"]
         self.enable_opt_decoding_mha = os.getenv("ENABLE_OPT_DECODE_MHA", "False").upper() in ["ON", "TRUE", "1"]
+        self.enable_flashinfer_prefilled = os.getenv("ENABLE_FLASHINFER_PREFILLED", "False").upper() in [
+            "ON",
+            "TRUE",
+            "1",
+        ]
         self.enable_flashinfer_decode_mla = os.getenv("ENABLE_FLASHINFER_DECODE_MLA", "False").upper() in [
             "ON",
             "TRUE",
@@ -220,22 +225,28 @@ def _context_attention_kernel_with_CC(
         out=None,
     ) -> torch.Tensor:
         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_kv_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            self.softmax_scale,
+        o_tensor = (
+            self.alloc_tensor((q.shape[0], q.shape[1], self.qk_nope_head_dim), dtype=q.dtype) if out is None else out
         )
+        if self.enable_flashinfer_prefilled:
+            k = torch.cat([k_nope, torch.repeat_interleave(k_rope, self.tp_q_head_num_, dim=-2)], dim=-1)
+            infer_state.prefill_wrapper.run(q, k, v, out=o_tensor)
+        else:
+            q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
+            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_kv_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(
@@ -249,20 +260,24 @@ def _context_attention_kernel_with_CC_fp8(
         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(
-            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_kv_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            self.softmax_scale,
-        )
+        if self.enable_flashinfer_prefilled:
+            k = torch.cat([k_nope, torch.repeat_interleave(k_rope, self.tp_q_head_num_, dim=-2)], dim=-1)
+            infer_state.prefill_wrapper.run(q, k, v, out=o_tensor)
+        else:
+            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_kv_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_origin(
@@ -378,7 +393,7 @@ def _token_gqa_decode_attention_flashdecoding(
             )
             return o_tensor
         elif self.enable_flashinfer_decode_mla:
-            infer_state.wrapper.run(
+            infer_state.decode_wrapper.run(
                 q_nope,
                 q_rope,
                 kv[:, :, : -self.qk_rope_head_dim],
diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py
@@ -40,6 +40,7 @@ def _init_some_value(self):
         self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim
         self.tp_q_head_num_ = self.config["num_attention_heads"] // self.world_size_
         self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5)
+        self.workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(self.tp_rank_)
         if self.config["rope_scaling"] is not None:
             rope_scaling = self.config["rope_scaling"]
             mscale_all_dim = rope_scaling.get("mscale_all_dim", 0)
diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py
@@ -179,8 +179,7 @@ def _fwd_kernel_calcu_index_and_block_seq(
     req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX)
     b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX
     b_start_loc = torch.arange(Z).cuda().int() * N_CTX
-    b_start_loc[0] = 0
-    b_req_idx = torch.arange(Z).cuda().int()
+    b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda()
     kv_starts = torch.cat([b_start_loc, b_start_loc[-1:] + b_seq_len[-1:]], dim=0)
 
     o = torch.zeros((Z, H, D_HEAD), dtype=dtype, device="cuda")
