大模型加速技术「TOPK基数排序-位运算的使用」

对 float 类型进行基数排序时,不能直接按十进制位处理,因为浮点数在内存中是以 IEEE 754 标准存储的。正确的方法是将 float 重新解释为整数(bitcast),然后利用整数基数排序,但必须修正符号位的影响,使其保持正确的数值顺序。

基数排序通常按 4-bit(16 个桶) 或 8-bit(256 个桶) 分段处理(称为“radix”或“pass”)。对 32 位整数,常用 4-pass × 8-bit。

利用下面的方法可以保证正数和负数的有序性。

IEEE 754 浮点数表示:| S (1 bit) | Exponent (8 bits) | Mantissa (23 bits) |

__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t { uint32_t bits = __float_as_uint(x); return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);}

原理:负数按位取反,首位变为0,并且满足单调性,负数越大,取反后也越大。正数首位变为1(无符号数)单调递增。这样正数和负数都可以在uint32_t中进行有效排序。

单调性的证明如下: