diff --git a/lightllm/common/basemodel/attention/fa3/fp.py b/lightllm/common/basemodel/attention/fa3/fp.py index 7ba00e5911..399970a056 100644 --- a/lightllm/common/basemodel/attention/fa3/fp.py +++ b/lightllm/common/basemodel/attention/fa3/fp.py @@ -158,6 +158,7 @@ class Fa3DecodeAttState(BaseDecodeAttState): b_att_seq_len: torch.Tensor = None # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。 decode_max_q_seq_len: int = None + decode_max_kv_seq_len: int = None causal: bool = None def init_state(self): @@ -225,6 +226,8 @@ def _init_page_table(self, b_att_req_idx: torch.Tensor): att_batch_size = b_att_req_idx.shape[0] model = self.backend.model actual_max_kv_len = self.infer_state.max_kv_seq_len + # Graph 捕获会将 infer_state.max_kv_seq_len 改为容量上限，提前保存真实长度用于 FA3 配置查找。 + self.decode_max_kv_seq_len = actual_max_kv_len page_table_width = actual_max_kv_len if model.graph is not None and model.graph.can_run( batch_size=self.infer_state.batch_size, @@ -295,8 +298,9 @@ def _normal_decode_att( page_table=self.page_table, cache_seqlens=self.b_att_seq_len, cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, + cu_seqlens_k_new=None, # KV 已提前写入缓存，此处不追加新的 K/V。 max_seqlen_q=self.decode_max_q_seq_len, + max_seqlen_k=self.decode_max_kv_seq_len, softmax_scale=sm_scale, causal=self.causal, window_size=window_size, diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py index e7ce66c774..e87ec63ea3 100644 --- a/lightllm/common/basemodel/attention/triton/fp.py +++ b/lightllm/common/basemodel/attention/triton/fp.py @@ -95,8 +95,11 @@ def _nomarl_prefill_att( @dataclasses.dataclass class TritonDecodeAttState(BaseDecodeAttState): b_mark_mtp_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None def init_state(self): + # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于 GQA decode 配置查找。 + self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: self.b_mark_mtp_shared_group = build_mtp_shared_group_markers( @@ -212,6 +215,7 @@ def _normal_decode_gqa_flash_decoding_att( infer_state=self.infer_state, cache_k=k, cache_v=v, + max_len_in_batch=self.decode_max_kv_seq_len, out=out, alloc_tensor_func=alloc_func, sliding_window=sliding_window, @@ -238,6 +242,7 @@ def _spec_decode_gqa_att( B_req_idx=self.infer_state.b_req_idx, b_seq_len=self.infer_state.b_seq_len, b_mark_shared_group=self.b_mark_mtp_shared_group, + max_kv_len=self.decode_max_kv_seq_len, alloc_tensor_func=alloc_func, ) diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py index 25199dc470..5a7c8b3d4d 100644 --- a/lightllm/common/basemodel/attention/triton/int4kv.py +++ b/lightllm/common/basemodel/attention/triton/int4kv.py @@ -115,8 +115,11 @@ def _groupsize_quant_prefill_att( @dataclasses.dataclass class Int4kvTritonDecodeAttState(BaseDecodeAttState): + decode_max_kv_seq_len: int = None + def init_state(self): - pass + # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于配置查找。 + self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len def copy_for_decode_cuda_graph(self, new_state: "Int4kvTritonDecodeAttState"): super().copy_for_decode_cuda_graph(new_state) @@ -166,5 +169,6 @@ def ppl_int4kv_decode_att( cache_k_scale=k_scale, cache_v=v, cache_v_scale=v_scale, + max_kv_seq_len=self.decode_max_kv_seq_len, alloc_tensor_func=alloc_func, ) diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py index f2b3371ccd..13deab9c0e 100644 --- a/lightllm/common/basemodel/attention/triton/int8kv.py +++ b/lightllm/common/basemodel/attention/triton/int8kv.py @@ -118,8 +118,11 @@ def _groupsize_quant_prefill_att( class Int8kvTritonDecodeAttState(BaseDecodeAttState): b_shared_seq_len: torch.Tensor = None b_mark_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None def init_state(self): + # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于普通 decode 配置查找。 + self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len if enable_diverse_mode_gqa_decode_fast_kernel(): self.b_mark_shared_group = build_diverse_shared_group_markers( b_shared_radix_node_id=self.infer_state.b_shared_radix_node_id, @@ -203,5 +206,6 @@ def normal_decode_att( cache_k_scale=k_scale, cache_v=v, cache_v_scale=v_scale, + max_len_in_batch=self.decode_max_kv_seq_len, alloc_tensor_func=alloc_func, ) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e80f2b552f..f1247e0ef4 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -38,11 +38,8 @@ ) from lightllm.common.basemodel.mtp_manager import MtpManager from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch -from lightllm.utils.envs_utils import ( - set_model_init_status, - enable_full_att_decode_tune, -) -from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.utils.envs_utils import set_model_init_status +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType from lightllm.utils.infer_utils import post_empty_cache from lightllm.utils.torch_memory_saver_utils import ( TorchMemorySaverWrapper, @@ -140,7 +137,6 @@ def __init__(self, kvargs): self._init_hidden_collector() self._autotune_warmup() - self._full_att_decode_autotune() self._init_padded_req() self._init_cudagraph() self._init_prefill_cuda_graph() @@ -308,60 +304,6 @@ 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 self.is_mtp_draft_model: - 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, - self.decode_att_backend1, - ] - 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 - - decode_batch_multiplier = self.mtp_manager.get_decode_batch_multiplier(self.is_mtp_draft_model) - cuda_graph_grow_step_size = self.mtp_manager.get_decode_cuda_graph_grow_step_size(self.is_mtp_draft_model) - cuda_graph_batch_sizes = CudaGraph.gen_cuda_graph_batch_sizes( - batch_step_size_before_split=cuda_graph_grow_step_size, - split_batch_size=self.args.graph_split_batch_size * decode_batch_multiplier, - batch_step_size_after_split=self.args.graph_grow_step_size * cuda_graph_grow_step_size, - max_batch_size=self.graph_max_batch_size, - tp_world_size=self.tp_world_size_, - ) - cuda_graph_batch_sizes = [ - batch_size for batch_size in cuda_graph_batch_sizes if batch_size % decode_batch_multiplier == 0 - ] - fa3_decode_autotune(self, cuda_graph_batch_sizes, batch_multiplier=decode_batch_multiplier) - return - def _init_custom(self): pass @@ -1155,7 +1097,7 @@ def autotune_layers(self): @torch.no_grad() @post_empty_cache def _autotune_warmup(self): - Autotuner.start_autotune_warmup() + Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) torch.distributed.barrier() warmup_lengths = [1, 4, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096] diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 5849cccf54..ceb79e22ca 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -9,6 +9,7 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType from lightllm.utils.torch_memory_saver_utils import ( TorchMemorySaverWrapper, MemoryTag, @@ -117,7 +118,9 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo): # 记录原始存在的变量 pure_para_set = set(vars(infer_state).keys()) torch.cuda.synchronize() - decode_func(copy.copy(infer_state)) + # 在正式捕获前调优 decode attention，退出作用域后再捕获选定的配置。 + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + decode_func(copy.copy(infer_state)) torch.cuda.synchronize() for param_name in set(vars(infer_state).keys()): if param_name not in pure_para_set: @@ -149,7 +152,8 @@ def _capture_decode_overlap( pure_para_set = set(vars(infer_state).keys()) pure_para_set1 = set(vars(infer_state1).keys()) torch.cuda.synchronize() - decode_func(copy.copy(infer_state), copy.copy(infer_state1)) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + decode_func(copy.copy(infer_state), copy.copy(infer_state1)) torch.cuda.synchronize() for para_name in set(vars(infer_state).keys()): if para_name not in pure_para_set: diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index a5ba656c9c..024be9f55c 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -15,7 +15,7 @@ quantize_fused_experts_input, ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd -from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair @@ -250,7 +250,7 @@ def prefilled_group_gemm( # A rank may receive no tokens during autotune warmup. Run one dummy token through # silu_and_mul_fwd so the empty rank matches the first kernel call made by non-empty ranks. # This branch does not synchronize additional calls caused by different positive chunk counts. - if Autotuner.is_autotune_warmup(): + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): N = w13_weight.shape[1] _gemm_out_a = torch.zeros((1, N), device=recv_x[0].device, dtype=hidden_dtype) _silu_out = torch.zeros((1, N // 2), device=recv_x[0].device, dtype=hidden_dtype) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py index 59a7d4f742..d121656c73 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py @@ -6,6 +6,7 @@ def gqa_token_decode_attention_flash_decoding( infer_state, cache_k: torch.Tensor, cache_v: torch.Tensor, + max_len_in_batch: int, out=None, alloc_tensor_func=torch.empty, sliding_window=(-1, -1), @@ -41,7 +42,7 @@ def gqa_token_decode_attention_flash_decoding( Req_to_tokens=infer_state.req_manager.req_to_token_indexs, B_req_idx=infer_state.b_req_idx, B_Seqlen=infer_state.b_seq_len, - max_len_in_batch=infer_state.max_kv_seq_len, + max_len_in_batch=max_len_in_batch, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, block_seq=BLOCK_SEQ, diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py index cae913b4bc..767cee9870 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py @@ -2,7 +2,8 @@ import triton import triton.language as tl from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @triton.jit @@ -138,11 +139,12 @@ def get_test_configs(): return configs -def get_static_key(q, k, block_seq): +def get_static_key(q, k, block_seq, sliding_window): key_params = { "gqa_group_size": int(q.shape[1] // k.shape[1]), "q_head_dim": int(q.shape[2]), "block_seq": block_seq, + "sliding_window": tuple(sliding_window), "out_dtype": str(q.dtype), } return key_params @@ -150,15 +152,79 @@ def get_static_key(q, k, block_seq): def get_run_key(q, max_len_in_batch): batch_size = q.shape[0] - return batch_size * 1000 * 1000 * 1000 + max_len_in_batch + # 正常执行使用调用方在 CPU 上保存的真实 KV 长度，不读取 GPU 长度张量或 Graph 的容量上限。 + max_kv_len = int(max_len_in_batch) + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) and get_triton_autotune_level() in [ + AutotuneLevel.ADAPTIVE_AUTOTUNE, + AutotuneLevel.FORCE_AUTOTUNE, + ]: + max_kv_len = get_decode_attn_autotune_seq_len() + # 调优和正常查找统一按 512 token 向上分桶，同一区间复用配置匹配结果。 + max_kv_len = (max_kv_len + 511) // 512 * 512 + return batch_size * 1000 * 1000 * 1000 + max_kv_len + + +def rebuild_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + Req_to_tokens: torch.Tensor, + B_req_idx: torch.Tensor, + B_Seqlen: torch.Tensor, + max_len_in_batch: int, + mid_out: torch.Tensor, + mid_out_logsumexp: torch.Tensor, + block_seq: int, + sliding_window=(-1, -1), + **kwargs, +): + # Graph 初始化时真实请求很短，Req_to_tokens 的宽度则是容量上限，都不代表期望调优的长度。 + # 仅在实际搜索配置前重建一次输入，构造开销不计入 benchmark；正常执行和 Graph 捕获使用原输入。 + batch_size = q.shape[0] + # 与调优时的 run key 共用该环境变量，默认 32768 token；实际计算保持精确长度，不做 512 分桶。 + max_len_in_batch = get_decode_attn_autotune_seq_len() + assert k.shape[0] == v.shape[0], "K/V caches must have the same number of tokens" + num_tokens = k.shape[0] + if num_tokens == 0: + raise ValueError("GQA decode autotuning requires a non-empty KV cache") + + # 新建每个请求到物理 token 的映射，不能直接扩展 B_Seqlen 后读取原映射中未初始化的条目。 + # 物理 token 充足时各请求使用不同位置，不足时取模循环复用，保证所有索引均落在 K/V 缓存内。 + # 复用已有 K/V 可以避免分配完整的长请求缓存，但可能提高 GPU 缓存命中率，影响调优的访存特征。 + Req_to_tokens = torch.arange(batch_size * max_len_in_batch, dtype=Req_to_tokens.dtype, device=Req_to_tokens.device) + Req_to_tokens = Req_to_tokens.remainder_(num_tokens).view(batch_size, max_len_in_batch) + # 新映射只有 batch_size 行，请求索引也必须重建，避免继续使用原全局请求表中的行号。 + B_req_idx = torch.arange(batch_size, dtype=B_req_idx.dtype, device=B_req_idx.device) + B_Seqlen = torch.full_like(B_Seqlen, max_len_in_batch) + + # 保留 Q、滑窗语义、BLOCK_SEQ 和中间缓冲区布局；一个 program 可循环处理多个 KV 块， + # 无需按调优长度扩容 mid_out。调优结束后 stage1 使用原始输入重新覆盖有效中间块， + # stage2 仍按相同 BLOCK_SEQ 和缓冲区中的 block_num 归约。 + return ( + q, + k, + v, + Req_to_tokens, + B_req_idx, + B_Seqlen, + max_len_in_batch, + mid_out, + mid_out_logsumexp, + block_seq, + sliding_window, + ), kwargs @autotune( - kernel_name="_fwd_kernel_gqa_flash_decode_stage1:v3", + kernel_name="_fwd_kernel_gqa_flash_decode_stage1:v4", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, configs_gen_func=get_test_configs, static_key_func=get_static_key, run_key_func=get_run_key, - mutates_args=["mid_out", "mid_out_logsumexp"], + rebuild_input_func=rebuild_inputs, + # stage1 对有效中间块执行覆盖写，候选配置不会读取已有输出，正式执行也会重新覆盖真实请求的有效块。 + # 不标记这两个大缓冲区，避免每个候选配置 benchmark 时反复 clone，增加显存峰值和拷贝开销。 + # mutates_args=["mid_out", "mid_out_logsumexp"], ) @torch.no_grad() def flash_decode_stage1( @@ -246,8 +312,6 @@ def flash_decode_stage1( if __name__ == "__main__": - from lightllm.utils.envs_utils import get_triton_autotune_level - if get_triton_autotune_level() != 2: raise Exception("you need set env LIGHTLLM_TRITON_AUTOTUNE_LEVEL=2 to start program.") @@ -258,11 +322,11 @@ def flash_decode_stage1( out_dtype = torch.bfloat16 batch_sizes = [1, 8, 16, 32, 64, 128] - decode_lengths = [1024, 2048, 8192, 16384] + decode_lengths = [get_decode_attn_autotune_seq_len()] q_head_num = gqa_group_size - Autotuner.start_autotune_warmup() + Autotuner.start_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) # autotuing kernel for batch_size in batch_sizes: for length in decode_lengths: diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py index 101a285e9d..13c7e8070e 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py @@ -13,6 +13,7 @@ def token_decode_attention_mtp_diverse_single_token( B_req_idx, b_seq_len, b_mark_shared_group, + max_kv_len: int, out=None, alloc_tensor_func=torch.empty, ): @@ -25,8 +26,6 @@ def token_decode_attention_mtp_diverse_single_token( else: o_tensor = out - max_kv_len = Req_to_tokens.shape[1] - if batch_size <= 16: block_num = 128 elif batch_size <= 64: @@ -57,5 +56,6 @@ def token_decode_attention_mtp_diverse_single_token( B_Seqlen=b_seq_len, out=o_tensor, block_n=BLOCK_N, + max_kv_len=max_kv_len, ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py index 9765d60aaf..d90aeb5bff 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py @@ -16,8 +16,9 @@ import triton import triton.language as tl from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel from lightllm.utils.device_utils import is_hopper +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level def get_test_configs(): @@ -50,9 +51,96 @@ def get_static_key(q, k, block_batch): def get_run_key(q, max_kv_len): batch_size = q.shape[0] + # 正常查找使用调用方在 Graph 改写长度上限前保存的真实 KV 长度，不读取 GPU 张量或请求表容量。 + max_kv_len = int(max_kv_len) + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) and get_triton_autotune_level() in [ + AutotuneLevel.ADAPTIVE_AUTOTUNE, + AutotuneLevel.FORCE_AUTOTUNE, + ]: + max_kv_len = get_decode_attn_autotune_seq_len() + # 调优和正常查找统一按 512 token 向上分桶；实际 benchmark 保留环境变量指定的精确长度。 + max_kv_len = (max_kv_len + 511) // 512 * 512 return batch_size * 1000 * 1000 * 1000 + max_kv_len +def rebuild_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + Req_to_tokens: torch.Tensor, + B_req_idx: torch.Tensor, + b_seq_len: torch.Tensor, + b_mark_shared_group: torch.Tensor, + max_kv_len: int, + mid_out: torch.Tensor, + mid_out_logsumexp: torch.Tensor, + block_batch: int, + **kwargs, +): + # Graph 初始化输入只有很短的 HOLD 请求，且每行独立成组，不能代表 MTP 的长 KV 共享计算。 + # 仅在实际搜索前构造一次调优输入，开销不计入 benchmark；正式执行和捕获仍使用原始输入。 + batch_size = q.shape[0] + max_kv_len = get_decode_attn_autotune_seq_len() + assert k.shape[0] == v.shape[0], "K/V caches must have the same number of tokens" + num_tokens = k.shape[0] + if num_tokens == 0: + raise ValueError("MTP decode autotuning requires a non-empty KV cache") + assert block_batch > 0, "block_batch must be positive" + + # 以 block_batch 为代表性共享组大小，尾组允许不足；KV 长度不改变分组方式。 + # 调优使用固定代表性分组，不从 HOLD 标记推断真实组大小，也不改动正常请求的动态分组。 + group_size = block_batch + # 组内长度递增且组末长度等于目标长度，目标长度必须能保证组首至少有一个可见 KV。 + if max_kv_len < min(batch_size, group_size): + raise ValueError("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN must be at least the largest MTP group size") + num_groups = (batch_size + group_size - 1) // group_size + + # 按每组实际行数生成列表，再构造 CPU tensor 并转回原设备；尾组可能不足 group_size。 + group_sizes = [min(group_size, batch_size - start) for start in range(0, batch_size, group_size)] + cpu_req_idx = torch.tensor( + [group_idx for group_idx, size in enumerate(group_sizes) for _ in range(size)], + dtype=B_req_idx.dtype, + device="cpu", + ) + # 同组请求共享同一行 KV 映射；例如目标长度 16384、组大小 3，长度为 [16382, 16383, 16384]。 + cpu_seq_len = torch.tensor( + [length for size in group_sizes for length in range(max_kv_len - size + 1, max_kv_len + 1)], + dtype=b_seq_len.dtype, + device="cpu", + ) + # 只有组末行标记组大小并启动计算；其他行保持 0，由组末行一次处理整组 Q。 + cpu_mark_shared_group = torch.tensor( + [size if offset == size - 1 else 0 for size in group_sizes for offset in range(size)], + dtype=b_mark_shared_group.dtype, + device="cpu", + ) + B_req_idx = cpu_req_idx.to(device=B_req_idx.device) + b_seq_len = cpu_seq_len.to(device=b_seq_len.device) + b_mark_shared_group = cpu_mark_shared_group.to(device=b_mark_shared_group.device) + + # 每组只需要一行已初始化的映射，不能扩展长度后读取原请求表中的未初始化条目。 + # 取模将所有物理索引限制在已有 K/V 池内，避免为长请求重新分配完整缓存；物理容量不足时， + # 不同位置会复用 K/V，可能提高 GPU 缓存命中率，因此调优的访存特征仍受现有缓存容量影响。 + Req_to_tokens = torch.arange(num_groups * max_kv_len, dtype=Req_to_tokens.dtype, device=Req_to_tokens.device) + Req_to_tokens = Req_to_tokens.remainder_(num_tokens).view(num_groups, max_kv_len) + + # 保留 Q/K/V、block_batch 和中间缓冲区布局；stage1 的每个 program 可循环处理多个 KV 块。 + # 选好配置后使用原始输入重新覆盖有效中间块，并返回对应的 BLOCK_N 供 stage2 归约。 + return ( + q, + k, + v, + Req_to_tokens, + B_req_idx, + b_seq_len, + b_mark_shared_group, + max_kv_len, + mid_out, + mid_out_logsumexp, + block_batch, + ), kwargs + + @triton.jit def _fwd_kernel_mtp_diverse_stage1_single_token( Q, @@ -174,11 +262,15 @@ def _fwd_kernel_mtp_diverse_stage1_single_token( @autotune( - kernel_name="_fwd_kernel_mtp_diverse_stage1_single_token:v2", + kernel_name="_fwd_kernel_mtp_diverse_stage1_single_token:v3", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, configs_gen_func=get_test_configs, static_key_func=get_static_key, run_key_func=get_run_key, - mutates_args=["mid_out", "mid_out_logsumexp"], + rebuild_input_func=rebuild_inputs, + # stage1 对有效中间块执行覆盖写，候选配置不会读取已有输出，正式执行也会重新覆盖真实请求的有效块。 + # 不标记这两个大缓冲区，避免每个候选配置 benchmark 时反复 clone，增加显存峰值和拷贝开销。 + # mutates_args=["mid_out", "mid_out_logsumexp"], ) def mtp_diverse_stage1_single_token( q: torch.Tensor, @@ -272,8 +364,6 @@ def mtp_diverse_stage1_single_token( if __name__ == "__main__": - from lightllm.utils.envs_utils import get_triton_autotune_level - if get_triton_autotune_level() != 2: raise Exception("you need set env LIGHTLLM_TRITON_AUTOTUNE_LEVEL=2 to start program.") @@ -283,7 +373,7 @@ def mtp_diverse_stage1_single_token( out_dtype = torch.bfloat16 batch_sizes = [1, 8, 16, 32, 64, 128] - decode_lengths = [32, 64, 128, 256, 512, 1024, 2048] + decode_lengths = [get_decode_attn_autotune_seq_len()] tp_world_size = 2 q_head_num = 64 // tp_world_size @@ -291,7 +381,7 @@ def mtp_diverse_stage1_single_token( gqa_group_size = q_head_num // k_head_num - Autotuner.start_autotune_warmup() + Autotuner.start_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) # autotuing kernel for batch_size in batch_sizes: for length in decode_lengths: diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py index 4c827ae5d9..b89b34ad14 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py @@ -10,7 +10,8 @@ import triton import triton.language as tl from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level def get_test_configs(): @@ -35,12 +36,42 @@ def get_static_key(mid_out, block_n, out): return key_params -def get_run_key(mid_out): +def get_run_key(mid_out, block_n, max_kv_len): batch_size_head = mid_out.shape[0] * mid_out.shape[1] - block_num = mid_out.shape[2] + # 正常查找使用外部保存的真实 KV 长度；只有 decode 调优时才使用重建输入的目标长度。 + max_kv_len = int(max_kv_len) + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) and get_triton_autotune_level() in [ + AutotuneLevel.ADAPTIVE_AUTOTUNE, + AutotuneLevel.FORCE_AUTOTUNE, + ]: + max_kv_len = get_decode_attn_autotune_seq_len() + # stage2 只归约 stage1 输出的有效分块，工作量达到缓冲区分块数后不再随 KV 长度增长。 + # 直接以有效分块数区分配置，不再按 512 token 分桶，以免短请求的 key 与实际归约次数不符。 + block_num = min(triton.cdiv(max_kv_len, block_n), mid_out.shape[2]) return batch_size_head * 1000 * 1000 * 1000 + block_num +def rebuild_inputs( + mid_out: torch.Tensor, + mid_out_logsumexp: torch.Tensor, + B_Seqlen: torch.Tensor, + out: torch.Tensor, + block_n: int, + max_kv_len: int, + **kwargs, +): + # 仅在实际搜索前构造一次输入，不计入 benchmark；正式执行和 Graph 捕获仍使用原始中间结果。 + max_kv_len = get_decode_attn_autotune_seq_len() + # stage2 按行独立归约，不需要共享组标记或 KV 映射，以目标长度代表各行的归约工作量。 + B_Seqlen = torch.full_like(B_Seqlen, max_kv_len, device="cpu").to(device=B_Seqlen.device) + + # mid_out 和 mid_out_logsumexp 是 stage1 产出的只读输入，stage2 不会修改它们。 + # 两个张量的 shape 已按最大分块数分配，调优时直接复用原始输入，避免无意义的重复分配。 + # block_n 必须沿用 stage1 实际返回的 BLOCK_N，它决定每行需要归约的分块数，不能独立改写。 + # benchmark 的最终输出写入继续由 mutates_args 隔离，原始 out 不会被调优覆盖。 + return (mid_out, mid_out_logsumexp, B_Seqlen, out, block_n, max_kv_len), kwargs + + @triton.jit def _fwd_kernel_mtp_diverse_stage2_single_token( B_Seqlen, @@ -108,10 +139,12 @@ def _fwd_kernel_mtp_diverse_stage2_single_token( @autotune( - kernel_name="_fwd_kernel_mtp_diverse_stage2_single_token:v2", + kernel_name="_fwd_kernel_mtp_diverse_stage2_single_token:v3", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, configs_gen_func=get_test_configs, static_key_func=get_static_key, run_key_func=get_run_key, + rebuild_input_func=rebuild_inputs, mutates_args=["out"], ) @torch.no_grad() @@ -121,6 +154,7 @@ def mtp_diverse_stage2_single_token( B_Seqlen: torch.Tensor, out: torch.Tensor, block_n: int, + max_kv_len: int, run_config: Optional[dict] = None, ): if not run_config: @@ -160,8 +194,6 @@ def mtp_diverse_stage2_single_token( if __name__ == "__main__": - from lightllm.utils.envs_utils import get_triton_autotune_level - if get_triton_autotune_level() != 2: raise Exception("you need set env LIGHTLLM_TRITON_AUTOTUNE_LEVEL=2 to start program.") @@ -171,8 +203,9 @@ def mtp_diverse_stage2_single_token( batch_sizes = [1, 8, 16, 32, 64, 128] q_head_num = 64 // tp_world_size + max_kv_len = get_decode_attn_autotune_seq_len() - Autotuner.start_autotune_warmup() + Autotuner.start_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) # autotuing kernel for batch_size in batch_sizes: for block_n in [16, 32, 64, 128]: @@ -188,7 +221,7 @@ def mtp_diverse_stage2_single_token( batch_size, q_head_num, block_num, q_head_dim, dtype=torch.bfloat16, device="cuda" ) mid_out_logsumexp = torch.randn(batch_size, q_head_num, block_num, dtype=torch.float32, device="cuda") - B_Seqlen = torch.full((batch_size,), 8196, dtype=torch.int32, device="cuda") + B_Seqlen = torch.full((batch_size,), max_kv_len, dtype=torch.int32, device="cuda") out = torch.zeros(batch_size, q_head_num, q_head_dim, dtype=out_dtype, device="cuda") mtp_diverse_stage2_single_token( @@ -197,6 +230,7 @@ def mtp_diverse_stage2_single_token( B_Seqlen=B_Seqlen, out=out, block_n=block_n, + max_kv_len=max_kv_len, ) Autotuner.end_autotune_warmup() diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py index 00e884352b..27cb97d9a0 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py @@ -2,7 +2,8 @@ import triton import triton.language as tl from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @triton.jit @@ -163,15 +164,84 @@ def get_static_key(q, k, k_scale, block_seq): def get_run_key(q, max_kv_seq_len): batch_size = q.shape[0] + # 正常查找使用 Graph 改写长度上限前保存的真实 KV 长度，不读取 GPU 张量或请求表容量。 + max_kv_seq_len = int(max_kv_seq_len) + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) and get_triton_autotune_level() in [ + AutotuneLevel.ADAPTIVE_AUTOTUNE, + AutotuneLevel.FORCE_AUTOTUNE, + ]: + max_kv_seq_len = get_decode_attn_autotune_seq_len() + # 调优和正常查找统一按 512 token 向上分桶；实际 benchmark 保留环境变量指定的精确长度。 + max_kv_seq_len = (max_kv_seq_len + 511) // 512 * 512 return batch_size * 1000 * 1000 * 1000 + max_kv_seq_len +def rebuild_inputs( + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + Req_to_tokens: torch.Tensor, + B_req_idx: torch.Tensor, + B_Seqlen: torch.Tensor, + max_kv_seq_len: int, + mid_out: torch.Tensor, + mid_out_logsumexp: torch.Tensor, + block_seq: int, + **kwargs, +): + # Graph 初始化请求很短，请求表宽度又是容量上限，两者都不能代表实际 decode 的计算量。 + # 仅在实际搜索前重建一次输入，不计入 benchmark；正式执行和 Graph 捕获仍使用原始输入。 + batch_size = q.shape[0] + max_kv_seq_len = get_decode_attn_autotune_seq_len() + # int4 数据和 scale 按同一物理 token 索引访问，四个缓存的 token 容量必须一致。 + assert ( + k.shape[0] == v.shape[0] == k_scale.shape[0] == v_scale.shape[0] + ), "K/V caches and scales must have the same number of tokens" + num_tokens = k.shape[0] + if num_tokens == 0: + raise ValueError("Int4 KV decode autotuning requires a non-empty KV cache") + + # 新建已初始化的映射，不能只拉长 B_Seqlen 后访问原请求表中未填充的位置。 + # 物理容量充足时各请求使用不同 token，不足时取模复用已有 K/V 和 scale，保证索引不越界。 + # 复用可避免重新分配完整的量化缓存，但可能提高 GPU 缓存命中率，影响调优的访存特征。 + Req_to_tokens = torch.arange(batch_size * max_kv_seq_len, dtype=Req_to_tokens.dtype, device=Req_to_tokens.device) + Req_to_tokens = Req_to_tokens.remainder_(num_tokens).view(batch_size, max_kv_seq_len) + # 新表只有 batch_size 行，因此请求索引也需重建，不能继续使用原全局请求表的行号。 + B_req_idx = torch.arange(batch_size, dtype=B_req_idx.dtype, device=B_req_idx.device) + B_Seqlen = torch.full_like(B_Seqlen, max_kv_seq_len) + + # 保留打包 int4 的 K/V、分组 scale 和原有 stride，不解量化或改变它们的存储布局。 + # 一个 program 可循环处理多个 BLOCK_SEQ，无需按目标长度扩容中间缓冲区。 + # BLOCK_N 只控制块内计算大小，stage2 仍按固定 BLOCK_SEQ 和原缓冲区分块数归约。 + # 调优后的正式执行会使用原始输入重新覆盖真实请求对应的有效中间块。 + return ( + q, + k, + k_scale, + v, + v_scale, + Req_to_tokens, + B_req_idx, + B_Seqlen, + max_kv_seq_len, + mid_out, + mid_out_logsumexp, + block_seq, + ), kwargs + + @autotune( - kernel_name="_fwd_kernel_flash_decode_stage1:v1", + kernel_name="_fwd_kernel_flash_decode_stage1:v2", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, configs_gen_func=get_test_configs, static_key_func=get_static_key, run_key_func=get_run_key, - mutates_args=["mid_out", "mid_out_logsumexp"], + rebuild_input_func=rebuild_inputs, + # stage1 对有效中间块执行覆盖写，候选配置不会读取已有输出，正式执行也会重新覆盖真实请求的有效块。 + # 不标记这两个大缓冲区，避免每个候选配置 benchmark 时反复 clone，增加显存峰值和拷贝开销。 + # mutates_args=["mid_out", "mid_out_logsumexp"], ) def int4kv_flash_decode_stage1( q, @@ -261,8 +331,6 @@ def int4kv_flash_decode_stage1( if __name__ == "__main__": - from lightllm.utils.envs_utils import get_triton_autotune_level - if get_triton_autotune_level() != 2: raise Exception("you need set env LIGHTLLM_TRITON_AUTOTUNE_LEVEL=2 to start program.") @@ -274,11 +342,11 @@ def int4kv_flash_decode_stage1( out_dtype = torch.bfloat16 batch_sizes = [1, 8, 16, 32, 64, 128] - decode_lengths = [1024, 2048, 8192, 16384] + decode_lengths = [get_decode_attn_autotune_seq_len()] q_head_num = gqa_group_size - Autotuner.start_autotune_warmup() + Autotuner.start_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) # autotuing kernel for batch_size in batch_sizes: for length in decode_lengths: diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py index 9521364ba6..936ca4dbbd 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py @@ -9,6 +9,7 @@ def token_decode_attention_flash_decoding( cache_k_scale: torch.Tensor, cache_v: torch.Tensor, cache_v_scale: torch.Tensor, + max_kv_seq_len: int, out: Optional[torch.Tensor] = None, alloc_tensor_func=torch.empty, ): @@ -44,7 +45,7 @@ def token_decode_attention_flash_decoding( Req_to_tokens=infer_state.req_manager.req_to_token_indexs, B_req_idx=infer_state.b_req_idx, B_Seqlen=infer_state.b_seq_len, - max_kv_seq_len=infer_state.max_kv_seq_len, + max_kv_seq_len=max_kv_seq_len, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, block_seq=BLOCK_SEQ, diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py index b61e8eace2..de74f0728d 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py @@ -13,6 +13,7 @@ def token_decode_attention_flash_decoding( cache_k_scale: torch.Tensor, cache_v: torch.Tensor, cache_v_scale: torch.Tensor, + max_len_in_batch: int, out: Optional[torch.Tensor] = None, alloc_tensor_func=torch.empty, ): @@ -48,7 +49,7 @@ def token_decode_attention_flash_decoding( Req_to_tokens=infer_state.req_manager.req_to_token_indexs, B_req_idx=infer_state.b_req_idx, B_seq_len=infer_state.b_seq_len, - max_len_in_batch=infer_state.max_kv_seq_len, + max_len_in_batch=max_len_in_batch, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, block_seq=BLOCK_SEQ, diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py index d241d31c07..a05db282d7 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py @@ -2,7 +2,8 @@ import triton import triton.language as tl from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @triton.jit @@ -169,15 +170,84 @@ def get_static_key(q, k, k_scale, block_seq): def get_run_key(q, max_len_in_batch): batch_size = q.shape[0] + # 正常查找使用 Graph 改写长度上限前保存的真实 KV 长度，不读取 GPU 张量或请求表容量。 + max_len_in_batch = int(max_len_in_batch) + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) and get_triton_autotune_level() in [ + AutotuneLevel.ADAPTIVE_AUTOTUNE, + AutotuneLevel.FORCE_AUTOTUNE, + ]: + max_len_in_batch = get_decode_attn_autotune_seq_len() + # 调优和正常查找统一按 512 token 向上分桶；实际 benchmark 保留环境变量指定的精确长度。 + max_len_in_batch = (max_len_in_batch + 511) // 512 * 512 return batch_size * 1000 * 1000 * 1000 + max_len_in_batch +def rebuild_inputs( + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + Req_to_tokens: torch.Tensor, + B_req_idx: torch.Tensor, + B_seq_len: torch.Tensor, + max_len_in_batch: int, + mid_out: torch.Tensor, + mid_out_logsumexp: torch.Tensor, + block_seq: int, + **kwargs, +): + # Graph 初始化请求很短，请求表宽度又是容量上限，两者都不能代表实际 decode 的计算量。 + # 仅在实际搜索前重建一次输入，不计入 benchmark；正式执行和 Graph 捕获仍使用原始输入。 + batch_size = q.shape[0] + max_len_in_batch = get_decode_attn_autotune_seq_len() + # int8 数据和 scale 按同一物理 token 索引访问，四个缓存的 token 容量必须一致。 + assert ( + k.shape[0] == v.shape[0] == k_scale.shape[0] == v_scale.shape[0] + ), "K/V caches and scales must have the same number of tokens" + num_tokens = k.shape[0] + if num_tokens == 0: + raise ValueError("Int8 KV decode autotuning requires a non-empty KV cache") + + # 新建已初始化的映射，不能只拉长 B_seq_len 后访问原请求表中未填充的位置。 + # 物理容量充足时各请求使用不同 token，不足时取模复用已有 K/V 和 scale，保证索引不越界。 + # 复用可避免重新分配完整的量化缓存，但可能提高 GPU 缓存命中率，影响调优的访存特征。 + Req_to_tokens = torch.arange(batch_size * max_len_in_batch, dtype=Req_to_tokens.dtype, device=Req_to_tokens.device) + Req_to_tokens = Req_to_tokens.remainder_(num_tokens).view(batch_size, max_len_in_batch) + # 新表只有 batch_size 行，因此请求索引也需重建，不能继续使用原全局请求表的行号。 + B_req_idx = torch.arange(batch_size, dtype=B_req_idx.dtype, device=B_req_idx.device) + B_seq_len = torch.full_like(B_seq_len, max_len_in_batch) + + # 保留 int8 K/V、分组 scale 和原有 stride，不解量化或改变它们的存储布局。 + # 一个 program 可循环处理多个 BLOCK_SEQ，无需按目标长度扩容中间缓冲区。 + # BLOCK_N 只控制块内计算大小，stage2 仍按固定 BLOCK_SEQ 和原缓冲区分块数归约。 + # 调优后的正式执行会使用原始输入重新覆盖真实请求对应的有效中间块。 + return ( + q, + k, + k_scale, + v, + v_scale, + Req_to_tokens, + B_req_idx, + B_seq_len, + max_len_in_batch, + mid_out, + mid_out_logsumexp, + block_seq, + ), kwargs + + @autotune( - kernel_name="_fwd_kernel_flash_decode_normal_stage1:v3", + kernel_name="_fwd_kernel_flash_decode_normal_stage1:v4", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, configs_gen_func=get_test_configs, static_key_func=get_static_key, run_key_func=get_run_key, - mutates_args=["mid_out", "mid_out_logsumexp"], + rebuild_input_func=rebuild_inputs, + # stage1 对有效中间块执行覆盖写，候选配置不会读取已有输出，正式执行也会重新覆盖真实请求的有效块。 + # 不标记这两个大缓冲区，避免每个候选配置 benchmark 时反复 clone，增加显存峰值和拷贝开销。 + # mutates_args=["mid_out", "mid_out_logsumexp"], ) def flash_decode_stage1( q: torch.Tensor, @@ -276,8 +346,6 @@ def flash_decode_stage1( if __name__ == "__main__": - from lightllm.utils.envs_utils import get_triton_autotune_level - if get_triton_autotune_level() != 2: raise Exception("you need set env LIGHTLLM_TRITON_AUTOTUNE_LEVEL=2 to start program.") @@ -289,11 +357,11 @@ def flash_decode_stage1( out_dtype = torch.bfloat16 batch_sizes = [1, 8, 16, 32, 64, 128] - decode_lengths = [1024, 2048, 8192, 16384] + decode_lengths = [get_decode_attn_autotune_seq_len()] q_head_num = gqa_group_size - Autotuner.start_autotune_warmup() + Autotuner.start_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) # autotuing kernel for batch_size in batch_sizes: for length in decode_lengths: diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 58d4d45514..ca39376bab 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -23,7 +23,7 @@ get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, ) -from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType from lightllm.utils.device_utils import is_sm100_gpu from lightllm.utils.sgl_utils import HAS_SGL_KERNEL from lightllm.utils.tensor_buffer_manager import TensorBufferManager @@ -320,7 +320,7 @@ def fused_experts_impl( # A rank may receive no tokens during autotune warmup. Run one dummy token through # silu_and_mul_fwd so the empty rank matches the first kernel call made by non-empty ranks. # This branch does not synchronize additional calls caused by different positive chunk counts. - if Autotuner.is_autotune_warmup(): + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): N = w1.shape[1] _gemm_out_a = torch.zeros((1, N), device=hidden_states.device, dtype=hidden_states.dtype) _silu_out = torch.zeros((1, N // 2), device=hidden_states.device, dtype=hidden_states.dtype) @@ -519,7 +519,7 @@ def chunked_expanded_moe_forward( # 中的分布式通信要求各 rank 进入 autotuning 的次数一致，否则容易发生通信错位。 # 所以只允许第一个 chunk 保持 autotuning；从第二个 chunk 开始临时关闭，循环结束 # 后再恢复进入函数时的 warmup 状态。零 token rank 的首次调用由外层特殊分支补齐。 - is_autotune_warmup = Autotuner.is_autotune_warmup() + is_autotune_warmup = Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) try: for chunk_index, chunk_start in enumerate(range(0, all_tokens, max_chunk_rows)): if is_autotune_warmup and chunk_index == 1: @@ -574,7 +574,7 @@ def workspace_quant_alloc(shape, dtype, device): workspace_manager.free(gemm_out_b) finally: if is_autotune_warmup: - Autotuner.start_autotune_warmup() + Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) ep_compact_metadata(recv_src_metadata) return gather_out diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py index 1c01cbd638..87cda6ce16 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py @@ -21,7 +21,7 @@ from lightllm.utils.sgl_utils import sgl_ops from typing import Callable, List, Optional, Tuple from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk -from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType def fused_topk( @@ -170,7 +170,7 @@ def select_experts( ######################################## warning ################################################## # here is used to match autotune feature, make topk_ids more random - if Autotuner.is_autotune_warmup(): + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): rand_gen = torch.Generator(device="cuda") rand_gen.manual_seed(router_logits.shape[0]) router_logits = torch.randn(size=router_logits.shape, generator=rand_gen, dtype=torch.float32, device="cuda") diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py index 2eb5ef5333..6fcef6c882 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py @@ -14,6 +14,8 @@ import torch import triton import triton.language as tl +from typing import Optional +from lightllm.common.triton_utils.autotuner import autotune, AutotuneKernelType # --------------------------------------------------------------------------- @@ -44,6 +46,7 @@ def _fused_recurrent_gated_delta_rule_fwd_kernel( V: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + NUM_STAGES: tl.constexpr, stride_q_tok: tl.constexpr, # token stride in q/k/v/a/b stride_k_tok: tl.constexpr, stride_v_tok: tl.constexpr, @@ -94,7 +97,7 @@ def _fused_recurrent_gated_delta_rule_fwd_kernel( p_h0 = p_h0 + i_hv * stride_state_hv + o_k[:, None] * V + o_v[None, :] b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - for i_t in range(0, T): + for i_t in tl.range(0, T, num_stages=NUM_STAGES): b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) @@ -137,6 +140,119 @@ def _fused_recurrent_gated_delta_rule_fwd_kernel( # --------------------------------------------------------------------------- +def get_test_configs(): + # BK 必须覆盖整个 K 维归约，只搜索 V 维分块、warp 数和循环流水级数；保留默认 BV=8/warps=1/stages=1。 + return [ + {"BV": bv, "num_warps": num_warps, "num_stages": num_stages} + for bv in [8, 16, 32, 64, 128] + for num_warps in [1, 2, 4, 8] + for num_stages in [1, 2, 3, 4] + ] + + +def get_static_key(q, v, initial_state, ssm_state_write_indices): + return { + "head_k_dim": q.shape[-1], + "head_v_dim": v.shape[-1], + "num_v_heads": v.shape[2], + "gva_group_size": v.shape[2] // q.shape[2], + "mtp_size": ssm_state_write_indices.shape[1], + "q_dtype": str(q.dtype), + "state_dtype": str(initial_state.dtype), + } + + +def get_run_key(q, ssm_state_write_indices): + # 线性 attention 的历史信息已压缩到固定大小的 SSM 状态，计算量不随历史 KV 长度增长。 + # head 数由 static key 区分，run key 只按本轮 Q token 数分桶，同桶内不同 MTP 分组共用配置。 + # 只读取形状，不读取 GPU 内容，保证正常查找可用于 Graph 捕获。 + num_q_tokens = q.shape[1] + # 按 4 个 MTP 请求的 token 宽度向上分桶，使相近长度复用配置。 + mtp_size = ssm_state_write_indices.shape[1] + token_bucket_size = 4 * mtp_size + return triton.cdiv(num_q_tokens, token_bucket_size) * token_bucket_size + + +def rebuild_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + ssm_state_write_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + a_raw: torch.Tensor, + b_raw: torch.Tensor, + **kwargs, +): + # Graph 初始化使用 HOLD 请求：固定 MTP 可能重复写同一状态槽，动态 MTP 则全部是空序列。 + # 直接计时会产生状态竞争或只测到提前返回，因此搜索前重建合法的 MTP 分组和互不重叠的状态索引。 + # 保留 Q/K/V、gate 的形状和 stride，以及 cu_seqlens 的长度，使 benchmark 与 run key 一致。 + num_tokens = q.shape[1] + num_seqs = cu_seqlens.shape[0] - 1 + mtp_size = ssm_state_write_indices.shape[1] + assert num_tokens > 0 and mtp_size > 0, "MTP autotuning requires non-empty tokens and state index rows" + active_seqs = triton.cdiv(num_tokens, mtp_size) + assert active_seqs <= num_seqs, "Not enough sequence rows for the MTP tokens" + num_state_slots = initial_state.shape[0] + # 有效 token 必须拥有互不重叠的状态槽，避免取余后不同请求并发读写同一槽位。 + assert num_tokens <= num_state_slots, "Not enough SSM state slots for autotuning" + + # 按完整 MTP 宽度分组，末组可更短，动态布局多余的行仍为空序列。 + # 例如 7 个 token、MTP 宽度 3、7 行序列，累积长度为 [0, 3, 6, 7, 7, 7, 7, 7]。 + cu_seqlens = torch.tensor( + [min(row * mtp_size, num_tokens) for row in range(num_seqs + 1)], + dtype=cu_seqlens.dtype, + device=cu_seqlens.device, + ) + # 按状态池总槽数取余，使读写索引都落在 [0, num_state_slots) 内。 + # 容量检查保证有效 token 的槽号不变且互不重叠；仅空序列及末组未使用的索引可能回绕。 + state_indices = [ + [(row * mtp_size + offset) % num_state_slots for offset in range(mtp_size)] for row in range(num_seqs) + ] + ssm_state_indices = torch.tensor(state_indices, dtype=ssm_state_indices.dtype, device=ssm_state_indices.device) + ssm_state_write_indices = ssm_state_indices.to(dtype=ssm_state_write_indices.dtype) + num_accepted_tokens = torch.ones_like(num_accepted_tokens) + + # 直接复用 initial_state，不新建或 clone 状态池，避免服务显存紧张时因额外分配触发 OOM。 + # benchmark 会原地覆盖池内前 num_tokens 个槽且不恢复；只能在启动阶段使用， + # 此时这些槽中不能有需要保留的请求状态。这里只重建分组和索引，不初始化状态数值。 + return ( + q, + k, + v, + initial_state, + cu_seqlens, + ssm_state_indices, + ssm_state_write_indices, + num_accepted_tokens, + A_log, + dt_bias, + a_raw, + b_raw, + ), kwargs + + +@autotune( + kernel_name="_mtp_fused_recurrent_gated_delta_rule_fwd_kernel:v1", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=get_test_configs, + static_key_func=get_static_key, + run_key_func=get_run_key, + rebuild_input_func=rebuild_inputs, + # 历史配置预热直接使用调用方的状态池，每额外执行一次都会推进 SSM 状态，影响正式计算。 + # 状态池很大，不适合通过 clone 保存/恢复，因此在所有阶段都关闭这类额外预热。 + warmup_all_exist_config=False, + # 不标记 initial_state 为可变参数：服务状态池很大，autotuner 在预热/调优时反复 clone + # 会额外占用大量显存，可能频繁触发 OOM，因此调优直接复用并原地更新传入的状态池。 + # 循环次数、访存地址和计算路径不依赖状态数值，理论上不影响调优对配置性能的选优。 + # 实际调优搜索仍会写入状态池，仅用于启动阶段尚无有效请求状态时。 + # mutates_args=["initial_state"], +) +@torch.no_grad() def mtp_fused_recurrent_gated_delta_rule( q: torch.Tensor, k: torch.Tensor, @@ -150,6 +266,7 @@ def mtp_fused_recurrent_gated_delta_rule( dt_bias: torch.Tensor, a_raw: torch.Tensor, b_raw: torch.Tensor, + run_config: Optional[dict] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Fused recurrent gated delta rule with fused gating (GDN layer). @@ -159,7 +276,7 @@ def mtp_fused_recurrent_gated_delta_rule( q: ``[1, T, H, K]`` queries. k: ``[1, T, H, K]`` keys. v: ``[1, T, HV, V]`` values (GVA when HV > H). - initial_state: ``[N, HV, K, V]`` initial SSM state. + initial_state: ``[num_slots, HV, K, V]`` SSM state pool, updated in-place. cu_seqlens: ``[N+1]`` int64 cumulative sequence lengths for the varlen (MTP verify) path. ssm_state_indices: ``[N, S+1]`` int32 slot indices (2D). @@ -170,10 +287,11 @@ def mtp_fused_recurrent_gated_delta_rule( dt_bias: ``[HV]`` per-head dt bias. a_raw: ``[T, HV]`` raw alpha. b_raw: ``[T, HV]`` raw beta. + run_config: Optional ``BV``, ``num_warps`` and ``num_stages`` configuration. Returns: ``(o, final_state)`` where ``o`` is ``[1, T, HV, V]`` and - ``final_state`` is ``[N, HV, K, V]``. + ``final_state`` is the same state pool as ``initial_state``. """ scale = k.shape[-1] ** -0.5 @@ -191,9 +309,11 @@ def mtp_fused_recurrent_gated_delta_rule( b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw) BK = triton.next_power_of_2(K) assert K == BK, f"K={K} must be a power of 2" - BV = min(triton.next_power_of_2(V), 8) - num_warps = 1 - num_stages = 3 + if run_config is None: + run_config = {"BV": 8, "num_warps": 1, "num_stages": 1} + BV = min(triton.next_power_of_2(V), run_config["BV"]) + num_warps = run_config["num_warps"] + num_stages = run_config["num_stages"] NV = triton.cdiv(V, BV) o = q.new_empty(v.shape) @@ -234,6 +354,7 @@ def mtp_fused_recurrent_gated_delta_rule( V=V, BK=BK, BV=BV, + NUM_STAGES=num_stages, stride_q_tok=stride_q_tok, stride_k_tok=stride_k_tok, stride_v_tok=stride_v_tok, @@ -250,7 +371,6 @@ def mtp_fused_recurrent_gated_delta_rule( SOFTPLUS_BETA=1.0, SOFTPLUS_THRESHOLD=20.0, num_warps=num_warps, - num_stages=num_stages, ) return o, final_state diff --git a/lightllm/common/triton_utils/autotuner.py b/lightllm/common/triton_utils/autotuner.py index 4cc6453d12..6752e73e5e 100644 --- a/lightllm/common/triton_utils/autotuner.py +++ b/lightllm/common/triton_utils/autotuner.py @@ -6,12 +6,14 @@ import torch.distributed as dist import random import collections +from contextlib import contextmanager +from enum import Enum from pathlib import Path from tqdm import tqdm from frozendict import frozendict from lightllm.utils.device_utils import get_current_device_name from lightllm.utils.log_utils import init_logger -from typing import Callable, List +from typing import Callable, List, Optional from lightllm.utils.envs_utils import get_triton_autotune_level from lightllm.common.kernel_config import KernelConfigs from lightllm.utils.dist_utils import get_global_world_size, get_global_rank, get_current_rank_in_node @@ -30,6 +32,12 @@ class AutotuneLevel: CLOSE_AUTOTUNE = 3 +class AutotuneKernelType(str, Enum): + GENERAL = "general" + # Includes full-attention and linear-attention decode kernels. + DECODE_ATTENTION = "decode_attention" + + def autotune( kernel_name: str, configs_gen_func: Callable[[], List], @@ -37,6 +45,9 @@ def autotune( run_key_func: Callable, run_key_distance_func: Callable = lambda run_key, config_key: abs(int(run_key) - int(config_key)), mutates_args: List[str] = [], + kernel_type: AutotuneKernelType = AutotuneKernelType.GENERAL, + rebuild_input_func: Optional[Callable] = None, + warmup_all_exist_config: bool = True, ): """Decorator that constructs and returns an Autotuner wrapper for a Triton kernel. @@ -56,6 +67,24 @@ def autotune( Defaults to ``abs(int(run_key) - int(config_key))``. mutates_args (List[str], optional): Names of arguments that can be mutated by the kernel. During benchmarking, defensive clones are made to avoid side effects. Defaults to ``[]``. + kernel_type (AutotuneKernelType, optional): Only a matching warmup phase benchmarks this kernel. + Other phases still execute it using cached configurations or its default configuration. + rebuild_input_func (Callable, optional): 调优输入重建回调，主要供 decode attention 算子使用。 + CUDA Graph 初始化时，输入的真实请求长度通常很短，无法代表实际 decode 场景的计算量， + 因此需要算子通过此回调自行重建输入，例如填入目标 KV 长度并构造对应的合法页表。 + 每次实际调优搜索前调用一次，接收算子的原始参数，返回用于计时的 ``(args, kwargs)``。 + 回调不应修改原始输入；缓存键、历史配置预热及最终执行仍使用原始参数。 + warmup_all_exist_config (bool, optional): 是否提前执行所有已有配置进行预热，默认 True。 + 设为 False 后，首次加载缓存和任何 warmup 阶段都不执行这一步，但仍正常加载、选择配置。 + 原地更新持久状态且无法低成本保存/恢复的算子应关闭，例如 MTP linear attention 的 + SSM 递推、原地追加 KV 或累加持久统计量的算子：每次预热都会额外推进或重复写入状态， + 可能改变后续正式计算的结果；把大型状态池加入 mutates_args 又会因 clone 增加显存占用， + 甚至触发 OOM。仅覆盖输出缓冲区，或可通过 mutates_args 完整保护输入的算子可保持默认值。 + 当前关闭该开关的特殊算子是 ``mtp_fused_recurrent_gated_delta_rule``，对应 autotune + ``kernel_name`` 为 ``_mtp_fused_recurrent_gated_delta_rule_fwd_kernel:v1``，可用这两个名字 + 查询实现、调用位置和缓存配置目录。 + 此开关只控制已有配置的额外预热，不关闭新配置的搜索、benchmark 内部的预热/计时和 + 最终正式执行；实际搜索仍需由调用方保证状态可以被反复更新，或提供相应的状态保护。 Returns: Callable: A callable object that wraps the original function and performs autotuning @@ -71,27 +100,48 @@ def decorator(fn: Callable) -> Callable: run_key_func=run_key_func, run_key_distance_func=run_key_distance_func, mutates_args=mutates_args, + kernel_type=kernel_type, + rebuild_input_func=rebuild_input_func, + warmup_all_exist_config=warmup_all_exist_config, ) return decorator class Autotuner: - _autotune_warmup: bool = False + _autotune_warmup_kernel_type: Optional[AutotuneKernelType] = None @staticmethod - def start_autotune_warmup(): - Autotuner._autotune_warmup = True + def start_autotune_warmup(kernel_type: AutotuneKernelType = AutotuneKernelType.GENERAL): + """Select the kernel category to tune; all distributed ranks must select the same phase.""" + Autotuner._autotune_warmup_kernel_type = AutotuneKernelType(kernel_type) return @staticmethod def end_autotune_warmup(): - Autotuner._autotune_warmup = False + Autotuner._autotune_warmup_kernel_type = None return @staticmethod - def is_autotune_warmup(): - return Autotuner._autotune_warmup + def is_autotune_warmup() -> bool: + """Report whether any warmup phase is active.""" + return Autotuner._autotune_warmup_kernel_type is not None + + @staticmethod + def is_kernel_autotune_warmup(kernel_type: AutotuneKernelType) -> bool: + """Report whether this kernel category is selected for warmup.""" + return Autotuner._autotune_warmup_kernel_type == AutotuneKernelType(kernel_type) + + @staticmethod + @contextmanager + def autotune_warmup(kernel_type: AutotuneKernelType = AutotuneKernelType.GENERAL): + """Restore the previous warmup phase on exit, including nested scopes and exceptions.""" + previous_type = Autotuner._autotune_warmup_kernel_type + Autotuner.start_autotune_warmup(kernel_type) + try: + yield + finally: + Autotuner._autotune_warmup_kernel_type = previous_type def __init__( self, @@ -102,10 +152,16 @@ def __init__( run_key_func: Callable, run_key_distance_func: Callable = lambda run_key, config_key: abs(int(run_key) - int(config_key)), mutates_args: List[str] = [], + kernel_type: AutotuneKernelType = AutotuneKernelType.GENERAL, + rebuild_input_func: Optional[Callable] = None, + warmup_all_exist_config: bool = True, ): self.configs_gen_func = configs_gen_func self.kernel_name = kernel_name + self.kernel_type = AutotuneKernelType(kernel_type) + self.rebuild_input_func = rebuild_input_func + self.warmup_all_exist_config = warmup_all_exist_config self.fn = fn self.static_key_func = static_key_func self.run_key_func = run_key_func @@ -141,6 +197,18 @@ def __call__(self, *args, **kwargs): if autotune_level == AutotuneLevel.CLOSE_AUTOTUNE: return self.fn(*args, **kwargs) + # decode attention 在多层中会重复调用，相同配置只需调优一次，避免强制调优拖慢启动。 + if self.kernel_type == AutotuneKernelType.DECODE_ATTENTION and autotune_level == AutotuneLevel.FORCE_AUTOTUNE: + autotune_level = AutotuneLevel.ADAPTIVE_AUTOTUNE + if not getattr(self, "_decode_force_autotune_logged", False): + logger.info( + f"Decode attention kernel {self.kernel_name}: FORCE_AUTOTUNE is treated as ADAPTIVE_AUTOTUNE " + "to avoid repeated tuning across layers and reduce startup time. Existing configs are reused. " + f"To retune, delete the cached config files in '{self.cache_dir}' before restarting " + "with LIGHTLLM_TRITON_AUTOTUNE_LEVEL=1 or 2." + ) + self._decode_force_autotune_logged = True + rank_id = 0 if not dist.is_initialized() else get_global_rank() world_size = 1 if not dist.is_initialized() else get_global_world_size() @@ -148,7 +216,8 @@ def __call__(self, *args, **kwargs): run_key = str(self._run_key(*args, **kwargs)) # Lazy load the cached configs in lightllm/common/triton_utils/autotune_kernel_configs - if self._try_load_cache(static_key) or Autotuner.is_autotune_warmup(): + # 先尝试加载缓存；关闭已有配置预热时仍须正常读取配置，不能用开关短路缓存加载。 + if (self._try_load_cache(static_key) or Autotuner.is_autotune_warmup()) and self.warmup_all_exist_config: all_configs = self.cached_configs.get(static_key, {}) for run_config in all_configs.values(): # warmup all configs @@ -165,10 +234,10 @@ def __call__(self, *args, **kwargs): ) self.cached_configs[static_key] = {} - if ( - autotune_level in [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE] - and Autotuner.is_autotune_warmup() - ): + if Autotuner.is_kernel_autotune_warmup(self.kernel_type) and autotune_level in [ + AutotuneLevel.ADAPTIVE_AUTOTUNE, + AutotuneLevel.FORCE_AUTOTUNE, + ]: need_tuning = (autotune_level == AutotuneLevel.FORCE_AUTOTUNE) or ( run_key not in self.cached_configs.get(static_key, {}) ) @@ -304,6 +373,10 @@ def _autotune(self, args, kwargs, static_key, run_key, rank_id, world_size): else: rank_tuning_configs = self.configs_gen_func() + # 仅为本次调优重建输入，构造开销不计入计时，最终执行仍使用调用方的原始输入。 + if self.rebuild_input_func is not None: + args, kwargs = self.rebuild_input_func(*args, **kwargs) + best_config = None best_time = float("inf") @@ -352,6 +425,9 @@ def _autotune(self, args, kwargs, static_key, run_key, rank_id, world_size): self.cached_configs[_static_key] = {} for _run_key, _config in _t_dict.items(): self.cached_configs[_static_key][_run_key] = _config + # 配置更新后，清除该 static_key 下缓存的旧匹配结果，避免继续使用旧配置，使新调优配置生效。 + # 新配置也可能改变其他 run_key 的最近邻选择，因此需要清除整个 static_key 的匹配缓存。 + self.fast_match_configs.pop(_static_key, None) # save configs to file if rank_id == 0: @@ -407,7 +483,11 @@ def _select_args(self, param_names, args, kwargs): if pos is not None and pos < len(args): values.append(args[pos]) else: - raise KeyError(f"Missing argument '{name}' required by key function") + # 可选参数也能参与 key；调用方省略时使用算子函数声明的默认值。 + parameter = inspect.signature(self.fn).parameters.get(name) + if parameter is None or parameter.default is inspect.Parameter.empty: + raise KeyError(f"Missing argument '{name}' required by key function") + values.append(parameter.default) return tuple(values) def _static_key(self, *args, **kwargs): diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 6f8973425f..7c7ac9fe48 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -319,6 +319,16 @@ def _launch_subprocesses(args: StartArgs): auto_configure_allreduce_flags_from_args(args) + # CUDA Graph 只需要覆盖调度器允许同时运行的请求数。配置得更大不会被真实请求使用， + # 反而会捕获无效的大 batch Graph 并额外占用显存，因此在全部参数调整完成后收敛到合法上限。 + # 关闭 CUDA Graph 时该参数不生效，保留用户原值。 + if not args.disable_cudagraph and args.graph_max_batch_size > args.running_max_req_size: + logger.warning( + f"graph_max_batch_size {args.graph_max_batch_size} exceeds running_max_req_size " + f"{args.running_max_req_size}; set graph_max_batch_size to {args.running_max_req_size}." + ) + args.graph_max_batch_size = args.running_max_req_size + # 校验用户已设置端口冲突（对齐原 PortManager 启动检查范围） ports_to_check = [args.port] if args.dp == 1 and args.nnodes > 1: diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index bded568bdc..acd2ac6711 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -162,15 +162,12 @@ def get_triton_autotune_level(): @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") +def get_decode_attn_autotune_seq_len() -> int: + """Decode attention 调优的代表性 KV 长度（token），默认 32768；调优时的 run key 按该长度分桶。""" + seq_len = int(os.getenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "32768")) + if seq_len <= 0: + raise ValueError("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN must be positive") + return seq_len g_model_init_done = False diff --git a/lightllm/utils/sgl_utils.py b/lightllm/utils/sgl_utils.py index c39e8bc142..475f3fd0c2 100644 --- a/lightllm/utils/sgl_utils.py +++ b/lightllm/utils/sgl_utils.py @@ -1,7 +1,8 @@ import torch +from typing import Optional, Tuple -from lightllm.common.triton_utils.autotuner import AutotuneLevel, Autotuner, autotune -from lightllm.utils.envs_utils import get_triton_autotune_level +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -61,43 +62,116 @@ def _flash_attn_kvcache_static_key(q, k_cache, v_cache, causal, window_size, sof } -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): +def _flash_attn_kvcache_run_key(page_table, max_seqlen_q, max_seqlen_k): 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]) + max_q_len = int(max_seqlen_q) + # 正常执行按调用方提供的真实 KV 长度查找配置，避免读取 CUDA 张量引入同步或使用 Graph 页表容量。 + max_kv_len = int(max_seqlen_k) + # 只有开启 decode attention 调优时才使用重建输入的目标长度。 + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) and get_triton_autotune_level() in [ + AutotuneLevel.ADAPTIVE_AUTOTUNE, + AutotuneLevel.FORCE_AUTOTUNE, + ]: + max_kv_len = get_decode_attn_autotune_seq_len() + # run key 中的 KV 长度向上取整到 512 的整数倍，同一区间复用配置匹配结果，减少重复搜索。 + max_kv_len = (max_kv_len + 511) // 512 * 512 return batch_size * 10_000_000_000_000 + max_q_len * 10_000_000 + max_kv_len +def _flash_attn_kvcache_rebuild_inputs( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + page_table: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + max_seqlen_q: int, + max_seqlen_k: int, + **kwargs, +): + # 本回调在每次实际搜索配置前执行一次，输入构造开销不计入候选配置的 benchmark 耗时。 + # Graph 初始化时的真实请求通常很短，而页表宽度是 Graph 能容纳的最大 KV 长度， + # 两者都不能直接代表期望调优的请求长度，因此需要单独构造用于计时的 KV 访问范围。 + # 返回的替换输入只用于本次调优；后续正常执行和 Graph 捕获仍使用调用方的原始输入。 + + # 页表每行对应一个 attention 请求。MTP 下一个请求可能包含多个 query token， + # 因此从页表取得请求数，不能直接把 Q 的 token 数当作 batch_size。 + batch_size = page_table.shape[0] + + # 用 LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN 指定代表性 KV 长度，单位是 token，默认 32768。 + # 调优阶段的 run_key 按同一个配置值分桶，保证配置对应实际 benchmark 的长度区间。 + kv_len = get_decode_attn_autotune_seq_len() + + # KV cache 布局为 [物理页数, 每页 token 数, KV head 数, head_dim]。 + # 将目标 token 数向上取整为所需页数，支持长度不整除 page_size，以及调优页表比原页表更宽或更窄。 + page_size = k_cache.shape[1] + max_pages = (kv_len + page_size - 1) // page_size + + # K/V 共用同一张页表，物理页数必须一致；至少需要一页才能生成合法索引并进行后续取模。 + assert k_cache.shape[0] == v_cache.shape[0], "K/V caches must have the same number of pages" + num_pages = k_cache.shape[0] + if num_pages == 0: + raise ValueError("FA3 autotuning requires a non-empty KV cache") + + # 复用模型已有的 K/V 存储，只分配较小的页表和长度张量，避免为每个模拟长请求另分配完整 KV。 + # 原页表可能只填充了占位请求所用的少量条目，不能仅增加 cache_seqlens 就直接读取后续条目。 + # 因此新建页表，并沿用原页表的 dtype/device；这里没有改写原页表或 K/V 数据。 + page_table = torch.arange(batch_size * max_pages, dtype=page_table.dtype, device=page_table.device) + + # 取模将页编号映射到 [0, num_pages)，再整理为每个请求一行、每行 max_pages 个页编号的页表。 + # 例如 batch_size=2、max_pages=4、num_pages=6，结果为 [[0, 1, 2, 3], [4, 5, 0, 1]]。 + # 物理页充足时，各请求使用不同页；不足时通过循环复用模拟多个长请求，避免越界。 + # 这种复用可能提高 GPU 缓存命中率， + # 因此调优时的访存特征与各请求拥有独立 KV 的真实场景存在差异。 + page_table = page_table.remainder_(num_pages).view(batch_size, max_pages) + + # FA3 根据 cache_seqlens 决定每个请求实际读取多少个 KV token，因此所有模拟请求都设为目标长度。 + # 这里保留精确的 token 数，而不是取整后的页容量，最后一页的多余位置不会算入有效 KV 长度。 + # 新建 int32 长度张量，避免修改原始短请求的长度并影响后续 Graph 捕获。 + cache_seqlens = torch.full((batch_size,), kv_len, dtype=torch.int32, device=k_cache.device) + + # Q、cu_seqlens_q 和 max_seqlen_q 保持原样，保留普通 decode、MTP 及不等长 query 分组的真实布局。 + # kwargs 原样透传 causal、window_size、softmax_scale 等设置，使计时采用与原调用一致的 attention 语义。 + # 返回参数顺序与 flash_attn_with_kvcache_autotune 的必填参数一致，仅替换 KV 页表和 KV 长度。 + return (q, k_cache, v_cache, cache_seqlens, page_table, cu_seqlens_q, max_seqlen_q, kv_len), kwargs + + @autotune( - kernel_name="sgl_fa3_kvcache_ns:v1", + kernel_name="sgl_fa3_kvcache_ns:v3", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, 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, + rebuild_input_func=_flash_attn_kvcache_rebuild_inputs, ) @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, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + page_table: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + max_seqlen_q: int, + max_seqlen_k: int, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + causal: bool = False, + window_size: Tuple[int, int] = (-1, -1), + softcap: float = 0.0, + num_splits: int = 0, + sinks: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + return_softmax_lse: bool = False, + run_config: Optional[dict] = None, ): + # KV 长度、页表及 query 布局由调用方显式提供；四维 batched Q 的 cu_seqlens_q 可以显式传 None。 + # max_seqlen_k 是 CPU 上已知的真实最大 KV token 数，仅用于配置查找，不传给底层 FA3 算子。 + # 此场景的 KV 已提前写入缓存，FA3 只读取已有缓存，不追加新的 K/V。 + # cu_seqlens_k_new 仅描述新增 K/V 的分段，因此必须为 None；实际 KV 长度由 cache_seqlens 指定。 + assert cu_seqlens_k_new is None, "cu_seqlens_k_new must be None when autotuning attention over an existing KV cache" + if run_config is None: run_config = {"num_splits": 0} @@ -118,67 +192,6 @@ def flash_attn_with_kvcache_autotune( sinks=sinks, k_descale=k_descale, v_descale=v_descale, - **kwargs, + softmax_scale=softmax_scale, + return_softmax_lse=return_softmax_lse, ) - - -def fa3_decode_autotune(model, cuda_graph_batch_sizes, batch_multiplier: int): - # 是否开启自动调优 - 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]) - for batch_size in cuda_graph_batch_sizes[::-1]: - assert batch_size % batch_multiplier == 0 - att_batch_size = batch_size // batch_multiplier - # 因为完整的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( - (batch_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) * batch_multiplier - 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=batch_multiplier, - 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 diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py new file mode 100644 index 0000000000..52b1c1cf2d --- /dev/null +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -0,0 +1,130 @@ +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import cuda_graph as cuda_graph_module +from lightllm.common.basemodel.cuda_graph import CudaGraph +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune + + +@pytest.fixture +def capture_env(monkeypatch): + env = SimpleNamespace(capturing=False, capture_count=0, graphs=[]) + env.args = SimpleNamespace( + enable_decode_microbatch_overlap=False, + enable_tpsp_mix_mode=False, + enable_torch_memory_saver=False, + ) + monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: env.args) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: None) + + class FakeGraph: + def __init__(self): + self.replays = 0 + env.graphs.append(self) + + def replay(self): + self.replays += 1 + + @contextmanager + def capture(graph, **kwargs): + assert not Autotuner.is_autotune_warmup() + env.capture_count += 1 + env.capturing = True + try: + yield + finally: + env.capturing = False + + monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) + monkeypatch.setattr(torch.cuda, "graph", capture) + return env + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("level", [0, 1, 2, 3]) +def test_decode_attention_tunes_before_capture_only(capture_env, tmp_path, monkeypatch, overlap, level): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + capture_env.args.enable_decode_microbatch_overlap = overlap + graph = CudaGraph(1, 2, 1, max_batch_size=2) + benchmarks = [] + + def make_kernel(kernel_type): + @autotune( + kernel_name=kernel_type.value, + kernel_type=kernel_type, + configs_gen_func=lambda: [{"block": 1}, {"block": 2}], + static_key_func=lambda: {}, + run_key_func=lambda size: size, + ) + def kernel(size, run_config=None): + return size + + def bench(size, run_config): + assert kernel_type == AutotuneKernelType.DECODE_ATTENTION + assert Autotuner.is_kernel_autotune_warmup(kernel_type) + assert not capture_env.capturing + benchmarks.append(run_config) + return 1.0 / run_config["block"] + + cache_dir = tmp_path / kernel_type.value + cache_dir.mkdir() + kernel._cache_dir = str(cache_dir) + monkeypatch.setattr(kernel, "_bench", bench) + return kernel, cache_dir / "configs.json" + + general, general_cache = make_kernel(AutotuneKernelType.GENERAL) + attention, attention_cache = make_kernel(AutotuneKernelType.DECODE_ATTENTION) + states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] + outputs = [object() for _ in states] + forward_calls = [] + + def decode_func(*infer_states): + forward_calls.append(capture_env.capturing) + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) == (not capture_env.capturing) + for original, state in zip(states, infer_states): + assert (state is original) == capture_env.capturing + assert not hasattr(state, "temporary_buffer") + state.temporary_buffer = object() + general(state.input_ids.shape[0]) + attention(state.input_ids.shape[0]) + return tuple(outputs) if overlap else outputs[0] + + result = graph.capture_decode(decode_func, *states) + assert result == (tuple(outputs) if overlap else outputs[0]) + assert forward_calls == [False, True] + assert not Autotuner.is_autotune_warmup() + assert capture_env.capture_count == 1 + assert capture_env.graphs[0].replays == 1 + expected_benchmarks = 0 + if level in [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]: + expected_benchmarks = 2 + assert len(benchmarks) == expected_benchmarks + assert attention_cache.exists() == (expected_benchmarks > 0) + assert not general_cache.exists() + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_warmup_failure_restores_phase_without_capturing(capture_env, overlap): + capture_env.args.enable_decode_microbatch_overlap = overlap + graph = CudaGraph(1, 2, 1, max_batch_size=2) + states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] + + def decode_func(*infer_states): + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + raise RuntimeError("decode warmup failed") + + with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): + with pytest.raises(RuntimeError, match="decode warmup failed"): + graph.capture_decode(decode_func, *states) + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert not Autotuner.is_autotune_warmup() + assert capture_env.capture_count == 0 + assert graph.graph == {} diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py index b25088790c..7c1af6bbc2 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py @@ -154,6 +154,7 @@ def mtp_diverse_attention( B_req_idx=b_req_idx, b_seq_len=b_seq_len, b_mark_shared_group=b_mark_shared_group, + max_kv_len=int(b_seq_len.max().item()), ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py new file mode 100644 index 0000000000..1a334f4868 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py @@ -0,0 +1,298 @@ +import collections +import inspect +import json +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel import cuda_graph as graph_module +from lightllm.common.basemodel.attention.triton import fp as state_module +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( + mtp_diverse_attn as decode_module, + stage1_single_token as stage1_module, + stage2_single_token as stage2_module, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) +def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage1_module.get_run_key(torch.empty(7, 8, 64), 2) == 7_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_lookup_uses_actual_length(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: + assert stage1_module.get_run_key(torch.empty(7, 8, 64), actual) == 7_000_000_000 + expected + + +def make_inputs(device="cpu", batch=7, block_batch=3, head_dim=128, token_count=17, table_width=16, block_num=2): + kv = torch.randn(token_count, 4, head_dim, dtype=torch.bfloat16, device=device) + return dict( + q=torch.randn(batch, 8, head_dim, dtype=kv.dtype, device=device), + k=kv[:, :2], + v=kv[:, 2:], + Req_to_tokens=torch.arange(6 * table_width, dtype=torch.int32, device=device).view(6, table_width) + % max(1, token_count), + B_req_idx=torch.full((batch,), 4, dtype=torch.int32, device=device), + b_seq_len=torch.full((batch,), 2, dtype=torch.int32, device=device), + b_mark_shared_group=torch.ones(batch, dtype=torch.int32, device=device), + max_kv_len=2, + mid_out=torch.full((batch, 8, block_num, head_dim), -100, dtype=kv.dtype, device=device), + mid_out_logsumexp=torch.full((batch, 8, block_num), -100, dtype=torch.float32, device=device), + block_batch=block_batch, + ) + + +def rebuild(inputs): + args, kwargs = stage1_module.rebuild_inputs(**inputs) + return inspect.signature(stage1_module.mtp_diverse_stage1_single_token.fn).bind(*args, **kwargs).arguments + + +@pytest.mark.parametrize("token_count", [3, 128]) +@pytest.mark.parametrize( + "batch, block_batch, length, req_ids, lengths, marks", + [ + (7, 3, 17, [0, 0, 0, 1, 1, 1, 2], [15, 16, 17, 15, 16, 17, 17], [0, 0, 3, 0, 0, 3, 1]), + (5, 4, 4, [0, 0, 0, 0, 1], [1, 2, 3, 4, 4], [0, 0, 0, 4, 1]), + (1, 4, 1, [0], [1], [1]), + ], +) +def test_rebuild_preserves_inputs_and_shared_prefixes( + monkeypatch, token_count, batch, block_batch, length, req_ids, lengths, marks +): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + inputs = make_inputs(batch=batch, block_batch=block_batch, token_count=token_count) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + rebuilt = rebuild(inputs) + assert rebuilt["B_req_idx"].tolist() == req_ids + assert rebuilt["b_seq_len"].tolist() == lengths + assert rebuilt["b_mark_shared_group"].tolist() == marks + assert rebuilt["max_kv_len"] == length + expected = torch.arange((max(req_ids) + 1) * length).view(-1, length) % token_count + torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), expected) + for name in ["q", "k", "v", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt["block_batch"] == block_batch + for name, value in snapshots.items(): + torch.testing.assert_close(inputs[name], value) + + +@pytest.mark.parametrize("batch, block_batch, length", [(7, 3, 2), (2, 4, 1)]) +def test_rebuild_rejects_length_shorter_than_group(monkeypatch, batch, block_batch, length): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + with pytest.raises(ValueError, match="at least the largest MTP group size"): + rebuild(make_inputs(batch=batch, block_batch=block_batch)) + + +def test_rebuild_rejects_empty_or_mismatched_cache(): + with pytest.raises(ValueError, match="non-empty KV cache"): + rebuild(make_inputs(token_count=0)) + inputs = make_inputs() + inputs["v"] = inputs["v"][:-1] + with pytest.raises(AssertionError, match="same number of tokens"): + rebuild(inputs) + + +def test_decode_passes_saved_length_and_selected_block_n(monkeypatch): + inputs = make_inputs() + infer_state = SimpleNamespace( + max_kv_seq_len=8193, + b_req_idx=inputs["B_req_idx"], + b_seq_len=inputs["b_seq_len"], + req_manager=SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]), + ) + model = SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), + ) + monkeypatch.setattr(state_module, "build_mtp_shared_group_markers", lambda *a, **kw: inputs["b_mark_shared_group"]) + state = state_module.TritonDecodeAttState(backend=SimpleNamespace(model=model), infer_state=infer_state) + state.init_state() + infer_state.max_kv_seq_len = 32768 + calls = [] + + def stage1(**kwargs): + calls.append(kwargs) + return 32 + + monkeypatch.setattr(decode_module, "mtp_diverse_stage1_single_token", stage1) + monkeypatch.setattr(decode_module, "mtp_diverse_stage2_single_token", lambda **kwargs: calls.append(kwargs)) + state.decode_att(q=inputs["q"], k=inputs["k"], v=inputs["v"]) + assert calls[0]["max_kv_len"] == 8193 + assert calls[0]["b_mark_shared_group"] is inputs["b_mark_shared_group"] + assert calls[1]["block_n"] == 32 + assert calls[1]["max_kv_len"] == 8193 + + +def reference_attention(inputs): + results = [] + for row, length in enumerate(inputs["b_seq_len"].tolist()): + indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() + with sdpa_kernel(SDPBackend.MATH): + results.append( + F.scaled_dot_product_attention( + inputs["q"][row].float().unsqueeze(1), + inputs["k"][indices].float().transpose(0, 1), + inputs["v"][indices].float().transpose(0, 1), + enable_gqa=True, + ).squeeze(1) + ) + return torch.stack(results) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [17, 8192, 16384]) +@pytest.mark.parametrize("block_batch", [1, 3, 4]) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_shared_groups_match_fp32_reference(monkeypatch, kv_len, block_batch, head_dim): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + inputs = rebuild(make_inputs("cuda", block_batch=block_batch, head_dim=head_dim, token_count=7 * kv_len)) + reference = reference_attention(inputs) + out = torch.empty_like(inputs["q"]) + # 仅两个 grid block，长 KV 必须循环多次；17 token 同时覆盖末块 mask 和无效分块的归约边界。 + for block_n in [16, 32, 64]: + selected_block_n = stage1_module.mtp_diverse_stage1_single_token.fn( + **inputs, + run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2, "warp_specialize": False}, + ) + stage2_module.mtp_diverse_stage2_single_token.fn( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["b_seq_len"], out, selected_block_n, kv_len + ) + torch.testing.assert_close(out.float(), reference, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8192, 16384]) +def test_full_autotune_before_graph_and_cached_reuse(tmp_path, monkeypatch, kv_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + kernel = stage1_module.mtp_diverse_stage1_single_token + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + assert kernel.mutates_args == [] + # 使用生产路径的小 batch 布局，让 stage2 在长请求下归约完整的 128 个分块。 + inputs = make_inputs("cuda", table_width=kv_len, token_count=3 * kv_len, block_num=128) + stage2 = stage2_module.mtp_diverse_stage2_single_token + stage2_cache = tmp_path / "stage2" + stage2_cache.mkdir() + monkeypatch.setattr(stage2, "_cache_dir", str(stage2_cache), raising=False) + monkeypatch.setattr(stage2, "cached_configs", {}) + monkeypatch.setattr(stage2, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(stage2, "warmuped_configs_set", set()) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + original_rebuild, original_bench = kernel.rebuild_input_func, kernel._bench + rebuild_count, benchmark_count, valid_benchmark_count = 0, 0, 0 + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + assert not torch.cuda.is_current_stream_capturing() + rebuild_count += 1 + return original_rebuild(*args, **kwargs) + + def checked_bench(*args, **kwargs): + nonlocal benchmark_count, valid_benchmark_count + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["max_kv_len"] == kv_len + assert rebuilt["b_seq_len"].tolist() == [kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len] + assert rebuilt["b_mark_shared_group"].tolist() == [0, 0, 3, 0, 0, 3, 1] + assert rebuilt["Req_to_tokens"].shape == (3, kv_len) + elapsed = original_bench(*args, **kwargs) + # 不可编译或资源不足的候选按 Autotuner 原有规则返回 inf；完整搜索必须选出有效配置。 + if math.isfinite(elapsed): + valid_benchmark_count += 1 + benchmark_count += 1 + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + monkeypatch.setattr( + graph_module, + "get_env_start_args", + lambda: SimpleNamespace( + enable_decode_microbatch_overlap=False, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False + ), + ) + out = torch.empty_like(inputs["q"]) + + def decode(state): + block_n = kernel(**inputs) + stage2(inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["b_seq_len"], out, block_n, inputs["max_kv_len"]) + return out + + graph = graph_module.CudaGraph(1, 7, 1, max_batch_size=7, max_len_in_batch=32768) + graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(7, device="cuda"))) + torch.cuda.synchronize() + assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) + assert valid_benchmark_count > 0 + # 两个 KV 的短请求沿用原 MTP 测试的 BF16 容差；下面的长请求仍使用更严格的 FP32 对照。 + torch.testing.assert_close(out.float(), reference_attention(inputs), atol=1e-2, rtol=1e-2) + for name, value in snapshots.items(): + if name not in ["mid_out", "mid_out_logsumexp"]: + torch.testing.assert_close(inputs[name], value) + cache_file = next(tmp_path.glob("*.json")) + saved = cache_file.read_bytes() + assert list(json.loads(saved)) == [str(7_000_000_000 + kv_len)] + assert json.loads(saved)[str(7_000_000_000 + kv_len)] is not None + stage2_configs = json.loads(next(stage2_cache.glob("*.json")).read_bytes()) + assert list(stage2_configs) == [str(7 * 8 * 1_000_000_000 + 128)] + assert next(iter(stage2_configs.values())) is not None + # 同一层配置以及从文件重新加载的配置，在 FORCE 模式下都不能再次搜索。 + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + kernel.cached_configs.clear() + kernel.fast_match_configs.clear() + kernel.warmuped_configs_set.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) + assert cache_file.read_bytes() == saved + # 原缓冲区改为真实长请求和共享组后回放，验证 Graph 没有捕获 benchmark 的重建输入。 + inputs["B_req_idx"].copy_(torch.tensor([1, 1, 1, 3, 3, 3, 2], device="cuda", dtype=torch.int32)) + inputs["b_mark_shared_group"].copy_(torch.tensor([0, 0, 3, 0, 0, 3, 1], device="cuda", dtype=torch.int32)) + inputs["b_seq_len"].copy_( + torch.tensor([kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len], device="cuda", dtype=torch.int32) + ) + graph.graph[7][0].replay() + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py new file mode 100644 index 0000000000..c7f2e4babe --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py @@ -0,0 +1,220 @@ +import collections +import inspect +import json +import math +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import cuda_graph as graph_module +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import stage2_single_token as stage2_module +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, capacity, expected", [(None, 128, 128), ("8193", 256, 129), ("65", 128, 2)]) +def test_tuning_key_uses_effective_reduction_blocks(monkeypatch, level, configured, capacity, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage2_module.get_run_key(torch.empty(2, 4, capacity, 64), 64, 2) == 8_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.CLOSE_AUTOTUNE), + ], +) +def test_lookup_counts_real_blocks_without_reading_gpu_lengths(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: level) + mid = torch.empty(2, 4, 128, 64) + for length, blocks in [(1, 1), (64, 1), (65, 2), (511, 8), (512, 8), (513, 9), (8192, 128), (16384, 128)]: + assert stage2_module.get_run_key(mid, 64, length) == 8_000_000_000 + blocks + + +def test_key_reuses_equal_work_and_distinguishes_capacity_limits(): + key = stage2_module.get_run_key + small, large = torch.empty(2, 4, 64, 64), torch.empty(2, 4, 128, 64) + assert key(small, 64, 513) == key(large, 64, 513) + assert key(small, 64, 8192) != key(large, 64, 8192) + assert key(large, 64, 8192) == key(large, 64, 16384) + + +def make_inputs(device="cpu", capacity=128, block_n=64, dtype=torch.bfloat16, head_dim=128): + # 模拟短请求：stage1 仅写入第一个分块，其余位置用 NaN 标记为未初始化。 + mid = torch.full((2, 4, capacity, head_dim), float("nan"), dtype=dtype, device=device) + lse = torch.full((2, 4, capacity), float("nan"), dtype=torch.float32, device=device) + mid[:, :, 0] = torch.randn(2, 4, head_dim, dtype=dtype, device=device) + lse[:, :, 0] = 0 + return dict( + mid_out=mid, + mid_out_logsumexp=lse, + B_Seqlen=torch.full((2,), 2, dtype=torch.int32, device=device), + out=torch.full((2, 4, head_dim), -100, dtype=dtype, device=device), + block_n=block_n, + max_kv_len=2, + ) + + +@pytest.mark.parametrize("length", [65, 8193, 16384]) +def test_rebuild_reuses_readonly_intermediate_inputs(monkeypatch, length): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + inputs = make_inputs() + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + args, kwargs = stage2_module.rebuild_inputs(**inputs) + rebuilt = inspect.signature(stage2_module.mtp_diverse_stage2_single_token.fn).bind(*args, **kwargs).arguments + assert rebuilt["max_kv_len"] == length + assert rebuilt["B_Seqlen"].tolist() == [length, length] + assert rebuilt["block_n"] == inputs["block_n"] + assert rebuilt["out"] is inputs["out"] + for name in ["mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt[name].shape == inputs[name].shape + assert rebuilt[name].stride() == inputs[name].stride() + assert rebuilt[name].dtype == inputs[name].dtype + assert rebuilt[name].device == inputs[name].device + assert rebuilt["B_Seqlen"] is not inputs["B_Seqlen"] + assert torch.isfinite(rebuilt["B_Seqlen"]).all() + for name, snapshot in snapshots.items(): + torch.testing.assert_close(inputs[name], snapshot, equal_nan=True) + + +def reduction_reference(inputs): + results = [] + for row, length in enumerate(inputs["B_Seqlen"].tolist()): + blocks = min((length + inputs["block_n"] - 1) // inputs["block_n"], inputs["mid_out"].shape[2]) + weights = torch.softmax(inputs["mid_out_logsumexp"][row, :, :blocks].float(), dim=-1) + results.append((weights.unsqueeze(-1) * inputs["mid_out"][row, :, :blocks].float()).sum(dim=1)) + return torch.stack(results) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("block_n", [16, 64]) +@pytest.mark.parametrize("capacity", [2, 128]) +@pytest.mark.parametrize("dtype,head_dim", [(torch.bfloat16, 128), (torch.float16, 64)]) +def test_all_configs_reduce_only_valid_blocks(block_n, capacity, dtype, head_dim): + lengths = [1, block_n, block_n + 1, capacity * block_n, capacity * block_n + 8192] + mid = torch.randn(len(lengths), 4, capacity, head_dim, dtype=dtype, device="cuda") + lse = torch.randn(len(lengths), 4, capacity, dtype=torch.float32, device="cuda") * 20 + for row, length in enumerate(lengths): + blocks = min((length + block_n - 1) // block_n, capacity) + mid[row, :, blocks:] = float("nan") + lse[row, :, blocks:] = float("nan") + inputs = dict( + mid_out=mid, + mid_out_logsumexp=lse, + B_Seqlen=torch.tensor(lengths, dtype=torch.int32, device="cuda"), + out=torch.empty(len(lengths), 4, head_dim, dtype=dtype, device="cuda"), + block_n=block_n, + max_kv_len=max(lengths), + ) + reference = reduction_reference(inputs) + for config in stage2_module.get_test_configs(): + stage2_module.mtp_diverse_stage2_single_token.fn(**inputs, run_config=config) + torch.testing.assert_close(inputs["out"].float(), reference, atol=2e-3, rtol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("length", [65, 8192, 16384]) +def test_full_tuning_graph_and_cache_reuse(tmp_path, monkeypatch, length): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + kernel = stage2_module.mtp_diverse_stage2_single_token + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + inputs = make_inputs("cuda") + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + original_rebuild, original_bench = kernel.rebuild_input_func, kernel._bench + rebuild_count, benchmark_count = 0, 0 + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + assert not torch.cuda.is_current_stream_capturing() + rebuild_count += 1 + return original_rebuild(*args, **kwargs) + + def checked_bench(*args, **kwargs): + nonlocal benchmark_count + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["B_Seqlen"].tolist() == [length, length] + assert rebuilt["max_kv_len"] == length and rebuilt["block_n"] == 64 + for name in ["mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt[name].shape == inputs[name].shape + elapsed = original_bench(*args, **kwargs) + assert math.isfinite(elapsed) + torch.testing.assert_close(inputs["out"], snapshots["out"]) + benchmark_count += 1 + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + monkeypatch.setattr( + graph_module, + "get_env_start_args", + lambda: SimpleNamespace( + enable_decode_microbatch_overlap=False, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False + ), + ) + + def decode(state): + kernel(**inputs) + return inputs["out"] + + graph = graph_module.CudaGraph(1, 2, 1, max_batch_size=2, max_len_in_batch=32768) + graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(2, device="cuda"))) + torch.cuda.synchronize() + assert rebuild_count == 1 and benchmark_count == 20 + torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) + for name in ["mid_out", "mid_out_logsumexp", "B_Seqlen"]: + torch.testing.assert_close(inputs[name], snapshots[name], equal_nan=True) + cache_file = next(tmp_path.glob("*.json")) + saved = cache_file.read_bytes() + expected_key = str(8_000_000_000 + min((length + 63) // 64, 128)) + assert list(json.loads(saved)) == [expected_key] + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + kernel.cached_configs.clear() + kernel.fast_match_configs.clear() + kernel.warmuped_configs_set.clear() + # 8K/16K 都归约 128 个块，切换代表性长度后也应复用同一个 key。 + if length >= 8192: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(24576 - length)) + get_decode_attn_autotune_seq_len.cache_clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == 20 + assert cache_file.read_bytes() == saved + inputs["mid_out"].normal_() + inputs["mid_out_logsumexp"].normal_() + inputs["B_Seqlen"].fill_(length) + graph.graph[2][0].replay() + torch.cuda.synchronize() + torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py new file mode 100644 index 0000000000..55ef9fcf37 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py @@ -0,0 +1,265 @@ +import collections +import inspect +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from frozendict import frozendict +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( + gqa_flash_decoding as decode_module, + gqa_flash_decoding_stage1 as stage1_module, +) +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( + flash_decode_stage2, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) +def test_tuning_key_uses_bucketed_configured_length(monkeypatch, level, configured, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + q = torch.empty(2, 8, 64) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage1_module.get_run_key(q, 2) == 2_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_non_tuning_key_uses_actual_length(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + q = torch.empty(2, 8, 64) + for actual, bucket in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: + assert stage1_module.get_run_key(q, actual) == 2_000_000_000 + bucket + + +def make_inputs(device="cpu", head_dim=64, token_count=17, table_width=16): + # K/V 是共享存储中的不连续视图；原请求索引也不是从 0 开始连续排列。 + kv = torch.randn(token_count, 4, head_dim, dtype=torch.bfloat16, device=device) + return dict( + q=torch.randn(2, 8, head_dim, dtype=kv.dtype, device=device), + k=kv[:, :2], + v=kv[:, 2:], + Req_to_tokens=torch.arange(5 * table_width, dtype=torch.int32, device=device).view(5, table_width) + % token_count, + B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), + B_Seqlen=torch.tensor([3, 2], dtype=torch.int32, device=device), + max_len_in_batch=3, + mid_out=torch.full((2, 8, 2, head_dim), -100, dtype=kv.dtype, device=device), + mid_out_logsumexp=torch.full((2, 8, 2), -100, dtype=torch.float32, device=device), + block_seq=256, + sliding_window=(-1, -1), + ) + + +def test_static_key_shares_heads_and_blocks_but_separates_windows(): + kernel = stage1_module.flash_decode_stage1 + inputs = make_inputs() + full_key = frozendict(kernel._static_key(**inputs)) + inputs.pop("sliding_window") + assert frozendict(kernel._static_key(**inputs)) == full_key + + inputs["q"] = torch.empty(2, 4, 64, dtype=torch.bfloat16) + inputs["k"] = torch.empty(17, 1, 64, dtype=torch.bfloat16) + inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) + assert frozendict(kernel._static_key(**inputs)) == full_key + + window_key = frozendict(kernel._static_key(**inputs, sliding_window=(511, 0))) + assert window_key != full_key + assert frozendict(kernel._static_key(**inputs, sliding_window=[511, 0])) == window_key + assert frozendict(kernel._static_key(**inputs, sliding_window=(1023, 0))) != window_key + + +@pytest.mark.parametrize("num_tokens", [3, 128]) +def test_rebuild_preserves_layout_and_bounds_indices(monkeypatch, num_tokens): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") + inputs = make_inputs(token_count=num_tokens) + inputs["sliding_window"] = (7, 0) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + + args, kwargs = stage1_module.rebuild_inputs(**inputs) + rebuilt = inspect.signature(stage1_module.flash_decode_stage1.fn).bind(*args, **kwargs).arguments + assert rebuilt["Req_to_tokens"].shape == (2, 17) + assert rebuilt["Req_to_tokens"].min() >= 0 and rebuilt["Req_to_tokens"].max() < num_tokens + assert rebuilt["B_req_idx"].tolist() == [0, 1] + assert rebuilt["B_Seqlen"].tolist() == [17, 17] + assert rebuilt["max_len_in_batch"] == 17 + for name in ["q", "k", "v", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt["block_seq"] == 256 + assert rebuilt["sliding_window"] == (7, 0) + for name, value in snapshots.items(): + torch.testing.assert_close(inputs[name], value) + + +def test_state_retains_actual_length_before_graph_capture(monkeypatch): + state = TritonDecodeAttState( + backend=SimpleNamespace( + model=SimpleNamespace( + is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0) + ) + ), + infer_state=SimpleNamespace(max_kv_seq_len=8193), + ) + state.init_state() + state.infer_state.max_kv_seq_len = 32768 + calls = [] + monkeypatch.setattr( + decode_module, "gqa_token_decode_attention_flash_decoding", lambda **kwargs: calls.append(kwargs) + ) + state.decode_att( + q=torch.empty(2, 8, 64), + k=torch.empty(16, 2, 64), + v=torch.empty(16, 2, 64), + alloc_func=lambda shape, dtype: torch.empty(shape, dtype=dtype), + ) + assert calls[0]["max_len_in_batch"] == 8193 + + +def reference_attention(inputs): + outputs = [] + for i, length in enumerate(inputs["B_Seqlen"].tolist()): + start = max(0, length - 1 - inputs["sliding_window"][0]) if inputs["sliding_window"][0] >= 0 else 0 + indices = inputs["Req_to_tokens"][inputs["B_req_idx"][i], start:length].long() + with sdpa_kernel(SDPBackend.MATH): + outputs.append( + F.scaled_dot_product_attention( + inputs["q"][i].float().unsqueeze(1), + inputs["k"][indices].float().transpose(0, 1), + inputs["v"][indices].float().transpose(0, 1), + enable_gqa=True, + ).squeeze(1) + ) + return torch.stack(outputs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8192, 16384]) +@pytest.mark.parametrize("head_dim", [64, 256]) +@pytest.mark.parametrize("sliding_window", [(-1, -1), (511, 0)]) +def test_stage1_stage2_long_kv_matches_reference(kv_len, head_dim, sliding_window): + inputs = make_inputs("cuda", head_dim=head_dim, token_count=kv_len * 2, table_width=kv_len) + inputs["B_Seqlen"] = torch.tensor([kv_len, kv_len - 1], dtype=torch.int32, device="cuda") + inputs["max_len_in_batch"] = kv_len + inputs["sliding_window"] = sliding_window + reference = reference_attention(inputs) + # 只有两个 grid block，长 KV 需要每个 program 循环多次；stage2 仍归约同一布局。 + output = torch.empty_like(inputs["q"]) + for block_n in [16, 64, 128]: + stage1_module.flash_decode_stage1.fn(**inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2}) + flash_decode_stage2( + inputs["mid_out"], + inputs["mid_out_logsumexp"], + inputs["B_Seqlen"], + output, + inputs["block_seq"], + sliding_window=sliding_window, + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8192, 16384]) +def test_autotune_rebuilds_once_and_graph_uses_original_inputs(tmp_path, monkeypatch, kv_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + kernel = stage1_module.flash_decode_stage1 + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + assert kernel.mutates_args == [] + monkeypatch.setattr( + kernel, + "configs_gen_func", + lambda: [{"BLOCK_N": n, "num_warps": 4, "num_stages": 2} for n in [16, 64, 128]], + ) + inputs = make_inputs("cuda") + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + reference = reference_attention(inputs) + rebuild = kernel.rebuild_input_func + benchmark = kernel._bench + rebuild_count = 0 + benchmark_count = 0 + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + rebuild_count += 1 + return rebuild(*args, **kwargs) + + def checked_bench(*args, **kwargs): + nonlocal benchmark_count + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["B_Seqlen"].tolist() == [kv_len, kv_len] + assert rebuilt["Req_to_tokens"].shape == (2, kv_len) + assert rebuilt["B_req_idx"].tolist() == [0, 1] + assert rebuilt["max_len_in_batch"] == kv_len + assert rebuilt["mid_out"] is inputs["mid_out"] + elapsed = benchmark(*args, **kwargs) + assert math.isfinite(elapsed) + benchmark_count += 1 + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == 3 + assert all(list(configs) == [str(2_000_000_000 + kv_len)] for configs in kernel.cached_configs.values()) + # FORCE 在 decode 阶段仍复用缓存，不因模型层重复调用而再次调优。 + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == 3 + + output = torch.empty_like(inputs["q"]) + + def decode(): + kernel(**inputs) + flash_decode_stage2( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"] + ) + + decode() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + decode() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and benchmark_count == 3 + for name, value in snapshots.items(): + if name not in ["mid_out", "mid_out_logsumexp"]: + torch.testing.assert_close(inputs[name], value) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py new file mode 100644 index 0000000000..655dd75d34 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py @@ -0,0 +1,293 @@ +import collections +import inspect +import json +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState +from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ( + int4kv_flash_decoding_stage1 as stage1_module, + ppl_int4kv_flash_decoding as decode_module, +) +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( + flash_decode_stage2, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +def make_inputs(device="cpu", head_dim=64, group_size=8, dtype=torch.bfloat16, token_count=17, block_num=2): + # 与内存管理器一致：K/V 和 scale 分别来自共享存储中的非连续视图。 + low = torch.randint(0, 15, (token_count, 4, head_dim // 2), device=device) + high = torch.randint(0, 15, low.shape, device=device) + kv = (low | (high << 4)).to(torch.int8) + scale = (torch.rand(token_count, 4, head_dim // group_size, device=device) * 0.1 + 0.02).to(dtype) + return dict( + q=torch.randn(2, 8, head_dim, dtype=dtype, device=device), + k=kv[:, :2], + k_scale=scale[:, :2], + v=kv[:, 2:], + v_scale=scale[:, 2:], + Req_to_tokens=torch.arange(5 * 16, dtype=torch.int32, device=device).view(5, 16) % max(1, token_count), + B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), + B_Seqlen=torch.tensor([3, 2], dtype=torch.int32, device=device), + max_kv_seq_len=3, + mid_out=torch.full((2, 8, block_num, head_dim), -100, dtype=dtype, device=device), + mid_out_logsumexp=torch.full((2, 8, block_num), -100, dtype=dtype, device=device), + block_seq=256, + ) + + +def rebuild(inputs): + args, kwargs = stage1_module.rebuild_inputs(**inputs) + return inspect.signature(stage1_module.int4kv_flash_decode_stage1.fn).bind(*args, **kwargs).arguments + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) +def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage1_module.get_run_key(torch.empty(2, 8, 64), 32768) == 2_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_lookup_uses_actual_length(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: + assert stage1_module.get_run_key(torch.empty(2, 8, 64), actual) == 2_000_000_000 + expected + + +@pytest.mark.parametrize("token_count", [3, 128]) +def test_rebuild_preserves_packed_cache_and_scales(monkeypatch, token_count): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") + inputs = make_inputs(token_count=token_count) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + rebuilt = rebuild(inputs) + torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), torch.arange(34).view(2, 17) % token_count) + assert rebuilt["B_req_idx"].tolist() == [0, 1] + assert rebuilt["B_Seqlen"].tolist() == [17, 17] + assert rebuilt["max_kv_seq_len"] == 17 + assert rebuilt["block_seq"] == inputs["block_seq"] + for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + for name, value in snapshots.items(): + torch.testing.assert_close(inputs[name], value) + + +@pytest.mark.parametrize("name", ["v", "k_scale", "v_scale"]) +def test_rebuild_rejects_inconsistent_token_capacity(name): + inputs = make_inputs() + inputs[name] = inputs[name][:-1] + with pytest.raises(AssertionError, match="same number of tokens"): + rebuild(inputs) + + +def test_rebuild_rejects_empty_cache(): + with pytest.raises(ValueError, match="non-empty KV cache"): + rebuild(make_inputs(token_count=0)) + + +def test_static_key_separates_quant_groups_without_head_or_block_counts(): + kernel = stage1_module.int4kv_flash_decode_stage1 + inputs = make_inputs() + key = kernel._static_key(**inputs) + inputs["q"] = inputs["q"][:, :4] + for name in ["k", "k_scale", "v", "v_scale"]: + inputs[name] = inputs[name][:, :1] + inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) + assert kernel._static_key(**inputs) == key + inputs["k_scale"] = inputs["k_scale"][:, :, :2] + assert kernel._static_key(**inputs) != key + + +def test_decode_passes_saved_length_to_stage1(monkeypatch): + inputs = make_inputs() + state = Int4kvTritonDecodeAttState( + backend=SimpleNamespace(), + infer_state=SimpleNamespace( + batch_size=2, + max_kv_seq_len=8193, + b_req_idx=inputs["B_req_idx"], + b_seq_len=inputs["B_Seqlen"], + req_manager=SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]), + ), + ) + state.init_state() + state.infer_state.max_kv_seq_len = 32768 + calls = [] + monkeypatch.setattr(stage1_module, "int4kv_flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) + stage2 = inspect.getmodule(flash_decode_stage2) + monkeypatch.setattr(stage2, "flash_decode_stage2", lambda *args: None) + state.decode_att( + inputs["q"], + (inputs["k"], inputs["k_scale"]), + (inputs["v"], inputs["v_scale"]), + alloc_func=lambda shape, dtype, device: torch.empty(shape, dtype=dtype), + ) + assert calls[0]["max_kv_seq_len"] == 8193 + assert calls[0]["block_seq"] == 256 + + +def reference_attention(inputs): + def dequant(packed, scale): + packed = packed.to(torch.uint8) + unpacked = torch.stack((packed & 15, packed >> 4), dim=-1).flatten(-2).float() - 7 + group_size = unpacked.shape[-1] // scale.shape[-1] + # 与算子一致，以 scale 的 dtype 产生解量化 K/V，再用 FP32 SDPA 验证 attention 本身。 + return (unpacked * scale.float().repeat_interleave(group_size, dim=-1)).to(scale.dtype).float() + + k, v = dequant(inputs["k"], inputs["k_scale"]), dequant(inputs["v"], inputs["v_scale"]) + outputs = [] + for row, length in enumerate(inputs["B_Seqlen"].tolist()): + indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() + with sdpa_kernel(SDPBackend.MATH): + outputs.append( + F.scaled_dot_product_attention( + inputs["q"][row].float().unsqueeze(1), + k[indices].transpose(0, 1), + v[indices].transpose(0, 1), + enable_gqa=True, + ).squeeze(1) + ) + return torch.stack(outputs) + + +def reduce_output(inputs): + output = torch.empty_like(inputs["q"]) + flash_decode_stage2(inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"]) + return output + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [257, 8192, 16384]) +@pytest.mark.parametrize( + "head_dim,group_size,dtype", [(64, 8, torch.bfloat16), (128, 32, torch.float16), (256, 8, torch.bfloat16)] +) +def test_long_kv_matches_dequantized_fp32_reference(monkeypatch, kv_len, head_dim, group_size, dtype): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + inputs = rebuild(make_inputs("cuda", head_dim, group_size, dtype, token_count=2 * kv_len)) + inputs["B_Seqlen"][1] -= 1 + reference = reference_attention(inputs) + # 两个中间分块迫使长请求循环处理多个 BLOCK_SEQ，同时验证末尾 mask。 + for block_n in [16, 32, 64, 128]: + stage1_module.int4kv_flash_decode_stage1.fn( + **inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2} + ) + output = reduce_output(inputs) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8193, 16384]) +def test_full_autotune_and_graph_reuse(tmp_path, monkeypatch, kv_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + kernel = stage1_module.int4kv_flash_decode_stage1 + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + assert kernel.mutates_args == [] + inputs = make_inputs("cuda", token_count=2 * kv_len, block_num=128) + inputs["Req_to_tokens"] = torch.arange(5 * kv_len, dtype=torch.int32, device="cuda").view(5, kv_len) % (2 * kv_len) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + reference = reference_attention(inputs) + rebuild_func, benchmark = kernel.rebuild_input_func, kernel._bench + rebuild_count, timings = 0, [] + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + rebuild_count += 1 + return rebuild_func(*args, **kwargs) + + def checked_bench(*args, **kwargs): + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["B_Seqlen"].tolist() == [kv_len, kv_len] + assert rebuilt["Req_to_tokens"].shape == (2, kv_len) + assert rebuilt["max_kv_seq_len"] == kv_len + for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + elapsed = benchmark(*args, **kwargs) + timings.append(elapsed) + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and len(timings) == 48 and any(math.isfinite(t) for t in timings) + run_key = str(2_000_000_000 + (kv_len + 511) // 512 * 512) + assert all(list(configs) == [run_key] for configs in kernel.cached_configs.values()) + assert len(list(tmp_path.glob("*.json"))) == 1 + assert run_key in json.loads(next(tmp_path.glob("*.json")).read_text()) + torch.testing.assert_close(reduce_output(inputs).float(), reference, atol=2e-3, rtol=2e-2) + + # 清空内存配置后从文件重载，FORCE 下仍不能跨层重复搜索。 + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and len(timings) == 48 + + output = torch.empty_like(inputs["q"]) + + def decode(): + kernel(**inputs) + flash_decode_stage2( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"] + ) + + decode() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + decode() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and len(timings) == 48 + for name, snapshot in snapshots.items(): + if name not in ["mid_out", "mid_out_logsumexp"]: + torch.testing.assert_close(inputs[name], snapshot) + + # 捕获后更新为真正的长请求，验证 Graph 读取原始输入缓冲区，stage2 仍按 BLOCK_SEQ 归约。 + inputs["B_Seqlen"].fill_(kv_len) + inputs["max_kv_seq_len"] = kv_len + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and len(timings) == 48 diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py index ec70d3c697..e1b6bc7758 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py @@ -115,6 +115,7 @@ def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(sha cache_k_scale=cache_k_scale, cache_v=cache_v, cache_v_scale=cache_v_scale, + max_len_in_batch=seq_len, alloc_tensor_func=alloc_tensor_func, ) # 运行 diverse 版本 diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py new file mode 100644 index 0000000000..bc96c4cc70 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py @@ -0,0 +1,298 @@ +import collections +import inspect +import json +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel.attention.triton import int8kv as state_module +from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( + int8kv_flash_decoding_stage1 as stage1_module, + int8kv_flash_decoding as decode_module, +) +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( + flash_decode_stage2, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.setattr(state_module, "enable_diverse_mode_gqa_decode_fast_kernel", lambda: False) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +def make_inputs( + device="cpu", head_dim=64, group_size=8, dtype=torch.bfloat16, token_count=17, block_num=2, scale_dtype=None +): + # 与内存管理器一致：K/V 和 scale 分别来自共享存储中的非连续视图。 + kv = torch.randint(-127, 128, (token_count, 4, head_dim), dtype=torch.int8, device=device) + scale = (torch.rand(token_count, 4, head_dim // group_size, device=device) * 0.006 + 0.002).to(scale_dtype or dtype) + return dict( + q=torch.randn(2, 8, head_dim, dtype=dtype, device=device), + k=kv[:, :2], + k_scale=scale[:, :2], + v=kv[:, 2:], + v_scale=scale[:, 2:], + Req_to_tokens=torch.arange(5 * 16, dtype=torch.int32, device=device).view(5, 16) % max(1, token_count), + B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), + B_seq_len=torch.tensor([3, 2], dtype=torch.int32, device=device), + max_len_in_batch=3, + mid_out=torch.full((2, 8, block_num, head_dim), -100, dtype=dtype, device=device), + mid_out_logsumexp=torch.full((2, 8, block_num), -100, dtype=torch.float32, device=device), + block_seq=256, + ) + + +def rebuild(inputs): + args, kwargs = stage1_module.rebuild_inputs(**inputs) + return inspect.signature(stage1_module.flash_decode_stage1.fn).bind(*args, **kwargs).arguments + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) +def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage1_module.get_run_key(torch.empty(2, 8, 64), 32768) == 2_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_lookup_uses_actual_length(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: + assert stage1_module.get_run_key(torch.empty(2, 8, 64), actual) == 2_000_000_000 + expected + + +@pytest.mark.parametrize("token_count", [3, 128]) +def test_rebuild_preserves_int8_cache_and_scales(monkeypatch, token_count): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") + inputs = make_inputs(token_count=token_count) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + rebuilt = rebuild(inputs) + torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), torch.arange(34).view(2, 17) % token_count) + assert rebuilt["B_req_idx"].tolist() == [0, 1] + assert rebuilt["B_seq_len"].tolist() == [17, 17] + assert rebuilt["max_len_in_batch"] == 17 + assert rebuilt["block_seq"] == inputs["block_seq"] + for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + for name, value in snapshots.items(): + torch.testing.assert_close(inputs[name], value) + + +@pytest.mark.parametrize("name", ["v", "k_scale", "v_scale"]) +def test_rebuild_rejects_inconsistent_token_capacity(name): + inputs = make_inputs() + inputs[name] = inputs[name][:-1] + with pytest.raises(AssertionError, match="same number of tokens"): + rebuild(inputs) + + +def test_rebuild_rejects_empty_cache(): + with pytest.raises(ValueError, match="non-empty KV cache"): + rebuild(make_inputs(token_count=0)) + + +def test_static_key_separates_quant_groups_without_head_or_block_counts(): + kernel = stage1_module.flash_decode_stage1 + inputs = make_inputs() + key = kernel._static_key(**inputs) + inputs["q"] = inputs["q"][:, :4] + for name in ["k", "k_scale", "v", "v_scale"]: + inputs[name] = inputs[name][:, :1] + inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) + assert kernel._static_key(**inputs) == key + inputs["k_scale"] = inputs["k_scale"][:, :, :2] + assert kernel._static_key(**inputs) != key + + +def test_decode_passes_saved_length_to_stage1(monkeypatch): + inputs = make_inputs() + state = Int8kvTritonDecodeAttState( + backend=SimpleNamespace(), + infer_state=SimpleNamespace( + batch_size=2, + max_kv_seq_len=8193, + b_req_idx=inputs["B_req_idx"], + b_seq_len=inputs["B_seq_len"], + req_manager=SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]), + ), + ) + state.init_state() + state.infer_state.max_kv_seq_len = 32768 + calls = [] + monkeypatch.setattr(decode_module, "flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) + monkeypatch.setattr(decode_module, "flash_decode_stage2", lambda **kwargs: None) + state.decode_att( + inputs["q"], + (inputs["k"], inputs["k_scale"]), + (inputs["v"], inputs["v_scale"]), + alloc_func=lambda shape, dtype, device: torch.empty(shape, dtype=dtype), + ) + assert calls[0]["max_len_in_batch"] == 8193 + assert calls[0]["block_seq"] == 256 + + +def reference_attention(inputs): + def dequant(quantized, scale): + group_size = quantized.shape[-1] // scale.shape[-1] + # int8 直接乘分组 scale，转换为 Q 的 dtype 后再用 FP32 SDPA 验证 attention。 + return (quantized.float() * scale.float().repeat_interleave(group_size, dim=-1)).to(inputs["q"].dtype).float() + + k, v = dequant(inputs["k"], inputs["k_scale"]), dequant(inputs["v"], inputs["v_scale"]) + outputs = [] + for row, length in enumerate(inputs["B_seq_len"].tolist()): + indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() + with sdpa_kernel(SDPBackend.MATH): + outputs.append( + F.scaled_dot_product_attention( + inputs["q"][row].float().unsqueeze(1), + k[indices].transpose(0, 1), + v[indices].transpose(0, 1), + enable_gqa=True, + ).squeeze(1) + ) + return torch.stack(outputs) + + +def reduce_output(inputs): + output = torch.empty_like(inputs["q"]) + flash_decode_stage2( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_seq_len"], output, inputs["block_seq"] + ) + return output + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [257, 8192, 16384]) +@pytest.mark.parametrize( + "head_dim,group_size,dtype,scale_dtype", + [ + (64, 8, torch.bfloat16, torch.bfloat16), + (128, 32, torch.float16, torch.float16), + (256, 8, torch.bfloat16, torch.bfloat16), + (128, 8, torch.bfloat16, torch.float32), + ], +) +def test_long_kv_matches_dequantized_fp32_reference(monkeypatch, kv_len, head_dim, group_size, dtype, scale_dtype): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + inputs = rebuild(make_inputs("cuda", head_dim, group_size, dtype, token_count=2 * kv_len, scale_dtype=scale_dtype)) + inputs["B_seq_len"][1] -= 1 + reference = reference_attention(inputs) + # 两个中间分块迫使长请求循环处理多个 BLOCK_SEQ，同时验证末尾 mask。 + for block_n in [16, 32, 64, 128]: + stage1_module.flash_decode_stage1.fn(**inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2}) + output = reduce_output(inputs) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8193, 16384]) +def test_full_autotune_and_graph_reuse(tmp_path, monkeypatch, kv_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + kernel = stage1_module.flash_decode_stage1 + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + assert kernel.mutates_args == [] + inputs = make_inputs("cuda", token_count=2 * kv_len, block_num=128) + inputs["Req_to_tokens"] = torch.arange(5 * kv_len, dtype=torch.int32, device="cuda").view(5, kv_len) % (2 * kv_len) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + reference = reference_attention(inputs) + rebuild_func, benchmark = kernel.rebuild_input_func, kernel._bench + rebuild_count, timings = 0, [] + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + rebuild_count += 1 + return rebuild_func(*args, **kwargs) + + def checked_bench(*args, **kwargs): + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["B_seq_len"].tolist() == [kv_len, kv_len] + assert rebuilt["Req_to_tokens"].shape == (2, kv_len) + assert rebuilt["max_len_in_batch"] == kv_len + for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + elapsed = benchmark(*args, **kwargs) + timings.append(elapsed) + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and len(timings) == 48 and any(math.isfinite(t) for t in timings) + run_key = str(2_000_000_000 + (kv_len + 511) // 512 * 512) + assert all(list(configs) == [run_key] for configs in kernel.cached_configs.values()) + assert len(list(tmp_path.glob("*.json"))) == 1 + assert run_key in json.loads(next(tmp_path.glob("*.json")).read_text()) + torch.testing.assert_close(reduce_output(inputs).float(), reference, atol=2e-3, rtol=2e-2) + + # 清空内存配置后从文件重载，FORCE 下仍不能跨层重复搜索。 + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and len(timings) == 48 + + output = torch.empty_like(inputs["q"]) + + def decode(): + kernel(**inputs) + flash_decode_stage2( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_seq_len"], output, inputs["block_seq"] + ) + + decode() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + decode() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and len(timings) == 48 + for name, snapshot in snapshots.items(): + if name not in ["mid_out", "mid_out_logsumexp"]: + torch.testing.assert_close(inputs[name], snapshot) + + # 捕获后更新为真正的长请求，验证 Graph 读取原始输入缓冲区，stage2 仍按 BLOCK_SEQ 归约。 + inputs["B_seq_len"].fill_(kv_len) + inputs["max_len_in_batch"] = kv_len + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and len(timings) == 48 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py new file mode 100644 index 0000000000..2f8544cf75 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py @@ -0,0 +1,374 @@ +import collections +import inspect +import json +import math + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.linear_att import ( + mtp_fused_recurrent as kernel_module, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import ( + Autotuner, + AutotuneKernelType, + AutotuneLevel, +) + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch, tmp_path): + torch.manual_seed(42) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr( + autotuner_module, + "get_triton_autotune_level", + lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE, + ) + kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + + +def make_inputs( + lengths=(3, 2, 0), + mtp_size=3, + device="cpu", + head_k_dim=64, + head_v_dim=64, + heads=2, + value_heads=4, + dtype=torch.bfloat16, + state_dtype=torch.float32, + separate_write=False, + padding=0, +): + num_seqs = len(lengths) + num_tokens = sum(lengths) + padding + # 模拟生产环境从 mixed QKV 和 gate 投影切分的非连续 token 视图。 + mixed = torch.randn( + num_tokens, + heads * head_k_dim * 2 + value_heads * head_v_dim, + dtype=dtype, + device=device, + ) + q, k, v = mixed.split([heads * head_k_dim, heads * head_k_dim, value_heads * head_v_dim], dim=-1) + gates = torch.randn(num_tokens, value_heads * 3, dtype=dtype, device=device) + indices = torch.arange(1, num_seqs * mtp_size + 1, dtype=torch.int32, device=device).view(num_seqs, mtp_size) + return dict( + q=q.view(1, num_tokens, heads, head_k_dim), + k=k.view(1, num_tokens, heads, head_k_dim), + v=v.view(1, num_tokens, value_heads, head_v_dim), + initial_state=torch.randn( + 2 * num_seqs * mtp_size + 5, + value_heads, + head_k_dim, + head_v_dim, + dtype=state_dtype, + device=device, + ) + * 0.25, + cu_seqlens=torch.tensor( + [0, *torch.tensor(lengths).cumsum(0).tolist()], + dtype=torch.int64, + device=device, + ), + ssm_state_indices=indices, + ssm_state_write_indices=(indices + num_seqs * mtp_size if separate_write else indices), + num_accepted_tokens=torch.tensor( + [i % mtp_size + 1 for i in range(num_seqs)], + dtype=torch.int32, + device=device, + ), + A_log=torch.randn(value_heads, dtype=torch.float32, device=device) * 0.1, + dt_bias=torch.randn(value_heads, dtype=torch.float32, device=device) * 0.1, + a_raw=gates[:, :value_heads], + b_raw=gates[:, value_heads : 2 * value_heads], + ) + + +def rebuild(inputs): + args, kwargs = kernel_module.rebuild_inputs(**inputs) + return inspect.signature(kernel_module.mtp_fused_recurrent_gated_delta_rule.fn).bind(*args, **kwargs).arguments + + +def reference(inputs): + state = inputs["initial_state"].clone() + output = torch.zeros_like(inputs["v"], dtype=torch.float32) + q, k, v = [inputs[name][0].float() for name in ["q", "k", "v"]] + group_size = v.shape[1] // q.shape[1] + q = (q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt()).repeat_interleave(group_size, dim=1) + k = (k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt()).repeat_interleave(group_size, dim=1) + q *= q.shape[-1] ** -0.5 + cumulative = inputs["cu_seqlens"].tolist() + for row, (start, end) in enumerate(zip(cumulative, cumulative[1:])): + if start == end: + continue + read_index = inputs["ssm_state_indices"][row, inputs["num_accepted_tokens"][row] - 1] + h = inputs["initial_state"][read_index].float().clone() + for token in range(start, end): + g = -inputs["A_log"].float().exp() * F.softplus(inputs["a_raw"][token].float() + inputs["dt_bias"].float()) + h *= g.exp()[:, None, None] + delta = (v[token] - (h * k[token, :, :, None]).sum(1)) * inputs["b_raw"][token].float().sigmoid()[:, None] + h += k[token, :, :, None] * delta[:, None, :] + output[0, token] = (h * q[token, :, :, None]).sum(1) + state[inputs["ssm_state_write_indices"][row, token - start]] = h.to(state.dtype) + return output, state + + +@pytest.mark.parametrize("lengths,mtp_size,padding", [((3, 3), 3, 0), ((0,) * 7, 3, 7), ((2, 1, 0), 4, 2)]) +def test_rebuild_uses_independent_state_slots(lengths, mtp_size, padding): + inputs = make_inputs(lengths=lengths, mtp_size=mtp_size, padding=padding) + snapshots = {name: value.clone() for name, value in inputs.items()} + rebuilt = rebuild(inputs) + tokens = inputs["q"].shape[1] + active_seqs = (tokens + mtp_size - 1) // mtp_size + assert rebuilt["cu_seqlens"].tolist() == [min(i * mtp_size, tokens) for i in range(len(lengths) + 1)] + assert rebuilt["initial_state"] is inputs["initial_state"] + assert rebuilt["num_accepted_tokens"].tolist() == [1] * len(lengths) + assert rebuilt["ssm_state_write_indices"][:active_seqs].flatten().tolist() == list(range(active_seqs * mtp_size)) + kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule + assert kernel._run_key(**inputs) == kernel._run_key(**rebuilt) + for name in ["q", "k", "v", "A_log", "dt_bias", "a_raw", "b_raw"]: + assert rebuilt[name] is inputs[name] + for name, snapshot in snapshots.items(): + torch.testing.assert_close(inputs[name], snapshot) + + +def test_rebuild_checks_active_state_capacity(): + inputs = make_inputs(lengths=(0,) * 7, padding=7) + # 7 个有效 token 只需 7 个槽；空序列及末组未使用的索引取余后也应在池内。 + inputs["initial_state"] = inputs["initial_state"][:7] + rebuilt = rebuild(inputs) + assert rebuilt["initial_state"] is inputs["initial_state"] + for name in ["ssm_state_indices", "ssm_state_write_indices"]: + indices = rebuilt[name] + assert torch.all((indices >= 0) & (indices < 7)) + assert indices.flatten()[:7].tolist() == list(range(7)) + assert indices[2].tolist() == [6, 0, 1] + # 槽数不足时仍拒绝调优，不能靠取余让有效请求的状态槽相互重叠。 + inputs["initial_state"] = inputs["initial_state"][:6] + with pytest.raises(AssertionError, match="Not enough SSM state slots"): + rebuild(inputs) + + +def test_keys_use_mtp_workload_and_state_dtype(monkeypatch): + # 历史上下文长度与递推计算量无关，不应读取 full attention 的调优长度环境变量。 + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule + inputs = make_inputs() + key = kernel._static_key(**inputs) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel._run_key(**inputs) == 12 + assert kernel._run_key(**make_inputs(lengths=(3, 3, 0))) == kernel._run_key(**inputs) + assert kernel._run_key(**make_inputs(lengths=(3, 2))) == kernel._run_key(**inputs) + assert kernel._run_key(**make_inputs(mtp_size=4)) == 16 + # 不同序列数共享 token 桶，11/12 个 token 共用一桶，第 13 个 token 进入下一桶。 + assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 2, 0))) == 12 + assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 3, 0))) == 12 + assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 3, 1))) == 24 + # head 数不再编码进 run key，但不同模型/TP 的 head 规模应使用不同静态配置文件。 + more_heads = make_inputs(heads=4, value_heads=8) + assert kernel._run_key(**more_heads) == kernel._run_key(**inputs) + assert kernel._static_key(**more_heads) != key + assert kernel._static_key(**make_inputs(state_dtype=torch.bfloat16)) != key + assert kernel._static_key(**make_inputs(mtp_size=4)) != key + assert kernel._static_key(**make_inputs(padding=1)) == key + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "lengths,mtp_size,head_k_dim,head_v_dim,heads,value_heads,dtype,state_dtype,separate_write", + [ + ((1, 0, 1), 1, 64, 64, 2, 8, torch.bfloat16, torch.bfloat16, False), + ((4, 2, 0), 4, 128, 128, 4, 8, torch.bfloat16, torch.float32, True), + ((2, 3, 1), 3, 64, 80, 2, 2, torch.float16, torch.float32, False), + ((3, 3), 3, 128, 128, 16, 16, torch.bfloat16, torch.float32, False), + ((3, 1), 3, 256, 128, 2, 4, torch.bfloat16, torch.bfloat16, True), + ], +) +def test_all_candidates_match_fp32( + lengths, + mtp_size, + head_k_dim, + head_v_dim, + heads, + value_heads, + dtype, + state_dtype, + separate_write, +): + inputs = make_inputs( + lengths, + mtp_size, + "cuda", + head_k_dim, + head_v_dim, + heads, + value_heads, + dtype, + state_dtype, + separate_write, + ) + output_ref, state_ref = reference(inputs) + state_before = inputs["initial_state"].clone() + kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule + for config in kernel_module.get_test_configs(): + inputs["initial_state"].copy_(state_before) + output, state = kernel(**inputs, run_config=config) + assert state is inputs["initial_state"] + torch.testing.assert_close(output.float(), output_ref, atol=2e-3, rtol=2e-2) + torch.testing.assert_close(state.float(), state_ref.float(), atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dynamic", [False, True]) +def test_autotune_reuses_state_cache_and_graph(monkeypatch, tmp_path, dynamic): + inputs = make_inputs( + lengths=(3, 2, 0) if dynamic else (3, 3), + device="cuda", + padding=2 if dynamic else 0, + ) + kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule + state_before = inputs["initial_state"].clone() + output_ref, state_ref = reference(inputs) + benchmark = kernel._bench + timings = [] + num_configs = len(kernel_module.get_test_configs()) + post_tuning_reference = None + + def checked_bench(*args, **kwargs): + nonlocal post_tuning_reference + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["cu_seqlens"].tolist() == ([0, 3, 6, 7] if dynamic else [0, 3, 6]) + assert rebuilt["initial_state"] is inputs["initial_state"] + elapsed = benchmark(*args, **kwargs) + timings.append(elapsed) + # benchmark 原地更新池内前 num_tokens 个槽，未使用的槽应保持不变。 + assert torch.isfinite(rebuilt["initial_state"]).all() + num_tokens = inputs["q"].shape[1] + assert not torch.equal(inputs["initial_state"][:num_tokens], state_before[:num_tokens]) + torch.testing.assert_close(inputs["initial_state"][num_tokens:], state_before[num_tokens:], atol=0, rtol=0) + if len(timings) == num_configs: + # 最后一次 benchmark 后的状态作为正式执行的初值，不再假定调优前后状态隔离。 + post_tuning_reference = reference(inputs) + return elapsed + + monkeypatch.setattr(kernel, "_bench", checked_bench) + # 普通 warmup 不搜索；decode 调优会原地更新状态，随后按原始请求布局正式执行。 + with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): + kernel(**inputs) + assert not timings + inputs["initial_state"].copy_(state_before) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + output, _ = kernel(**inputs) + assert len(timings) == num_configs and all(math.isfinite(t) for t in timings) + tuned_output_ref, tuned_state_ref = post_tuning_reference + valid_tokens = int(inputs["cu_seqlens"][-1]) + torch.testing.assert_close( + output[:, :valid_tokens].float(), + tuned_output_ref[:, :valid_tokens], + atol=2e-3, + rtol=2e-2, + ) + torch.testing.assert_close(inputs["initial_state"], tuned_state_ref, atol=2e-3, rtol=2e-2) + config_path = next(tmp_path.glob("*.json")) + assert str(kernel._run_key(**inputs)) in json.loads(config_path.read_text()) + + # 关闭已有配置预热后，文件重载和 FORCE 复用都只能正式执行一次，即使在 decode warmup 阶段。 + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr( + autotuner_module, + "get_triton_autotune_level", + lambda: AutotuneLevel.FORCE_AUTOTUNE, + ) + inputs["initial_state"].copy_(state_before) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert len(timings) == num_configs + torch.testing.assert_close(inputs["initial_state"], state_ref, atol=2e-3, rtol=2e-2) + assert not kernel.warmuped_configs_set + # 预热结束后恢复测试请求的初始状态，正常调用及 Graph 回放每次只推进一次。 + inputs["initial_state"].copy_(state_before) + kernel(**inputs) + torch.testing.assert_close(inputs["initial_state"], state_ref, atol=2e-3, rtol=2e-2) + + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output, graph_state = kernel(**inputs) + inputs["initial_state"].copy_(state_before) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + graph_output[:, :valid_tokens].float(), + output_ref[:, :valid_tokens], + atol=2e-3, + rtol=2e-2, + ) + torch.testing.assert_close(graph_state, state_ref, atol=2e-3, rtol=2e-2) + # 同一 Graph 更新变长分组和接受位置，确保回放依赖原始输入而非调优构造的元数据。 + inputs["cu_seqlens"].copy_(torch.tensor([0, 1, 3, 3] if dynamic else [0, 2, 3], device="cuda")) + inputs["num_accepted_tokens"].fill_(2) + inputs["initial_state"].copy_(state_before) + new_output_ref, new_state_ref = reference(inputs) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(graph_output[:, :3].float(), new_output_ref[:, :3], atol=2e-3, rtol=2e-2) + torch.testing.assert_close(graph_state, new_state_ref, atol=2e-3, rtol=2e-2) + assert len(timings) == num_configs + + # 改变序列数和 token 数但仍在同一桶内，decode warmup 应复用已有配置，不再搜索。 + same_bucket_inputs = make_inputs(lengths=(3,), device="cuda") + assert kernel._static_key(**same_bucket_inputs) == kernel._static_key(**inputs) + assert kernel._run_key(**same_bucket_inputs) == kernel._run_key(**inputs) + same_output_ref, same_state_ref = reference(same_bucket_inputs) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + same_output, same_state = kernel(**same_bucket_inputs) + assert len(timings) == num_configs + torch.testing.assert_close(same_output.float(), same_output_ref, atol=2e-3, rtol=2e-2) + torch.testing.assert_close(same_state, same_state_ref, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("level", [0, 1, 2]) +def test_first_decode_loads_history_without_advancing_state_twice(monkeypatch, tmp_path, level): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + inputs = make_inputs(lengths=(3, 3), device="cuda") + # 模拟两个刚完成 prefill 的请求，从各自的 canonical 状态槽开始执行首次 decode。 + inputs["ssm_state_indices"] = torch.arange(6, dtype=torch.int32, device="cuda").view(2, 3) + inputs["ssm_state_write_indices"] = inputs["ssm_state_indices"] + inputs["num_accepted_tokens"].fill_(1) + output_ref, state_ref = reference(inputs) + kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule + config = {"BV": 8, "num_warps": 1, "num_stages": 1} + filename = autotuner_module.KernelConfigs.get_config_file_name(kernel._static_key(**inputs)) + cache_file = tmp_path / filename + cache_file.write_text( + json.dumps( + { + str(kernel._run_key(**inputs)): config, + "24": {"BV": 16, "num_warps": 2, "num_stages": 1}, + } + ) + ) + calls = [] + fn = kernel.fn + + def count_calls(*args, **kwargs): + calls.append(kwargs.get("run_config")) + return fn(*args, **kwargs) + + monkeypatch.setattr(kernel, "fn", count_calls) + assert not Autotuner.is_autotune_warmup() + output, state = kernel(**inputs) + assert calls == [config] + assert state is inputs["initial_state"] + assert not kernel.warmuped_configs_set + torch.testing.assert_close(output.float(), output_ref, atol=2e-3, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py index 357975fbb3..3ff34601d0 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py +++ b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py @@ -35,6 +35,7 @@ def test_int8kv_decode_state_rebuilds_diverse_metadata(monkeypatch): monkeypatch.setattr(int8kv_module, "enable_diverse_mode_gqa_decode_fast_kernel", lambda: True) monkeypatch.setattr(diverse_utils, "get_diverse_max_batch_shared_group_size", lambda: 3) infer_state = SimpleNamespace( + max_kv_seq_len=8192, b_shared_seq_len=torch.tensor([8, 8, 5, 0], dtype=torch.int32, device="cuda"), b_shared_radix_node_id=torch.tensor([10, 10, 20, -1], dtype=torch.int64, device="cuda"), ) diff --git a/unit_tests/common/test_autotuner.py b/unit_tests/common/test_autotuner.py new file mode 100644 index 0000000000..11f96245ad --- /dev/null +++ b/unit_tests/common/test_autotuner.py @@ -0,0 +1,427 @@ +import json +from contextlib import nullcontext + +import pytest +import torch +from frozendict import frozendict + +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") + + +def make_kernel(tmp_path, monkeypatch, name, kernel_type=None, warmup_all_exist_config=True): + calls = [] + benchmarks = [] + options = {} if kernel_type is None else {"kernel_type": kernel_type} + + @autotune( + kernel_name=name, + configs_gen_func=lambda: [{"block": 1}, {"block": 2}], + static_key_func=lambda: {"dtype": "test"}, + run_key_func=lambda size: size, + warmup_all_exist_config=warmup_all_exist_config, + **options, + ) + def kernel(size, run_config=None): + calls.append((size, run_config)) + return run_config + + def bench(size, run_config): + benchmarks.append((size, run_config)) + return 1.0 / run_config["block"] + + cache_dir = tmp_path / name + cache_dir.mkdir() + kernel._cache_dir = str(cache_dir) + monkeypatch.setattr(kernel, "_bench", bench) + return kernel, calls, benchmarks, cache_dir / "configs.json" + + +def test_key_defaults_preserve_explicit_arguments_and_required_checks(): + @autotune( + kernel_name="key_defaults", + configs_gen_func=lambda: [], + static_key_func=lambda window: {"window": window}, + run_key_func=lambda size: size, + ) + def kernel(size, window=(-1, -1), run_config=None): + return run_config + + assert kernel._static_key(8) == {"window": (-1, -1)} + assert kernel._static_key(8, (511, 0)) == {"window": (511, 0)} + assert kernel._static_key(8, window=(1023, 0)) == {"window": (1023, 0)} + assert kernel._static_key(8, window=None) == {"window": None} + with pytest.raises(KeyError, match="Missing argument 'size'"): + kernel._run_key() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_two_warmup_phases_only_persist_matching_kernel_configs(tmp_path, monkeypatch, level): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + general, _, general_benchmarks, general_cache = make_kernel(tmp_path, monkeypatch, "general") + decode, decode_calls, decode_benchmarks, decode_cache = make_kernel( + tmp_path, monkeypatch, "decode", AutotuneKernelType.DECODE_ATTENTION + ) + + with Autotuner.autotune_warmup(): + assert general(8) == {"block": 2} + assert decode(8) is None + assert len(general_benchmarks) == 2 + assert decode_benchmarks == [] + assert decode_calls == [(8, None)] + assert not decode_cache.exists() + general_cache_before = general_cache.read_bytes() + + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert general(16) == {"block": 2} + assert decode(16) == {"block": 2} + assert len(general_benchmarks) == 2 + assert len(decode_benchmarks) == 2 + assert general_cache.read_bytes() == general_cache_before + assert json.loads(general_cache.read_text()) == {"8": {"block": 2}} + assert json.loads(decode_cache.read_text()) == {"16": {"block": 2}} + + +def test_other_phase_warms_history_without_tuning_and_later_uses_new_config(tmp_path, monkeypatch): + decode, calls, benchmarks, cache_file = make_kernel( + tmp_path, monkeypatch, "decode", AutotuneKernelType.DECODE_ATTENTION + ) + cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) + cache_before = cache_file.read_bytes() + + with Autotuner.autotune_warmup(): + assert decode(16) == {"block": 1} + # Historical configs are warmed in either phase, but only the matching phase searches new configs. + assert calls == [(16, {"block": 1}), (16, {"block": 2}), (16, {"block": 1})] + assert benchmarks == [] + assert cache_file.read_bytes() == cache_before + + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert decode(16) == {"block": 2} + assert len(benchmarks) == 2 + assert decode(16) == {"block": 2} + assert json.loads(cache_file.read_text())["16"] == {"block": 2} + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_excluded_kernel_does_not_enter_distributed_tuning(tmp_path, monkeypatch, level): + general, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(autotuner_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(autotuner_module, "get_global_world_size", lambda: 2) + + def unexpected_collective(*args, **kwargs): + pytest.fail("An excluded kernel must not enter autotuning collectives") + + monkeypatch.setattr(general, "_get_autotune_group", unexpected_collective) + monkeypatch.setattr(autotuner_module.dist, "all_gather_object", unexpected_collective) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert general(16) is None + assert calls == [(16, None)] + assert benchmarks == [] + assert not cache_file.exists() + + +@pytest.mark.parametrize("level", [0, 1, 2, 3]) +@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_matching_phase_respects_kernel_autotune_policy(tmp_path, monkeypatch, level, kernel_type): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + kernel, _, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, kernel_type.value, kernel_type) + cache_file.write_text(json.dumps({"16": {"block": 1}})) + cache_before = cache_file.read_bytes() + with Autotuner.autotune_warmup(kernel_type): + result = kernel(16) + if level == AutotuneLevel.FORCE_AUTOTUNE and kernel_type == AutotuneKernelType.GENERAL: + assert result == {"block": 2} + assert len(benchmarks) == 2 + assert json.loads(cache_file.read_text()) == {"16": {"block": 2}} + else: + assert result == (None if level == AutotuneLevel.CLOSE_AUTOTUNE else {"block": 1}) + assert benchmarks == [] + assert cache_file.read_bytes() == cache_before + + +@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_force_autotune_reuses_decode_configs_across_layers(tmp_path, monkeypatch, kernel_type): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + kernel, _, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, kernel_type.value, kernel_type) + with Autotuner.autotune_warmup(kernel_type): + for size in [8, 16]: + benchmarks.clear() + assert kernel(size) == {"block": 2} + assert len(benchmarks) == 2 # A new run key still needs tuning. + for _ in range(3): + assert kernel(size) == {"block": 2} + assert len(benchmarks) == (2 if kernel_type == AutotuneKernelType.DECODE_ATTENTION else 8) + assert json.loads(cache_file.read_text()) == {"8": {"block": 2}, "16": {"block": 2}} + + +def test_history_is_warmed_on_load_or_during_autotune_warmup(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1}), (16, {"block": 2}), (16, {"block": 1})] + calls.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + calls.clear() + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + assert benchmarks == [] + + +def test_repeated_configs_are_skipped_after_success(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + config = {"block": 1, "warps": 4} + cache_file.write_text(json.dumps({"8": config, "16": config, "32": {"block": 2}})) + with Autotuner.autotune_warmup(): + assert kernel(16) == config + assert calls == [(16, config), (16, {"block": 2}), (16, config)] + + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == config + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == config + assert kernel(16) == config + assert calls == [(16, config)] * 3 + assert benchmarks == [] + + +@pytest.mark.parametrize("level", [0, 1, 2]) +@pytest.mark.parametrize("phase", [None, AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_disabled_history_warmup_loads_cache_without_extra_state_updates(tmp_path, monkeypatch, level, phase): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + calls = [] + + @autotune( + kernel_name="stateful_kernel", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=lambda: [{"block": 1}, {"block": 2}], + static_key_func=lambda: {}, + run_key_func=lambda state: state.numel(), + warmup_all_exist_config=False, + ) + def kernel(state, run_config=None): + calls.append(run_config) + state.add_(run_config["block"]) + return state + + kernel._cache_dir = str(tmp_path) + cache_file = tmp_path / "configs.json" + cache_file.write_text(json.dumps({"4": {"block": 1}, "8": {"block": 2}})) + cache_before = cache_file.read_bytes() + + def unexpected_warmup(*args, **kwargs): + pytest.fail("Disabling history warmup must skip its execution and argument cloning") + + monkeypatch.setattr(kernel, "kernel_warmup", unexpected_warmup) + monkeypatch.setattr(kernel, "_mutate_args_clone", unexpected_warmup) + state = torch.zeros(4) + scope = nullcontext() if phase is None else Autotuner.autotune_warmup(phase) + with scope: + for _ in range(2): + assert kernel(state) is state + assert calls == [{"block": 1}, {"block": 1}] + torch.testing.assert_close(state, torch.full((4,), 2.0)) + assert cache_file.read_bytes() == cache_before + assert not kernel.warmuped_configs_set + + +@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_disabled_history_warmup_still_searches_new_configs(tmp_path, monkeypatch, kernel_type): + kernel, calls, benchmarks, cache_file = make_kernel( + tmp_path, monkeypatch, kernel_type.value, kernel_type, warmup_all_exist_config=False + ) + with Autotuner.autotune_warmup(kernel_type): + for size in [8, 16]: + calls.clear() + assert kernel(size) == {"block": 2} + assert calls == [(size, {"block": 2})] + calls.clear() + assert kernel(size) == {"block": 2} + assert calls == [(size, {"block": 2})] + assert len(benchmarks) == 4 + assert json.loads(cache_file.read_text()) == {"8": {"block": 2}, "16": {"block": 2}} + assert not kernel.warmuped_configs_set + + +def test_failed_warmup_retries_only_during_autotune_warmup(tmp_path, monkeypatch): + kernel, calls, _, cache_file = make_kernel(tmp_path, monkeypatch, "general") + cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) + + def shape_sensitive_kernel(size, run_config=None): + calls.append((size, run_config)) + if size < 16 and run_config["block"] == 2: + raise RuntimeError("This config requires a larger input") + return run_config + + monkeypatch.setattr(kernel, "fn", shape_sensitive_kernel) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(8) == {"block": 1} + assert calls == [(8, {"block": 1}), (8, {"block": 2}), (8, {"block": 1})] + calls.clear() + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + calls.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 2}), (16, {"block": 1})] + calls.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + + +def test_new_configs_are_warmed_on_later_calls_and_state_is_not_persisted(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + cache_file.write_text(json.dumps({"8": {"block": 1}})) + with Autotuner.autotune_warmup(): + assert kernel(8) == {"block": 1} + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2})] + assert len(benchmarks) == 2 + assert json.loads(cache_file.read_text()) == {"8": {"block": 1}, "16": {"block": 2}} + + calls.clear() + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2})] + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2}), (16, {"block": 2})] + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2})] + + # A new autotuner loading the same file must warm its configs again. + reloaded, reload_calls, _, _ = make_kernel(tmp_path, monkeypatch, "reloaded") + reloaded._cache_dir = str(cache_file.parent) + assert reloaded(16) == {"block": 2} + assert len(reload_calls) == 3 + assert {call[1]["block"] for call in reload_calls[:2]} == {1, 2} + reload_calls.clear() + with Autotuner.autotune_warmup(): + assert reloaded(16) == {"block": 2} + assert reload_calls == [(16, {"block": 2})] + + +def test_warmup_preserves_mutated_input_and_skips_repeated_execution(): + executions = [] + + @autotune( + kernel_name="mutating_kernel", + configs_gen_func=lambda: [{"block": 1}], + static_key_func=lambda: {}, + run_key_func=lambda state: state.numel(), + mutates_args=["state"], + ) + def kernel(state, run_config=None): + executions.append(run_config) + state.add_(run_config["block"]) + + state = torch.zeros(4) + static_key = frozendict({}) + kernel.kernel_warmup(static_key, state, run_config={"block": 1}) + torch.testing.assert_close(state, torch.zeros(4)) + + kernel.kernel_warmup(static_key, state, run_config={"block": 1}) + assert executions == [{"block": 1}] + torch.testing.assert_close(state, torch.zeros(4)) + + +def test_explicit_config_bypasses_tuning(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + with Autotuner.autotune_warmup(): + assert kernel(16, run_config={"block": 3}) == {"block": 3} + assert calls == [(16, {"block": 3})] + assert benchmarks == [] + assert not cache_file.exists() + + +@pytest.mark.parametrize("keyword_input", [False, True]) +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_rebuild_inputs_only_for_benchmarking(tmp_path, monkeypatch, keyword_input, level): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + rebuilds = [] + executions = [] + benchmarks = [] + original = torch.zeros(2) + + def rebuild(state, size, run_config=None): + rebuilds.append(state) + return (torch.ones(size),), {"size": size} + + @autotune( + kernel_name="rebuild_inputs", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=lambda: [{"block": 1}, {"block": 2}], + static_key_func=lambda state: {"input_size": state.numel()}, + run_key_func=lambda size: size, + rebuild_input_func=rebuild, + ) + def kernel(state, size, run_config=None): + executions.append(state) + return state + + def bench(state, size, run_config): + benchmarks.append(state) + torch.testing.assert_close(state, torch.ones(8)) + assert size == 8 + return 1.0 / run_config["block"] + + kernel._cache_dir = str(tmp_path) + monkeypatch.setattr(kernel, "_bench", bench) + args, kwargs = ((), {"state": original, "size": 8}) if keyword_input else ((original, 8), {}) + + assert kernel(*args, **kwargs) is original + with Autotuner.autotune_warmup(): + assert kernel(*args, **kwargs) is original + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(*args, **kwargs, run_config={"block": 3}) is original + assert not rebuilds + assert kernel(*args, **kwargs) is original + assert kernel(*args, **kwargs) is original # Exact cache hit only warms historical configs. + + assert kernel(*args, **kwargs) is original + assert len(rebuilds) == 1 and rebuilds[0] is original + assert len(benchmarks) == 2 and benchmarks[0] is benchmarks[1] + assert all(state is original for state in executions) + assert kernel.cached_configs == {frozendict({"input_size": 2}): {"8": {"block": 2}}} + torch.testing.assert_close(original, torch.zeros(2)) + + +def test_default_api_and_nested_phase_restore_after_exception(): + assert not Autotuner.is_autotune_warmup() + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + Autotuner.start_autotune_warmup() + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + with pytest.raises(RuntimeError, match="test failure"): + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert Autotuner.is_autotune_warmup() + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + with Autotuner.autotune_warmup(): + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + raise RuntimeError("test failure") + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + Autotuner.end_autotune_warmup() + assert not Autotuner.is_autotune_warmup() + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) diff --git a/unit_tests/utils/test_sgl_utils.py b/unit_tests/utils/test_sgl_utils.py new file mode 100644 index 0000000000..2d326acbe0 --- /dev/null +++ b/unit_tests/utils/test_sgl_utils.py @@ -0,0 +1,278 @@ +import collections +import inspect +import math +from types import SimpleNamespace + +import pytest +import torch +from frozendict import frozendict + +from lightllm.common.basemodel.attention.fa3 import fp as fa3_module +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner +from lightllm.utils import sgl_utils +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_seq_len_environment(monkeypatch): + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("seq_len, expected_len", [(None, 32768), ("8192", 8192), ("16384", 16384), ("8193", 8704)]) +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_fa3_run_key_uses_configured_length_during_tuning(monkeypatch, seq_len, expected_len, level): + if seq_len is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", seq_len) + monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + for graph_capacity in [16, 32768]: + page_table = torch.empty(2, graph_capacity, dtype=torch.int32) + assert ( + sgl_utils._flash_attn_kvcache_run_key(page_table, 3, 2) + == 2 * 10_000_000_000_000 + 3 * 10_000_000 + expected_len + ) + + +@pytest.mark.parametrize( + "phase, level", + [ + (None, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_fa3_run_key_uses_actual_length_without_tuning(monkeypatch, phase, level): + # 未调优时不应读取该环境变量，即使它无效也不影响正常配置查找。 + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: level) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + page_table = torch.empty(2, 32768, dtype=torch.int32) + for actual_len, expected_len in [(1, 512), (512, 512), (513, 1024), (8192, 8192), (16384, 16384)]: + assert ( + sgl_utils._flash_attn_kvcache_run_key(page_table, 3, actual_len) + == 2 * 10_000_000_000_000 + 3 * 10_000_000 + expected_len + ) + + +def test_fa3_runtime_matches_cached_config_by_actual_length(monkeypatch): + kernel = sgl_utils.flash_attn_with_kvcache_autotune + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(kernel, "fn", lambda **kwargs: kwargs["run_config"]) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + inputs = dict( + q=torch.empty(2, 4, 8), + k_cache=torch.empty(1, 1, 2, 8), + v_cache=torch.empty(1, 1, 2, 8), + page_table=torch.empty(2, 32768, dtype=torch.int32), + max_seqlen_q=1, + causal=True, + window_size=(-1, -1), + softcap=0.0, + sinks=None, + k_descale=None, + v_descale=None, + ) + static_key = frozendict(kernel._static_key(**inputs)) + base_key = 2 * 10_000_000_000_000 + 10_000_000 + monkeypatch.setattr( + kernel, + "cached_configs", + {static_key: {str(base_key + 8192): {"num_splits": 16}, str(base_key + 16384): {"num_splits": 32}}}, + ) + + assert kernel(**inputs, max_seqlen_k=8192) == {"num_splits": 16} + assert kernel(**inputs, max_seqlen_k=8191) == {"num_splits": 16} + assert kernel(**inputs, max_seqlen_k=15000) == {"num_splits": 32} + assert set(kernel.fast_match_configs[static_key]) == {str(base_key + 8192), str(base_key + 15360)} + + +def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch): + graph = SimpleNamespace(can_run=lambda **kwargs: True, graph_max_len_in_batch=32768) + model = SimpleNamespace( + graph=graph, + req_manager=SimpleNamespace(req_to_token_indexs=torch.zeros(2, 8192, dtype=torch.int32)), + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), + ) + state = fa3_module.Fa3DecodeAttState( + backend=SimpleNamespace( + model=model, + uses_causal_attention=lambda: True, + uses_dynamic_spec_verify_layout=lambda: False, + get_page_table_view=lambda att_batch_size, max_kv_len, microbatch_index: torch.zeros( + att_batch_size, max_kv_len, dtype=torch.int32 + ), + ), + infer_state=SimpleNamespace( + batch_size=2, + max_kv_seq_len=8192, + microbatch_index=0, + b_req_idx=torch.tensor([0, 1], dtype=torch.int32), + b_seq_len=torch.tensor([4096, 8192], dtype=torch.int32), + b1_cu_q_seq_len=torch.tensor([0, 1, 2], dtype=torch.int32), + b1_cu_kv_seq_len=torch.tensor([0, 4096, 12288], dtype=torch.int32), + ), + ) + monkeypatch.setattr(fa3_module, "page_table_copy", lambda **kwargs: None) + state.init_state() + # 模拟 CudaGraph._capture_decode 对 infer_state 的修改。 + state.infer_state.max_kv_seq_len = graph.graph_max_len_in_batch + calls = [] + + def attention(**kwargs): + calls.append(kwargs) + return kwargs["q"] + + monkeypatch.setattr(fa3_module, "flash_attn_with_kvcache_autotune", attention) + state.decode_att(torch.empty(2, 4, 8), torch.empty(1, 2, 8), torch.empty(1, 2, 8)) + assert calls[0]["max_seqlen_k"] == 8192 + assert calls[0]["page_table"].shape == (2, 32768) + + +@pytest.mark.parametrize("seq_len", ["0", "-1", "invalid"]) +def test_invalid_decode_autotune_seq_len_is_rejected(monkeypatch, seq_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", seq_len) + with pytest.raises(ValueError): + get_decode_attn_autotune_seq_len() + + +@pytest.mark.parametrize("page_size, num_pages", [(1, 64), (1, 3), (256, 3)]) +@pytest.mark.parametrize("query_lengths", [[1, 1], [3, 3], [0, 3]]) +def test_fa3_rebuilds_valid_kv_metadata_without_changing_originals(monkeypatch, page_size, num_pages, query_lengths): + kv_len = 17 + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + max_pages = (kv_len + page_size - 1) // page_size + q = torch.randn(sum(query_lengths), 4, 8) + k = torch.randn(num_pages, page_size, 2, 8) + v = torch.randn_like(k) + page_table = torch.full((2, 8), -1, dtype=torch.int32) + seq_lens = torch.full((2,), 2, dtype=torch.int32) + cu_q = torch.tensor([0, query_lengths[0], sum(query_lengths)], dtype=torch.int32) + original = (q, k, v, seq_lens, page_table, cu_q, max(query_lengths), 2) + snapshots = [tensor.clone() for tensor in original[:6]] + options = { + "cu_seqlens_k_new": None, + "causal": True, + "window_size": (16, 0), + "softmax_scale": 0.25, + "return_softmax_lse": False, + } + + args, kwargs = sgl_utils._flash_attn_kvcache_rebuild_inputs(*original, **options) + + assert all(args[i] is original[i] for i in [0, 1, 2, 5]) + assert args[6:] == (max(query_lengths), kv_len) + assert kwargs == options + assert args[4].shape == (2, max_pages) + assert args[4].min() >= 0 and args[4].max() < num_pages + if num_pages >= 2 * max_pages: + assert args[4].unique().numel() == 2 * max_pages + torch.testing.assert_close(args[3], torch.full_like(seq_lens, kv_len)) + for tensor, snapshot in zip(original[:6], snapshots): + torch.testing.assert_close(tensor, snapshot) + + # Batched Q does not require cumulative query or KV lengths. + batched_q = torch.randn(2, 1, 4, 8) + args, kwargs = sgl_utils._flash_attn_kvcache_rebuild_inputs( + q=batched_q, + k_cache=k, + v_cache=v, + page_table=page_table, + cache_seqlens=seq_lens, + cu_seqlens_q=None, + max_seqlen_q=1, + max_seqlen_k=2, + ) + assert args[0] is batched_q + assert args[5:] == (None, 1, kv_len) + + +@pytest.mark.parametrize("kv_len", [8192, 16384]) +@pytest.mark.parametrize("query_lengths", [[1, 1], [3, 3], [0, 3]]) +def test_fa3_autotunes_long_kv_then_captures_original_inputs(tmp_path, monkeypatch, kv_len, query_lengths): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("FA3 requires a Hopper GPU") + if sgl_utils.flash_attn_with_kvcache is None: + pytest.skip("sgl_kernel FA3 is unavailable") + + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + kernel = sgl_utils.flash_attn_with_kvcache_autotune + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + + q = torch.randn(sum(query_lengths), 4, 64, device="cuda", dtype=torch.bfloat16) + k = torch.randn(2 * kv_len, 1, 2, 64, device="cuda", dtype=q.dtype) + v = torch.randn_like(k) + inputs = dict( + q=q, + k_cache=k, + v_cache=v, + page_table=torch.zeros(2, 32768, device="cuda", dtype=torch.int32), + cache_seqlens=torch.full((2,), 2, device="cuda", dtype=torch.int32), + cu_seqlens_q=torch.tensor([0, query_lengths[0], sum(query_lengths)], device="cuda", dtype=torch.int32), + cu_seqlens_k_new=None, + max_seqlen_q=max(query_lengths), + max_seqlen_k=2, + causal=True, + window_size=(-1, -1), + softcap=0.0, + sinks=None, + k_descale=None, + v_descale=None, + ) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + reference = kernel.fn(**inputs) + benchmark = kernel._bench + timings = [] + + def checked_bench(*args, **kwargs): + bound = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert bound["cache_seqlens"].tolist() == [kv_len, kv_len] + assert bound["page_table"].shape == (2, kv_len) + assert bound["cu_seqlens_k_new"] is None + assert bound["max_seqlen_k"] == kv_len + assert bound["q"] is q + elapsed = benchmark(*args, **kwargs) + assert math.isfinite(elapsed), f"FA3 benchmark failed for {kwargs['run_config']}" + timings.append(elapsed) + return elapsed + + monkeypatch.setattr(kernel, "_bench", checked_bench) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + output = kernel(**inputs) + assert len(timings) == 3 + expected_key = str(2 * 10_000_000_000_000 + max(query_lengths) * 10_000_000 + kv_len) + assert kernel._run_key(**inputs) == 2 * 10_000_000_000_000 + max(query_lengths) * 10_000_000 + 512 + assert all(list(configs) == [expected_key] for configs in kernel.cached_configs.values()) + torch.testing.assert_close(output, reference) + + def unexpected_rebuild(*args, **kwargs): + pytest.fail("Cached execution and CUDA Graph capture must use the original inputs") + + monkeypatch.setattr(kernel, "rebuild_input_func", unexpected_rebuild) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = kernel(**inputs) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured_output, reference) + for name, snapshot in snapshots.items(): + torch.testing.assert_close(inputs[name], snapshot) diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py index aaa3e05268..fa44f6c5a6 100644 --- a/unit_tests/utils/test_speculative_utils.py +++ b/unit_tests/utils/test_speculative_utils.py @@ -165,6 +165,7 @@ def test_triton_mtp_decode_state_builds_group_markers(): backend=SimpleNamespace(model=model), infer_state=SimpleNamespace( b_req_idx=torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda"), + max_kv_seq_len=4, ), )