Gahow Wang
7ebdd7c552
From-scratch CUDA kernel for gpt-oss expert weights: one thread per packed byte decodes 2 FP4 (E2M1) codes, applies the per-32-block E8M0 scale (2^(e-127)), and writes BF16 transposed into [IN, OUT] (IN = nblk*32) so it drops straight into x @ W. dequant_mxfp4() wrapper takes raw GpuBuffers (uint8 is not an xserv Tensor dtype). mxfp4-check bin dequants layer-0 expert-0 on GPU and matches tools/mxfp4_probe.py exactly: [0, 0, 0, -0.0625, 0, -0, -0.015625, -0.03125] This lets experts stay MXFP4-resident on GPU (13GB, fits one 32GB card) and be dequantized to a BF16 scratch right before each expert GEMM, instead of holding 36GB of BF16 or uploading experts per token. Loader plumbing + GPU MoE decode use it next. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
74 lines
2.9 KiB
Plaintext
74 lines
2.9 KiB
Plaintext
// MXFP4 dequantization (gpt-oss expert weights) -> BF16.
|
|
//
|
|
// gpt-oss stores each expert MLP weight in OCP Microscaling FP4:
|
|
// blocks: uint8 [OUT, nblk, 16] — each 16-byte row packs 32 FP4 codes
|
|
// (low nibble = even elem, high = odd)
|
|
// scales: uint8 [OUT, nblk] — one E8M0 (8-bit exponent) per 32-elem block
|
|
// value = fp4_e2m1[code] * 2^(scale - 127)
|
|
// The contraction (input) dim is IN = nblk * 32.
|
|
//
|
|
// We emit BF16 in [IN, OUT] (row-major) layout — i.e. the transpose of the
|
|
// natural [OUT, IN] — so it drops straight into y = x[1,IN] @ W[IN,OUT] without
|
|
// a separate transpose. This matches the offline BF16 path (gptoss_dequant.py),
|
|
// keeping the two routes numerically identical.
|
|
|
|
#include <cuda_bf16.h>
|
|
#include <cstdint>
|
|
|
|
// FP4 E2M1 code -> value (OCP MX). 16 codes: sign, 2-bit exp, 1-bit mantissa.
|
|
__constant__ float kFp4[16] = {
|
|
0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f,
|
|
-0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f,
|
|
};
|
|
|
|
// One thread per packed byte: OUT * nblk * 16 threads. Each decodes 2 elements.
|
|
__global__ void mxfp4_dequant_kernel(
|
|
const uint8_t* __restrict__ blocks, // [OUT, nblk, 16]
|
|
const uint8_t* __restrict__ scales, // [OUT, nblk]
|
|
__nv_bfloat16* __restrict__ out, // [IN, OUT], IN = nblk*32
|
|
int out_dim, int nblk
|
|
) {
|
|
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
|
|
long long total = (long long)out_dim * nblk * 16;
|
|
if (idx >= total) return;
|
|
|
|
int b = idx & 15; // byte within block (0..15)
|
|
long long t = idx >> 4; // (out_row, block)
|
|
int block = t % nblk;
|
|
int out_row = t / nblk;
|
|
|
|
uint8_t byte = blocks[idx];
|
|
int code_lo = byte & 0x0F;
|
|
int code_hi = (byte >> 4) & 0x0F;
|
|
|
|
// E8M0 scale: 2^(e - 127). e==0 is the canonical zero scale.
|
|
int e = scales[(long long)out_row * nblk + block];
|
|
float scale = exp2f((float)e - 127.0f);
|
|
|
|
int in_lo = block * 32 + 2 * b; // even element index in the block
|
|
int in_hi = in_lo + 1; // odd element
|
|
|
|
// transposed write: out[in_idx, out_row] = out[in_idx*out_dim + out_row]
|
|
out[(long long)in_lo * out_dim + out_row] = __float2bfloat16(kFp4[code_lo] * scale);
|
|
out[(long long)in_hi * out_dim + out_row] = __float2bfloat16(kFp4[code_hi] * scale);
|
|
}
|
|
|
|
extern "C" {
|
|
|
|
// Dequantize one expert. `blocks`/`scales` point at this expert's slice.
|
|
// Output `out` must hold IN*OUT bf16 (IN = nblk*32). Runs on `stream`.
|
|
void launch_mxfp4_dequant_bf16(
|
|
const void* blocks, const void* scales, void* out,
|
|
int out_dim, int nblk, void* stream
|
|
) {
|
|
long long total = (long long)out_dim * nblk * 16;
|
|
int threads = 256;
|
|
int grid = (int)((total + threads - 1) / threads);
|
|
mxfp4_dequant_kernel<<<grid, threads, 0, (cudaStream_t)stream>>>(
|
|
(const uint8_t*)blocks, (const uint8_t*)scales,
|
|
(__nv_bfloat16*)out, out_dim, nblk
|
|
);
|
|
}
|
|
|
|
} // extern "C"
|