obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-24 22:50:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use std::bit_cast for type punning in CPU kernels (#45089)

Browse Source 
Signed-off-by: Yuanyuan Chen <[email protected]>
Co-authored-by: Li, Jiang <[email protected]>
This commit is contained in:
Yuanyuan Chen
2026-06-10 22:07:44 -07:00
committed by GitHub
co-authored by Li, Jiang <[email protected]>
parent 43914dd743
commit 0b995f8609
3 changed files with 36 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cpu/cpu_types_riscv_impl.hpp
+20 -39
View File
@@ -9,10 +9,14 @@
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <torch/all.h>
#include "float_convert.hpp"
namespace vec_op {
// FP8 KV cache is not supported on RISC-V. These tag types and the
@@ -245,8 +249,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[8];
for (int i = 0; i < 8; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
tmp[i] = bf16_to_float(u16[i]);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
}
@@ -256,9 +259,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
u16[i] = float_to_bf16(tmp[i]);
}
}
void save(void* ptr, int elem_num) const {
@@ -266,9 +267,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
u16[i] = float_to_bf16(tmp[i]);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -277,10 +276,8 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
float_to_bf16(tmp[i]);
}
}
};
@@ -292,8 +289,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[16];
for (int i = 0; i < 16; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
tmp[i] = bf16_to_float(u16[i]);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
}
@@ -306,9 +302,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
u16[i] = float_to_bf16(tmp[i]);
}
}
void save(void* ptr, int elem_num) const {
@@ -316,9 +310,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
u16[i] = float_to_bf16(tmp[i]);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -327,10 +319,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
float_to_bf16(tmp[i]);
}
}
};
@@ -343,8 +333,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[32];
for (int i = 0; i < 32; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
tmp[i] = bf16_to_float(u16[i]);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
}
@@ -371,9 +360,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
u16[i] = float_to_bf16(tmp[i]);
}
}
@@ -382,9 +369,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
u16[i] = float_to_bf16(tmp[i]);
}
}
@@ -394,10 +379,8 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
float_to_bf16(tmp[i]);
}
}
};
@@ -985,9 +968,7 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
#else
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
uint32_t val;
std::memcpy(&val, &v, 4);
*reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
*reinterpret_cast<uint16_t*>(ptr) = float_to_bf16(v);
}
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
csrc/cpu/cpu_types_vxe.hpp
+3 -2
View File
@@ -3,7 +3,9 @@
#define CPU_TYPES_VXE_HPP
#include <vecintrin.h>
#include <bit>
#include <cmath>
#include <cstdint>
#include <limits>
#include <torch/all.h>
namespace vec_op {
@@ -817,8 +819,7 @@ inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
// intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
// produce incorrect results for some inputs. Process each of the 4 vectors
// separately.
uint32_t in;
std::memcpy(&in, &v, sizeof(in));
uint32_t in = std::bit_cast<uint32_t>(v);
uint32_t s = (in & 0x80000000) >> 16; // Sign
uint32_t e = (in & 0x7F800000) >> 23; // Exponent
csrc/cpu/float_convert.hpp
+13 -15
View File
@@ -1,14 +1,15 @@
#pragma once
static float bf16_to_float(uint16_t bf16) {
#include <bit>
#include <cstdint>
inline float bf16_to_float(uint16_t bf16) {
uint32_t bits = static_cast<uint32_t>(bf16) << 16;
float fp32;
std::memcpy(&fp32, &bits, sizeof(fp32));
return fp32;
return std::bit_cast<float>(bits);
}
static uint16_t float_to_bf16(float fp32) {
uint32_t bits;
std::memcpy(&bits, &fp32, sizeof(fp32));
inline uint16_t float_to_bf16(float fp32) {
uint32_t bits = std::bit_cast<uint32_t>(fp32);
return static_cast<uint16_t>(bits >> 16);
}
@@ -18,14 +19,13 @@ static uint16_t float_to_bf16(float fp32) {
* Codes below copied from
* https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab
*************************************************/
static uint16_t float_to_fp16(float fp32) {
inline uint16_t float_to_fp16(float fp32) {
uint16_t fp16;
unsigned x;
unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1;
unsigned sign, exponent, mantissa;
std::memcpy(&x, &fp32, sizeof(fp32));
uint32_t x = std::bit_cast<uint32_t>(fp32);
u = (x & 0x7fffffff);
// Get rid of +NaN/-NaN case first.
@@ -77,12 +77,11 @@ static uint16_t float_to_fp16(float fp32) {
return fp16;
}
static float fp16_to_float(uint16_t fp16) {
inline float fp16_to_float(uint16_t fp16) {
unsigned sign = ((fp16 >> 15) & 1);
unsigned exponent = ((fp16 >> 10) & 0x1f);
unsigned mantissa = ((fp16 & 0x3ff) << 13);
int temp;
float fp32;
uint32_t temp;
if (exponent == 0x1f) { /* NaN or Inf */
mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0);
exponent = 0xff;
@@ -101,6 +100,5 @@ static float fp16_to_float(uint16_t fp16) {
exponent += 0x70;
}
temp = ((sign << 31) | (exponent << 23) | mantissa);
std::memcpy(&fp32, &temp, sizeof(temp));
return fp32;
return std::bit_cast<float>(temp);
}