TensorRT-LLM 0.5.0 源码之二十八

cudaTypeUtils

#include "tensorrt_llm/common/cudaBf16Fallbacks.cuh"
#include "tensorrt_llm/common/cudaBf16Wrapper.h"
#include "tensorrt_llm/common/cudaFp8Utils.h"
#include <cuda.h>
#include <cuda_fp16.h>
#if ENABLE_BF16
#include <cuda_bf16.h>
#endif

ldg

template <typename T>
inline __device__ T ldg(const T* val)
{
    return __ldg(val);
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 ldg(const __nv_bfloat162* val)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
    return val[0];
#else
    return __ldg(val);
#endif
}

template <>
inline __device__ __nv_bfloat16 ldg(const __nv_bfloat16* val)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
    return val[0];
#else
    return __ldg(val);
#endif
}
#endif // ENABLE_BF16

这是一个用于在 CUDA核函数 中高效读取只读全局内存 的封装函数。它的核心作用是利用GPU的只读数据缓存来提升内存访问性能。

下面这个表格能帮你快速抓住要点：

特性	解释
核心功能	封装CUDA内置函数 `__ldg`，通过只读数据缓存（Read-Only Cache）读取全局内存数据。
性能优势	减少对通用L1缓存的压力，避免缓存污染，尤其在多线程访问相同数据时能提高带宽利用率和性能。
适用场景	主要用于处理只读的全局内存数据，例如大型查表、常量参数、输入数据流等。
关键限制	需要计算能力（Compute Capability）≥ 3.5的GPU。在新架构（如Ampere、Hopper）上，其优化效果可能减弱或与普通访问无异。

💡 深入理解其工作机理与价值

• 它是如何工作的：在没有只读缓存的情况下，从全局内存读取的数据会进入通用的L1和L2缓存。如果这些数据只读但被频繁访问，它们会占用大量可写缓存空间，可能“挤占”掉那些需要频繁写入的数据，导致缓存效率下降（这被称为“缓存污染”）。__ldg 指令让数据通过一个独立的只读数据通道被加载到SM（流多处理器）上的只读缓存中。这片缓存是专门为只读数据设计的，从而避免了与可写缓存区域的竞争，提升了整体缓存利用效率。
• 为什么需要这样封装：CUDA内置的 __ldg 本身就是一个编译器内置函数（intrinsic）。这里用模板函数对其进行封装，主要有两个好处：

1. 类型安全：模板能自动推导类型，使代码更通用且安全，你无需关心底层类型转换。
2. 可读性与可维护性：ldg(val) 的写法比 __ldg(val) 在语法上更直观。此外，如果需要在未来架构上移除或修改此优化，只需修改这个封装函数即可，非常方便。

⚠️ 重要注意事项

• 架构差异：这个函数的优化效果与硬件架构紧密相关。在Kepler和Maxwell等老架构上，独立的只读缓存能带来显著的性能提升。然而，在Volta、Turing及之后的现代架构（如Ampere、Hopper）中，缓存体系已经统一，所有加载操作默认都通过统一的L1缓存，此时 __ldg 的作用已经减弱，甚至与普通的指针访问没有区别。
• 编译要求：要使用 __ldg，在编译CUDA代码时必须指定足够的计算能力。例如，对于计算能力为6.1的GTX 1060显卡，需要在CMakeLists.txt中设置 -arch=compute_61 -code=sm_61。

💎 总结

你提供的这个模板函数是一个针对只读全局内存访问的性能优化封装。在支持的GPU架构上，对大规模只读数据使用它能带来可观的性能收益。但在较新的GPU上，其必要性可能降低。

TypeConverter

// Get type2 from type or vice versa (applied to half and bfloat16)
template <typename T>
struct TypeConverter
{
    using Type = half2;
}; // keep for generality

template <>
struct TypeConverter<half2>
{
    using Type = half;
};

template <>
struct TypeConverter<half>
{
    using Type = half2;
};

#if ENABLE_BF16
template <>
struct TypeConverter<__nv_bfloat162>
{
    using Type = __nv_bfloat16;
};

template <>
struct TypeConverter<__nv_bfloat16>
{
    using Type = __nv_bfloat162;
};
#endif // ENABLE_BF16

hadd2

// Defined math operations (bfloat16 fallback to fp32 when it is not supported)
template <typename T>
inline __device__ T hadd2(T a, T b)
{
    return __hadd2(a, b);
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 hadd2(__nv_bfloat162 a, __nv_bfloat162 b)
{
    return bf16hadd2(a, b);
}
#endif // ENABLE_BF16

add

template <typename T>
inline __device__ T add(T a, T b)
{
    return a + b;
}

template <>
inline __device__ half2 add(half2 a, half2 b)
{
    return __hadd2(a, b);
}

template <>
inline __device__ half add(half a, half b)
{
    return __hadd(a, b);
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b)
{
    return bf16hadd2(a, b);
}

template <>
inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b)
{
    return bf16hadd(a, b);
}

inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, float b)
{
    return bf16hadd(a, __float2bfloat16(b));
}
#endif // ENABLE_BF16

// applies to all 4 values addition
template <typename T>
inline __device__ T add(T a, T b, T c)
{
    return a + b + c;
}

#if ENABLE_BF16
inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c)
{
    return bf16hadd(a, b, c);
}

inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c)
{
    return bf16hadd2(a, b, c);
}
#endif // ENABLE_BF16

// applies to all 4 values addition
template <typename T>
inline __device__ T add(T a, T b, T c, T d)
{
    return (T) ((float) a + (float) b + (float) c + (float) d);
}

#if ENABLE_BF16
inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c, __nv_bfloat16 d)
{
    return bf16hadd(a, b, c, d);
}
#endif // ENABLE_BF16

hsub2

template <typename T>
inline __device__ T hsub2(T a, T b)
{
    return __hsub2(a, b);
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 hsub2(__nv_bfloat162 a, __nv_bfloat162 b)
{
    return bf16hsub2(a, b);
}
#endif // ENABLE_BF16

hmul2

template <typename T>
inline __device__ T hmul2(T a, T b)
{
    return __hmul2(a, b);
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 hmul2(__nv_bfloat162 a, __nv_bfloat162 b)
{
    return bf16hmul2(a, b);
}
#endif // ENABLE_BF16

template <typename T>
inline __device__ T hmul2(T a, T b, T c)
{
    return a * b * c;
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 hmul2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c)
{
    return bf16hmul2(a, b, c);
}
#endif // ENABLE_BF16

mul

template <typename T>
inline __device__ T mul(T a, T b, T c)
{
    return a * b * c;
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c)
{
    return bf16hmul(a, b, c);
}

inline __device__ __nv_bfloat162 mul(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c)
{
    return bf16hmul2(a, b, c);
}
#endif // ENABLE_BF16

fma

template <typename T>
inline __device__ T fma(T a, T b, T c, T d)
{
    return a * b * c + d;
}

#if ENABLE_BF16
inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c, __nv_bfloat162 d)
{
    return bf16hfma2(a, b, c, d);
}
#endif // ENABLE_BF16

template <typename T>
inline __device__ T fma(T a, T b, T c)
{
    return a * b + c;
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c)
{
    return bf16hfma2(a, b, c);
}

template <>
inline __device__ __nv_bfloat16 fma(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c)
{
    return bf16hfma(a, b, c);
}
#endif // ENABLE_BF16

hexp2

template <typename T>
inline __device__ T hexp2(T a)
{
    return h2exp(a);
}

#if ENABLE_BF16
template <>
inline __device__ __nv_bfloat162 hexp2(__nv_bfloat162 a)
{
    return bf16exp2(a);
}
#endif // ENABLE_BF16

cuda_cast

template <typename T_OUT, typename T_IN>
__device__ inline T_OUT cuda_cast(T_IN val)
{
    return val;
}

template <>
__device__ inline float2 cuda_cast<float2, int2>(int2 val)
{
    return make_float2(val.x, val.y);
}

template <>
__device__ inline float2 cuda_cast<float2, float>(float val)
{
    return make_float2(val, val);
}

template <>
__device__ inline float2 cuda_cast<float2, half2>(half2 val)
{
    return __half22float2(val);
}

template <>
__device__ inline half2 cuda_cast<half2, float2>(float2 val)
{
    return __float22half2_rn(val);
}

template <>
__device__ inline half2 cuda_cast<half2, float>(float val)
{
    return __float2half2_rn(val);
}

template <>
__device__ inline half2 cuda_cast<half2, half>(half val)
{
    return __half2half2(val);
}

template <>
__device__ inline int8_t cuda_cast<int8_t, half>(half val)
{
    union
    {
        int8_t int8[2];
        int16_t int16;
    };

    union
    {
        half fp16;
        int16_t int16_in;
    };

    fp16 = val;
asm volatile("cvt.rni.sat.s8.f16 %0, %1;" : "=h"(int16) : "h"(int16_in));
    return int8[0];
}

template <>
__device__ inline int16_t cuda_cast<int16_t, half2>(half2 val)
{
    union
    {
        int8_t int8[2];
        int16_t int16;
    };

    int8[0] = cuda_cast<int8_t>(val.x);
    int8[1] = cuda_cast<int8_t>(val.y);
    return int16;
}

template <>
__device__ inline int8_t cuda_cast<int8_t, float>(float val)
{
    union
    {
        int8_t int8[2];
        int16_t int16;
    };

asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val));
    return int8[0];
}

template <>
__device__ inline int16_t cuda_cast<int16_t, float2>(float2 val)
{
    union
    {
        int8_t int8[2];
        int16_t int16;
    };

    int8[0] = cuda_cast<int8_t>(val.x);
    int8[1] = cuda_cast<int8_t>(val.y);
    return int16;
}

template <>
__device__ inline half2 cuda_cast<half2, int16_t>(int16_t val)
{
    union
    {
        int8_t int8[2];
        int16_t int16;
    };

    int16 = val;
    return make_half2(int8[0], int8[1]);
}

template <>
__device__ inline float2 cuda_cast<float2, int16_t>(int16_t val)
{
    union
    {
        int8_t int8[2];
        int16_t int16;
    };

    int16 = val;
    return make_float2(int8[0], int8[1]);
}

#ifdef ENABLE_BF16
template <>
__device__ inline __nv_bfloat16 cuda_cast(int32_t val)
{
    return static_cast<float>(val);
}

template <>
__device__ inline __nv_bfloat16 cuda_cast(int8_t val)
{
    return static_cast<float>(val);
}

template <>
__device__ inline int8_t cuda_cast(__nv_bfloat16 val)
{
    return static_cast<float>(val);
}

template <>
__device__ inline float cuda_cast<float, __nv_bfloat16>(__nv_bfloat16 val)
{
    return __bfloat162float(val);
}

template <>
__device__ inline float2 cuda_cast<float2, __nv_bfloat162>(__nv_bfloat162 val)
{
    return bf1622float2(val);
}

template <>
__device__ inline half cuda_cast<half, __nv_bfloat16>(__nv_bfloat16 val)
{
    return __float2half(__bfloat162float(val));
}

template <>
__device__ inline int16_t cuda_cast<int16_t, __nv_bfloat162>(__nv_bfloat162 val)
{
    return bf1622int16(val);
}

template <>
__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val)
{
    return __float2bfloat16(val);
}

template <>
__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val)
{
    return __float2bfloat16(__half2float(val));
}

template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val)
{
    return bf162bf162(val);
}

template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val)
{
    return __float2bfloat162_rn(val);
}

template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val)
{
    return float22bf162(val);
}

template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val)
{
    union
    {
        int8_t int8[2];
        int16_t int16;
    };

    int16 = val;
    __nv_bfloat162 res;
    res.x = cuda_cast<__nv_bfloat16>(int8[0]);
    res.y = cuda_cast<__nv_bfloat16>(int8[1]);
    return res;
}

template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val)
{
    return float22bf162(__half22float2(val));
}

#endif // ENABLE BF16

这行代码是CUDA编程中使用的PTX内联汇编指令，用于在GPU上执行一个非常特定的数据转换操作。它的功能是将一个32位的浮点数(f32)按照特定的舍入和饱和处理规则，转换为8位的有符号整数(s8)，并最终存储到一个16位的整型变量中。

下面这个表格详细解释了指令中每个部分的作用：

组成部分	说明
`asm volatile`	`asm` 声明内联汇编。`volatile` 告诉编译器不要优化或移动此段汇编代码。
`"cvt.rni.sat.s8.f32 %0, %1;"`	核心的PTX转换指令。`%1`是输入操作数，`%0`是输出操作数。
└── `cvt`	表示转换（convert）操作。
└── `.rni`	舍入模式：向最接近的偶数取整（round to nearest even）。
└── `.sat`	饱和处理：若结果超出目标类型范围，则钳制在边界值（如`127`或`-128`）。
└── `.s8.f32`	定义转换方向：从32位浮点数（`.f32`）到8位有符号整数（`.s8`）。
`: "=h"(int16)`	输出操作数部分。将汇编输出（`%0`）关联到C变量`int16`。`"=h"`约束表示使用16位寄存器（如`ax`的低16位）。
`: "f"(val)`	输入操作数部分。将C变量`val`的值作为输入（`%1`）传递给汇编。`"f"`约束告知编译器使用浮点寄存器。

💡 指令执行流程与注意事项

这条指令的执行逻辑如下：

1. 读取输入：从由"f"约束指定的浮点寄存器中取出变量val的32位浮点数值。
2. 转换计算：执行cvt.rni.sat.s8.f32指令。

• 首先，按照向最接近偶数取整的模式将浮点数舍入为整数。
• 然后，检查该整数值是否在8位有符号整数范围内（-128 到 127）。如果超出，则使用饱和处理（.sat），将其钳制到最接近的边界值（-128或127）。

3. 处理输出：将得到的8位有符号整数结果，通过"=h"约束存入一个16位寄存器的低8位。变量int16最终获得的是一个16位的整数，但其低8位包含了有效的转换结果，高8位可能是未定义的。

⚠️ 一个重要提示

虽然转换结果是一个8位值（.s8），但它被存储到了一个16位的整型变量（int16）中。这意味着，当你使用int16的值时，需要意识到只有它的低8位是有效的。高8位的内容可能是不确定的。如果你需要确保得到一个纯粹的8位整数，可能需要进行额外的位掩码操作，例如 int8_result = int16 & 0xFF;。

#ifdef ENABLE_FP8
template <>
__device__ inline float2 cuda_cast<float2, __nv_fp8x2_e4m3>(__nv_fp8x2_e4m3 val)
{
    return bf1622float2(fp8x2_e4m3_to_bfloat2(&val));
}

template <>
__device__ inline half2 cuda_cast<half2, __nv_fp8x2_e4m3>(__nv_fp8x2_e4m3 val)
{
    return fp8x2_e4m3_to_half2(&val);
}

template <>
__device__ inline __nv_fp8x2_e4m3 cuda_cast<__nv_fp8x2_e4m3, float2>(float2 val)
{
    return __nv_fp8x2_e4m3(bf1622float2(float22bf162(val)));
}

template <>
__device__ inline __nv_fp8x2_e4m3 cuda_cast<__nv_fp8x2_e4m3, half2>(half2 val)
{
    return __nv_fp8x2_e4m3(cuda_cast<float2>(val));
}

template <>
__device__ inline __nv_fp8x2_e4m3 cuda_cast<__nv_fp8x2_e4m3, __nv_bfloat162>(__nv_bfloat162 val)
{
    return __nv_fp8x2_e4m3(cuda_cast<float2>(val));
}

template <>
__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, half>(half val)
{
    return __nv_fp8_e4m3(val);
}

template <>
__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, __nv_bfloat16>(__nv_bfloat16 val)
{
    return __nv_fp8_e4m3(val);
}

template <>
__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, float>(float val)
{
    return __nv_fp8_e4m3(val);
}

template <>
__device__ inline float cuda_cast<float, __nv_fp8_e4m3>(__nv_fp8_e4m3 val)
{
    return (float) val;
}

template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_fp8x2_e4m3>(__nv_fp8x2_e4m3 val)
{
    return fp8x2_e4m3_to_bfloat2(&val);
}

template <>
__device__ inline int8_t cuda_cast<int8_t, __nv_fp8_e4m3>(__nv_fp8_e4m3 val)
{
    // no impl
    return 0;
}

template <>
__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, int8_t>(int8_t val)
{
    return cuda_cast<__nv_fp8_e4m3>(cuda_cast<__nv_bfloat16>(cuda_cast<float>(val)));
}

#endif // ENABLE_FP8

你观察到的这种“迂回”的转换路径 cuda_cast<__nv_fp8_e4m3>(cuda_cast<__nv_bfloat16>(cuda_cast<float>(val)));，其核心目的在于确保数值精度，并遵循一条安全、通用的转换规则。

下面这个表格详细拆解了每一步的动机和作用。

🔄 转换路径详解

转换步骤	转换函数示例	作用与原因分析
第1步：提升到高精度	`cuda_cast<float>(val)`	将`int8_t`整数先转换为`float`。这是因为`float`拥有极高的动态范围和精度，可以无损地表示`int8_t`的所有整数值（-128到127），为后续所有操作提供一个稳定且无精度损失的起点。
第2步：过渡到中等精度	`cuda_cast<__nv_bfloat16>(...)` )	将`float`转换为`__nv_bfloat16`。这是一个可选的精度过渡步骤。`bfloat16`虽然精度低于`float`，但其表示范围与`float`相当，能较好地容纳来自`float`的数值，同时其存储大小更接近目标类型`fp8`。
第3步：量化到目标精度	`cuda_cast<__nv_fp8_e4m3>(...)` )	最终将`__nv_bfloat16`转换为`__nv_fp8_e4m3`。这是实际的量化步骤，将数据精度降低到目标格式。由于`bfloat16`到`fp8`的转换是直接支持的，这一步通常有高效的硬件指令实现。

💡 为何不直接转换？

你可能会问，为什么不直接从int8_t转换到__nv_fp8_e4m3呢？这主要是因为：

1. 缺乏直接的硬件指令支持：GPU（特别是CUDA核心和Tensor Core）可能没有为这种特定且不常见的类型转换（整型到非标准低精度浮点）设计专用的硬件电路。强制直接转换可能需要复杂的位操作，效率低下且容易出错。
2. 数值精度与安全性的权衡：如表格所示，通过float甚至__nv_bfloat16进行中转，相当于为转换过程设置了“缓冲带”。尤其是在有float参与的情况下，可以最大限度地保留原始数值信息，避免在一步到位的直接转换中可能发生的不可预见的舍入误差或溢出问题。这是一种谨慎且通用的做法。
3. 实现上的通用性与清晰性：编写转换函数时，优先利用已有的、经过充分测试的基础转换（如int8_t到float，float到bfloat16，bfloat16到fp8）来组合成复杂转换，代码会更清晰、更易于维护，也能保证在不同硬件平台和计算架构上行为的一致性。

💎 总结

总而言之，这条看似曲折的转换路径，实际上是深度学习和高性能计算中一种常见的设计哲学体现：在追求极致效率（使用FP8）的同时，通过高精度中间媒介来保障数值转换的可靠性和准确性。这是一种在性能和精度之间取得的精心权衡。

cuda_abs

template <typename T>
__device__ inline T cuda_abs(T val)
{
    assert(false);
    return {};
}

template <>
__device__ inline float cuda_abs(float val)
{
    return fabs(val);
}

template <>
__device__ inline float2 cuda_abs(float2 val)
{
    return make_float2(fabs(val.x), fabs(val.y));
}

template <>
__device__ inline half cuda_abs(half val)
{
    return __habs(val);
}

template <>
__device__ inline half2 cuda_abs(half2 val)
{
    return __habs2(val);
}

#ifdef ENABLE_BF16

#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)
template <>
__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val)
{
    return __habs(val);
}

template <>
__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val)
{
    return __habs2(val);
}
#endif

#endif // ENABLE_FP16

cuda_sum

template <typename To, typename Ti>
__device__ inline To cuda_sum(Ti val)
{
    return cuda_cast<To>(val);
};

template <typename To>
__device__ inline To cuda_sum(float2 val)
{
    return cuda_cast<To>(val.x + val.y);
};

// Unary maximum: compute the max of a vector type
template <typename To, typename Ti>
__device__ inline To cuda_max(Ti val)
{
    return cuda_cast<To>(val);
};

template <>
__device__ inline float cuda_max(float2 val)
{
    return fmaxf(val.x, val.y);
}

template <>
__device__ inline half cuda_max(half2 val)
{
    return __hmax(val.x, val.y);
}

#ifdef ENABLE_BF16
template <>
__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val)
{
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
    return __hmax(val.x, val.y);
#endif
}
#endif

// Binary maximum: compute the max of two scalar types
template <typename T>
__device__ inline T cuda_max(T val1, T val2)
{
    return (val1 > val2) ? val1 : val2;
}

参考文献

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

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

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