feat: support page size variable for flashinfer

Author: niushengxiao

lightllm/models/llama/flashinfer_struct.py

@@ -3,16 +3,21 @@
 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.envs_utils import get_env_start_args, get_page_size
 from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
 
 
+def cdiv(a, b):
+    return (a + b - 1) // b
+
+
 class LlamaFlashInferStateInfo(LlamaInferStateInfo):
     def __init__(self):
         super().__init__()
         self.prefill_wrapper = None
         self.decode_wrapper = None
         self.flashinfer_extra_state = None
+        self.page_size = get_page_size()
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
         super().init_some_extra_state(model, input_ids)
@@ -22,29 +27,41 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
 
         if not self.is_prefill:
             if get_env_start_args().enable_flashinfer_decode:
-                self.kv_last_page_len_buffer = torch.full(
-                    (self.batch_size,), 1, dtype=torch.int32, device=input_ids.device
-                )
+                self.kv_last_page_len = torch.full((self.batch_size,), 1, dtype=torch.int32, device=input_ids.device)
+                length = cdiv(self.flashinfer_extra_state.max_seq_length, self.page_size)
                 if self.batch_size <= model.graph_max_batch_size:
                     self.kv_indices = self.flashinfer_extra_state.kv_indices_buffer[self.microbatch_index][
-                        : self.batch_size * self.flashinfer_extra_state.max_seq_length
+                        : self.batch_size * length
                     ]
                 else:
                     self.kv_indices = torch.empty(
-                        self.batch_size * self.flashinfer_extra_state.max_seq_length,
+                        self.batch_size * length,
                         dtype=torch.int32,
                         device=input_ids.device,
                     )
 
-                repack_kv_index(
-                    self.req_manager.req_to_token_indexs,
-                    self.b_req_idx,
-                    self.b_seq_len,
-                    self.b_start_loc,
-                    self.max_len_in_batch,
-                    self.kv_indices,
-                )
                 self.kv_starts = self.b1_cu_kv_seq_len.int()
+                if "page_size_variable" in model.mode:
+                    b_page_len = cdiv(self.b_seq_len, self.page_size)
+                    self.kv_starts[1:] = b_page_len.cumsum(0)
+                    self.kv_last_page_len = self.b_seq_len - (b_page_len - 1) * self.page_size
+                    repack_kv_index(
+                        self.req_manager.req_to_page_indexs,
+                        self.b_req_idx,
+                        b_page_len,
+                        self.kv_starts[:-1],
+                        cdiv(self.max_kv_seq_len, self.page_size),
+                        self.kv_indices,
+                    )
+                else:
+                    repack_kv_index(
+                        self.req_manager.req_to_token_indexs,
+                        self.b_req_idx,
+                        self.b_seq_len,
+                        self.b_start_loc,
+                        self.max_kv_seq_len,
+                        self.kv_indices,
+                    )
                 if self.decode_wrapper is None:
                     self.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
                         self.flashinfer_extra_state.workspace_buffer,
@@ -53,16 +70,16 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                         use_tensor_cores=True,
                         paged_kv_indptr_buffer=self.kv_starts,
                         paged_kv_indices_buffer=self.kv_indices,
-                        paged_kv_last_page_len_buffer=self.kv_last_page_len_buffer,
+                        paged_kv_last_page_len_buffer=self.kv_last_page_len,
                     )
                     self.decode_wrapper.plan(
                         self.kv_starts,
                         self.kv_indices,
-                        self.kv_last_page_len_buffer,
+                        self.kv_last_page_len,
                         self.flashinfer_extra_state.tp_q_head_num,
                         self.flashinfer_extra_state.tp_kv_head_num,
                         self.flashinfer_extra_state.head_dim,
-                        1,
+                        self.page_size,
                         q_data_type=self.flashinfer_extra_state.q_data_type,
                         kv_data_type=self.flashinfer_extra_state.kv_data_type,
                         non_blocking=True,
@@ -72,19 +89,33 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                 q_starts = self.b1_cu_q_seq_len.int()
                 kv_starts = self.b1_cu_kv_seq_len.int()
                 kv_last_page_len = torch.full((self.batch_size,), 1, dtype=torch.int32, device=input_ids.device)
+                length = cdiv(self.flashinfer_extra_state.max_seq_length, self.page_size)
                 kv_indices = torch.empty(
-                    self.batch_size * self.flashinfer_extra_state.max_seq_length,
+                    self.batch_size * length,
                     dtype=torch.int32,
                     device=input_ids.device,
                 )
-                repack_kv_index(
-                    self.req_manager.req_to_token_indexs,
-                    self.b_req_idx,
-                    self.b_seq_len,
-                    kv_starts[:-1],
-                    self.max_kv_seq_len,
-                    kv_indices,
-                )
+                if "page_size_variable" in model.mode:
+                    b_page_len = cdiv(self.b_seq_len, self.page_size)
+                    kv_starts[1:] = b_page_len.cumsum(0)
+                    kv_last_page_len = self.b_seq_len - (b_page_len - 1) * self.page_size
+                    repack_kv_index(
+                        self.req_manager.req_to_page_indexs,
+                        self.b_req_idx,
+                        b_page_len,
+                        kv_starts[:-1],
+                        cdiv(self.max_kv_seq_len, self.page_size),
+                        kv_indices,
+                    )
+                else:
+                    repack_kv_index(
+                        self.req_manager.req_to_token_indexs,
+                        self.b_req_idx,
+                        self.b_seq_len,
+                        kv_starts[:-1],
+                        self.max_kv_seq_len,
+                        kv_indices,
+                    )
                 self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(
                     self.flashinfer_extra_state.workspace_buffer,
                     qo_indptr_buf=q_starts,
@@ -100,7 +131,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                     self.flashinfer_extra_state.tp_q_head_num,
                     self.flashinfer_extra_state.tp_kv_head_num,
                     self.flashinfer_extra_state.head_dim,
-                    1,
+                    self.page_size,
                     causal=True,
                     pos_encoding_mode="NONE",
                     logits_soft_cap=0.0,
@@ -115,11 +146,11 @@ def copy_for_cuda_graph(self, new_infer_state):
             self.decode_wrapper.plan(
                 new_infer_state.kv_starts,
                 new_infer_state.kv_indices,
-                new_infer_state.kv_last_page_len_buffer,
+                new_infer_state.kv_last_page_len,
                 new_infer_state.flashinfer_extra_state.tp_q_head_num,
                 new_infer_state.flashinfer_extra_state.tp_kv_head_num,
                 new_infer_state.flashinfer_extra_state.head_dim,
-                1,
+                self.page_size,
                 q_data_type=new_infer_state.flashinfer_extra_state.q_data_type,
                 kv_data_type=new_infer_state.flashinfer_extra_state.kv_data_type,
                 non_blocking=True,

lightllm/models/llama/layer_infer/transformer_layer_infer.py

@@ -107,9 +107,16 @@ def _bind_attention(self):
                 raise Exception(f"Unsupported mode for fa3 backend: {self.mode}")
             return
         elif get_env_start_args().enable_flashinfer_prefill:
-            self._context_attention_kernel = partial(
-                LlamaTransformerLayerInfer._context_attention_flashinfer_kernel, self
-            )
+            if "page_size_variable" in self.mode:
+                self._context_attention_kernel = partial(
+                    LlamaTransformerLayerInfer._paged_context_attention_flashinfer_kernel, self
+                )
+            elif not self.mode:
+                self._context_attention_kernel = partial(
+                    LlamaTransformerLayerInfer._context_attention_flashinfer_kernel, self
+                )
+            else:
+                raise Exception(f"Unsupported mode for flashinfer backend: {self.mode}")
         else:
             self._context_attention_kernel = partial(LlamaTransformerLayerInfer._context_attention_kernel, self)
         if "ppl_int8kv" in self.mode:
@@ -174,6 +181,12 @@ def _bind_attention(self):
             self._copy_kv_to_mem_cache = partial(
                 LlamaTransformerLayerInfer._copy_kv_to_mem_cache_with_calibration, self
             )
+        elif "page_size_variable" in self.mode:
+            assert get_env_start_args().enable_flashinfer_prefill and get_env_start_args().enable_flashinfer_decode
+            self._token_attention_kernel = partial(
+                LlamaTransformerLayerInfer._paged_token_decode_attention_flashinfer, self
+            )
+            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
         elif not self.mode:
             if get_env_start_args().enable_flashinfer_decode:
                 self._token_attention_kernel = partial(
@@ -274,6 +287,21 @@ def _context_attention_flashinfer_kernel(
         )
         return o_tensor
 
+    def _paged_context_attention_flashinfer_kernel(
+        self, q, kv, infer_state: LlamaFlashInferStateInfo, layer_weight, out=None
+    ) -> torch.Tensor:
+        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
+        page_size = get_page_size()
+        kv = infer_state.mem_manager.kv_buffer[self.layer_num_].view(
+            -1, page_size, 2 * self.tp_k_head_num_, self.head_dim_
+        )
+        infer_state.prefill_wrapper.run(
+            q.view(q.shape[0], -1, self.head_dim_),
+            (kv[:, :, : self.tp_k_head_num_, :], kv[:, :, self.tp_k_head_num_ :, :]),
+            out=o_tensor.view(q.shape[0], -1, self.head_dim_),
+        )
+        return o_tensor
+
     def _context_attention_kernel(
         self, q, kv, infer_state: LlamaInferStateInfo, layer_weight, out=None
     ) -> torch.Tensor:
@@ -587,6 +615,24 @@ def _token_decode_attention_flashinfer(self, q, infer_state: LlamaFlashInferStat
         )
         return o_tensor
 
+    def _paged_token_decode_attention_flashinfer(
+        self, q, infer_state: LlamaFlashInferStateInfo, layer_weight, out=None
+    ):
+        batch_size = infer_state.batch_size
+        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
+
+        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
+        page_size = get_page_size()
+        kv = infer_state.mem_manager.kv_buffer[self.layer_num_].view(
+            -1, page_size, 2 * self.tp_k_head_num_, self.head_dim_
+        )
+        infer_state.decode_wrapper.run(
+            q.view(calcu_shape1),
+            (kv[:, :, : self.tp_k_head_num_, :], kv[:, :, self.tp_k_head_num_ :, :]),
+            out=o_tensor.view(calcu_shape1),
+        )
+        return o_tensor
+
     def _token_decode_attention_normal(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
         total_token_num = infer_state.total_token_num
         batch_size = infer_state.batch_size

lightllm/utils/envs_utils.py

@@ -153,7 +153,7 @@ def get_kv_quant_calibration_inference_count():
 def get_page_size():
     try:
         args = get_env_start_args()
-        return int(os.getenv("PAGE_SIZE", 4)) if "page_size_variable" in args.mode else 1
+        return int(os.getenv("PAGE_SIZE", 64)) if "page_size_variable" in args.mode else 1
     except:
         return 1
 

unit_tests/models/llama/test_context_flashattention_nopad.py

@@ -10,7 +10,6 @@
     context_attention_fwd_no_prompt_cache,
 )
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.common.req_manager import ReqManager
 
 logger = init_logger(__name__)
 
@@ -56,8 +55,6 @@ def test_context_attention_fwd(batch, seqlen, q_heads, kv_heads, head_dim):
     infer_state.batch_size = Z
     infer_state.max_len_in_batch = N_CTX
     infer_state.total_token_num = Z * N_CTX
-    infer_state.req_manager = ReqManager(Z, N_CTX, None)
-    infer_state.req_manager.req_to_token_indexs = req_to_token_indexs
     infer_state.b_req_idx = b_req_idx
     infer_state.b_seq_len = b_seq_len
     infer_state.b_ready_cache_len = b_ready_cache_len
@@ -73,7 +70,7 @@ def test_context_attention_fwd(batch, seqlen, q_heads, kv_heads, head_dim):
         infer_state.b_seq_len,
         infer_state.b_ready_cache_len,
         infer_state.max_len_in_batch,
-        infer_state.req_manager.req_to_token_indexs,
+        req_to_token_indexs,
     )
 
     batch_size = Z