feat: page attention for fa4

Author: niushengxiao

lightllm/common/basemodel/attention/fa4/fp.py

@@ -2,16 +2,17 @@
 import torch
 
 from ..base_att import AttControl
-from ..fa3.fp import Fa3AttBackend, Fa3PrefillAttState, Fa3DecodeAttState
+from ..paged_fa3.fp import PagedFa3AttBackend, PagedFa3PrefillAttState, PagedFa3DecodeAttState
 from lightllm.utils.fa4_utils import (
     ensure_fa4_available,
     ensure_fa4_supported_gpu,
     flash_attn_varlen_func,
+    sm90_fa4_paged_kv_tile_n,
     unwrap_fa4_output,
 )
 
 
-class Fa4AttBackend(Fa3AttBackend):
+class Fa4AttBackend(PagedFa3AttBackend):
     def __init__(self, model):
         ensure_fa4_available()
         ensure_fa4_supported_gpu()
@@ -29,20 +30,7 @@ def _sm90_fa4_paged_kv_tile_n(
     head_dim_v: int,
     window_size: tuple[int, int],
 ) -> int | None:
-    major, _minor = torch.cuda.get_device_capability()
-    if major != 9:
-        return None
-
-    is_local = window_size != (-1, -1)
-    if head_dim <= 64:
-        return 128
-    if head_dim <= 96:
-        return 128 if is_local else 144
-    if head_dim <= 128:
-        return 128
-    if head_dim <= 192:
-        return 96 if is_local else (128 if head_dim_v <= 128 else 112)
-    return 64 if is_local else 80
+    return sm90_fa4_paged_kv_tile_n(head_dim=head_dim, head_dim_v=head_dim_v, window_size=window_size)
 
 
 def _ensure_fa4_paged_kv_supported(
@@ -67,7 +55,7 @@ def _ensure_fa4_paged_kv_supported(
 
 
 @dataclasses.dataclass
-class Fa4PrefillAttState(Fa3PrefillAttState):
+class Fa4PrefillAttState(PagedFa3PrefillAttState):
     def _nomarl_prefill_att(
         self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty
     ) -> torch.Tensor:
@@ -84,12 +72,12 @@ def _nomarl_prefill_att(
         head_dim = q.shape[-1]
         head_dim_v = v.shape[-1]
         softmax_scale = 1.0 / (head_dim ** 0.5)
-        _ensure_fa4_paged_kv_supported(head_dim, head_dim_v, window_size, page_size=1)
+        _ensure_fa4_paged_kv_supported(head_dim, head_dim_v, window_size, page_size=self.backend.page_size)
 
         out = flash_attn_varlen_func(
             q=q,
-            k=k.view(k.shape[0], 1, k.shape[1], k.shape[2]),
-            v=v.view(v.shape[0], 1, v.shape[1], v.shape[2]),
+            k=k.view(-1, self.backend.page_size, k.shape[1], k.shape[2]),
+            v=v.view(-1, self.backend.page_size, v.shape[1], v.shape[2]),
             cu_seqlens_q=self.cu_seqlens_q,
             seqused_k=self.infer_state.b_seq_len.int(),
             max_seqlen_q=self.infer_state.max_q_seq_len,
@@ -106,7 +94,7 @@ def _nomarl_prefill_att(
 
 
 @dataclasses.dataclass
-class Fa4DecodeAttState(Fa3DecodeAttState):
+class Fa4DecodeAttState(PagedFa3DecodeAttState):
     def _normal_decode_att(
         self,
         q: torch.Tensor,
@@ -128,12 +116,12 @@ def _normal_decode_att(
         head_dim = q.shape[-1]
         head_dim_v = v.shape[-1]
         softmax_scale = 1.0 / (head_dim ** 0.5)
-        _ensure_fa4_paged_kv_supported(head_dim, head_dim_v, window_size, page_size=1)
+        _ensure_fa4_paged_kv_supported(head_dim, head_dim_v, window_size, page_size=self.backend.page_size)
 
         out = flash_attn_varlen_func(
             q=q,
-            k=k.view(k.shape[0], 1, k.shape[1], k.shape[2]),
-            v=v.view(v.shape[0], 1, v.shape[1], v.shape[2]),
+            k=k.view(-1, self.backend.page_size, k.shape[1], k.shape[2]),
+            v=v.view(-1, self.backend.page_size, v.shape[1], v.shape[2]),
             cu_seqlens_q=self.cu_seqlens_q,
             seqused_k=self.b_att_seq_len.int(),
             max_seqlen_q=self.decode_max_q_seq_len,

lightllm/server/api_start.py

@@ -11,7 +11,7 @@
 from .embed_cache.manager import start_cache_manager
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.envs_utils import set_env_start_args, set_unique_server_name, get_unique_server_name
-from lightllm.utils.envs_utils import get_lightllm_gunicorn_keep_alive, get_page_size
+from lightllm.utils.envs_utils import get_lightllm_gunicorn_keep_alive, get_page_size, set_page_size
 from .detokenization.manager import start_detokenization_process
 from .router.manager import start_router_process
 from lightllm.utils.process_check import is_process_active
@@ -29,6 +29,20 @@
 logger = init_logger(__name__)
 
 
+def _auto_set_fa4_page_size(args, requested_backends):
+    if "fa4" not in requested_backends or "PAGE_SIZE" in os.environ:
+        return
+
+    from lightllm.utils.fa4_utils import infer_fa4_page_size
+
+    page_size = infer_fa4_page_size(args.model_dir)
+    if page_size is None:
+        return
+
+    set_page_size(page_size)
+    logger.info(f"auto set PAGE_SIZE={page_size} for FA4 backend")
+
+
 def setup_signal_handlers(http_server_process, process_manager):
     def signal_handler(sig, frame):
         if sig == signal.SIGINT:
@@ -205,6 +219,11 @@ def normal_or_p_d_start(args):
                 f"{sorted(allowed_ep_att_backends)}; flashinfer is not supported."
             )
 
+    llm_requested_backends = list(args.llm_prefill_att_backend) + list(args.llm_decode_att_backend)
+    requested_backends = llm_requested_backends + list(args.vit_att_backend)
+    if "fa4" in requested_backends:
+        _auto_set_fa4_page_size(args, llm_requested_backends)
+
     # mtp params check
     if args.mtp_mode is not None:
         assert args.mtp_draft_model_dir is not None

lightllm/utils/envs_utils.py

@@ -170,6 +170,11 @@ def get_page_size():
     return int(os.getenv("PAGE_SIZE", 1))
 
 
+def set_page_size(page_size: int):
+    os.environ["PAGE_SIZE"] = str(page_size)
+    get_page_size.cache_clear()
+
+
 g_model_init_done = False
 
 

lightllm/utils/fa4_utils.py

@@ -42,5 +42,41 @@ def ensure_fa4_supported_gpu() -> None:
         )
 
 
+def sm90_fa4_paged_kv_tile_n(head_dim: int, head_dim_v: int, window_size: tuple[int, int] = (-1, -1)) -> int | None:
+    major, _minor = torch.cuda.get_device_capability()
+    if major != 9:
+        return None
+
+    is_local = window_size != (-1, -1)
+    if head_dim <= 64:
+        return 128
+    if head_dim <= 96:
+        return 128 if is_local else 144
+    if head_dim <= 128:
+        return 128
+    if head_dim <= 192:
+        return 96 if is_local else (128 if head_dim_v <= 128 else 112)
+    return 64 if is_local else 80
+
+
+def infer_fa4_page_size(model_dir: str) -> int | None:
+    from transformers.configuration_utils import PretrainedConfig
+
+    model_cfg, _ = PretrainedConfig.get_config_dict(model_dir)
+    llm_config = model_cfg.get("text_config", model_cfg)
+
+    head_dim = llm_config.get("head_dim")
+    if head_dim is None:
+        head_dim = llm_config["hidden_size"] // llm_config["num_attention_heads"]
+    head_dim_v = llm_config.get("v_head_dim", head_dim)
+
+    window_size = (-1, -1)
+    sliding_window = llm_config.get("sliding_window", None)
+    if sliding_window is not None and not llm_config.get("full_attention_interval", None):
+        window_size = (sliding_window - 1, sliding_window - 1)
+
+    return sm90_fa4_paged_kv_tile_n(head_dim=head_dim, head_dim_v=head_dim_v, window_size=window_size)
+
+
 def unwrap_fa4_output(output):
     return output[0] if isinstance(output, tuple) else output