diff --git a/README.md b/README.md index 265eee4..9dc32f8 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,9 @@ using Gemm = typename Kernel::Gemm; |-------|---------|-------| | `DequantW8A8` | `int32_acc * scale_token[m] * scale_channel[n]` | INT8 GEMM → bf16 output | | `DequantW8A8Biased` | `... + bias[n]` | Same with per-channel bias | +| `DequantRoPE` | dequant → RoPE rotation | K4 W8A8: Q/K projections | +| `DequantSwiGLU` | dequant → `silu(gate) * up` on adjacent pairs | K2 W8A8: FFN (LLaMA-style) | +| `DequantGeGLU` | dequant → `gelu(gate) * up` on adjacent pairs | K2 W8A8: FFN (Gemma-style) | INT8 GEMM uses `int8_t` A/B inputs, `int32_t` accumulator, `bf16` output, and `AlignmentAB=32` for 256-bit INT8 loads via the `XE_8x16x32_S32S8S8S32_TT` MMA atom. @@ -282,6 +285,90 @@ each expert gets its own scale vector and SwiGLU activation without extra launch | DeepSeek-V3 | 7168 | 2048 | 256 | 8 | 4096×7168 | | DBRX | 6144 | 10752 | 16 | 4 | 21504×6144 | +## W8A8 INT8 Quantization Pipeline + +W8A8 quantizes both weights (W) and activations (A) to INT8, then dequantizes inside the GEMM epilogue — the INT32 accumulator never touches global memory. Compared to BF16 this doubles arithmetic intensity and reaches the higher INT8 TOPS ceiling on BMG-G31. + +### High-level kernel structs + +Three ready-to-use kernel structs in `include/xe-fuse/kernels/`: + +```cpp +#include "xe-fuse/kernels/gemm_dequant_w8a8.hpp" +#include "xe-fuse/kernels/gemm_dequant_rope.hpp" +#include "xe-fuse/kernels/gemm_dequant_swiglu.hpp" + +// K0/K1: plain dequant (O and V projections) +using K0 = xe_fuse::GemmDequantW8A8<>; +auto evt = K0::make_evt_args(scale_token, M, scale_channel, N); + +// K4: dequant + RoPE (Q projection) +using K4 = xe_fuse::GemmDequantRoPE<>; +auto stride_cs = cutlass::make_cute_packed_stride(K4::StrideCosSin{}, make_shape(M, N, L)); +auto evt = K4::make_evt_args(scale_token, M, scale_channel, N, cos_sin, stride_cs); + +// K2: dequant + SwiGLU (FFN — LLaMA/Mistral/Qwen) +// For GeGLU (Gemma), use GemmDequantGeGLU<> instead +using K2 = xe_fuse::GemmDequantSwiGLU<>; +auto evt = K2::make_evt_args(scale_token, M, scale_channel, N); +``` + +Strides must be derived from the kernel's own types — CUTLASS XE StrideB has a compile-time leading element that differs from StrideA: + +```cpp +using StrideA = typename K0::Gemm::GemmKernel::StrideA; +using StrideB = typename K0::Gemm::GemmKernel::StrideB; +auto sA = cutlass::make_cute_packed_stride(StrideA{}, make_shape(M, K, L)); +auto sB = cutlass::make_cute_packed_stride(StrideB{}, make_shape(N, K, L)); // N first +``` + +### Activation quantization + +`launch_compute_rstd_and_quantize` combines RMSNorm and INT8 quantization in a single 3-pass subgroup kernel. The resulting `scale_token[m]` absorbs the RMSNorm reciprocal std — the GEMM epilogue multiplies by it directly with no separate normalization kernel needed: + +```cpp +#include "xe-fuse/kernels/compute_rstd.hpp" + +// x[M,N] → x_i8[M,N] + scale_token[M] +// scale_token[m] = max_n|x*rstd| / 127 (absorbs rstd into quant scale) +xe_fuse::launch_compute_rstd_and_quantize(q, x, x_i8, scale_token, M, N, L, eps); +``` + +### Standalone INT8 ops (for baselines) + +```cpp +#include "xe-fuse/standalone/ops.hpp" +#include "xe-fuse/standalone/vllm_ops.hpp" + +// Naive: separate dequant then op +xe_fuse::standalone::dequant_w8a8(q, out, i32_acc, scale_token, scale_channel, M, N, L); + +// vllm_int8_equiv: merged dequant + op in one kernel +xe_fuse::vllm_equiv::dequant_and_rotary_embedding(q, out, i32_acc, st, sc, cos_sin, M, N); +xe_fuse::vllm_equiv::dequant_and_silu_mul(q, out, i32_acc, st, sc, I, M); +xe_fuse::vllm_equiv::dequant_and_gelu_mul(q, out, i32_acc, st, sc, I, M); +``` + +### Three-way pipeline benchmark + +`generate_pipeline.py --int8-mode w8a8` generates a benchmark comparing three implementations across the full attention+FFN pipeline: + +| Variant | Description | +|---------|-------------| +| `XE_W8A8_FUSED` | INT8 GEMM with fused dequant epilogue — INT32 acc stays in registers | +| `VLLM_INT8_EQUIV` | INT8 GEMM (INT32 to DRAM) + merged dequant+op kernels | +| `NAIVE_INT8` | INT8 GEMM (INT32 to DRAM) + separate dequant + separate ops | + +```bash +# Generate C++ for a model preset +python3 autotune/generate_pipeline.py --preset llama3_8b --int8-mode w8a8 -o /tmp/w8a8.cpp + +# Or run via sbatch (correctness check + benchmarks at multiple sequence lengths): +sbatch tests/run_w8a8_pipeline.sh llama3_8b +``` + +At `--iterations=0` the binary runs a per-kernel float reference comparison and exits without benchmarking. + ## Autotune Tile Selection xe-fuse includes a tile shape auto-selector based on empirical sweep data (792 configurations diff --git a/autotune/generate_pipeline.py b/autotune/generate_pipeline.py index b770223..6b05c9c 100644 --- a/autotune/generate_pipeline.py +++ b/autotune/generate_pipeline.py @@ -28,8 +28,12 @@ def generate_pipeline_cpp( preset_name: str = "custom", seq_len: int = 2048, autotune: bool = False, + int8_mode: str = "none", ) -> str: - template_path = Path(__file__).parent / "pipeline_template.cpp.j2" + if int8_mode == "w8a8": + template_path = Path(__file__).parent / "pipeline_w8a8_template.cpp.j2" + else: + template_path = Path(__file__).parent / "pipeline_template.cpp.j2" if not template_path.exists(): print(f"ERROR: Template not found at {template_path}", file=sys.stderr) @@ -78,6 +82,7 @@ def generate_pipeline_cpp( ffn_activation=config["ffn_activation"], gated_ffn=config["gated_ffn"], pipeline_desc=" -> ".join(pipeline_parts), + int8_mode=int8_mode, **tile_vars, ) @@ -102,6 +107,12 @@ def main(): default=2048, help="Target sequence length for tile selection (default: 2048)", ) + parser.add_argument( + "--int8-mode", + choices=["none", "w8a8"], + default="none", + help="INT8 quantization mode: none=BF16 pipeline (default), w8a8=W8A8 INT8 pipeline", + ) parser.add_argument( "--list-presets", action="store_true", @@ -128,7 +139,8 @@ def main(): return code = generate_pipeline_cpp( - config, preset_name, seq_len=args.seq_len, autotune=args.autotune + config, preset_name, seq_len=args.seq_len, autotune=args.autotune, + int8_mode=args.int8_mode, ) Path(args.output).parent.mkdir(parents=True, exist_ok=True) @@ -141,6 +153,7 @@ def main(): print( f" Dims: H={config['H']}, H_kv={config['H_kv']}, I={config['I']}, N_ffn={n_ffn}" ) + print(f" Mode: {args.int8_mode.upper() if args.int8_mode != 'none' else 'BF16'}") print(f" Q/K: {'K4 (RMSNorm+RoPE)' if config['use_rope'] else 'K1 (RMSNorm)'}") print(f" FFN: {config['ffn_activation']}") if args.autotune: diff --git a/autotune/pipeline_w8a8_template.cpp.j2 b/autotune/pipeline_w8a8_template.cpp.j2 new file mode 100644 index 0000000..cb49441 --- /dev/null +++ b/autotune/pipeline_w8a8_template.cpp.j2 @@ -0,0 +1,692 @@ +// Auto-generated xe-fuse W8A8 pipeline benchmark +// Model: {{ model_name }} +// Generated: {{ timestamp }} +// Architecture: H={{ default_H }}, H_kv={{ default_H_kv }}, I={{ default_I }} +// Q/K kernel: {{ "K4_W8A8 (DequantRoPE)" if use_rope else "K1_W8A8 (DequantW8A8)" }} +// FFN activation: {{ ffn_activation }} +// +// Three-way comparison: +// XE_W8A8_FUSED — INT8 GEMM with fused dequant epilogue (INT32 acc stays in register) +// VLLM_INT8_EQUIV — INT8 GEMM (INT32 to DRAM) + merged dequant+op kernels +// NAIVE_INT8 — INT8 GEMM (INT32 to DRAM) + separate dequant + separate ops + +#include "xe-fuse/kernels/gemm_dequant_w8a8.hpp" +#include "xe-fuse/kernels/gemm_dequant_swiglu.hpp" +#include "xe-fuse/kernels/compute_rstd.hpp" +#include "xe-fuse/standalone/ops.hpp" +#include "xe-fuse/standalone/vllm_ops.hpp" +{% if use_rope %} +#include "xe-fuse/kernels/gemm_dequant_rope.hpp" +{% endif %} + +#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 "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" + +#include "sycl_common.hpp" +#include "helper.h" + +#include +#include +#include +#include +#include +#include + +using namespace cute; +using bf16 = cutlass::bfloat16_t; + +// ============================================================ +// Kernel type aliases +// ============================================================ + +// K4_W8A8: Q projection — INT8 GEMM + dequant (+ RoPE if applicable) +{% if use_rope %} +using K4_W8A8 = xe_fuse::GemmDequantRoPE<>; +{% else %} +using K4_W8A8 = xe_fuse::GemmDequantW8A8<>; +{% endif %} + +// K1_W8A8: V projection — INT8 GEMM + dequant +using K1_W8A8 = xe_fuse::GemmDequantW8A8<>; + +// K0_W8A8: O projection — INT8 GEMM + dequant +using K0_W8A8 = xe_fuse::GemmDequantW8A8<>; + +// K2_W8A8: FFN projection — INT8 GEMM + dequant + activation +{% if ffn_activation == "swiglu" %} +using K2_W8A8 = xe_fuse::GemmDequantSwiGLU<>; +{% elif ffn_activation == "geglu" %} +using K2_W8A8 = xe_fuse::GemmDequantGeGLU<>; +{% else %} +using K2_W8A8 = xe_fuse::GemmDequantW8A8<>; +{% endif %} + +// BareInt8Gemm: INT8→INT32 GEMM with no epilogue fusion (baseline) +// AlignmentCD=4: 128-bit loads for int32_t (32 bits × 4 = 128 bits) +namespace bare_i8 { +using ElementA = int8_t; +using ElementB = int8_t; +using ElementD = int32_t; +using ElementAcc = int32_t; +using ElementCompute = float; +using LayoutA = cutlass::layout::RowMajor; +using LayoutB = cutlass::layout::RowMajor; +using StrideC = cute::Stride, int64_t>; +using StrideD = cute::Stride, int64_t>; +using TileShape = cute::Shape; + +using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + TileShape, cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAcc, ElementCompute, + ElementD, StrideC, 4, + ElementD, StrideD, 4, + cutlass::epilogue::collective::EpilogueScheduleAuto +>::CollectiveOp; + +using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 32, + ElementB, LayoutB, 32, + 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; +} // namespace bare_i8 + +// ============================================================ + +struct Options { + int m = 2048; + int h = {{ default_H }}; + int hkv = {{ default_H_kv }}; + int inter = {{ default_I }}; + int l = 1; + int iterations = 0; + float eps = 1e-6f; + + void parse(int argc, const char** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", m, 2048); + cmd.get_cmd_line_argument("h", h, {{ default_H }}); + cmd.get_cmd_line_argument("hkv", hkv, {{ default_H_kv }}); + cmd.get_cmd_line_argument("inter",inter,{{ default_I }}); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("iterations", iterations, 0); + } +}; + +// Per-column INT8 weight quantization (offline, not benchmarked) +// scale[n] = max_k|W[k,n]| / 127; W_i8[k,n] = round(clamp(W[k,n]/scale[n], -128,127)) +static void quantize_weight_cols(sycl::queue& q, + bf16 const* W_bf16, int8_t* W_i8, float* scale, + int K, int N) { + q.parallel_for(sycl::range<1>(N), [=](sycl::id<1> idx) { + int n = static_cast(idx[0]); + float mx = 0.f; + for (int k = 0; k < K; ++k) + mx = sycl::fmax(mx, sycl::fabs(static_cast(W_bf16[k * N + n]))); + scale[n] = mx / 127.f + 1e-8f; + }); + q.wait(); + q.parallel_for(sycl::range<1>(static_cast(K) * N), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int n = static_cast(i % N); + float v = sycl::round(static_cast(W_bf16[i]) / scale[n]); + W_i8[i] = static_cast(sycl::fmin(sycl::fmax(v, -128.f), 127.f)); + }); + q.wait(); +} + +// ─── Float reference helpers for numerical correctness ─────────────────────── +// B stored [K,N] RowMajor: B[k,n] = B[k*N+n] (matches CUTLASS XE StrideB convention) +// Fills acc[L,M,N] = sum_k float(A[L,M,K]) * float(B[K,N]) +static void ref_float_gemm_acc(sycl::queue& q, float* acc, + int8_t const* A, int8_t const* B, int M, int N, int K, int L) { + q.parallel_for(sycl::range<1>(static_cast(M) * N * L), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int l = static_cast(i / (M * N)); + int m = static_cast((i / N) % M); + int n = static_cast(i % N); + float sum = 0.f; + for (int k = 0; k < K; ++k) + sum += static_cast(A[l * M * K + m * K + k]) + * static_cast(B[k * N + n]); + acc[i] = sum; + }); +} + +// out[l,m,n] = bf16(acc[l,m,n] * scale_tok[l*M+m] * scale_ch[l*N+n]) +static void ref_apply_dequant(sycl::queue& q, bf16* out, float const* acc, + float const* scale_tok, float const* scale_ch, int M, int N, int L) { + q.parallel_for(sycl::range<1>(static_cast(M) * N * L), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int l = static_cast(i / (M * N)); + int m = static_cast((i / N) % M); + int n = static_cast(i % N); + out[i] = static_cast(acc[i] * scale_tok[l * M + m] * scale_ch[l * N + n]); + }); +} + +// Dequant + RoPE (adjacent pairs: even=cos, odd=sin rotation) +// Reads acc (float, no in-place race) and writes bf16 to out +static void ref_apply_dequant_rope(sycl::queue& q, bf16* out, + float const* acc, float const* scale_tok, float const* scale_ch, + float const* cos_sin, int M, int N, int L) { + q.parallel_for(sycl::range<1>(static_cast(M) * N * L), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int l = static_cast(i / (M * N)); + int m = static_cast((i / N) % M); + int n = static_cast(i % N); + int64_t base = static_cast(l) * M * N; + int even_n = n & ~1, odd_n = even_n + 1; + if (odd_n >= N) { + out[i] = static_cast(acc[i] * scale_tok[l * M + m] * scale_ch[l * N + n]); + return; + } + float x_even = acc[base + m * N + even_n] * scale_tok[l * M + m] * scale_ch[l * N + even_n]; + float x_odd = acc[base + m * N + odd_n] * scale_tok[l * M + m] * scale_ch[l * N + odd_n]; + float cos_v = cos_sin[base + m * N + even_n]; + float sin_v = cos_sin[base + m * N + odd_n]; + out[i] = static_cast((n & 1) == 0 + ? x_even * cos_v + x_odd * sin_v + : -x_even * sin_v + x_odd * cos_v); + }); +} + +// Dequant + SwiGLU: adjacent pairs (even=gate, odd=up) +// out[m, 2k] = out[m, 2k+1] = silu(dequant[m,2k]) * dequant[m,2k+1] +static void ref_apply_dequant_swiglu(sycl::queue& q, bf16* out, + float const* acc, float const* scale_tok, float const* scale_ch, + int M, int N, int L) { + q.parallel_for(sycl::range<1>(static_cast(M) * N * L), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int l = static_cast(i / (M * N)); + int m = static_cast((i / N) % M); + int n = static_cast(i % N); + int64_t base = static_cast(l) * M * N; + int even_n = n & ~1, odd_n = even_n + 1; + float gate = acc[base + m * N + even_n] * scale_tok[l * M + m] * scale_ch[l * N + even_n]; + float up = acc[base + m * N + odd_n] * scale_tok[l * M + m] * scale_ch[l * N + odd_n]; + out[i] = static_cast(gate / (1.f + sycl::exp(-gate)) * up); + }); +} + +// Dequant + GeGLU: same layout as SwiGLU but uses GELU activation +static void ref_apply_dequant_geglu(sycl::queue& q, bf16* out, + float const* acc, float const* scale_tok, float const* scale_ch, + int M, int N, int L) { + q.parallel_for(sycl::range<1>(static_cast(M) * N * L), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int l = static_cast(i / (M * N)); + int m = static_cast((i / N) % M); + int n = static_cast(i % N); + int64_t base = static_cast(l) * M * N; + int even_n = n & ~1, odd_n = even_n + 1; + float gate = acc[base + m * N + even_n] * scale_tok[l * M + m] * scale_ch[l * N + even_n]; + float up = acc[base + m * N + odd_n] * scale_tok[l * M + m] * scale_ch[l * N + odd_n]; + float gelu = gate * 0.5f * (1.f + sycl::erf(gate * 0.7071067811865476f)); + out[i] = static_cast(gelu * up); + }); +} + +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, H = opts.h, H_kv = opts.hkv, I = opts.inter, L = opts.l; +{% if gated_ffn %} + int N_ffn = 2 * I; +{% else %} + int N_ffn = I; +{% endif %} + auto q = compat::get_default_queue(); + + std::cout << "============================================================" << std::endl; + std::cout << "xe-fuse W8A8: {{ model_name }}" << std::endl; + std::cout << " M=" << M << " H=" << H << " H_kv=" << H_kv + << " I=" << I << " N_ffn=" << N_ffn << std::endl; + std::cout << "============================================================" << std::endl; + + // === Buffer sizes === + size_t mh = static_cast(M) * H * L; + size_t mhkv = static_cast(M) * H_kv * L; + size_t hh = static_cast(H) * H * L; + size_t h_hkv = static_cast(H) * H_kv * L; + size_t h_nffn = static_cast(H) * N_ffn * L; + size_t m_nffn = static_cast(M) * N_ffn * L; + size_t max_mn = std::max({mh, mhkv, m_nffn}); + + // === BF16 activations + weights === + cutlass::DeviceAllocation x(mh); + cutlass::DeviceAllocation W_q(hh), W_v(h_hkv), W_o(hh), W_ffn(h_nffn); + cutlass::DeviceAllocation attn_out(mh); + cutlass::DeviceAllocation gamma(static_cast(H) * L); +{% if use_rope %} + cutlass::DeviceAllocation cos_sin(mh); +{% endif %} + + // === INT8 weights + per-column scales === + cutlass::DeviceAllocation W_q_i8(hh), W_v_i8(h_hkv); + cutlass::DeviceAllocation W_o_i8(hh), W_ffn_i8(h_nffn); + cutlass::DeviceAllocation scale_q_ch(static_cast(H) * L); + cutlass::DeviceAllocation scale_v_ch(static_cast(H_kv) * L); + cutlass::DeviceAllocation scale_o_ch(static_cast(H) * L); + cutlass::DeviceAllocation scale_ffn_ch(static_cast(N_ffn) * L); + + // === INT8 activations + per-token scales === + cutlass::DeviceAllocation x_i8(mh), attn_i8(mh), raw_sum_i8(mh); + cutlass::DeviceAllocation scale_tok1(static_cast(M) * L); + cutlass::DeviceAllocation scale_tok_attn(static_cast(M) * L); + cutlass::DeviceAllocation scale_tok_raw(static_cast(M) * L); + + // === W8A8 pipeline outputs === + cutlass::DeviceAllocation Q_out(mh), V_out(mhkv), o_proj(mh); + cutlass::DeviceAllocation residual1(mh), raw_sum(mh), ffn_out(m_nffn); + + // === INT32 scratch for bare baselines === + cutlass::DeviceAllocation i32_scratch(max_mn); + + // === Random init === + initialize_block(x, 2023); + initialize_block(W_q, 2024); + initialize_block(W_v, 2025); + initialize_block(W_o, 2026); + initialize_block(W_ffn, 2027); + initialize_block(attn_out,2028); + { + std::vector hg(static_cast(H) * L); + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.5f, 1.5f); + for (auto& v : hg) v = dist(rng); + q.memcpy(gamma.get(), hg.data(), hg.size() * sizeof(float)).wait(); + } +{% if use_rope %} + { + std::vector hcs(mh); + std::mt19937 rng(2024); + std::uniform_real_distribution dist(-1.f, 1.f); + for (auto& v : hcs) v = dist(rng); + q.memcpy(cos_sin.get(), hcs.data(), hcs.size() * sizeof(float)).wait(); + } +{% endif %} + compat::wait(); + + // === Offline weight quantization === + std::cout << "\n=== Quantizing weights ===" << std::endl; + quantize_weight_cols(q, W_q.get(), W_q_i8.get(), scale_q_ch.get(), H, H); + quantize_weight_cols(q, W_v.get(), W_v_i8.get(), scale_v_ch.get(), H, H_kv); + quantize_weight_cols(q, W_o.get(), W_o_i8.get(), scale_o_ch.get(), H, H); + quantize_weight_cols(q, W_ffn.get(), W_ffn_i8.get(), scale_ffn_ch.get(), H, N_ffn); + + // === Activation quantization (RMSNorm + per-token INT8) === + xe_fuse::launch_compute_rstd_and_quantize(q, x.get(), x_i8.get(), scale_tok1.get(), M, H, L, opts.eps); + xe_fuse::launch_compute_rstd_and_quantize(q, attn_out.get(), attn_i8.get(), scale_tok_attn.get(), M, H, L, opts.eps); + compat::wait(); + + // === Strides — use kernel-derived types to match CUTLASS XE conventions === + // A strides (activations): RowMajor M×K → (K, 1, M*K) + // B strides (weights): RowMajor — CUTLASS XE uses (1, K, N*K) for B + // Use GemmKernel::StrideA/StrideB to get correct types, same as existing tests. + using StrideA_F = typename K0_W8A8::Gemm::GemmKernel::StrideA; // all W8A8 kernels share layout + using StrideB_F = typename K0_W8A8::Gemm::GemmKernel::StrideB; + using StrideC_F = K0_W8A8::StrideC; + using StrideD_F = K0_W8A8::StrideD; + using StrideA_B = typename bare_i8::GemmKernel::StrideA; + using StrideB_B = typename bare_i8::GemmKernel::StrideB; + using StrideD_B = bare_i8::StrideD; + + // A: M×K activations (row-major, shape = M, K, L) + auto sA_mh = cutlass::make_cute_packed_stride(StrideA_F{}, make_shape(M, H, L)); + auto sA_mh_b = cutlass::make_cute_packed_stride(StrideA_B{}, make_shape(M, H, L)); + + // B: N×K weights (shape = N, K, L — N is output dim, K is input dim) + auto sB_hh = cutlass::make_cute_packed_stride(StrideB_F{}, make_shape(H, H, L)); + auto sB_hkv = cutlass::make_cute_packed_stride(StrideB_F{}, make_shape(H_kv, H, L)); + auto sB_hffn = cutlass::make_cute_packed_stride(StrideB_F{}, make_shape(N_ffn, H, L)); + auto sB_hh_b = cutlass::make_cute_packed_stride(StrideB_B{}, make_shape(H, H, L)); + auto sB_hkv_b = cutlass::make_cute_packed_stride(StrideB_B{}, make_shape(H_kv, H, L)); + auto sB_hffn_b= cutlass::make_cute_packed_stride(StrideB_B{}, make_shape(N_ffn, H, L)); + + // D: M×N output (row-major) + auto sD_mh = cutlass::make_cute_packed_stride(StrideD_F{}, make_shape(M, H, L)); + auto sD_mhkv = cutlass::make_cute_packed_stride(StrideD_F{}, make_shape(M, H_kv, L)); + auto sD_mffn = cutlass::make_cute_packed_stride(StrideD_F{}, make_shape(M, N_ffn, L)); + auto sD_mh_b = cutlass::make_cute_packed_stride(StrideD_B{}, make_shape(M, H, L)); + auto sD_mhkv_b= cutlass::make_cute_packed_stride(StrideD_B{}, make_shape(M, H_kv, L)); + auto sD_mffn_b= cutlass::make_cute_packed_stride(StrideD_B{}, make_shape(M, N_ffn, L)); + + // ============================================================ + // W8A8 FUSED GEMM ops (workspaces kept alive for the entire run) + // ============================================================ + typename K4_W8A8::Gemm k4_op; + cutlass::device_memory::allocation k4_ws; + { +{% if use_rope %} + auto k4_stride_cs = cutlass::make_cute_packed_stride(K4_W8A8::StrideCosSin{}, make_shape(M, H, L)); + auto k4_evt = K4_W8A8::make_evt_args(scale_tok1.get(), M, scale_q_ch.get(), H, cos_sin.get(), k4_stride_cs); +{% else %} + auto k4_evt = K4_W8A8::make_evt_args(scale_tok1.get(), M, scale_q_ch.get(), H); +{% endif %} + typename K4_W8A8::Gemm::GemmKernel::EpilogueArguments epi{k4_evt, nullptr, sD_mh, Q_out.get(), sD_mh}; + typename K4_W8A8::Gemm::GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H, H, L}, + {x_i8.get(), sA_mh, W_q_i8.get(), sB_hh}, + epi, hw_info}; + k4_ws = cutlass::device_memory::allocation(K4_W8A8::Gemm::get_workspace_size(args)); + CUTLASS_CHECK(k4_op.can_implement(args)); + CUTLASS_CHECK(k4_op.initialize(args, k4_ws.get())); + } + + typename K1_W8A8::Gemm k1_op; + cutlass::device_memory::allocation k1_ws; + { + auto k1_evt = K1_W8A8::make_evt_args(scale_tok1.get(), M, scale_v_ch.get(), H_kv); + typename K1_W8A8::Gemm::GemmKernel::EpilogueArguments epi{k1_evt, nullptr, sD_mhkv, V_out.get(), sD_mhkv}; + typename K1_W8A8::Gemm::GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H_kv, H, L}, + {x_i8.get(), sA_mh, W_v_i8.get(), sB_hkv}, + epi, hw_info}; + k1_ws = cutlass::device_memory::allocation(K1_W8A8::Gemm::get_workspace_size(args)); + CUTLASS_CHECK(k1_op.can_implement(args)); + CUTLASS_CHECK(k1_op.initialize(args, k1_ws.get())); + } + + typename K0_W8A8::Gemm k0_op; + cutlass::device_memory::allocation k0_ws; + { + auto k0_evt = K0_W8A8::make_evt_args(scale_tok_attn.get(), M, scale_o_ch.get(), H); + typename K0_W8A8::Gemm::GemmKernel::EpilogueArguments epi{k0_evt, nullptr, sD_mh, o_proj.get(), sD_mh}; + typename K0_W8A8::Gemm::GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H, H, L}, + {attn_i8.get(), sA_mh, W_o_i8.get(), sB_hh}, + epi, hw_info}; + k0_ws = cutlass::device_memory::allocation(K0_W8A8::Gemm::get_workspace_size(args)); + CUTLASS_CHECK(k0_op.can_implement(args)); + CUTLASS_CHECK(k0_op.initialize(args, k0_ws.get())); + } + + typename K2_W8A8::Gemm k2_op; + cutlass::device_memory::allocation k2_ws; + { + auto k2_evt = K2_W8A8::make_evt_args(scale_tok_raw.get(), M, scale_ffn_ch.get(), N_ffn); + typename K2_W8A8::Gemm::GemmKernel::EpilogueArguments epi{k2_evt, nullptr, sD_mffn, ffn_out.get(), sD_mffn}; + typename K2_W8A8::Gemm::GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N_ffn, H, L}, + {raw_sum_i8.get(), sA_mh, W_ffn_i8.get(), sB_hffn}, + epi, hw_info}; + k2_ws = cutlass::device_memory::allocation(K2_W8A8::Gemm::get_workspace_size(args)); + CUTLASS_CHECK(k2_op.can_implement(args)); + CUTLASS_CHECK(k2_op.initialize(args, k2_ws.get())); + } + + // Helper: run full W8A8 fused pipeline once + auto run_fused = [&]() { + xe_fuse::launch_compute_rstd_and_quantize(q, x.get(), x_i8.get(), scale_tok1.get(), M, H, L, opts.eps); + k4_op.run(); k1_op.run(); compat::wait(); + k0_op.run(); compat::wait(); + // Residual add: residual1 = o_proj + x + q.memcpy(residual1.get(), o_proj.get(), mh * sizeof(bf16)).wait(); + xe_fuse::standalone::add_residual(q, residual1.get(), x.get(), M, H, L); + // Copy pre-gamma sum, apply gamma, then RMSNorm+quant for FFN input + q.memcpy(raw_sum.get(), residual1.get(), mh * sizeof(bf16)).wait(); + xe_fuse::standalone::scale_cols(q, raw_sum.get(), gamma.get(), M, H, L); + xe_fuse::launch_compute_rstd_and_quantize(q, raw_sum.get(), raw_sum_i8.get(), scale_tok_raw.get(), M, H, L, opts.eps); + k2_op.run(); compat::wait(); + }; + + // ============================================================ + // Bare INT8 GEMM ops (for baselines) + // ============================================================ + using BareGemm = bare_i8::Gemm; + using BareK = bare_i8::GemmKernel; + + auto make_bare_args = [&](int Mm, int Nn, int Kk, + int8_t const* A, StrideA_B sA, + int8_t const* B, StrideB_B sB, + int32_t* D, StrideD_B sD) + -> typename BareK::Arguments { + typename bare_i8::CollectiveEpilogue::Arguments epi_args; + epi_args.thread.alpha = float(1); + epi_args.thread.beta = float(0); + epi_args.ptr_C = nullptr; + epi_args.dC = sD; + epi_args.ptr_D = D; + epi_args.dD = sD; + return typename BareK::Arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, {Mm, Nn, Kk, L}, + {A, sA, B, sB}, + epi_args, hw_info}; + }; + + BareGemm bq_op, bv_op, bo_op, bffn_op; + cutlass::device_memory::allocation bq_ws, bv_ws, bo_ws, bffn_ws; + { + auto args = make_bare_args(M, H, H, x_i8.get(), sA_mh_b, W_q_i8.get(), sB_hh_b, i32_scratch.get(), sD_mh_b); + bq_ws = cutlass::device_memory::allocation(BareGemm::get_workspace_size(args)); + CUTLASS_CHECK(bq_op.can_implement(args)); + CUTLASS_CHECK(bq_op.initialize(args, bq_ws.get())); + } + { + auto args = make_bare_args(M, H_kv, H, x_i8.get(), sA_mh_b, W_v_i8.get(), sB_hkv_b, i32_scratch.get(), sD_mhkv_b); + bv_ws = cutlass::device_memory::allocation(BareGemm::get_workspace_size(args)); + CUTLASS_CHECK(bv_op.can_implement(args)); + CUTLASS_CHECK(bv_op.initialize(args, bv_ws.get())); + } + { + auto args = make_bare_args(M, H, H, attn_i8.get(), sA_mh_b, W_o_i8.get(), sB_hh_b, i32_scratch.get(), sD_mh_b); + bo_ws = cutlass::device_memory::allocation(BareGemm::get_workspace_size(args)); + CUTLASS_CHECK(bo_op.can_implement(args)); + CUTLASS_CHECK(bo_op.initialize(args, bo_ws.get())); + } + { + auto args = make_bare_args(M, N_ffn, H, raw_sum_i8.get(), sA_mh_b, W_ffn_i8.get(), sB_hffn_b, i32_scratch.get(), sD_mffn_b); + bffn_ws = cutlass::device_memory::allocation(BareGemm::get_workspace_size(args)); + CUTLASS_CHECK(bffn_op.can_implement(args)); + CUTLASS_CHECK(bffn_op.initialize(args, bffn_ws.get())); + } + + // Helper: run full vllm_int8_equiv pipeline once + auto run_vllm = [&]() { + xe_fuse::launch_compute_rstd_and_quantize(q, x.get(), x_i8.get(), scale_tok1.get(), M, H, L, opts.eps); + bq_op.run(); compat::wait(); +{% if use_rope %} + xe_fuse::vllm_equiv::dequant_and_rotary_embedding(q, Q_out.get(), i32_scratch.get(), scale_tok1.get(), scale_q_ch.get(), cos_sin.get(), M, H); +{% else %} + xe_fuse::standalone::dequant_w8a8(q, Q_out.get(), i32_scratch.get(), scale_tok1.get(), scale_q_ch.get(), M, H, L); +{% endif %} + bv_op.run(); compat::wait(); + xe_fuse::standalone::dequant_w8a8(q, V_out.get(), i32_scratch.get(), scale_tok1.get(), scale_v_ch.get(), M, H_kv, L); + bo_op.run(); compat::wait(); + // vllm uses fused dequant+residual in one pass (merged kernel) + xe_fuse::standalone::dequant_w8a8(q, o_proj.get(), i32_scratch.get(), scale_tok_attn.get(), scale_o_ch.get(), M, H, L); + q.memcpy(residual1.get(), o_proj.get(), mh * sizeof(bf16)).wait(); + xe_fuse::standalone::add_residual(q, residual1.get(), x.get(), M, H, L); + q.memcpy(raw_sum.get(), residual1.get(), mh * sizeof(bf16)).wait(); + xe_fuse::standalone::scale_cols(q, raw_sum.get(), gamma.get(), M, H, L); + xe_fuse::launch_compute_rstd_and_quantize(q, raw_sum.get(), raw_sum_i8.get(), scale_tok_raw.get(), M, H, L, opts.eps); + bffn_op.run(); compat::wait(); +{% if ffn_activation == "swiglu" %} + xe_fuse::vllm_equiv::dequant_and_silu_mul(q, ffn_out.get(), i32_scratch.get(), scale_tok_raw.get(), scale_ffn_ch.get(), I, M); +{% elif ffn_activation == "geglu" %} + xe_fuse::vllm_equiv::dequant_and_gelu_mul(q, ffn_out.get(), i32_scratch.get(), scale_tok_raw.get(), scale_ffn_ch.get(), I, M); +{% else %} + xe_fuse::standalone::dequant_w8a8(q, ffn_out.get(), i32_scratch.get(), scale_tok_raw.get(), scale_ffn_ch.get(), M, N_ffn, L); +{% endif %} + compat::wait(); + }; + + // Helper: run full naive_int8 pipeline once + auto run_naive = [&]() { + xe_fuse::launch_compute_rstd_and_quantize(q, x.get(), x_i8.get(), scale_tok1.get(), M, H, L, opts.eps); + bq_op.run(); compat::wait(); + xe_fuse::standalone::dequant_w8a8(q, Q_out.get(), i32_scratch.get(), scale_tok1.get(), scale_q_ch.get(), M, H, L); + bv_op.run(); compat::wait(); + xe_fuse::standalone::dequant_w8a8(q, V_out.get(), i32_scratch.get(), scale_tok1.get(), scale_v_ch.get(), M, H_kv, L); + bo_op.run(); compat::wait(); + xe_fuse::standalone::dequant_w8a8(q, o_proj.get(), i32_scratch.get(), scale_tok_attn.get(), scale_o_ch.get(), M, H, L); + q.memcpy(residual1.get(), o_proj.get(), mh * sizeof(bf16)).wait(); + xe_fuse::standalone::add_residual(q, residual1.get(), x.get(), M, H, L); + q.memcpy(raw_sum.get(), residual1.get(), mh * sizeof(bf16)).wait(); + xe_fuse::standalone::scale_cols(q, raw_sum.get(), gamma.get(), M, H, L); + xe_fuse::launch_compute_rstd_and_quantize(q, raw_sum.get(), raw_sum_i8.get(), scale_tok_raw.get(), M, H, L, opts.eps); + bffn_op.run(); compat::wait(); + xe_fuse::standalone::dequant_w8a8(q, ffn_out.get(), i32_scratch.get(), scale_tok_raw.get(), scale_ffn_ch.get(), M, N_ffn, L); +{% if ffn_activation == "swiglu" %} + xe_fuse::standalone::swiglu(q, ffn_out.get(), M, N_ffn, L); +{% elif ffn_activation == "geglu" %} + xe_fuse::standalone::geglu(q, ffn_out.get(), M, N_ffn, L); +{% endif %} + compat::wait(); + }; + + // ============================================================ + // Correctness: float reference GEMM comparison (rtol=0.15) + // run_fused() populates x_i8/scale_tok1, attn_i8/scale_tok_attn, + // raw_sum_i8/scale_tok_raw which we then re-use as reference inputs. + // ============================================================ + std::cout << "\n=== Correctness (float reference, rtol=0.15) ===" << std::endl; + run_fused(); + bool all_ok = true; + + { + size_t max_ref = std::max({mh, mhkv, m_nffn}); + cutlass::DeviceAllocation f_acc(max_ref); + cutlass::DeviceAllocation ref_out(max_ref); + + // K1_W8A8: V projection x_i8[M,H] @ W_v_i8[H,H_kv] * scale_tok1 * scale_v_ch + ref_float_gemm_acc(q, f_acc.get(), x_i8.get(), W_v_i8.get(), M, H_kv, H, L); + q.wait(); + ref_apply_dequant(q, ref_out.get(), f_acc.get(), scale_tok1.get(), scale_v_ch.get(), M, H_kv, L); + q.wait(); + { + bool ok = cutlass::reference::device::BlockCompareRelativelyEqual( + ref_out.get(), V_out.get(), mhkv, + static_cast(0.15f), static_cast(0.05f)); + std::cout << " K1_W8A8 V: " << (ok ? "PASS (rtol=0.15)" : "FAIL") << std::endl; + all_ok &= ok; + } + + // K0_W8A8: O projection attn_i8[M,H] @ W_o_i8[H,H] * scale_tok_attn * scale_o_ch + ref_float_gemm_acc(q, f_acc.get(), attn_i8.get(), W_o_i8.get(), M, H, H, L); + q.wait(); + ref_apply_dequant(q, ref_out.get(), f_acc.get(), scale_tok_attn.get(), scale_o_ch.get(), M, H, L); + q.wait(); + { + bool ok = cutlass::reference::device::BlockCompareRelativelyEqual( + ref_out.get(), o_proj.get(), mh, + static_cast(0.15f), static_cast(0.05f)); + std::cout << " K0_W8A8 O: " << (ok ? "PASS (rtol=0.15)" : "FAIL") << std::endl; + all_ok &= ok; + } + + // K4_W8A8: Q projection x_i8[M,H] @ W_q_i8[H,H] * scale_tok1 * scale_q_ch [+RoPE] + ref_float_gemm_acc(q, f_acc.get(), x_i8.get(), W_q_i8.get(), M, H, H, L); + q.wait(); +{% if use_rope %} + ref_apply_dequant_rope(q, ref_out.get(), f_acc.get(), + scale_tok1.get(), scale_q_ch.get(), cos_sin.get(), M, H, L); +{% else %} + ref_apply_dequant(q, ref_out.get(), f_acc.get(), scale_tok1.get(), scale_q_ch.get(), M, H, L); +{% endif %} + q.wait(); + { + bool ok = cutlass::reference::device::BlockCompareRelativelyEqual( + ref_out.get(), Q_out.get(), mh, + static_cast(0.15f), static_cast(0.05f)); + std::cout << " K4_W8A8 Q: " << (ok ? "PASS (rtol=0.15)" : "FAIL") << std::endl; + all_ok &= ok; + } + + // K2_W8A8: FFN raw_sum_i8[M,H] @ W_ffn_i8[H,N_ffn] * scale_tok_raw * scale_ffn_ch + act + ref_float_gemm_acc(q, f_acc.get(), raw_sum_i8.get(), W_ffn_i8.get(), M, N_ffn, H, L); + q.wait(); +{% if ffn_activation == "swiglu" %} + ref_apply_dequant_swiglu(q, ref_out.get(), f_acc.get(), + scale_tok_raw.get(), scale_ffn_ch.get(), M, N_ffn, L); +{% elif ffn_activation == "geglu" %} + ref_apply_dequant_geglu(q, ref_out.get(), f_acc.get(), + scale_tok_raw.get(), scale_ffn_ch.get(), M, N_ffn, L); +{% else %} + ref_apply_dequant(q, ref_out.get(), f_acc.get(), scale_tok_raw.get(), scale_ffn_ch.get(), M, N_ffn, L); +{% endif %} + q.wait(); + { + bool ok = cutlass::reference::device::BlockCompareRelativelyEqual( + ref_out.get(), ffn_out.get(), m_nffn, + static_cast(0.15f), static_cast(0.05f)); + std::cout << " K2_W8A8 FFN: " << (ok ? "PASS (rtol=0.15)" : "FAIL") << std::endl; + all_ok &= ok; + } + } // ref buffers freed here + + std::cout << "Correctness: " << (all_ok ? "ALL PASS" : "SOME FAILED") << std::endl; + + if (opts.iterations == 0) return all_ok ? 0 : 1; + + // ============================================================ + // Benchmark + // ============================================================ + double flops = 2.0 * M * L * + (static_cast(H) * H + + static_cast(H_kv) * H + + static_cast(H) * H + + static_cast(N_ffn) * H); + + GPU_Clock timer; + + // --- XE_W8A8_FUSED --- + run_fused(); // warmup + timer.start(); + for (int i = 0; i < opts.iterations; ++i) run_fused(); + float fused_s = timer.seconds() / opts.iterations; + + // --- VLLM_INT8_EQUIV --- + run_vllm(); // warmup + timer.start(); + for (int i = 0; i < opts.iterations; ++i) run_vllm(); + float vllm_s = timer.seconds() / opts.iterations; + + // --- NAIVE_INT8 --- + run_naive(); // warmup + timer.start(); + for (int i = 0; i < opts.iterations; ++i) run_naive(); + float naive_s = timer.seconds() / opts.iterations; + + // ============================================================ + // Report + // ============================================================ + printf("\nXE_W8A8_FUSED: %.4f ms %.3f TOp/s\n", + fused_s * 1000, flops * 1e-12 / fused_s); + printf("VLLM_INT8_EQUIV: %.4f ms %.3f TOp/s\n", + vllm_s * 1000, flops * 1e-12 / vllm_s); + printf("NAIVE_INT8: %.4f ms %.3f TOp/s\n", + naive_s * 1000, flops * 1e-12 / naive_s); + + std::cout << "\n=== STRUCTURED OUTPUT ===" << std::endl; + std::cout << "PIPELINE: {{ preset_name }}" << std::endl; + std::cout << "MODEL: {{ model_name }}" << std::endl; + printf("DIMS: M=%d H=%d H_kv=%d I=%d N_ffn=%d\n", M, H, H_kv, I, N_ffn); + std::cout << "CORRECTNESS: " << (all_ok ? "RTOL_PASS" : "RTOL_FAIL") << std::endl; + printf("W8A8_FUSED: %.4f ms %.3f TOp/s\n", fused_s * 1000, flops * 1e-12 / fused_s); + printf("VLLM_INT8_EQUIV: %.4f ms %.3f TOp/s\n", vllm_s * 1000, flops * 1e-12 / vllm_s); + printf("NAIVE_INT8: %.4f ms %.3f TOp/s\n", naive_s * 1000, flops * 1e-12 / naive_s); + printf("W8A8_VS_NAIVE_INT8: %.2fx\n", naive_s / fused_s); + printf("W8A8_VS_VLLM_INT8: %.2fx\n", vllm_s / fused_s); + printf("VLLM_VS_NAIVE: %.2fx\n", naive_s / vllm_s); + + return all_ok ? 0 : 1; +} diff --git a/autotune/quantize_weights_quaRot.py b/autotune/quantize_weights_quaRot.py new file mode 100644 index 0000000..dc79ac0 --- /dev/null +++ b/autotune/quantize_weights_quaRot.py @@ -0,0 +1,357 @@ +#!/usr/bin/env python3 +""" +Offline weight preprocessing for QuaRot-style W8A8 quantization. + +Implements the weight rotation + per-column INT8 quantization described in: + QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs (Ashkboos et al., 2024) + https://arxiv.org/abs/2404.00456 + +For each weight matrix W [K, N]: + 1. Apply random Hadamard rotation block-wise (group_size K dimension, group_size N dimension) + 2. Per-column quantize the rotated matrix to INT8 with symmetric scaling + +The rotation redistributes activation outliers uniformly across dimensions, +enabling INT8 quantization with near-BF16 accuracy on large models (70B+). + +Usage: + python3 quantize_weights_quaRot.py --preset llama3_8b --group-size 128 + python3 quantize_weights_quaRot.py --input weights.npy --output weights_q.npz --seed 42 + +Output (.npz): + W_i8: int8 [K, N] — quantized rotated weight + scale_channel: float [N] — per-column dequant scale (max/127) + rotation_seed: uint64 scalar — random seed for the Hadamard ±1 diagonal + group_size: int scalar — rotation group size used +""" + +import argparse +import sys +import numpy as np +from pathlib import Path + +try: + from model_presets import MODEL_PRESETS +except ImportError: + MODEL_PRESETS = {} + + +# ── Walsh-Hadamard Transform (WHT) ─────────────────────────────────────────── + +def wht(x: np.ndarray) -> np.ndarray: + """ + Fast Walsh-Hadamard Transform (unnormalized) applied to the last axis of x. + Last dimension must be a power of 2. + Uses the butterfly (Cooley-Tukey) algorithm: O(n log n). + """ + n = x.shape[-1] + assert n > 0 and (n & (n - 1)) == 0, f"WHT requires power-of-2 size, got {n}" + h = n + while h > 1: + h //= 2 + x = x.copy() + x[..., :h], x[..., h:2*h] = (x[..., :h] + x[..., h:2*h], + x[..., :h] - x[..., h:2*h]) + # Reshape so the butterfly operates on non-contiguous blocks + # Full standard WHT butterfly (operates on all pairs): + x = np.reshape(x, x.shape[:-1] + (n // (h * 2), 2, h)) + lo = x[..., 0, :] + x[..., 1, :] + hi = x[..., 0, :] - x[..., 1, :] + x = np.stack([lo, hi], axis=-2).reshape(x.shape[:-3] + (n,)) + return x + + +def wht_correct(x: np.ndarray) -> np.ndarray: + """ + Reference WHT using the recursive Hadamard matrix definition. + Slower but unambiguously correct — used for validation. + """ + n = x.shape[-1] + if n == 1: + return x.copy() + H = np.array([[1, 1], [1, -1]], dtype=float) + # Build H_n via Kronecker product + H_n = np.array([[1.0]]) + size = 1 + while size < n: + H_n = np.kron(H_n, H) + size *= 2 + return (x @ H_n.T.astype(x.dtype)) + + +def wht_fast(x: np.ndarray) -> np.ndarray: + """ + Iterative in-place WHT. x shape: (..., n), n = power of 2. + Returns a new array with the transform applied along the last axis. + """ + x = x.copy().astype(np.float32) + n = x.shape[-1] + step = 1 + while step < n: + for i in range(0, n, step * 2): + lo = x[..., i:i+step].copy() + hi = x[..., i+step:i+2*step].copy() + x[..., i:i+step] = lo + hi + x[..., i+step:i+2*step] = lo - hi + step *= 2 + return x + + +# ── Nearest power of 2 ─────────────────────────────────────────────────────── + +def next_pow2(n: int) -> int: + p = 1 + while p < n: + p *= 2 + return p + + +def pad_to_pow2(x: np.ndarray, axis: int) -> tuple: + """Pad x along axis to the next power of 2. Returns (padded, original_size).""" + n = x.shape[axis] + p = next_pow2(n) + if p == n: + return x, n + pad_width = [(0, 0)] * x.ndim + pad_width[axis] = (0, p - n) + return np.pad(x, pad_width), n + + +# ── Block-wise Hadamard rotation ────────────────────────────────────────────── + +def hadamard_rotate_matrix(W: np.ndarray, + group_size: int, + seed: int) -> np.ndarray: + """ + Apply random Hadamard rotation to weight matrix W [K, N]. + + For each K-group of rows and N-group of columns: + W_rot[k_blk*g:(k_blk+1)*g, n_blk*g:(n_blk+1)*g] + = D_k @ H_g @ W_block @ H_g.T @ D_n / g + + where D_k, D_n are random ±1 diagonal matrices derived from `seed`. + H_g is the unnormalized Hadamard matrix of size g. + Normalization: / g keeps the Frobenius norm invariant. + + If K or N are not divisible by group_size, the matrix is padded with zeros, + rotated, and then un-padded. + + Returns W_rot [K, N] as float32. + """ + K, N = W.shape + g = group_size + + # Pad to multiple of group_size + K_pad = ((K + g - 1) // g) * g + N_pad = ((N + g - 1) // g) * g + W_padded = np.zeros((K_pad, N_pad), dtype=np.float32) + W_padded[:K, :N] = W.astype(np.float32) + + rng = np.random.default_rng(seed) + n_k_blks = K_pad // g + n_n_blks = N_pad // g + + W_rot = W_padded.copy() + + for k_blk in range(n_k_blks): + k0, k1 = k_blk * g, (k_blk + 1) * g + # Random ±1 diagonal for K dimension + d_k = rng.choice([-1.0, 1.0], size=g).astype(np.float32) + for n_blk in range(n_n_blks): + n0, n1 = n_blk * g, (n_blk + 1) * g + d_n = rng.choice([-1.0, 1.0], size=g).astype(np.float32) + block = W_padded[k0:k1, n0:n1].copy() + # Apply: D_k @ H_g @ block @ H_g.T @ D_n / g + block = (d_k[:, None] * block) # D_k @ block + block = wht_fast(block) # H_g @ (D_k @ block) + block = (block * d_n[None, :]) # result @ D_n + block = wht_fast(block.T).T # result @ H_g.T (WHT on columns) + block /= g # normalize + W_rot[k0:k1, n0:n1] = block + + return W_rot[:K, :N] + + +# ── Per-column quantization ─────────────────────────────────────────────────── + +def quantize_columns(W_rot: np.ndarray) -> tuple: + """ + Symmetric per-column INT8 quantization. + + scale_channel[n] = max(|W_rot[:,n]|) / 127 + W_i8[k,n] = round(clamp(W_rot[k,n] / scale_channel[n], -128, 127)) + + Returns (W_i8: int8 [K,N], scale_channel: float32 [N]). + """ + max_abs = np.abs(W_rot).max(axis=0, keepdims=True) # [1, N] + scale_channel = (max_abs / 127.0).squeeze(0) # [N] + scale_channel = np.where(scale_channel == 0, 1e-8, scale_channel) + + W_scaled = W_rot / scale_channel[None, :] + W_i8 = np.round(np.clip(W_scaled, -128.0, 127.0)).astype(np.int8) + return W_i8, scale_channel.astype(np.float32) + + +# ── Verification helpers ────────────────────────────────────────────────────── + +def quantization_error(W_orig: np.ndarray, W_i8: np.ndarray, + scale_channel: np.ndarray) -> dict: + """Compute relative quantization error stats.""" + W_dequant = W_i8.astype(np.float32) * scale_channel[None, :] + err = np.abs(W_dequant - W_orig.astype(np.float32)) + ref = np.abs(W_orig.astype(np.float32)).mean() + 1e-8 + return { + "mean_abs_err": float(err.mean()), + "max_abs_err": float(err.max()), + "relative_err_pct": float(err.mean() / ref * 100), + } + + +# ── Main pipeline ───────────────────────────────────────────────────────────── + +def process_weight_matrix(W: np.ndarray, + group_size: int = 128, + seed: int = 42, + verbose: bool = True) -> dict: + """ + Rotate and quantize a single weight matrix. + + Returns dict with keys: W_i8, scale_channel, rotation_seed, group_size, + quantization_error (stats for logging) + """ + K, N = W.shape + if verbose: + print(f" Shape: [{K}, {N}], group_size={group_size}, seed={seed}") + + W_rot = hadamard_rotate_matrix(W, group_size=group_size, seed=seed) + W_i8, scale_channel = quantize_columns(W_rot) + err = quantization_error(W_rot, W_i8, scale_channel) + + if verbose: + print(f" Quant error: mean={err['mean_abs_err']:.4f}, " + f"max={err['max_abs_err']:.4f}, " + f"relative={err['relative_err_pct']:.2f}%") + + return { + "W_i8": W_i8, + "scale_channel": scale_channel, + "rotation_seed": np.uint64(seed), + "group_size": np.int32(group_size), + "quant_error": err, + } + + +def generate_transformer_weights(H: int, H_kv: int, I: int, + group_size: int = 128, + base_seed: int = 42, + dtype: np.dtype = np.float32) -> dict: + """ + Generate random weight matrices for the 5-GEMM xe-fuse pipeline and + apply QuaRot quantization to each. + + Returns a dict of {name: {W_i8, scale_channel, ...}} for: + W_q [H, H], W_k [H, H_kv], W_v [H, H_kv], W_o [H, H], W_ffn [H, 2*I] + """ + rng = np.random.default_rng(base_seed) + + weight_shapes = { + "W_q": (H, H), + "W_k": (H, H_kv), + "W_v": (H, H_kv), + "W_o": (H, H), + "W_ffn": (H, 2 * I), + } + + results = {} + for i, (name, shape) in enumerate(weight_shapes.items()): + print(f"\n[{name}]") + W = rng.standard_normal(shape).astype(dtype) + # Scale to typical LLM weight magnitude + W *= 0.02 + results[name] = process_weight_matrix(W, + group_size=group_size, + seed=base_seed + i, + verbose=True) + results[name]["shape"] = shape + return results + + +def save_quantized_weights(results: dict, output_path: str) -> None: + """Save all quantized weight matrices to a .npz archive.""" + arrays = {} + for name, data in results.items(): + arrays[f"{name}_i8"] = data["W_i8"] + arrays[f"{name}_scale"] = data["scale_channel"] + arrays[f"{name}_seed"] = np.array(data["rotation_seed"]) + arrays[f"{name}_group_size"] = np.array(data["group_size"]) + np.savez_compressed(output_path, **arrays) + print(f"\nSaved to: {output_path}.npz") + + +# ── CLI ─────────────────────────────────────────────────────────────────────── + +def main(): + ap = argparse.ArgumentParser( + description="QuaRot offline weight rotation and INT8 quantization for xe-fuse") + ap.add_argument("--preset", choices=list(MODEL_PRESETS.keys()) if MODEL_PRESETS else [], + help="Model preset (uses model_presets.py dims)") + ap.add_argument("--input", type=str, + help="Path to .npy weight matrix [K, N] (if not using preset)") + ap.add_argument("--output", type=str, default="weights_quaRot", + help="Output .npz file path (without extension)") + ap.add_argument("--group-size", type=int, default=128, + help="Hadamard rotation group size (must be power of 2, default: 128)") + ap.add_argument("--seed", type=int, default=42, + help="Base random seed for Hadamard diagonal signs") + ap.add_argument("--verify", action="store_true", + help="Run WHT correctness check before processing") + args = ap.parse_args() + + if args.verify: + print("Verifying WHT implementation...") + for n in [4, 8, 16, 32, 128]: + x = np.random.default_rng(0).standard_normal((4, n)).astype(np.float32) + ref = wht_correct(x) + fast = wht_fast(x) + max_err = np.abs(ref - fast).max() + status = "OK" if max_err < 1e-4 else "FAIL" + print(f" WHT n={n:4d}: max_err={max_err:.2e} [{status}]") + print() + + if args.input: + W = np.load(args.input).astype(np.float32) + if W.ndim != 2: + sys.exit(f"Expected 2D weight matrix, got shape {W.shape}") + print(f"Processing single weight matrix from {args.input}") + result = process_weight_matrix(W, group_size=args.group_size, seed=args.seed) + np.savez_compressed(args.output, + W_i8=result["W_i8"], + scale_channel=result["scale_channel"], + rotation_seed=np.array(result["rotation_seed"]), + group_size=np.array(result["group_size"])) + print(f"Saved: {args.output}.npz") + + elif args.preset and MODEL_PRESETS: + config = MODEL_PRESETS[args.preset] + H = config["H"] + H_kv = config["H_kv"] + I = config["I"] + print(f"Preset: {args.preset} H={H}, H_kv={H_kv}, I={I}") + results = generate_transformer_weights(H, H_kv, I, + group_size=args.group_size, + base_seed=args.seed) + save_quantized_weights(results, args.output) + print("\nQuantization summary:") + for name, data in results.items(): + err = data["quant_error"] + print(f" {name:8s} {str(data['shape']):18s} " + f"relative_err={err['relative_err_pct']:.2f}%") + + else: + ap.print_help() + print("\nExample:") + print(" python3 quantize_weights_quaRot.py --preset llama3_8b --group-size 128") + print(" python3 quantize_weights_quaRot.py --input W_q.npy --output W_q_quaRot --verify") + + +if __name__ == "__main__": + main() diff --git a/include/xe-fuse/builder/epilogue_builder.hpp b/include/xe-fuse/builder/epilogue_builder.hpp index 095a3cd..9cfa374 100644 --- a/include/xe-fuse/builder/epilogue_builder.hpp +++ b/include/xe-fuse/builder/epilogue_builder.hpp @@ -43,6 +43,7 @@ #include #include "xe-fuse/visitors/xe_elementwise_compute.hpp" +#include "xe-fuse/visitors/xe_hadamard_compute.hpp" #include "xe-fuse/visitors/xe_pairwise_compute.hpp" #include "xe-fuse/visitors/xe_rope_compute.hpp" #include "xe-fuse/visitors/xe_scalerows_compute.hpp" @@ -256,6 +257,40 @@ using DequantW8A8Biased = Add< RowBroadcast<2, TileShape, ElementBias, ElementCompute>, ElementCompute, ElementCompute>; +// DequantRoPE: dequant(int32_acc) + RoPE — K4_W8A8 (Q/K projections) +// scale_token[m] absorbs both the per-token quantization range and the RMSNorm +// reciprocal std, so no separate RMSNorm multiply is needed in the epilogue. +// Tree: XeRoPEComputeTwoChild( DequantW8A8, AuxLoad ) +template +using DequantRoPE = RoPEComposed, + ElementCosSin>; + +// DequantSwiGLU: dequant(int32_acc) + SwiGLU — K2_W8A8 (FFN, SwiGLU models) +// Tree: XePairwiseCompute( DequantW8A8 ) +template +using DequantSwiGLU = SwiGLU>; + +// DequantGeGLU: dequant(int32_acc) + GeGLU — K2_W8A8 (FFN, Gemma-style models) +// Tree: XePairwiseCompute( DequantW8A8 ) +template +using DequantGeGLU = GeGLU>; + +// HadamardOutput: apply WHT to the output of InnerEVT +// +// Used as the final epilogue step in K0_W8A8 (O-projection) for QuaRot: +// K0_W8A8 epilogue: DequantW8A8(acc) → add_residual → gamma → HadamardOutput → BF16 +// +// The rotated BF16 activations written to DRAM are ready for the next layer's +// launch_compute_rstd_and_quantize without a separate Hadamard kernel. +// +// Tree: XeEVT, InnerEVT> +template +using HadamardOutput = EVT, InnerEVT>; + // ============================================================ // Auxiliary Store — write intermediate results to a buffer // ============================================================ diff --git a/include/xe-fuse/kernels/compute_rstd.hpp b/include/xe-fuse/kernels/compute_rstd.hpp index 2506ced..3a2b04d 100644 --- a/include/xe-fuse/kernels/compute_rstd.hpp +++ b/include/xe-fuse/kernels/compute_rstd.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include namespace xe_fuse { @@ -47,4 +48,80 @@ void launch_compute_rstd( }); } +// Combined RMSNorm + INT8 quantization kernel. +// +// Computes per-row: +// rstd[m] = rsqrt( mean_n(X[m,n]^2) + eps ) +// normed[m,n] = X[m,n] * rstd[m] +// scale_token[m] = max_n(|normed[m,n]|) / 127 (per-token quant scale) +// quant_out[m,n] = round(normed[m,n] / scale_token[m]) clamped to [-128, 127] +// +// The combined scale_token[m] encodes both the RMSNorm reciprocal std and the +// per-token quantization range. The W8A8 GEMM epilogue uses this combined scale +// directly via ColBroadcast, so no separate RMSNorm multiply is needed. +// +// Three sub-group passes per row: +// Pass 1 — reduce sum_sq → compute rstd +// Pass 2 — reduce max_abs of (X * rstd) → compute scale_token +// Pass 3 — write clamped INT8 values and scale_token[m] +template +void launch_compute_rstd_and_quantize( + sycl::queue& q, + ElementInput const* input_ptr, + int8_t* quant_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) { + int row = item.get_group(0); + int lane = item.get_local_id(0); + + // ── Pass 1: reduce sum_sq ────────────────────────────────────────── + 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); + // Broadcast rstd to all lanes via group_broadcast + rstd = sycl::group_broadcast(sg, rstd, 0); + + // ── Pass 2: reduce 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; // epsilon guards against all-zero rows + scale_tok = sycl::group_broadcast(sg, scale_tok, 0); + + // ── Pass 3: quantize and write outputs ───────────────────────────── + for (int col = lane; col < N; col += SG_SIZE) { + float normed = static_cast(input_ptr[row * N + col]) * rstd; + 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_dequant_rope.hpp b/include/xe-fuse/kernels/gemm_dequant_rope.hpp new file mode 100644 index 0000000..6d84ed0 --- /dev/null +++ b/include/xe-fuse/kernels/gemm_dequant_rope.hpp @@ -0,0 +1,179 @@ +#pragma once + +// K4_W8A8: gemm_dequant_rope — D = RoPE( dequant(A_i8 @ B_i8) ) +// +// INT8×INT8 GEMM with W8A8 dequantization and RoPE fused in the epilogue. +// Used for Q and K projections in the quantized transformer forward pass. +// +// scale_token[m] absorbs both the per-token RMSNorm reciprocal std and the +// per-token quantization range, so no separate RMSNorm multiply is needed: +// scale_token[m] = quant_scale[m] * rstd[m] +// +// EVT tree: +// XeEVT +// >, +// XeRowBroadcast<1, scale_channel> // * scale_channel[n] +// >, +// XeAuxLoad // child 1: interleaved cos/sin +// > + +#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 + +#include "xe-fuse/visitors/xe_rope_compute.hpp" + +namespace xe_fuse { + +template < + typename ElementA_ = int8_t, + typename ElementB_ = int8_t, + typename ElementD_ = cutlass::bfloat16_t, + typename ElementScale_ = float, // type for scale_token and scale_channel + typename ElementCosSin_ = float, + typename ElementAcc_ = int32_t, + typename ElementCompute_ = float, + typename TileShape_ = cute::Shape +> +struct GemmDequantRoPE { + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementD = ElementD_; + using ElementScale = ElementScale_; + using ElementCosSin = ElementCosSin_; + 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 StrideCosSin = cute::Stride, int64_t>; + + // INT8 requires 256-bit aligned loads: 32 elements of int8_t = 256 bits + static constexpr int AlignmentAB = 32; + static constexpr int AlignmentCD = 8; + + // ── Inner dequant tree: int32_acc * scale_token[m] ────────────────────── + 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 + >; + + // int32_acc * scale_token[m] + using InnerDequant = cutlass::epilogue::fusion::XeEVT; + + // ── Outer dequant: InnerDequant * scale_channel[n] ────────────────────── + using ChannelScaleBroadcast = cutlass::epilogue::fusion::XeRowBroadcast< + 0, TileShape, ElementScale, ElementCompute, + cute::Stride, cute::Int<1>, int64_t>, + 128 / cutlass::sizeof_bits_v + >; + + // (acc * scale_token[m]) * scale_channel[n] + using DequantTree = cutlass::epilogue::fusion::XeEVT; + + // ── cos/sin table load ─────────────────────────────────────────────────── + using CosSinLoad = cutlass::epilogue::fusion::XeAuxLoad< + ElementCosSin, StrideCosSin, void, + 128 / cutlass::sizeof_bits_v, true, true + >; + + // ── Root: RoPE( dequant_result, cos_sin ) ─────────────────────────────── + using RoPENode = XeRoPEComputeTwoChild; + + 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 CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, AlignmentAB, + ElementB, LayoutB, AlignmentAB, + 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; + + // Build EVT arguments. + // scale_token[m]: per-token scale absorbing both quant range and RMSNorm rstd + // scale_channel[n]: per-channel weight quantization scale + static typename EVT::Arguments make_evt_args( + ElementScale const* scale_token_ptr, int M, + ElementScale const* scale_channel_ptr, int N, + ElementCosSin const* cos_sin_ptr, + StrideCosSin stride_cos_sin) { + + // Inner dequant: XeEVT + typename Accum::Arguments accum_args{}; + + typename TokenScaleBroadcast::Arguments token_scale_args; + token_scale_args.ptr_col = scale_token_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}; + + // Outer dequant: XeEVT + typename ChannelScaleBroadcast::Arguments channel_scale_args; + channel_scale_args.ptr_row = scale_channel_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}; + + // cos/sin AuxLoad + typename CosSinLoad::Arguments cos_sin_args; + cos_sin_args.ptr_aux = cos_sin_ptr; + cos_sin_args.null_default = ElementCosSin(0); + cos_sin_args.dAux = stride_cos_sin; + + // Root RoPE node + typename RoPENode::Arguments rope_args{}; + + return {dequant_args, cos_sin_args, rope_args}; + } +}; + +} // namespace xe_fuse diff --git a/include/xe-fuse/kernels/gemm_dequant_swiglu.hpp b/include/xe-fuse/kernels/gemm_dequant_swiglu.hpp new file mode 100644 index 0000000..63950c9 --- /dev/null +++ b/include/xe-fuse/kernels/gemm_dequant_swiglu.hpp @@ -0,0 +1,270 @@ +#pragma once + +// K2_W8A8: gemm_dequant_swiglu — D = SwiGLU( dequant(A_i8 @ B_i8) ) +// +// INT8×INT8 GEMM with W8A8 dequantization and SwiGLU fused in the epilogue. +// Used for the FFN (gate+up projection) in the quantized transformer forward +// pass for SwiGLU-based models (LLaMA 3, Mistral, Qwen 2.5, etc.). +// +// The GEMM produces M×N_ffn output where N_ffn = 2*I (interleaved gate and +// up projections). After dequantization, adjacent pairs are fed to SwiGLU: +// output[m, 2k] = silu(dequant[m, 2k]) * dequant[m, 2k+1] +// output[m, 2k+1] = silu(dequant[m, 2k]) * dequant[m, 2k+1] (same value) +// +// The N→N/2 contraction (writing only I output columns) is handled by the +// caller; both even and odd lanes carry the same SwiGLU value here. +// +// EVT tree: +// XeEVT, // root: SwiGLU +// XeEVT +// >, +// XeRowBroadcast<1, scale_channel> // * scale_channel[n] +// > +// > + +#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 + +#include "xe-fuse/visitors/xe_pairwise_compute.hpp" + +namespace xe_fuse { + +template < + typename ElementA_ = int8_t, + typename ElementB_ = int8_t, + typename ElementD_ = cutlass::bfloat16_t, + typename ElementScale_ = float, // type for scale_token and scale_channel + typename ElementAcc_ = int32_t, + typename ElementCompute_ = float, + typename TileShape_ = cute::Shape +> +struct GemmDequantSwiGLU { + 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; + + // ── Inner dequant: int32_acc * scale_token[m] ─────────────────────────── + 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; + + // ── Outer dequant: InnerDequant * scale_channel[n] ────────────────────── + 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; + + // ── Root: SwiGLU( dequant_result ) ────────────────────────────────────── + 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 CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, AlignmentAB, + ElementB, LayoutB, AlignmentAB, + 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( + ElementScale const* scale_token_ptr, int M, + ElementScale const* scale_channel_ptr, int N) { + + typename Accum::Arguments accum_args{}; + + typename TokenScaleBroadcast::Arguments token_scale_args; + token_scale_args.ptr_col = scale_token_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_channel_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}; + } +}; + +// GeGLU variant for Gemma-style models +template < + typename ElementA_ = int8_t, + typename ElementB_ = int8_t, + typename ElementD_ = cutlass::bfloat16_t, + typename ElementScale_ = float, + typename ElementAcc_ = int32_t, + typename ElementCompute_ = float, + typename TileShape_ = cute::Shape +> +struct GemmDequantGeGLU { + 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; + + 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 GeGLUVisitor = 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 CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, AlignmentAB, + ElementB, LayoutB, AlignmentAB, + 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( + ElementScale const* scale_token_ptr, int M, + ElementScale const* scale_channel_ptr, int N) { + + typename Accum::Arguments accum_args{}; + + typename TokenScaleBroadcast::Arguments token_scale_args; + token_scale_args.ptr_col = scale_token_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_channel_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 GeGLUVisitor::Arguments geglu_args{}; + + return {dequant_args, geglu_args}; + } +}; + +} // namespace xe_fuse diff --git a/include/xe-fuse/kernels/gemm_dequant_w8a8.hpp b/include/xe-fuse/kernels/gemm_dequant_w8a8.hpp new file mode 100644 index 0000000..2db1196 --- /dev/null +++ b/include/xe-fuse/kernels/gemm_dequant_w8a8.hpp @@ -0,0 +1,136 @@ +#pragma once + +// K1_W8A8 / K0_W8A8: plain INT8 GEMM + W8A8 dequantization to BF16. +// +// EVT tree: +// XeEVT +// >, +// XeRowBroadcast<1, scale_channel> +// > +// +// Used for: +// K1_W8A8 — V projection: x_i8 @ W_v_i8 → dequant → BF16 +// K0_W8A8 — O projection: attn_i8 @ W_o_i8 → dequant → BF16 +// (residual add + gamma applied via standalone ops afterwards) + +#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_ = int8_t, + typename ElementB_ = int8_t, + typename ElementD_ = cutlass::bfloat16_t, + typename ElementScale_ = float, + typename ElementAcc_ = int32_t, + typename ElementCompute_ = float, + typename TileShape_ = cute::Shape +> +struct GemmDequantW8A8 { + 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; + + 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; + + 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 CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, AlignmentAB, + ElementB, LayoutB, AlignmentAB, + 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( + ElementScale const* scale_token_ptr, int M, + ElementScale const* scale_channel_ptr, int N) { + + typename Accum::Arguments accum_args{}; + + typename TokenScaleBroadcast::Arguments token_scale_args; + token_scale_args.ptr_col = scale_token_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_channel_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}; + } +}; + +} // namespace xe_fuse diff --git a/include/xe-fuse/standalone/ops.hpp b/include/xe-fuse/standalone/ops.hpp index b9d49f1..d9c8830 100644 --- a/include/xe-fuse/standalone/ops.hpp +++ b/include/xe-fuse/standalone/ops.hpp @@ -283,4 +283,74 @@ inline void select_logits(sycl::queue& q, bf16 const* logits, int const* targets }); } +// ── INT8 quantization and dequantization ───────────────────────────────────── + +// Quantize BF16 normalized activations to INT8 using a precomputed rstd vector. +// Used in the naive_int8 baseline where rstd already exists from a prior +// compute_rstd call. +// +// normed[m,n] = input[m,n] * rstd[m] +// scale_tok[m] = max_n(|normed[m,n]|) / 127 (written to scale_token_out) +// quant_out[m,n] = round(normed[m,n] / scale_tok[m]) clamped to [-128, 127] +// +// Note: this is a 2-pass kernel (max_abs then quantize). For a fused single-pass +// version that also computes rstd, use launch_compute_rstd_and_quantize(). +inline void quantize_activations(sycl::queue& q, + bf16 const* input, + float const* rstd, + int8_t* quant_out, + float* scale_token_out, + int M, int N, int L) { + int m_val = M, n_val = N; + // Pass 1: compute scale_token[m] = max(|input[m,n]*rstd[m]|) / 127 + q.parallel_for(sycl::range<1>(static_cast(M) * L), [=](sycl::id<1> idx) { + int64_t row_idx = idx[0]; + int batch = static_cast(row_idx / m_val); + int m = static_cast(row_idx % m_val); + int64_t base = static_cast(batch) * m_val * n_val + static_cast(m) * n_val; + float r = rstd[row_idx]; + float mx = 0.f; + for (int n = 0; n < n_val; ++n) + mx = sycl::fmax(mx, sycl::fabs(static_cast(input[base + n]) * r)); + scale_token_out[row_idx] = mx / 127.f + 1e-8f; + }); + q.wait(); + + // Pass 2: quantize + int64_t total = static_cast(M) * N * L; + q.parallel_for(sycl::range<1>(total), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int row_idx = static_cast(i / n_val); + float normed = static_cast(input[i]) * rstd[row_idx]; + float qval = sycl::round(normed / scale_token_out[row_idx]); + qval = sycl::fmin(sycl::fmax(qval, -128.f), 127.f); + quant_out[i] = static_cast(qval); + }); +} + +// Dequantize INT32 GEMM accumulator to BF16 using per-token and per-channel scales. +// Naive baseline: called after a bare INT8 GEMM that wrote INT32 output. +// +// out[m,n] = bf16( int32_acc[m,n] * scale_token[m] * scale_channel[n] ) +inline void dequant_w8a8(sycl::queue& q, + bf16* out, + int32_t const* acc, + float const* scale_token, + float const* scale_channel, + int M, int N, int L) { + int64_t total = static_cast(M) * N * L; + int m_val = M, n_val = N; + q.parallel_for(sycl::range<1>(total), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int col = static_cast(i % n_val); + int row_idx = static_cast(i / n_val); + int m = row_idx % m_val; + int batch = row_idx / m_val; + float val = static_cast(acc[i]) + * scale_token[batch * m_val + m] + * scale_channel[batch * n_val + col]; + out[i] = static_cast(val); + }); +} + } // namespace xe_fuse::standalone diff --git a/include/xe-fuse/standalone/vllm_ops.hpp b/include/xe-fuse/standalone/vllm_ops.hpp index 27b079d..bfd3b68 100644 --- a/include/xe-fuse/standalone/vllm_ops.hpp +++ b/include/xe-fuse/standalone/vllm_ops.hpp @@ -1,6 +1,13 @@ #pragma once // vllm-equivalent standalone kernels for xe-fuse comparison benchmarks. +// Includes merged INT8 dequant + op kernels for the vllm_int8_equiv comparison. +// These represent what a well-implemented INT8 inference engine would do: +// one kernel that reads the INT32 GEMM accumulator and applies dequant + activation +// without a separate DRAM round-trip for the dequant output. +// +// xe-fuse W8A8 goes one step further: the INT32 accumulator never reaches DRAM +// (dequant + activation happen directly in the GEMM epilogue registers). // // These re-implement the algorithmic patterns from vllm-xpu-kernels // (csrc/layernorm.cpp, csrc/activation.cpp, csrc/pos_encoding_kernels.cpp) @@ -12,6 +19,7 @@ // vllm: bare GEMM → separate fused standalone kernel → bare GEMM → ... // xe-fuse: GEMM + epilogue fusion (ops run on register data) +#include #include #include "cutlass/bfloat16.h" @@ -220,4 +228,123 @@ inline void rotary_embedding(sycl::queue& q, bf16* query, bf16* key, }); } +// ── INT8 merged dequant + op kernels (vllm_int8_equiv comparison) ──────────── +// +// These read INT32 GEMM accumulator output from DRAM, apply W8A8 dequantization +// and the activation op in one pass, and write BF16 output. +// +// Compared to naive_int8 (separate dequant_w8a8 + separate op kernel): +// - One kernel launch instead of two +// - INT32 accumulator read once instead of write+read +// Compared to xe-fuse W8A8: +// - INT32 accumulator still reaches DRAM (written by bare INT8 GEMM) +// - 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) +inline void dequant_and_silu_mul(sycl::queue& q, + bf16* out, + int32_t const* acc, + float const* scale_token, + float const* scale_channel, + int d, int M, int L = 1) { + int N = 2 * d; + int wg_size = std::min(d, 1024); + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(static_cast(M) * L * wg_size, wg_size), + [=](sycl::nd_item<1> item) { + int grp = item.get_group(0); + int l = grp / M; + int row = grp % M; + int lid = item.get_local_id(0); + int lsz = item.get_local_range(0); + int64_t in_base = (static_cast(l) * M + row) * N; + float st = scale_token[l * M + row]; + + for (int i = lid; i < d; i += lsz) { + float gate = static_cast(acc[in_base + i]) + * st * scale_channel[l * N + i]; + float up = static_cast(acc[in_base + d + i]) + * st * scale_channel[l * N + d + i]; + float silu_gate = gate / (1.0f + sycl::exp(-gate)); + bf16 result = static_cast(silu_gate * up); + out[in_base + i] = result; + out[in_base + d + i] = result; + } + }); + }); +} + +// Merged: dequant INT32 → BF16, then apply GeGLU +inline void dequant_and_gelu_mul(sycl::queue& q, + bf16* out, + int32_t const* acc, + float const* scale_token, + float const* scale_channel, + int d, int M, int L = 1) { + int N = 2 * d; + int wg_size = std::min(d, 1024); + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(static_cast(M) * L * wg_size, wg_size), + [=](sycl::nd_item<1> item) { + int grp = item.get_group(0); + int l = grp / M; + int row = grp % M; + int lid = item.get_local_id(0); + int lsz = item.get_local_range(0); + int64_t in_base = (static_cast(l) * M + row) * N; + float st = scale_token[l * M + row]; + + for (int i = lid; i < d; i += lsz) { + float gate = static_cast(acc[in_base + i]) + * st * scale_channel[l * N + i]; + float up = static_cast(acc[in_base + d + i]) + * st * scale_channel[l * N + d + i]; + float gelu_gate = gate * 0.5f * (1.0f + sycl::erf(gate * 0.7071067811865475f)); + bf16 result = static_cast(gelu_gate * up); + out[in_base + i] = result; + out[in_base + d + i] = result; + } + }); + }); +} + +// Merged: dequant INT32 → BF16, then apply NeoX RoPE in-place. +// Input: INT32 accumulator [L, M, N], scale_token[L*M], scale_channel[L*N] +// cos_sin_cache: [L, M, N] interleaved cos/sin +// Output: BF16 [L, M, N] with RoPE applied +inline void dequant_and_rotary_embedding(sycl::queue& q, + bf16* out, + int32_t const* acc, + float const* scale_token, + float const* scale_channel, + float const* cos_sin_cache, + int M, int N, int L = 1) { + q.parallel_for(sycl::range<1>(static_cast(M) * N * L), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int l = static_cast(i / (M * N)); + int row = static_cast((i / N) % M); + int col = static_cast(i % N); + int64_t base = static_cast(l) * M * N + static_cast(row) * N; + float st = scale_token[l * M + row]; + + int even_col = col & ~1; + int odd_col = even_col + 1; + if (odd_col >= N) { + out[i] = static_cast(static_cast(acc[i]) * st * scale_channel[l * N + col]); + return; + } + float x_even = static_cast(acc[base + even_col]) * st * scale_channel[l * N + even_col]; + float x_odd = static_cast(acc[base + odd_col]) * st * scale_channel[l * N + odd_col]; + float cos_val = cos_sin_cache[base + even_col]; + float sin_val = cos_sin_cache[base + odd_col]; + out[i] = static_cast((col & 1) == 0 + ? x_even * cos_val + x_odd * sin_val + : -x_even * sin_val + x_odd * cos_val); + }); +} + } // namespace xe_fuse::vllm_equiv diff --git a/include/xe-fuse/visitors/xe_hadamard_compute.hpp b/include/xe-fuse/visitors/xe_hadamard_compute.hpp new file mode 100644 index 0000000..0dde08f --- /dev/null +++ b/include/xe-fuse/visitors/xe_hadamard_compute.hpp @@ -0,0 +1,101 @@ +#pragma once + +// XeHadamardCompute — Epilogue visitor for Walsh-Hadamard Transform (WHT). +// +// Applies an in-register WHT of size GroupSize to the output of a child EVT subtree, +// using sub-group shuffles (shfl_xor_sync) across lanes — the same mechanism used by +// XeRoPECompute and XePairwiseCompute. +// +// For GroupSize=16 (= XeGPU sub-group size): 4 butterfly stages, all intra-SG, no SLM. +// Stage k: shfl_xor(mask = 2^k); bit k of lane_id → a-lane (+) or b-lane (-) +// Normalized by 1/sqrt(GroupSize) so the transform is orthonormal. +// +// Usage in QuaRot W8A8: +// Wrap the K0 (O-projection) output with HadamardOutput so the BF16 +// activations written to DRAM are already rotated — the next layer's +// launch_compute_rstd_and_quantize sees rotated values without an extra kernel. +// +// Current support: GroupSize <= 16 (single sub-group). GroupSize > 16 requires +// inter-SG SLM staging and is deferred to a future extension. + +#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp" +#include "cutlass/gpu_generics.h" + +namespace xe_fuse { + +template +struct XeHadamardCompute : cutlass::epilogue::fusion::Sm90VisitorImpl<> { + + static constexpr int GroupSize = GroupSize_; + static_assert(GroupSize >= 2 && (GroupSize & (GroupSize - 1)) == 0, + "GroupSize must be a power of 2"); + static_assert(GroupSize <= 16, + "GroupSize > 16 requires inter-subgroup SLM staging (not yet implemented)"); + + // Number of butterfly stages = log2(GroupSize) + static constexpr int kStages = []() constexpr { + int s = 0, g = GroupSize; + while (g > 1) { g >>= 1; ++s; } + return s; + }(); + + using Sm90VisitorImpl<>::Sm90VisitorImpl; + + struct ConsumerStoreCallbacks : cutlass::epilogue::fusion::EmptyConsumerStoreCallbacks { + + // Takes the output of a child EVT subtree (frg_input) and applies WHT. + // frg_acc is passed through but not used — same pattern as XePairwiseCompute. + template + CUTLASS_DEVICE cutlass::Array + visit(cutlass::Array const& frg_acc, + int epi_v, int epi_m, int epi_n, + cutlass::Array const& frg_input) { + + cutlass::Array result; + + auto sg = sycl::ext::oneapi::this_work_item::get_sub_group(); + uint32_t lane_id = sg.get_local_linear_id(); + + // Normalization: 1 / sqrt(GroupSize). Computed as a float constant. + // For GroupSize=16: 0.25f; GroupSize=8: ~0.354f; GroupSize=4: 0.5f. + const float kInvSqrtG = sycl::rsqrt(static_cast(GroupSize)); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < FragmentSize; ++i) { + float val = static_cast(frg_input[i]); + + // Apply log2(GroupSize) butterfly stages. + // + // At stage k (xor_mask = 2^k): + // - Each lane exchanges its value with its xor-partner. + // - bit k of lane_id == 0 → "a-lane": output = my_val + partner_val + // - bit k of lane_id == 1 → "b-lane": output = partner_val - my_val + // + // This is the iterative Hadamard butterfly, which matches the Kronecker + // product definition H_n = H_2 ⊗ H_{n/2}. + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kStages; ++stage) { + uint32_t mask = 1u << stage; + uint32_t val_bits = reinterpret_cast(val); + uint32_t partner_bits = shfl_xor_sync(0xFFFFFFFF, val_bits, mask, GroupSize); + float partner_val = reinterpret_cast(partner_bits); + + bool is_b_lane = (lane_id >> stage) & 1u; + val = is_b_lane ? (partner_val - val) : (val + partner_val); + } + + result[i] = static_cast(val * kInvSqrtG); + } + + return result; + } + }; + + template + CUTLASS_DEVICE auto + get_consumer_store_callbacks(cutlass::epilogue::fusion::ConsumerStoreArgs const& args) { + return ConsumerStoreCallbacks{}; + } +}; + +} // namespace xe_fuse diff --git a/tests/test_hadamard_visitor.cpp b/tests/test_hadamard_visitor.cpp new file mode 100644 index 0000000..3b593db --- /dev/null +++ b/tests/test_hadamard_visitor.cpp @@ -0,0 +1,174 @@ +// xe-fuse test: XeHadamardCompute<16> correctness +// +// Directly instantiates the XeHadamardCompute visitor inside a SYCL kernel +// and compares each group of 16 outputs to a CPU Walsh-Hadamard Transform +// reference. +// +// The test kernel uses nd_range<1> with local size 16 (one sub-group per group), +// which matches how shfl_xor_sync operates inside the visitor (uses get_nd_item<1>). +// +// Correctness check: relative error < 1e-5 for all elements. +// In practice the WHT butterfly on float32 is near-exact for values < 2^20. + +#include "xe-fuse/visitors/xe_hadamard_compute.hpp" + +#include "cutlass/array.h" +#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 + +// ── CPU WHT reference ───────────────────────────────────────────────────────── + +// Iterative in-place WHT on x[0..n-1], normalized by 1/sqrt(n). +// Matches the butterfly convention in XeHadamardCompute: +// stage k (stride = 2^k): +// index with bit-k == 0 → even lane: a + b +// index with bit-k == 1 → odd lane: a - b +static void wht_reference(float* x, int n) { + for (int step = 1; step < n; step <<= 1) { + for (int i = 0; i < n; i += step * 2) { + for (int j = 0; j < step; ++j) { + float a = x[i + j]; + float b = x[i + step + j]; + x[i + j] = a + b; + x[i + step + j] = a - b; + } + } + } + float inv_sqrt_n = 1.0f / std::sqrt(static_cast(n)); + for (int i = 0; i < n; ++i) x[i] *= inv_sqrt_n; +} + +// ── Device kernel ───────────────────────────────────────────────────────────── + +// Each work-group has 16 work items (= 1 sub-group on XeGPU B70). +// Each group processes 16 consecutive float elements. +static sycl::event launch_hadamard_kernel(float const* d_input, + float* d_output, + int num_groups) +{ + return compat::get_default_queue().submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(static_cast(num_groups) * 16, 16), + [=](sycl::nd_item<1> item) { + int group_id = static_cast(item.get_group(0)); + int lane = static_cast(item.get_local_id(0)); + int idx = group_id * 16 + lane; + + // Single-element fragment + constexpr int kFragSize = 1; + cutlass::Array frg_input; + frg_input[0] = d_input[idx]; + + cutlass::Array frg_acc; + frg_acc[0] = 0.0f; + + // Apply XeHadamardCompute<16>: 4-stage butterfly across the sub-group + xe_fuse::XeHadamardCompute<16>::ConsumerStoreCallbacks cb{}; + auto result = cb.visit(frg_acc, /*epi_v=*/0, /*epi_m=*/0, /*epi_n=*/0, frg_input); + + d_output[idx] = result[0]; + } + ); + }); +} + +// ── Options ────────────────────────────────────────────────────────────────── + +struct Options { + int num_groups = 4096; // number of 16-element WHT groups (total_elems = groups * 16) + int iterations = 100; + int verify = 1; + + void parse(int argc, char const** argv) { + cutlass::CommandLine cmd(argc, argv); + cmd.get_cmd_line_argument("groups", num_groups, 4096); + cmd.get_cmd_line_argument("iterations", iterations, 100); + cmd.get_cmd_line_argument("verify", verify, 1); + } +}; + +// ── Main ───────────────────────────────────────────────────────────────────── + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + int total_elems = opts.num_groups * 16; + + // Initialize random host input in [-1, 1] + std::vector h_input(total_elems); + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (auto& v : h_input) v = dist(rng); + + // CPU WHT reference: apply separately to each group of 16 + std::vector h_ref = h_input; + for (int g = 0; g < opts.num_groups; ++g) + wht_reference(h_ref.data() + g * 16, 16); + + // Device buffers + cutlass::device_memory::allocation d_input(total_elems); + cutlass::device_memory::allocation d_output(total_elems); + + compat::get_default_queue().memcpy(d_input.get(), h_input.data(), + total_elems * sizeof(float)); + compat::wait(); + + // Run (warm-up + result) + launch_hadamard_kernel(d_input.get(), d_output.get(), opts.num_groups); + compat::wait(); + + // Correctness check + bool passed = true; + if (opts.verify) { + std::vector h_output(total_elems); + compat::get_default_queue().memcpy(h_output.data(), d_output.get(), + total_elems * sizeof(float)); + compat::wait(); + + float max_abs_err = 0.0f; + float max_rel_err = 0.0f; + int num_fail = 0; + + for (int i = 0; i < total_elems; ++i) { + float abs_err = std::abs(h_output[i] - h_ref[i]); + float ref_mag = std::abs(h_ref[i]) + 1e-6f; + float rel_err = abs_err / ref_mag; + max_abs_err = std::max(max_abs_err, abs_err); + max_rel_err = std::max(max_rel_err, rel_err); + if (rel_err > 1e-5f) ++num_fail; + } + + passed = (num_fail == 0); + printf("Correctness: %s\n", passed ? "PASSED" : "FAILED"); + printf(" num_groups=%d total_elems=%d\n", opts.num_groups, total_elems); + printf(" max_abs_err=%.2e max_rel_err=%.2e failures=%d\n", + max_abs_err, max_rel_err, num_fail); + } + + // Benchmark + if (opts.iterations > 0) { + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) + launch_hadamard_kernel(d_input.get(), d_output.get(), opts.num_groups); + compat::wait(); + + float time_ms = timer.milliseconds() / static_cast(opts.iterations); + float gb = static_cast(total_elems) * sizeof(float) * 2.0f / 1e9f; + float bw_gbs = gb / (time_ms * 1e-3f); + + printf("Throughput: %.3f ms/iter %.1f GB/s (%d iters)\n", + time_ms, bw_gbs, opts.iterations); + } + + return passed ? 0 : 1; +} diff --git a/tests/test_k2_w8a8.cpp b/tests/test_k2_w8a8.cpp new file mode 100644 index 0000000..ba3975b --- /dev/null +++ b/tests/test_k2_w8a8.cpp @@ -0,0 +1,223 @@ + +// xe-fuse test: K2_W8A8 — gemm_dequant_swiglu +// D = SwiGLU( dequant(A_i8 @ B_i8) ) +// INT8×INT8 GEMM with W8A8 dequantization and SwiGLU fused in a single epilogue. +// +// Reference: +// acc_f32[m,n] = sum_k( float(A_i8[m,k]) * float(B_i8[k,n]) ) +// dequant[m,n] = acc_f32[m,n] * scale_token[m] * scale_channel[n] +// D[m, 2i] = silu(dequant[m, 2i]) * dequant[m, 2i+1] +// D[m, 2i+1] = silu(dequant[m, 2i]) * dequant[m, 2i+1] (same value) + +#include "xe-fuse/kernels/gemm_dequant_swiglu.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 + +using namespace cute; + +struct Options { + int m = 512, n = 28672, k = 4096, l = 1; // LLaMA 3 8B FFN dims + 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 K2W8A8 = xe_fuse::GemmDequantSwiGLU<>; +using GemmOp = K2W8A8::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(K2W8A8::StrideC{}, make_shape(M, N, L)); + auto stride_D = cutlass::make_cute_packed_stride(K2W8A8::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_token(static_cast(M) * L); + cutlass::DeviceAllocation block_scale_channel(static_cast(N) * L); + + 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); + + { + std::vector h_A(static_cast(M) * K * L); + std::vector h_B(static_cast(K) * N * L); + std::mt19937 rng_a(2001), rng_b(2002); + std::uniform_int_distribution dist(-64, 63); + for (auto& v : h_A) v = static_cast(dist(rng_a)); + for (auto& v : h_B) v = static_cast(dist(rng_b)); + compat::get_default_queue().memcpy(block_A.get(), h_A.data(), h_A.size() * sizeof(int8_t)); + compat::get_default_queue().memcpy(block_B.get(), h_B.data(), h_B.size() * sizeof(int8_t)); + } + + { + std::vector h_st(static_cast(M) * L); + std::vector h_sc(static_cast(N) * L); + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.001f, 0.01f); + for (auto& v : h_st) v = dist(rng); + for (auto& v : h_sc) v = dist(rng); + compat::get_default_queue().memcpy(block_scale_token.get(), h_st.data(), h_st.size() * sizeof(float)); + compat::get_default_queue().memcpy(block_scale_channel.get(), h_sc.data(), h_sc.size() * sizeof(float)); + } + compat::wait(); + + auto evt_args = K2W8A8::make_evt_args( + block_scale_token.get(), M, + block_scale_channel.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) { + // Convert INT8 → float32 + { + const int8_t* a_ptr = block_A.get(); + float* af_ptr = block_A_f32.get(); + int64_t na = static_cast(M) * K * L; + compat::get_default_queue().parallel_for(sycl::range<1>(na), [=](sycl::id<1> idx) { + af_ptr[idx[0]] = static_cast(a_ptr[idx[0]]); + }); + const int8_t* b_ptr = block_B.get(); + float* bf_ptr = block_B_f32.get(); + int64_t nb = static_cast(K) * N * L; + compat::get_default_queue().parallel_for(sycl::range<1>(nb), [=](sycl::id<1> idx) { + bf_ptr[idx[0]] = static_cast(b_ptr[idx[0]]); + }); + compat::wait(); + } + + // Reference GEMM in float + { + 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 / (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(); + } + + // Apply dequant + SwiGLU reference + // SwiGLU: out[2i] = out[2i+1] = silu(gate) * up + // gate = dequant[2i], up = dequant[2i+1] + { + const float* acc = block_acc_f32.get(); + const float* st = block_scale_token.get(); + const float* sc = block_scale_channel.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 / (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] * st[batch * M_ + row] * sc[batch * N_ + col]; + ref[i] = static_cast(v); + return; + } + float gate = acc[base + row * N_ + even_col] + * st[batch * M_ + row] * sc[batch * N_ + even_col]; + float up = acc[base + row * N_ + odd_col] + * st[batch * M_ + row] * sc[batch * N_ + odd_col]; + float silu_gate = gate / (1.f + sycl::exp(-gate)); + ref[i] = static_cast(silu_gate * up); + } + ); + compat::wait(); + } + + 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 tops = (2.0 * M * N * K * L) * 1e-12; + std::cout << "Problem Size: " << M << 'x' << N << 'x' << K << 'x' << L << std::endl; + printf("xe-fuse K2_W8A8 (INT8 GEMM+Dequant+SwiGLU): [%4.3f]TOp/s (%6.4f)ms\n", + tops / time_s, time_s * 1000); + } + + return 0; +} diff --git a/tests/test_k4_w8a8.cpp b/tests/test_k4_w8a8.cpp new file mode 100644 index 0000000..f8c9c2e --- /dev/null +++ b/tests/test_k4_w8a8.cpp @@ -0,0 +1,254 @@ + +// xe-fuse test: K4_W8A8 — gemm_dequant_rope +// D = RoPE( dequant(A_i8 @ B_i8) ) +// INT8×INT8 GEMM with W8A8 dequantization and RoPE fused in a single epilogue. +// +// Reference: +// acc_f32[m,n] = sum_k( float(A_i8[m,k]) * float(B_i8[k,n]) ) +// dequant[m,n] = acc_f32[m,n] * scale_token[m] * scale_channel[n] +// D[m,n] = RoPE( dequant[m,n], cos_sin ) + +#include "xe-fuse/kernels/gemm_dequant_rope.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 + +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 K4W8A8 = xe_fuse::GemmDequantRoPE<>; +using GemmOp = K4W8A8::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(K4W8A8::StrideC{}, make_shape(M, N, L)); + auto stride_D = cutlass::make_cute_packed_stride(K4W8A8::StrideD{}, make_shape(M, N, L)); + auto stride_cs = cutlass::make_cute_packed_stride(K4W8A8::StrideCosSin{}, 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_token(static_cast(M) * L); + cutlass::DeviceAllocation block_scale_channel(static_cast(N) * L); + cutlass::DeviceAllocation block_cos_sin(static_cast(M) * N * L); + + // Temporaries for the reference GEMM computation + 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 INT8 inputs with values in [-64, 63] to avoid INT32 accumulator overflow + { + std::vector h_A(static_cast(M) * K * L); + std::vector h_B(static_cast(K) * N * L); + std::mt19937 rng_a(1001), rng_b(1002); + std::uniform_int_distribution dist(-64, 63); + for (auto& v : h_A) v = static_cast(dist(rng_a)); + for (auto& v : h_B) v = static_cast(dist(rng_b)); + compat::get_default_queue().memcpy(block_A.get(), h_A.data(), h_A.size() * sizeof(int8_t)); + compat::get_default_queue().memcpy(block_B.get(), h_B.data(), h_B.size() * sizeof(int8_t)); + } + + // Per-token and per-channel quantization scales (simulate realistic LLM values) + { + std::vector h_st(static_cast(M) * L); + std::vector h_sc(static_cast(N) * L); + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.001f, 0.01f); + for (auto& v : h_st) v = dist(rng); + for (auto& v : h_sc) v = dist(rng); + compat::get_default_queue().memcpy(block_scale_token.get(), h_st.data(), h_st.size() * sizeof(float)); + compat::get_default_queue().memcpy(block_scale_channel.get(), h_sc.data(), h_sc.size() * sizeof(float)); + } + + // cos_sin with realistic RoPE frequencies + { + std::vector h_cs(static_cast(M) * N * L); + for (int batch = 0; batch < L; ++batch) { + for (int m = 0; m < M; ++m) { + for (int k_pair = 0; k_pair < N / 2; ++k_pair) { + float freq = 1.0f / std::pow(10000.0f, 2.0f * k_pair / static_cast(N)); + float angle = static_cast(m) * freq; + size_t base = static_cast(batch) * M * N + static_cast(m) * N; + h_cs[base + 2 * k_pair] = std::cos(angle); + h_cs[base + 2 * k_pair + 1] = std::sin(angle); + } + } + } + compat::get_default_queue().memcpy(block_cos_sin.get(), h_cs.data(), h_cs.size() * sizeof(float)); + } + compat::wait(); + + // Run K4_W8A8 kernel + auto evt_args = K4W8A8::make_evt_args( + block_scale_token.get(), M, + block_scale_channel.get(), N, + block_cos_sin.get(), stride_cs); + + 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) { + // Step 1: Convert INT8 inputs to float32 on device + { + const int8_t* a_ptr = block_A.get(); + float* af_ptr = block_A_f32.get(); + int64_t total_a = static_cast(M) * K * L; + compat::get_default_queue().parallel_for(sycl::range<1>(total_a), [=](sycl::id<1> idx) { + af_ptr[idx[0]] = static_cast(a_ptr[idx[0]]); + }); + + const int8_t* b_ptr = block_B.get(); + float* bf_ptr = block_B_f32.get(); + int64_t total_b = static_cast(K) * N * L; + compat::get_default_queue().parallel_for(sycl::range<1>(total_b), [=](sycl::id<1> idx) { + bf_ptr[idx[0]] = static_cast(b_ptr[idx[0]]); + }); + compat::wait(); + } + + // Step 2: Reference GEMM in float (slow O(MNK) kernel, small dims only) + { + 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 / (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(); + } + + // Step 3: Apply dequant + RoPE to produce reference output + { + const float* acc = block_acc_f32.get(); + const float* st = block_scale_token.get(); + const float* sc = block_scale_channel.get(); + const float* cs = block_cos_sin.get(); + auto* ref_ptr = 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 / (M_ * N_)); + int64_t base = static_cast(batch) * M_ * N_; + + float val = acc[i] * st[batch * M_ + row] * sc[batch * N_ + col]; + + // RoPE: even col = x*cos + x_odd*sin; odd col = -x_even*sin + x*cos + int even_col = col & ~1; + int odd_col = even_col + 1; + if (odd_col >= N_) { + ref_ptr[i] = static_cast(val); + return; + } + float x_even = acc[base + row * N_ + even_col] + * st[batch * M_ + row] * sc[batch * N_ + even_col]; + float x_odd = acc[base + row * N_ + odd_col] + * st[batch * M_ + row] * sc[batch * N_ + odd_col]; + float cos_val = cs[base + row * N_ + even_col]; + float sin_val = cs[base + row * N_ + odd_col]; + float out; + if ((col & 1) == 0) + out = x_even * cos_val + x_odd * sin_val; + else + out = -x_even * sin_val + x_odd * cos_val; + ref_ptr[i] = static_cast(out); + } + ); + compat::wait(); + } + + 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 tops = (2.0 * M * N * K * L) * 1e-12; + std::cout << "Problem Size: " << M << 'x' << N << 'x' << K << 'x' << L << std::endl; + printf("xe-fuse K4_W8A8 (INT8 GEMM+Dequant+RoPE): [%4.3f]TOp/s (%6.4f)ms\n", + tops / time_s, time_s * 1000); + } + + return 0; +}