From 60d0fd175e2457574dee94577f12992eba67c67c Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Mon, 7 Sep 2026 15:11:17 -0700 Subject: [PATCH 1/2] Add OpenAI SwiGLU benchmarks Add CUDA, HIP, OpenMP, and SYCL implementations of the fused SwiGLU activation with low-precision input and output modes. --- src/swiglu-oai-cuda/CMakeLists.txt | 8 + src/swiglu-oai-cuda/Makefile | 88 ++++ src/swiglu-oai-cuda/main.cu | 814 ++++++++++++++++++++++++++++ src/swiglu-oai-cuda/reference.h | 672 ++++++++++++++++++++++++ src/swiglu-oai-hip/CMakeLists.txt | 8 + src/swiglu-oai-hip/Makefile | 88 ++++ src/swiglu-oai-hip/main.cu | 818 +++++++++++++++++++++++++++++ src/swiglu-oai-hip/reference.h | 173 ++++++ src/swiglu-oai-omp/CMakeLists.txt | 8 + src/swiglu-oai-omp/Makefile | 62 +++ src/swiglu-oai-omp/Makefile.aomp | 66 +++ src/swiglu-oai-omp/Makefile.nvc | 63 +++ src/swiglu-oai-omp/main.cpp | 131 +++++ src/swiglu-oai-sycl/CMakeLists.txt | 8 + src/swiglu-oai-sycl/Makefile | 102 ++++ src/swiglu-oai-sycl/main.cpp | 788 +++++++++++++++++++++++++++ 16 files changed, 3897 insertions(+) create mode 100644 src/swiglu-oai-cuda/CMakeLists.txt create mode 100644 src/swiglu-oai-cuda/Makefile create mode 100644 src/swiglu-oai-cuda/main.cu create mode 100644 src/swiglu-oai-cuda/reference.h create mode 100644 src/swiglu-oai-hip/CMakeLists.txt create mode 100644 src/swiglu-oai-hip/Makefile create mode 100644 src/swiglu-oai-hip/main.cu create mode 100644 src/swiglu-oai-hip/reference.h create mode 100644 src/swiglu-oai-omp/CMakeLists.txt create mode 100644 src/swiglu-oai-omp/Makefile create mode 100644 src/swiglu-oai-omp/Makefile.aomp create mode 100644 src/swiglu-oai-omp/Makefile.nvc create mode 100644 src/swiglu-oai-omp/main.cpp create mode 100644 src/swiglu-oai-sycl/CMakeLists.txt create mode 100644 src/swiglu-oai-sycl/Makefile create mode 100644 src/swiglu-oai-sycl/main.cpp diff --git a/src/swiglu-oai-cuda/CMakeLists.txt b/src/swiglu-oai-cuda/CMakeLists.txt new file mode 100644 index 000000000..068c01a7d --- /dev/null +++ b/src/swiglu-oai-cuda/CMakeLists.txt @@ -0,0 +1,8 @@ +# swiglu-oai-cuda/CMakeLists.txt + +add_hecbench_benchmark( + NAME swiglu-oai + MODEL cuda + SOURCES main.cu + CATEGORIES algorithms +) diff --git a/src/swiglu-oai-cuda/Makefile b/src/swiglu-oai-cuda/Makefile new file mode 100644 index 000000000..0ebda35a6 --- /dev/null +++ b/src/swiglu-oai-cuda/Makefile @@ -0,0 +1,88 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = nvcc +OPTIMIZE = yes +DEBUG = no +# __nv_fp8_e4m3 (cuda_fp8.h) needs CUDA >= 11.8; __nv_fp4x2_e2m1 +# (cuda_fp4.h) needs CUDA >= 12.8. FP4 conversion is native on +# architecture-specific Blackwell targets such as sm_100a; other targets use +# CUDA's emulation path. sm_90 keeps this benchmark runnable on Hopper while +# exercising native FP8 support. +ARCH = sm_90 +LAUNCHER ?= + +# Run options. The default run covers five representative production paths. +# Override the shape or repeat count on the make command line, for example: +# make run ROWS=128 DIM=256 REPEAT=10 +ROWS ?= 8192 +DIM ?= 1024 +REPEAT ?= 100 + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cu + +obj = $(source:.cu=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Xcompiler -Wall -arch=$(ARCH) + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g -DDEBUG + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cu reference.h Makefile + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +.PHONY: clean run + +run: $(program) + @run_case() { \ + output="$$($(LAUNCHER) ./$(program) "$$@" 2>&1)"; status=$$?; \ + printf '%s\n' "$$output"; \ + if [ $$status -ne 0 ] || \ + ! printf '%s\n' "$$output" | awk '$$0 == "PASS" { pass=1 } END { exit !pass }'; then \ + return 1; \ + fi; \ + }; \ + set -e; \ + # Standard LLM inference: BF16 activation, no bias, no output quantization. \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 none none; \ + # Bias-bearing FP16 model. \ + run_case $(ROWS) $(DIM) $(REPEAT) fp16 none fp16; \ + # vLLM/SGLang-style fused SwiGLU followed by FP8 activation quantization. \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 fp8 none; \ + # Blackwell-style fused SwiGLU followed by MXFP4 activation quantization. \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 mxfp4 none; \ + # Low-precision MoE path with MXFP8 input, FP8 output, and BF16 bias. \ + run_case $(ROWS) $(DIM) $(REPEAT) mxfp8 fp8 bf16 diff --git a/src/swiglu-oai-cuda/main.cu b/src/swiglu-oai-cuda/main.cu new file mode 100644 index 000000000..bcfa980a8 --- /dev/null +++ b/src/swiglu-oai-cuda/main.cu @@ -0,0 +1,814 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(__CUDACC_VER_MAJOR__) || !defined(__CUDACC_VER_MINOR__) +#error "swiglu-oai-cuda must be compiled with the NVIDIA CUDA compiler" +#elif (__CUDACC_VER_MAJOR__ < 12) || \ + (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8) +#error "swiglu-oai-cuda requires CUDA Toolkit 12.8 or newer for FP4 support" +#endif + +#include +#include "reference.h" + +#define CUDA_CHECK(call) \ + do { \ + const cudaError_t error_ = (call); \ + if (error_ != cudaSuccess) { \ + fprintf(stderr, "CUDA error at %s:%d: %s failed: %s\n", \ + __FILE__, __LINE__, #call, cudaGetErrorString(error_)); \ + return EXIT_FAILURE; \ + } \ + } while (0) + +static constexpr int WarpSize = 32; +static_assert(MX_BLOCK == WarpSize, + "The MXFP4 kernel requires one 32-lane warp per MX block"); +static_assert(sizeof(size_t) >= 8, + "This benchmark requires 64-bit size_t indexing"); + +// --------------------------------------------------------------------------- +// Native low-precision helpers, shared by host and device +// --------------------------------------------------------------------------- + +__host__ __device__ __forceinline__ +__nv_fp4x2_e2m1 fp4x2_pack(const float lo, const float hi) +{ + __nv_fp4x2_e2m1 packed; + packed.__x = __nv_cvt_float2_to_fp4x2( + make_float2(lo, hi), __NV_E2M1, cudaRoundNearest); + return packed; +} + +__host__ __device__ __forceinline__ +float fp4x2_unpack(const __nv_fp4x2_storage_t byte, const int hi) +{ + const __half2_raw h2 = __nv_cvt_fp4x2_to_halfraw2(byte, __NV_E2M1); + __half_raw h; + h.x = hi ? h2.y : h2.x; + return __half2float(__half(h)); +} + +// FP32 reserves exponent 0 for zero/subnormals while E8M0 treats exponent 0 as 2^{-127} +// TODO: b = 0xFF +__host__ __device__ __forceinline__ +float e8m0_decode(const uint8_t b) +{ + const uint32_t bits = b ? static_cast(b) << 23 : 0x00400000u; +#ifdef __CUDA_ARCH__ + return __uint_as_float(bits); +#else + float value; + memcpy(&value, &bits, sizeof(value)); + return value; +#endif +} + +// expf(-alpha*gate) overflows for large negative gate. +__device__ __forceinline__ float dev_sigmoid(const float x) +{ + if (x >= 0.f) { + return 1.f / (1.f + expf(-x)); + } else { + const float ex = expf(x); + return ex / (1.f + ex); + } +} + +// Input load: fetch the native element and undo the quantization scale. + +template +__device__ __forceinline__ float load_x_raw(const void* __restrict__ X, + const size_t i) +{ + if constexpr (STORE == XS_F32) { + return static_cast(X)[i]; + } else if constexpr (STORE == XS_F16) { + return __half2float(static_cast(X)[i]); + } else if constexpr (STORE == XS_BF16) { + return __bfloat162float(static_cast(X)[i]); + } else if constexpr (STORE == XS_FP8) { + return static_cast(static_cast(X)[i]); + } else { + return fp4x2_unpack(static_cast(X)[i >> 1].__x, + static_cast(i & 1)); + } +} + +// Load one adjacent gate/linear pair. FP4 stores the pair in one byte and FP8 +// stores it in one aligned 16-bit word, so one vector conversion decodes both +// values. Other formats retain their ordinary scalar loads. +template +__device__ __forceinline__ float2 load_x_pair( + const void* __restrict__ X, + const uint8_t* __restrict__ xscales, + const float xtscale, + const size_t i, + const size_t scaleIdx) +{ + float2 pair; + if constexpr (STORE == XS_FP4) { + const __nv_fp4x2_storage_t packed = + static_cast(X)[i >> 1].__x; + const __half2_raw h2 = __nv_cvt_fp4x2_to_halfraw2(packed, __NV_E2M1); + __half_raw lo, hi; + lo.x = h2.x; + hi.x = h2.y; + pair = make_float2(__half2float(__half(lo)), __half2float(__half(hi))); + } else if constexpr (STORE == XS_FP8) { + const __nv_fp8x2_storage_t packed = + static_cast(X)[i >> 1].__x; + const __half2_raw h2 = __nv_cvt_fp8x2_to_halfraw2(packed, __NV_E4M3); + __half_raw lo, hi; + lo.x = h2.x; + hi.x = h2.y; + pair = make_float2(__half2float(__half(lo)), __half2float(__half(hi))); + } else { + pair = make_float2(load_x_raw(X, i), + load_x_raw(X, i + 1)); + } + + float scale = 1.f; + if constexpr (SCALE == XSC_TENSOR) scale = xtscale; + else if constexpr (SCALE == XSC_MX) + scale = e8m0_decode(xscales[scaleIdx]); + pair.x *= scale; + pair.y *= scale; + return pair; +} + +// Bias is a compile-time choice. In the BS_NONE case the +// compiler replaces this whole function with 0 and removes the bias pointer +// and global-memory load. FP16/BF16 bias is converted to float before the +// activation math. BS_F32 exists for the accuracy-only test path. +template +__device__ __forceinline__ float load_bias(const void* __restrict__ B, + const int d) +{ + if constexpr (BIAS == BS_NONE) { + (void)B; + (void)d; + return 0.f; + } else if constexpr (BIAS == BS_F16) { + return __half2float(static_cast(B)[d]); + } else if constexpr (BIAS == BS_BF16) { + return __bfloat162float(static_cast(B)[d]); + } else { + return static_cast(B)[d]; + } +} + +// Compute one output value from one adjacent gate/linear pair. +// +// Each input row contains 2*dim values in this order: +// +// gate[0], linear[0], gate[1], linear[1], ... +// +// Therefore output column d reads input columns 2*d and 2*d+1. Low-precision +// input values are first converted to float and multiplied by either their +// tensor-wide scale or their 32-value MX-block scale. Bias is then added in +// float. The gate is capped only above, while the linear value is capped on +// both sides: +// +// y = gate * sigmoid(alpha * gate) * (linear + 1) +// +// This helper is shared by every output path. +template +__device__ __forceinline__ float swiglu_elem( + const KernelArgs& a, const int64_t row, const int64_t d) +{ + const int64_t index = row * a.dim + d; + const int64_t xIndex = 2 * index; + // 2*d is even and 2*d+1 is odd, so the pair can never be split across an + // MX_BLOCK boundary: one scale lookup covers both halves. + int64_t scaleIdx = 0; + if constexpr (SCALE == XSC_MX) { + // LLM dimensions are normally multiples of 16. In that case input MX + // blocks have no row padding, so the output's flat index maps directly to + // its input scale. Retain the row-aware path for arbitrary dimensions. + scaleIdx = (a.dim % (MX_BLOCK / 2) == 0) + ? index / (MX_BLOCK / 2) + : row * a.xBlocksPerRow + (2 * d) / MX_BLOCK; + } + + const float2 pair = load_x_pair( + a.X, a.Xs, a.xScale, xIndex, scaleIdx); + float gate = pair.x + load_bias(a.B, 2 * d); + float linear = pair.y + load_bias(a.B, 2 * d + 1); + + gate = fminf(gate, a.limit); + linear = fminf(fmaxf(linear, -a.limit), a.limit); + return gate * dev_sigmoid(a.alpha * gate) * (linear + 1.f); +} + +/* + Description of the fused kernel: + + Each input row contains interleaved gate and linear values: + + gate[0], linear[0], gate[1], linear[1], ... + + For output column d, it computes: + + gate = min(input[2*d] + bias[2*d], limit) + linear = clamp(input[2*d+1] + bias[2*d+1], -limit, limit) + output[d] = gate × sigmoid(alpha × gate) × (linear + 1) + + Thus an input shaped [rows, 2 × dim] produces [rows, dim]. + + Input loading + + The kernel supports FP32, FP16, BF16, FP8, and FP4 storage. + + 1. It loads the stored value. + 2. FP8/FP4 values are converted to float. + 3.1 For tensor-scaled formats, it multiplies by one tensor-wide scale. + 3.2 For MX formats, it multiplies by the E8M0 scale belonging to that 32-value block. + 4. All activation arithmetic then happens in FP32. + + FP4 stores two values per byte, so the unpacking helper selects either the low or high value. + + Bias handling + + Bias is selected at compile time: + + • none: the compiler completely removes the bias load. + • fp16: loaded and converted to float. + • bf16: loaded and converted to float. + • fp32: loaded directly, but used only for accuracy testing. + + This avoids a runtime branch for every output element. + + FP32/FP8 output path + + A single fused kernel serves every output path (matching the HIP and SYCL + ports); one CUDA thread computes one output element. + + The thread maps its tiled index to a row and column, loads the gate/linear pair, evaluates SwiGLU, and + then: + + • writes the float directly for FP32 output, or + • divides by the tensor-wide output scale, rounds to E4M3, and writes FP8. + + Threads are tiled as (row, block of blockDim.x columns) rather than laid out flat, so that an MXFP4 group of + 32 values always lands on one 32-lane warp. A grid-stride loop over int64 tile ids lets each warp process + additional groups when the tensor exceeds the launched grid, and keeps the index and guard 64-bit safe. + + MXFP4 output path + + MXFP4 requires one shared scale for every group of 32 output values. Therefore, one 32-thread CUDA warp owns + one MXFP4 group. + + 1. Lane 0 computes value 0, lane 1 computes value 1, and so forth. + 2. Missing values in a partial final group are represented as zero. + 3. The warp finds the largest absolute value using five shuffle-reduction steps: offsets 16, 8, 4, 2, and + 1. + 4. Lane 0 converts that maximum’s exponent into an E8M0 scale byte and stores it. + 5. Lane 0 broadcasts the scale through warp registers. + 6. Every value is divided by the shared scale and rounded to E2M1. + 7. Even-numbered lanes pack adjacent pairs: + • lane 0 packs values 0 and 1 + • lane 2 packs values 2 and 3 + • … + • lane 30 packs values 30 and 31 + + The 16 even lanes therefore write 16 bytes containing 32 FP4 values. +*/ + +template +__global__ void swiglu_oai_kernel(const KernelArgs a) +{ + const int64_t tilesPerRow = + (static_cast(a.dim) + blockDim.x - 1) / blockDim.x; + const int64_t totalTiles = static_cast(a.rows) * tilesPerRow; + const unsigned int mask = 0xffffffffu; + + for (int64_t tile = blockIdx.x; tile < totalTiles; tile += gridDim.x) { + const int64_t row = tile / tilesPerRow; + const int64_t d = + (tile % tilesPerRow) * static_cast(blockDim.x) + threadIdx.x; + const bool active = (d < static_cast(a.dim)); + + const float y = + active ? swiglu_elem(a, row, d) : 0.f; + + if constexpr (OUT == OUT_NONE) { + if (active) + static_cast(a.Y)[row * static_cast(a.dim) + d] = y; + } else if constexpr (OUT == OUT_FP8) { + if (active) + static_cast<__nv_fp8_e4m3*>(a.Y)[ + row * static_cast(a.dim) + d].__x = + __nv_cvt_float_to_fp8(y / a.yScale, __NV_SATFINITE, __NV_E4M3); + } else { // OUT_MXFP4 + // amax over the 32 lanes of this mx block; inactive lanes contribute 0. + float amax = active ? fabsf(y) : 0.f; + #pragma unroll + for (int m = 1; m < WarpSize; m <<= 1) + amax = fmaxf(amax, __shfl_xor_sync(mask, amax, m, WarpSize)); + + const int lane = threadIdx.x & (WarpSize - 1); + // Dividing by E2M1's largest power-of-two exponent before conversion + // chooses the block scale. E8M0 round-toward-zero is floor(log2) for + // positive values, exactly matching reference.h's scale rule. E8M0 has + // no zero, so retain byte 127 (scale 1) for an all-zero/NaN block. + const uint8_t sb = (amax > 0.f) + ? __nv_cvt_float_to_e8m0( + amax * 0.25f, __NV_SATFINITE, cudaRoundZero) + : 127; + const float s = e8m0_decode(sb); + + // Lane 0 of the block owns the scale byte. Its column is in range + // whenever the block exists, so an inactive lane 0 means nothing to do. + if (lane == 0 && active) + a.Ys[row * a.yBlocksPerRow + (d / MX_BLOCK)] = sb; + + // Two fp4 values share a byte, so the even lane of each pair packs its + // own value with its right-hand neighbour's (width 2 keeps adjacent + // pairs independent) and performs the single store. + const float paired = __shfl_down_sync(mask, y, 1, 2); + if (active && (d & 1) == 0) { + const float hi = + (d + 1 < static_cast(a.dim)) ? paired : 0.f; + static_cast<__nv_fp4x2_e2m1*>(a.Y)[row * a.yRowStrideBytes + (d >> 1)] = + fp4x2_pack(y / s, hi / s); + } + } + } +} + +// --------------------------------------------------------------------------- +// Launch dispatch. The switch runs on the host and costs nothing next to a +// kernel launch, so the timing loop can call it directly. +// --------------------------------------------------------------------------- + +template +static void launch_swiglu_out(OutQuant out, dim3 grids, dim3 blocks, + const KernelArgs& a) { + switch (out) { + case OUT_NONE: + swiglu_oai_kernel + <<>>(a); + break; + case OUT_FP8: + swiglu_oai_kernel + <<>>(a); + break; + case OUT_MXFP4: + swiglu_oai_kernel + <<>>(a); + break; + default: + break; + } +} + +template +static void launch_swiglu_bias(InDtype in, OutQuant oq, + dim3 grids, dim3 blocks, const KernelArgs& a) +{ + switch (in) { + case IN_FP32: + launch_swiglu_out(oq, grids, blocks, a); break; + case IN_FP16: + launch_swiglu_out(oq, grids, blocks, a); break; + case IN_BF16: + launch_swiglu_out(oq, grids, blocks, a); break; + case IN_FP8: + launch_swiglu_out(oq, grids, blocks, a); break; + case IN_MXFP8: + launch_swiglu_out(oq, grids, blocks, a); break; + case IN_FP4: + launch_swiglu_out(oq, grids, blocks, a); break; + case IN_MXFP4: + launch_swiglu_out(oq, grids, blocks, a); break; + default: break; + } +} + +static void launch_swiglu(InDtype in, OutQuant oq, BiasStore bias, + dim3 grids, dim3 blocks, const KernelArgs& a) +{ + switch (bias) { + case BS_NONE: launch_swiglu_bias(in, oq, grids, blocks, a); break; + case BS_F16: launch_swiglu_bias(in, oq, grids, blocks, a); break; + case BS_BF16: launch_swiglu_bias(in, oq, grids, blocks, a); break; + case BS_F32: launch_swiglu_bias(in, oq, grids, blocks, a); break; + default: break; + } +} + +// --------------------------------------------------------------------------- +// Host packing: turn the dequantized floats reference.h handed back into the +// native storage the kernel reads. Dividing by the same scale reference.h +// multiplied by lands on an exactly representable value, so the native +// conversion here is lossless and host and device see identical numbers. +// --------------------------------------------------------------------------- + +static void pack_input(InDtype dtype, const float* xq, size_t n, + size_t rowLen, const uint8_t* scales, float tscale, + void* dst) +{ + const XStore store = x_store_of(dtype); + const XScale sc = x_scale_of(dtype); + const size_t bpr = mx_blocks_per_row(rowLen); + + // Scale for element `i`, using the same [row][block] layout reference.h + // filled in. + auto scale_at = [&](size_t i) -> float { + if (sc == XSC_MX) { + const size_t r = i / rowLen; + const size_t c = i - r * rowLen; + return e8m0_decode(scales[r * bpr + c / MX_BLOCK]); + } + return tscale; + }; + + if (store == XS_FP4) { + // rowLen is 2*dim and therefore even, so a packed pair never straddles a + // row, and 2j/2j+1 never straddle an MX block either: one scale for both. + __nv_fp4x2_e2m1* p = static_cast<__nv_fp4x2_e2m1*>(dst); + for (size_t j = 0; j < n / 2; j++) { + const size_t i = 2 * j; + const float s = scale_at(i); + p[j] = fp4x2_pack(xq[i] / s, xq[i + 1] / s); + } + return; + } + + for (size_t i = 0; i < n; i++) { + switch (store) { + case XS_F32: + static_cast(dst)[i] = xq[i]; + break; + case XS_F16: + static_cast<__half*>(dst)[i] = __float2half(xq[i]); + break; + case XS_BF16: + static_cast<__nv_bfloat16*>(dst)[i] = __float2bfloat16(xq[i]); + break; + default: // XS_FP8 + static_cast<__nv_fp8_e4m3*>(dst)[i] = __nv_fp8_e4m3(xq[i] / scale_at(i)); + break; + } + } +} + +static int bias_elem_bytes(BiasStore bias) +{ + if (bias == BS_F16 || bias == BS_BF16) return 2; + if (bias == BS_F32) return 4; + return 0; +} + +// Pack the bias in its production storage type and replace the FP32 reference +// values with the exact values that the device will decode. +static void pack_bias(BiasStore bias, float* values, size_t n, void* dst) +{ + if (bias == BS_NONE) { + memset(values, 0, n * sizeof(float)); + return; + } + if (bias == BS_F16) { + __half* packed = static_cast<__half*>(dst); + for (size_t i = 0; i < n; i++) { + packed[i] = __float2half(values[i]); + values[i] = __half2float(packed[i]); + } + return; + } + if (bias == BS_BF16) { + __nv_bfloat16* packed = static_cast<__nv_bfloat16*>(dst); + for (size_t i = 0; i < n; i++) { + packed[i] = __float2bfloat16(values[i]); + values[i] = __bfloat162float(packed[i]); + } + return; + } + memcpy(dst, values, n * sizeof(float)); +} + +// Undo the device's output quantization so the result can be compared with +// reference.h's dequantized reference elementwise. +static void unpack_output(OutQuant q, int rows, int dim, size_t blocksPerRow, + const void* src, const uint8_t* yscales, + float youtscale, float* Y) +{ + const size_t n = static_cast(rows) * dim; + + if (q == OUT_NONE) { + memcpy(Y, src, n * sizeof(float)); + return; + } + + if (q == OUT_FP8) { + const __nv_fp8_e4m3* p = static_cast(src); + for (size_t i = 0; i < n; i++) Y[i] = float(p[i]) * youtscale; + return; + } + + const __nv_fp4x2_e2m1* p = static_cast(src); + const size_t rowBytes = mxfp4_row_bytes(dim); + for (size_t i = 0; i < static_cast(rows); i++) { + for (size_t b = 0; b < blocksPerRow; b++) { + const float s = e8m0_decode(yscales[i * blocksPerRow + b]); + const int d0 = static_cast(b * MX_BLOCK); + const int nk = (d0 + MX_BLOCK < dim) ? MX_BLOCK : (dim - d0); + const __nv_fp4x2_e2m1* blk = p + i * rowBytes + b * (MX_BLOCK / 2); + for (int k = 0; k < nk; k++) { + Y[i * dim + d0 + k] = fp4x2_unpack(blk[k / 2].__x, k & 1) * s; + } + } + } +} + +static void usage(const char* prog) +{ + printf("Usage: %s " + "[in_dtype] [out_quant] [bias_dtype]\n", prog); + printf(" in_dtype : fp32 fp16 bf16 fp8 mxfp8 fp4 mxfp4 (default fp32)\n"); + printf(" out_quant : none fp8 mxfp4 (default none)\n"); + printf(" bias_dtype: none fp16 bf16 fp32 (default none; fp32 is accuracy-only)\n"); +} + +int main(int argc, char* argv[]) +{ + if (argc < 4 || argc > 7) { + usage(argv[0]); + return 1; + } + + const int rows = atoi(argv[1]); + const int dim = atoi(argv[2]); + const int repeat = atoi(argv[3]); + + if (rows <= 0 || dim <= 0) { + fprintf(stderr, "Error: rows and dimension must be positive (got %d and %d)\n", + rows, dim); + return 1; + } + + // The average kernel time is divided by the repeat count. + if (repeat <= 0) { + fprintf(stderr, "Error: repeat count must be positive (got %d)\n", repeat); + return 1; + } + + // Assert the rounding rules directly. The end-to-end comparison cannot catch + // a tie-policy bug -- a tie sits exactly on a rounding boundary, where + // values_match() grants slack by design -- so this is what guards it. + if (!self_test_rounding()) return 1; + + const InDtype indt = (argc > 4) ? parse_in_dtype(argv[4]) : IN_FP32; + if (indt == IN_INVALID) { + fprintf(stderr, "Error: invalid input dtype '%s' " + "(valid: fp32 fp16 bf16 fp8 mxfp8 fp4 mxfp4)\n", argv[4]); + return 1; + } + + const OutQuant oq = (argc > 5) ? parse_out_quant(argv[5]) : OUT_NONE; + if (oq == OUT_INVALID) { + fprintf(stderr, "Error: invalid output quantization '%s' " + "(valid: none fp8 mxfp4)\n", argv[5]); + return 1; + } + + const BiasStore bias = + (argc > 6) ? parse_bias_store(argv[6]) : BS_NONE; + if (static_cast(bias) < static_cast(BS_NONE) || bias > BS_F32) { + fprintf(stderr, "Error: invalid bias dtype '%s' " + "(valid: none fp16 bf16 fp32)\n", argv[6]); + return 1; + } + + printf("Shape of input tensor: ( %d %d ) in=%s out=%s bias=%s\n", + rows, 2 * dim, in_dtype_name(indt), out_quant_name(oq), + bias_store_name(bias)); + + const size_t y_nelems = static_cast(rows) * dim; + const size_t x_nelems = y_nelems * 2; + const size_t b_nelems = static_cast(2) * dim; // bias + + const size_t x_bytes = x_nelems * sizeof(float); + const size_t y_bytes = y_nelems * sizeof(float); + const size_t b_bytes = b_nelems * sizeof(float); + const size_t b_store_bytes = b_nelems * bias_elem_bytes(bias); + + // Native input storage. Input MX scales are [row][block] over rows of + // 2*dim elements, with blocks cut inside a row. + const size_t x_rowlen = b_nelems; // 2*dim + const size_t x_bpr = mx_blocks_per_row(x_rowlen); + const size_t x_store_bytes = (in_elem_bits(indt) == 4) + ? (x_nelems + 1) / 2 + : x_nelems * (in_elem_bits(indt) / 8); + const size_t x_nscales = + in_is_mx(indt) ? static_cast(rows) * x_bpr : 0; + + // mxfp4 output rows are padded to a whole number of MX blocks so no byte is + // ever shared between two rows; mxfp4_row_bytes() is the layout every + // backend agrees on. + const size_t blocksPerRow = mx_blocks_per_row(dim); + const size_t y_row_bytes = mxfp4_row_bytes(dim); + const size_t y_nscales = + out_is_mx(oq) ? static_cast(rows) * blocksPerRow : 0; + const size_t y_store_bytes = (oq == OUT_NONE) ? y_bytes + : (oq == OUT_FP8) ? y_nelems + : static_cast(rows) * y_row_bytes; + + float *X = static_cast(malloc(x_bytes)); + float *B = static_cast(malloc(b_bytes)); + float *Xq = static_cast(malloc(x_bytes)); + float *Y = static_cast(malloc(y_bytes)); + float *Y_ref = static_cast(malloc(y_bytes)); + float *Y_pre = static_cast(malloc(y_bytes)); + + void *Xp = malloc(x_store_bytes); + void *Bp = (bias == BS_NONE) ? nullptr : malloc(b_store_bytes); + void *Yp = malloc(y_store_bytes); + + uint8_t *Xs = static_cast(malloc(x_nscales ? x_nscales : 1)); + uint8_t *Ys = static_cast(malloc(y_nscales ? y_nscales : 1)); + uint8_t *Ys_ref = + static_cast(malloc(y_nscales ? y_nscales : 1)); + + if (!X || !B || !Xq || !Y || !Y_ref || !Y_pre || + !Xp || (bias != BS_NONE && !Bp) || !Yp || + !Xs || !Ys || !Ys_ref) { + fprintf(stderr, "Error: host allocation failed\n"); + return EXIT_FAILURE; + } + + std::default_random_engine generator(123); + // both clamped and unclamped elements are exercised. + std::uniform_real_distribution distribution(-12.f, 12.f); + + for (size_t i = 0; i < x_nelems; i++) { + X[i] = distribution(generator); + } + for (size_t i = 0; i < b_nelems; i++) { + B[i] = distribution(generator); + } + pack_bias(bias, B, b_nelems, Bp); + + // Quantize FP32 values and then pack them + float xtscale = 1.f; + quantize_dequantize_input(indt, X, x_nelems, x_rowlen, Xq, Xs, &xtscale); + pack_input(indt, Xq, x_nelems, x_rowlen, Xs, xtscale, Xp); + + void *d_X, *d_Y, *d_B = nullptr; + uint8_t *d_Xs, *d_Ys; + CUDA_CHECK(cudaMalloc(reinterpret_cast(&d_X), x_store_bytes)); + CUDA_CHECK(cudaMemcpy(d_X, Xp, x_store_bytes, cudaMemcpyHostToDevice)); + if (bias != BS_NONE) { + CUDA_CHECK(cudaMalloc(reinterpret_cast(&d_B), b_store_bytes)); + CUDA_CHECK(cudaMemcpy(d_B, Bp, b_store_bytes, cudaMemcpyHostToDevice)); + } + CUDA_CHECK(cudaMalloc( + reinterpret_cast(&d_Xs), x_nscales ? x_nscales : 1)); + if (x_nscales) + CUDA_CHECK(cudaMemcpy(d_Xs, Xs, x_nscales, cudaMemcpyHostToDevice)); + + CUDA_CHECK(cudaMalloc(reinterpret_cast(&d_Y), y_store_bytes)); + CUDA_CHECK(cudaMalloc( + reinterpret_cast(&d_Ys), y_nscales ? y_nscales : 1)); + + ComputeSwigluOAI(rows, dim, kSwigluAlpha, kSwigluLimit, Xq, B, Y_ref); + + // The fp8 fused quant uses a single per-tensor scale. A kernel cannot see + // the whole tensor in its epilogue, so the scale is computed up front and + // handed to the kernel, which is how aiter's per-tensor quant works too. + // This is the same amax/448 reference.h derives internally. + float youtscale = 1.f; + if (oq == OUT_FP8) { + float amax = 0.f; + for (size_t i = 0; i < y_nelems; i++) { + const float a = fabsf(Y_ref[i]); + if (a > amax) amax = a; + } + youtscale = (amax > 0.f) ? (amax / kE4M3Max) : 1.f; + } + + // values_match() needs the reference as it was BEFORE the output quant: once + // quantized every value sits exactly on a representable code, which says + // nothing about how close the activation came to a rounding boundary. + memcpy(Y_pre, Y_ref, y_bytes); + quantize_dequantize_output(oq, rows, dim, Y_ref, Ys_ref); + + const int block_size = 256; + // Threads are tiled (row, block of block_size columns) so an MXFP4 group of + // 32 values always lands on one warp. The grid-stride loop then covers any + // excess, keeping grid.x inside its hardware limit for very large tensors. + const int64_t tilesPerRow = + (static_cast(dim) + block_size - 1) / block_size; + const int64_t totalTiles = static_cast(rows) * tilesPerRow; + const int64_t max_blocks = 1 << 20; + size_t nblocks = static_cast( + totalTiles < max_blocks ? totalTiles : max_blocks); + + // Kernel launch arguments + KernelArgs a; + a.rows = rows; + a.dim = dim; + a.alpha = kSwigluAlpha; + a.limit = kSwigluLimit; + a.X = d_X; + a.Xs = d_Xs; + a.xScale = xtscale; + a.B = d_B; + a.Y = d_Y; + a.Ys = d_Ys; + a.yScale = youtscale; + a.xBlocksPerRow = static_cast(x_bpr); + a.yRowStrideBytes = static_cast(y_row_bytes); + a.yBlocksPerRow = static_cast(blocksPerRow); + const dim3 grids(static_cast(nblocks)); + const dim3 blocks(block_size); + + // check correctness before benchmarking + launch_swiglu(indt, oq, bias, grids, blocks, a); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaMemcpy(Yp, d_Y, y_store_bytes, cudaMemcpyDeviceToHost)); + + if (y_nscales) + CUDA_CHECK(cudaMemcpy(Ys, d_Ys, y_nscales, cudaMemcpyDeviceToHost)); + unpack_output(oq, rows, dim, blocksPerRow, Yp, Ys, youtscale, Y); + + const float tol = 1e-3f; + bool ok = true; + float maxdiff = 0.f; + size_t worst = 0; + for (size_t i = 0; i < y_nelems; i++) { + float scale = youtscale; + if (oq == OUT_MXFP4) { + const size_t r = i / dim; + const size_t b = (i - r * dim) / MX_BLOCK; + scale = e8m0_to_float(Ys_ref[r * blocksPerRow + b]); + } + const float diff = fabsf(Y[i] - Y_ref[i]); + if (diff > maxdiff) { + maxdiff = diff; + worst = i; + } + if (!values_match(Y[i], Y_ref[i], Y_pre[i], oq, scale, tol)) + ok = false; + } + printf("%s\n", ok ? "PASS" : "FAIL"); + if (!ok) { + size_t scale_mismatch = 0; + for (size_t i = 0; i < y_nscales; ++i) + if (Ys[i] != Ys_ref[i]) ++scale_mismatch; + fprintf(stderr, " max |device - reference| = %g at %zu " + "(device %g, reference %g, tol %g)\n", + maxdiff, worst, Y[worst], Y_ref[worst], tol); + if (y_nscales) + fprintf(stderr, " mismatching e8m0 output scales: %zu of %zu\n", + scale_mismatch, y_nscales); + } + + if (bias == BS_F32) { + printf("FP32 bias is accuracy-only; performance timing skipped.\n"); + } else { + CUDA_CHECK(cudaDeviceSynchronize()); + auto start = std::chrono::steady_clock::now(); + + for (int i = 0; i < repeat; i++) { + launch_swiglu(indt, oq, bias, grids, blocks, a); + } + + CUDA_CHECK(cudaDeviceSynchronize()); + auto end = std::chrono::steady_clock::now(); + const auto time = std::chrono::duration_cast( + end - start).count(); + printf("Average execution time of SwiGLU (OAI) kernel: %f (us)\n", + (time * 1e-3f) / repeat); + } + + CUDA_CHECK(cudaFree(d_X)); + if (d_B != nullptr) CUDA_CHECK(cudaFree(d_B)); + CUDA_CHECK(cudaFree(d_Y)); + CUDA_CHECK(cudaFree(d_Xs)); + CUDA_CHECK(cudaFree(d_Ys)); + + free(X); + free(B); + free(Xq); + free(Y); + free(Y_ref); + free(Y_pre); + free(Xp); + free(Bp); + free(Yp); + free(Xs); + free(Ys); + free(Ys_ref); + + return EXIT_SUCCESS; +} diff --git a/src/swiglu-oai-cuda/reference.h b/src/swiglu-oai-cuda/reference.h new file mode 100644 index 000000000..36e153f41 --- /dev/null +++ b/src/swiglu-oai-cuda/reference.h @@ -0,0 +1,672 @@ +// The gate and linear halves are interleaved along the last dimension, as they +// are in gpt-oss: element 2*d is the gate, element 2*d+1 is the linear term. +// Both come out of a fused up-projection, so the bias is added here +// +// gate = min(x[2*d] + b[2*d], limit) +// linear = clamp(x[2*d+1] + b[2*d+1], -limit, limit) +// y[d] = gate * sigmoid(alpha * gate) * (linear + 1) +// +// Note the asymmetry: the gate is clamped from above only, the linear term +// from both sides. That is the reference behaviour, not an oversight. It also +// means `gate` is unbounded below, which is why the sigmoid below is the +// numerically stable form -- and why device kernels must use the stable form +// too rather than the naive 1/(1+exp(-x)). +// +// Two independent precision axes are modelled, mirroring aiter's fused MoE +// stage-1 kernel: +// +// input -- the gate/up tensor is stored as fp32/fp16/bf16/fp8/mxfp8/ +// fp4/mxfp4 and dequantized before the activation. +// output -- the activation result is quantized to fp8 or mxfp4 in the +// kernel epilogue (aiter calls this the fused quant). +// +// The "mx" formats are OCP microscaling: one e8m0 (power-of-two) scale per +// MX_BLOCK elements along the last dimension. The plain fp8/fp4 formats use a +// single per-tensor scale, matching aiter's QuantType::per_Tensor. + +#ifndef REFERENCE_H +#define REFERENCE_H + +#include +#include +#include +#include +#include + +constexpr float kSwigluAlpha = 1.702f; +constexpr float kSwigluLimit = 7.f; + +// OCP microscaling block length. +constexpr int MX_BLOCK = 32; + +enum InDtype { + IN_FP32 = 0, + IN_FP16, + IN_BF16, + IN_FP8, + IN_MXFP8, + IN_FP4, + IN_MXFP4, + IN_INVALID +}; + +enum OutQuant : int { OUT_NONE = 0, OUT_FP8, OUT_MXFP4, OUT_INVALID }; + +// Physical input storage and its independent scaling policy. Keeping these +// axes separate lets every backend share one load implementation between, +// for example, FP8 and MXFP8. +enum XStore { XS_F32 = 0, XS_F16, XS_BF16, XS_FP8, XS_FP4 }; +enum XScale { XSC_NONE = 0, XSC_TENSOR, XSC_MX }; + +enum BiasStore { BS_NONE = 0, BS_F16, BS_BF16, BS_F32 }; + +struct KernelArgs { + int rows; + int dim; + float alpha; + float limit; + const void* X; // packed gate/up tensor + const uint8_t* Xs; // e8m0 scales [row][block], mx inputs only + float xScale; // per-tensor scale, plain fp8/fp4 only + const void* B; // optional packed bias + void* Y; // packed activation + uint8_t* Ys; // e8m0 scales [row][block], mxfp4 output only + float yScale; // per-tensor scale, fp8 output only + int64_t xBlocksPerRow; + int64_t yRowStrideBytes; + int64_t yBlocksPerRow; +}; + +inline InDtype parse_in_dtype(const char* s) { + if (!strcmp(s, "fp32")) return IN_FP32; + if (!strcmp(s, "fp16")) return IN_FP16; + if (!strcmp(s, "bf16")) return IN_BF16; + if (!strcmp(s, "fp8")) return IN_FP8; + if (!strcmp(s, "mxfp8")) return IN_MXFP8; + if (!strcmp(s, "fp4")) return IN_FP4; + if (!strcmp(s, "mxfp4")) return IN_MXFP4; + return IN_INVALID; +} + +inline OutQuant parse_out_quant(const char* s) { + if (!strcmp(s, "none")) return OUT_NONE; + if (!strcmp(s, "fp8")) return OUT_FP8; + if (!strcmp(s, "mxfp4")) return OUT_MXFP4; + return OUT_INVALID; +} + +inline BiasStore parse_bias_store(const char* s) { + if (!strcmp(s, "none")) return BS_NONE; + if (!strcmp(s, "fp16")) return BS_F16; + if (!strcmp(s, "bf16")) return BS_BF16; + if (!strcmp(s, "fp32")) return BS_F32; + return (BiasStore)-1; +} + +inline const char* bias_store_name(BiasStore b) { + switch (b) { + case BS_NONE: return "none"; + case BS_F16: return "fp16"; + case BS_BF16: return "bf16"; + case BS_F32: return "fp32-accuracy"; + default: return "invalid"; + } +} + +inline XStore x_store_of(InDtype t) { + switch (t) { + case IN_FP32: return XS_F32; + case IN_FP16: return XS_F16; + case IN_BF16: return XS_BF16; + case IN_FP8: + case IN_MXFP8: return XS_FP8; + default: return XS_FP4; + } +} + +inline XScale x_scale_of(InDtype t) { + if (t == IN_MXFP8 || t == IN_MXFP4) return XSC_MX; + if (t == IN_FP8 || t == IN_FP4) return XSC_TENSOR; + return XSC_NONE; +} + +inline const char* in_dtype_name(InDtype t) { + static const char* n[] = {"fp32", "fp16", "bf16", "fp8", + "mxfp8", "fp4", "mxfp4", "invalid"}; + return n[t]; +} + +// MXFP8 ? +inline const char* out_quant_name(OutQuant q) { + static const char* n[] = {"none", "fp8", "mxfp4", "invalid"}; + return n[q]; +} + +// True when the format carries one e8m0 scale per MX_BLOCK elements. +inline bool in_is_mx(InDtype t) { return t == IN_MXFP8 || t == IN_MXFP4; } +inline bool out_is_mx(OutQuant q) { return q == OUT_MXFP4; } + +// True when the format carries a single per-tensor scale. +inline bool in_is_scaled(InDtype t) { return t == IN_FP8 || t == IN_FP4; } + +// Bits of storage per element. fp4 returns 4; callers pack two elements per +// byte. Only IN_INVALID returns 0. +inline int in_elem_bits(InDtype t) { + switch (t) { + case IN_FP32: return 32; + case IN_FP16: + case IN_BF16: return 16; + case IN_FP8: + case IN_MXFP8: return 8; + case IN_FP4: + case IN_MXFP4: return 4; + default: return 0; + } +} + +// --------------------------------------------------------------------------- +// Scalar format conversions, host side. These define ground truth; device +// kernels use the vendor's native types and must agree with these. +// --------------------------------------------------------------------------- + +// Largest finite magnitude of each element format, used to derive MX scales. +constexpr float kE4M3Max = 448.f; // OCP e4m3 +constexpr float kE2M1Max = 6.f; // e2m1 +// Exponent of the largest power of two <= the format max. +constexpr int kE4M3MaxExp = 8; // 2^8 = 256 <= 448 +constexpr int kE2M1MaxExp = 2; // 2^2 = 4 <= 6 + +inline float round_to_fp16(float x) { + // Round-to-nearest-even through the 16-bit binary16 encoding. + uint32_t u; + memcpy(&u, &x, 4); + uint32_t sign = (u >> 16) & 0x8000u; + int32_t exp = (int32_t)((u >> 23) & 0xff) - 127 + 15; + uint32_t man = u & 0x7fffffu; + + // Infinities and NaNs are representable in binary16, so pass them through + // unchanged. (Decoding them with the normal-number formula below would be + // wrong: exponent 0x1f is reserved, and doing so turned +inf into 65536.) + if (((u >> 23) & 0xff) == 0xff) return x; + + uint32_t h; + if (exp >= 0x1f) { + h = sign | 0x7c00u; // overflow to inf + } else if (exp <= 0) { + if (exp < -10) { + h = sign; // underflow to zero + } else { + man |= 0x800000u; + uint32_t shift = (uint32_t)(14 - exp); + uint32_t sub = man >> shift; + uint32_t rem = man & ((1u << shift) - 1u); + uint32_t halfv = 1u << (shift - 1); + if (rem > halfv || (rem == halfv && (sub & 1u))) sub++; + h = sign | sub; + } + } else { + uint32_t sub = man >> 13; + uint32_t rem = man & 0x1fffu; + if (rem > 0x1000u || (rem == 0x1000u && (sub & 1u))) { + sub++; + if (sub == 0x400u) { + sub = 0; + exp++; + if (exp >= 0x1f) return sign ? -INFINITY : INFINITY; + } + } + h = sign | ((uint32_t)exp << 10) | sub; + } + + // Decode back to float. + uint32_t s = (h & 0x8000u) << 16; + uint32_t e = (h >> 10) & 0x1fu; + uint32_t m = h & 0x3ffu; + float r; + if (e == 0) { + if (m == 0) { + memcpy(&r, &s, 4); + return r; + } + float v = (float)m * 5.9604644775390625e-8f; // 2^-24 + r = s ? -v : v; + return r; + } + if (e == 0x1f) return s ? -INFINITY : INFINITY; + uint32_t o = s | ((e + 112u) << 23) | (m << 13); + memcpy(&r, &o, 4); + return r; +} + +inline float round_to_bf16(float x) { + uint32_t u; + memcpy(&u, &x, 4); + if (((u >> 23) & 0xff) == 0xff) return x; // inf/nan pass through + uint32_t lsb = (u >> 16) & 1u; + uint32_t rounded = u + 0x7fffu + lsb; // round to nearest even + rounded &= 0xffff0000u; + float r; + memcpy(&r, &rounded, 4); + return r; +} + +// Round a float to the nearest representable OCP e4m3 value (saturating). +inline float round_to_e4m3(float x) { + if (!(x == x)) return x; // NaN + float a = fabsf(x); + float sign = (x < 0.f) ? -1.f : 1.f; + if (a >= kE4M3Max) return sign * kE4M3Max; + if (a == 0.f) return x; + + int e; + frexpf(a, &e); // a in [0.5,1) * 2^e => unbiased exponent is e-1 + e -= 1; + const int minNormExp = -6; // e4m3 bias 7 => smallest normal 2^-6 + if (e < minNormExp) e = minNormExp; + // 3 mantissa bits: quantum is 2^(e-3). Ties go to even, matching OCP and the + // hardware convert instructions -- roundf() would break them away from zero + // and disagree with any backend that rounds in hardware. + float q = ldexpf(1.f, e - 3); + const float n = a / q; + const float fl = floorf(n); + const float frac = n - fl; + float rn; + if (frac > 0.5f) rn = fl + 1.f; + else if (frac < 0.5f) rn = fl; + else rn = (fmodf(fl, 2.f) == 0.f) ? fl : fl + 1.f; // tie -> even + float r = rn * q; + if (r > kE4M3Max) r = kE4M3Max; + return sign * r; +} + +// Round a float to the nearest representable e2m1 value (saturating). +// Representable magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. +inline float round_to_e2m1(float x) { + static const float lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + float a = fabsf(x); + float sign = (x < 0.f) ? -1.f : 1.f; + // Strictly above the 4/6 midpoint saturates. Exactly 5.0 is a tie and must + // fall through to the search below, which sends it to 4 (even mantissa) as + // round-to-nearest-even requires. + if (a > 5.f) return sign * 6.f; + int best = 0; + float bestd = fabsf(a - lut[0]); + for (int i = 1; i < 8; i++) { + float d = fabsf(a - lut[i]); + // Ties round to even index, matching round-to-nearest-even on the + // mantissa bit. + if (d < bestd || (d == bestd && (i % 2) == 0)) { + bestd = d; + best = i; + } + } + return sign * lut[best]; +} + +// e8m0: an unsigned power-of-two scale, biased by 127. Returns the biased +// exponent byte for a block whose largest magnitude is amax, given the +// exponent of the element format's maximum. +inline uint8_t e8m0_scale_byte(float amax, int elemMaxExp) { + if (!(amax > 0.f)) return 127; // 2^0 for an all-zero block + int e; + frexpf(amax, &e); + e -= 1; // unbiased exponent of amax + int s = e - elemMaxExp + 127; + if (s < 0) s = 0; + if (s > 254) s = 254; + return (uint8_t)s; +} + +inline float e8m0_to_float(uint8_t b) { return ldexpf(1.f, (int)b - 127); } + +// --------------------------------------------------------------------------- +// Input preparation. +// +// Produces the exact float values the device will see after it dequantizes, +// so host and device operate on identical inputs and any mismatch is a real +// kernel bug rather than a rounding difference in the data itself. +// +// `scales` receives one e8m0 byte per MX_BLOCK elements for the mx formats, +// and is left untouched otherwise. `tensor_scale` receives the per-tensor +// scale for the plain fp8/fp4 formats and 1.0f otherwise. +// --------------------------------------------------------------------------- + +// Number of MX blocks in one row of `rowLen` elements. Blocks are cut inside a +// row and never straddle two rows, so the last block of a row may be short. +// The input and the output paths use the same rule. +inline int64_t mx_blocks_per_row(int64_t rowLen) { + return (rowLen + MX_BLOCK - 1) / MX_BLOCK; +} + +// Row stride, in bytes, of a packed mxfp4 buffer whose logical row is `dim` +// elements. Rows are padded out to a whole number of MX blocks so a byte is +// never shared between two rows -- that keeps nibble writes race-free and, +// more importantly, makes every backend agree on the layout. +inline int64_t mxfp4_row_bytes(int64_t dim) { + return mx_blocks_per_row(dim) * (MX_BLOCK / 2); +} + +// `rowLen` is the logical row length of `in` (2*dim for the gate/up tensor). +// MX scales are laid out as [row][block]. +inline void quantize_dequantize_input( + InDtype dtype, + const float* in, + int64_t nelems, + int64_t rowLen, + float* out, + uint8_t* scales, + float* tensor_scale) +{ + *tensor_scale = 1.f; + + switch (dtype) { + case IN_FP32: + for (int64_t i = 0; i < nelems; i++) out[i] = in[i]; + return; + case IN_FP16: + for (int64_t i = 0; i < nelems; i++) out[i] = round_to_fp16(in[i]); + return; + case IN_BF16: + for (int64_t i = 0; i < nelems; i++) out[i] = round_to_bf16(in[i]); + return; + case IN_FP8: + case IN_FP4: { + const float emax = (dtype == IN_FP8) ? kE4M3Max : kE2M1Max; + float amax = 0.f; + for (int64_t i = 0; i < nelems; i++) { + float a = fabsf(in[i]); + if (a > amax) amax = a; + } + const float s = (amax > 0.f) ? (amax / emax) : 1.f; + *tensor_scale = s; + for (int64_t i = 0; i < nelems; i++) { + float q = (dtype == IN_FP8) ? round_to_e4m3(in[i] / s) + : round_to_e2m1(in[i] / s); + out[i] = q * s; + } + return; + } + case IN_MXFP8: + case IN_MXFP4: { + const int maxExp = (dtype == IN_MXFP8) ? kE4M3MaxExp : kE2M1MaxExp; + const int64_t bpr = mx_blocks_per_row(rowLen); + const int64_t nrows = (rowLen > 0) ? (nelems / rowLen) : 0; + for (int64_t r = 0; r < nrows; r++) { + const int64_t rowBeg = r * rowLen; + for (int64_t b = 0; b < bpr; b++) { + const int64_t beg = rowBeg + b * MX_BLOCK; + const int64_t rowEnd = rowBeg + rowLen; + const int64_t end = (beg + MX_BLOCK < rowEnd) ? beg + MX_BLOCK : rowEnd; + float amax = 0.f; + for (int64_t i = beg; i < end; i++) { + float a = fabsf(in[i]); + if (a > amax) amax = a; + } + const uint8_t sb = e8m0_scale_byte(amax, maxExp); + scales[r * bpr + b] = sb; + const float s = e8m0_to_float(sb); + for (int64_t i = beg; i < end; i++) { + float q = (dtype == IN_MXFP8) ? round_to_e4m3(in[i] / s) + : round_to_e2m1(in[i] / s); + out[i] = q * s; + } + } + } + return; + } + default: + return; + } +} + +// --------------------------------------------------------------------------- +// The activation itself. +// --------------------------------------------------------------------------- + +// Numerically stable sigmoid: expf(-x) overflows for large negative x, and the +// gate is clamped only from above, so x here really can be very negative. +inline float sigmoid(const float x) { + if (x >= 0) { + return 1.f / (1.f + expf(-x)); + } else { + const float exp_x = expf(x); + return exp_x / (1.f + exp_x); + } +} + +// Xdata and Bdata are already-dequantized float values. +inline void ComputeSwigluOAI( + const int rows, + const int dim, + const float alpha, + const float limit, + const float* Xdata, + const float* Bdata, + float* Ydata) +{ + for (int i = 0; i < rows; ++i) { + const int64_t xOffset = (int64_t)i * 2 * dim; + const int64_t yOffset = (int64_t)i * dim; + for (int d = 0; d < dim; ++d) { + float gate = Xdata[xOffset + 2 * d] + Bdata[2 * d]; + float linear = Xdata[xOffset + 2 * d + 1] + Bdata[2 * d + 1]; + + gate = fminf(gate, limit); + linear = fminf(fmaxf(linear, -limit), limit); + + Ydata[yOffset + d] = gate * sigmoid(alpha * gate) * (linear + 1.f); + } + } +} + +// --------------------------------------------------------------------------- +// Output quantization (aiter's fused stage-1 epilogue). +// +// Quantizes in place, leaving dequantized float values behind so the result +// can be compared against the device output elementwise. mxfp4 blocks run +// along the last dimension and never straddle a row. +// --------------------------------------------------------------------------- + +inline void quantize_dequantize_output( + OutQuant q, + int rows, + int dim, + float* Ydata, + uint8_t* scales) +{ + if (q == OUT_NONE) return; + + if (q == OUT_FP8) { + // Per-tensor scale, as aiter does for a non-microscaled fused quant. + const int64_t n = (int64_t)rows * dim; + float amax = 0.f; + for (int64_t i = 0; i < n; i++) { + float a = fabsf(Ydata[i]); + if (a > amax) amax = a; + } + const float s = (amax > 0.f) ? (amax / kE4M3Max) : 1.f; + for (int64_t i = 0; i < n; i++) Ydata[i] = round_to_e4m3(Ydata[i] / s) * s; + return; + } + + // OUT_MXFP4 + const int64_t blocksPerRow = mx_blocks_per_row(dim); + for (int i = 0; i < rows; i++) { + for (int64_t b = 0; b < blocksPerRow; b++) { + const int64_t beg = (int64_t)i * dim + b * MX_BLOCK; + const int64_t rowEnd = (int64_t)i * dim + dim; + const int64_t end = (beg + MX_BLOCK < rowEnd) ? beg + MX_BLOCK : rowEnd; + float amax = 0.f; + for (int64_t k = beg; k < end; k++) { + float a = fabsf(Ydata[k]); + if (a > amax) amax = a; + } + const uint8_t sb = e8m0_scale_byte(amax, kE2M1MaxExp); + scales[(int64_t)i * blocksPerRow + b] = sb; + const float s = e8m0_to_float(sb); + for (int64_t k = beg; k < end; k++) + Ydata[k] = round_to_e2m1(Ydata[k] / s) * s; + } + } +} + +// Additional slack for a QUANTIZED output element. +// +// Host and device quantize with the same rules, so they normally agree +// exactly. But an element sitting within a ULP of a rounding boundary can be +// pushed across it by a 1-ULP expf difference, and that costs a whole +// quantum -- which at these magnitudes is around 4.0, not 0.005. Comparing +// quantized output against a small absolute tolerance is therefore really an +// exact-match test, and it fails by a huge margin when it fails at all. +// +// Allowing exactly one quantum keeps the check meaningful (a genuinely wrong +// kernel is off by far more than one step, as the injected-fault runs +// confirm) without being one unlucky dataset away from a spurious failure. +// `scale` is the e8m0 block scale for mxfp4, or the per-tensor scale for fp8. +// +// The slack is GATED on the reference value actually sitting near a rounding +// boundary. Granting it unconditionally made the fp8 path blind to uniform +// relative errors below one quantum -- and since e4m3 has 3 mantissa bits, a +// quantum is 6.25%-12.5% of the value, so a 2% or 5% error in the activation +// went undetected (measured on the HIP port; it only began failing at ~7%). +// +// Gating fixes that because the two cases live five orders of magnitude apart: +// a 1-ULP expf difference lands ~1.6e-6 of a quantum from a boundary, while a +// 2% error lands ~0.16 of a quantum away. A threshold of 1e-3 quanta forgives +// the former and catches the latter. +// +// This is why values_match() needs the PRE-quantization reference: once +// quantized, every reference value sits exactly on a representable code and is +// therefore always half a quantum from a boundary, which carries no +// information about whether a flip was plausible. +// +// NOTE -- the naive sigmoid is benign HERE, and no input range can catch it. +// The stable form above is good practice and remains what this contract asks +// for. But for this particular expression the naive 1/(1+expf(-x)) is not +// actually wrong: when -alpha*gate overflows, expf returns +inf and +// 1/(1+inf) is exactly 0 -- a finite zero, never NaN -- and by that point the +// true sigmoid is already denormal-small. Measured divergence in y: exactly 0 +// at |gate| = 52, 5.0e-43 at |gate| = 60, and exactly 0 again beyond |gate| = +// 80 where both forms give -0. So widening the data range does NOT expose the +// difference; an earlier version of this comment claimed |gate| > 52 would, +// and that was wrong. The naive form can only be found by reading the code, +// not by any output check this benchmark can perform. +inline float allowed_error(OutQuant q, float ref, float scale) { + if (q == OUT_NONE) return 0.f; + + if (q == OUT_MXFP4) { + // e2m1 steps are not uniform: 0.5 below 2, then 1.0, then 2.0 up to the + // 6 ceiling. Use the LOCAL step -- returning the smallest one would reject + // a legitimate flip at the 4/6 boundary, which is the very case this slack + // exists to absorb. + const float a = fabsf(ref / scale); + if (a < 2.f) return 0.5f * scale; + if (a < 4.f) return 1.f * scale; + return 2.f * scale; + } + + // e4m3: the step is 2^(e-3) at the exponent of the scaled magnitude. + const float a = fabsf(ref / scale); + int e = -6; + if (a > 0.f) { + frexpf(a, &e); + e -= 1; + if (e < -6) e = -6; + } + return ldexpf(1.f, e - 3) * scale; +} + +// Direct conformance check on the rounding rules. +// +// The end-to-end output comparison CANNOT catch a tie-policy bug, and no +// choice of threshold would let it: a tie sits at distance 0 from a rounding +// boundary, which is the deepest part of the region where values_match() +// grants slack, so "tie-rule disagreement" and "benign ULP noise at a +// boundary" are the same geometry. Exact ties are also vanishingly rare in +// random data -- measured, about 1 e4m3 tie per 16.7M elements -- so even an +// exact-match comparison would only catch one by luck. +// +// Backends therefore assert the rules directly. This is what guards against a +// port that hands ties to a hardware convert with different semantics. +// Returns true if every rule holds, and prints the first failure otherwise. +inline bool self_test_rounding() { + struct Case { const char* what; float got; float want; }; + const Case cases[] = { + // e4m3 ties go to even, not away from zero. + {"round_to_e4m3(1.0625)", round_to_e4m3(1.0625f), 1.0f}, + {"round_to_e4m3(1.1875)", round_to_e4m3(1.1875f), 1.25f}, + {"round_to_e4m3(1.09)", round_to_e4m3(1.09f), 1.125f}, + {"round_to_e4m3(1e30)", round_to_e4m3(1e30f), kE4M3Max}, + // e2m1 ties go to even. 5.0 is the 4-vs-6 midpoint and must give 4. + {"round_to_e2m1(5.0)", round_to_e2m1(5.0f), 4.0f}, + {"round_to_e2m1(3.5)", round_to_e2m1(3.5f), 4.0f}, + {"round_to_e2m1(2.5)", round_to_e2m1(2.5f), 2.0f}, + {"round_to_e2m1(0.75)", round_to_e2m1(0.75f), 1.0f}, + {"round_to_e2m1(0.25)", round_to_e2m1(0.25f), 0.0f}, + {"round_to_e2m1(100)", round_to_e2m1(100.f), 6.0f}, + {"round_to_e2m1(-5.0)", round_to_e2m1(-5.0f), -4.0f}, + }; + for (unsigned i = 0; i < sizeof(cases) / sizeof(cases[0]); i++) { + if (cases[i].got != cases[i].want) { + fprintf(stderr, "reference.h self-test FAILED: %s = %g, expected %g\n", + cases[i].what, cases[i].got, cases[i].want); + return false; + } + } + // fp16 must preserve inf and NaN rather than decoding them as normals. + const float pinf = round_to_fp16(INFINITY); + const float qnan = round_to_fp16(NAN); + if (!(pinf > 0.f) || !(pinf == pinf * 2.f) || qnan == qnan) { + fprintf(stderr, "reference.h self-test FAILED: round_to_fp16 inf/NaN " + "(+inf -> %g, NaN -> %g)\n", pinf, qnan); + return false; + } + return true; +} + +// Was the pre-quantization reference value close enough to a rounding boundary +// that a ULP-level difference in the activation could plausibly have pushed the +// device to the neighbouring code? Distances are measured in units of the local +// quantum, so the threshold is dimensionless. +inline bool near_rounding_boundary(OutQuant q, float pre, float scale) { + if (q == OUT_NONE) return false; + const float kBoundaryQuanta = 1e-3f; + const float a = fabsf(pre / scale); + + if (q == OUT_MXFP4) { + // Midpoints between the e2m1 magnitudes {0,.5,1,1.5,2,3,4,6}. The quantum + // is not uniform, so each boundary carries its own local step. + static const float bnd[7] = {0.25f, 0.75f, 1.25f, 1.75f, 2.5f, 3.5f, 5.f}; + static const float step[7] = {0.5f, 0.5f, 0.5f, 0.5f, 1.f, 1.f, 2.f}; + for (int i = 0; i < 7; i++) + if (fabsf(a - bnd[i]) <= kBoundaryQuanta * step[i]) return true; + return false; + } + + // e4m3: boundaries are the odd half-multiples of the local quantum. + int e = -6; + if (a > 0.f) { + frexpf(a, &e); + e -= 1; + if (e < -6) e = -6; + } + const float qs = ldexpf(1.f, e - 3); + const float n = a / qs; + return fabsf(n - floorf(n) - 0.5f) <= kBoundaryQuanta; +} + +// Does a device value match the reference for this element? +// dev -- device result, dequantized +// ref -- reference result after quantize_dequantize_output +// ref_prequant -- reference result BEFORE that call (keep a copy) +// scale -- e8m0 block scale for mxfp4, per-tensor scale for fp8 +inline bool values_match(float dev, float ref, float ref_prequant, + OutQuant q, float scale, float tol) { + const float diff = fabsf(dev - ref); + if (diff <= tol) return true; + if (q == OUT_NONE) return false; + return near_rounding_boundary(q, ref_prequant, scale) && + diff <= tol + allowed_error(q, ref, scale); +} + +#endif diff --git a/src/swiglu-oai-hip/CMakeLists.txt b/src/swiglu-oai-hip/CMakeLists.txt new file mode 100644 index 000000000..5c392c144 --- /dev/null +++ b/src/swiglu-oai-hip/CMakeLists.txt @@ -0,0 +1,8 @@ +# swiglu-oai-hip/CMakeLists.txt + +add_hecbench_benchmark( + NAME swiglu-oai + MODEL hip + SOURCES main.cu + CATEGORIES algorithms +) diff --git a/src/swiglu-oai-hip/Makefile b/src/swiglu-oai-hip/Makefile new file mode 100644 index 000000000..a81969483 --- /dev/null +++ b/src/swiglu-oai-hip/Makefile @@ -0,0 +1,88 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = hipcc +OPTIMIZE = yes +DEBUG = no +LAUNCHER ?= + +ROWS ?= 8192 +DIM ?= 1024 +REPEAT ?= 100 + +# E8M0 conversion API selection: +# auto - use HIP version detection in main.cu (default) +# yes - force __hip_cvt_float_to_e8m0 (for vendor backports) +# no - force the portable exponent-bit fallback +HIP_E8M0_CONVERSION ?= auto + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cu + +obj = $(source:.cu=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall + +ifeq ($(HIP_E8M0_CONVERSION),yes) + CFLAGS += -DSWIGLU_HAS_HIP_E8M0_CONVERSION=1 +else ifeq ($(HIP_E8M0_CONVERSION),no) + CFLAGS += -DSWIGLU_HAS_HIP_E8M0_CONVERSION=0 +else ifneq ($(HIP_E8M0_CONVERSION),auto) + $(error HIP_E8M0_CONVERSION must be auto, yes, or no) +endif + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g -DDEBUG + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cu reference.h ../swiglu-oai-cuda/reference.h Makefile + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +.PHONY: clean run + +run: $(program) + @run_case() { \ + output="$$($(LAUNCHER) ./$(program) "$$@" 2>&1)"; status=$$?; \ + printf '%s\n' "$$output"; \ + if [ $$status -ne 0 ] || \ + ! printf '%s\n' "$$output" | awk '$$0 == "PASS" { pass=1 } END { exit !pass }'; then \ + return 1; \ + fi; \ + }; \ + set -e; \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 none none; \ + run_case $(ROWS) $(DIM) $(REPEAT) fp16 none fp16; \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 fp8 none; \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 mxfp4 none; \ + run_case $(ROWS) $(DIM) $(REPEAT) mxfp8 fp8 bf16 diff --git a/src/swiglu-oai-hip/main.cu b/src/swiglu-oai-hip/main.cu new file mode 100644 index 000000000..99f9b7508 --- /dev/null +++ b/src/swiglu-oai-hip/main.cu @@ -0,0 +1,818 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "reference.h" + +#ifndef SWIGLU_HAS_HIP_E8M0_CONVERSION +// __hip_cvt_float_to_e8m0 first shipped in the TheRock 7.11 line. Allow an +// explicit build definition to handle vendor backports or unusual versioning. +#if (HIP_VERSION_MAJOR > 7) || \ + (HIP_VERSION_MAJOR == 7 && HIP_VERSION_MINOR >= 11) +#define SWIGLU_HAS_HIP_E8M0_CONVERSION 1 +#else +#define SWIGLU_HAS_HIP_E8M0_CONVERSION 0 +#endif +#endif + +#if !defined(__HIPCC__) +#error "swiglu-oai-hip must be compiled with hipcc" +#elif HIP_VERSION_MAJOR < 7 +#error "swiglu-oai-hip requires ROCm 7.0 or newer for hip_fp4.h" +#else +#include +#endif + + +#define HIP_CHECK(call) \ + do { \ + const hipError_t error_ = (call); \ + if (error_ != hipSuccess) { \ + fprintf(stderr, "HIP error at %s:%d: %s failed: %s\n", \ + __FILE__, __LINE__, #call, hipGetErrorString(error_)); \ + return EXIT_FAILURE; \ + } \ + } while (0) + +static_assert(sizeof(size_t) >= 8, + "This benchmark requires 64-bit size_t indexing"); + +// --------------------------------------------------------------------------- +// Architecture-matched E4M3 storage. +// +// CDNA3 gfx940/941/942 uses its native FNUZ format (bias 8, max 240), and the +// HIP-local reference extends the shared reference with matching FNUZ rules. +// gfx950, gfx1200/1201, and gfx1250 use native OCP E4M3 (bias 7, max 448). +// CDNA1/2 have no native FP8 conversion and use the software OCP codec below. +// Raw byte storage keeps host packing independent from device-only +// architecture macros. +// --------------------------------------------------------------------------- + +__host__ __device__ __forceinline__ float bits_to_float(uint32_t u) { + return __builtin_bit_cast(float, u); +} + +__host__ __device__ __forceinline__ uint32_t float_to_bits(float f) { + return __builtin_bit_cast(uint32_t, f); +} + +__host__ __device__ __forceinline__ float ocp_e4m3_decode(uint8_t x) { +#if defined(__HIP_DEVICE_COMPILE__) && \ + (defined(__gfx950__) || defined(__gfx1200__) || \ + defined(__gfx1201__) || defined(__gfx1250__)) + __hip_fp8_e4m3 v; + v.__x = x; + return static_cast(v); +#else + const uint32_t sign = static_cast(x & 0x80u) << 24; + const uint32_t mag = x & 0x7fu; + if (mag == 0) return bits_to_float(sign); + if (mag < 8) { + const float value = static_cast(mag) * 0.001953125f; // 2^-9 + return sign ? -value : value; + } + if (mag == 0x7f) return NAN; + const uint32_t exponent = mag >> 3; + const uint32_t mantissa = mag & 7u; + return bits_to_float(sign | ((exponent + 120u) << 23) | + (mantissa << 20)); +#endif +} + +__host__ __device__ __forceinline__ uint8_t ocp_e4m3_encode(float x) { +#if defined(__HIP_DEVICE_COMPILE__) && \ + (defined(__gfx950__) || defined(__gfx1200__) || \ + defined(__gfx1201__) || defined(__gfx1250__)) + return __hip_fp8_e4m3(x).__x; +#else + const uint32_t u = float_to_bits(x); + const uint8_t sign = static_cast((u >> 24) & 0x80u); + const uint32_t fexp = (u >> 23) & 0xffu; + const float a = fabsf(x); + if (fexp == 0xffu) + return (u & 0x7fffffu) ? static_cast(sign | 0x7fu) + : static_cast(sign | 0x7eu); + if (a >= 448.f) return static_cast(sign | 0x7eu); + if (a < 0.015625f) { // E4M3 subnormals, quantum 2^-9 + const float scaled = a * 512.f; + uint32_t code = static_cast(scaled); + const float fraction = scaled - static_cast(code); + if (fraction > 0.5f || (fraction == 0.5f && (code & 1u))) ++code; + return static_cast(sign | code); + } + + int exponent = static_cast(fexp) - 127; + const uint32_t significand = 0x800000u | (u & 0x7fffffu); + uint32_t rounded = significand >> 20; + const uint32_t remainder = significand & 0xfffffu; + if (remainder > 0x80000u || + (remainder == 0x80000u && (rounded & 1u))) { + ++rounded; + if (rounded == 16u) { + rounded = 8u; + ++exponent; + } + } + uint32_t code = static_cast(exponent + 7) * 8u + (rounded - 8u); + if (code > 0x7eu) code = 0x7eu; + return static_cast(sign | code); +#endif +} + +__device__ __forceinline__ float device_e4m3_decode(uint8_t x) { +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + __hip_fp8_e4m3_fnuz value; + value.__x = x; + return static_cast(value); +#else + return ocp_e4m3_decode(x); +#endif +} + +__device__ __forceinline__ uint8_t device_e4m3_encode(float x) { +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + return __hip_fp8_e4m3_fnuz(x).__x; +#else + return ocp_e4m3_encode(x); +#endif +} + +static_assert(sizeof(__hip_fp4x2_e2m1) == 1, "two fp4 values must pack into one byte"); + +// --------------------------------------------------------------------------- +// Device-side helpers +// --------------------------------------------------------------------------- + +// expf(-alpha*gate) overflows for large negative gate. +__device__ __forceinline__ float dev_sigmoid(float x) { + if (x >= 0.f) { + return 1.f / (1.f + expf(-x)); + } else { + const float ex = expf(x); + return ex / (1.f + ex); + } +} + +// FP32 reserves exponent 0 for zero/subnormals while E8M0 treats exponent 0 as 2^{-127} +// TODO: b = 0xFF +__device__ __forceinline__ float e8m0_decode(uint8_t b) { + return bits_to_float( + b ? static_cast(b) << 23 : 0x00400000u); +} + +// Input load: fetch the native element and undo the quantization scale. + +template +__device__ __forceinline__ float load_x_raw( + const void* __restrict__ X, const int64_t idx) { + if constexpr (STORE == XS_F32) { + return static_cast(X)[idx]; + } else if constexpr (STORE == XS_F16) { + return __half2float(static_cast(X)[idx]); + } else if constexpr (STORE == XS_BF16) { + return __bfloat162float(static_cast(X)[idx]); + } else if constexpr (STORE == XS_FP8) { + return device_e4m3_decode(static_cast(X)[idx]); + } else { + __hip_fp4_e2m1 v; + v.__x = static_cast<__hip_fp4_storage_t>( + (static_cast(X)[idx >> 1] >> + ((idx & 1) << 2)) & 0xf); + return static_cast(v); + } +} + +// Load one adjacent gate/linear pair. FP4 stores the pair in one byte and FP8 +// stores it in one aligned 16-bit word, so one vector conversion decodes both +// values. Other formats retain their ordinary scalar loads. +template +__device__ __forceinline__ float2 load_x_pair( + const void* __restrict__ X, const uint8_t* __restrict__ Xs, + const float tensorScale, const int64_t idx, const int64_t scaleIdx) { + float2 pair; + if constexpr (STORE == XS_FP4) { + __hip_fp4x2_e2m1 packed; + packed.__x = static_cast(X)[idx >> 1]; + pair = static_cast(packed); + } else if constexpr (STORE == XS_FP8) { +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + __hip_fp8x2_e4m3_fnuz packed; +#else + __hip_fp8x2_e4m3 packed; +#endif + packed.__x = static_cast(X)[idx >> 1]; + pair = static_cast(packed); + } else { + pair = make_float2( + load_x_raw(X, idx), load_x_raw(X, idx + 1)); + } + float scale = 1.f; + if constexpr (SCALE == XSC_TENSOR) scale = tensorScale; + else if constexpr (SCALE == XSC_MX) + scale = e8m0_decode(Xs[scaleIdx]); + pair.x *= scale; + pair.y *= scale; + return pair; +} + +// Bias is a compile-time choice. In the BS_NONE case the +// compiler replaces this whole function with 0 and removes the bias pointer +// and global-memory load. FP16/BF16 bias is converted to float before the +// activation math. BS_F32 exists for the accuracy-only test path. +template +__device__ __forceinline__ float load_bias(const void* __restrict__ B, + const int d) { + if constexpr (BIAS == BS_NONE) { + (void)B; + (void)d; + return 0.f; + } else if constexpr (BIAS == BS_F16) { + return __half2float(static_cast(B)[d]); + } else if constexpr (BIAS == BS_BF16) { + return __bfloat162float(static_cast(B)[d]); + } else { + return static_cast(B)[d]; + } +} + +template +__device__ __forceinline__ float swiglu_elem( + const KernelArgs& a, const int64_t row, const int64_t d) { + const int64_t index = row * a.dim + d; + const int64_t xOffset = 2 * index; + int64_t scaleIdx = 0; + if constexpr (SCALE == XSC_MX) { + scaleIdx = (a.dim % (MX_BLOCK / 2) == 0) + ? index / (MX_BLOCK / 2) + : row * a.xBlocksPerRow + (2 * d) / MX_BLOCK; + } + const float2 pair = + load_x_pair(a.X, a.Xs, a.xScale, xOffset, scaleIdx); + float gate = pair.x + load_bias(a.B, 2 * d); + float linear = pair.y + load_bias(a.B, 2 * d + 1); + + gate = fminf(gate, a.limit); + linear = fminf(fmaxf(linear, -a.limit), a.limit); + return gate * dev_sigmoid(a.alpha * gate) * (linear + 1.f); +} + +// One thread per output element. Threads are tiled (row, block of blockDim.x +// columns) rather than laid out flat so that an MXFP4 block always lands on a +// 32-lane logical subgroup and never straddles a row. Shuffle width=MX_BLOCK +// makes that subgroup one full wave when warpSize is 32 and half a wave when +// warpSize is 64. The tile loop is a grid-stride loop over int64 tile ids, so +// both the index and the guard stay 64-bit safe no matter how large rows*dim +// gets. +template +__global__ void swiglu_oai_kernel(const KernelArgs a) { + const int64_t tilesPerRow = + (static_cast(a.dim) + blockDim.x - 1) / blockDim.x; + const int64_t totalTiles = static_cast(a.rows) * tilesPerRow; + + for (int64_t tile = blockIdx.x; tile < totalTiles; tile += gridDim.x) { + const int64_t row = tile / tilesPerRow; + const int64_t d = + (tile % tilesPerRow) * static_cast(blockDim.x) + threadIdx.x; + const bool active = (d < static_cast(a.dim)); + + const float y = + active ? swiglu_elem(a, row, d) : 0.f; + + if constexpr (OUT == OUT_NONE) { + if (active) + static_cast(a.Y)[row * static_cast(a.dim) + d] = y; + } else if constexpr (OUT == OUT_FP8) { + if (active) + static_cast(a.Y)[ + row * static_cast(a.dim) + d] = + device_e4m3_encode(y / a.yScale); + } else { // OUT_MXFP4 + // amax over the 32 lanes of this mx block; inactive lanes contribute 0. + float amax = active ? fabsf(y) : 0.f; + for (int m = 1; m < MX_BLOCK; m <<= 1) + amax = fmaxf(amax, __shfl_xor(amax, m, MX_BLOCK)); + + // Dividing by E2M1's largest power-of-two exponent before conversion + // chooses the block scale. E8M0 round-toward-zero is floor(log2) for + // positive values, exactly matching reference.h's scale rule. E8M0 has + // no zero, so retain byte 127 (scale 1) for an all-zero/NaN block. +#if SWIGLU_HAS_HIP_E8M0_CONVERSION + const uint8_t sb = (amax > 0.f) + ? __hip_cvt_float_to_e8m0( + amax * 0.25f, __HIP_SATFINITE, hipRoundZero) + : 127; +#else + // Older HIP headers have no E8M0 conversion API. For a positive normal + // float, its biased IEEE-754 exponent byte minus kE2M1MaxExp is exactly + // the desired E8M0 byte; subnormals clamp to byte 0. + const int scaleByte = + static_cast((__float_as_uint(amax) >> 23) & 0xffu) - + kE2M1MaxExp; + const uint8_t sb = !(amax > 0.f) ? 127 + : scaleByte < 0 ? 0 + : scaleByte > 254 ? 254 + : static_cast(scaleByte); +#endif + const float s = e8m0_decode(sb); + + // Lane 0 of the block owns the scale. If its column is out of range the + // whole block is, so there is nothing to write. + if ((threadIdx.x & (MX_BLOCK - 1)) == 0 && active) + a.Ys[row * a.yBlocksPerRow + (d / MX_BLOCK)] = sb; + + // Two fp4 values share a byte. The even lane collects its neighbour's + // float and converts the pair together, matching CUDA's packed path. + const float paired = __shfl_down(y, 1, 2); + if (active && (d & 1) == 0) { + const float hi = + (d + 1 < static_cast(a.dim)) ? paired : 0.f; + static_cast(a.Y)[ + row * a.yRowStrideBytes + (d >> 1)] = + __hip_fp4x2_e2m1(make_float2(y / s, hi / s)).__x; + } + } + } +} + +template +static void launch_swiglu_out(OutQuant out, dim3 grids, dim3 blocks, + const KernelArgs& a) { + switch (out) { + case OUT_NONE: + swiglu_oai_kernel<<>>(a); + break; + case OUT_FP8: + swiglu_oai_kernel<<>>(a); + break; + case OUT_MXFP4: + swiglu_oai_kernel<<>>(a); + break; + default: break; + } +} + +template +static void launch_swiglu_bias(InDtype in, OutQuant out, + dim3 grids, dim3 blocks, const KernelArgs& a) { + switch (in) { + case IN_FP32: + launch_swiglu_out(out, grids, blocks, a); break; + case IN_FP16: + launch_swiglu_out(out, grids, blocks, a); break; + case IN_BF16: + launch_swiglu_out(out, grids, blocks, a); break; + case IN_FP8: + launch_swiglu_out(out, grids, blocks, a); break; + case IN_MXFP8: + launch_swiglu_out(out, grids, blocks, a); break; + case IN_FP4: + launch_swiglu_out(out, grids, blocks, a); break; + case IN_MXFP4: + launch_swiglu_out(out, grids, blocks, a); break; + default: break; + } +} + +static void launch_swiglu(InDtype in, OutQuant out, BiasStore bias, + dim3 grids, dim3 blocks, const KernelArgs& a) { + switch (bias) { + case BS_NONE: + launch_swiglu_bias(in, out, grids, blocks, a); break; + case BS_F16: + launch_swiglu_bias(in, out, grids, blocks, a); break; + case BS_BF16: + launch_swiglu_bias(in, out, grids, blocks, a); break; + case BS_F32: + launch_swiglu_bias(in, out, grids, blocks, a); break; + default: break; + } +} + +// --------------------------------------------------------------------------- +// Host packing / unpacking +// --------------------------------------------------------------------------- + +// Bytes needed to store `n` elements of the given input format. +static int64_t in_storage_bytes(InDtype dt, int64_t n) { + const int bits = in_elem_bits(dt); + return (bits == 4) ? (n + 1) / 2 : n * (bits / 8); +} + +// quantize_dequantize_input() hands back dequantized floats; each is exactly +// representable in the target format, so re-encoding here is lossless and the +// device sees bit-identical values to the host reference. `rowLen` is the +// logical row length (2*dim); mx scales are indexed [row][block] because the +// blocks are cut inside a row. +static void pack_input(HipFp8Format fp8_format, InDtype dt, const float* xq, + const uint8_t* xs, float ts, int64_t n, + int64_t rowLen, void* dst) { + const XStore store = x_store_of(dt); + const XScale scaling = x_scale_of(dt); + const int64_t bpr = mx_blocks_per_row(rowLen); + + auto scale_at = [&](int64_t i) -> float { + if (scaling == XSC_MX) { + const int64_t row = i / rowLen; + const int64_t col = i - row * rowLen; + return e8m0_to_float(xs[row * bpr + col / MX_BLOCK]); + } + return ts; + }; + + if (store == XS_FP4) { + uint8_t* p = static_cast(dst); + for (int64_t i = 0; i < n; i += 2) { + const float scale = scale_at(i); + const float hi = (i + 1 < n) ? xq[i + 1] / scale : 0.f; + p[i >> 1] = + __hip_fp4x2_e2m1(make_float2(xq[i] / scale, hi)).__x; + } + return; + } + + for (int64_t i = 0; i < n; ++i) { + switch (store) { + case XS_F32: + static_cast(dst)[i] = xq[i]; + break; + case XS_F16: + static_cast<__half*>(dst)[i] = __float2half(xq[i]); + break; + case XS_BF16: + static_cast<__hip_bfloat16*>(dst)[i] = __float2bfloat16(xq[i]); + break; + default: + static_cast(dst)[i] = + hip_fp8_encode(xq[i] / scale_at(i), fp8_format); + break; + } + } +} + +static int bias_elem_bytes(BiasStore bias) { + if (bias == BS_F16 || bias == BS_BF16) return 2; + if (bias == BS_F32) return 4; + return 0; +} + +static void pack_bias(BiasStore bias, float* values, size_t n, void* dst) { + if (bias == BS_NONE) { + memset(values, 0, n * sizeof(float)); + return; + } + if (bias == BS_F16) { + __half* packed = static_cast<__half*>(dst); + for (size_t i = 0; i < n; ++i) { + packed[i] = __float2half(values[i]); + values[i] = __half2float(packed[i]); + } + return; + } + if (bias == BS_BF16) { + __hip_bfloat16* packed = static_cast<__hip_bfloat16*>(dst); + for (size_t i = 0; i < n; ++i) { + packed[i] = __float2bfloat16(values[i]); + values[i] = __bfloat162float(packed[i]); + } + return; + } + memcpy(dst, values, n * sizeof(float)); +} + +static int64_t out_storage_bytes(OutQuant q, int rows, int dim) { + switch (q) { + case OUT_NONE: + return static_cast(rows) * dim * sizeof(float); + case OUT_FP8: + return static_cast(rows) * dim; + case OUT_MXFP4: + return static_cast(rows) * mxfp4_row_bytes(dim); + default: return 0; + } +} + +// Decode what the kernel stored back into floats so it can be compared with +// the (also dequantized) host reference elementwise. +static void unpack_output(HipFp8Format fp8_format, OutQuant q, + int rows, int dim, const void* src, + const uint8_t* scales, float ts, float* dst) { + if (q == OUT_NONE) { + memcpy(dst, src, static_cast(rows) * dim * sizeof(float)); + return; + } + if (q == OUT_FP8) { + const uint8_t* p = static_cast(src); + for (int64_t i = 0; i < static_cast(rows) * dim; i++) + dst[i] = hip_fp8_decode(p[i], fp8_format) * ts; + return; + } + // OUT_MXFP4 + const int64_t blocksPerRow = mx_blocks_per_row(dim); + const int64_t rowStride = mxfp4_row_bytes(dim); + const uint8_t* p = static_cast(src); + for (int64_t r = 0; r < rows; r++) { + for (int64_t d = 0; d < dim; d++) { + __hip_fp4_e2m1 v; + v.__x = static_cast<__hip_fp4_storage_t>( + (p[r * rowStride + (d >> 1)] >> ((d & 1) << 2)) & 0xf); + dst[r * dim + d] = + static_cast(v) * + e8m0_to_float(scales[r * blocksPerRow + (d / MX_BLOCK)]); + } + } +} + +static void usage(const char* prog) { + printf("Usage: %s " + "[in_dtype] [out_quant] [bias_dtype]\n", prog); + printf(" in_dtype : fp32 | fp16 | bf16 | fp8 | mxfp8 | fp4 | mxfp4 (default fp32)\n"); + printf(" out_quant: none | fp8 | mxfp4 (default none)\n"); + printf(" bias_dtype: none | fp16 | bf16 | fp32 " + "(default none; fp32 is accuracy-only)\n"); +} + +int main(int argc, char* argv[]) +{ + if (argc < 4 || argc > 7) { + usage(argv[0]); + return 1; + } + + const int rows = atoi(argv[1]); + const int dim = atoi(argv[2]); + const int repeat = atoi(argv[3]); + + if (rows <= 0 || dim <= 0) { + fprintf(stderr, "Error: rows and dimension must be positive (got %d and %d)\n", + rows, dim); + return 1; + } + + // The average kernel time is divided by the repeat count. + if (repeat <= 0) { + fprintf(stderr, "Error: repeat count must be positive (got %d)\n", repeat); + return 1; + } + + // Assert the rounding rules directly. The end-to-end comparison cannot catch + // a tie-policy bug -- a tie sits exactly on a rounding boundary, where + // values_match() grants slack by design -- so this is what guards it. + if (!self_test_hip_rounding()) return 1; + + const InDtype in_dtype = (argc > 4) ? parse_in_dtype(argv[4]) : IN_FP32; + if (in_dtype == IN_INVALID) { + fprintf(stderr, "Error: unknown input dtype '%s' " + "(valid: fp32 fp16 bf16 fp8 mxfp8 fp4 mxfp4)\n", argv[4]); + return 1; + } + + const OutQuant out_quant = (argc > 5) ? parse_out_quant(argv[5]) : OUT_NONE; + if (out_quant == OUT_INVALID) { + fprintf(stderr, "Error: unknown output quantization '%s' " + "(valid: none fp8 mxfp4)\n", argv[5]); + return 1; + } + + const BiasStore bias = + (argc > 6) ? parse_bias_store(argv[6]) : BS_NONE; + if (static_cast(bias) < static_cast(BS_NONE) || bias > BS_F32) { + fprintf(stderr, "Error: invalid bias dtype '%s' " + "(valid: none fp16 bf16 fp32)\n", argv[6]); + return 1; + } + + int device; + hipDeviceProp_t device_properties; + HIP_CHECK(hipGetDevice(&device)); + HIP_CHECK(hipGetDeviceProperties(&device_properties, device)); + const HipFp8Format fp8_format = + hip_arch_uses_fnuz(device_properties.gcnArchName) + ? HIP_FP8_FNUZ : HIP_FP8_OCP; + + printf("Shape of input tensor: ( %d %d ) in=%s out=%s bias=%s fp8=%s\n", + rows, 2 * dim, in_dtype_name(in_dtype), out_quant_name(out_quant), + bias_store_name(bias), hip_fp8_format_name(fp8_format)); + + const int64_t x_nelems = static_cast(rows) * 2 * dim; + const int64_t y_nelems = static_cast(rows) * dim; + const int64_t b_nelems = static_cast(2) * dim; + + const int64_t b_bytes = b_nelems * sizeof(float); + const int64_t b_store_bytes = b_nelems * bias_elem_bytes(bias); + const int64_t x_store_bytes = in_storage_bytes(in_dtype, x_nelems); + const int64_t y_store_bytes = out_storage_bytes(out_quant, rows, dim); + + // MX scales are [row][block] on both axes; blocks are cut inside a row. + const int64_t x_blocksPerRow = + mx_blocks_per_row(2 * static_cast(dim)); + const int64_t x_nscales = in_is_mx(in_dtype) + ? static_cast(rows) * x_blocksPerRow : 0; + const int64_t y_blocksPerRow = mx_blocks_per_row(dim); + const int64_t y_nscales = out_is_mx(out_quant) + ? static_cast(rows) * y_blocksPerRow : 0; + + float *X = static_cast(malloc(x_nelems * sizeof(float))); + float *Xq = static_cast(malloc(x_nelems * sizeof(float))); + float *B = static_cast(malloc(b_bytes)); + float *Y = static_cast(malloc(y_nelems * sizeof(float))); + float *Y_ref = static_cast(malloc(y_nelems * sizeof(float))); + // The reference activation before the epilogue quantization. values_match() + // needs it: once quantized, every reference value sits exactly on a + // representable code and so is always half a quantum from a boundary, which + // says nothing about whether a rounding flip was plausible. + float *Y_pre = static_cast(malloc(y_nelems * sizeof(float))); + void *Xpacked = malloc(x_store_bytes); + void *Bpacked = (bias == BS_NONE) ? nullptr : malloc(b_store_bytes); + void *Ypacked = malloc(y_store_bytes); + uint8_t *Xscales = + static_cast(malloc(x_nscales ? x_nscales : 1)); + uint8_t *Yscales = + static_cast(malloc(y_nscales ? y_nscales : 1)); + uint8_t *Yscales_ref = + static_cast(malloc(y_nscales ? y_nscales : 1)); + + if (!X || !Xq || !B || !Y || !Y_ref || !Y_pre || !Xpacked || + (bias != BS_NONE && !Bpacked) || !Ypacked || !Xscales || + !Yscales || !Yscales_ref) { + fprintf(stderr, "Error: host allocation failed\n"); + return EXIT_FAILURE; + } + + std::default_random_engine generator(123); + // The range straddles the clamp limit so both clamped and unclamped + // elements are exercised. + std::uniform_real_distribution distribution(-12.f, 12.f); + + for (int64_t i = 0; i < x_nelems; i++) { + X[i] = distribution(generator); + } + for (int64_t i = 0; i < b_nelems; i++) { + B[i] = distribution(generator); + } + pack_bias(bias, B, static_cast(b_nelems), Bpacked); + + // Ground truth for the input axis: Xq holds the exact values the kernel must + // reconstruct after dequantization. + float xTensorScale = 1.f; + quantize_dequantize_input_hip( + fp8_format, in_dtype, X, x_nelems, 2 * static_cast(dim), + Xq, Xscales, &xTensorScale); + pack_input(fp8_format, in_dtype, Xq, Xscales, xTensorScale, + x_nelems, 2 * static_cast(dim), Xpacked); + + void *d_X, *d_Y, *d_B = nullptr; + uint8_t *d_Xs, *d_Ys; + HIP_CHECK(hipMalloc(&d_X, x_store_bytes)); + if (bias != BS_NONE) { + HIP_CHECK(hipMalloc(&d_B, b_store_bytes)); + HIP_CHECK(hipMemcpy(d_B, Bpacked, b_store_bytes, hipMemcpyHostToDevice)); + } + HIP_CHECK(hipMalloc(&d_Y, y_store_bytes)); + HIP_CHECK(hipMalloc(&d_Xs, x_nscales ? x_nscales : 1)); + HIP_CHECK(hipMalloc(&d_Ys, y_nscales ? y_nscales : 1)); + + HIP_CHECK(hipMemcpy(d_X, Xpacked, x_store_bytes, hipMemcpyHostToDevice)); + if (x_nscales) + HIP_CHECK(hipMemcpy(d_Xs, Xscales, x_nscales, hipMemcpyHostToDevice)); + + // Host reference: activation, then the epilogue quantization. + ComputeSwigluOAI(rows, dim, kSwigluAlpha, kSwigluLimit, Xq, B, Y_ref); + + // A per-tensor fp8 output scale depends on the amax of the whole activation, + // which a single-pass kernel cannot know; as in aiter it is supplied by the + // caller. It is derived here exactly as quantize_dequantize_output() derives + // it internally, so both sides quantize with the same scale. + float yTensorScale = 1.f; + if (out_quant == OUT_FP8) { + float amax = 0.f; + for (int64_t i = 0; i < y_nelems; i++) { + const float v = fabsf(Y_ref[i]); + if (v > amax) amax = v; + } + const float fp8_max = + fp8_format == HIP_FP8_FNUZ ? kE4M3FnuzMax : kE4M3Max; + yTensorScale = (amax > 0.f) ? (amax / fp8_max) : 1.f; + } + memcpy(Y_pre, Y_ref, static_cast(y_nelems) * sizeof(float)); + quantize_dequantize_output_hip( + fp8_format, out_quant, rows, dim, Y_ref, Yscales_ref); + + const int block_size = 256; + const int64_t tilesPerRow = (dim + block_size - 1) / block_size; + const int64_t totalTiles = static_cast(rows) * tilesPerRow; + // The grid-stride loop lets the launch geometry stay inside the 32-bit grid + // limit while the tile index, the element index and the guard remain int64. + const int64_t maxBlocks = 1 << 20; + dim3 grids(static_cast( + totalTiles < maxBlocks ? totalTiles : maxBlocks)); + dim3 blocks (block_size); + + KernelArgs args; + args.rows = rows; + args.dim = dim; + args.alpha = kSwigluAlpha; + args.limit = kSwigluLimit; + args.X = d_X; + args.Xs = d_Xs; + args.xScale = xTensorScale; + args.B = d_B; + args.Y = d_Y; + args.Ys = d_Ys; + args.yScale = yTensorScale; + args.xBlocksPerRow = x_blocksPerRow; + args.yRowStrideBytes = mxfp4_row_bytes(dim); + args.yBlocksPerRow = y_blocksPerRow; + + // check correctness before benchmarking + launch_swiglu(in_dtype, out_quant, bias, grids, blocks, args); + HIP_CHECK(hipGetLastError()); + HIP_CHECK(hipMemcpy(Ypacked, d_Y, y_store_bytes, hipMemcpyDeviceToHost)); + if (y_nscales) + HIP_CHECK(hipMemcpy(Yscales, d_Ys, y_nscales, hipMemcpyDeviceToHost)); + unpack_output(fp8_format, out_quant, rows, dim, Ypacked, Yscales, + yTensorScale, Y); + + // The tolerance is the baseline for the unquantized path; values_match() adds + // one quantum of slack for a quantized element -- but only when the + // pre-quantization reference sat near a rounding boundary -- using that + // element's e8m0 block scale (mxfp4) or the per-tensor scale (fp8). + const float tol = 1e-3f; + bool ok = true; + float maxdiff = 0.f; + int64_t worst = 0; + for (int64_t r = 0; r < rows; r++) { + for (int64_t d = 0; d < dim; d++) { + const int64_t i = r * dim + d; + float scale = 1.f; + if (out_quant == OUT_FP8) + scale = yTensorScale; + else if (out_quant == OUT_MXFP4) + scale = e8m0_to_float(Yscales_ref[r * y_blocksPerRow + d / MX_BLOCK]); + + const float diff = fabsf(Y[i] - Y_ref[i]); + if (diff > maxdiff) { maxdiff = diff; worst = i; } + if (!hip_values_match(fp8_format, Y[i], Y_ref[i], Y_pre[i], + out_quant, scale, tol)) + ok = false; + } + } + printf("%s\n", ok ? "PASS" : "FAIL"); + if (!ok) { + int64_t scale_mismatch = 0; + for (int64_t i = 0; i < y_nscales; i++) + if (Yscales[i] != Yscales_ref[i]) scale_mismatch++; + fprintf(stderr, " max |device - reference| = %g at %ld (device %g, reference %g, tol %g)\n", + maxdiff, static_cast(worst), Y[worst], Y_ref[worst], tol); + if (y_nscales) + fprintf(stderr, " mismatching e8m0 output scales: %ld of %ld\n", + static_cast(scale_mismatch), + static_cast(y_nscales)); + } + + if (bias == BS_F32) { + printf("FP32 bias is accuracy-only; performance timing skipped.\n"); + } else { + HIP_CHECK(hipDeviceSynchronize()); + auto start = std::chrono::steady_clock::now(); + + for (int i = 0; i < repeat; i++) { + launch_swiglu(in_dtype, out_quant, bias, grids, blocks, args); + } + + HIP_CHECK(hipDeviceSynchronize()); + auto end = std::chrono::steady_clock::now(); + auto time = std::chrono::duration_cast( + end - start).count(); + printf("Average execution time of SwiGLU (OAI) kernel: %f (us)\n", + (time * 1e-3f) / repeat); + } + + HIP_CHECK(hipFree(d_X)); + if (d_B != nullptr) HIP_CHECK(hipFree(d_B)); + HIP_CHECK(hipFree(d_Y)); + HIP_CHECK(hipFree(d_Xs)); + HIP_CHECK(hipFree(d_Ys)); + + free(X); + free(Xq); + free(B); + free(Y); + free(Y_ref); + free(Y_pre); + free(Xpacked); + free(Bpacked); + free(Ypacked); + free(Xscales); + free(Yscales); + free(Yscales_ref); + + return EXIT_SUCCESS; +} diff --git a/src/swiglu-oai-hip/reference.h b/src/swiglu-oai-hip/reference.h new file mode 100644 index 000000000..2b4601f7a --- /dev/null +++ b/src/swiglu-oai-hip/reference.h @@ -0,0 +1,173 @@ +#ifndef SWIGLU_OAI_HIP_REFERENCE_H +#define SWIGLU_OAI_HIP_REFERENCE_H + +#include "../swiglu-oai-cuda/reference.h" +#include + +enum HipFp8Format { HIP_FP8_OCP = 0, HIP_FP8_FNUZ }; + +constexpr float kE4M3FnuzMax = 240.f; +constexpr int kE4M3FnuzMaxExp = 7; + +inline bool hip_arch_uses_fnuz(const char* arch) { + return !strncmp(arch, "gfx940", 6) || + !strncmp(arch, "gfx941", 6) || + !strncmp(arch, "gfx942", 6); +} + +inline const char* hip_fp8_format_name(HipFp8Format format) { + return format == HIP_FP8_FNUZ ? "fnuz" : "ocp"; +} + +inline float round_to_e4m3_fnuz(float x) { + if (!(x == x)) return x; + const float a = fabsf(x); + const float sign = (x < 0.f) ? -1.f : 1.f; + if (a >= kE4M3FnuzMax) return sign * kE4M3FnuzMax; + if (a == 0.f) return 0.f; // FNUZ has one unsigned zero. + + int e; + frexpf(a, &e); + e -= 1; + if (e < -7) e = -7; // FNUZ E4M3 uses exponent bias 8. + const float quantum = ldexpf(1.f, e - 3); + const float n = a / quantum; + const float lower = floorf(n); + const float fraction = n - lower; + float rounded; + if (fraction > 0.5f) + rounded = lower + 1.f; + else if (fraction < 0.5f) + rounded = lower; + else + rounded = (fmodf(lower, 2.f) == 0.f) ? lower : lower + 1.f; + + float result = rounded * quantum; + if (result > kE4M3FnuzMax) result = kE4M3FnuzMax; + return result == 0.f ? 0.f : sign * result; +} + +inline uint8_t hip_fp8_encode(float x, HipFp8Format format) { + return format == HIP_FP8_FNUZ ? __hip_fp8_e4m3_fnuz(x).__x + : __hip_fp8_e4m3(x).__x; +} + +inline float hip_fp8_decode(uint8_t x, HipFp8Format format) { + if (format == HIP_FP8_FNUZ) { + __hip_fp8_e4m3_fnuz value; + value.__x = x; + return (float)value; + } + __hip_fp8_e4m3 value; + value.__x = x; + return (float)value; +} + +inline void quantize_dequantize_input_hip( + HipFp8Format format, InDtype dtype, const float* input, int64_t n, + int64_t row_len, float* output, uint8_t* scales, float* tensor_scale) { + if (format != HIP_FP8_FNUZ || + (dtype != IN_FP8 && dtype != IN_MXFP8)) { + quantize_dequantize_input( + dtype, input, n, row_len, output, scales, tensor_scale); + return; + } + + *tensor_scale = 1.f; + if (dtype == IN_FP8) { + float amax = 0.f; + for (int64_t i = 0; i < n; ++i) + amax = fmaxf(amax, fabsf(input[i])); + const float scale = amax > 0.f ? amax / kE4M3FnuzMax : 1.f; + *tensor_scale = scale; + for (int64_t i = 0; i < n; ++i) + output[i] = round_to_e4m3_fnuz(input[i] / scale) * scale; + return; + } + + const int64_t blocks_per_row = mx_blocks_per_row(row_len); + const int64_t rows = n / row_len; + for (int64_t row = 0; row < rows; ++row) { + for (int64_t block = 0; block < blocks_per_row; ++block) { + const int64_t begin = row * row_len + block * MX_BLOCK; + const int64_t row_end = (row + 1) * row_len; + const int64_t end = begin + MX_BLOCK < row_end + ? begin + MX_BLOCK : row_end; + float amax = 0.f; + for (int64_t i = begin; i < end; ++i) + amax = fmaxf(amax, fabsf(input[i])); + const uint8_t scale_byte = + e8m0_scale_byte(amax, kE4M3FnuzMaxExp); + scales[row * blocks_per_row + block] = scale_byte; + const float scale = e8m0_to_float(scale_byte); + for (int64_t i = begin; i < end; ++i) + output[i] = round_to_e4m3_fnuz(input[i] / scale) * scale; + } + } +} + +inline void quantize_dequantize_output_hip( + HipFp8Format format, OutQuant quant, int rows, int dim, + float* output, uint8_t* scales) { + if (format != HIP_FP8_FNUZ || quant != OUT_FP8) { + quantize_dequantize_output(quant, rows, dim, output, scales); + return; + } + + const int64_t n = (int64_t)rows * dim; + float amax = 0.f; + for (int64_t i = 0; i < n; ++i) + amax = fmaxf(amax, fabsf(output[i])); + const float scale = amax > 0.f ? amax / kE4M3FnuzMax : 1.f; + for (int64_t i = 0; i < n; ++i) + output[i] = round_to_e4m3_fnuz(output[i] / scale) * scale; +} + +inline bool hip_values_match( + HipFp8Format format, float device, float reference, + float reference_prequant, OutQuant quant, float scale, float tolerance) { + if (format != HIP_FP8_FNUZ || quant != OUT_FP8) + return values_match(device, reference, reference_prequant, + quant, scale, tolerance); + + const float difference = fabsf(device - reference); + if (difference <= tolerance) return true; + + const float a = fabsf(reference_prequant / scale); + int exponent = -7; + if (a > 0.f) { + frexpf(a, &exponent); + exponent -= 1; + if (exponent < -7) exponent = -7; + } + const float quantum = ldexpf(1.f, exponent - 3); + const float n = a / quantum; + const bool near_boundary = + fabsf(n - floorf(n) - 0.5f) <= 1e-3f; + return near_boundary && + difference <= tolerance + quantum * scale; +} + +inline bool self_test_hip_rounding() { + if (!self_test_rounding()) return false; + struct Case { float input; float expected; }; + const Case cases[] = { + {1.0625f, 1.f}, + {1.1875f, 1.25f}, + {1.09f, 1.125f}, + {1e30f, kE4M3FnuzMax}, + {-0.f, 0.f}, + }; + for (const Case& test : cases) { + const float result = round_to_e4m3_fnuz(test.input); + if (result != test.expected || signbit(result) != signbit(test.expected)) { + fprintf(stderr, + "HIP reference self-test FAILED: FNUZ round(%g)=%g, expected %g\n", + test.input, result, test.expected); + return false; + } + } + return true; +} + +#endif diff --git a/src/swiglu-oai-omp/CMakeLists.txt b/src/swiglu-oai-omp/CMakeLists.txt new file mode 100644 index 000000000..7fdba6884 --- /dev/null +++ b/src/swiglu-oai-omp/CMakeLists.txt @@ -0,0 +1,8 @@ +# swiglu-oai-omp/CMakeLists.txt + +add_hecbench_benchmark( + NAME swiglu-oai + MODEL omp + SOURCES main.cpp + CATEGORIES algorithms +) diff --git a/src/swiglu-oai-omp/Makefile b/src/swiglu-oai-omp/Makefile new file mode 100644 index 000000000..9ce19b4c0 --- /dev/null +++ b/src/swiglu-oai-omp/Makefile @@ -0,0 +1,62 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = icpx +OPTIMIZE = yes +DEBUG = no +DEVICE = gpu +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../swiglu-oai-cuda + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(DEVICE),gpu) + CFLAGS +=-fiopenmp -fopenmp-targets=spir64 -D__STRICT_ANSI__ +else + CFLAGS +=-qopenmp +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cpp ../swiglu-oai-cuda/reference.h + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 8192 1024 100 diff --git a/src/swiglu-oai-omp/Makefile.aomp b/src/swiglu-oai-omp/Makefile.aomp new file mode 100644 index 000000000..11347caf9 --- /dev/null +++ b/src/swiglu-oai-omp/Makefile.aomp @@ -0,0 +1,66 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = clang++ +OPTIMIZE = yes +DEBUG = no +DEVICE = gpu +ARCH = gfx906 +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../swiglu-oai-cuda + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(DEVICE),gpu) + CFLAGS += -target x86_64-pc-linux-gnu \ + -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa \ + -Xopenmp-target=amdgcn-amd-amdhsa \ + -march=$(ARCH) +else + CFLAGS +=-fopenmp +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cpp ../swiglu-oai-cuda/reference.h + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 8192 1024 100 diff --git a/src/swiglu-oai-omp/Makefile.nvc b/src/swiglu-oai-omp/Makefile.nvc new file mode 100644 index 000000000..6a166e5a6 --- /dev/null +++ b/src/swiglu-oai-omp/Makefile.nvc @@ -0,0 +1,63 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = nvc++ +OPTIMIZE = yes +DEBUG = no +DEVICE = gpu +SM ?= cc70 +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -Wall -I../swiglu-oai-cuda + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(DEVICE),gpu) + CFLAGS +=-Minfo -mp=gpu -gpu=$(SM) +else + CFLAGS += +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cpp ../swiglu-oai-cuda/reference.h + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 8192 1024 100 diff --git a/src/swiglu-oai-omp/main.cpp b/src/swiglu-oai-omp/main.cpp new file mode 100644 index 000000000..10afbaba6 --- /dev/null +++ b/src/swiglu-oai-omp/main.cpp @@ -0,0 +1,131 @@ +#include +#include +#include +#include +#include +#include +#include +#include "reference.h" + +#define block_size 256 + +void swiglu_oai_kernel( + const int rows, + const int dim, + const float alpha, + const float limit, + const float* Xdata, + const float* Bdata, + float* Ydata) +{ + #pragma omp target teams distribute parallel for num_threads(block_size) + for (int index = 0; index < rows * dim; index++) { + const int i = index / dim; + const int d = index % dim; + + const int64_t xOffset = (int64_t)i * 2 * dim; + + float gate = Xdata[xOffset + 2 * d] + Bdata[2 * d]; + float linear = Xdata[xOffset + 2 * d + 1] + Bdata[2 * d + 1]; + + gate = fminf(gate, limit); + linear = fminf(fmaxf(linear, -limit), limit); + + const float s = 1.f / (1.f + expf(-alpha * gate)); + Ydata[(int64_t)i * dim + d] = gate * s * (linear + 1.f); + } +} + +int main(int argc, char* argv[]) +{ + if (argc != 4) { + printf("Usage: %s \n", argv[0]); + return 1; + } + + const int rows = atoi(argv[1]); + const int dim = atoi(argv[2]); + const int repeat = atoi(argv[3]); + + if (rows <= 0 || dim <= 0) { + fprintf(stderr, "Error: rows and dimension must be positive (got %d and %d)\n", + rows, dim); + return 1; + } + + // The average kernel time is divided by the repeat count. + if (repeat <= 0) { + fprintf(stderr, "Error: repeat count must be positive (got %d)\n", repeat); + return 1; + } + + // Assert the rounding rules directly. The end-to-end comparison cannot catch + // a tie-policy bug -- a tie sits exactly on a rounding boundary, where + // values_match() grants slack by design -- so this is what guards it. + if (!self_test_rounding()) return 1; + + printf("Shape of input tensor: ( %d %d )\n", rows, 2 * dim); + + const uint64_t x_nelems = (uint64_t)rows * 2 * dim; + const uint64_t y_nelems = (uint64_t)rows * dim; + const uint64_t b_nelems = (uint64_t)2 * dim; + + const uint64_t x_bytes = x_nelems * sizeof(float); + const uint64_t y_bytes = y_nelems * sizeof(float); + const uint64_t b_bytes = b_nelems * sizeof(float); + + float *X = (float*) malloc (x_bytes); + float *B = (float*) malloc (b_bytes); + float *Y = (float*) malloc (y_bytes); + float *Y_ref = (float*) malloc (y_bytes); + + std::default_random_engine generator(123); + // The range straddles the clamp limit so both clamped and unclamped + // elements are exercised. + std::uniform_real_distribution distribution(-12.f, 12.f); + + for (uint64_t i = 0; i < x_nelems; i++) { + X[i] = distribution(generator); + } + for (uint64_t i = 0; i < b_nelems; i++) { + B[i] = distribution(generator); + } + + ComputeSwigluOAI(rows, dim, kSwigluAlpha, kSwigluLimit, X, B, Y_ref); + + #pragma omp target data map(to: X[0:x_nelems], B[0:b_nelems]) \ + map(from: Y[0:y_nelems]) + { + // check correctness before benchmarking + swiglu_oai_kernel(rows, dim, kSwigluAlpha, kSwigluLimit, X, B, Y); + + #pragma omp target update from (Y[0:y_nelems]) + + bool ok = true; + for (uint64_t i = 0; i < y_nelems; i++) { + if (fabsf(Y[i] - Y_ref[i]) > 1e-3f) { + ok = false; + break; + } + } + printf("%s\n", ok ? "PASS" : "FAIL"); + + auto start = std::chrono::steady_clock::now(); + + for (int i = 0; i < repeat; i++) { + swiglu_oai_kernel(rows, dim, kSwigluAlpha, kSwigluLimit, X, B, Y); + } + + auto end = std::chrono::steady_clock::now(); + auto time = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of SwiGLU (OAI) kernel: %f (us)\n", + (time * 1e-3f) / repeat); + } + + free(X); + free(B); + free(Y); + free(Y_ref); + + return 0; +} diff --git a/src/swiglu-oai-sycl/CMakeLists.txt b/src/swiglu-oai-sycl/CMakeLists.txt new file mode 100644 index 000000000..66ce44b39 --- /dev/null +++ b/src/swiglu-oai-sycl/CMakeLists.txt @@ -0,0 +1,8 @@ +# swiglu-oai-sycl/CMakeLists.txt + +add_hecbench_benchmark( + NAME swiglu-oai + MODEL sycl + SOURCES main.cpp + CATEGORIES algorithms +) diff --git a/src/swiglu-oai-sycl/Makefile b/src/swiglu-oai-sycl/Makefile new file mode 100644 index 000000000..8f8d697f2 --- /dev/null +++ b/src/swiglu-oai-sycl/Makefile @@ -0,0 +1,102 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = clang++ +OPTIMIZE = yes +DEBUG = no +LAUNCHER ?= + +GPU = yes +CUDA = no +CUDA_ARCH ?= sm_90 +HIP = no +HIP_ARCH ?= gfx908 +#GCC_TOOLCHAIN = "/auto/software/gcc/x86_64/gcc-9.1.0/" + +# Run options. The default run mirrors the representative CUDA/HIP paths. +ROWS ?= 8192 +DIM ?= 1024 +REPEAT ?= 100 + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../swiglu-oai-cuda -fsycl + +ifneq ($(strip $(GCC_TOOLCHAIN)),) + CFLAGS += --gcc-toolchain=$(GCC_TOOLCHAIN) +endif + +ifeq ($(VENDOR), AdaptiveCpp) + CFLAGS_TMP := $(CFLAGS) + CFLAGS = $(filter-out -fsycl, $(CFLAGS_TMP)) +endif + +# Linker Flags +LDFLAGS = + +ifeq ($(CUDA), yes) + CFLAGS += -fsycl-targets=nvptx64-nvidia-cuda \ + -Xsycl-target-backend --cuda-gpu-arch=$(CUDA_ARCH) +endif + +ifeq ($(HIP), yes) + CFLAGS += -fsycl-targets=amdgcn-amd-amdhsa \ + -Xsycl-target-backend --offload-arch=$(HIP_ARCH) +endif + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g -DDEBUG + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(GPU),yes) + CFLAGS +=-DUSE_GPU +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(source) ../swiglu-oai-cuda/reference.h Makefile + $(CC) $(CFLAGS) $(source) -o $@ $(LDFLAGS) + +clean: + rm -rf $(program) $(obj) + +.PHONY: clean run + +run: $(program) + @run_case() { \ + output="$$($(LAUNCHER) ./$(program) "$$@" 2>&1)"; status=$$?; \ + printf '%s\n' "$$output"; \ + if [ $$status -ne 0 ] || \ + ! printf '%s\n' "$$output" | awk '$$0 == "PASS" { pass=1 } END { exit !pass }'; then \ + return 1; \ + fi; \ + }; \ + set -e; \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 none none; \ + run_case $(ROWS) $(DIM) $(REPEAT) fp16 none fp16; \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 fp8 none; \ + run_case $(ROWS) $(DIM) $(REPEAT) bf16 mxfp4 none; \ + run_case $(ROWS) $(DIM) $(REPEAT) mxfp8 fp8 bf16 diff --git a/src/swiglu-oai-sycl/main.cpp b/src/swiglu-oai-sycl/main.cpp new file mode 100644 index 000000000..bf713afba --- /dev/null +++ b/src/swiglu-oai-sycl/main.cpp @@ -0,0 +1,788 @@ +// SwiGLU as used by the OpenAI (gpt-oss) MoE, implementing the two-axis +// precision contract defined in ../swiglu-oai-cuda/reference.h: the gate/up +// tensor is stored in one of fp32/fp16/bf16/fp8/mxfp8/fp4/mxfp4 and is +// dequantized by the kernel, and the activation result is optionally quantized +// to fp8 or mxfp4 in the epilogue. +// +#include +#include +#include +#include +#include +#include +#include +#include "reference.h" + +using sycl::ext::oneapi::bfloat16; + +// --------------------------------------------------------------------------- +// Device-side format conversions. These mirror the host helpers in reference.h +// step for step so that the two sides make identical rounding decisions. +// --------------------------------------------------------------------------- +namespace dev { + +// e8m0: biased power-of-two scale. Mirrors e8m0_to_float(). +inline float e8m0_decode(uint8_t b) { + const uint32_t bits = + b ? static_cast(b) << 23 : 0x00400000u; + return sycl::bit_cast(bits); +} + +// Mirrors e8m0_scale_byte(). +inline uint8_t e8m0_scale_byte(float amax, int elemMaxExp) { + if (!(amax > 0.f)) return 127; // 2^0 for an all-zero block + int e; + sycl::frexp(amax, &e); + e -= 1; // unbiased exponent of amax + int s = e - elemMaxExp + 127; + if (s < 0) s = 0; + if (s > 254) s = 254; + return static_cast(s); +} + +// Round half to even, as round_to_e4m3() does. Hardware convert instructions +// break ties this way, so roundf()-style half-away-from-zero would disagree. +inline float rne(float n) { + const float fl = sycl::floor(n); + const float frac = n - fl; + if (frac > 0.5f) return fl + 1.f; + if (frac < 0.5f) return fl; + return (sycl::fmod(fl, 2.f) == 0.f) ? fl : fl + 1.f; // tie -> even +} + +// OCP e4m3 (bias 7, no infinities) decode. The encoder below saturates at 448, +// so the 0x7f/0xff NaN pattern is never stored and needs no special case here. +inline float e4m3_to_float(uint8_t v) { + const int exp = (v >> 3) & 0xf; + const int mant = v & 0x7; + const float sgn = (v & 0x80u) ? -1.f : 1.f; + if (exp == 0) // subnormal: mant * 2^-9 + return sgn * sycl::ldexp(static_cast(mant), -9); + return sgn * sycl::ldexp(static_cast(8 + mant), exp - 10); +} + +// e2m1 decode. Representable magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. +inline float e2m1_to_float(uint8_t v) { + const float lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + const float m = lut[v & 0x7]; + return (v & 0x8u) ? -m : m; +} + +// Encode to OCP e4m3, rounding exactly as round_to_e4m3() does. A negative zero +// encodes as +0, which is numerically identical. +inline uint8_t f32_to_e4m3(float x) { + const uint8_t sign = (x < 0.f) ? 0x80u : 0x00u; + const float a = sycl::fabs(x); + if (a >= kE4M3Max) return sign | 0x7eu; // 448, largest finite magnitude + if (a == 0.f) return sign; + + int e; + sycl::frexp(a, &e); + e -= 1; + const int minNormExp = -6; // e4m3 bias 7 => smallest normal 2^-6 + if (e < minNormExp) e = minNormExp; + int r = static_cast( + rne(a / sycl::ldexp(1.f, e - 3))); // 3 mantissa bits + if (r >= 16) { r = 8; e += 1; } + if (e == minNormExp && r < 8) + return sign | static_cast(r); // subnormal + return sign | static_cast(((e + 7) << 3) | (r - 8)); +} + +// Encode to e2m1, mirroring round_to_e2m1() including its round-to-even tie +// rule -- note the strict `> 5.f`, which lets the 4/6 tie at exactly 5 fall +// through to the search and land on 4 (even mantissa). +inline uint8_t f32_to_e2m1(float x) { + const float lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + const uint8_t sign = (x < 0.f) ? 0x8u : 0x0u; + const float a = sycl::fabs(x); + if (a > 5.f) return sign | 7u; + int best = 0; + float bestd = sycl::fabs(a - lut[0]); + for (int i = 1; i < 8; i++) { + const float d = sycl::fabs(a - lut[i]); + if (d < bestd || (d == bestd && (i % 2) == 0)) { + bestd = d; + best = i; + } + } + return sign | static_cast(best); +} + +inline float dev_sigmoid(const float x) { + if (x >= 0.f) { + return 1.f / (1.f + sycl::exp(-x)); + } else { + const float exp_x = sycl::exp(x); + return exp_x / (1.f + exp_x); + } +} + +} // namespace dev + +template +inline float load_x_raw(const void *X, int64_t index) { + if constexpr (STORE == XS_F32) { + return static_cast(X)[index]; + } else if constexpr (STORE == XS_F16) { + return static_cast(static_cast(X)[index]); + } else if constexpr (STORE == XS_BF16) { + return static_cast(static_cast(X)[index]); + } else if constexpr (STORE == XS_FP8) { + return dev::e4m3_to_float(static_cast(X)[index]); + } else { + const uint8_t byte = static_cast(X)[index >> 1]; + return dev::e2m1_to_float( + (index & 1) ? static_cast(byte >> 4) + : static_cast(byte & 0xf)); + } +} + +template +inline sycl::float2 load_x_pair( + const void *X, const uint8_t *Xs, float tensor_scale, + int64_t index, int64_t scale_index) { + sycl::float2 pair( + load_x_raw(X, index), load_x_raw(X, index + 1)); + float scale = 1.f; + if constexpr (SCALE == XSC_TENSOR) scale = tensor_scale; + else if constexpr (SCALE == XSC_MX) + scale = dev::e8m0_decode(Xs[scale_index]); + pair *= scale; + return pair; +} + +// The activation for one output element, computed in fp32 whatever the storage +// format is. Gate and linear are interleaved along the last dimension. +template +inline float load_bias(const KernelArgs &a, int64_t d) { + if constexpr (BIAS == BS_NONE) { + (void)a; + (void)d; + return 0.f; + } else if constexpr (BIAS == BS_F16) { + return static_cast(static_cast(a.B)[d]); + } else if constexpr (BIAS == BS_BF16) { + return static_cast(static_cast(a.B)[d]); + } else { + return static_cast(a.B)[d]; + } +} + +template +inline float swiglu_elem(const KernelArgs &a, int64_t row, int64_t d) { + const int64_t index = row * static_cast(a.dim) + d; + const int64_t x_offset = 2 * index; + int64_t scale_index = 0; + if constexpr (SCALE == XSC_MX) { + scale_index = (a.dim % (MX_BLOCK / 2) == 0) + ? index / (MX_BLOCK / 2) + : row * a.xBlocksPerRow + (2 * d) / MX_BLOCK; + } + const sycl::float2 pair = + load_x_pair(a.X, a.Xs, a.xScale, x_offset, scale_index); + float gate = pair.x() + load_bias(a, 2 * d); + float linear = pair.y() + load_bias(a, 2 * d + 1); + + gate = sycl::fmin(gate, a.limit); + linear = sycl::fmin(sycl::fmax(linear, -a.limit), a.limit); + + return gate * dev::dev_sigmoid(a.alpha * gate) * (linear + 1.f); +} + +template +void swiglu_oai_kernel(sycl::nd_item<1> item, const KernelArgs a) +{ + const int64_t tiles_per_row = + (static_cast(a.dim) + item.get_local_range(0) - 1) / + item.get_local_range(0); + const int64_t total_tiles = + static_cast(a.rows) * tiles_per_row; + const int64_t group_stride = + static_cast(item.get_group_range(0)); + for (int64_t tile = static_cast(item.get_group(0)); + tile < total_tiles; tile += group_stride) { + const int64_t row = tile / tiles_per_row; + const int64_t d = (tile % tiles_per_row) * + static_cast(item.get_local_range(0)) + + static_cast(item.get_local_id(0)); + const bool active = d < a.dim; + const float y = + active ? swiglu_elem(a, row, d) : 0.f; + + if constexpr (OUT == OUT_NONE) { + if (active) + static_cast(a.Y)[ + row * static_cast(a.dim) + d] = y; + } else if constexpr (OUT == OUT_FP8) { + if (active) + static_cast(a.Y)[ + row * static_cast(a.dim) + d] = + dev::f32_to_e4m3(y / a.yScale); + } else { + const sycl::sub_group sg = item.get_sub_group(); + float amax = active ? sycl::fabs(y) : 0.f; + for (int mask = 1; mask < MX_BLOCK; mask <<= 1) + amax = sycl::fmax( + amax, sycl::permute_group_by_xor(sg, amax, mask)); + const uint8_t sb = dev::e8m0_scale_byte(amax, kE2M1MaxExp); + const float s = dev::e8m0_decode(sb); + const int lane = + static_cast(sg.get_local_linear_id()) & (MX_BLOCK - 1); + + if (lane == 0 && active) + a.Ys[row * a.yBlocksPerRow + d / MX_BLOCK] = sb; + + const float paired = sycl::shift_group_left(sg, y, 1); + if (active && lane % 2 == 0) { + const float hi = (d + 1 < a.dim) ? paired : 0.f; + const uint8_t lower_code = dev::f32_to_e2m1(y / s); + const uint8_t upper_code = dev::f32_to_e2m1(hi / s); + static_cast(a.Y)[row * a.yRowStrideBytes + (d >> 1)] = + static_cast(lower_code | (upper_code << 4)); + } + } + } +} + +template +class SwigluOAIKernel; + +template +static void submit_swiglu_kernel(sycl::queue &q, sycl::nd_range<1> launch, + const KernelArgs &a) { + q.submit([&](sycl::handler &cgh) { + cgh.parallel_for>( + launch, [=](sycl::nd_item<1> item) { + swiglu_oai_kernel(item, a); + }); + }); +} + +template +static void submit_swiglu_out(sycl::queue &q, OutQuant out, + sycl::nd_range<1> launch, const KernelArgs &a) { + switch (out) { + case OUT_NONE: + submit_swiglu_kernel( + q, launch, a); break; + case OUT_FP8: + submit_swiglu_kernel( + q, launch, a); break; + case OUT_MXFP4: + submit_swiglu_kernel( + q, launch, a); break; + default: break; + } +} + +template +static void submit_swiglu_input(sycl::queue &q, InDtype in, OutQuant out, + sycl::nd_range<1> launch, + const KernelArgs &a) { + switch (in) { + case IN_FP32: + submit_swiglu_out( + q, out, launch, a); break; + case IN_FP16: + submit_swiglu_out( + q, out, launch, a); break; + case IN_BF16: + submit_swiglu_out( + q, out, launch, a); break; + case IN_FP8: + submit_swiglu_out( + q, out, launch, a); break; + case IN_MXFP8: + submit_swiglu_out( + q, out, launch, a); break; + case IN_FP4: + submit_swiglu_out( + q, out, launch, a); break; + case IN_MXFP4: + submit_swiglu_out( + q, out, launch, a); break; + default: break; + } +} + +template +static void submit_swiglu(sycl::queue &q, InDtype in, OutQuant out, + BiasStore bias, sycl::nd_range<1> launch, + const KernelArgs &a) { +#define SUBMIT_BIAS(B) \ + submit_swiglu_input(q, in, out, launch, a) + switch (bias) { + case BS_NONE: SUBMIT_BIAS(BS_NONE); break; + case BS_F16: SUBMIT_BIAS(BS_F16); break; + case BS_BF16: SUBMIT_BIAS(BS_BF16); break; + case BS_F32: SUBMIT_BIAS(BS_F32); break; + default: break; + } +#undef SUBMIT_BIAS +} + +// --------------------------------------------------------------------------- +// Host-side packing and unpacking. The encoders round the same way reference.h +// does, so the value the kernel decodes is bit-identical to the value +// quantize_dequantize_input() left behind. +// --------------------------------------------------------------------------- + +static float host_rne(float n) { // round half to even, as round_to_e4m3() does + const float fl = floorf(n); + const float frac = n - fl; + if (frac > 0.5f) return fl + 1.f; + if (frac < 0.5f) return fl; + return (fmodf(fl, 2.f) == 0.f) ? fl : fl + 1.f; +} + +static uint8_t f32_to_e4m3_bits(float v) { + const uint8_t sign = (v < 0.f) ? 0x80u : 0x00u; + const float a = fabsf(v); + if (a >= kE4M3Max) return sign | 0x7eu; // 448, largest finite magnitude + if (a == 0.f) return sign; + + int e; + frexpf(a, &e); + e -= 1; + if (e < -6) e = -6; + int r = static_cast(host_rne(a / ldexpf(1.f, e - 3))); + if (r >= 16) { r = 8; e += 1; } + if (e == -6 && r < 8) + return sign | static_cast(r); // subnormal + return sign | static_cast(((e + 7) << 3) | (r - 8)); +} + +static float e4m3_bits_to_float(uint8_t v) { + const int exp = (v >> 3) & 0xf; + const int mant = v & 0x7; + const float sgn = (v & 0x80u) ? -1.f : 1.f; + if (exp == 0) return sgn * ldexpf(static_cast(mant), -9); + return sgn * ldexpf(static_cast(8 + mant), exp - 10); +} + +static const float kE2M1Lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + +static uint8_t f32_to_e2m1_nibble(float v) { + const uint8_t sign = (v < 0.f) ? 0x8u : 0x0u; + const float a = fabsf(v); + if (a > 5.f) return sign | 7u; // strict: exactly 5 is a tie, handled below + int best = 0; + float bestd = fabsf(a - kE2M1Lut[0]); + for (int i = 1; i < 8; i++) { + const float d = fabsf(a - kE2M1Lut[i]); + if (d < bestd || (d == bestd && (i % 2) == 0)) { + bestd = d; + best = i; + } + } + return sign | static_cast(best); +} + +static float e2m1_nibble_to_float(uint8_t v) { + const float m = kE2M1Lut[v & 0x7]; + return (v & 0x8u) ? -m : m; +} + +// `Xdq` holds the dequantized values the device must reproduce; dividing by the +// element's scale recovers the code that quantize_dequantize_input() chose. +// MX scales are [row][block] over rows of `rowLen` elements. +static void pack_input(InDtype dtype, const float *Xdq, int64_t nelems, + int64_t rowLen, int64_t bpr, const uint8_t *scales, + float tensor_scale, void *store) +{ + const XStore xstore = x_store_of(dtype); + const XScale scaling = x_scale_of(dtype); + + auto scale_at = [&](int64_t i) -> float { + if (scaling == XSC_MX) { + const int64_t row = i / rowLen; + const int64_t col = i - row * rowLen; + return e8m0_to_float(scales[row * bpr + col / MX_BLOCK]); + } + return tensor_scale; + }; + + if (xstore == XS_FP4) { + uint8_t *p = static_cast(store); + for (int64_t i = 0; i < nelems; i += 2) { + const float scale = scale_at(i); + const uint8_t lo = f32_to_e2m1_nibble(Xdq[i] / scale); + const uint8_t hi = (i + 1 < nelems) + ? f32_to_e2m1_nibble(Xdq[i + 1] / scale) : 0; + p[i >> 1] = static_cast(lo | (hi << 4)); + } + return; + } + + for (int64_t i = 0; i < nelems; ++i) { + switch (xstore) { + case XS_F32: + static_cast(store)[i] = Xdq[i]; + break; + case XS_F16: + static_cast(store)[i] = + static_cast(Xdq[i]); + break; + case XS_BF16: + static_cast(store)[i] = static_cast(Xdq[i]); + break; + default: + static_cast(store)[i] = + f32_to_e4m3_bits(Xdq[i] / scale_at(i)); + break; + } + } +} + +static int bias_elem_bytes(BiasStore bias) { + if (bias == BS_F16 || bias == BS_BF16) return 2; + if (bias == BS_F32) return 4; + return 0; +} + +// Store bias in the requested production type and update the FP32 reference +// values to exactly what the device conversion will produce. +static void pack_bias(BiasStore bias, float *values, int64_t n, void *store) { + if (bias == BS_NONE) { + memset(values, 0, static_cast(n) * sizeof(float)); + return; + } + if (bias == BS_F16) { + sycl::half *p = static_cast(store); + for (int64_t i = 0; i < n; ++i) { + p[i] = static_cast(values[i]); + values[i] = static_cast(p[i]); + } + return; + } + if (bias == BS_BF16) { + bfloat16 *p = static_cast(store); + for (int64_t i = 0; i < n; ++i) { + p[i] = static_cast(values[i]); + values[i] = static_cast(p[i]); + } + return; + } + memcpy(store, values, static_cast(n) * sizeof(float)); +} + +// Decodes the quantized device output back to float using the scales the device +// itself produced, so the comparison sees exactly what the kernel wrote. +static void unpack_output(OutQuant q, int rows, int dim, const void *storage, + const uint8_t *scales, float tensor_scale, float *Y) +{ + if (q == OUT_NONE) { + memcpy(Y, storage, static_cast(rows) * dim * sizeof(float)); + return; + } + const uint8_t *store = static_cast(storage); + if (q == OUT_FP8) { + const int64_t n = static_cast(rows) * dim; + for (int64_t i = 0; i < n; i++) + Y[i] = e4m3_bits_to_float(store[i]) * tensor_scale; + return; + } + + // OUT_MXFP4 + const int64_t bpr = mx_blocks_per_row(dim); + const int64_t rowBytes = mxfp4_row_bytes(dim); + for (int64_t r = 0; r < rows; r++) { + for (int64_t b = 0; b < bpr; b++) { + const int64_t beg = b * MX_BLOCK; + const int64_t n = ((beg + MX_BLOCK < dim) ? beg + MX_BLOCK : dim) - beg; + const float s = e8m0_to_float(scales[r * bpr + b]); + const uint8_t *src = store + r * rowBytes + b * (MX_BLOCK / 2); + for (int64_t t = 0; t < n; t++) { + const uint8_t nibble = (t & 1) ? (src[t >> 1] >> 4) : (src[t >> 1] & 0xf); + Y[r * dim + beg + t] = e2m1_nibble_to_float(nibble) * s; + } + } + } +} + +int main(int argc, char* argv[]) +{ + if (argc < 4 || argc > 7) { + printf("Usage: %s " + "[in_dtype] [out_quant] [bias_dtype]\n", argv[0]); + printf(" in_dtype : fp32 (default), fp16, bf16, fp8, mxfp8, fp4, mxfp4\n"); + printf(" out_quant: none (default), fp8, mxfp4\n"); + printf(" bias_dtype: none (default), fp16, bf16, fp32 " + "(fp32 is accuracy-only)\n"); + return 1; + } + + const int rows = atoi(argv[1]); + const int dim = atoi(argv[2]); + const int repeat = atoi(argv[3]); + + if (rows <= 0 || dim <= 0) { + fprintf(stderr, "Error: rows and dimension must be positive (got %d and %d)\n", + rows, dim); + return 1; + } + + // The average kernel time is divided by the repeat count. + if (repeat <= 0) { + fprintf(stderr, "Error: repeat count must be positive (got %d)\n", repeat); + return 1; + } + + // Assert the rounding rules directly. The end-to-end comparison cannot catch + // a tie-policy bug -- a tie sits exactly on a rounding boundary, where + // values_match() grants slack by design -- so this is what guards it. + if (!self_test_rounding()) return 1; + + const InDtype in_dtype = (argc > 4) ? parse_in_dtype(argv[4]) : IN_FP32; + if (in_dtype == IN_INVALID) { + fprintf(stderr, "Error: unknown input dtype '%s' " + "(valid: fp32 fp16 bf16 fp8 mxfp8 fp4 mxfp4)\n", argv[4]); + return 1; + } + + const OutQuant out_quant = (argc > 5) ? parse_out_quant(argv[5]) : OUT_NONE; + if (out_quant == OUT_INVALID) { + fprintf(stderr, "Error: unknown output quantization '%s' " + "(valid: none fp8 mxfp4)\n", argv[5]); + return 1; + } + + const BiasStore bias = + (argc > 6) ? parse_bias_store(argv[6]) : BS_NONE; + if (static_cast(bias) < static_cast(BS_NONE) || bias > BS_F32) { + fprintf(stderr, "Error: invalid bias dtype '%s' " + "(valid: none fp16 bf16 fp32)\n", argv[6]); + return 1; + } + + printf("Shape of input tensor: ( %d %d ) in=%s out=%s bias=%s\n", + rows, 2 * dim, in_dtype_name(in_dtype), out_quant_name(out_quant), + bias_store_name(bias)); + + const int64_t x_row_len = static_cast(2) * dim; + const int64_t x_nelems = static_cast(rows) * x_row_len; + const int64_t y_nelems = static_cast(rows) * dim; + const int64_t b_nelems = x_row_len; + + const size_t x_bytes = x_nelems * sizeof(float); + const size_t y_bytes = y_nelems * sizeof(float); + const size_t b_bytes = b_nelems * sizeof(float); + const size_t b_store_bytes = + static_cast(b_nelems) * + static_cast(bias_elem_bytes(bias)); + + // Bytes of the storage representation the kernel actually reads. fp4 packs + // two elements per byte. + const size_t x_store_bytes = + static_cast( + (x_nelems * in_elem_bits(in_dtype) + 7) / 8); + + // MX blocks are cut inside a row on both axes; the input row is 2*dim long. + const int64_t x_blocks_per_row = mx_blocks_per_row(x_row_len); + const int64_t y_blocks_per_row = mx_blocks_per_row(dim); + const int64_t x_nscales = + in_is_mx(in_dtype) + ? static_cast(rows) * x_blocks_per_row : 0; + const int64_t y_nscales = + out_is_mx(out_quant) + ? static_cast(rows) * y_blocks_per_row : 0; + const int64_t y_row_bytes = mxfp4_row_bytes(dim); + + // fp8 stores one byte per element; mxfp4 pads each row to whole MX blocks. + const size_t y_store_bytes = (out_quant == OUT_NONE) ? y_bytes + : (out_quant == OUT_FP8) ? static_cast(y_nelems) + : static_cast( + static_cast(rows) * y_row_bytes); + + float *X = static_cast(malloc(x_bytes)); + float *X_dq = static_cast(malloc(x_bytes)); + float *B = static_cast(malloc(b_bytes)); + float *Y = static_cast(malloc(y_bytes)); + float *Y_ref = static_cast(malloc(y_bytes)); + // values_match() needs the reference from before the output quant: once + // quantized every reference value sits exactly on a code, which says nothing + // about whether a ULP-level difference could have flipped it. + float *Y_pre = static_cast(malloc(y_bytes)); + void *X_store = malloc(x_store_bytes); + void *B_store = (bias == BS_NONE) ? nullptr : malloc(b_store_bytes); + void *Y_store = malloc(y_store_bytes); + uint8_t *x_scales = + static_cast(malloc(x_nscales ? x_nscales : 1)); + uint8_t *y_scales = + static_cast(malloc(y_nscales ? y_nscales : 1)); + uint8_t *y_scales_ref = + static_cast(malloc(y_nscales ? y_nscales : 1)); + + if (!X || !X_dq || !B || !Y || !Y_ref || !Y_pre || !X_store || + (bias != BS_NONE && !B_store) || !Y_store || !x_scales || + !y_scales || !y_scales_ref) { + fprintf(stderr, "Error: host allocation failed\n"); + return EXIT_FAILURE; + } + + std::default_random_engine generator(123); + // The range straddles the clamp limit so both clamped and unclamped + // elements are exercised. + std::uniform_real_distribution distribution(-12.f, 12.f); + + for (int64_t i = 0; i < x_nelems; i++) { + X[i] = distribution(generator); + } + for (int64_t i = 0; i < b_nelems; i++) { + B[i] = distribution(generator); + } + pack_bias(bias, B, b_nelems, B_store); + + // The exact float values the device must see once it has dequantized. + float x_scale = 1.f; + quantize_dequantize_input(in_dtype, X, x_nelems, x_row_len, X_dq, x_scales, + &x_scale); + pack_input(in_dtype, X_dq, x_nelems, x_row_len, x_blocks_per_row, x_scales, + x_scale, X_store); + + ComputeSwigluOAI(rows, dim, kSwigluAlpha, kSwigluLimit, X_dq, B, Y_ref); + memcpy(Y_pre, Y_ref, y_bytes); + + // aiter's fused per-tensor quant scales by a factor calibrated ahead of time + // rather than reducing over the whole tensor inside the kernel, so the fp8 + // output scale is derived here (from the amax quantize_dequantize_output() + // itself uses) and passed in. + float y_scale = 1.f; + if (out_quant == OUT_FP8) { + float amax = 0.f; + for (int64_t i = 0; i < y_nelems; i++) { + const float a = fabsf(Y_ref[i]); + if (a > amax) amax = a; + } + y_scale = (amax > 0.f) ? (amax / kE4M3Max) : 1.f; + } + + quantize_dequantize_output(out_quant, rows, dim, Y_ref, y_scales_ref); + +#ifdef USE_GPU + sycl::queue q(sycl::gpu_selector_v, sycl::property::queue::in_order()); +#else + sycl::queue q(sycl::cpu_selector_v, sycl::property::queue::in_order()); +#endif + + void *d_X = sycl::malloc_device(x_store_bytes, q); + void *d_B = (bias == BS_NONE) + ? nullptr : sycl::malloc_device(b_store_bytes, q); + void *d_Y = sycl::malloc_device(y_store_bytes, q); + uint8_t *d_Xscales = + sycl::malloc_device(x_nscales ? x_nscales : 1, q); + uint8_t *d_Yscales = + sycl::malloc_device(y_nscales ? y_nscales : 1, q); + + if (!d_X || (bias != BS_NONE && !d_B) || !d_Y || + !d_Xscales || !d_Yscales) { + fprintf(stderr, "Error: device allocation failed\n"); + return EXIT_FAILURE; + } + + q.memcpy(d_X, X_store, x_store_bytes); + if (d_B) q.memcpy(d_B, B_store, b_store_bytes); + if (x_nscales) q.memcpy(d_Xscales, x_scales, x_nscales); + + KernelArgs args; + args.rows = rows; + args.dim = dim; + args.alpha = kSwigluAlpha; + args.limit = kSwigluLimit; + args.X = d_X; + args.Xs = d_Xscales; + args.xScale = x_scale; + args.B = d_B; + args.Y = d_Y; + args.Ys = d_Yscales; + args.yScale = y_scale; + args.xBlocksPerRow = x_blocks_per_row; + args.yRowStrideBytes = y_row_bytes; + args.yBlocksPerRow = y_blocks_per_row; + const int block_size = 256; + const int64_t tiles_per_row = (dim + block_size - 1) / block_size; + const int64_t total_tiles = static_cast(rows) * tiles_per_row; + const int64_t max_groups = 1 << 20; + const int64_t groups = + (total_tiles < max_groups) ? total_tiles : max_groups; + sycl::range<1> gws(static_cast(groups) * block_size); + sycl::range<1> lws (block_size); + sycl::nd_range<1> launch(gws, lws); + + // check correctness before benchmarking + submit_swiglu(q, in_dtype, out_quant, bias, launch, args); + + q.memcpy(Y_store, d_Y, y_store_bytes); + if (y_nscales) q.memcpy(y_scales, d_Yscales, y_nscales); + q.wait_and_throw(); + unpack_output(out_quant, rows, dim, Y_store, y_scales, y_scale, Y); + + const float tol = 1e-3f; + + bool ok = true; + float maxdiff = 0.f; + int64_t worst = 0; + for (int64_t i = 0; i < y_nelems; i++) { + // The block scale sets the size of one quantization step for this element. + const float scale = out_is_mx(out_quant) + ? e8m0_to_float(y_scales_ref[(i / dim) * y_blocks_per_row + + (i % dim) / MX_BLOCK]) + : y_scale; + const float diff = fabsf(Y[i] - Y_ref[i]); + if (diff > maxdiff) { + maxdiff = diff; + worst = i; + } + if (!values_match(Y[i], Y_ref[i], Y_pre[i], out_quant, scale, tol)) + ok = false; + } + printf("%s\n", ok ? "PASS" : "FAIL"); + if (!ok) { + int64_t scale_mismatch = 0; + for (int64_t i = 0; i < y_nscales; ++i) + if (y_scales[i] != y_scales_ref[i]) ++scale_mismatch; + fprintf(stderr, " max |device - reference| = %g at %ld " + "(device %g, reference %g, tol %g)\n", + maxdiff, static_cast(worst), Y[worst], Y_ref[worst], tol); + if (y_nscales) + fprintf(stderr, " mismatching e8m0 output scales: %ld of %ld\n", + static_cast(scale_mismatch), + static_cast(y_nscales)); + } + + if (bias == BS_F32) { + printf("FP32 bias is accuracy-only; performance timing skipped.\n"); + } else { + q.wait_and_throw(); + auto start = std::chrono::steady_clock::now(); + for (int i = 0; i < repeat; i++) + submit_swiglu(q, in_dtype, out_quant, bias, launch, args); + q.wait_and_throw(); + auto end = std::chrono::steady_clock::now(); + auto time = std::chrono::duration_cast( + end - start).count(); + printf("Average execution time of SwiGLU (OAI) kernel: %f (us)\n", + (time * 1e-3f) / repeat); + } + + free(X); + free(X_dq); + free(B); + free(Y); + free(Y_ref); + free(Y_pre); + free(X_store); + free(B_store); + free(Y_store); + free(x_scales); + free(y_scales); + free(y_scales_ref); + sycl::free(d_X, q); + if (d_B) sycl::free(d_B, q); + sycl::free(d_Y, q); + sycl::free(d_Xscales, q); + sycl::free(d_Yscales, q); + + return 0; +} From 8d102f08a49df6f4657fd85490ed0ce62e18a148 Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Thu, 10 Sep 2026 06:24:31 -0700 Subject: [PATCH 2/2] Address review: portability of the SwiGLU (OAI) benchmarks - register swiglu-oai in the master CMake benchmark list so the four model directories are actually configured - move the CUDA `make run` case annotations out of the backslash-continued recipe, where the first comment swallowed every following run_case - build on CUDA toolkits older than 12.8 and ROCm older than 7.0 by falling back to software E2M1/E8M0 codecs instead of #error - SYCL: request a sub-group of MX_BLOCK lanes for the MXFP4 epilogue where the device offers it, keep the plain kernel for wider-only sub-groups (wave64), and skip devices that may pick a narrower one - OpenMP: iterate the fused kernel over a 64-bit element count so large shapes no longer overflow rows * dim as int Co-authored-by: Cursor --- src/CMakeLists.txt | 1 + src/swiglu-oai-cuda/Makefile | 17 +++--- src/swiglu-oai-cuda/main.cu | 110 ++++++++++++++++++++++++++++------- src/swiglu-oai-hip/main.cu | 90 +++++++++++++++++++++++----- src/swiglu-oai-omp/main.cpp | 12 ++-- src/swiglu-oai-sycl/Makefile | 7 ++- src/swiglu-oai-sycl/main.cpp | 81 ++++++++++++++++++++------ 7 files changed, 252 insertions(+), 66 deletions(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 374c7e71a..17f532a6c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -548,6 +548,7 @@ set(HECBENCH_POC_BENCHMARKS silu softmax-fused softmax-online + swiglu-oai sundials ssm storeKVCache diff --git a/src/swiglu-oai-cuda/Makefile b/src/swiglu-oai-cuda/Makefile index 0ebda35a6..1378fc637 100644 --- a/src/swiglu-oai-cuda/Makefile +++ b/src/swiglu-oai-cuda/Makefile @@ -6,8 +6,9 @@ CC = nvcc OPTIMIZE = yes DEBUG = no -# __nv_fp8_e4m3 (cuda_fp8.h) needs CUDA >= 11.8; __nv_fp4x2_e2m1 -# (cuda_fp4.h) needs CUDA >= 12.8. FP4 conversion is native on +# __nv_fp8_e4m3 (cuda_fp8.h) needs CUDA >= 11.8; the native FP4 conversions +# (cuda_fp4.h) need CUDA >= 12.8, and older toolkits fall back to the software +# E2M1/E8M0 codecs in main.cu. FP4 conversion is native on # architecture-specific Blackwell targets such as sm_100a; other targets use # CUDA's emulation path. sm_90 keeps this benchmark runnable on Hopper while # exercising native FP8 support. @@ -66,6 +67,13 @@ clean: .PHONY: clean run +# +# Standard LLM inference: BF16 activation, no bias, no output quantization. +# Bias-bearing FP16 model. +# vLLM/SGLang-style fused SwiGLU followed by FP8 activation quantization. +# Blackwell-style fused SwiGLU followed by MXFP4 activation quantization. +# Low-precision MoE path with MXFP8 input, FP8 output, and BF16 bias. +# run: $(program) @run_case() { \ output="$$($(LAUNCHER) ./$(program) "$$@" 2>&1)"; status=$$?; \ @@ -76,13 +84,8 @@ run: $(program) fi; \ }; \ set -e; \ - # Standard LLM inference: BF16 activation, no bias, no output quantization. \ run_case $(ROWS) $(DIM) $(REPEAT) bf16 none none; \ - # Bias-bearing FP16 model. \ run_case $(ROWS) $(DIM) $(REPEAT) fp16 none fp16; \ - # vLLM/SGLang-style fused SwiGLU followed by FP8 activation quantization. \ run_case $(ROWS) $(DIM) $(REPEAT) bf16 fp8 none; \ - # Blackwell-style fused SwiGLU followed by MXFP4 activation quantization. \ run_case $(ROWS) $(DIM) $(REPEAT) bf16 mxfp4 none; \ - # Low-precision MoE path with MXFP8 input, FP8 output, and BF16 bias. \ run_case $(ROWS) $(DIM) $(REPEAT) mxfp8 fp8 bf16 diff --git a/src/swiglu-oai-cuda/main.cu b/src/swiglu-oai-cuda/main.cu index bcfa980a8..2bd170b7c 100644 --- a/src/swiglu-oai-cuda/main.cu +++ b/src/swiglu-oai-cuda/main.cu @@ -11,12 +11,24 @@ #if !defined(__CUDACC_VER_MAJOR__) || !defined(__CUDACC_VER_MINOR__) #error "swiglu-oai-cuda must be compiled with the NVIDIA CUDA compiler" -#elif (__CUDACC_VER_MAJOR__ < 12) || \ - (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8) -#error "swiglu-oai-cuda requires CUDA Toolkit 12.8 or newer for FP4 support" +#elif !defined(CUDART_VERSION) || (CUDART_VERSION < 11080) +#error "swiglu-oai-cuda requires CUDA Toolkit 11.8 or newer for FP8 support" #endif +// cuda_fp4.h and __nv_cvt_float_to_e8m0() first ship with CUDA Toolkit 12.8. +// Older toolkits keep every other format and take the software E2M1/E8M0 +// codecs below, which make the same rounding decisions as reference.h. +#ifndef SWIGLU_HAS_NATIVE_FP4 +#if CUDART_VERSION >= 12080 +#define SWIGLU_HAS_NATIVE_FP4 1 +#else +#define SWIGLU_HAS_NATIVE_FP4 0 +#endif +#endif + +#if SWIGLU_HAS_NATIVE_FP4 #include +#endif #include "reference.h" #define CUDA_CHECK(call) \ @@ -39,22 +51,80 @@ static_assert(sizeof(size_t) >= 8, // Native low-precision helpers, shared by host and device // --------------------------------------------------------------------------- +#if !SWIGLU_HAS_NATIVE_FP4 +// Software E2M1 codec for toolkits without cuda_fp4.h. Representable +// magnitudes are 0, 0.5, 1, 1.5, 2, 3, 4, 6; the nearest-value search with its +// tie-to-even-index rule mirrors round_to_e2m1() in reference.h, including the +// strict `> 5.f` that lets the 4/6 tie land on 4. +__host__ __device__ __forceinline__ +uint8_t soft_e2m1_encode(const float x) +{ + const float lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + const uint8_t sign = (x < 0.f) ? 0x8u : 0x0u; + const float a = fabsf(x); + if (a > 5.f) return sign | 7u; + int best = 0; + float bestd = fabsf(a - lut[0]); + for (int i = 1; i < 8; i++) { + const float d = fabsf(a - lut[i]); + if (d < bestd || (d == bestd && (i % 2) == 0)) { + bestd = d; + best = i; + } + } + return sign | static_cast(best); +} + __host__ __device__ __forceinline__ -__nv_fp4x2_e2m1 fp4x2_pack(const float lo, const float hi) +float soft_e2m1_decode(const uint8_t nibble) { - __nv_fp4x2_e2m1 packed; - packed.__x = __nv_cvt_float2_to_fp4x2( + const float lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + const float m = lut[nibble & 0x7u]; + return (nibble & 0x8u) ? -m : m; +} +#endif + +// Two E2M1 values packed into one byte, low nibble first. Raw byte storage +// keeps the host packing independent of the toolkit's FP4 type. +__host__ __device__ __forceinline__ +uint8_t fp4x2_pack(const float lo, const float hi) +{ +#if SWIGLU_HAS_NATIVE_FP4 + return __nv_cvt_float2_to_fp4x2( make_float2(lo, hi), __NV_E2M1, cudaRoundNearest); - return packed; +#else + return static_cast(soft_e2m1_encode(lo) | + (soft_e2m1_encode(hi) << 4)); +#endif } __host__ __device__ __forceinline__ -float fp4x2_unpack(const __nv_fp4x2_storage_t byte, const int hi) +float fp4x2_unpack(const uint8_t byte, const int hi) { +#if SWIGLU_HAS_NATIVE_FP4 const __half2_raw h2 = __nv_cvt_fp4x2_to_halfraw2(byte, __NV_E2M1); __half_raw h; h.x = hi ? h2.y : h2.x; return __half2float(__half(h)); +#else + return soft_e2m1_decode(hi ? (byte >> 4) : (byte & 0xfu)); +#endif +} + +// Round-toward-zero E8M0 encode of amax / 2^kE2M1MaxExp, which is the MXFP4 +// block scale. For a positive normal float the biased IEEE-754 exponent byte +// minus kE2M1MaxExp is exactly the wanted byte; subnormals clamp to 0. +__device__ __forceinline__ +uint8_t mxfp4_scale_byte(const float amax) +{ + if (!(amax > 0.f)) return 127; // E8M0 has no zero; 2^0 for an empty block +#if SWIGLU_HAS_NATIVE_FP4 + return __nv_cvt_float_to_e8m0(amax * 0.25f, __NV_SATFINITE, cudaRoundZero); +#else + const int s = static_cast((__float_as_uint(amax) >> 23) & 0xffu) - + kE2M1MaxExp; + return s < 0 ? 0 : s > 254 ? 254 : static_cast(s); +#endif } // FP32 reserves exponent 0 for zero/subnormals while E8M0 treats exponent 0 as 2^{-127} @@ -98,7 +168,7 @@ __device__ __forceinline__ float load_x_raw(const void* __restrict__ X, } else if constexpr (STORE == XS_FP8) { return static_cast(static_cast(X)[i]); } else { - return fp4x2_unpack(static_cast(X)[i >> 1].__x, + return fp4x2_unpack(static_cast(X)[i >> 1], static_cast(i & 1)); } } @@ -116,13 +186,16 @@ __device__ __forceinline__ float2 load_x_pair( { float2 pair; if constexpr (STORE == XS_FP4) { - const __nv_fp4x2_storage_t packed = - static_cast(X)[i >> 1].__x; + const uint8_t packed = static_cast(X)[i >> 1]; +#if SWIGLU_HAS_NATIVE_FP4 const __half2_raw h2 = __nv_cvt_fp4x2_to_halfraw2(packed, __NV_E2M1); __half_raw lo, hi; lo.x = h2.x; hi.x = h2.y; pair = make_float2(__half2float(__half(lo)), __half2float(__half(hi))); +#else + pair = make_float2(fp4x2_unpack(packed, 0), fp4x2_unpack(packed, 1)); +#endif } else if constexpr (STORE == XS_FP8) { const __nv_fp8x2_storage_t packed = static_cast(X)[i >> 1].__x; @@ -320,10 +393,7 @@ __global__ void swiglu_oai_kernel(const KernelArgs a) // chooses the block scale. E8M0 round-toward-zero is floor(log2) for // positive values, exactly matching reference.h's scale rule. E8M0 has // no zero, so retain byte 127 (scale 1) for an all-zero/NaN block. - const uint8_t sb = (amax > 0.f) - ? __nv_cvt_float_to_e8m0( - amax * 0.25f, __NV_SATFINITE, cudaRoundZero) - : 127; + const uint8_t sb = mxfp4_scale_byte(amax); const float s = e8m0_decode(sb); // Lane 0 of the block owns the scale byte. Its column is in range @@ -338,7 +408,7 @@ __global__ void swiglu_oai_kernel(const KernelArgs a) if (active && (d & 1) == 0) { const float hi = (d + 1 < static_cast(a.dim)) ? paired : 0.f; - static_cast<__nv_fp4x2_e2m1*>(a.Y)[row * a.yRowStrideBytes + (d >> 1)] = + static_cast(a.Y)[row * a.yRowStrideBytes + (d >> 1)] = fp4x2_pack(y / s, hi / s); } } @@ -435,7 +505,7 @@ static void pack_input(InDtype dtype, const float* xq, size_t n, if (store == XS_FP4) { // rowLen is 2*dim and therefore even, so a packed pair never straddles a // row, and 2j/2j+1 never straddle an MX block either: one scale for both. - __nv_fp4x2_e2m1* p = static_cast<__nv_fp4x2_e2m1*>(dst); + uint8_t* p = static_cast(dst); for (size_t j = 0; j < n / 2; j++) { const size_t i = 2 * j; const float s = scale_at(i); @@ -515,16 +585,16 @@ static void unpack_output(OutQuant q, int rows, int dim, size_t blocksPerRow, return; } - const __nv_fp4x2_e2m1* p = static_cast(src); + const uint8_t* p = static_cast(src); const size_t rowBytes = mxfp4_row_bytes(dim); for (size_t i = 0; i < static_cast(rows); i++) { for (size_t b = 0; b < blocksPerRow; b++) { const float s = e8m0_decode(yscales[i * blocksPerRow + b]); const int d0 = static_cast(b * MX_BLOCK); const int nk = (d0 + MX_BLOCK < dim) ? MX_BLOCK : (dim - d0); - const __nv_fp4x2_e2m1* blk = p + i * rowBytes + b * (MX_BLOCK / 2); + const uint8_t* blk = p + i * rowBytes + b * (MX_BLOCK / 2); for (int k = 0; k < nk; k++) { - Y[i * dim + d0 + k] = fp4x2_unpack(blk[k / 2].__x, k & 1) * s; + Y[i * dim + d0 + k] = fp4x2_unpack(blk[k / 2], k & 1) * s; } } } diff --git a/src/swiglu-oai-hip/main.cu b/src/swiglu-oai-hip/main.cu index 99f9b7508..fab7a7b0d 100644 --- a/src/swiglu-oai-hip/main.cu +++ b/src/swiglu-oai-hip/main.cu @@ -25,9 +25,20 @@ #if !defined(__HIPCC__) #error "swiglu-oai-hip must be compiled with hipcc" -#elif HIP_VERSION_MAJOR < 7 -#error "swiglu-oai-hip requires ROCm 7.0 or newer for hip_fp4.h" +#endif + +#ifndef SWIGLU_HAS_HIP_FP4 +// hip_fp4.h first ships with ROCm 7.0. Older installations keep every other +// format and take the software E2M1 codec below, which makes the same +// rounding decisions as reference.h. +#if HIP_VERSION_MAJOR >= 7 +#define SWIGLU_HAS_HIP_FP4 1 #else +#define SWIGLU_HAS_HIP_FP4 0 +#endif +#endif + +#if SWIGLU_HAS_HIP_FP4 #include #endif @@ -145,7 +156,53 @@ __device__ __forceinline__ uint8_t device_e4m3_encode(float x) { #endif } +#if SWIGLU_HAS_HIP_FP4 static_assert(sizeof(__hip_fp4x2_e2m1) == 1, "two fp4 values must pack into one byte"); +#endif + +// --------------------------------------------------------------------------- +// E2M1 codec. Representable magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. The +// software path is used when hip_fp4.h is unavailable; its nearest-value +// search with a tie-to-even-index rule mirrors round_to_e2m1() in +// reference.h, including the strict `> 5.f` that lets the 4/6 tie land on 4. +// --------------------------------------------------------------------------- + +__host__ __device__ __forceinline__ float e2m1_decode(uint8_t nibble) { +#if SWIGLU_HAS_HIP_FP4 + __hip_fp4_e2m1 v; + v.__x = static_cast<__hip_fp4_storage_t>(nibble & 0xfu); + return static_cast(v); +#else + const float lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + const float m = lut[nibble & 0x7u]; + return (nibble & 0x8u) ? -m : m; +#endif +} + +// Two E2M1 values packed into one byte, low nibble first. +__host__ __device__ __forceinline__ uint8_t e2m1x2_encode(float lo, float hi) { +#if SWIGLU_HAS_HIP_FP4 + return __hip_fp4x2_e2m1(make_float2(lo, hi)).__x; +#else + auto encode = [](float x) -> uint8_t { + const float lut[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + const uint8_t sign = (x < 0.f) ? 0x8u : 0x0u; + const float a = fabsf(x); + if (a > 5.f) return sign | 7u; + int best = 0; + float bestd = fabsf(a - lut[0]); + for (int i = 1; i < 8; i++) { + const float d = fabsf(a - lut[i]); + if (d < bestd || (d == bestd && (i % 2) == 0)) { + bestd = d; + best = i; + } + } + return sign | static_cast(best); + }; + return static_cast(encode(lo) | (encode(hi) << 4)); +#endif +} // --------------------------------------------------------------------------- // Device-side helpers @@ -182,11 +239,8 @@ __device__ __forceinline__ float load_x_raw( } else if constexpr (STORE == XS_FP8) { return device_e4m3_decode(static_cast(X)[idx]); } else { - __hip_fp4_e2m1 v; - v.__x = static_cast<__hip_fp4_storage_t>( - (static_cast(X)[idx >> 1] >> - ((idx & 1) << 2)) & 0xf); - return static_cast(v); + return e2m1_decode(static_cast( + (static_cast(X)[idx >> 1] >> ((idx & 1) << 2)) & 0xf)); } } @@ -199,9 +253,15 @@ __device__ __forceinline__ float2 load_x_pair( const float tensorScale, const int64_t idx, const int64_t scaleIdx) { float2 pair; if constexpr (STORE == XS_FP4) { - __hip_fp4x2_e2m1 packed; - packed.__x = static_cast(X)[idx >> 1]; - pair = static_cast(packed); + const uint8_t packed = static_cast(X)[idx >> 1]; +#if SWIGLU_HAS_HIP_FP4 + __hip_fp4x2_e2m1 native; + native.__x = packed; + pair = static_cast(native); +#else + pair = make_float2(e2m1_decode(packed & 0xfu), + e2m1_decode(packed >> 4)); +#endif } else if constexpr (STORE == XS_FP8) { #if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) __hip_fp8x2_e4m3_fnuz packed; @@ -336,7 +396,7 @@ __global__ void swiglu_oai_kernel(const KernelArgs a) { (d + 1 < static_cast(a.dim)) ? paired : 0.f; static_cast(a.Y)[ row * a.yRowStrideBytes + (d >> 1)] = - __hip_fp4x2_e2m1(make_float2(y / s, hi / s)).__x; + e2m1x2_encode(y / s, hi / s); } } } @@ -432,8 +492,7 @@ static void pack_input(HipFp8Format fp8_format, InDtype dt, const float* xq, for (int64_t i = 0; i < n; i += 2) { const float scale = scale_at(i); const float hi = (i + 1 < n) ? xq[i + 1] / scale : 0.f; - p[i >> 1] = - __hip_fp4x2_e2m1(make_float2(xq[i] / scale, hi)).__x; + p[i >> 1] = e2m1x2_encode(xq[i] / scale, hi); } return; } @@ -520,11 +579,10 @@ static void unpack_output(HipFp8Format fp8_format, OutQuant q, const uint8_t* p = static_cast(src); for (int64_t r = 0; r < rows; r++) { for (int64_t d = 0; d < dim; d++) { - __hip_fp4_e2m1 v; - v.__x = static_cast<__hip_fp4_storage_t>( + const uint8_t nibble = static_cast( (p[r * rowStride + (d >> 1)] >> ((d & 1) << 2)) & 0xf); dst[r * dim + d] = - static_cast(v) * + e2m1_decode(nibble) * e8m0_to_float(scales[r * blocksPerRow + (d / MX_BLOCK)]); } } diff --git a/src/swiglu-oai-omp/main.cpp b/src/swiglu-oai-omp/main.cpp index 10afbaba6..2688fbd54 100644 --- a/src/swiglu-oai-omp/main.cpp +++ b/src/swiglu-oai-omp/main.cpp @@ -18,12 +18,14 @@ void swiglu_oai_kernel( const float* Bdata, float* Ydata) { + const int64_t nelems = (int64_t)rows * dim; + #pragma omp target teams distribute parallel for num_threads(block_size) - for (int index = 0; index < rows * dim; index++) { - const int i = index / dim; - const int d = index % dim; + for (int64_t index = 0; index < nelems; index++) { + const int64_t i = index / dim; + const int64_t d = index % dim; - const int64_t xOffset = (int64_t)i * 2 * dim; + const int64_t xOffset = i * 2 * dim; float gate = Xdata[xOffset + 2 * d] + Bdata[2 * d]; float linear = Xdata[xOffset + 2 * d + 1] + Bdata[2 * d + 1]; @@ -32,7 +34,7 @@ void swiglu_oai_kernel( linear = fminf(fmaxf(linear, -limit), limit); const float s = 1.f / (1.f + expf(-alpha * gate)); - Ydata[(int64_t)i * dim + d] = gate * s * (linear + 1.f); + Ydata[i * dim + d] = gate * s * (linear + 1.f); } } diff --git a/src/swiglu-oai-sycl/Makefile b/src/swiglu-oai-sycl/Makefile index 8f8d697f2..106aefc57 100644 --- a/src/swiglu-oai-sycl/Makefile +++ b/src/swiglu-oai-sycl/Makefile @@ -57,6 +57,10 @@ endif ifeq ($(HIP), yes) CFLAGS += -fsycl-targets=amdgcn-amd-amdhsa \ -Xsycl-target-backend --offload-arch=$(HIP_ARCH) + # The mxfp4 kernel is also instantiated with a required sub-group size of + # 32, which amdgcn ignores. main() only selects that instantiation on + # devices that report the size as supported. + CFLAGS += -Wno-incorrect-sub-group-size endif # Debug Flags @@ -90,7 +94,8 @@ run: $(program) output="$$($(LAUNCHER) ./$(program) "$$@" 2>&1)"; status=$$?; \ printf '%s\n' "$$output"; \ if [ $$status -ne 0 ] || \ - ! printf '%s\n' "$$output" | awk '$$0 == "PASS" { pass=1 } END { exit !pass }'; then \ + ! printf '%s\n' "$$output" | \ + awk '$$0 == "PASS" || /^Skipped: / { ok=1 } END { exit !ok }'; then \ return 1; \ fi; \ }; \ diff --git a/src/swiglu-oai-sycl/main.cpp b/src/swiglu-oai-sycl/main.cpp index bf713afba..3082cf5e8 100644 --- a/src/swiglu-oai-sycl/main.cpp +++ b/src/swiglu-oai-sycl/main.cpp @@ -8,8 +8,10 @@ #include #include #include +#include #include #include +#include #include #include "reference.h" @@ -248,10 +250,30 @@ void swiglu_oai_kernel(sycl::nd_item<1> item, const KernelArgs a) template class SwigluOAIKernel; +template +class SwigluOAIKernelSG32; + +// `force_sg32` asks for a sub-group of exactly MX_BLOCK lanes. The MXFP4 +// epilogue needs a sub-group of at least that width: the XOR masks stay below +// MX_BLOCK, so on a wider sub-group each aligned MX_BLOCK-lane slice reduces +// its own block and elects its own lane 0, but a narrower one would spread a +// block over several sub-groups that each write a different scale for it. +// Devices offering MX_BLOCK take the annotated kernel; devices whose sizes are +// all wider (wave64, say) take the plain one, and main() rejects the rest. template static void submit_swiglu_kernel(sycl::queue &q, sycl::nd_range<1> launch, - const KernelArgs &a) { + const KernelArgs &a, bool force_sg32) { q.submit([&](sycl::handler &cgh) { + if constexpr (OUT == OUT_MXFP4) { + if (force_sg32) { + cgh.parallel_for>( + launch, [=](sycl::nd_item<1> item) + [[sycl::reqd_sub_group_size(MX_BLOCK)]] { + swiglu_oai_kernel(item, a); + }); + return; + } + } cgh.parallel_for>( launch, [=](sycl::nd_item<1> item) { swiglu_oai_kernel(item, a); @@ -261,17 +283,18 @@ static void submit_swiglu_kernel(sycl::queue &q, sycl::nd_range<1> launch, template static void submit_swiglu_out(sycl::queue &q, OutQuant out, - sycl::nd_range<1> launch, const KernelArgs &a) { + sycl::nd_range<1> launch, const KernelArgs &a, + bool force_sg32) { switch (out) { case OUT_NONE: submit_swiglu_kernel( - q, launch, a); break; + q, launch, a, force_sg32); break; case OUT_FP8: submit_swiglu_kernel( - q, launch, a); break; + q, launch, a, force_sg32); break; case OUT_MXFP4: submit_swiglu_kernel( - q, launch, a); break; + q, launch, a, force_sg32); break; default: break; } } @@ -279,29 +302,29 @@ static void submit_swiglu_out(sycl::queue &q, OutQuant out, template static void submit_swiglu_input(sycl::queue &q, InDtype in, OutQuant out, sycl::nd_range<1> launch, - const KernelArgs &a) { + const KernelArgs &a, bool force_sg32) { switch (in) { case IN_FP32: submit_swiglu_out( - q, out, launch, a); break; + q, out, launch, a, force_sg32); break; case IN_FP16: submit_swiglu_out( - q, out, launch, a); break; + q, out, launch, a, force_sg32); break; case IN_BF16: submit_swiglu_out( - q, out, launch, a); break; + q, out, launch, a, force_sg32); break; case IN_FP8: submit_swiglu_out( - q, out, launch, a); break; + q, out, launch, a, force_sg32); break; case IN_MXFP8: submit_swiglu_out( - q, out, launch, a); break; + q, out, launch, a, force_sg32); break; case IN_FP4: submit_swiglu_out( - q, out, launch, a); break; + q, out, launch, a, force_sg32); break; case IN_MXFP4: submit_swiglu_out( - q, out, launch, a); break; + q, out, launch, a, force_sg32); break; default: break; } } @@ -309,9 +332,9 @@ static void submit_swiglu_input(sycl::queue &q, InDtype in, OutQuant out, template static void submit_swiglu(sycl::queue &q, InDtype in, OutQuant out, BiasStore bias, sycl::nd_range<1> launch, - const KernelArgs &a) { + const KernelArgs &a, bool force_sg32) { #define SUBMIT_BIAS(B) \ - submit_swiglu_input(q, in, out, launch, a) + submit_swiglu_input(q, in, out, launch, a, force_sg32) switch (bias) { case BS_NONE: SUBMIT_BIAS(BS_NONE); break; case BS_F16: SUBMIT_BIAS(BS_F16); break; @@ -666,6 +689,28 @@ int main(int argc, char* argv[]) sycl::queue q(sycl::cpu_selector_v, sycl::property::queue::in_order()); #endif + // The MXFP4 epilogue reduces an MX block within one sub-group, so every + // sub-group the device may hand the kernel has to be at least MX_BLOCK + // lanes wide. When MX_BLOCK itself is available the kernel asks for exactly + // that; otherwise the device only offers wider sub-groups and the plain + // kernel is safe. A device that can pick a narrower size cannot run the + // path correctly, so it is skipped rather than silently mis-quantized. + bool force_sg32 = false; + if (out_quant == OUT_MXFP4) { + const std::vector sg_sizes = + q.get_device().get_info(); + force_sg32 = std::find(sg_sizes.begin(), sg_sizes.end(), + static_cast(MX_BLOCK)) != sg_sizes.end(); + const size_t smallest = sg_sizes.empty() + ? 0 : *std::min_element(sg_sizes.begin(), sg_sizes.end()); + if (!force_sg32 && smallest < static_cast(MX_BLOCK)) { + printf("Skipped: mxfp4 output needs a sub-group of at least %d lanes, " + "which %s does not guarantee\n", MX_BLOCK, + q.get_device().get_info().c_str()); + return EXIT_SUCCESS; + } + } + void *d_X = sycl::malloc_device(x_store_bytes, q); void *d_B = (bias == BS_NONE) ? nullptr : sycl::malloc_device(b_store_bytes, q); @@ -711,7 +756,8 @@ int main(int argc, char* argv[]) sycl::nd_range<1> launch(gws, lws); // check correctness before benchmarking - submit_swiglu(q, in_dtype, out_quant, bias, launch, args); + submit_swiglu(q, in_dtype, out_quant, bias, launch, args, + force_sg32); q.memcpy(Y_store, d_Y, y_store_bytes); if (y_nscales) q.memcpy(y_scales, d_Yscales, y_nscales); @@ -757,7 +803,8 @@ int main(int argc, char* argv[]) q.wait_and_throw(); auto start = std::chrono::steady_clock::now(); for (int i = 0; i < repeat; i++) - submit_swiglu(q, in_dtype, out_quant, bias, launch, args); + submit_swiglu(q, in_dtype, out_quant, bias, launch, args, + force_sg32); q.wait_and_throw(); auto end = std::chrono::steady_clock::now(); auto time = std::chrono::duration_cast(