mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 23:43:44 +08:00
Daniel Koepping
57 lines
2.8 KiB
C++
57 lines
2.8 KiB
C++
#include <hip/hip_runtime.h>
|
|
#include <hip/hip_bf16.h>
|
|
#include <hip/hip_fp8.h>
|
|
|
|
typedef short short2v __attribute__((ext_vector_type(2)));
|
|
typedef unsigned uint4v __attribute__((ext_vector_type(4)));
|
|
typedef unsigned short ushort4v __attribute__((ext_vector_type(4)));
|
|
|
|
extern "C" __global__ __launch_bounds__(256) void gptoss_combine_input_quantize_backward(
|
|
__hip_bfloat16 *__restrict__ out, unsigned char *__restrict__ q, unsigned char *__restrict__ scales,
|
|
const __hip_bfloat16 *__restrict__ dout, const int *__restrict__ src_row, const float *__restrict__ weights) {
|
|
const int row = blockIdx.x, group = row / ROWS, tid = threadIdx.x;
|
|
// Four lanes own one MX block, with aligned eight-value loads and stores per lane.
|
|
const int src = src_row[row], lane = tid % 4;
|
|
for (int block = tid / 4; block < 96; block += 64) {
|
|
const int d = block * 32 + lane * 8;
|
|
const long long index = (long long)row * 3072 + d;
|
|
// Neither padding nor unused rows may read dout or weights.
|
|
if (src < 0 || d >= 2880) {
|
|
*reinterpret_cast<uint4v*>(&out[index]) = uint4v{};
|
|
*reinterpret_cast<ushort4v*>(&q[index]) = ushort4v{};
|
|
if (lane == 0) scales[(long long)row * 96 + block] = 0;
|
|
continue;
|
|
}
|
|
const int token = group * TOKENS + src / 4;
|
|
const float weight = weights[token * 4 + src % 4];
|
|
const uint4v input = *reinterpret_cast<const uint4v*>(&dout[(long long)token * 2880 + d]);
|
|
uint4v output;
|
|
float2 values[4];
|
|
float amax = 0.0f;
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; i++) {
|
|
const unsigned bits = input[i];
|
|
const float2 dy = make_float2(__builtin_bit_cast(float, bits << 16), __builtin_bit_cast(float, bits & 0xffff0000u));
|
|
const float2 value = __bfloat1622float2(__float22bfloat162_rn(make_float2(dy.x * weight, dy.y * weight)));
|
|
values[i] = value;
|
|
const __hip_bfloat162 rounded_pair = __float22bfloat162_rn(value);
|
|
output[i] = (unsigned)__bfloat16_as_ushort(rounded_pair.x) | ((unsigned)__bfloat16_as_ushort(rounded_pair.y) << 16);
|
|
amax = fmaxf(amax, fmaxf(fabsf(value.x), fabsf(value.y)));
|
|
}
|
|
*reinterpret_cast<uint4v*>(&out[index]) = output;
|
|
#pragma unroll
|
|
for (int offset = 2; offset > 0; offset >>= 1) amax = fmaxf(amax, __shfl_down(amax, offset, 4));
|
|
amax = __shfl(amax, 0, 4);
|
|
const unsigned e8 = min((__builtin_bit_cast(unsigned, amax) >> 23) & 255u, 254u);
|
|
if (lane == 0) scales[(long long)row * 96 + block] = e8;
|
|
ushort4v quantized;
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; i++) {
|
|
short2v packed = {0, 0};
|
|
packed = __builtin_amdgcn_cvt_scalef32_pk_fp8_f32(packed, values[i].x, values[i].y, __builtin_bit_cast(float, e8 << 23), false);
|
|
quantized[i] = __builtin_bit_cast(unsigned short, packed[0]);
|
|
}
|
|
*reinterpret_cast<ushort4v*>(&q[index]) = quantized;
|
|
}
|
|
}
|