// 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"