Wzj demo

lightllm/common/basemodel/triton_kernel/pad_decode_inputs.py

@@ -0,0 +1,127 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _pad_decode_inputs_kernel(
+    input_ids,
+    b_req_idx,
+    b_mtp_index,
+    b_seq_len,
+    b_position_delta,
+    b_shared_seq_len,
+    b_shared_radix_node_id,
+    padded_input_ids,
+    padded_b_req_idx,
+    padded_b_mtp_index,
+    padded_b_seq_len,
+    padded_b_position_delta,
+    padded_b_shared_seq_len,
+    padded_b_shared_radix_node_id,
+    batch_size,
+    new_batch_size,
+    hold_request_id,
+    HAS_B_POSITION_DELTA: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    source_mask = offsets < batch_size
+    output_mask = offsets < new_batch_size
+
+    # 保留已有请求，新增 HOLD 请求使用 token id 1、seq_len 2。
+    token_ids = tl.load(input_ids + offsets, mask=source_mask, other=1)
+    tl.store(padded_input_ids + offsets, token_ids, mask=output_mask)
+
+    req_idx = tl.load(b_req_idx + offsets, mask=source_mask, other=hold_request_id)
+    tl.store(padded_b_req_idx + offsets, req_idx, mask=output_mask)
+
+    mtp_index = tl.load(b_mtp_index + offsets, mask=source_mask, other=0)
+    tl.store(padded_b_mtp_index + offsets, mtp_index, mask=output_mask)
+
+    seq_len = tl.load(b_seq_len + offsets, mask=source_mask, other=2)
+    tl.store(padded_b_seq_len + offsets, seq_len, mask=output_mask)
+
+    # 位置偏移按需补齐；HOLD 请求不携带位置偏移和共享缓存。
+    if HAS_B_POSITION_DELTA:
+        position_delta = tl.load(b_position_delta + offsets, mask=source_mask, other=0)
+        tl.store(padded_b_position_delta + offsets, position_delta, mask=output_mask)
+
+    shared_seq_len = tl.load(b_shared_seq_len + offsets, mask=source_mask, other=0)
+    tl.store(padded_b_shared_seq_len + offsets, shared_seq_len, mask=output_mask)
+
+    shared_radix_node_id = tl.load(b_shared_radix_node_id + offsets, mask=source_mask, other=-1)
+    tl.store(padded_b_shared_radix_node_id + offsets, shared_radix_node_id, mask=output_mask)
+
+
+@torch.no_grad()
+def pad_decode_inputs(
+    input_ids: torch.Tensor,
+    b_req_idx: torch.Tensor,
+    b_mtp_index: torch.Tensor,
+    b_seq_len: torch.Tensor,
+    b_position_delta: Optional[torch.Tensor],
+    b_shared_seq_len: torch.Tensor,
+    b_shared_radix_node_id: torch.Tensor,
+    new_batch_size: int,
+    hold_request_id: int,
+):
+    """一次 kernel 调用将 decode Token 和请求元数据补齐到 new_batch_size。
+
+    调用方保证输入为同一 CUDA 设备上的连续一维 Tensor，各字段按请求对齐，
+    且 new_batch_size 大于现有请求数。b_position_delta 为 None 时保持 None。
+    返回顺序与七个 Tensor 输入一致。
+    """
+    batch_size = b_req_idx.shape[0]
+
+    padded_input_ids = torch.empty((new_batch_size,), dtype=input_ids.dtype, device=input_ids.device)
+    padded_b_req_idx = torch.empty((new_batch_size,), dtype=b_req_idx.dtype, device=b_req_idx.device)
+    padded_b_mtp_index = torch.empty((new_batch_size,), dtype=b_mtp_index.dtype, device=b_mtp_index.device)
+    padded_b_seq_len = torch.empty((new_batch_size,), dtype=b_seq_len.dtype, device=b_seq_len.device)
+    padded_b_position_delta = None
+    if b_position_delta is not None:
+        padded_b_position_delta = torch.empty(
+            (new_batch_size,), dtype=b_position_delta.dtype, device=b_position_delta.device
+        )
+    padded_b_shared_seq_len = torch.empty(
+        (new_batch_size,), dtype=b_shared_seq_len.dtype, device=b_shared_seq_len.device
+    )
+    padded_b_shared_radix_node_id = torch.empty(
+        (new_batch_size,), dtype=b_shared_radix_node_id.dtype, device=b_shared_radix_node_id.device
+    )
+
+    block_size = 256
+    grid = (triton.cdiv(new_batch_size, block_size),)
+    _pad_decode_inputs_kernel[grid](
+        input_ids=input_ids,
+        b_req_idx=b_req_idx,
+        b_mtp_index=b_mtp_index,
+        b_seq_len=b_seq_len,
+        b_position_delta=b_position_delta,
+        b_shared_seq_len=b_shared_seq_len,
+        b_shared_radix_node_id=b_shared_radix_node_id,
+        padded_input_ids=padded_input_ids,
+        padded_b_req_idx=padded_b_req_idx,
+        padded_b_mtp_index=padded_b_mtp_index,
+        padded_b_seq_len=padded_b_seq_len,
+        padded_b_position_delta=padded_b_position_delta,
+        padded_b_shared_seq_len=padded_b_shared_seq_len,
+        padded_b_shared_radix_node_id=padded_b_shared_radix_node_id,
+        batch_size=batch_size,
+        new_batch_size=new_batch_size,
+        hold_request_id=hold_request_id,
+        HAS_B_POSITION_DELTA=b_position_delta is not None,
+        BLOCK_SIZE=block_size,
+        num_warps=4,
+    )
+    return (
+        padded_input_ids,
+        padded_b_req_idx,
+        padded_b_mtp_index,
+        padded_b_seq_len,
+        padded_b_position_delta,
+        padded_b_shared_seq_len,
+        padded_b_shared_radix_node_id,
+    )

lightllm/common/basemodel/triton_kernel/pad_prefill_inputs.py

@@ -0,0 +1,113 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _pad_prefill_inputs_kernel(
+    input_ids,
+    b_req_idx,
+    b_mtp_index,
+    b_seq_len,
+    b_ready_cache_len,
+    b_is_decode_req,
+    padded_input_ids,
+    padded_b_req_idx,
+    padded_b_mtp_index,
+    padded_b_seq_len,
+    padded_b_ready_cache_len,
+    padded_b_is_decode_req,
+    handle_token_num,
+    padded_token_num,
+    batch_size,
+    hold_request_id,
+    BLOCK_SIZE: tl.constexpr,
+):
+    block_start = tl.program_id(0) * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+    # Token 按实际补齐数量扩展，新增位置填入 token id 1。
+    token_ids = tl.load(input_ids + offsets, mask=offsets < handle_token_num, other=1)
+    tl.store(padded_input_ids + offsets, token_ids, mask=offsets < handle_token_num + padded_token_num)
+
+    # 请求元数据只追加一个 HOLD 请求；后续 Token block 不再处理元数据。
+    if block_start < batch_size + 1:
+        source_mask = offsets < batch_size
+        output_mask = offsets < batch_size + 1
+
+        req_idx = tl.load(b_req_idx + offsets, mask=source_mask, other=hold_request_id)
+        tl.store(padded_b_req_idx + offsets, req_idx, mask=output_mask)
+
+        mtp_index = tl.load(b_mtp_index + offsets, mask=source_mask, other=0)
+        tl.store(padded_b_mtp_index + offsets, mtp_index, mask=output_mask)
+
+        seq_len = tl.load(b_seq_len + offsets, mask=source_mask, other=padded_token_num)
+        tl.store(padded_b_seq_len + offsets, seq_len, mask=output_mask)
+
+        ready_cache_len = tl.load(b_ready_cache_len + offsets, mask=source_mask, other=0)
+        tl.store(padded_b_ready_cache_len + offsets, ready_cache_len, mask=output_mask)
+
+        is_decode_req = tl.load(b_is_decode_req + offsets, mask=source_mask, other=False)
+        tl.store(padded_b_is_decode_req + offsets, is_decode_req, mask=output_mask)
+
+
+@torch.no_grad()
+def pad_prefill_inputs(
+    input_ids: torch.Tensor,
+    b_req_idx: torch.Tensor,
+    b_mtp_index: torch.Tensor,
+    b_seq_len: torch.Tensor,
+    b_ready_cache_len: torch.Tensor,
+    b_is_decode_req: torch.Tensor,
+    new_handle_token_num: int,
+    hold_request_id: int,
+):
+    """一次 kernel 调用补齐 prefill Token，并追加一个 HOLD 请求。
+
+    调用方保证输入为同一 CUDA 设备上的连续一维 Tensor，元数据按请求对齐，
+    且 new_handle_token_num 大于现有 Token 数。返回顺序与六个 Tensor 输入一致。
+    """
+    handle_token_num = input_ids.shape[0]
+    padded_token_num = new_handle_token_num - handle_token_num
+    batch_size = b_req_idx.shape[0]
+    new_batch_size = batch_size + 1
+
+    padded_input_ids = torch.empty((new_handle_token_num,), dtype=input_ids.dtype, device=input_ids.device)
+    padded_b_req_idx = torch.empty((new_batch_size,), dtype=b_req_idx.dtype, device=b_req_idx.device)
+    padded_b_mtp_index = torch.empty((new_batch_size,), dtype=b_mtp_index.dtype, device=b_mtp_index.device)
+    padded_b_seq_len = torch.empty((new_batch_size,), dtype=b_seq_len.dtype, device=b_seq_len.device)
+    padded_b_ready_cache_len = torch.empty(
+        (new_batch_size,), dtype=b_ready_cache_len.dtype, device=b_ready_cache_len.device
+    )
+    padded_b_is_decode_req = torch.empty((new_batch_size,), dtype=b_is_decode_req.dtype, device=b_is_decode_req.device)
+
+    block_size = 256
+    grid = (triton.cdiv(max(new_handle_token_num, new_batch_size), block_size),)
+    _pad_prefill_inputs_kernel[grid](
+        input_ids=input_ids,
+        b_req_idx=b_req_idx,
+        b_mtp_index=b_mtp_index,
+        b_seq_len=b_seq_len,
+        b_ready_cache_len=b_ready_cache_len,
+        b_is_decode_req=b_is_decode_req,
+        padded_input_ids=padded_input_ids,
+        padded_b_req_idx=padded_b_req_idx,
+        padded_b_mtp_index=padded_b_mtp_index,
+        padded_b_seq_len=padded_b_seq_len,
+        padded_b_ready_cache_len=padded_b_ready_cache_len,
+        padded_b_is_decode_req=padded_b_is_decode_req,
+        handle_token_num=handle_token_num,
+        padded_token_num=padded_token_num,
+        batch_size=batch_size,
+        hold_request_id=hold_request_id,
+        BLOCK_SIZE=block_size,
+        num_warps=4,
+    )
+    return (
+        padded_input_ids,
+        padded_b_req_idx,
+        padded_b_mtp_index,
+        padded_b_seq_len,
+        padded_b_ready_cache_len,
+        padded_b_is_decode_req,
+    )

lightllm/common/cuda_wrapper.py

@@ -16,13 +16,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""This file is a pure Python wrapper for the cudart library.
-It avoids the need to compile a separate shared library, and is
-convenient for use when we just need to call a few functions.
-"""
+"""Python cudart wrapper with asynchronous pinned-source lifetime management."""
 
 import ctypes
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
+from threading import Lock
 from typing import Any, Dict, List, Optional
 
 # this line makes it possible to directly load `libcudart.so` using `ctypes`
@@ -40,6 +39,125 @@
 cudaMemcpyKind = ctypes.c_int
 
 
+class PinnedH2DBatchCopier:
+    """每个推理进程共享批量 H2D 状态，统一管理异步源保活与异常回退。"""
+
+    _native_batch_available = True
+    # 原生 batch 禁止 legacy 默认流，仅在默认流调用时按设备延迟创建上传流。
+    _upload_streams: Dict[torch.device, torch.cuda.Stream] = {}
+    # 只在原生 H2D 成功提交、需要后台保活时创建线程池，导入和 copy_ 回退均不创建。
+    _release_pool: Optional[ThreadPoolExecutor] = None
+    _release_pool_lock = Lock()
+
+    @classmethod
+    def copy_async(cls, dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor]):
+        """按当前流顺序批量异步 H2D，并保护 pinned 源直到传输结束。
+
+        源由 PyTorch pinned allocator 申请，各对 Tensor 连续且形状和 dtype 一致。
+        默认流的原生 batch 适配封装在内部；调用方只遵循当前流顺序，不阻塞 CPU。
+        上传完成前不能主动改写源，目标须保持有效到传输结束。
+        """
+        if not dst_gpu_tensors:
+            return
+        stream = torch.cuda.current_stream()
+        if cls._native_batch_available and stream.cuda_stream == 0:
+            # 原生 API 不接受 legacy 默认流；上传流先等待当前流，提交后再交回依赖。
+            upload_stream = cls._get_upload_stream(stream.device)
+            upload_stream.wait_stream(stream)
+            with torch.cuda.stream(upload_stream):
+                cls._copy_on_stream(dst_gpu_tensors, src_pinned_tensors, upload_stream)
+            stream.wait_stream(upload_stream)
+        else:
+            cls._copy_on_stream(dst_gpu_tensors, src_pinned_tensors, stream)
+
+    @classmethod
+    def _get_upload_stream(cls, device: torch.device) -> torch.cuda.Stream:
+        if device not in cls._upload_streams:
+            cls._upload_streams[device] = torch.cuda.Stream(device=device)
+        return cls._upload_streams[device]
+
+    @classmethod
+    def _copy_on_stream(
+        cls, dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor], stream: torch.cuda.Stream
+    ):
+        """在已选定的当前流提交，沿用原生保活和 PyTorch 回退。"""
+        if cls._native_batch_available and cls._try_native_batch_h2d(dst_gpu_tensors, src_pinned_tensors, stream):
+            cls._protect_pinned_sources(src_pinned_tensors, stream)
+            return
+
+        # PyTorch copy_ 已经登记 host allocator，回退路径无需额外保活或登记。
+        for dst, src in zip(dst_gpu_tensors, src_pinned_tensors):
+            dst.copy_(src, non_blocking=True)
+
+    @classmethod
+    def _try_native_batch_h2d(
+        cls, dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor], stream: torch.cuda.Stream
+    ) -> bool:
+        """尝试原生批量提交；接口准备或提交失败时禁用该路径，让调用方回退到 copy_。"""
+        try:
+            from cuda.bindings import runtime as cuda_runtime
+
+            attrs = cuda_runtime.cudaMemcpyAttributes()
+            # Stream 模式让源读取遵循上传流顺序；提交成功后由完成事件和后台任务保活。
+            attrs.srcAccessOrder = cuda_runtime.cudaMemcpySrcAccessOrder.cudaMemcpySrcAccessOrderStream
+            dst_ptrs = [tensor.data_ptr() for tensor in dst_gpu_tensors]
+            src_ptrs = [tensor.data_ptr() for tensor in src_pinned_tensors]
+            byte_sizes = [tensor.numel() * tensor.element_size() for tensor in src_pinned_tensors]
+            batch_copy = cuda_runtime.cudaMemcpyBatchAsync
+        except Exception as error:
+            cls._native_batch_available = False
+            logger.warning(
+                "Native batch H2D API unavailable; "
+                f"disabling native batch copy and falling back to PyTorch copy_: {error}"
+            )
+            return False
+
+        try:
+            (error,) = batch_copy(dst_ptrs, src_ptrs, byte_sizes, len(byte_sizes), [attrs], [0], 1, stream.cuda_stream)
+            if error != cuda_runtime.cudaError_t.cudaSuccess:
+                raise RuntimeError(f"cudaMemcpyBatchAsync failed: {error.name}")
+        except Exception as error:
+            cls._native_batch_available = False
+            logger.warning(
+                "Native batch H2D submission failed; "
+                f"disabling native batch copy and falling back to PyTorch copy_: {error}"
+            )
+            return False
+        return True
+
+    @classmethod
+    def _protect_pinned_sources(cls, src_pinned_tensors: List[torch.Tensor], stream: torch.cuda.Stream):
+        """每批 raw H2D 在上传流记录一个完成事件，由复用线程池接管源 Tensor 引用。"""
+        # blocking=True 让后台线程休眠等待；Event.synchronize 会释放 GIL，不阻塞主推理线程。
+        # 事件排在 raw batch 之后，覆盖此前所有源读取；调用方的 wait_stream 负责 GPU 数据依赖。
+        event = torch.cuda.Event(blocking=True)
+        event.record(stream)
+        # 必须复制容器：调用方可能立即清空自己的列表，切片本身仍持有整块 pinned storage。
+        cls._get_release_pool().submit(cls._wait_and_release_pinned_sources, event, list(src_pinned_tensors))
+
+    @classmethod
+    def _get_release_pool(cls) -> ThreadPoolExecutor:
+        if cls._release_pool is None:
+            # 首次创建加锁，避免多个提交线程同时建立等待线程池；后续调用直接复用。
+            with cls._release_pool_lock:
+                if cls._release_pool is None:
+                    cls._release_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="lightllm-pinned-release")
+        return cls._release_pool
+
+    @classmethod
+    def _wait_and_release_pinned_sources(cls, event: torch.cuda.Event, sources: List[torch.Tensor]):
+        """只管理 CPU 源的生命周期，不修改请求状态，也不向流提交新的 GPU 工作。"""
+        event.synchronize()
+        # 普通工作线程可安全释放 Tensor；不受 CUDA host callback 禁止调用 CUDA API 的限制。
+        # 不把 sources 返回给 Future，避免任务结束后仍被结果对象持有。
+        sources.clear()
+
+
+def cuda_memcpy_pinned_h2d_batch_async(dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor]):
+    """按当前流顺序批量异步 H2D，默认流适配、源内存保活和回退由封装统一管理。"""
+    PinnedH2DBatchCopier.copy_async(dst_gpu_tensors, src_pinned_tensors)
+
+
 class cudaIpcMemHandle_t(ctypes.Structure):
     _fields_ = [("internal", ctypes.c_byte * 128)]