Function _tensor_copy_3dim

lightllm/utils/custom_kernel_utis.py:56–82

(
    in_ptr,
    in_stride_0,
    in_stride_1,
    in_stride_2,
    out_ptr,
    out_stride_0,
    out_stride_1,
    out_stride_2,
    head_num,
    head_dim,
    total_len,
    BLOCK_N: tl.constexpr,
)

@triton.jit
def _tensor_copy_3dim(
    in_ptr,
    in_stride_0,
    in_stride_1,
    in_stride_2,
    out_ptr,
    out_stride_0,
    out_stride_1,
    out_stride_2,
    head_num,
    head_dim,
    total_len,
    BLOCK_N: tl.constexpr,
):
    start_index = tl.program_id(0)
    grid_num = tl.num_programs(0)

    offs_d = tl.arange(0, BLOCK_N)
    for cur_index in range(start_index, total_len, step=grid_num):
        for cur_head in tl.range(head_num, num_stages=3):
            in_tensor = tl.load(
                in_ptr + in_stride_0 * cur_index + in_stride_1 * cur_head + offs_d, mask=offs_d < head_dim, other=0
            )
            tl.store(
                out_ptr + out_stride_0 * cur_index + out_stride_1 * cur_head + offs_d, in_tensor, mask=offs_d < head_dim
            )
    return


@torch.no_grad()