perf: batch KV page allocation and fuse token and KV index gathering

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

@@ -3,7 +3,6 @@
 import torch
 import triton
 from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
-from typing import Optional, TYPE_CHECKING
 from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune
 from lightllm.common.basemodel.triton_kernel.fa3_utils import (
     build_dynamic_spec_fa3_decode_params,

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

@@ -3,7 +3,7 @@
 import torch
 import triton
 from ..base_att import BasePrefillAttState, BaseDecodeAttState, AttControl
-from typing import Optional, TYPE_CHECKING, Tuple
+from typing import Tuple
 from lightllm.utils.sgl_utils import flash_attn_with_kvcache
 from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params, page_table_copy
 from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor

lightllm/common/basemodel/batch_objs.py

@@ -32,6 +32,8 @@ class ModelInput:
     # Decode 逐行携带的 radix node 标识。相同 id 表示请求引用同一个共享
     # radix node；该 id 只用于重建 diverse attention 的 b_mark_shared_group。
     b_shared_radix_node_id: torch.Tensor = None
+    # GPU write slots are resolved after padding, from the final row positions.
+    mem_indexes: torch.Tensor = None
     is_prefill: bool = False
     b_ready_cache_len: torch.Tensor = None
     # Request/row-aligned MRoPE position offset. It is decode-only; prefill
@@ -189,9 +191,6 @@ class ModelOutput:
     # attached here. ModelOutput therefore owns a stable output view instead
     # of the mutable collector used while the forward is still running.
     mtp_collector: Optional[ModelMtpOutputCollector] = None
-    # 用于判断 mem_indexes 是否成功写入 req manager 中的事件对象。
-    prefill_mem_indexes_ready_event: torch.Event = None
-
     # prompt_logics 用于在开启 return_all_prompt_logics 模式（如 enable_prompt_logprobs）时，
     # 保存整个 prefill 阶段每一个 token 位置对应的 logits（而非仅最后一个位置的 logits）。
     # 此时 logits 依然只保存每个请求最后一个位置的 logits，prompt_logics 为可选项，仅在

lightllm/common/basemodel/cuda_graph.py

@@ -1,4 +1,3 @@
-import os
 import torch
 import torch.distributed as dist
 import copy
@@ -8,13 +7,10 @@
 from typing import Optional
 from lightllm.utils.log_utils import init_logger
 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.basemodel.triton_kernel.copy_kv_index_to_req import build_kv_indexes_and_input_ids
 from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType
-from lightllm.utils.torch_memory_saver_utils import (
-    TorchMemorySaverWrapper,
-    MemoryTag,
-)
+from lightllm.utils.torch_memory_saver_utils import TorchMemorySaverWrapper
 from .infer_struct import InferStateInfo
 
 
@@ -286,6 +282,12 @@ def warmup(self, model):
                 multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)],
                 **model._gen_special_model_input(batch_size),
             )
+            # Warm the ordinary decode gather path as well as the explicit-ID path below.
+            _, model_input.input_ids = build_kv_indexes_and_input_ids(
+                model_input,
+                model.req_manager.req_to_token_indexs,
+                model.req_manager.req_sampling_params_manager.req_to_next_token_ids,
+            )
             model_output: ModelOutput = model.forward(model_input)
             del model_output
             del input_ids
@@ -344,6 +346,11 @@ def warmup_overlap(self, model):
                     multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)],
                     **model._gen_special_model_input(batch_size),
                 )
+                _, micro_batch.input_ids = build_kv_indexes_and_input_ids(
+                    micro_batch,
+                    model.req_manager.req_to_token_indexs,
+                    model.req_manager.req_sampling_params_manager.req_to_next_token_ids,
+                )
                 decode_batches.append(micro_batch)
                 del micro_batch
 

lightllm/common/basemodel/prefill_cuda_graph.py

@@ -1,14 +1,11 @@
-import os
 import torch
-import copy
 import bisect
 import triton
 from typing import List, Tuple
 from typing import Optional
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
-from lightllm.distributed import dist_group_manager
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
 from .infer_struct import InferStateInfo
 from .cuda_graph import CudaGraph