diff --git a/include/xe-fuse/builder/epilogue_builder.hpp b/include/xe-fuse/builder/epilogue_builder.hpp index 9cfa374..438bc5e 100644 --- a/include/xe-fuse/builder/epilogue_builder.hpp +++ b/include/xe-fuse/builder/epilogue_builder.hpp @@ -137,6 +137,29 @@ template using AddResidual = Add, ElementCompute, ElementCompute>; +// GateAcc: gate[m] * acc — per-token gate (ColBroadcast) applied to accumulator. +// gate[m] is a scalar per output row; used before residual add in DiT blocks. +template +using GateAcc = Mul, + Acc, + ElementCompute, ElementCompute>; + +// GateResidualGamma: gamma[n] * (gate[m] * acc + residual) — K0g pattern. +// Extends K0a with a per-token gate that modulates the GEMM output before +// the residual add. Used in FLUX.2-style DiT single-block transformer blocks. +template +using GateResidualGamma = ScaleCols< + Add, + AuxLoad, + ElementCompute, ElementCompute>, + TileShape, ElementGamma, ElementCompute>; + // ============================================================ // Pairwise Operations — lane-shuffle-based pair computations // ============================================================ @@ -279,6 +302,31 @@ template using DequantGeGLU = GeGLU>; +// DequantFP8: float_acc * scale_a[m] * scale_b[n] → bf16 +// FP8×FP8 GEMM dequantization. Identical EVT structure to DequantW8A8 but +// ElementAcc is float (FP8 upcasts to FP16 before XMX, accumulates in float). +// scale_a[m] = per-token input scale, scale_b[n] = per-channel weight scale. +template +using DequantFP8 = Mul< + Mul, + ElementCompute, ElementCompute>, + RowBroadcast<0, TileShape, ElementScale, ElementCompute>, + ElementCompute, ElementCompute>; + +// DequantFP8SwiGLU: SwiGLU( float_acc * scale_a[m] * scale_b[n] ) → bf16 +// FP8 FFN kernel (K2_FP8): gate+up projection with dequant and SwiGLU fused. +template +using DequantFP8SwiGLU = SwiGLU>; + +// DequantFP8GeGLU: GeGLU( float_acc * scale_a[m] * scale_b[n] ) → bf16 +// FP8 FFN kernel for Gemma-style models. +template +using DequantFP8GeGLU = GeGLU>; + // HadamardOutput: apply WHT to the output of InnerEVT // // Used as the final epilogue step in K0_W8A8 (O-projection) for QuaRot: diff --git a/include/xe-fuse/kernels/compute_rstd.hpp b/include/xe-fuse/kernels/compute_rstd.hpp index 312341a..e5649ee 100644 --- a/include/xe-fuse/kernels/compute_rstd.hpp +++ b/include/xe-fuse/kernels/compute_rstd.hpp @@ -3,6 +3,8 @@ #include #include +#include "cutlass/bfloat16.h" + namespace xe_fuse { // Standalone rstd reduction kernel. @@ -124,4 +126,72 @@ void launch_compute_rstd_and_quantize( }); } +// Dual-output RMSNorm + INT8 quantization kernel. +// +// Same three-pass algorithm as launch_compute_rstd_and_quantize but also +// writes a BF16 normed output for the residual path. Use when both a normed +// BF16 value (residual path) and an INT8 quantized value (next GEMM input) +// are needed from the same input. +template +void launch_norm_quantize_dual( + sycl::queue& q, + ElementInput const* input_ptr, + int8_t* quant_out_ptr, + ElementNormed* normed_out_ptr, + float* scale_token_ptr, + int M, int N, int L, + float eps = 1e-6f) +{ + constexpr int SG_SIZE = 16; + int work_groups = M * L; + + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(static_cast(work_groups) * SG_SIZE, SG_SIZE), + [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SG_SIZE)]] { + int row = item.get_group(0); + int lane = item.get_local_id(0); + + // ── Pass 1: sum_sq → rstd ────────────────────────────────────────── + float sum_sq = 0.f; + for (int col = lane; col < N; col += SG_SIZE) { + float v = static_cast(input_ptr[row * N + col]); + sum_sq += v * v; + } + auto sg = item.get_sub_group(); + for (int off = SG_SIZE / 2; off > 0; off /= 2) + sum_sq += sycl::shift_group_left(sg, sum_sq, off); + + float rstd = sycl::rsqrt(sum_sq / static_cast(N) + eps); + rstd = sycl::group_broadcast(sg, rstd, 0); + + // ── Pass 2: max_abs of normalized values ────────────────────────── + float max_abs = 0.f; + for (int col = lane; col < N; col += SG_SIZE) { + float normed = static_cast(input_ptr[row * N + col]) * rstd; + max_abs = sycl::fmax(max_abs, sycl::fabs(normed)); + } + for (int off = SG_SIZE / 2; off > 0; off /= 2) + max_abs = sycl::fmax(max_abs, sycl::shift_group_left(sg, max_abs, off)); + + float scale_tok = max_abs / 127.f + 1e-8f; + scale_tok = sycl::group_broadcast(sg, scale_tok, 0); + + // ── Pass 3: write INT8 quantized + BF16 normed ──────────────────── + for (int col = lane; col < N; col += SG_SIZE) { + float normed = static_cast(input_ptr[row * N + col]) * rstd; + normed_out_ptr[row * N + col] = static_cast(normed); + + float qval = sycl::round(normed / scale_tok); + qval = sycl::fmin(sycl::fmax(qval, -128.f), 127.f); + quant_out_ptr[row * N + col] = static_cast(qval); + } + + if (lane == 0) + scale_token_ptr[row] = scale_tok; + } + ); + }); +} + } // namespace xe_fuse diff --git a/include/xe-fuse/kernels/gemm_fp8_dequant.hpp b/include/xe-fuse/kernels/gemm_fp8_dequant.hpp new file mode 100644 index 0000000..edca564 --- /dev/null +++ b/include/xe-fuse/kernels/gemm_fp8_dequant.hpp @@ -0,0 +1,279 @@ +#pragma once + +// FP8 GEMM epilogue fusion kernels for Intel Xe / BMG-G31. +// +// BMG FP8 implementation note: +// BMG-G31 has no native FP8 XMX instruction. The mainloop upcasts +// float_e4m3_t/float_e5m2_t to FP16 before the XMX16 MMA, which +// accumulates in float. Requires a VNNI layout workaround for 8-bit loads +// (sycl-tla issue #357). +// +// Scale convention: +// D[m,n] = epilogue( acc[m,n] * scale_a[m] * scale_b[n] ) +// scale_a[m]: per-token input scale (M values, ColBroadcast) +// scale_b[n]: per-channel weight scale (N values, RowBroadcast) +// For per-tensor quant: pass uniform arrays (all M copies of scale_a, etc.) +// +// Kernels provided: +// GemmFP8Dequant — FP8×FP8 → dequant → BF16 (K0_FP8) +// GemmFP8DequantSwiGLU — FP8×FP8 → dequant → SwiGLU → BF16 (K2_FP8) + +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/xe_epilogue.hpp" +#include "cutlass/epilogue/fusion/xe_callbacks.hpp" +#include "cutlass/gemm/collective/collective_mma.hpp" +#include "cutlass/gemm/device/gemm_universal.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/float8.h" +#include "cutlass/fp8_to_fp16.h" + +#include +#include + +#include "xe-fuse/visitors/xe_pairwise_compute.hpp" + +namespace xe_fuse { + +using namespace cute; + +// ───────────────────────────────────────────────────────────────────────────── +// GemmFP8Dequant — K0_FP8 +// D[m,n] = bf16( acc[m,n] * scale_a[m] * scale_b[n] ) +// ───────────────────────────────────────────────────────────────────────────── +template < + typename ElementA_ = cutlass::float_e4m3_t, + typename ElementB_ = cutlass::float_e4m3_t, + typename ElementD_ = cutlass::bfloat16_t, + typename ElementScale_ = float, + typename ElementAcc_ = float, + typename ElementCompute_ = float, + typename TileShape_ = cute::Shape +> +struct GemmFP8Dequant { + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementD = ElementD_; + using ElementScale = ElementScale_; + using ElementAcc = ElementAcc_; + using ElementCompute = ElementCompute_; + using TileShape = TileShape_; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::RowMajor; + + using StrideC = cute::Stride, int64_t>; + using StrideD = cute::Stride, int64_t>; + + static constexpr int AlignmentAB = 32; // FP8 = 1 byte; 32 elements = 32 B + static constexpr int AlignmentCD = 8; // BF16 = 2 bytes; 8 elements = 16 B + static constexpr int PipelineStages = 2; + + // ── MMA atom: FP16 MMA with float accumulator (FP8 upcasted to FP16) ──── + using TiledMma = + typename TiledMMAHelper, + Layout, + Layout, Stride<_4,_1,_0>>>::TiledMMA; + + // ── Mainloop: same dispatch as W8A8, FP8 types instead of INT8 ────────── + using GEMMDispatchPolicy = cutlass::gemm::MainloopIntelW8A8; + + using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma< + GEMMDispatchPolicy, TileShape, + ElementA, cutlass::gemm::TagToStrideA_t, + ElementB, cutlass::gemm::TagToStrideB_t, + TiledMma, + XE_2D_U8x32x32_LD_N, void, void, cute::identity, + XE_2D_U8x32x32_LD_V, void, void, cute::identity + >; + + // ── EVT: acc[float] * scale_a[m] * scale_b[n] → bf16 ─────────────────── + using Accum = cutlass::epilogue::fusion::XeAccFetch; + + using TokenScaleBroadcast = cutlass::epilogue::fusion::XeColBroadcast< + 0, TileShape, ElementScale, ElementCompute, + cute::Stride, cute::Int<0>, int64_t>, + 128 / cutlass::sizeof_bits_v + >; + + using MulCompute = cutlass::epilogue::fusion::XeCompute< + cutlass::multiplies, ElementCompute, ElementCompute, + cutlass::FloatRoundStyle::round_to_nearest + >; + + using InnerDequant = cutlass::epilogue::fusion::XeEVT; + + using ChannelScaleBroadcast = cutlass::epilogue::fusion::XeRowBroadcast< + 0, TileShape, ElementScale, ElementCompute, + cute::Stride, cute::Int<1>, int64_t>, + 128 / cutlass::sizeof_bits_v + >; + + using EVT = cutlass::epilogue::fusion::XeEVT; + + // ── Epilogue via CollectiveBuilder ─────────────────────────────────────── + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + TileShape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAcc, ElementCompute, + ElementD, StrideC, AlignmentCD, + ElementD, StrideD, AlignmentCD, + cutlass::epilogue::collective::EpilogueScheduleAuto, + EVT + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, CollectiveMainloop, CollectiveEpilogue>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + static typename EVT::Arguments make_evt_args( + ElementScale const* scale_a_ptr, int M, + ElementScale const* scale_b_ptr, int N) { + + typename Accum::Arguments accum_args{}; + + typename TokenScaleBroadcast::Arguments token_scale_args; + token_scale_args.ptr_col = scale_a_ptr; + token_scale_args.null_default = ElementScale(1); + token_scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)}; + + typename MulCompute::Arguments inner_mul_args{}; + typename InnerDequant::Arguments inner_args{accum_args, token_scale_args, inner_mul_args}; + + typename ChannelScaleBroadcast::Arguments channel_scale_args; + channel_scale_args.ptr_row = scale_b_ptr; + channel_scale_args.null_default = ElementScale(1); + channel_scale_args.dRow = {cute::Int<0>{}, cute::Int<1>{}, static_cast(N)}; + + typename MulCompute::Arguments outer_mul_args{}; + return {inner_args, channel_scale_args, outer_mul_args}; + } +}; + +// ───────────────────────────────────────────────────────────────────────────── +// GemmFP8DequantSwiGLU — K2_FP8 +// D[m,n] = SwiGLU( acc[m,n] * scale_a[m] * scale_b[n] ) +// For FFN gate+up projections; output is N columns wide — each pair (2i, 2i+1) +// carries the same silu(gate)*up value; caller consumes only N/2 (even columns). +// ───────────────────────────────────────────────────────────────────────────── +template < + typename ElementA_ = cutlass::float_e4m3_t, + typename ElementB_ = cutlass::float_e4m3_t, + typename ElementD_ = cutlass::bfloat16_t, + typename ElementScale_ = float, + typename ElementAcc_ = float, + typename ElementCompute_ = float, + typename TileShape_ = cute::Shape +> +struct GemmFP8DequantSwiGLU { + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementD = ElementD_; + using ElementScale = ElementScale_; + using ElementAcc = ElementAcc_; + using ElementCompute = ElementCompute_; + using TileShape = TileShape_; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::RowMajor; + + using StrideC = cute::Stride, int64_t>; + using StrideD = cute::Stride, int64_t>; + + static constexpr int AlignmentAB = 32; + static constexpr int AlignmentCD = 8; + static constexpr int PipelineStages = 2; + + using TiledMma = + typename TiledMMAHelper, + Layout, + Layout, Stride<_4,_1,_0>>>::TiledMMA; + + using GEMMDispatchPolicy = cutlass::gemm::MainloopIntelW8A8; + + using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma< + GEMMDispatchPolicy, TileShape, + ElementA, cutlass::gemm::TagToStrideA_t, + ElementB, cutlass::gemm::TagToStrideB_t, + TiledMma, + XE_2D_U8x32x32_LD_N, void, void, cute::identity, + XE_2D_U8x32x32_LD_V, void, void, cute::identity + >; + + // ── EVT: SwiGLU( acc * scale_a[m] * scale_b[n] ) ─────────────────────── + using Accum = cutlass::epilogue::fusion::XeAccFetch; + + using TokenScaleBroadcast = cutlass::epilogue::fusion::XeColBroadcast< + 0, TileShape, ElementScale, ElementCompute, + cute::Stride, cute::Int<0>, int64_t>, + 128 / cutlass::sizeof_bits_v + >; + + using MulCompute = cutlass::epilogue::fusion::XeCompute< + cutlass::multiplies, ElementCompute, ElementCompute, + cutlass::FloatRoundStyle::round_to_nearest + >; + + using InnerDequant = cutlass::epilogue::fusion::XeEVT; + + using ChannelScaleBroadcast = cutlass::epilogue::fusion::XeRowBroadcast< + 0, TileShape, ElementScale, ElementCompute, + cute::Stride, cute::Int<1>, int64_t>, + 128 / cutlass::sizeof_bits_v + >; + + using DequantTree = cutlass::epilogue::fusion::XeEVT; + + using SwiGLUVisitor = XePairwiseCompute; + using EVT = cutlass::epilogue::fusion::XeEVT; + + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + TileShape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAcc, ElementCompute, + ElementD, StrideC, AlignmentCD, + ElementD, StrideD, AlignmentCD, + cutlass::epilogue::collective::EpilogueScheduleAuto, + EVT + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, CollectiveMainloop, CollectiveEpilogue>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + static typename EVT::Arguments make_evt_args( + ElementScale const* scale_a_ptr, int M, + ElementScale const* scale_b_ptr, int N) { + + typename Accum::Arguments accum_args{}; + + typename TokenScaleBroadcast::Arguments token_scale_args; + token_scale_args.ptr_col = scale_a_ptr; + token_scale_args.null_default = ElementScale(1); + token_scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)}; + + typename MulCompute::Arguments inner_mul_args{}; + typename InnerDequant::Arguments inner_args{accum_args, token_scale_args, inner_mul_args}; + + typename ChannelScaleBroadcast::Arguments channel_scale_args; + channel_scale_args.ptr_row = scale_b_ptr; + channel_scale_args.null_default = ElementScale(1); + channel_scale_args.dRow = {cute::Int<0>{}, cute::Int<1>{}, static_cast(N)}; + + typename MulCompute::Arguments outer_mul_args{}; + typename DequantTree::Arguments dequant_args{inner_args, channel_scale_args, outer_mul_args}; + + typename SwiGLUVisitor::Arguments swiglu_args{}; + return {dequant_args, swiglu_args}; + } +}; + +} // namespace xe_fuse diff --git a/include/xe-fuse/kernels/gemm_gate_residual_norm.hpp b/include/xe-fuse/kernels/gemm_gate_residual_norm.hpp new file mode 100644 index 0000000..f35db57 --- /dev/null +++ b/include/xe-fuse/kernels/gemm_gate_residual_norm.hpp @@ -0,0 +1,164 @@ +#pragma once + +// K0g: gemm_gate_residual_norm -- D[m,n] = gamma[n] * (gate[m] * acc[m,n] + residual[m,n]) +// +// Extends K0a (residual add + gamma) with a per-token gate scalar that +// modulates the GEMM accumulator before the residual add. Used in DiT-style +// transformer blocks (e.g. FLUX.2) where a timestep-dependent gate vector +// controls how much of the GEMM output contributes to the residual stream. +// +// EVT tree: +// XeEVT, +// ResidualLoad>> +// +// gate[m] is a per-token scalar (ColBroadcast over M rows). +// gamma[n] is the RMSNorm weight pre-folded with rstd, same convention as K0a. + +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/xe_epilogue.hpp" +#include "cutlass/epilogue/fusion/xe_callbacks.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/collective/collective_mma.hpp" + +#include + +namespace xe_fuse { + +template < + typename ElementA_ = cutlass::bfloat16_t, + typename ElementB_ = cutlass::bfloat16_t, + typename ElementD_ = cutlass::bfloat16_t, + typename ElementResidual_ = cutlass::bfloat16_t, + typename ElementGate_ = float, + typename ElementGamma_ = float, + typename ElementAcc_ = float, + typename ElementCompute_ = float, + typename TileShape_ = cute::Shape +> +struct GemmGateResidualGamma { + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementD = ElementD_; + using ElementResidual = ElementResidual_; + using ElementGate = ElementGate_; + using ElementGamma = ElementGamma_; + using ElementAcc = ElementAcc_; + using ElementCompute = ElementCompute_; + using TileShape = TileShape_; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::RowMajor; + + using StrideC = cute::Stride, int64_t>; + using StrideD = cute::Stride, int64_t>; + using StrideResidual = cute::Stride, int64_t>; + + // gate[m] x acc + using Accum = cutlass::epilogue::fusion::XeAccFetch; + + using GateBroadcast = cutlass::epilogue::fusion::XeColBroadcast< + 0, TileShape, ElementGate, ElementCompute, + cute::Stride, cute::Int<0>, int64_t>, + 128 / cutlass::sizeof_bits_v + >; + + using MulCompute = cutlass::epilogue::fusion::XeCompute< + cutlass::multiplies, ElementCompute, ElementCompute, + cutlass::FloatRoundStyle::round_to_nearest + >; + + using GatedAcc = cutlass::epilogue::fusion::XeEVT; + + // gate_acc + residual + using ResidualLoad = cutlass::epilogue::fusion::XeAuxLoad< + ElementResidual, StrideResidual, void, + 128 / cutlass::sizeof_bits_v, true, true + >; + + using AddCompute = cutlass::epilogue::fusion::XeCompute< + cutlass::plus, ElementCompute, ElementCompute, + cutlass::FloatRoundStyle::round_to_nearest + >; + + using GatedResidual = cutlass::epilogue::fusion::XeEVT; + + // gamma[n] x (gate_acc + residual) + using GammaBroadcast = cutlass::epilogue::fusion::XeRowBroadcast< + 0, TileShape, ElementGamma, ElementCompute, + cute::Stride, cute::Int<1>, int64_t>, + 128 / cutlass::sizeof_bits_v + >; + + using EVT = cutlass::epilogue::fusion::XeEVT; + + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + TileShape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAcc, ElementCompute, + ElementD, StrideC, 8, + ElementD, StrideD, 8, + cutlass::epilogue::collective::EpilogueScheduleAuto, + EVT + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 8, + ElementB, LayoutB, 8, + ElementAcc, + TileShape, + cute::Shape, + cutlass::gemm::collective::StageCountAuto, + cutlass::gemm::collective::KernelScheduleAuto + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, CollectiveMainloop, CollectiveEpilogue>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + static typename EVT::Arguments make_evt_args( + ElementGate const* gate_ptr, int M, + ElementResidual const* residual_ptr, StrideResidual stride_residual, + ElementGamma const* gamma_ptr, int N) { + + typename GateBroadcast::Arguments gate_args; + gate_args.ptr_col = gate_ptr; + gate_args.null_default = ElementGate(1); + gate_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)}; + + typename Accum::Arguments accum_args{}; + typename MulCompute::Arguments gate_mul_args{}; + typename GatedAcc::Arguments gated_acc_args{gate_args, accum_args, gate_mul_args}; + + typename ResidualLoad::Arguments residual_args; + residual_args.ptr_aux = residual_ptr; + residual_args.null_default = ElementResidual(0); + residual_args.dAux = stride_residual; + + typename AddCompute::Arguments add_args{}; + typename GatedResidual::Arguments gated_residual_args{gated_acc_args, residual_args, add_args}; + + typename GammaBroadcast::Arguments gamma_args; + gamma_args.ptr_row = gamma_ptr; + gamma_args.null_default = ElementGamma(1); + gamma_args.dRow = {cute::Int<0>{}, cute::Int<1>{}, static_cast(N)}; + + typename MulCompute::Arguments outer_mul_args{}; + + return {gamma_args, gated_residual_args, outer_mul_args}; + } +}; + +} // namespace xe_fuse diff --git a/include/xe-fuse/kernels/qk_norm_rope.hpp b/include/xe-fuse/kernels/qk_norm_rope.hpp new file mode 100644 index 0000000..574b969 --- /dev/null +++ b/include/xe-fuse/kernels/qk_norm_rope.hpp @@ -0,0 +1,118 @@ +#pragma once + +// Standalone fused QK normalization + RoPE kernel. +// +// Applies per-head RMSNorm followed by RoPE rotation in a single pass, +// eliminating the intermediate write-and-read between separate norm and +// RoPE kernels. Used for Q and K tensors in attention blocks of models +// that apply per-head normalization (Gemma 3, Chameleon, Qwen2.5-VL, +// FLUX.2 single-block path). +// +// Input layout: [M, num_heads * head_dim] (M = batch * seq_len) +// cos_sin layout: [M, head_dim] (interleaved: even = cos, odd = sin) +// gamma layout: [head_dim] (per-dim RMSNorm weight; nullable) +// +// One SYCL workgroup = one (token, head) pair. +// SG_SIZE = 16 lanes; head_dim must be a multiple of 16. +// +// Two passes per workgroup: +// Pass 1 — reduce sum_sq across head_dim → rstd +// Pass 2 — normalize (× rstd × gamma), then apply RoPE via shfl_xor + +#include +#include + +#include "cutlass/detail/helper_macros.hpp" +#include "cutlass/gpu_generics.h" + +namespace xe_fuse { + +template +void launch_qk_norm_rope( + sycl::queue& q, + ElementInput const* input_ptr, + ElementInput* output_ptr, + ElementGamma const* gamma_ptr, // [head_dim]; nullptr = no per-dim scale + ElementCosSin const* cos_sin_ptr, // [M, head_dim] interleaved + int M, + int num_heads, + int head_dim, + float eps = 1e-6f) +{ + constexpr int SG_SIZE = 16; + + int work_groups = M * num_heads; + int total_head_dim = num_heads * head_dim; + int elems_per_lane = head_dim / SG_SIZE; // e.g. head_dim=128 → 8 + + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>( + static_cast(work_groups) * SG_SIZE, + static_cast(SG_SIZE)), + [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SG_SIZE)]] { + int wg = static_cast(item.get_group(0)); + int tok = wg / num_heads; + int head = wg % num_heads; + int lane = static_cast(item.get_local_id(0)); + + // Base offset into the full [M, total_head_dim] tensor for this head + int row_base = tok * total_head_dim + head * head_dim; + // cos_sin row offset for this token + int cs_base = tok * head_dim; + + auto sg = item.get_sub_group(); + + // ── Pass 1: reduce sum_sq over head_dim ────────────────────────── + float sum_sq = 0.f; + for (int j = 0; j < elems_per_lane; ++j) { + int d = lane + j * SG_SIZE; + float v = static_cast(input_ptr[row_base + d]); + sum_sq += v * v; + } + for (int off = SG_SIZE / 2; off > 0; off /= 2) + sum_sq += sycl::shift_group_left(sg, sum_sq, off); + + float rstd = sycl::rsqrt(sum_sq / static_cast(head_dim) + eps); + rstd = sycl::group_broadcast(sg, rstd, 0); + + // ── Pass 2: normalize × gamma, then RoPE via shfl_xor ──────────── + bool is_even = (lane & 1) == 0; + + for (int j = 0; j < elems_per_lane; ++j) { + int d = lane + j * SG_SIZE; + + float v = static_cast(input_ptr[row_base + d]) * rstd; + if (gamma_ptr) + v *= static_cast(gamma_ptr[d]); + + // RoPE: interleave-shuffle with adjacent lane + uint32_t my_bits = reinterpret_cast(v); + uint32_t partner_bits = shfl_xor_sync(0xFFFFFFFF, my_bits, 1, 16); + float my_val = reinterpret_cast(my_bits); + float partner_val = reinterpret_cast(partner_bits); + + float cs_val = static_cast(cos_sin_ptr[cs_base + d]); + uint32_t cs_bits = reinterpret_cast(cs_val); + uint32_t partner_cs_bits = shfl_xor_sync(0xFFFFFFFF, cs_bits, 1, 16); + float partner_cs = reinterpret_cast(partner_cs_bits); + + float cos_val = is_even ? cs_val : partner_cs; + float sin_val = is_even ? partner_cs : cs_val; + + float out; + if (is_even) + out = my_val * cos_val + partner_val * sin_val; + else + out = -partner_val * sin_val + my_val * cos_val; + + output_ptr[row_base + d] = static_cast(out); + } + } + ); + }); +} + +} // namespace xe_fuse diff --git a/include/xe-fuse/kernels/swiglu_requant.hpp b/include/xe-fuse/kernels/swiglu_requant.hpp new file mode 100644 index 0000000..d274105 --- /dev/null +++ b/include/xe-fuse/kernels/swiglu_requant.hpp @@ -0,0 +1,92 @@ +#pragma once + +// Standalone SwiGLU + INT8 requantization kernel. +// +// Reads the BF16 output of a SwiGLU GEMM epilogue (K2/K8), applies SwiGLU +// if the input is the raw gate+up interleaved tensor, and writes: +// - INT8 quantized output for the downstream W8A8 GEMM (down projection) +// - per-token scale factors scale_token[m] +// +// +// Two usage modes controlled by ApplySwiGLU: +// ApplySwiGLU = false (default): input is already post-SwiGLU BF16 [M, I]. +// Just quantize it to INT8. +// ApplySwiGLU = true: input is interleaved gate+up BF16 [M, 2*I]. +// Apply SwiGLU then quantize; output is [M, I]. +// +// Two sub-group passes per row: +// Pass 1 -- SwiGLU (if ApplySwiGLU) + reduce max_abs -> scale_token +// Pass 2 -- write INT8 clamped output + +#include +#include + +namespace xe_fuse { + +template +void launch_swiglu_requant( + sycl::queue& q, + ElementInput const* input_ptr, // [M, N_in] N_in = I (post-SwiGLU) or 2*I (interleaved) + int8_t* quant_out_ptr, + float* scale_token_ptr, + int M, + int N_in, // input columns + int L = 1) +{ + constexpr int SG_SIZE = 16; + // Output columns: I = N_in when not applying SwiGLU, N_in/2 otherwise + int N_out = ApplySwiGLU ? N_in / 2 : N_in; + int work_groups = M * L; + + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(static_cast(work_groups) * SG_SIZE, SG_SIZE), + [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SG_SIZE)]] { + int row = item.get_group(0); + int lane = item.get_local_id(0); + auto sg = item.get_sub_group(); + + // ── Pass 1: reduce max_abs (with optional SwiGLU) ────────────────── + float max_abs = 0.f; + for (int col = lane; col < N_out; col += SG_SIZE) { + float val; + if constexpr (ApplySwiGLU) { + float gate = static_cast(input_ptr[row * N_in + col * 2]); + float up = static_cast(input_ptr[row * N_in + col * 2 + 1]); + float silu_gate = gate / (1.0f + sycl::exp(-gate)); + val = silu_gate * up; + } else { + val = static_cast(input_ptr[row * N_in + col]); + } + max_abs = sycl::fmax(max_abs, sycl::fabs(val)); + } + for (int off = SG_SIZE / 2; off > 0; off /= 2) + max_abs = sycl::fmax(max_abs, sycl::shift_group_left(sg, max_abs, off)); + + float scale_tok = max_abs / 127.f + 1e-8f; + scale_tok = sycl::group_broadcast(sg, scale_tok, 0); + + // ── Pass 2: write INT8 output ──────────────────────────────────────── + for (int col = lane; col < N_out; col += SG_SIZE) { + float val; + if constexpr (ApplySwiGLU) { + float gate = static_cast(input_ptr[row * N_in + col * 2]); + float up = static_cast(input_ptr[row * N_in + col * 2 + 1]); + float silu_gate = gate / (1.0f + sycl::exp(-gate)); + val = silu_gate * up; + } else { + val = static_cast(input_ptr[row * N_in + col]); + } + float qval = sycl::round(val / scale_tok); + qval = sycl::fmin(sycl::fmax(qval, -128.f), 127.f); + quant_out_ptr[row * N_out + col] = static_cast(qval); + } + + if (lane == 0) + scale_token_ptr[row] = scale_tok; + } + ); + }); +} + +} // namespace xe_fuse diff --git a/include/xe-fuse/standalone/ops.hpp b/include/xe-fuse/standalone/ops.hpp index d9c8830..97c519e 100644 --- a/include/xe-fuse/standalone/ops.hpp +++ b/include/xe-fuse/standalone/ops.hpp @@ -353,4 +353,50 @@ inline void dequant_w8a8(sycl::queue& q, }); } +// Quantize already-normalized BF16 values to INT8. +// (Use after rms_norm when INT8 output is also needed — 2-pass: max then write.) +// scale_token_out[m] = max_n(|input[m,n]|) / 127 +// quant_out[m,n] = round(input[m,n] / scale_token) clamped to [-128, 127] +inline void quantize_bf16_to_int8(sycl::queue& q, + bf16 const* input, + int8_t* quant_out, + float* scale_token_out, + int M, int N, int L = 1) { + constexpr int SG_SIZE = 16; + int work_groups = M * L; + + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(static_cast(work_groups) * SG_SIZE, SG_SIZE), + [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SG_SIZE)]] { + int row = item.get_group(0); + int lane = item.get_local_id(0); + auto sg = item.get_sub_group(); + + // ── Pass 1: reduce max_abs ──────────────────────────────────────── + float max_abs = 0.f; + for (int col = lane; col < N; col += SG_SIZE) { + float v = static_cast(input[row * N + col]); + max_abs = sycl::fmax(max_abs, sycl::fabs(v)); + } + for (int off = SG_SIZE / 2; off > 0; off /= 2) + max_abs = sycl::fmax(max_abs, sycl::shift_group_left(sg, max_abs, off)); + + float scale_tok = max_abs / 127.f + 1e-8f; + scale_tok = sycl::group_broadcast(sg, scale_tok, 0); + + // ── Pass 2: write INT8 ──────────────────────────────────────────── + for (int col = lane; col < N; col += SG_SIZE) { + float v = static_cast(input[row * N + col]); + float qval = sycl::round(v / scale_tok); + qval = sycl::fmin(sycl::fmax(qval, -128.f), 127.f); + quant_out[row * N + col] = static_cast(qval); + } + + if (lane == 0) + scale_token_out[row] = scale_tok; + }); + }); +} + } // namespace xe_fuse::standalone diff --git a/include/xe-fuse/standalone/vllm_ops.hpp b/include/xe-fuse/standalone/vllm_ops.hpp index c056f70..32be34f 100644 --- a/include/xe-fuse/standalone/vllm_ops.hpp +++ b/include/xe-fuse/standalone/vllm_ops.hpp @@ -22,6 +22,8 @@ #include #include #include "cutlass/bfloat16.h" +#include "cutlass/detail/helper_macros.hpp" +#include "cutlass/gpu_generics.h" namespace xe_fuse::vllm_equiv { @@ -239,8 +241,9 @@ inline void rotary_embedding(sycl::queue& q, bf16* query, bf16* key, // - xe-fuse keeps it in registers throughout the GEMM epilogue // Merged: dequant INT32 → BF16, then apply SwiGLU -// Input layout: [L, M, 2*d] INT32 (gate interleaved with up) -// Output layout: [L, M, 2*d] BF16 (both lanes carry the same silu(gate)*up value) +// Input layout: [L, M, 2*d] INT32 (gate and up interleaved: even=gate, odd=up) +// Output layout: [L, M, 2*d] BF16 — both even and odd positions at index i carry +// silu(gate[i]) * up[i]; the caller reads only N/2 columns (even). inline void dequant_and_silu_mul(sycl::queue& q, bf16* out, int32_t const* acc, @@ -345,4 +348,115 @@ inline void dequant_and_rotary_embedding(sycl::queue& q, }); } +// Per-head RMSNorm for QK tensors. +// Input/output layout: [M, num_heads * head_dim] +// gamma: [head_dim] nullable (float) +// One workgroup per (token, head) pair. SG_SIZE=16. +inline void rms_norm_per_head(sycl::queue& q, bf16* out, bf16 const* input, + float const* gamma, // nullable + int M, int num_heads, int head_dim, + float eps = 1e-6f) { + constexpr int SG_SIZE = 16; + int work_groups = M * num_heads; + int total_head_dim = num_heads * head_dim; + int elems_per_lane = head_dim / SG_SIZE; + + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>( + static_cast(work_groups) * SG_SIZE, + static_cast(SG_SIZE)), + [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SG_SIZE)]] { + int wg = static_cast(item.get_group(0)); + int tok = wg / num_heads; + int head = wg % num_heads; + int lane = static_cast(item.get_local_id(0)); + + int row_base = tok * total_head_dim + head * head_dim; + auto sg = item.get_sub_group(); + + // ── Pass 1: reduce sum_sq over head_dim ────────────────────────── + float sum_sq = 0.f; + for (int j = 0; j < elems_per_lane; ++j) { + int d = lane + j * SG_SIZE; + float v = static_cast(input[row_base + d]); + sum_sq += v * v; + } + for (int off = SG_SIZE / 2; off > 0; off /= 2) + sum_sq += sycl::shift_group_left(sg, sum_sq, off); + + float rstd = sycl::rsqrt(sum_sq / static_cast(head_dim) + eps); + rstd = sycl::group_broadcast(sg, rstd, 0); + + // ── Pass 2: normalize × gamma (if present), write bf16 ─────────── + for (int j = 0; j < elems_per_lane; ++j) { + int d = lane + j * SG_SIZE; + float v = static_cast(input[row_base + d]) * rstd; + if (gamma != nullptr) + v *= gamma[d]; + out[row_base + d] = static_cast(v); + } + }); + }); +} + +// RoPE with interleaved cos/sin format matching xe-fuse convention. +// Input/output: [M, num_heads * head_dim] in-place (reads from 'input', writes to 'out'). +// cos_sin: [M, head_dim] interleaved: even index = cos, odd index = sin. +// One workgroup per (token, head). SG_SIZE=16. +inline void rope_interleaved(sycl::queue& q, bf16* out, bf16 const* input, + float const* cos_sin, + int M, int num_heads, int head_dim) { + constexpr int SG_SIZE = 16; + int work_groups = M * num_heads; + int total_head_dim = num_heads * head_dim; + int elems_per_lane = head_dim / SG_SIZE; + + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>( + static_cast(work_groups) * SG_SIZE, + static_cast(SG_SIZE)), + [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SG_SIZE)]] { + int wg = static_cast(item.get_group(0)); + int tok = wg / num_heads; + int head = wg % num_heads; + int lane = static_cast(item.get_local_id(0)); + + int row_base = tok * total_head_dim + head * head_dim; + int cs_base = tok * head_dim; + + bool is_even = (lane & 1) == 0; + + for (int j = 0; j < elems_per_lane; ++j) { + int d = lane + j * SG_SIZE; + + float v = static_cast(input[row_base + d]); + + // RoPE: interleave-shuffle with adjacent lane + uint32_t my_bits = reinterpret_cast(v); + uint32_t partner_bits = shfl_xor_sync(0xFFFFFFFF, my_bits, 1, 16); + float my_val = reinterpret_cast(my_bits); + float partner_val = reinterpret_cast(partner_bits); + + float cs_val = cos_sin[cs_base + d]; + uint32_t cs_bits = reinterpret_cast(cs_val); + uint32_t partner_cs_bits = shfl_xor_sync(0xFFFFFFFF, cs_bits, 1, 16); + float partner_cs = reinterpret_cast(partner_cs_bits); + + float cos_val = is_even ? cs_val : partner_cs; + float sin_val = is_even ? partner_cs : cs_val; + + float result; + if (is_even) + result = my_val * cos_val + partner_val * sin_val; + else + result = -partner_val * sin_val + my_val * cos_val; + + out[row_base + d] = static_cast(result); + } + }); + }); +} + } // namespace xe_fuse::vllm_equiv diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5cb5afe..3a830a3 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -14,3 +14,11 @@ xe_fuse_add_executable(test_k8 test_k8.cpp) xe_fuse_add_executable(test_k9 test_k9.cpp) xe_fuse_add_executable(test_k9_fused test_k9_fused.cpp) xe_fuse_add_executable(test_k9_fused_atomic test_k9_fused_atomic.cpp) + +xe_fuse_add_executable(test_qk_norm_rope test_qk_norm_rope.cpp) +xe_fuse_add_executable(test_k0g test_k0g.cpp) +xe_fuse_add_executable(test_fp8_k2 test_fp8_k2.cpp) + +xe_fuse_add_executable(bench_qk_norm_rope bench_qk_norm_rope.cpp) +xe_fuse_add_executable(bench_norm_quantize bench_norm_quantize.cpp) +xe_fuse_add_executable(bench_swiglu_requant bench_swiglu_requant.cpp) diff --git a/tests/bench_norm_quantize.cpp b/tests/bench_norm_quantize.cpp new file mode 100644 index 0000000..7129312 --- /dev/null +++ b/tests/bench_norm_quantize.cpp @@ -0,0 +1,173 @@ +// xe-fuse bench: norm + INT8 quantize +// Compares fusion strategies for RMSNorm followed by INT8 quantization. +// +// XE_FUSE_DUAL: launch_norm_quantize_dual (1 kernel -> BF16 normed + INT8) +// XE_FUSE_ORIG: launch_compute_rstd_and_quantize (1 kernel -> INT8 only) +// VLLM_EQUIV: vllm_equiv::rms_norm + quantize_bf16_to_int8 (2 kernels) +// NAIVE: compute_rstd + quantize_activations (3 kernels: rstd + 2-pass quant) +// +// oneDNN: no native fused RMSNorm+quantize; would require >=2 primitives. + +#include "xe-fuse/kernels/compute_rstd.hpp" +#include "xe-fuse/standalone/ops.hpp" +#include "xe-fuse/standalone/vllm_ops.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "sycl_common.hpp" +#include "helper.h" +#include +#include +#include + +using bf16 = cutlass::bfloat16_t; + +struct Options { + int M = 512; + int N = 4096; + int L = 1; + int iterations = 100; + + void parse(int argc, char const** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", M, 512); + cmd.get_cmd_line_argument("n", N, 4096); + cmd.get_cmd_line_argument("l", L, 1); + cmd.get_cmd_line_argument("iterations", iterations, 100); + } +}; + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + int M = opts.M; + int N = opts.N; + int L = opts.L; + int iters = opts.iterations; + + int64_t total = static_cast(M) * N * L; + + // ── Shared input buffer ────────────────────────────────────────────────── + cutlass::DeviceAllocation block_input(total); + initialize_block(block_input, 2025); + + // bf16 weight for vllm_equiv::rms_norm (identity = 1.0) + cutlass::DeviceAllocation block_weight(N); + { + std::vector hw(N, static_cast(1.0f)); + compat::get_default_queue().memcpy(block_weight.get(), hw.data(), + N * sizeof(bf16)); + } + compat::wait(); + + // ── Per-approach output buffers ────────────────────────────────────────── + // XE_FUSE_DUAL + cutlass::DeviceAllocation quant_dual(total); + cutlass::DeviceAllocation normed_dual(total); + cutlass::DeviceAllocation scale_dual(static_cast(M) * L); + + // XE_FUSE_ORIG + cutlass::DeviceAllocation quant_orig(total); + cutlass::DeviceAllocation scale_orig(static_cast(M) * L); + + // VLLM_EQUIV + cutlass::DeviceAllocation normed_vllm(total); + cutlass::DeviceAllocation quant_vllm(total); + cutlass::DeviceAllocation scale_vllm(static_cast(M) * L); + + // NAIVE + cutlass::DeviceAllocation rstd_naive(static_cast(M) * L); + cutlass::DeviceAllocation quant_naive(total); + cutlass::DeviceAllocation scale_naive(static_cast(M) * L); + + sycl::queue q = compat::get_default_queue(); + + // ── Warm up ────────────────────────────────────────────────────────────── + constexpr int warmup = 5; + + for (int i = 0; i < warmup; ++i) + xe_fuse::launch_norm_quantize_dual(q, block_input.get(), + quant_dual.get(), normed_dual.get(), scale_dual.get(), M, N, L); + + for (int i = 0; i < warmup; ++i) + xe_fuse::launch_compute_rstd_and_quantize(q, block_input.get(), + quant_orig.get(), scale_orig.get(), M, N, L); + + for (int i = 0; i < warmup; ++i) { + xe_fuse::vllm_equiv::rms_norm(q, normed_vllm.get(), block_input.get(), + block_weight.get(), M * L, N); + xe_fuse::standalone::quantize_bf16_to_int8(q, normed_vllm.get(), + quant_vllm.get(), scale_vllm.get(), M, N, L); + } + + for (int i = 0; i < warmup; ++i) { + xe_fuse::standalone::compute_rstd(q, rstd_naive.get(), block_input.get(), + M, N, L); + xe_fuse::standalone::quantize_activations(q, block_input.get(), + rstd_naive.get(), quant_naive.get(), scale_naive.get(), M, N, L); + } + compat::wait(); + + // ── Benchmark XE_FUSE_DUAL ─────────────────────────────────────────────── + GPU_Clock timer; + timer.start(); + for (int i = 0; i < iters; ++i) + xe_fuse::launch_norm_quantize_dual(q, block_input.get(), + quant_dual.get(), normed_dual.get(), scale_dual.get(), M, N, L); + compat::wait(); + float time_dual = timer.seconds() / iters; + // read bf16 input + write bf16 normed + write int8 + double bytes_dual = (double)total * (sizeof(bf16) + sizeof(bf16) + sizeof(int8_t)); + + // ── Benchmark XE_FUSE_ORIG ─────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) + xe_fuse::launch_compute_rstd_and_quantize(q, block_input.get(), + quant_orig.get(), scale_orig.get(), M, N, L); + compat::wait(); + float time_orig = timer.seconds() / iters; + // read bf16 input (3 passes) + write int8 + double bytes_orig = (double)total * (3 * sizeof(bf16) + sizeof(int8_t)); + + // ── Benchmark VLLM_EQUIV ───────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) { + xe_fuse::vllm_equiv::rms_norm(q, normed_vllm.get(), block_input.get(), + block_weight.get(), M * L, N); + xe_fuse::standalone::quantize_bf16_to_int8(q, normed_vllm.get(), + quant_vllm.get(), scale_vllm.get(), M, N, L); + } + compat::wait(); + float time_vllm = timer.seconds() / iters; + // rms_norm: read input + write normed; quant: read normed (2 passes) + write int8 + double bytes_vllm = (double)total * (sizeof(bf16) + sizeof(bf16) + 2 * sizeof(bf16) + sizeof(int8_t)); + + // ── Benchmark NAIVE ─────────────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) { + xe_fuse::standalone::compute_rstd(q, rstd_naive.get(), block_input.get(), + M, N, L); + xe_fuse::standalone::quantize_activations(q, block_input.get(), + rstd_naive.get(), quant_naive.get(), scale_naive.get(), M, N, L); + } + compat::wait(); + float time_naive = timer.seconds() / iters; + // compute_rstd: read input; quantize pass1: read input; quantize pass2: read input + write int8 + double bytes_naive = (double)total * (3 * sizeof(bf16) + sizeof(int8_t)); + + // ── Print results ───────────────────────────────────────────────────────── + printf("\n=== Norm + INT8 Quantize: M=%d N=%d ===\n", M, N); + printf("XE_FUSE_DUAL (1 kernel, BF16+INT8 out): [%.3f]GB/s (%.4f)ms\n", + bytes_dual * 1e-9 / time_dual, time_dual * 1000.f); + printf("XE_FUSE_ORIG (1 kernel, INT8 only): [%.3f]GB/s (%.4f)ms\n", + bytes_orig * 1e-9 / time_orig, time_orig * 1000.f); + printf("VLLM_EQUIV (2 kernels): [%.3f]GB/s (%.4f)ms\n", + bytes_vllm * 1e-9 / time_vllm, time_vllm * 1000.f); + printf("NAIVE (3 kernels): [%.3f]GB/s (%.4f)ms\n", + bytes_naive * 1e-9 / time_naive, time_naive * 1000.f); + printf("Speedup XE_FUSE_DUAL vs VLLM_EQUIV: %.2fx\n", time_vllm / time_dual); + printf("Speedup XE_FUSE_ORIG vs NAIVE: %.2fx\n", time_naive / time_orig); + + return 0; +} diff --git a/tests/bench_qk_norm_rope.cpp b/tests/bench_qk_norm_rope.cpp new file mode 100644 index 0000000..261d91c --- /dev/null +++ b/tests/bench_qk_norm_rope.cpp @@ -0,0 +1,206 @@ +// xe-fuse bench: QK norm + RoPE +// Compares three fusion strategies for per-head RMSNorm + RoPE on Q/K tensors. +// +// XE_FUSE: launch_qk_norm_rope (1 kernel, fused) +// VLLM_EQUIV: rms_norm_per_head + rope_interleaved (2 kernels, intermediate DRAM write) +// NAIVE: compute_rstd_per_head + scale_rows + rope (3 separate kernels, 2 DRAM writes) +// +// oneDNN has no native QK-norm+RoPE primitive; the VLLM_EQUIV path represents +// what an oneDNN-based implementation would require at minimum. + +#include "xe-fuse/kernels/qk_norm_rope.hpp" +#include "xe-fuse/kernels/compute_rstd.hpp" +#include "xe-fuse/standalone/ops.hpp" +#include "xe-fuse/standalone/vllm_ops.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "sycl_common.hpp" +#include "helper.h" +#include +#include +#include + +using bf16 = cutlass::bfloat16_t; + +struct Options { + int seq_len = 4096; + int num_heads = 32; + int head_dim = 128; + int batch = 1; + int iterations = 100; + + void parse(int argc, char const** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("seq_len", seq_len, 4096); + cmd.get_cmd_line_argument("num_heads", num_heads, 32); + cmd.get_cmd_line_argument("head_dim", head_dim, 128); + cmd.get_cmd_line_argument("batch", batch, 1); + cmd.get_cmd_line_argument("iterations", iterations, 100); + } +}; + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + int M = opts.batch * opts.seq_len; + int num_heads = opts.num_heads; + int head_dim = opts.head_dim; + int iters = opts.iterations; + + int64_t total_size = static_cast(M) * num_heads * head_dim; + int64_t cos_sin_size = static_cast(M) * head_dim; + // cos_sin_flat: [M, num_heads * head_dim] — per-token values broadcast per head + int64_t cos_sin_flat_size = total_size; + + // ── Shared read-only buffers ───────────────────────────────────────────── + cutlass::DeviceAllocation block_input(total_size); + cutlass::DeviceAllocation block_gamma(head_dim); + cutlass::DeviceAllocation block_cos_sin(cos_sin_size); + + initialize_block(block_input, 2025); + + // gamma: uniform (0.5, 1.5) + { + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.5f, 1.5f); + std::vector h(head_dim); + for (auto& v : h) v = dist(rng); + compat::get_default_queue().memcpy(block_gamma.get(), h.data(), + head_dim * sizeof(float)); + } + + // cos_sin: realistic RoPE frequencies [M, head_dim] interleaved + { + std::vector h(cos_sin_size); + for (int tok = 0; tok < M; ++tok) { + for (int k = 0; k < head_dim / 2; ++k) { + float freq = 1.0f / std::pow(10000.f, 2.f * k / static_cast(head_dim)); + float angle = static_cast(tok) * freq; + int base = tok * head_dim; + h[base + 2 * k] = std::cos(angle); + h[base + 2 * k + 1] = std::sin(angle); + } + } + compat::get_default_queue().memcpy(block_cos_sin.get(), h.data(), + cos_sin_size * sizeof(float)); + } + + // cos_sin_flat: [M, num_heads * head_dim] — broadcast per head for NAIVE rope + cutlass::DeviceAllocation block_cos_sin_flat(cos_sin_flat_size); + { + std::vector h_cs(cos_sin_size); + compat::get_default_queue().memcpy(h_cs.data(), block_cos_sin.get(), + cos_sin_size * sizeof(float)).wait(); + std::vector h_flat(cos_sin_flat_size); + for (int tok = 0; tok < M; ++tok) { + for (int hd = 0; hd < num_heads; ++hd) { + for (int d = 0; d < head_dim; ++d) { + h_flat[tok * num_heads * head_dim + hd * head_dim + d] = + h_cs[tok * head_dim + d]; + } + } + } + compat::get_default_queue().memcpy(block_cos_sin_flat.get(), h_flat.data(), + cos_sin_flat_size * sizeof(float)); + } + compat::wait(); + + // ── Per-approach output / working buffers ──────────────────────────────── + // XE_FUSE + cutlass::DeviceAllocation out_xe(total_size); + + // VLLM_EQUIV: uses normed intermediate + cutlass::DeviceAllocation normed_vllm(total_size); + cutlass::DeviceAllocation out_vllm(total_size); + + // NAIVE: uses normed working copy (modified in-place by scale_rows), + // rstd [M * num_heads], separate rope output + cutlass::DeviceAllocation normed_naive(total_size); + cutlass::DeviceAllocation rstd_naive(static_cast(M) * num_heads); + cutlass::DeviceAllocation out_naive(total_size); + + // Pre-copy input to normed_naive (scale_rows operates in-place) + compat::get_default_queue().memcpy(normed_naive.get(), block_input.get(), + total_size * sizeof(bf16)); + compat::wait(); + + sycl::queue q = compat::get_default_queue(); + + // ── Warm up ────────────────────────────────────────────────────────────── + constexpr int warmup = 5; + + for (int i = 0; i < warmup; ++i) + xe_fuse::launch_qk_norm_rope(q, block_input.get(), out_xe.get(), + block_gamma.get(), block_cos_sin.get(), M, num_heads, head_dim); + + for (int i = 0; i < warmup; ++i) { + xe_fuse::vllm_equiv::rms_norm_per_head(q, normed_vllm.get(), block_input.get(), + block_gamma.get(), M, num_heads, head_dim); + xe_fuse::vllm_equiv::rope_interleaved(q, out_vllm.get(), normed_vllm.get(), + block_cos_sin.get(), M, num_heads, head_dim); + } + + for (int i = 0; i < warmup; ++i) { + xe_fuse::launch_compute_rstd(q, block_input.get(), rstd_naive.get(), + M * num_heads, head_dim, 1); + xe_fuse::standalone::scale_rows(q, normed_naive.get(), rstd_naive.get(), + M * num_heads, head_dim, 1); + xe_fuse::standalone::rope(q, out_naive.get(), normed_naive.get(), + block_cos_sin_flat.get(), M, num_heads * head_dim, 1); + } + compat::wait(); + + // ── Benchmark XE_FUSE ──────────────────────────────────────────────────── + GPU_Clock timer; + timer.start(); + for (int i = 0; i < iters; ++i) + xe_fuse::launch_qk_norm_rope(q, block_input.get(), out_xe.get(), + block_gamma.get(), block_cos_sin.get(), M, num_heads, head_dim); + compat::wait(); + float time_xe = timer.seconds() / iters; + double bytes_xe = 2.0 * total_size * sizeof(bf16); + + // ── Benchmark VLLM_EQUIV ───────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) { + xe_fuse::vllm_equiv::rms_norm_per_head(q, normed_vllm.get(), block_input.get(), + block_gamma.get(), M, num_heads, head_dim); + xe_fuse::vllm_equiv::rope_interleaved(q, out_vllm.get(), normed_vllm.get(), + block_cos_sin.get(), M, num_heads, head_dim); + } + compat::wait(); + float time_vllm = timer.seconds() / iters; + // read + intermediate write + read (final write implicit) + double bytes_vllm = 3.0 * total_size * sizeof(bf16); + + // ── Benchmark NAIVE ─────────────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) { + xe_fuse::launch_compute_rstd(q, block_input.get(), rstd_naive.get(), + M * num_heads, head_dim, 1); + xe_fuse::standalone::scale_rows(q, normed_naive.get(), rstd_naive.get(), + M * num_heads, head_dim, 1); + xe_fuse::standalone::rope(q, out_naive.get(), normed_naive.get(), + block_cos_sin_flat.get(), M, num_heads * head_dim, 1); + } + compat::wait(); + float time_naive = timer.seconds() / iters; + // rstd read + scale read+write + rope read+write + double bytes_naive = 4.0 * total_size * sizeof(bf16); + + // ── Print results ───────────────────────────────────────────────────────── + printf("\n=== QK Norm + RoPE: M=%d num_heads=%d head_dim=%d ===\n", + M, num_heads, head_dim); + printf("XE_FUSE (1 kernel, fused): [%.3f]GB/s (%.4f)ms\n", + bytes_xe * 1e-9 / time_xe, time_xe * 1000.f); + printf("VLLM_EQUIV (2 kernels, norm+rope): [%.3f]GB/s (%.4f)ms\n", + bytes_vllm * 1e-9 / time_vllm, time_vllm * 1000.f); + printf("NAIVE (3 kernels, rstd+scale+rope): [%.3f]GB/s (%.4f)ms\n", + bytes_naive * 1e-9 / time_naive, time_naive * 1000.f); + printf("Speedup XE_FUSE vs VLLM_EQUIV: %.2fx\n", time_vllm / time_xe); + printf("Speedup XE_FUSE vs NAIVE: %.2fx\n", time_naive / time_xe); + + return 0; +} diff --git a/tests/bench_swiglu_requant.cpp b/tests/bench_swiglu_requant.cpp new file mode 100644 index 0000000..8bc2db9 --- /dev/null +++ b/tests/bench_swiglu_requant.cpp @@ -0,0 +1,184 @@ +// xe-fuse bench: SwiGLU + INT8 requant +// Compares fusion strategies for SwiGLU activation followed by INT8 quantization. +// +// XE_FUSE_FUSED: launch_swiglu_requant (1 kernel, interleaved gate+up -> INT8) +// XE_FUSE_POST: launch_swiglu_requant (1 kernel, post-SwiGLU BF16 -> INT8) +// VLLM_EQUIV: vllm_equiv::silu_and_mul + quantize_bf16_to_int8 (2 kernels) +// NAIVE: standalone::swiglu + compute_rstd + quantize_activations (3 kernels) +// +// oneDNN: no native SwiGLU+quantize primitive. + +#include "xe-fuse/kernels/swiglu_requant.hpp" +#include "xe-fuse/standalone/ops.hpp" +#include "xe-fuse/standalone/vllm_ops.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "sycl_common.hpp" +#include "helper.h" +#include +#include +#include + +using bf16 = cutlass::bfloat16_t; + +struct Options { + int M = 512; + int N_ffn = 28672; // LLaMA 3 8B FFN intermediate size (gate+up concatenated) + int L = 1; + int iterations = 100; + + void parse(int argc, char const** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", M, 512); + cmd.get_cmd_line_argument("n_ffn", N_ffn, 28672); + cmd.get_cmd_line_argument("l", L, 1); + cmd.get_cmd_line_argument("iterations", iterations, 100); + } +}; + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + int M = opts.M; + int N_ffn = opts.N_ffn; + int L = opts.L; + int iters = opts.iterations; + int d = N_ffn / 2; // post-SwiGLU intermediate size + + int64_t total_gate_up = static_cast(M) * N_ffn * L; + int64_t total_out = static_cast(M) * d * L; + + sycl::queue q = compat::get_default_queue(); + + // ── Shared interleaved gate+up input (for XE_FUSE_FUSED and VLLM_EQUIV) ── + cutlass::DeviceAllocation block_gate_up(total_gate_up); + initialize_block(block_gate_up, 2025); + + // Pre-compute post-SwiGLU BF16 for XE_FUSE_POST and NAIVE sub-steps. + cutlass::DeviceAllocation block_post_swiglu(total_out); + xe_fuse::vllm_equiv::silu_and_mul(q, + block_post_swiglu.get(), block_gate_up.get(), d, M * L); + compat::wait(); + + // ── Per-approach output buffers ────────────────────────────────────────── + // XE_FUSE_FUSED + cutlass::DeviceAllocation quant_fused(total_out); + cutlass::DeviceAllocation scale_fused(static_cast(M) * L); + + // XE_FUSE_POST + cutlass::DeviceAllocation quant_post(total_out); + cutlass::DeviceAllocation scale_post(static_cast(M) * L); + + // VLLM_EQUIV: intermediate post-swiglu bf16 + cutlass::DeviceAllocation swiglu_vllm(total_out); + cutlass::DeviceAllocation quant_vllm(total_out); + cutlass::DeviceAllocation scale_vllm(static_cast(M) * L); + + // NAIVE: swiglu step on gate+up in-place working copy; rstd+quant on post_swiglu + cutlass::DeviceAllocation swiglu_work_naive(total_gate_up); + cutlass::DeviceAllocation rstd_naive(static_cast(M) * L); + cutlass::DeviceAllocation quant_naive(total_out); + cutlass::DeviceAllocation scale_naive(static_cast(M) * L); + + // Initialize NAIVE swiglu working copy from gate+up input + q.memcpy(swiglu_work_naive.get(), block_gate_up.get(), + total_gate_up * sizeof(bf16)); + compat::wait(); + + // ── Warm up ────────────────────────────────────────────────────────────── + constexpr int warmup = 5; + + for (int i = 0; i < warmup; ++i) + xe_fuse::launch_swiglu_requant(q, block_gate_up.get(), + quant_fused.get(), scale_fused.get(), M * L, N_ffn); + + for (int i = 0; i < warmup; ++i) + xe_fuse::launch_swiglu_requant(q, block_post_swiglu.get(), + quant_post.get(), scale_post.get(), M * L, d); + + for (int i = 0; i < warmup; ++i) { + xe_fuse::vllm_equiv::silu_and_mul(q, swiglu_vllm.get(), block_gate_up.get(), + d, M * L); + xe_fuse::standalone::quantize_bf16_to_int8(q, swiglu_vllm.get(), + quant_vllm.get(), scale_vllm.get(), M, d, L); + } + + for (int i = 0; i < warmup; ++i) { + xe_fuse::standalone::swiglu(q, swiglu_work_naive.get(), M * L, N_ffn, 1); + xe_fuse::standalone::compute_rstd(q, rstd_naive.get(), + swiglu_work_naive.get(), M, d, L); + xe_fuse::standalone::quantize_activations(q, swiglu_work_naive.get(), + rstd_naive.get(), quant_naive.get(), scale_naive.get(), M, d, L); + } + compat::wait(); + + // ── Benchmark XE_FUSE_FUSED ────────────────────────────────────────────── + GPU_Clock timer; + timer.start(); + for (int i = 0; i < iters; ++i) + xe_fuse::launch_swiglu_requant(q, block_gate_up.get(), + quant_fused.get(), scale_fused.get(), M * L, N_ffn); + compat::wait(); + float time_fused = timer.seconds() / iters; + // read gate+up [M, N_ffn] bf16 + write [M, d] int8 + double bytes_fused = (double)M * L * (N_ffn * sizeof(bf16) + d * sizeof(int8_t)); + + // ── Benchmark XE_FUSE_POST ─────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) + xe_fuse::launch_swiglu_requant(q, block_post_swiglu.get(), + quant_post.get(), scale_post.get(), M * L, d); + compat::wait(); + float time_post = timer.seconds() / iters; + // read post-swiglu [M, d] bf16 (2 passes) + write [M, d] int8 + double bytes_post = (double)M * L * d * (2 * sizeof(bf16) + sizeof(int8_t)); + + // ── Benchmark VLLM_EQUIV ───────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) { + xe_fuse::vllm_equiv::silu_and_mul(q, swiglu_vllm.get(), block_gate_up.get(), + d, M * L); + xe_fuse::standalone::quantize_bf16_to_int8(q, swiglu_vllm.get(), + quant_vllm.get(), scale_vllm.get(), M, d, L); + } + compat::wait(); + float time_vllm = timer.seconds() / iters; + // silu_and_mul: read [M, N_ffn] + write [M, d]; quant: read [M, d] (2 passes) + write [M, d] int8 + double bytes_vllm = (double)M * L * (N_ffn * sizeof(bf16) + d * sizeof(bf16) + + 2 * d * sizeof(bf16) + d * sizeof(int8_t)); + + // ── Benchmark NAIVE ─────────────────────────────────────────────────────── + timer.start(); + for (int i = 0; i < iters; ++i) { + // swiglu: reads [M, N_ffn] in-place, writes SwiGLU result to first [M, d] + xe_fuse::standalone::swiglu(q, swiglu_work_naive.get(), M * L, N_ffn, 1); + // rstd: reads first [M, d] of swiglu output + xe_fuse::standalone::compute_rstd(q, rstd_naive.get(), + swiglu_work_naive.get(), M, d, L); + // quantize: reads [M, d] twice + writes [M, d] int8 + xe_fuse::standalone::quantize_activations(q, swiglu_work_naive.get(), + rstd_naive.get(), quant_naive.get(), scale_naive.get(), M, d, L); + } + compat::wait(); + float time_naive = timer.seconds() / iters; + // swiglu: 2*N_ffn; rstd: d; quant: 2*d read + d int8 write + double bytes_naive = (double)M * L * (2 * N_ffn * sizeof(bf16) + d * sizeof(bf16) + + 2 * d * sizeof(bf16) + d * sizeof(int8_t)); + + // ── Print results ───────────────────────────────────────────────────────── + printf("\n=== SwiGLU + INT8 Requant: M=%d N_ffn=%d (d=%d) ===\n", M, N_ffn, d); + printf("XE_FUSE_FUSED (1 kernel, gate+up->INT8): [%.3f]GB/s (%.4f)ms\n", + bytes_fused * 1e-9 / time_fused, time_fused * 1000.f); + printf("XE_FUSE_POST (1 kernel, swiglu->INT8): [%.3f]GB/s (%.4f)ms\n", + bytes_post * 1e-9 / time_post, time_post * 1000.f); + printf("VLLM_EQUIV (2 kernels): [%.3f]GB/s (%.4f)ms\n", + bytes_vllm * 1e-9 / time_vllm, time_vllm * 1000.f); + printf("NAIVE (3 kernels): [%.3f]GB/s (%.4f)ms\n", + bytes_naive * 1e-9 / time_naive, time_naive * 1000.f); + printf("Speedup XE_FUSE_FUSED vs VLLM_EQUIV: %.2fx\n", time_vllm / time_fused); + printf("Speedup XE_FUSE_FUSED vs NAIVE: %.2fx\n", time_naive / time_fused); + + return 0; +} diff --git a/tests/test_fp8_k2.cpp b/tests/test_fp8_k2.cpp new file mode 100644 index 0000000..e2edcec --- /dev/null +++ b/tests/test_fp8_k2.cpp @@ -0,0 +1,240 @@ +// xe-fuse test: K2_FP8 — gemm_fp8_dequant_swiglu +// D = SwiGLU( dequant(A_f8 @ B_f8) ) +// +// FP8×FP8 GEMM on BMG-G31: FP8 inputs are upcasted to FP16 before the XMX16 +// MMA, which accumulates in float. Dequantization applies per-token scale_a[m] +// and per-channel scale_b[n] in the epilogue. SwiGLU follows. +// +// Reference: +// acc[m,n] = sum_k( float(A_f8[m,k]) * float(B_f8[k,n]) ) +// dequant[m,n] = acc[m,n] * scale_a[m] * scale_b[n] +// D[m, 2i] = silu(dequant[m, 2i]) * dequant[m, 2i+1] +// D[m, 2i+1] = silu(dequant[m, 2i]) * dequant[m, 2i+1] +// +// Verification tolerance is looser than W8A8 (FP8 range ≈ ±448 vs INT8 ±127). + +#include "xe-fuse/kernels/gemm_fp8_dequant.hpp" + +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/device/tensor_compare.h" + +#include "sycl_common.hpp" +#include "helper.h" + +#include +#include +#include +#include + +using namespace cute; + +struct Options { + int m = 512, n = 28672, k = 4096, l = 1; + int iterations = 100; + int verify = 1; + + void parse(int argc, char const** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", m, 512); + cmd.get_cmd_line_argument("n", n, 28672); + cmd.get_cmd_line_argument("k", k, 4096); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("iterations", iterations, 100); + cmd.get_cmd_line_argument("verify", verify, 1); + } +}; + +using K2FP8 = xe_fuse::GemmFP8DequantSwiGLU<>; +using GemmOp = K2FP8::Gemm; +using ElementFP8 = K2FP8::ElementA; // float_e4m3_t + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + int M = opts.m, N = opts.n, K = opts.k, L = opts.l; + + using StrideA = typename GemmOp::GemmKernel::StrideA; + using StrideB = typename GemmOp::GemmKernel::StrideB; + + auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, make_shape(M, K, L)); + auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, make_shape(N, K, L)); + auto stride_C = cutlass::make_cute_packed_stride(K2FP8::StrideC{}, make_shape(M, N, L)); + auto stride_D = cutlass::make_cute_packed_stride(K2FP8::StrideD{}, make_shape(M, N, L)); + + cutlass::DeviceAllocation block_A(static_cast(M) * K * L); + cutlass::DeviceAllocation block_B(static_cast(K) * N * L); + cutlass::DeviceAllocation block_D(static_cast(M) * N * L); + cutlass::DeviceAllocation block_ref_D(static_cast(M) * N * L); + + cutlass::DeviceAllocation block_scale_a(static_cast(M) * L); + cutlass::DeviceAllocation block_scale_b(static_cast(N) * L); + + // Float working buffers for reference GEMM + cutlass::DeviceAllocation block_A_f32(static_cast(M) * K * L); + cutlass::DeviceAllocation block_B_f32(static_cast(K) * N * L); + cutlass::DeviceAllocation block_acc_f32(static_cast(M) * N * L); + + // Initialize FP8 A and B on host with values in [-4, 4] (well within E4M3 range) + { + std::mt19937 rng_a(2001), rng_b(2002); + std::uniform_real_distribution dist(-4.f, 4.f); + + std::vector h_A(static_cast(M) * K * L); + std::vector h_B(static_cast(K) * N * L); + for (auto& v : h_A) v = ElementFP8(dist(rng_a)); + for (auto& v : h_B) v = ElementFP8(dist(rng_b)); + + compat::get_default_queue().memcpy(block_A.get(), h_A.data(), h_A.size() * sizeof(ElementFP8)); + compat::get_default_queue().memcpy(block_B.get(), h_B.data(), h_B.size() * sizeof(ElementFP8)); + } + + // Scales: small positive floats in [0.001, 0.01] + { + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.001f, 0.01f); + + std::vector h_sa(static_cast(M) * L); + std::vector h_sb(static_cast(N) * L); + for (auto& v : h_sa) v = dist(rng); + for (auto& v : h_sb) v = dist(rng); + + compat::get_default_queue().memcpy(block_scale_a.get(), h_sa.data(), h_sa.size() * sizeof(float)); + compat::get_default_queue().memcpy(block_scale_b.get(), h_sb.data(), h_sb.size() * sizeof(float)); + } + compat::wait(); + + // ── Run kernel ──────────────────────────────────────────────────────────── + auto evt_args = K2FP8::make_evt_args( + block_scale_a.get(), M, + block_scale_b.get(), N); + + typename GemmOp::GemmKernel::EpilogueArguments epilogue_args{ + evt_args, nullptr, stride_C, block_D.get(), stride_D + }; + + typename GemmOp::GemmKernel::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, L}, + {block_A.get(), stride_A, block_B.get(), stride_B}, + epilogue_args, hw_info + }; + + GemmOp gemm_op; + size_t workspace_size = GemmOp::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + CUTLASS_CHECK(gemm_op.can_implement(arguments)); + CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get())); + CUTLASS_CHECK(gemm_op.run()); + compat::wait(); + + // ── Reference ───────────────────────────────────────────────────────────── + if (opts.verify) { + // FP8 → float32 upcast on GPU + { + const ElementFP8* a_ptr = block_A.get(); + const ElementFP8* b_ptr = block_B.get(); + float* af = block_A_f32.get(); + float* bf = block_B_f32.get(); + int64_t na = static_cast(M) * K * L; + int64_t nb = static_cast(K) * N * L; + + compat::get_default_queue().parallel_for(sycl::range<1>(na), [=](sycl::id<1> idx) { + af[idx[0]] = static_cast(a_ptr[idx[0]]); + }); + compat::get_default_queue().parallel_for(sycl::range<1>(nb), [=](sycl::id<1> idx) { + bf[idx[0]] = static_cast(b_ptr[idx[0]]); + }); + compat::wait(); + } + + // Float GEMM reference + { + const float* af = block_A_f32.get(); + const float* bf = block_B_f32.get(); + float* acc = block_acc_f32.get(); + int M_ = M, N_ = N, K_ = K, L_ = L; + compat::get_default_queue().parallel_for( + sycl::range<1>(static_cast(M) * N * L), + [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int col = static_cast(i % N_); + int row = static_cast((i / N_) % M_); + int batch = static_cast(i / ((int64_t)M_ * N_)); + float sum = 0.f; + for (int k = 0; k < K_; ++k) + sum += af[batch * M_ * K_ + row * K_ + k] + * bf[batch * K_ * N_ + k * N_ + col]; + acc[i] = sum; + } + ); + compat::wait(); + } + + // Dequant + SwiGLU reference + { + const float* acc = block_acc_f32.get(); + const float* sa = block_scale_a.get(); + const float* sb = block_scale_b.get(); + auto* ref = block_ref_D.get(); + int M_ = M, N_ = N, L_ = L; + compat::get_default_queue().parallel_for( + sycl::range<1>(static_cast(M) * N * L), + [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int col = static_cast(i % N_); + int row = static_cast((i / N_) % M_); + int batch = static_cast(i / ((int64_t)M_ * N_)); + int64_t base = static_cast(batch) * M_ * N_; + + int even_col = col & ~1; + int odd_col = even_col + 1; + if (odd_col >= N_) { + float v = acc[i] * sa[batch * M_ + row] * sb[batch * N_ + col]; + ref[i] = static_cast(v); + return; + } + float gate = acc[base + row * N_ + even_col] + * sa[batch * M_ + row] * sb[batch * N_ + even_col]; + float up = acc[base + row * N_ + odd_col] + * sa[batch * M_ + row] * sb[batch * N_ + odd_col]; + float silu_gate = gate / (1.f + sycl::exp(-gate)); + ref[i] = static_cast(silu_gate * up); + } + ); + compat::wait(); + } + + // FP8 rounding introduces ~0.5 ULP per multiply; tolerate 15% relative error + bool passed = cutlass::reference::device::BlockCompareRelativelyEqual( + block_ref_D.get(), block_D.get(), block_D.size(), + static_cast(0.15f), static_cast(0.05f)); + + std::cout << "Disposition: " << (passed ? "Passed" : "Failed") << std::endl; + if (!passed) return 1; + } else { + std::cout << "Disposition is skipped." << std::endl; + } + + if (opts.iterations > 0) { + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) gemm_op.run(); + compat::wait(); + + float time_s = timer.seconds() / opts.iterations; + double tflops = (2.0 * M * N * K * L) * 1e-12; + printf("Problem: %dx%dx%dx%d\n", M, N, K, L); + printf("xe-fuse K2_FP8 (FP8 GEMM+Dequant+SwiGLU): [%4.3f]TFlop/s (%6.4f)ms\n", + tflops / time_s, time_s * 1000); + } + + return 0; +} diff --git a/tests/test_k0g.cpp b/tests/test_k0g.cpp new file mode 100644 index 0000000..d3d6449 --- /dev/null +++ b/tests/test_k0g.cpp @@ -0,0 +1,178 @@ +// xe-fuse test: K0g -- gemm_gate_residual_norm +// D[m,n] = gamma[n] * (gate[m] * acc[m,n] + residual[m,n]) + +#include "xe-fuse/kernels/gemm_gate_residual_norm.hpp" + +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/device/gemm_complex.h" +#include "cutlass/util/reference/device/tensor_compare.h" + +#include "sycl_common.hpp" +#include "helper.h" + +#include +#include + +using namespace cute; + +struct Options { + int m = 512, n = 4096, k = 4096, l = 1; + int iterations = 100; + int verify = 1; + + void parse(int argc, char const** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", m, 512); + cmd.get_cmd_line_argument("n", n, 4096); + cmd.get_cmd_line_argument("k", k, 4096); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("iterations", iterations, 100); + cmd.get_cmd_line_argument("verify", verify, 1); + } +}; + +using K0g = xe_fuse::GemmGateResidualGamma<>; +using GemmOp = K0g::Gemm; + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + int M = opts.m, N = opts.n, K = opts.k, L = opts.l; + + using StrideA = typename GemmOp::GemmKernel::StrideA; + using StrideB = typename GemmOp::GemmKernel::StrideB; + + auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, make_shape(M, K, L)); + auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, make_shape(N, K, L)); + auto stride_C = cutlass::make_cute_packed_stride(K0g::StrideC{}, make_shape(M, N, L)); + auto stride_D = cutlass::make_cute_packed_stride(K0g::StrideD{}, make_shape(M, N, L)); + auto stride_residual = cutlass::make_cute_packed_stride(K0g::StrideResidual{}, make_shape(M, N, L)); + + cutlass::DeviceAllocation block_A(static_cast(M) * K * L); + cutlass::DeviceAllocation block_B(static_cast(K) * N * L); + cutlass::DeviceAllocation block_D(static_cast(M) * N * L); + cutlass::DeviceAllocation block_ref_D(static_cast(M) * N * L); + cutlass::DeviceAllocation block_residual(static_cast(M) * N * L); + cutlass::DeviceAllocation block_gate(static_cast(M) * L); + cutlass::DeviceAllocation block_gamma(static_cast(N)); + + initialize_block(block_A, 2023); + initialize_block(block_B, 2022); + initialize_block(block_residual, 2021); + + // gate: uniform (0.0, 1.0) to mimic timestep-dependent gating + { + std::mt19937 rng(77); + std::uniform_real_distribution dist(0.0f, 1.0f); + std::vector h(static_cast(M) * L); + for (auto& v : h) v = dist(rng); + compat::get_default_queue().memcpy(block_gate.get(), h.data(), h.size() * sizeof(float)); + } + + // gamma: uniform (0.5, 1.5) + { + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.5f, 1.5f); + std::vector h(static_cast(N)); + for (auto& v : h) v = dist(rng); + compat::get_default_queue().memcpy(block_gamma.get(), h.data(), h.size() * sizeof(float)); + } + compat::wait(); + + auto evt_args = K0g::make_evt_args( + block_gate.get(), M, + block_residual.get(), stride_residual, + block_gamma.get(), N); + + typename GemmOp::GemmKernel::EpilogueArguments epilogue_args{ + evt_args, nullptr, stride_C, block_D.get(), stride_D + }; + + typename GemmOp::GemmKernel::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, L}, + {block_A.get(), stride_A, block_B.get(), stride_B}, + epilogue_args, hw_info + }; + + GemmOp gemm_op; + size_t workspace_size = GemmOp::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + CUTLASS_CHECK(gemm_op.can_implement(arguments)); + CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get())); + CUTLASS_CHECK(gemm_op.run()); + compat::wait(); + + if (opts.verify) { + cutlass::DeviceAllocation block_gemm_f32(static_cast(M) * N * L); + cutlass::TensorRef ref_A(block_A.get(), cutlass::layout::RowMajor::packed({M, K})); + cutlass::TensorRef ref_B(block_B.get(), cutlass::layout::RowMajor::packed({K, N})); + cutlass::TensorRef ref_C_f32(block_gemm_f32.get(), cutlass::layout::RowMajor::packed({M, N})); + cutlass::TensorRef ref_D_f32(block_gemm_f32.get(), cutlass::layout::RowMajor::packed({M, N})); + + cutlass::reference::device::GemmComplex( + {M, N, K}, float(1), ref_A, cutlass::ComplexTransform::kNone, + ref_B, cutlass::ComplexTransform::kNone, float(0), + ref_C_f32, ref_D_f32, float(0), L, M * K, K * N, M * N, M * N); + compat::wait(); + + { + auto* ref_ptr = block_ref_D.get(); + auto* acc_ptr = block_gemm_f32.get(); + auto* res_ptr = block_residual.get(); + auto* gate_ptr = block_gate.get(); + auto* gamma_ptr = block_gamma.get(); + int n_val = N, m_val = M; + + compat::get_default_queue().parallel_for( + sycl::range<1>(static_cast(M) * N * L), + [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int col = static_cast(i % n_val); + int row = static_cast((i / n_val) % m_val); + int batch = static_cast(i / (static_cast(m_val) * n_val)); + + float acc = acc_ptr[i]; + float residual = static_cast(res_ptr[i]); + float gate = gate_ptr[batch * m_val + row]; + float gamma = gamma_ptr[col]; + + ref_ptr[i] = static_cast(gamma * (gate * acc + residual)); + } + ); + } + compat::wait(); + + bool passed = cutlass::reference::device::BlockCompareRelativelyEqual( + block_ref_D.get(), block_D.get(), block_D.size(), + K0g::ElementD(0.05f), K0g::ElementD(0.05f)); + + std::cout << "Disposition: " << (passed ? "Passed" : "Failed") << std::endl; + if (!passed) return 1; + } else { + std::cout << "Disposition is skipped." << std::endl; + } + + if (opts.iterations > 0) { + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) gemm_op.run(); + compat::wait(); + + float time_s = timer.seconds() / opts.iterations; + double tflops = (2.0 * M * N * K * L) * 1e-12; + std::cout << "Problem Size: " << M << 'x' << N << 'x' << K << 'x' << L << std::endl; + printf("xe-fuse K0g (GEMM+Gate+Residual+Gamma): [%4.3f]TFlop/s (%6.4f)ms\n", + tflops / time_s, time_s * 1000); + } + + return 0; +} diff --git a/tests/test_qk_norm_rope.cpp b/tests/test_qk_norm_rope.cpp new file mode 100644 index 0000000..bc6724a --- /dev/null +++ b/tests/test_qk_norm_rope.cpp @@ -0,0 +1,163 @@ +// xe-fuse test: QK norm + RoPE +// launch_qk_norm_rope: per-head RMSNorm followed by RoPE in a single kernel pass. + +#include "xe-fuse/kernels/qk_norm_rope.hpp" + +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/reference/device/tensor_compare.h" + +#include "sycl_common.hpp" +#include "helper.h" + +#include +#include +#include +#include + +using ElementQ = cutlass::bfloat16_t; +using ElementF = float; + +struct Options { + int seq_len = 4096; + int num_heads = 32; + int head_dim = 128; + int batch = 1; + int iterations = 100; + int verify = 1; + + void parse(int argc, char const** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("seq_len", seq_len, 4096); + cmd.get_cmd_line_argument("num_heads", num_heads, 32); + cmd.get_cmd_line_argument("head_dim", head_dim, 128); + cmd.get_cmd_line_argument("batch", batch, 1); + cmd.get_cmd_line_argument("iterations", iterations, 100); + cmd.get_cmd_line_argument("verify", verify, 1); + } +}; + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + int M = opts.batch * opts.seq_len; + int num_heads = opts.num_heads; + int head_dim = opts.head_dim; + + size_t input_size = static_cast(M) * num_heads * head_dim; + size_t cos_sin_size = static_cast(M) * head_dim; + size_t gamma_size = static_cast(head_dim); + + cutlass::DeviceAllocation block_in(input_size); + cutlass::DeviceAllocation block_out(input_size); + cutlass::DeviceAllocation block_ref(input_size); + cutlass::DeviceAllocation block_gamma(gamma_size); + cutlass::DeviceAllocation block_cos_sin(cos_sin_size); + + initialize_block(block_in, 2025); + sycl::queue q = compat::get_default_queue(); + + // gamma: uniform (0.5, 1.5) + { + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.5f, 1.5f); + std::vector h(gamma_size); + for (auto& v : h) v = dist(rng); + compat::get_default_queue().memcpy(block_gamma.get(), h.data(), h.size() * sizeof(float)); + } + + // cos_sin: realistic RoPE frequencies + { + std::vector h(cos_sin_size); + for (int tok = 0; tok < M; ++tok) { + for (int k = 0; k < head_dim / 2; ++k) { + float freq = 1.0f / std::pow(10000.f, 2.f * k / static_cast(head_dim)); + float angle = static_cast(tok) * freq; + int base = tok * head_dim; + h[base + 2 * k] = std::cos(angle); + h[base + 2 * k + 1] = std::sin(angle); + } + } + compat::get_default_queue().memcpy(block_cos_sin.get(), h.data(), h.size() * sizeof(float)); + } + compat::wait(); + + xe_fuse::launch_qk_norm_rope( + q, block_in.get(), block_out.get(), + block_gamma.get(), block_cos_sin.get(), + M, num_heads, head_dim); + compat::wait(); + + if (opts.verify) { + // CPU reference + std::vector h_in(input_size); + std::vector h_gamma(gamma_size); + std::vector h_cos_sin(cos_sin_size); + compat::get_default_queue().memcpy(h_in.data(), block_in.get(), input_size * sizeof(ElementQ)).wait(); + compat::get_default_queue().memcpy(h_gamma.data(), block_gamma.get(), gamma_size * sizeof(float)).wait(); + compat::get_default_queue().memcpy(h_cos_sin.data(), block_cos_sin.get(), cos_sin_size * sizeof(float)).wait(); + + std::vector h_ref(input_size); + constexpr float eps = 1e-6f; + std::vector normed(head_dim); + + for (int tok = 0; tok < M; ++tok) { + for (int h = 0; h < num_heads; ++h) { + int row_base = tok * num_heads * head_dim + h * head_dim; + int cs_base = tok * head_dim; + + // RMSNorm + float sum_sq = 0.f; + for (int d = 0; d < head_dim; ++d) { + float v = static_cast(h_in[row_base + d]); + sum_sq += v * v; + } + float rstd = 1.f / std::sqrt(sum_sq / head_dim + eps); + + // normalize + gamma + for (int d = 0; d < head_dim; ++d) + normed[d] = static_cast(h_in[row_base + d]) * rstd * h_gamma[d]; + + // RoPE + for (int d = 0; d < head_dim; d += 2) { + float x0 = normed[d], x1 = normed[d + 1]; + float cos_v = h_cos_sin[cs_base + d]; + float sin_v = h_cos_sin[cs_base + d + 1]; + h_ref[row_base + d] = static_cast( x0 * cos_v + x1 * sin_v); + h_ref[row_base + d + 1] = static_cast(-x0 * sin_v + x1 * cos_v); + } + } + } + compat::get_default_queue().memcpy(block_ref.get(), h_ref.data(), input_size * sizeof(ElementQ)).wait(); + + bool passed = cutlass::reference::device::BlockCompareRelativelyEqual( + block_ref.get(), block_out.get(), block_out.size(), + ElementQ(0.05f), ElementQ(0.05f)); + + std::cout << "Disposition: " << (passed ? "Passed" : "Failed") << std::endl; + if (!passed) return 1; + } else { + std::cout << "Disposition is skipped." << std::endl; + } + + if (opts.iterations > 0) { + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) + xe_fuse::launch_qk_norm_rope( + q, block_in.get(), block_out.get(), + block_gamma.get(), block_cos_sin.get(), + M, num_heads, head_dim); + compat::wait(); + + float time_s = timer.seconds() / opts.iterations; + double bytes = 2.0 * input_size * sizeof(ElementQ); // 1 read + 1 write + printf("Problem: M=%d num_heads=%d head_dim=%d\n", M, num_heads, head_dim); + printf("qk_norm_rope: [%4.3f]GB/s (%6.4f)ms\n", + bytes * 1e-9 / time_s, time_s * 1000); + } + + return 0; +}