feat: opt autotuner

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

@@ -3,7 +3,7 @@
 from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
 from typing import Optional, TYPE_CHECKING
 from lightllm.utils.dist_utils import get_current_device_id
-from lightllm.utils.sgl_utils import flash_attn_with_kvcache
+from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy
 from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor
@@ -222,7 +222,7 @@ def _normal_decode_att(
         k_descale, v_descale = None, None  # disable quantization
         Lq = q.shape[-1]
         sm_scale = 1.0 / (Lq ** 0.5)
-        o = flash_attn_with_kvcache(
+        o = flash_attn_with_kvcache_autotune(
             q=q,
             k_cache=k.view(k.shape[0], 1, k.shape[1], k.shape[2]),
             v_cache=v.view(v.shape[0], 1, v.shape[1], v.shape[2]),

lightllm/common/basemodel/basemodel.py

@@ -28,9 +28,12 @@
 from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type, get_added_mtp_kv_layer_num
 from lightllm.distributed.communication_op import dist_group_manager
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
-from lightllm.common.triton_utils.autotuner import AutotuneLevel
 from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch
-from lightllm.utils.envs_utils import set_model_init_status, enable_diverse_mode_gqa_decode_fast_kernel
+from lightllm.utils.envs_utils import (
+    set_model_init_status,
+    enable_diverse_mode_gqa_decode_fast_kernel,
+    enable_full_att_decode_tune,
+)
 from lightllm.common.triton_utils.autotuner import Autotuner
 from lightllm.utils.infer_utils import post_empty_cache
 from .attention import get_prefill_att_backend_class, get_decode_att_backend_class
@@ -126,6 +129,7 @@ def __init__(self, kvargs):
             logger.info(f"use decode att backend1: {self.decode_att_backend1.__class__.__name__}")
 
         self._autotune_warmup()
+        self._full_att_decode_autotune()
         self._init_padded_req()
         self._init_cudagraph()
         self._init_prefill_cuda_graph()
@@ -286,6 +290,52 @@ def _init_prefill_cuda_graph(self):
             else:
                 self.prefill_graph.warmup(self)
 
+    @final
+    @torch.no_grad()
+    @post_empty_cache
+    def _full_att_decode_autotune(self):
+        """
+        Warm up / autotune FA3 full-attention decode ``num_splits`` before CUDA Graph capture.
+
+        Runs only when all of the following hold:
+          - CUDA Graph is enabled (``disable_cudagraph`` is False)
+          - this is the main model (MTP draft models are skipped)
+          - ``ENABLE_FULL_ATT_DECODE_TUNE`` is set to 1/ON/TRUE (default off)
+          - decode attention backend is ``Fa3AttBackend``
+
+        Candidate batch sizes follow the same schedule as CUDA Graph capture.
+        Actual benchmarking is delegated to ``fa3_decode_autotune`` in ``sgl_utils``.
+        """
+        if self.disable_cudagraph:
+            return
+        # Only tune on the main model; MTP draft models skip this path.
+        if getattr(self, "is_mtp_draft_model", False):
+            return
+
+        # Opt-in switch for FA3 full-attention decode num_splits tuning.
+        # Set ENABLE_FULL_ATT_DECODE_TUNE=1/ON/TRUE to enable; default is off.
+        if not enable_full_att_decode_tune():
+            return
+
+        # Only Fa3AttBackend decode path needs this num_splits warmup.
+        decode_backends = [
+            self.decode_att_backend,
+            getattr(self, "decode_att_backend1", None),
+        ]
+        if not any(
+            backend is not None and backend.__class__.__name__ == "Fa3AttBackend" for backend in decode_backends
+        ):
+            return
+
+        from lightllm.utils.sgl_utils import fa3_decode_autotune
+
+        cuda_graph_batch_sizes = CudaGraph.gen_cuda_graph_batch_sizes(
+            max_batch_size=self.graph_max_batch_size,
+            tp_world_size=self.tp_world_size_,
+        )
+        fa3_decode_autotune(self, cuda_graph_batch_sizes)
+        return
+
     def _init_custom(self):
         pass
 
@@ -1050,7 +1100,7 @@ def _autotune_warmup(self):
         Autotuner.start_autotune_warmup()
         torch.distributed.barrier()
 
-        warmup_lengths = [1, 8, 16, 32, 64, 100, 128, 256, 1024, 2048, 4096]
+        warmup_lengths = [1, 4, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096]
 
         if self.batch_max_tokens not in warmup_lengths:
             warmup_lengths.append(self.batch_max_tokens)

lightllm/common/basemodel/cuda_graph.py

@@ -17,38 +17,49 @@
 class CudaGraph:
     # CudaGraph forward pass for the decoding stage.
 
-    def __init__(self, max_batch_size=8, max_len_in_batch=8192, tp_world_size: int = 1):
-        self.graph = {}
-        self.tp_world_size = tp_world_size
-        self.mempool = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None
-        self.args = get_env_start_args()
-        self.mtp_step = self.args.mtp_step
-        self.max_batch_size = max_batch_size
-        self.graph_max_len_in_batch = max_len_in_batch
-        self.enable_decode_microbatch_overlap = self.args.enable_decode_microbatch_overlap
+    @staticmethod
+    def gen_cuda_graph_batch_sizes(max_batch_size=8, tp_world_size: int = 1):
+        args = get_env_start_args()
+        mtp_size = args.mtp_step + 1
 
         # gen cuda graph batch_sizes
         # cuda graph gen for batch size = [1, 2, 3, ..., graph_split_batch_size]
         # and [graph_split_batch_size + graph_grow_step_size,
         # if the mtp_step is not 0, then the batch_sizes will be multiply of (mtp_step + 1)
 
-        graph_split_batch_size = self.args.graph_split_batch_size * (self.mtp_step + 1)
-        graph_grow_step_size = self.args.graph_grow_step_size * (self.mtp_step + 1)
+        graph_split_batch_size = args.graph_split_batch_size * mtp_size
+        graph_grow_step_size = args.graph_grow_step_size * mtp_size
 
-        batch_sizes = [i * (self.mtp_step + 1) for i in range(1, self.args.graph_split_batch_size + 1)]
+        batch_sizes = [i * mtp_size for i in range(1, args.graph_split_batch_size + 1)]
         for _batch_size in range(graph_split_batch_size + graph_grow_step_size, max_batch_size, graph_grow_step_size):
             batch_sizes.append(_batch_size)
 
         batch_sizes = list(set([e for e in batch_sizes if e < max_batch_size]))
         batch_sizes.append(max_batch_size)
         batch_sizes.sort()
-        if self.args.enable_tpsp_mix_mode:
-            batch_sizes = [triton.cdiv(e, self.tp_world_size) * self.tp_world_size for e in batch_sizes]
+        if args.enable_tpsp_mix_mode:
+            batch_sizes = [triton.cdiv(e, tp_world_size) * tp_world_size for e in batch_sizes]
             batch_sizes = list(set(batch_sizes))
             batch_sizes.sort()
 
-        self.cuda_graph_batch_sizes = batch_sizes
-        assert batch_sizes[-1] == self.max_batch_size
+        assert batch_sizes[-1] == max_batch_size
+        return batch_sizes
+
+    def __init__(self, max_batch_size=8, max_len_in_batch=8192, tp_world_size: int = 1):
+        self.graph = {}
+        self.tp_world_size = tp_world_size
+        self.mempool = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None
+        self.args = get_env_start_args()
+        self.mtp_step = self.args.mtp_step
+        self.max_batch_size = max_batch_size
+        self.graph_max_len_in_batch = max_len_in_batch
+        self.enable_decode_microbatch_overlap = self.args.enable_decode_microbatch_overlap
+
+        self.cuda_graph_batch_sizes = self.gen_cuda_graph_batch_sizes(
+            max_batch_size=max_batch_size,
+            tp_world_size=tp_world_size,
+        )
+        assert self.cuda_graph_batch_sizes[-1] == self.max_batch_size
         logger.info(f"cuda graph batch_sizes: {self.cuda_graph_batch_sizes}")
 
     def can_run(self, batch_size, max_len_in_batch):

lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py

@@ -455,7 +455,7 @@ def moe_align2(token_num_mul_topk_num: int, exports_token_num: torch.Tensor, blo
     out tensor is a tensor that contain block schduel infos tensor.
     """
     max_num_tokens_padded = token_num_mul_topk_num + exports_token_num.shape[0] * (block_m - 1)
-    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_m)
+    max_num_m_blocks = min(token_num_mul_topk_num, triton.cdiv(max_num_tokens_padded, block_m))
     # first is expert, second is m_index, third is token_start_index
     mblocks_to_tuple_info = torch.empty((max_num_m_blocks, 3), dtype=torch.int32, device="cuda")
 

lightllm/utils/envs_utils.py

@@ -165,6 +165,18 @@ def get_triton_autotune_level():
     return int(os.getenv("LIGHTLLM_TRITON_AUTOTUNE_LEVEL", 0))
 
 
+@lru_cache(maxsize=None)
+def enable_full_att_decode_tune() -> bool:
+    """
+    Whether to run FA3 full-attention decode num_splits warmup/autotune at model init.
+
+    Env: ENABLE_FULL_ATT_DECODE_TUNE
+      - ON / TRUE / 1: enable
+      - otherwise (default False): skip this operator-specific tuning
+    """
+    return enable_env_vars("ENABLE_FULL_ATT_DECODE_TUNE")
+
+
 g_model_init_done = False
 
 

lightllm/utils/sgl_utils.py

@@ -1,3 +1,7 @@
+import torch
+
+from lightllm.common.triton_utils.autotuner import AutotuneLevel, Autotuner, autotune
+from lightllm.utils.envs_utils import get_triton_autotune_level
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -30,3 +34,154 @@
         "sgl_kernel is not installed, or the installed version did not support fa3. \
         Try to upgrade it."
     )
+
+
+def _flash_attn_kvcache_num_splits_configs():
+    return [{"num_splits": num_splits} for num_splits in [0, 16, 32]]
+
+
+def _flash_attn_kvcache_static_key(q, k_cache, v_cache, causal, window_size, softcap, sinks, k_descale, v_descale):
+    return {
+        "qd": str(q.dtype),
+        "kd": str(k_cache.dtype),
+        "vd": str(v_cache.dtype),
+        "qh": int(q.shape[-2]),
+        "kh": int(k_cache.shape[-2]),
+        "hd": int(q.shape[-1]),
+        "vh": int(v_cache.shape[-1]),
+        "pb": int(k_cache.shape[-3]),  # page size
+        "c": int(bool(causal)),
+        "wl": int(window_size[0]),
+        "wr": int(window_size[1]),
+        "sc": int(softcap > 0.0),
+        "sk": int(sinks is not None),
+        "has_k_descale": k_descale is not None,
+        "has_v_descale": v_descale is not None,
+        "sgl": getattr(sgl_ops, "__version__", "unknown"),
+    }
+
+
+def _flash_attn_max_q_len(q, max_seqlen_q):
+    return int(max_seqlen_q if max_seqlen_q is not None else q.shape[1] if q.dim() >= 4 else q.shape[0])
+
+
+def _flash_attn_kvcache_run_key(q, page_table, max_seqlen_q):
+    batch_size = int(page_table.shape[0])
+    max_q_len = _flash_attn_max_q_len(q, max_seqlen_q)
+    max_kv_len = int(page_table.shape[1])
+    return batch_size * 10_000_000_000_000 + max_q_len * 10_000_000 + max_kv_len
+
+
+@autotune(
+    kernel_name="sgl_fa3_kvcache_ns:v1",
+    configs_gen_func=_flash_attn_kvcache_num_splits_configs,
+    static_key_func=_flash_attn_kvcache_static_key,
+    run_key_func=_flash_attn_kvcache_run_key,
+)
+@torch.no_grad()
+def flash_attn_with_kvcache_autotune(
+    q,
+    k_cache,
+    v_cache,
+    cache_seqlens=None,
+    page_table=None,
+    cu_seqlens_q=None,
+    cu_seqlens_k_new=None,
+    max_seqlen_q=None,
+    causal=False,
+    window_size=(-1, -1),
+    softcap=0.0,
+    num_splits=0,
+    sinks=None,
+    k_descale=None,
+    v_descale=None,
+    run_config=None,
+    **kwargs,
+):
+    if run_config is None:
+        run_config = {"num_splits": 0}
+
+    num_splits = run_config["num_splits"]
+    return flash_attn_with_kvcache(
+        q=q,
+        k_cache=k_cache,
+        v_cache=v_cache,
+        cache_seqlens=cache_seqlens,
+        page_table=page_table,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k_new=cu_seqlens_k_new,
+        max_seqlen_q=max_seqlen_q,
+        causal=causal,
+        window_size=window_size,
+        softcap=softcap,
+        num_splits=num_splits,
+        sinks=sinks,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        **kwargs,
+    )
+
+
+def fa3_decode_autotune(model, cuda_graph_batch_sizes):
+    # 是否开启自动调优
+    if get_triton_autotune_level() not in [
+        AutotuneLevel.ADAPTIVE_AUTOTUNE,
+        AutotuneLevel.FORCE_AUTOTUNE,
+    ]:
+        return
+
+    Autotuner.start_autotune_warmup()
+    try:
+        max_kv_len = int(model.graph_max_len_in_batch)
+        if max_kv_len <= 0:
+            return
+
+        k, v = model.mem_manager.get_att_input_params(layer_index=0)
+        k_cache = k.view(k.shape[0], 1, k.shape[1], k.shape[2])
+        v_cache = v.view(v.shape[0], 1, v.shape[1], v.shape[2])
+        q_head_num = int(model.config["num_attention_heads"]) // model.tp_world_size_
+        head_dim = int(k.shape[-1])
+        mtp_size = model.args.mtp_step + 1
+
+        for batch_size in cuda_graph_batch_sizes[::-1]:
+            att_batch_size = batch_size // mtp_size
+            if att_batch_size <= 0:
+                continue
+            # 因为完整的kv空间可能无法装下所有token，所以在tuning的时候，所有token都使用相同的kv空间。
+            # 保证tuning的时候不会出现大的问题。
+            kv_range = torch.arange(att_batch_size * max_kv_len, dtype=torch.int32, device=k.device) % max_kv_len
+            k[kv_range].zero_()
+            v[kv_range].zero_()
+
+            q = torch.zeros(
+                (att_batch_size * mtp_size, q_head_num, head_dim),
+                dtype=model.data_type,
+                device=k.device,
+            )
+            page_table = kv_range.view(att_batch_size, max_kv_len)
+            cache_seqlens = torch.full((att_batch_size,), max_kv_len, dtype=torch.int32, device=k.device)
+            cu_seqlens_q = torch.arange(att_batch_size + 1, dtype=torch.int32, device=k.device) * mtp_size
+            cu_seqlens_k = torch.arange(att_batch_size + 1, dtype=torch.int32, device=k.device) * max_kv_len
+            softmax_scale = 1.0 / (head_dim ** 0.5)
+
+            flash_attn_with_kvcache_autotune(
+                q=q,
+                k_cache=k_cache,
+                v_cache=v_cache,
+                page_table=page_table,
+                cache_seqlens=cache_seqlens,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k_new=cu_seqlens_k,
+                max_seqlen_q=mtp_size,
+                softmax_scale=softmax_scale,
+                causal=True,
+                window_size=(-1, -1),
+                softcap=0.0,
+                k_descale=None,
+                v_descale=None,
+                return_softmax_lse=False,
+                sinks=None,
+            )
+    finally:
+        Autotuner.end_autotune_warmup()
+    return