feat: add flashinfer prefilled operator in the attention module

niushengxiao · niushengxiao · commit 24d6fc414aa9 · 2025-02-21T14:15:46.000+08:00
diff --git a/lightllm/models/deepseek2/infer_struct.py b/lightllm/models/deepseek2/infer_struct.py
@@ -3,15 +3,19 @@
 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.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 +24,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,27 +54,23 @@ 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,
-                )
-                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.softmax_scale = model.softmax_scale
-                self.q_data_type = model.data_type
-                self.kv_data_type = model.data_type
+                if not self.wrapper:
+                    workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(input_ids.device)
+                    self.wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
+                        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,
+                    )
                 self.wrapper.plan(
                     self.q_indptr,
                     self.kv_starts,
                     self.kv_indices,
                     self.b_seq_len,
-                    self.head_num,
+                    self.tp_q_head_num,
                     self.kv_lora_rank,
                     self.qk_rope_head_dim,
                     1,
@@ -67,9 +79,39 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                     self.q_data_type,
                     self.kv_data_type,
                 )
-
-        if self.is_prefill:
+        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.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
+
+                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 not self.wrapper:
+                    workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0)
+                    self.wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper(workspace_buffer, "NHD")
+                self.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()
@@ -95,7 +137,7 @@ def copy_for_cuda_graph(self, new_infer_state):
                 self.kv_starts,
                 self.kv_indices,
                 self.b_seq_len,
-                self.head_num,
+                self.tp_q_head_num,
                 self.kv_lora_rank,
                 self.qk_rope_head_dim,
                 1,
diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py
@@ -68,6 +68,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",
@@ -223,22 +228,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.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(
@@ -252,20 +263,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.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(
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")
