TensorRT-LLM 0.5.0 源码之三十

memoryUtils

inline bool checkIfFileExist(const std::string& file_path)
{
    std::ifstream in(file_path, std::ios::in | std::ios::binary);
    if (in.is_open())
    {
        in.close();
        return true;
    }
    return false;
}

// The following functions implement conversion of multi-dimensional indices to an index in a flat array.
// The shape of the Tensor dimensions is passed as one array (`dims`), the indices are given as individual arguments.
// For examples on how to use these functions, see their tests `test_memory_utils.cu`.
// All of these functions can be evaluated at compile time by recursive template expansion.

template <typename TDim, typename T>
__inline__ __host__ __device__ std::enable_if_t<std::is_pointer<TDim>::value, T>constexpr flat_index(
    T const& acc, TDim dims, T const& index)
{
    assert(index < dims[0]);
    return acc * dims[0] + index;
}

template <typename TDim, typename T, typename... TArgs>
__inline__ __host__ __device__ std::enable_if_t<std::is_pointer<TDim>::value, T>constexpr flat_index(
    T const& acc, TDim dims, T const& index, TArgs... indices)
{
    assert(index < dims[0]);
    return flat_index(acc * dims[0] + index, dims + 1, indices...);
}

template <typename TDim, typename T>
__inline__ __host__ __device__ std::enable_if_t<std::is_pointer<TDim>::value, T>constexpr flat_index(
    [[maybe_unused]] TDim dims, T const& index)
{
    assert(index < dims[0]);
    return index;
}

template <typename TDim, typename T, typename... TArgs>
__inline__ __host__ __device__ std::enable_if_t<std::is_pointer<TDim>::value, T>constexpr flat_index(
    TDim dims, T const& index, TArgs... indices)
{
    assert(index < dims[0]);
    return flat_index(index, dims + 1, indices...);
}

template <unsigned skip = 0, typename T, std::size_t N, typename... TIndices>
__inline__ __host__ __device__ T constexpr flat_index(std::array<T, N> const& dims, T const& index, TIndices... indices)
{
    static_assert(skip < N);
    static_assert(sizeof...(TIndices) < N - skip, "Number of indices exceeds number of dimensions");
    return flat_index(&dims[skip], index, indices...);
}

template <unsigned skip = 0, typename T, std::size_t N, typename... TIndices>
__inline__ __host__ __device__ T constexpr flat_index(
    T const& acc, std::array<T, N> const& dims, T const& index, TIndices... indices)
{
    static_assert(skip < N);
    static_assert(sizeof...(TIndices) < N - skip, "Number of indices exceeds number of dimensions");
    return flat_index(acc, &dims[skip], index, indices...);
}

template <unsigned skip = 0, typename T, std::size_t N, typename... TIndices>
__inline__ __host__ __device__ T constexpr flat_index(T const (&dims)[N], T const& index, TIndices... indices)
{
    static_assert(skip < N);
    static_assert(sizeof...(TIndices) < N - skip, "Number of indices exceeds number of dimensions");
    return flat_index(static_cast<T const*>(dims) + skip, index, indices...);
}

template <unsigned skip = 0, typename T, std::size_t N, typename... TIndices>
__inline__ __host__ __device__ T constexpr flat_index(
    T const& acc, T const (&dims)[N], T const& index, TIndices... indices)
{
    static_assert(skip < N);
    static_assert(sizeof...(TIndices) < N - skip, "Number of indices exceeds number of dimensions");
    return flat_index(acc, static_cast<T const*>(dims) + skip, index, indices...);
}

这段代码实现了一个多维索引到一维扁平索引的编译时计算函数，支持多种容器类型和维度跳过功能。下面是我的详细解析：

函数设计目的

这个flat_index函数集的主要目的是将多维索引转换为一维线性索引，这在数组扁平化存储中非常常见。函数使用constexpr关键字确保在编译时进行计算优化，并通过SFINAE技术提供灵活的接口重载。

核心函数解析

1. 指针维度版本（基础实现）

template <typename TDim, typename T>
__inline__ __host__ __device__ std::enable_if_t<std::is_pointer<TDim>::value, T>constexpr flat_index(
    T const& acc, TDim dims, T const& index)
{
    assert(index < dims[0]);
    return acc * dims[0] + index;
}

功能：递归终止情况，处理最后一个维度

• acc：累积的索引值
• dims：维度数组指针
• index：当前维度索引
• 计算：acc × 当前维度大小 + 当前索引

2. 可变参数递归版本

template <typename TDim, typename T, typename... TArgs>
__inline__ __host__ __device__ std::enable_if_t<std::is_pointer<TDim>::value, T>constexpr flat_index(
    T const& acc, TDim dims, T const& index, TArgs... indices)
{
    assert(index < dims[0]);
    return flat_index(acc * dims[0] + index, dims + 1, indices...);
}

功能：递归处理多个维度索引

• 通过参数包TArgs... indices接受可变数量的索引
• 每次递归将维度指针后移（dims + 1）

关键技术特性

1. 编译时计算优化

使用constexpr确保索引计算在编译时完成，提高运行时效率：

constexpr float x = exp(5, 3);  // 编译时计算

2. 跨平台支持

__host__ __device__修饰符使函数既可在CPU也可在GPU上执行，适合CUDA编程。

3. 类型安全

通过std::enable_if_t<std::is_pointer<TDim>::value, T>确保只有指针类型才能匹配该模板。

容器适配器版本

std::array适配器

template <unsigned skip = 0, typename T, std::size_t N, typename... TIndices>
__inline__ __host__ __device__ T constexpr flat_index(std::array<T, N> const& dims, T const& index, TIndices... indices)
{
    static_assert(skip < N);
    static_assert(sizeof...(TIndices) < N - skip, "Number of indices exceeds number of dimensions");
    return flat_index(&dims[skip], index, indices...);
}

特性：

• 支持维度跳过（skip参数）
• 编译时静态断言检查维度边界
• 将std::array转换为指针后调用核心实现

C风格数组适配器

template <unsigned skip = 0, typename T, std::size_t N, typename... TIndices>
__inline__ __host__ __device__ T constexpr flat_index(T const (&dims)[N], T const& index, TIndices... indices)
{
    static_assert(skip < N);
    static_assert(sizeof...(TIndices) < N - skip, "Number of indices exceeds number of dimensions");
    return flat_index(static_cast<T const*>(dims) + skip, index, indices...);
}

使用示例

假设3维数组维度为[2, 3, 4]：

int dims[] = {2, 3, 4};
int flat_idx = flat_index(dims, 1, 2, 3);  // 计算结果：1×3×4 + 2×4 + 3 = 23

设计优势

1. 零运行时开销：全部编译时计算完成
2. 类型安全：完善的静态断言和SFINAE控制
3. 灵活性：支持多种容器类型和维度跳过
4. 跨平台：CPU/GPU通用支持

这个设计体现了现代C++元编程的强大能力，通过模板元编程在编译期完成复杂的多维索引计算。

// These are simpler functions for multi-dimensional index conversion. Indices and dimensions are passed as individual
// arguments. These functions are more suitable for usage inside kernels than the corresponding flat_index functions
// which require arrays as arguments. Usage examples can be found in `test_memory_utils.cu`. The functions can be
// evaluated at compile time.

template <typename T>
__inline__ __host__ __device__ T constexpr flat_index2(T const& index_0, T const& index_1, T const& dim_1)
{
    assert(index_1 < dim_1);
    return index_0 * dim_1 + index_1;
}

template <typename T>
__inline__ __host__ __device__ T constexpr flat_index3(
    T const& index_0, T const& index_1, T const& index_2, T const& dim_1, T const& dim_2)
{
    assert(index_2 < dim_2);
    return flat_index2(index_0, index_1, dim_1) * dim_2 + index_2;
}

template <typename T>
__inline__ __host__ __device__ T constexpr flat_index4(T const& index_0, T const& index_1, T const& index_2,
    T const& index_3, T const& dim_1, T const& dim_2, T const& dim_3)
{
    assert(index_3 < dim_3);
    return flat_index3(index_0, index_1, index_2, dim_1, dim_2) * dim_3 + index_3;
}

template <typename T>
__inline__ __host__ __device__ T constexpr flat_index5(T const& index_0, T const& index_1, T const& index_2,
    T const& index_3, T const& index_4, T const& dim_1, T const& dim_2, T const& dim_3, T const& dim_4)
{
    assert(index_4 < dim_4);
    return flat_index4(index_0, index_1, index_2, index_3, dim_1, dim_2, dim_3) * dim_4 + index_4;
}

// stride 是每个维度对应的 stride值，直接用于shape -> offset 计算
template <typename T>
__inline__ __host__ __device__ T constexpr flat_index_strided3(
    T const& index_0, T const& index_1, T const& index_2, T const& stride_1, T const& stride_2)
{
    assert(index_1 < stride_1 / stride_2);
    assert(index_2 < stride_2);
    return index_0 * stride_1 + index_1 * stride_2 + index_2;
}

template <typename T>
__inline__ __host__ __device__ T constexpr flat_index_strided4(T const& index_0, T const& index_1, T const& index_2,
    T const& index_3, T const& stride_1, T const& stride_2, T const& stride_3)
{
    assert(index_1 < stride_1 / stride_2);
    assert(index_2 < stride_2 / stride_3);
    assert(index_3 < stride_3);
    return index_0 * stride_1 + index_1 * stride_2 + index_2 * stride_3 + index_3;
}

这段代码实现了一套将多维索引映射到一维线性索引的实用函数，其核心原理是基于行优先存储顺序和步长计算。这种转换在处理多维数组的线性化存储（如在C++数组、CUDA内核或数值计算库中）时非常常见。下面我将详细解析其工作原理、函数设计和关键技术特点。

多维索引线性化原理

基本概念：行优先顺序

多维数组在内存中通常以一维形式存储。代码采用行优先顺序，即最左侧索引（index_0）变化最慢，最右侧索引变化最快。线性偏移的计算本质是一个混合进制数转换过程，每个维度的尺寸或步长作为进制基数。

• 维度版本（如 flat_index3）：直接使用各维度大小（dim_1、dim_2）计算偏移。例如，三维索引 (i, j, k) 对应的线性索引为 i * (dim_1 * dim_2) + j * dim_2 + k。这里，dim_1 是第二维的大小，dim_2 是第三维的大小，而第一维的大小 dim_0 不直接参与计算，因为其影响已通过更高维度的乘积体现。
• 步长版本（如 flat_index_strided3）：使用预计算的步长参数（stride_1、stride_2），表示在对应维度上移动一个索引位置时线性地址的增量。步长允许处理非紧凑或特殊内存布局的数组（如填充数组或转置矩阵）。例如，三维索引 (i, j, k) 的线性索引为 i * stride_1 + j * stride_2 + k * 1，其中 stride_1 通常等于 dim_1 * dim_2，stride_2 等于 dim_2。

函数设计解析

1. 维度显式版本（flat_index2 到 flat_index5）

这类函数通过递归组合实现维度扩展，每个函数添加一个新维度的计算：

• flat_index2：基础二维转换，公式为 index_0 * dim_1 + index_1。这对应于二维数组的行优先布局。
• flat_index3：通过调用 flat_index2 计算前两维的偏移 (index_0 * dim_1 + index_1)，再乘以第三维大小 dim_2 并加上 index_2，得到 (index_0 * dim_1 + index_1) * dim_2 + index_2。等价于 index_0 * (dim_1 * dim_2) + index_1 * dim_2 + index_2。
• 更高维函数（如 flat_index4）采用相同模式递归扩展，确保代码复用和一致性。

2. 步长显式版本（flat_index_strided3 和 flat_index_strided4）

步长版本更灵活，直接使用内存布局参数：

• flat_index_strided3：公式为 index_0 * stride_1 + index_1 * stride_2 + index_2。其中：

• stride_1 是最高维的步长（例如三维数组中表示一整个二维平面的元素数量）。
• stride_2 是次高维的步长（例如三维数组中表示一行的元素数量）。
• 最低维步长隐式为 1（即 index_2 的系数）。

• 断言的作用：例如 assert(index_1 < stride_1 / stride_2) 验证 index_1 不超过第二维的实际大小（因为 stride_1 / stride_2 等于第二维尺寸）。这提供了运行时安全性检查。

3. 关键技术特性

• 编译时计算：所有函数标记为 constexpr，允许编译器在编译期计算索引偏移，消除运行时开销。例如，如果维度和索引是编译期常量，结果将直接内联优化。
• 跨平台支持：__host__ __device__ 修饰符确保函数可在 CPU 和 GPU（如 CUDA）上执行，适用于异构计算场景。这种索引计算在 CUDA 线程映射中非常常见。
• 类型泛化：模板参数 T 支持任意整数类型（如 int、size_t），增强代码通用性。

应用场景与示例

多维数组访问

假设三维数组维度为 [D][H][W]（深度、高度、宽度），索引 (d, h, w) 的线性偏移为：

// 使用维度版本
int linear_idx = flat_index3(d, h, w, H, W); // 等价于 d*(H*W) + h*W + w

// 使用步长版本（预计算步长）
int stride_1 = H * W; // 深度步长
int stride_2 = W;      // 高度步长
int linear_idx = flat_index_strided3(d, h, w, stride_1, stride_2);

CUDA 线程索引映射

在 CUDA 编程中，类似公式用于将多维线程坐标（blockIdx、threadIdx）转换为全局线性线程 ID。例如，三维 Grid 和一维 Block 的线程 ID 计算为：

// 与 flat_index3 逻辑相似：blockIdx 对应 index_0，threadIdx 对应 index_1/index_2
int tid = blockIdx.x * blockDim.x + threadIdx.x; // 一维情况
// 多维扩展类似 flat_index 的递归组合

设计优势与局限性

• 优势：

• 清晰易懂：直接映射行优先数学公式，代码可读性强。
• 零开销抽象：constexpr 和内联优化确保无运行时性能损失。
• 灵活性：步长版本支持非标准内存布局（如跨步数组或填充对齐）。

• 局限性：

• 维度限制：函数需显式定义每个维度（如最多到 flat_index5），不如可变参数模板灵活。
• 静态断言缺失：依赖运行时 assert，无法在编译时检查所有维度边界（但 constexpr 上下文中期望索引为常量）。

总结

这套函数的核心原理是通过维度乘积累积或步长加权和将多维索引转换为线性地址。其设计体现了高效的内存访问模式，与 CUDA 线程索引计算、数组扁平化存储等场景高度契合。通过编译时计算和泛型编程，它在保证性能的同时提供了类型安全的索引转换工具。

deviceMalloc

template <typename T>
void deviceMalloc(T** ptr, size_t size, bool is_random_initialize)
{
    check_cuda_error(cudaMalloc((void**) (ptr), sizeof(T) * size));
    if (is_random_initialize)
    {
        cudaRandomUniform(*ptr, size);
    }
}


template void deviceMalloc(float** ptr, size_t size, bool is_random_initialize);
template void deviceMalloc(half** ptr, size_t size, bool is_random_initialize);
#ifdef ENABLE_BF16
template void deviceMalloc(__nv_bfloat16** ptr, size_t size, bool is_random_initialize);
#endif
template void deviceMalloc(uint16_t** ptr, size_t size, bool is_random_initialize);
template void deviceMalloc(int** ptr, size_t size, bool is_random_initialize);
template void deviceMalloc(bool** ptr, size_t size, bool is_random_initialize);
template void deviceMalloc(char** ptr, size_t size, bool is_random_initialize);
template void deviceMalloc(int8_t** ptr, size_t size, bool is_random_initialize);
#ifdef ENABLE_FP8
template void deviceMalloc(__nv_fp8_e4m3** ptr, size_t size, bool is_random_initialize);
#endif

template <typename T>
void deviceMemSetZero(T* ptr, size_t size)
{
    check_cuda_error(cudaMemset(static_cast<void*>(ptr), 0, sizeof(T) * size));
}

template void deviceMemSetZero(float* ptr, size_t size);
template void deviceMemSetZero(half* ptr, size_t size);
template void deviceMemSetZero(int* ptr, size_t size);
template void deviceMemSetZero(uint32_t* ptr, size_t size);
template void deviceMemSetZero(bool* ptr, size_t size);
#ifdef ENABLE_FP8
template void deviceMemSetZero(__nv_fp8_e4m3* ptr, size_t size);
#endif
#ifdef ENABLE_BF16
template void deviceMemSetZero(__nv_bfloat16* ptr, size_t size);
#endif

template <typename T>
void deviceFree(T*& ptr)
{
    if (ptr != NULL)
    {
        check_cuda_error(cudaFree(ptr));
        ptr = NULL;
    }
}

template void deviceFree(float*& ptr);
template void deviceFree(half*& ptr);
#ifdef ENABLE_BF16
template void deviceFree(__nv_bfloat16*& ptr);
#endif
template void deviceFree(unsigned short*& ptr);
template void deviceFree(int*& ptr);
template void deviceFree(bool*& ptr);
template void deviceFree(char*& ptr);
template void deviceFree(int8_t*& ptr);
#ifdef ENABLE_FP8
template void deviceFree(__nv_fp8_e4m3*& ptr);
#endif

template <typename T>
void deviceFill(T* devptr, size_t size, T value, cudaStream_t stream)
{
    T* arr = new T[size];
    std::fill(arr, arr + size, value);
    check_cuda_error(cudaMemcpyAsync(devptr, arr, sizeof(T) * size, cudaMemcpyHostToDevice, stream));
    delete[] arr;
}

template void deviceFill(float* devptr, size_t size, float value, cudaStream_t stream);
template void deviceFill(half* devptr, size_t size, half value, cudaStream_t stream);
#ifdef ENABLE_BF16
template void deviceFill(__nv_bfloat16* devptr, size_t size, __nv_bfloat16 value, cudaStream_t stream);
#endif
template void deviceFill(int* devptr, size_t size, int value, cudaStream_t stream);
template void deviceFill(bool* devptr, size_t size, bool value, cudaStream_t stream);

template <typename T>
void cudaD2Hcpy(T* tgt, const T* src, const size_t size)
{
    check_cuda_error(cudaMemcpy(tgt, src, sizeof(T) * size, cudaMemcpyDeviceToHost));
}

template void cudaD2Hcpy(float* tgt, const float* src, size_t size);
template void cudaD2Hcpy(half* tgt, const half* src, size_t size);
#ifdef ENABLE_BF16
template void cudaD2Hcpy(__nv_bfloat16* tgt, const __nv_bfloat16* src, size_t size);
#endif
template void cudaD2Hcpy(int* tgt, const int* src, size_t size);
template void cudaD2Hcpy(bool* tgt, const bool* src, size_t size);
#ifdef ENABLE_FP8
template void cudaD2Hcpy(__nv_fp8_e4m3* tgt, const __nv_fp8_e4m3* src, size_t size);
#endif
template void cudaD2Hcpy(unsigned long long* tgt, const unsigned long long* src, size_t size);
template void cudaD2Hcpy(unsigned int* tgt, const unsigned int* src, size_t size);
template void cudaD2Hcpy(int8_t* tgt, const int8_t* src, size_t size);

template <typename T>
void cudaH2Dcpy(T* tgt, const T* src, const size_t size)
{
    check_cuda_error(cudaMemcpy(tgt, src, sizeof(T) * size, cudaMemcpyHostToDevice));
}

template void cudaH2Dcpy(float* tgt, const float* src, size_t size);
template void cudaH2Dcpy(half* tgt, const half* src, size_t size);
#ifdef ENABLE_BF16
template void cudaH2Dcpy(__nv_bfloat16* tgt, const __nv_bfloat16* src, size_t size);
#endif
template void cudaH2Dcpy(int* tgt, const int* src, size_t size);
template void cudaH2Dcpy(bool* tgt, const bool* src, size_t size);
#ifdef ENABLE_FP8
template void cudaH2Dcpy(__nv_fp8_e4m3* tgt, const __nv_fp8_e4m3* src, size_t size);
#endif
template void cudaH2Dcpy(unsigned long long* tgt, const unsigned long long* src, size_t size);
template void cudaH2Dcpy(unsigned int* tgt, const unsigned int* src, size_t size);
template void cudaH2Dcpy(int8_t* tgt, const int8_t* src, size_t size);

template <typename T>
void cudaD2Dcpy(T* tgt, const T* src, const size_t size)
{
    check_cuda_error(cudaMemcpy(tgt, src, sizeof(T) * size, cudaMemcpyDeviceToDevice));
}

template void cudaD2Dcpy(float* tgt, const float* src, size_t size);
template void cudaD2Dcpy(half* tgt, const half* src, size_t size);
#ifdef ENABLE_BF16
template void cudaD2Dcpy(__nv_bfloat16* tgt, const __nv_bfloat16* src, size_t size);
#endif
template void cudaD2Dcpy(int* tgt, const int* src, size_t size);
template void cudaD2Dcpy(bool* tgt, const bool* src, size_t size);
template void cudaD2Dcpy(int8_t* tgt, const int8_t* src, size_t size);
#ifdef ENABLE_FP8
template void cudaD2Dcpy(__nv_fp8_e4m3* tgt, const __nv_fp8_e4m3* src, size_t size);
#endif
template void cudaD2Dcpy(unsigned long long* tgt, const unsigned long long* src, size_t size);

template <typename T_OUT, typename T_IN>
__global__ void cudaCast(T_OUT* dst, T_IN* src, const size_t size)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = (T_OUT) ((float) (src[tid]));
    }
}

template <typename T_OUT, typename T_IN>
void invokeCudaCast(T_OUT* dst, T_IN const* const src, const size_t size, cudaStream_t stream)
{
    cudaCast<<<256, 256, 0, stream>>>(dst, src, size);
}

template void invokeCudaCast(float* dst, half const* const src, const size_t size, cudaStream_t stream);
#ifdef ENABLE_BF16
template void invokeCudaCast(float* dst, __nv_bfloat16 const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(__nv_bfloat16* dst, float const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(__nv_bfloat16* dst, half const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(half* dst, __nv_bfloat16 const* const src, const size_t size, cudaStream_t stream);
#endif
#ifdef ENABLE_FP8
template void invokeCudaCast(float* dst, __nv_fp8_e4m3 const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(
    __nv_bfloat16* dst, __nv_fp8_e4m3 const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(half* dst, __nv_fp8_e4m3 const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(__nv_fp8_e4m3* dst, float const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(
    __nv_fp8_e4m3* dst, __nv_bfloat16 const* const src, const size_t size, cudaStream_t stream);
template void invokeCudaCast(__nv_fp8_e4m3* dst, half const* const src, const size_t size, cudaStream_t stream);
#endif

template <typename T>
void cudaAutoCpy(T* tgt, const T* src, const size_t size, cudaStream_t stream)
{
    if (stream != NULL)
    {
        check_cuda_error(cudaMemcpyAsync(tgt, src, sizeof(T) * size, cudaMemcpyDefault, stream));
    }
    else
    {
        check_cuda_error(cudaMemcpy(tgt, src, sizeof(T) * size, cudaMemcpyDefault));
    }
}

template void cudaAutoCpy(float* tgt, const float* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(half* tgt, const half* src, size_t size, cudaStream_t stream);
#ifdef ENABLE_BF16
template void cudaAutoCpy(__nv_bfloat16* tgt, const __nv_bfloat16* src, size_t size, cudaStream_t stream);
#endif
template void cudaAutoCpy(int* tgt, const int* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(bool* tgt, const bool* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(int8_t* tgt, const int8_t* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(uint8_t* tgt, const uint8_t* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(uint32_t* tgt, const uint32_t* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(unsigned long long* tgt, const unsigned long long* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(char* tgt, const char* src, size_t size, cudaStream_t stream);

template void cudaAutoCpy(float const** tgt, float const* const* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(half const** tgt, half const* const* src, size_t size, cudaStream_t stream);
#ifdef ENABLE_BF16
template void cudaAutoCpy(__nv_bfloat16 const** tgt, __nv_bfloat16 const* const* src, size_t size, cudaStream_t stream);
#endif
template void cudaAutoCpy(int const** tgt, int const* const* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(bool const** tgt, bool const* const* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(int8_t const** tgt, int8_t const* const* src, size_t size, cudaStream_t stream);
template void cudaAutoCpy(
    unsigned long long const** tgt, unsigned long long const* const* src, size_t size, cudaStream_t stream);

template <typename T>
__global__ void cuda_random_uniform_kernel(T* buffer, const size_t size, const int seq_offset)
{
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    curandState_t local_state;
    curand_init((unsigned long long int) 1337, idx + seq_offset, 0, &local_state);
    for (size_t index = idx; index < size; index += blockDim.x * gridDim.x)
    {
        buffer[index] = (T) (curand_uniform(&local_state) * 0.2f - 0.1f);
    }
}

template <>
__global__ void cuda_random_uniform_kernel<int>(int* buffer, const size_t size, const int seq_offset)
{
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    curandState_t local_state;
    curand_init((float) 1337.f, idx + seq_offset, 0, &local_state);
    for (size_t index = idx; index < size; index += blockDim.x * gridDim.x)
    {
        buffer[index] = curand(&local_state);
    }
}

template <>
__global__ void cuda_random_uniform_kernel<bool>(bool* buffer, const size_t size, const int seq_offset)
{
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    curandState_t local_state;
    curand_init((float) 1337.f, idx + seq_offset, 0, &local_state);
    for (size_t index = idx; index < size; index += blockDim.x * gridDim.x)
    {
        buffer[index] = (curand(&local_state) % 2 == 0);
    }
}

template <>
__global__ void cuda_random_uniform_kernel<char>(char* buffer, const size_t size, const int seq_offset)
{
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    curandState_t local_state;
    curand_init((float) 1337.f, idx + seq_offset, 0, &local_state);
    for (size_t index = idx; index < size; index += blockDim.x * gridDim.x)
    {
        buffer[index] = curand(&local_state) % 0xFF;
    }
}

template <typename T>
void cudaRandomUniform(T* buffer, const size_t size)
{
    static int seq_offset = 0;
    cuda_random_uniform_kernel<T><<<256, 256>>>(buffer, size, seq_offset);
    seq_offset += 256 * 256;
}

template void cudaRandomUniform(float* buffer, const size_t size);
template void cudaRandomUniform(half* buffer, const size_t size);
#ifdef ENABLE_BF16
template void cudaRandomUniform(__nv_bfloat16* buffer, const size_t size);
#endif
template void cudaRandomUniform(int* buffer, const size_t size);
template void cudaRandomUniform(bool* buffer, const size_t size);
template void cudaRandomUniform(char* buffer, const size_t size);
#ifdef ENABLE_FP8
template void cudaRandomUniform(__nv_fp8_e4m3* buffer, const size_t size);
#endif

// loads data from binary file. If it succeeds, returns a non-empty vector. If loading fails or
// the product of the elements in shape is 0, this function will return an empty vector.
template <typename T>
std::vector<T> loadWeightFromBinHelper(std::vector<size_t> shape, std::string filename)
{
    if (shape.size() > 2)
    {
        printf("[ERROR] shape should have less than two dims \n");
        return std::vector<T>();
    }
    size_t dim0 = shape[0], dim1 = 1;
    if (shape.size() == 2)
    {
        dim1 = shape[1];
    }
    size_t size = dim0 * dim1;
    if (size == 0)
    {
        TLLM_LOG_WARNING("shape is zero, skip loading weight from file %s \n", filename.c_str());
        return std::vector<T>();
    }

    std::vector<T> host_array(size);
    std::ifstream in(filename, std::ios::in | std::ios::binary);
    if (!in.is_open())
    {
        TLLM_LOG_WARNING("file %s cannot be opened, loading model fails! \n", filename.c_str());
        return std::vector<T>();
    }

    size_t loaded_data_size = sizeof(T) * size;
    in.seekg(0, in.end);
    in.seekg(0, in.beg);

    TLLM_LOG_DEBUG("Read " + std::to_string(loaded_data_size) + " bytes from " + filename);
    in.read((char*) host_array.data(), loaded_data_size);

    size_t in_get_size = in.gcount();
    if (in_get_size != loaded_data_size)
    {
        TLLM_LOG_WARNING("file %s only has %ld, but request %ld, loading model fails! \n", filename.c_str(),
            in_get_size, loaded_data_size);
        return std::vector<T>();
    }
    in.close();
    // If we succeed, return an array with values.
    return host_array;
}

template <typename T, typename T_IN>
int loadWeightFromBinFunc(T* ptr, std::vector<size_t> shape, std::string filename)
{
    std::vector<T_IN> host_array = loadWeightFromBinHelper<T_IN>(shape, filename);

    if (host_array.empty())
    {
        return 0;
    }

    if (std::is_same<T, T_IN>::value == true)
    {
        cudaH2Dcpy(ptr, (T*) host_array.data(), host_array.size());
    }
    else
    {
        T_IN* ptr_2 = nullptr;
        deviceMalloc(&ptr_2, host_array.size(), false);
        cudaH2Dcpy(ptr_2, host_array.data(), host_array.size());
        invokeCudaD2DcpyConvert(ptr, ptr_2, host_array.size());
        deviceFree(ptr_2);
    }
    return 0;
}

template int loadWeightFromBinFunc<float, float>(float* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<half, float>(half* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<float, half>(float* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<half, half>(half* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<int8_t, int8_t>(int8_t* ptr, std::vector<size_t> shape, std::string filename);
#ifdef ENABLE_BF16
template int loadWeightFromBinFunc<__nv_bfloat16, float>(
    __nv_bfloat16* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<__nv_bfloat16, half>(
    __nv_bfloat16* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<float, __nv_bfloat16>(float* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<half, __nv_bfloat16>(half* ptr, std::vector<size_t> shape, std::string filename);
template int loadWeightFromBinFunc<__nv_bfloat16, __nv_bfloat16>(
    __nv_bfloat16* ptr, std::vector<size_t> shape, std::string filename);
#endif // ENABLE_BF16
template int loadWeightFromBinFunc<int, int>(int* ptr, std::vector<size_t> shape, std::string filename);
#ifdef ENABLE_FP8
template int loadWeightFromBinFunc<__nv_fp8_e4m3, float>(
    __nv_fp8_e4m3* ptr, std::vector<size_t> shape, std::string filename);
#endif // ENABLE_FP8

template <typename T>
int loadWeightFromBin(T* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type)
{
    switch (model_file_type)
    {
    case TRTLLMCudaDataType::FP32: loadWeightFromBinFunc<T, float>(ptr, shape, filename); break;
    case TRTLLMCudaDataType::FP16: loadWeightFromBinFunc<T, half>(ptr, shape, filename); break;
    case TRTLLMCudaDataType::INT8: loadWeightFromBinFunc<T, int8_t>(ptr, shape, filename); break;
#ifdef ENABLE_BF16
    case TRTLLMCudaDataType::BF16: loadWeightFromBinFunc<T, __nv_bfloat16>(ptr, shape, filename); break;
#endif
#ifdef ENABLE_FP8
    case TRTLLMCudaDataType::FP8: loadWeightFromBinFunc<T, float>(ptr, shape, filename); break;
#endif
    default: TLLM_LOG_ERROR("Does not support TRTLLMCudaDataType=%d", model_file_type); TLLM_CHECK(false);
    }
    return 0;
}

template <>
int loadWeightFromBin(int* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type)
{
    loadWeightFromBinFunc<int, int>(ptr, shape, filename);
    return 0;
}

template int loadWeightFromBin(
    float* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type);
template int loadWeightFromBin(
    half* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type);
template int loadWeightFromBin(
    int8_t* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type);
#ifdef ENABLE_BF16
template int loadWeightFromBin(
    __nv_bfloat16* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type);
#endif
#ifdef ENABLE_FP8
template int loadWeightFromBin(
    __nv_fp8_e4m3* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type);
#endif
template int loadWeightFromBin(
    int* ptr, std::vector<size_t> shape, std::string filename, TRTLLMCudaDataType model_file_type);

template <typename T_IN, typename T_OUT>
__global__ void cudaD2DcpyConvert(T_OUT* dst, const T_IN* src, const size_t size)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = cuda_cast<T_OUT>(src[tid]);
    }
}

template <typename T_IN, typename T_OUT>
void invokeCudaD2DcpyConvert(T_OUT* tgt, const T_IN* src, const size_t size, cudaStream_t stream)
{
    cudaD2DcpyConvert<<<256, 256, 0, stream>>>(tgt, src, size);
}

template void invokeCudaD2DcpyConvert(int8_t* tgt, const float* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(float* tgt, const int8_t* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(float* tgt, const int* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(half* tgt, const int* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(float* tgt, const float* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(half* tgt, const float* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(float* tgt, const half* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(uint32_t* tgt, const int* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(int* tgt, const uint32_t* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(int* tgt, const float* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(int* tgt, const half* src, const size_t size, cudaStream_t stream);

#ifdef ENABLE_BF16
template void invokeCudaD2DcpyConvert(__nv_bfloat16* tgt, const float* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(__nv_bfloat16* tgt, const int* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(float* tgt, const __nv_bfloat16* src, const size_t size, cudaStream_t stream);
template void invokeCudaD2DcpyConvert(int* tgt, const __nv_bfloat16* src, const size_t size, cudaStream_t stream);
#endif // ENABLE_BF16

template <typename T_IN, typename T_OUT>
__global__ void cudaD2DScaleCpyConvert(
    T_OUT* dst, const T_IN* src, const float* scale, bool invert_scale, const size_t size)
{
    const float scale_value = invert_scale ? 1.0f / scale[0] : scale[0];
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = cuda_cast<T_OUT>(cuda_cast<float>(src[tid]) * scale_value);
    }
}

template <typename T_IN, typename T_OUT>
void invokeCudaD2DScaleCpyConvert(
    T_OUT* tgt, const T_IN* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream)
{
    cudaD2DScaleCpyConvert<<<256, 256, 0, stream>>>(tgt, src, scale, invert_scale, size);
}

// clang-format off
template void invokeCudaD2DScaleCpyConvert(float* tgt, const int32_t* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream);
template void invokeCudaD2DScaleCpyConvert(int32_t* tgt, const float* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream);
template void invokeCudaD2DScaleCpyConvert(half* tgt, const int32_t* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream);
template void invokeCudaD2DScaleCpyConvert(int32_t* tgt, const half* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream);
#ifdef ENABLE_BF16
template void invokeCudaD2DScaleCpyConvert(__nv_bfloat16* tgt, const int32_t* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream);
template void invokeCudaD2DScaleCpyConvert(int32_t* tgt, const __nv_bfloat16* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream);
#endif  // ENABLE_BF16
#ifdef ENABLE_FP8
template void invokeCudaD2DScaleCpyConvert(float* tgt, const __nv_fp8_e4m3* src, const float* scale, bool invert_scale, const size_t size, cudaStream_t stream);
#endif  // ENABLE_FP8
// clang-format on

void invokeCudaD2DcpyHalf2Float(float* dst, half* src, const size_t size, cudaStream_t stream)
{
    invokeCudaD2DcpyConvert(dst, src, size, stream);
}

void invokeCudaD2DcpyFloat2Half(half* dst, float* src, const size_t size, cudaStream_t stream)
{
    invokeCudaD2DcpyConvert(dst, src, size, stream);
}

template <typename T>
void saveToBinary(const T* ptr, const size_t size, std::string filename)
{

    std::vector<T> h_ptr(size);
    cudaD2Hcpy(h_ptr.data(), ptr, size);
    std::vector<float> float_ptr(size);
    for (size_t i = 0; i < size; i++)
    {
        float_ptr[i] = (float) h_ptr[i];
    }

    std::ofstream out(filename, std::ios::out | std::ios::binary);
    TLLM_CHECK_WITH_INFO(out.is_open(), "Fail to open file " + filename);

    out.write((char*) float_ptr.data(), size * sizeof(float));
}

template void saveToBinary(const float* ptr, const size_t size, std::string filename);
template void saveToBinary(const half* ptr, const size_t size, std::string filename);
#ifdef ENABLE_BF16
template void saveToBinary(const __nv_bfloat16* ptr, const size_t size, std::string filename);
#endif // ENABLE_BF16

template <>
void saveToBinary(const int* ptr, const size_t size, std::string filename)
{
    std::vector<int> h_ptr(size);
    cudaD2Hcpy(h_ptr.data(), ptr, size);
    std::ofstream out(filename, std::ios::out | std::ios::binary);
    TLLM_CHECK_WITH_INFO(out.is_open(), "Fail to open file " + filename);
    out.write((char*) h_ptr.data(), size * sizeof(int));
}

template <typename T_IN, typename T_fake_type>
__global__ void fakeCast(T_IN* input_ptr, const size_t size)
{
    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += blockDim.x * gridDim.x)
    {
        T_fake_type tmp_val = (T_fake_type) ((float) input_ptr[i]);
        input_ptr[i] = (T_IN) ((float) tmp_val);
    }
}

template <typename T_IN, typename T_fake_type>
void invokeFakeCast(T_IN* input_ptr, const size_t size, cudaStream_t stream)
{
    dim3 block(256);
    dim3 grid((size + 255) / 256);
    fakeCast<T_IN, T_fake_type><<<grid, block, 0, stream>>>(input_ptr, size);
}

#ifdef ENABLE_FP8
__global__ void cudaD2Dcpyfp82Float(float* dst, __nv_fp8_e4m3* src, const size_t size)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = (float) (src[tid]);
    }
}

void invokeCudaD2Dcpyfp82Float(float* dst, __nv_fp8_e4m3* src, const size_t size, cudaStream_t stream)
{
    cudaD2Dcpyfp82Float<<<256, 256, 0, stream>>>(dst, src, size);
}

__global__ void cudaD2Dcpyfp82Half(half* dst, __nv_fp8_e4m3* src, const size_t size)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = (half) ((float) (src[tid]));
    }
}

void invokeCudaD2Dcpyfp82Half(half* dst, __nv_fp8_e4m3* src, const size_t size, cudaStream_t stream)
{
    cudaD2Dcpyfp82Half<<<256, 256, 0, stream>>>(dst, src, size);
}

__global__ void cudaD2DcpyFloat2fp8(__nv_fp8_e4m3* dst, float* src, const size_t size)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = (__nv_fp8_e4m3) src[tid];
    }
}

void invokeCudaD2DcpyFloat2fp8(__nv_fp8_e4m3* dst, float* src, const size_t size, cudaStream_t stream)
{
    cudaD2DcpyFloat2fp8<<<256, 256, 0, stream>>>(dst, src, size);
}

__global__ void cudaD2DcpyHalf2fp8(__nv_fp8_e4m3* dst, half* src, const size_t size)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = (__nv_fp8_e4m3) src[tid];
    }
}

void invokeCudaD2DcpyHalf2fp8(__nv_fp8_e4m3* dst, half* src, const size_t size, cudaStream_t stream)
{
    cudaD2DcpyHalf2fp8<<<256, 256, 0, stream>>>(dst, src, size);
}

__global__ void cudaD2DcpyBfloat2fp8(__nv_fp8_e4m3* dst, __nv_bfloat16* src, const size_t size)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < size; tid += blockDim.x * gridDim.x)
    {
        dst[tid] = (__nv_fp8_e4m3) src[tid];
    }
}

void invokeCudaD2DcpyBfloat2fp8(__nv_fp8_e4m3* dst, __nv_bfloat16* src, const size_t size, cudaStream_t stream)
{
    cudaD2DcpyBfloat2fp8<<<256, 256, 0, stream>>>(dst, src, size);
}

#endif // ENABLE_FP8

// [dim0, dim1] -> [dim1, dim0]
template <typename T_OUT, typename T_IN>
__global__ void transpose(T_OUT* dst, T_IN* src, const size_t dim0, const size_t dim1)
{
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < dim0 * dim1; tid += blockDim.x * gridDim.x)
    {
        const size_t src_col_id = tid % dim1;
        const size_t src_row_id = tid / dim1;
        dst[src_col_id * dim0 + src_row_id] = (T_OUT) (src[tid]);
    }
}

template <typename T>
void invokeInPlaceTranspose(T* data, T* workspace, const size_t dim0, const size_t dim1)
{
    // copy data to workspace, and then transpose from workspace to data
    cudaD2Dcpy(workspace, data, dim0 * dim1);
    transpose<<<256, 256>>>(data, workspace, dim0, dim1);
}


#ifdef ENABLE_FP8
template void invokeInPlaceTranspose(
    __nv_fp8_e4m3* data, __nv_fp8_e4m3* workspace, const size_t dim0, const size_t dim1);
#endif // ENABLE_FP8
#ifdef ENABLE_BF16
template void invokeInPlaceTranspose(
    __nv_bfloat16* data, __nv_bfloat16* workspace, const size_t dim0, const size_t dim1);
#endif // ENABLE_BF16
template void invokeInPlaceTranspose(float* data, float* workspace, const size_t dim0, const size_t dim1);

template <typename T_OUT, typename T_IN>
__global__ void transpose0213(
    T_OUT* dst, T_IN* src, const size_t dim0, const size_t dim1, const size_t dim2, const size_t dim3)
{
    // src permutation: [0, 1, 2, 3]
    // dst permutation: [0, 2, 1, 3]
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < dim0 * dim1 * dim2 * dim3;
         tid += blockDim.x * gridDim.x)
    {
        size_t tmp_idx = tid;
        const size_t dim_3_idx = tmp_idx % dim3;
        tmp_idx = (tmp_idx - dim_3_idx) / dim3;
        const size_t dim_2_idx = tmp_idx % dim2;
        tmp_idx = (tmp_idx - dim_2_idx) / dim2;
        const size_t dim_1_idx = tmp_idx % dim1;
        tmp_idx = (tmp_idx - dim_1_idx) / dim1;
        const size_t dim_0_idx = tmp_idx % dim0;
        dst[dim_0_idx * dim1 * dim2 * dim3 + dim_2_idx * dim1 * dim3 + dim_1_idx * dim3 + dim_3_idx] = src[tid];
    }
}

template <typename T>
void invokeInPlaceTranspose0213(
    T* data, T* workspace, const size_t dim0, const size_t dim1, const size_t dim2, const size_t dim3)
{
    // copy data to workspace, and then transpose from workspace to data
    // Note that this kernel is used for pre-processing and not very efficient.
    cudaD2Dcpy(workspace, data, dim0 * dim1 * dim2 * dim3);
    transpose0213<<<256, 256>>>(data, workspace, dim0, dim1, dim2, dim3);
}

#ifdef ENABLE_FP8
template void invokeInPlaceTranspose0213(__nv_fp8_e4m3* data, __nv_fp8_e4m3* workspace, const size_t dim0,
    const size_t dim1, const size_t dim2, const size_t dim3);
#endif // ENABLE_FP8
#ifdef ENABLE_BF16
template void invokeInPlaceTranspose0213(__nv_bfloat16* data, __nv_bfloat16* workspace, const size_t dim0,
    const size_t dim1, const size_t dim2, const size_t dim3);
#endif // ENABLE_BF16
template void invokeInPlaceTranspose0213(
    float* data, float* workspace, const size_t dim0, const size_t dim1, const size_t dim2, const size_t dim3);

template <typename T_OUT, typename T_IN>
__global__ void transpose102(T_OUT* dst, T_IN* src, const size_t dim0, const size_t dim1, const size_t dim2)
{
    // src permutation: [0, 1, 2]
    // dst permutation: [1, 0, 2]
    for (size_t tid = threadIdx.x + blockIdx.x * blockDim.x; tid < dim0 * dim1 * dim2; tid += blockDim.x * gridDim.x)
    {
        size_t tmp_idx = tid;
        const size_t dim_2_idx = tmp_idx % dim2;
        tmp_idx = (tmp_idx - dim_2_idx) / dim2;
        const size_t dim_1_idx = tmp_idx % dim1;
        tmp_idx = (tmp_idx - dim_1_idx) / dim1;
        const size_t dim_0_idx = tmp_idx % dim0;
        dst[dim_1_idx * dim0 * dim2 + dim_0_idx * dim2 + dim_2_idx] = src[tid];
    }
}

template <typename T>
void invokeInPlaceTranspose102(T* data, T* workspace, const size_t dim0, const size_t dim1, const size_t dim2)
{
    // copy data to workspace, and then transpose from workspace to data
    // Note that this kernel is used for pre-processing and not very efficient.
    cudaD2Dcpy(workspace, data, dim0 * dim1 * dim2);
    transpose102<<<256, 256>>>(data, workspace, dim0, dim1, dim2);
}

#ifdef ENABLE_FP8
template void invokeInPlaceTranspose102(
    __nv_fp8_e4m3* data, __nv_fp8_e4m3* workspace, const size_t dim0, const size_t dim1, const size_t dim2);
#endif // ENABLE_FP8
#ifdef ENABLE_BF16
template void invokeInPlaceTranspose102(
    __nv_bfloat16* data, __nv_bfloat16* workspace, const size_t dim0, const size_t dim1, const size_t dim2);
#endif // ENABLE_BF16
template void invokeInPlaceTranspose102(
    float* data, float* workspace, const size_t dim0, const size_t dim1, const size_t dim2);

template <typename T>
void __global__ multiplyScale(T* tensor, float scale, const size_t size)
{
    for (size_t index = threadIdx.x + blockIdx.x * blockDim.x; index < size; index += blockDim.x * gridDim.x)
    {
        tensor[index] = (T) (((float) tensor[index]) * scale);
    }
}

template <typename T>
void invokeMultiplyScale(T* tensor, float scale, const size_t size, cudaStream_t stream)
{
    int block = 256;
    int grid = (size + 255) / 256;
    multiplyScale<<<grid, block, 0, stream>>>(tensor, scale, size);
}

template void invokeMultiplyScale(float* tensor, float scale, const size_t size, cudaStream_t stream);
template void invokeMultiplyScale(half* tensor, float scale, const size_t size, cudaStream_t stream);
#ifdef ENABLE_BF16
template void invokeMultiplyScale(__nv_bfloat16* tensor, float scale, const size_t size, cudaStream_t stream);
#endif
#ifdef ENABLE_FP8
template void invokeMultiplyScale(__nv_fp8_e4m3* tensor, float scale, const size_t size, cudaStream_t stream);
#endif

template <typename T>
void __global__ divideScale(T* tensor, float scale, const size_t size)
{
    for (size_t index = threadIdx.x + blockIdx.x * blockDim.x; index < size; index += blockDim.x * gridDim.x)
    {
        tensor[index] = (T) (((float) tensor[index]) / scale);
    }
}

template <typename T>
void invokeDivideScale(T* tensor, float scale, const size_t size, cudaStream_t stream)
{
    int block = 256;
    int grid = (size + 255) / 256;
    divideScale<<<grid, block, 0, stream>>>(tensor, scale, size);
}

template void invokeDivideScale(float* tensor, float scale, const size_t size, cudaStream_t stream);
template void invokeDivideScale(half* tensor, float scale, const size_t size, cudaStream_t stream);
#ifdef ENABLE_BF16
template void invokeDivideScale(__nv_bfloat16* tensor, float scale, const size_t size, cudaStream_t stream);
#endif
#ifdef ENABLE_FP8
template void invokeDivideScale(__nv_fp8_e4m3* tensor, float scale, const size_t size, cudaStream_t stream);
#endif

size_t cuda_datatype_size(TRTLLMCudaDataType dt)
{
    static const std::unordered_map<TRTLLMCudaDataType, size_t> sizes{
        {TRTLLMCudaDataType::FP32, sizeof(float)}, {TRTLLMCudaDataType::FP16, sizeof(half)}
#ifdef ENABLE_BF16
        ,
        {TRTLLMCudaDataType::BF16, sizeof(__nv_bfloat16)}
#endif
    };

    return sizes.at(dt);
}

template <typename T>
__global__ void check_range(const T* buffer, size_t size, T min, T max, bool* d_within_range)
{
    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += blockDim.x * gridDim.x)
    {
        const T val = buffer[i];
        if (val < min || val > max)
        {
            *d_within_range = false;
        }
    }
}

template <typename T>
bool invokeCheckRange(const T* buffer, const size_t size, T min, T max, bool* d_within_range, cudaStream_t stream)
{
    cudaMemsetAsync(d_within_range, true, sizeof(bool), stream);

    dim3 block(256);
    dim3 grid((size + 255) / 256);
    check_range<T><<<grid, block, 0, stream>>>(buffer, size, min, max, d_within_range);

    bool result;
    cudaD2Hcpy(&result, d_within_range, 1);
    return result;
}

template bool invokeCheckRange<int>(
    const int* buffer, const size_t size, int min, int max, bool* d_within_range, cudaStream_t stream);

参考文献

• https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/common/memoryUtils.h

点个「赞」+「在看」❤️

让我们知道这份文字有温暖到你，也是我们持续创作的最大动力！