feat: add protected code

Author: niushengxiao

lightllm/models/deepseek2/flashattention_infer_struct.py

@@ -16,6 +16,7 @@ class Deepseek2FlashAttentionStateInfo(Deepseek2InferStateInfo):
 
     def __init__(self):
         super().__init__()
+        self.page_size = get_page_size()
 
     @classmethod
     def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
@@ -43,19 +44,18 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
             self.cu_seqlens_q = self.b1_cu_q_seq_len
             self.cu_seqlens_k = self.b1_cu_kv_seq_len
             max_seq_len_k = self.max_kv_seq_len
-            page_size = get_page_size()
             if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
-                length = cdiv(model.graph_max_len_in_batch, page_size)
+                length = cdiv(model.graph_max_len_in_batch, self.page_size)
                 page_buffer = Deepseek2FlashAttentionStateInfo.get_page_table_buffer(model.graph_max_batch_size, length)
                 self.page_table = page_buffer[self.microbatch_index][: self.batch_size * length].reshape(
                     self.batch_size, length
                 )
             else:
-                length = cdiv(self.max_len_in_batch, page_size)
+                length = cdiv(self.max_len_in_batch, self.page_size)
                 self.page_table = torch.empty((self.batch_size, length), dtype=torch.int32).to(input_ids.device)
 
             if "page_size_variable" in model.mode:
-                length = cdiv(max_seq_len_k, page_size)
+                length = cdiv(max_seq_len_k, self.page_size)
                 self.page_table[:, :length].copy_(model.req_manager.req_to_page_indexs[self.b_req_idx, :length])
                 self.page_table[:, length:].fill_(0)
             else:

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -26,7 +26,7 @@
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 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, get_page_size
+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
 from lightllm.utils.sgl_utils import flash_attn_varlen_func, flash_attn_with_kvcache, merge_state_v2
@@ -590,12 +590,11 @@ def _token_gqa_decode_attention_flashattention(
     def _token_gqa_decode_attention_flashattention_paged(
         self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):
-        page_size = get_page_size()
         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, page_size, 1, self.qk_rope_head_dim)
-        kv_nope = kv[:, :, : -self.qk_rope_head_dim].reshape(-1, page_size, 1, self.kv_lora_rank)
+        k_rope = kv[:, :, -self.qk_rope_head_dim :].reshape(-1, infer_state.page_size, 1, self.qk_rope_head_dim)
+        kv_nope = kv[:, :, : -self.qk_rope_head_dim].reshape(-1, infer_state.page_size, 1, self.kv_lora_rank)
         k_descale, v_descale = None, None
         o_tensor = flash_attn_with_kvcache(
             q=q_rope,
@@ -639,7 +638,6 @@ def _token_gqa_decode_attention_flashinfer(
     def _token_gqa_decode_attention_flashinfer_paged(
         self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):
-        page_size = get_page_size()
         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)
 
@@ -649,8 +647,8 @@ def _token_gqa_decode_attention_flashinfer_paged(
         infer_state.decode_wrapper.run(
             q_nope,
             q_rope,
-            kv[:, :, : -self.qk_rope_head_dim].reshape(-1, page_size, 1, self.kv_lora_rank),
-            kv[:, :, -self.qk_rope_head_dim :].reshape(-1, page_size, 1, self.qk_rope_head_dim),
+            kv[:, :, : -self.qk_rope_head_dim].reshape(-1, infer_state.page_size, 1, self.kv_lora_rank),
+            kv[:, :, -self.qk_rope_head_dim :].reshape(-1, infer_state.page_size, 1, self.qk_rope_head_dim),
             out=o_tensor,
             return_lse=False,
         )

lightllm/models/llama/flashattention_infer_struct.py

@@ -18,6 +18,7 @@ class FlashAttentionStateInfo(LlamaInferStateInfo):
 
     def __init__(self):
         super().__init__()
+        self.page_size = get_page_size()
 
     @classmethod
     def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
@@ -33,7 +34,7 @@ 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()
-            length = cdiv(self.max_seq_len, get_page_size())
+            length = cdiv(self.max_seq_len, self.page_size)
             self.page_table = torch.empty((self.batch_size, length), dtype=torch.int32, device=input_ids.device)
             if "page_size_variable" in model.mode:
                 self.page_table.copy_(model.req_manager.req_to_page_indexs[self.b_req_idx, :length])
@@ -45,17 +46,16 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
             self.cu_seqlens_k = self.b1_cu_kv_seq_len.int()
             max_seq_len_k = self.max_kv_seq_len
             if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
-                page_size = get_page_size()
-                length = cdiv(model.graph_max_len_in_batch, page_size)
+                length = cdiv(model.graph_max_len_in_batch, self.page_size)
                 page_buffer = FlashAttentionStateInfo.get_page_table_buffer(model.graph_max_batch_size, length)
                 self.page_table = page_buffer[self.microbatch_index][: self.batch_size * length].reshape(
                     self.batch_size, length
                 )
             else:
-                length = cdiv(self.max_len_in_batch, get_page_size())
+                length = cdiv(self.max_len_in_batch, self.page_size)
                 self.page_table = torch.empty((self.batch_size, length), dtype=torch.int32, device=input_ids.device)
 
-            length = cdiv(max_seq_len_k, get_page_size())
+            length = cdiv(max_seq_len_k, self.page_size)
             if "page_size_variable" in model.mode:
                 self.page_table[:, :length].copy_(model.req_manager.req_to_page_indexs[self.b_req_idx, :length])
             else:

lightllm/models/llama/layer_infer/transformer_layer_infer.py

@@ -27,7 +27,7 @@
 from lightllm.models.llama.triton_kernel.ppl_quant_copy_kv import destindex_copy_dequantize_kv
 from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor
 from lightllm.utils.log_utils import init_logger
-from lightllm.utils.envs_utils import get_env_start_args, get_page_size
+from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.light_utils import HAS_LIGHTLLM_KERNEL, light_ops
 from lightllm.common.basemodel.triton_kernel.q_per_head_fp8_quant import q_per_head_fp8_quant
 from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
@@ -291,9 +291,8 @@ 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_
+            -1, infer_state.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_),
@@ -356,13 +355,12 @@ def _context_attention_kernel_ppl_int8kv(
     def _paged_context_attention_flashattention(
         self, q, kv, infer_state: FlashAttentionStateInfo, layer_weight, out=None
     ):
-        page_size = get_page_size()
         cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, page_size, self.tp_k_head_num_, self.head_dim_
+            -1, infer_state.page_size, self.tp_k_head_num_, self.head_dim_
         )
         cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
             :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, page_size, self.tp_v_head_num_, self.head_dim_)
+        ].reshape(-1, infer_state.page_size, self.tp_v_head_num_, self.head_dim_)
         q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
         k_descale, v_descale = None, None  # disable quantization
         Lq = q.shape[-1]
@@ -622,9 +620,8 @@ def _paged_token_decode_attention_flashinfer(
         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_
+            -1, infer_state.page_size, 2 * self.tp_k_head_num_, self.head_dim_
         )
         infer_state.decode_wrapper.run(
             q.view(calcu_shape1),
@@ -914,13 +911,12 @@ def _token_decode_attention_gqa_flashdecoding_vsm(
     def _paged_token_decode_attention_flashattention(
         self, q, infer_state: FlashAttentionStateInfo, layer_weight, out=None
     ):
-        page_size = get_page_size()
         cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, page_size, self.tp_k_head_num_, self.head_dim_
+            -1, infer_state.page_size, self.tp_k_head_num_, self.head_dim_
         )
         cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
             :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, page_size, self.tp_v_head_num_, self.head_dim_)
+        ].reshape(-1, infer_state.page_size, self.tp_v_head_num_, self.head_dim_)
         q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
         k_descale, v_descale = None, None  # disable quantization
         Lq = q.shape[-1]

lightllm/server/api_cli.py

@@ -165,7 +165,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         nargs="+",
         help="""Model mode: [triton_int8kv | ppl_int8kv | ppl_fp16 | triton_flashdecoding
                         | triton_gqa_attention | triton_gqa_flashdecoding | triton_fp8kv | offline_calibration_fp8kv
-                        | export_fp8kv_calibration
+                        | export_fp8kv_calibration | page_size_variable
                         triton_flashdecoding mode is for long context, current support llama llama2 qwen;
                         triton_gqa_attention and triton_gqa_flashdecoding is fast kernel for model which use GQA;
                         triton_int8kv mode use int8 to store kv cache, can increase token capacity, use triton kernel;
@@ -177,6 +177,8 @@ def make_argument_parser() -> argparse.ArgumentParser:
                         Calibration need to disable cudagraph and use fa3 or flashinfer backend.
                         ppl_int8kv mode use int8 to store kv cache, and use ppl fast kernel;
                         ppl_fp16 mode use ppl fast fp16 decode attention kernel;
+                        page_size_variable allow to use page size > 1, use PAGE_SIZE env to set page size,
+                        page_size_variable only support fa3 and flashinfer backend for now
                         you need to read source code to make sure the supported detail mode for all models""",
     )
     parser.add_argument(

lightllm/server/api_start.py

@@ -125,6 +125,13 @@ def normal_or_p_d_start(args):
             "--enable_flashinfer_prefill and --enable_flashinfer_decode"
         )
         assert args.disable_cudagraph is True, "export_fp8kv_calibration mode need disable cudagraph"
+    if "page_size_variable" in args.mode:
+        assert args.enable_fa3 is True or (
+            args.enable_flashinfer_prefill is True and args.enable_flashinfer_decode is True
+        ), (
+            "page_size_variable mode need enable fa3 or flashinfer, add --enable_fa3 or "
+            "--enable_flashinfer_prefill and --enable_flashinfer_decode"
+        )
 
     # 部分模式还不能支持与高级动态调度算法协同，to do.
     if args.diverse_mode: