From 85046a389895a4ac62e7539b924ee4063d58dc59 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 29 Jul 2026 22:32:36 -0700 Subject: [PATCH 1/3] [executorch][cuda] Optimize short-query INT4 matvec kernels Pull Request resolved: https://github.com/pytorch/executorch/pull/21473 Add compile-time M=2/3/4 fused short-query kernels while preserving the existing M=1 path. The shared ROWS template was reviewed and benchmarked by KernelAgent; KA found no faster safe rewrite, so the validated template is retained. A100 kernel microbenchmarks versus main: - M1: 0.10226 -> 0.10228 ms (performance neutral; no regression) - M2: 0.196 -> 0.127 ms (1.54x) - M3: 0.289 -> 0.159 ms (1.82x) - M4: 0.383 -> 0.196 ms (1.95x) Correctness tests cover M=1,2,3,4 with three randomized seeds. ptxas reports zero spills. ghstack-source-id: 408535989 @exported-using-ghexport Differential Revision: [D114032326](https://our.internmc.facebook.com/intern/diff/D114032326/) --- backends/cuda/runtime/shims/int4_plain_mm.cuh | 323 +++++++++++----- .../cuda/runtime/shims/tests/CMakeLists.txt | 30 ++ .../shims/tests/benchmark_int4_plain_mm.cu | 362 ++++++++++++++++++ 3 files changed, 621 insertions(+), 94 deletions(-) create mode 100644 backends/cuda/runtime/shims/tests/benchmark_int4_plain_mm.cu diff --git a/backends/cuda/runtime/shims/int4_plain_mm.cuh b/backends/cuda/runtime/shims/int4_plain_mm.cuh index 35144ad617d..be38d8634d9 100644 --- a/backends/cuda/runtime/shims/int4_plain_mm.cuh +++ b/backends/cuda/runtime/shims/int4_plain_mm.cuh @@ -136,136 +136,213 @@ __global__ void quantize_activations_q8_kernel( // those 32 group codes are contiguous => a single coalesced load. // --------------------------------------------------------------------------- -__global__ void __launch_bounds__(MV_THREADS) - int4_w4a8_matvec_coalesced_kernel( - const uint8_t* __restrict__ qdata, - const uint8_t* __restrict__ w_scale_t, // [N, n_groups] uint8 codes - const __half* __restrict__ w_scale_step, // [N, n_super] fp16 - const uint8_t* __restrict__ w_zero_t, // [N, n_groups] uint8 codes - const __half* __restrict__ w_zero_point_step, // [N, n_super] fp16 - const Q8Block* __restrict__ q8, - __nv_bfloat16* __restrict__ out, - int32_t N, - int32_t K, - int32_t gs_shift, - int32_t n_groups, - int32_t n_super) { - const int32_t n = blockIdx.x * MV_NWARPS + threadIdx.y; - const int32_t m = blockIdx.y; - if (n >= N) - return; +__device__ __forceinline__ uint32_t int4_uint4_at(uint4 value, int32_t index) { + switch (index) { + case 0: + return value.x; + case 1: + return value.y; + case 2: + return value.z; + default: + return value.w; + } +} +template +__device__ __forceinline__ void int4_w4a8_matvec_coalesced_body( + const uint8_t* __restrict__ qdata, + const uint8_t* __restrict__ w_scale_t, + const __half* __restrict__ w_scale_step, + const uint8_t* __restrict__ w_zero_t, + const __half* __restrict__ w_zero_point_step, + const Q8Block* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t n, + int32_t N, + int32_t K, + int32_t gs_shift, + int32_t n_groups, + int32_t n_super) { const int32_t K_half = K / 2; const int32_t lane_id = threadIdx.x; const int32_t n_q8_blocks = K / Q8_BLOCK_SIZE; - const uint8_t* qrow = qdata + static_cast(n) * K_half; const uint8_t* scale_row = w_scale_t + static_cast(n) * n_groups; const __half* scale_step_row = w_scale_step + static_cast(n) * n_super; const uint8_t* zero_row = w_zero_t + static_cast(n) * n_groups; - // Per-256 fp16 zero step (z_pack): decoded via the SAME 8-lane leader - // broadcast as the scale step (both packed into one 32-bit shuffle word - // below), so the dp4a dot products stay bit-identical to the scale-only - // kernel. zero = zero_code * zero_point_step[super-block]. const __half* zero_point_step_row = w_zero_point_step + static_cast(n) * n_super; - const Q8Block* q8_row = q8 + static_cast(m) * n_q8_blocks; const uint4* qrow16 = reinterpret_cast(qrow); const int32_t K_half_16 = K_half / 16; - - float sum = 0.0f; - - // T3: within a warp iteration the 32 lanes cover groups i0..i0+31 = 4 - // consecutive super-blocks, split into 8-lane subgroups (lanes 8s..8s+7 share - // super-block b = g >> sb_shift). Only each subgroup leader (lane_id % 8 == 0) - // loads + converts the two fp16 steps, PACKS them into one 32-bit word, and - // __shfl-broadcasts that single word to the 7 followers. 8x fewer step loads, - // ONE shuffle (same count as the scale-only baseline), register-only (no smem - // => no occupancy cliff). - const int32_t sb_shift = SUPER_BLOCK_SHIFT - gs_shift; // group g -> super-block - const int32_t leader = lane_id & ~7; // base lane of this 8-lane subgroup - - // Warp-aligned trip count so ALL 32 lanes execute the same number of - // iterations and therefore all reach the __shfl_sync every iteration (a - // full-mask shuffle deadlocks if some lanes exit the loop early — which - // happens when K_half_16 < 32, e.g. tiny test shapes). Out-of-range lanes do a - // safe dummy load (index 0) and contribute 0 to the accumulation. + const int32_t sb_shift = SUPER_BLOCK_SHIFT - gs_shift; + const int32_t leader = lane_id & ~7; const int32_t n_iters = ((K_half_16 + MV_WARP_SIZE - 1) / MV_WARP_SIZE) * MV_WARP_SIZE; + float sums[ROWS] = {}; for (int32_t it = 0; it < n_iters; it += MV_WARP_SIZE) { - int32_t i = it + lane_id; - bool active = i < K_half_16; - int32_t i_safe = active ? i : 0; - - uint4 packed16 = __ldg(&qrow16[i_safe]); - int32_t k_base = i_safe * 32; - uint32_t words[4] = {packed16.x, packed16.y, packed16.z, packed16.w}; - - // Group index for this uint4 (constant across its 4 dp4a words at gs=32). - int32_t g = k_base >> gs_shift; - // Subgroup leader packs BOTH per-256 fp16 steps (scale low16, zero high16) - // into one 32-bit word and broadcasts it once; followers unpack. All lanes - // reach this shuffle (warp-aligned loop), so the full mask is safe. + const int32_t i = it + lane_id; + const bool active = i < K_half_16; + const int32_t i_safe = active ? i : 0; + const uint4 packed16 = __ldg(&qrow16[i_safe]); + const int32_t k_base = i_safe * 32; + const uint32_t words[4] = { + packed16.x, packed16.y, packed16.z, packed16.w}; + const int32_t g = k_base >> gs_shift; + uint32_t steps_packed = 0; if (lane_id == leader) { - int32_t sb = g >> sb_shift; - unsigned short s_bits = __half_as_ushort(__ldg(&scale_step_row[sb])); - unsigned short z_bits = __half_as_ushort(__ldg(&zero_point_step_row[sb])); + const int32_t sb = g >> sb_shift; + const unsigned short s_bits = + __half_as_ushort(__ldg(&scale_step_row[sb])); + const unsigned short z_bits = + __half_as_ushort(__ldg(&zero_point_step_row[sb])); steps_packed = static_cast(s_bits) | (static_cast(z_bits) << 16); } steps_packed = __shfl_sync(0xffffffff, steps_packed, leader); - if (!active) + if (!active) { continue; - float scale_step = __half2float( + } + + const float scale_step = __half2float( __ushort_as_half(static_cast(steps_packed & 0xFFFF))); - float zero_point_step = __half2float( + const float zero_point_step = __half2float( __ushort_as_half(static_cast(steps_packed >> 16))); - // Effective per-group scale/zero (one coalesced code byte each per group). - float ws = static_cast(__ldg(&scale_row[g])) * scale_step; - float wz = static_cast(__ldg(&zero_row[g])) * zero_point_step; - - // One uint4 (32 weights) maps to exactly one Q8 activation block (32 - // activations), i.e. q8_block_idx == i. Load the whole block with two - // vectorized uint4 loads (+ one scale load) instead of eight scalar int32 - // loads. ae.{x,y,z,w} == qs_even[0:4],[4:8],[8:12],[12:16] == a_even for - // w=0..3 (same for ao/qs_odd) -> bit-identical to the scalar path. - const Q8Block* qb = &q8_row[i]; - uint4 ae = *reinterpret_cast(qb->qs_even); - uint4 ao = *reinterpret_cast(qb->qs_odd); - float a_scale = qb->d; - const uint32_t a_even[4] = {ae.x, ae.y, ae.z, ae.w}; - const uint32_t a_odd[4] = {ao.x, ao.y, ao.z, ao.w}; + const float ws = static_cast(__ldg(&scale_row[g])) * scale_step; + const float wz = static_cast(__ldg(&zero_row[g])) * zero_point_step; + uint4 activations_even[ROWS]; + uint4 activations_odd[ROWS]; + float activation_scales[ROWS]; #pragma unroll - for (int32_t w = 0; w < 4; w++) { - uint32_t packed = words[w]; - - int32_t vi_lo = packed & 0x0F0F0F0F; - int32_t vi_hi = (packed >> 4) & 0x0F0F0F0F; + for (int32_t row = 0; row < ROWS; ++row) { + const Q8Block* qb = + q8 + static_cast(row) * n_q8_blocks + i_safe; + activations_even[row] = + *reinterpret_cast(qb->qs_even); + activations_odd[row] = + *reinterpret_cast(qb->qs_odd); + activation_scales[row] = qb->d; + } - int32_t dp = __dp4a(vi_lo, static_cast(a_even[w]), 0); - dp = __dp4a(vi_hi, static_cast(a_odd[w]), dp); +#pragma unroll + for (int32_t w = 0; w < 4; ++w) { + const uint32_t packed = words[w]; + const int32_t vi_lo = packed & 0x0F0F0F0F; + const int32_t vi_hi = (packed >> 4) & 0x0F0F0F0F; - int32_t a_sum8 = __dp4a(0x01010101, static_cast(a_even[w]), 0); - a_sum8 = __dp4a(0x01010101, static_cast(a_odd[w]), a_sum8); +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + const uint32_t a_even = int4_uint4_at(activations_even[row], w); + const uint32_t a_odd = int4_uint4_at(activations_odd[row], w); + int32_t dp = __dp4a(vi_lo, static_cast(a_even), 0); + dp = __dp4a(vi_hi, static_cast(a_odd), dp); + int32_t a_sum8 = + __dp4a(0x01010101, static_cast(a_even), 0); + a_sum8 = + __dp4a(0x01010101, static_cast(a_odd), a_sum8); + sums[row] += ws * activation_scales[row] * + (static_cast(dp) - wz * static_cast(a_sum8)); + } + } + } - sum += ws * a_scale * - (static_cast(dp) - wz * static_cast(a_sum8)); + for (int32_t offset = MV_WARP_SIZE / 2; offset > 0; offset >>= 1) { +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + sums[row] += __shfl_xor_sync(0xffffffff, sums[row], offset); } } - for (int offset = MV_WARP_SIZE / 2; offset > 0; offset >>= 1) - sum += __shfl_xor_sync(0xffffffff, sum, offset); + if (lane_id == 0) { +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + out[static_cast(row) * N + n] = + __float2bfloat16(sums[row]); + } + } +} - if (lane_id == 0) - out[static_cast(m) * N + n] = __float2bfloat16(sum); +__global__ void __launch_bounds__(MV_THREADS) + int4_w4a8_matvec_coalesced_kernel( + const uint8_t* __restrict__ qdata, + const uint8_t* __restrict__ w_scale_t, + const __half* __restrict__ w_scale_step, + const uint8_t* __restrict__ w_zero_t, + const __half* __restrict__ w_zero_point_step, + const Q8Block* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t N, + int32_t K, + int32_t gs_shift, + int32_t n_groups, + int32_t n_super) { + const int32_t n = blockIdx.x * MV_NWARPS + threadIdx.y; + const int32_t m = blockIdx.y; + if (n >= N) { + return; + } + int4_w4a8_matvec_coalesced_body<1>( + qdata, + w_scale_t, + w_scale_step, + w_zero_t, + w_zero_point_step, + q8 + static_cast(m) * (K / Q8_BLOCK_SIZE), + out + static_cast(m) * N, + n, + N, + K, + gs_shift, + n_groups, + n_super); } -// --------------------------------------------------------------------------- +#define DEFINE_INT4_MULTIROW_KERNEL(ROWS) \ + __global__ void __launch_bounds__(MV_THREADS) \ + int4_w4a8_matvec_m##ROWS##_coalesced_kernel( \ + const uint8_t* __restrict__ qdata, \ + const uint8_t* __restrict__ w_scale_t, \ + const __half* __restrict__ w_scale_step, \ + const uint8_t* __restrict__ w_zero_t, \ + const __half* __restrict__ w_zero_point_step, \ + const Q8Block* __restrict__ q8, \ + __nv_bfloat16* __restrict__ out, \ + int32_t N, \ + int32_t K, \ + int32_t gs_shift, \ + int32_t n_groups, \ + int32_t n_super) { \ + const int32_t n = blockIdx.x * MV_NWARPS + threadIdx.y; \ + if (n >= N) { \ + return; \ + } \ + int4_w4a8_matvec_coalesced_body( \ + qdata, \ + w_scale_t, \ + w_scale_step, \ + w_zero_t, \ + w_zero_point_step, \ + q8, \ + out, \ + n, \ + N, \ + K, \ + gs_shift, \ + n_groups, \ + n_super); \ + } + +DEFINE_INT4_MULTIROW_KERNEL(2) +DEFINE_INT4_MULTIROW_KERNEL(3) +DEFINE_INT4_MULTIROW_KERNEL(4) + +#undef DEFINE_INT4_MULTIROW_KERNEL + // Persistent Q8 buffer (lazy init, not thread-safe — single-stream only). // Freed at process exit via a static guard so leak detectors stay quiet; the // CUDA runtime would otherwise reclaim it on teardown anyway. @@ -307,7 +384,7 @@ static Q8Block* get_q8_buffer(size_t needed) { // Main entry point // --------------------------------------------------------------------------- -void _int4_plain_mm_cuda( +inline void _int4_plain_mm_cuda( const Tensor& A, // [M, K] bf16 const Tensor& qdata, // [N, K//2] uint8 const Tensor& scale, // [N, K//gs] uint8 codes @@ -383,6 +460,60 @@ void _int4_plain_mm_cuda( int32_t n_groups = static_cast(scale.size(1)); int32_t n_super = static_cast(scale_step.size(1)); + if (M == 4) { + dim3 m4_grid((N + MV_NWARPS - 1) / MV_NWARPS); + int4_w4a8_matvec_m4_coalesced_kernel<<>>( + reinterpret_cast(qdata.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(scale_step.data_ptr()), + reinterpret_cast(zero.data_ptr()), + reinterpret_cast(zero_point_step.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + gs_shift, + n_groups, + n_super); + return; + } + + if (M == 3) { + dim3 m3_grid((N + MV_NWARPS - 1) / MV_NWARPS); + int4_w4a8_matvec_m3_coalesced_kernel<<>>( + reinterpret_cast(qdata.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(scale_step.data_ptr()), + reinterpret_cast(zero.data_ptr()), + reinterpret_cast(zero_point_step.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + gs_shift, + n_groups, + n_super); + return; + } + + if (M == 2) { + dim3 m2_grid((N + MV_NWARPS - 1) / MV_NWARPS); + int4_w4a8_matvec_m2_coalesced_kernel<<>>( + reinterpret_cast(qdata.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(scale_step.data_ptr()), + reinterpret_cast(zero.data_ptr()), + reinterpret_cast(zero_point_step.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + gs_shift, + n_groups, + n_super); + return; + } + int4_w4a8_matvec_coalesced_kernel<<>>( reinterpret_cast(qdata.data_ptr()), reinterpret_cast(scale.data_ptr()), @@ -391,7 +522,11 @@ void _int4_plain_mm_cuda( reinterpret_cast(zero_point_step.data_ptr()), q8_buf, reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), - N, K, gs_shift, n_groups, n_super); + N, + K, + gs_shift, + n_groups, + n_super); } } // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/tests/CMakeLists.txt b/backends/cuda/runtime/shims/tests/CMakeLists.txt index 986e54556b1..950fd8f1d91 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -74,6 +74,36 @@ foreach(test_name ${CUDA_SHIM_TESTS}) add_test(NAME ${test_name} COMMAND ${test_name}) endforeach() +add_executable(benchmark_int4_plain_mm benchmark_int4_plain_mm.cu) +target_include_directories( + benchmark_int4_plain_mm PRIVATE ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT} + ${CUDAToolkit_INCLUDE_DIRS} +) +target_compile_definitions(benchmark_int4_plain_mm PRIVATE CUDA_AVAILABLE=1) +set_property(TARGET benchmark_int4_plain_mm PROPERTY CUDA_ARCHITECTURES 80) +target_compile_options( + benchmark_int4_plain_mm PRIVATE $<$:-Xptxas=-v> +) +target_link_libraries( + benchmark_int4_plain_mm PRIVATE aoti_cuda_shims executorch_core CUDA::cudart +) +add_test(NAME benchmark_int4_plain_mm_m1 + COMMAND benchmark_int4_plain_mm --M=1 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int4_plain_mm_m2 + COMMAND benchmark_int4_plain_mm --M=2 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int4_plain_mm_m3 + COMMAND benchmark_int4_plain_mm --M=3 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int4_plain_mm_m4 + COMMAND benchmark_int4_plain_mm --M=4 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) + foreach(test_name ${CUDA_KERNEL_TESTS}) add_executable(${test_name} ${test_name}.cpp) diff --git a/backends/cuda/runtime/shims/tests/benchmark_int4_plain_mm.cu b/backends/cuda/runtime/shims/tests/benchmark_int4_plain_mm.cu new file mode 100644 index 00000000000..2a0f6e913c0 --- /dev/null +++ b/backends/cuda/runtime/shims/tests/benchmark_int4_plain_mm.cu @@ -0,0 +1,362 @@ +/* Copyright (c) Meta Platforms, Inc. and affiliates. */ + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuda_shims = executorch::backends::cuda; + +namespace { + +#define CUDA_CHECK(expr) \ + do { \ + cudaError_t err__ = (expr); \ + if (err__ != cudaSuccess) { \ + std::fprintf( \ + stderr, \ + "CUDA error %s:%d: %s\n", \ + __FILE__, \ + __LINE__, \ + cudaGetErrorString(err__)); \ + std::exit(1); \ + } \ + } while (0) + + +int32_t log2_pow2_host(int32_t v) { + int32_t r = 0; + while (v > 1) { + v >>= 1; + ++r; + } + return r; +} + +uint16_t float_to_bf16(float x) { + uint32_t bits; + std::memcpy(&bits, &x, sizeof(bits)); + return static_cast(bits >> 16); +} + +void fill_case( + int64_t M, + int64_t N, + int64_t K, + int64_t gs, + uint32_t seed, + std::vector& A, + std::vector& qdata, + std::vector& scale, + std::vector& scale_step, + std::vector& zero, + std::vector& zero_step) { + std::mt19937 rng(seed); + std::uniform_real_distribution adist(-2.0f, 2.0f); + std::uniform_int_distribution qdist(0, 15); + std::uniform_int_distribution codedist(1, 15); + std::uniform_real_distribution sdist(0.003f, 0.035f); + std::uniform_real_distribution zdist(0.25f, 1.25f); + + A.resize(M * K); + qdata.assign(N * (K / 2), 0); + scale.resize(N * (K / gs)); + scale_step.resize(N * (K / 256)); + zero.resize(N * (K / gs)); + zero_step.resize(N * (K / 256)); + + for (auto& x : A) { + x = float_to_bf16(adist(rng)); + } + for (auto& x : scale) { + x = static_cast(codedist(rng)); + } + for (auto& x : zero) { + x = static_cast(codedist(rng)); + } + for (auto& x : scale_step) { + x = __half_as_ushort(__float2half(sdist(rng))); + } + for (auto& x : zero_step) { + x = __half_as_ushort(__float2half(zdist(rng))); + } + + std::vector u(K); + for (int64_t n = 0; n < N; ++n) { + uint8_t* row = qdata.data() + n * (K / 2); + for (int64_t k = 0; k < K; ++k) { + u[k] = static_cast(qdist(rng)); + } + for (int64_t k = 0; k < K; k += 2) { + row[k / 2] = static_cast((u[k] & 0xF) | ((u[k + 1] & 0xF) << 4)); + } + } +} + +template +float time_ms(Fn fn, int warmup, int iterations) { + for (int i = 0; i < warmup; ++i) { + fn(); + } + CUDA_CHECK(cudaDeviceSynchronize()); + cudaEvent_t start, stop; + CUDA_CHECK(cudaEventCreate(&start)); + CUDA_CHECK(cudaEventCreate(&stop)); + CUDA_CHECK(cudaEventRecord(start)); + for (int i = 0; i < iterations; ++i) { + fn(); + } + CUDA_CHECK(cudaEventRecord(stop)); + CUDA_CHECK(cudaEventSynchronize(stop)); + float elapsed = 0.0f; + CUDA_CHECK(cudaEventElapsedTime(&elapsed, start, stop)); + CUDA_CHECK(cudaEventDestroy(start)); + CUDA_CHECK(cudaEventDestroy(stop)); + return elapsed / iterations; +} + +void* device_alloc_copy(const void* src, size_t bytes) { + void* dst = nullptr; + CUDA_CHECK(cudaMalloc(&dst, bytes)); + CUDA_CHECK(cudaMemcpy(dst, src, bytes, cudaMemcpyHostToDevice)); + return dst; +} + +int64_t arg_value(int argc, char** argv, const char* name, int64_t fallback) { + std::string key = std::string("--") + name + "="; + for (int i = 1; i < argc; ++i) { + std::string arg(argv[i]); + if (arg.rfind(key, 0) == 0) { + return std::stoll(arg.substr(key.size())); + } + } + return fallback; +} + +} // namespace + +int main(int argc, char** argv) { + const int64_t M = arg_value(argc, argv, "M", 2); + const int64_t K = arg_value(argc, argv, "K", 6656); + const int64_t N = arg_value(argc, argv, "N", 39936); + const int64_t gs = arg_value(argc, argv, "gs", 32); + const int64_t warmup = arg_value(argc, argv, "warmup", 30); + const int64_t iterations = arg_value(argc, argv, "iters", 200); + const int64_t seeds = arg_value(argc, argv, "seeds", 3); + + if (M < 1 || M > 4 || K % 256 != 0 || K % 32 != 0 || gs < 32 || (gs & (gs - 1))) { + std::fprintf(stderr, "unsupported shape M=%lld N=%lld K=%lld gs=%lld\n", M, N, K, gs); + return 2; + } + + CUDA_CHECK(cudaSetDevice(0)); + int32_t gs_shift = log2_pow2_host(static_cast(gs)); + int32_t n_groups = static_cast(K / gs); + int32_t n_super = static_cast(K / 256); + int32_t n_q8_blocks = static_cast(K / cuda_shims::Q8_BLOCK_SIZE); + dim3 q8_grid((n_q8_blocks + cuda_shims::MV_NWARPS - 1) / cuda_shims::MV_NWARPS, M); + dim3 block(cuda_shims::MV_WARP_SIZE, cuda_shims::MV_NWARPS); + dim3 ref_grid((N + cuda_shims::MV_NWARPS - 1) / cuda_shims::MV_NWARPS, M); + dim3 cand_grid((N + cuda_shims::MV_NWARPS - 1) / cuda_shims::MV_NWARPS); + + std::printf( + "shape M=%lld N=%lld K=%lld gs=%lld q8_blocks=%d warmup=%lld iters=%lld seeds=%lld\n", + M, + N, + K, + gs, + n_q8_blocks, + warmup, + iterations, + seeds); + + double ref_total = 0.0; + double cand_total = 0.0; + for (int64_t seed_idx = 0; seed_idx < seeds; ++seed_idx) { + std::vector hA; + std::vector hqdata; + std::vector hscale; + std::vector hscale_step; + std::vector hzero; + std::vector hzero_step; + fill_case( + M, + N, + K, + gs, + 2027 + seed_idx * 17, + hA, + hqdata, + hscale, + hscale_step, + hzero, + hzero_step); + + auto* dA = static_cast<__nv_bfloat16*>(device_alloc_copy(hA.data(), hA.size() * sizeof(uint16_t))); + auto* dqdata = static_cast(device_alloc_copy(hqdata.data(), hqdata.size() * sizeof(uint8_t))); + auto* dscale = static_cast(device_alloc_copy(hscale.data(), hscale.size() * sizeof(uint8_t))); + auto* dscale_step = static_cast<__half*>(device_alloc_copy(hscale_step.data(), hscale_step.size() * sizeof(uint16_t))); + auto* dzero = static_cast(device_alloc_copy(hzero.data(), hzero.size() * sizeof(uint8_t))); + auto* dzero_step = static_cast<__half*>(device_alloc_copy(hzero_step.data(), hzero_step.size() * sizeof(uint16_t))); + cuda_shims::Q8Block* dq8 = nullptr; + __nv_bfloat16* dref = nullptr; + __nv_bfloat16* dcand = nullptr; + CUDA_CHECK(cudaMalloc(&dq8, M * n_q8_blocks * sizeof(cuda_shims::Q8Block))); + CUDA_CHECK(cudaMalloc(&dref, M * N * sizeof(__nv_bfloat16))); + CUDA_CHECK(cudaMalloc(&dcand, M * N * sizeof(__nv_bfloat16))); + + cuda_shims::quantize_activations_q8_kernel<<>>(dA, dq8, static_cast(K)); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + + cuda_shims::int4_w4a8_matvec_coalesced_kernel<<>>( + dqdata, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dref, + static_cast(N), + static_cast(K), + gs_shift, + n_groups, + n_super); + if (M == 3) { + cuda_shims::int4_w4a8_matvec_m3_coalesced_kernel<<>>( + dqdata, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dcand, + static_cast(N), + static_cast(K), + gs_shift, + n_groups, + n_super); + } else if (M == 2) { + cuda_shims::int4_w4a8_matvec_m2_coalesced_kernel<<>>( + dqdata, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dcand, + static_cast(N), + static_cast(K), + gs_shift, + n_groups, + n_super); + } else if (M == 4) { + cuda_shims::int4_w4a8_matvec_m4_coalesced_kernel<<>>( + dqdata, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dcand, + static_cast(N), + static_cast(K), + gs_shift, + n_groups, + n_super); + } else { + cuda_shims::int4_w4a8_matvec_coalesced_kernel<<>>( + dqdata, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dcand, + static_cast(N), + static_cast(K), + gs_shift, + n_groups, + n_super); + } + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + + std::vector href(M * N); + std::vector hcand(M * N); + CUDA_CHECK(cudaMemcpy(href.data(), dref, href.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaMemcpy(hcand.data(), dcand, hcand.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + size_t mismatches = 0; + for (size_t i = 0; i < href.size(); ++i) { + if (href[i] != hcand[i]) { + if (mismatches < 8) { + std::printf("mismatch seed=%lld idx=%zu ref=0x%04x cand=0x%04x\n", seed_idx, i, href[i], hcand[i]); + } + ++mismatches; + } + } + if (mismatches != 0) { + std::printf("bitwise=FAIL mismatches=%zu/%zu\n", mismatches, href.size()); + return 1; + } + + float q8_ms = time_ms([&]() { + cuda_shims::quantize_activations_q8_kernel<<>>(dA, dq8, static_cast(K)); + }, warmup, iterations); + float ref_ms = time_ms([&]() { + cuda_shims::int4_w4a8_matvec_coalesced_kernel<<>>( + dqdata, dscale, dscale_step, dzero, dzero_step, dq8, dref, static_cast(N), static_cast(K), gs_shift, n_groups, n_super); + }, warmup, iterations); + float cand_ms = time_ms([&]() { + if (M == 3) { + cuda_shims::int4_w4a8_matvec_m3_coalesced_kernel<<>>( + dqdata, dscale, dscale_step, dzero, dzero_step, dq8, dcand, static_cast(N), static_cast(K), gs_shift, n_groups, n_super); + } else if (M == 2) { + cuda_shims::int4_w4a8_matvec_m2_coalesced_kernel<<>>( + dqdata, dscale, dscale_step, dzero, dzero_step, dq8, dcand, static_cast(N), static_cast(K), gs_shift, n_groups, n_super); + } else if (M == 4) { + cuda_shims::int4_w4a8_matvec_m4_coalesced_kernel<<>>( + dqdata, dscale, dscale_step, dzero, dzero_step, dq8, dcand, static_cast(N), static_cast(K), gs_shift, n_groups, n_super); + } else { + cuda_shims::int4_w4a8_matvec_coalesced_kernel<<>>( + dqdata, dscale, dscale_step, dzero, dzero_step, dq8, dcand, static_cast(N), static_cast(K), gs_shift, n_groups, n_super); + } + }, warmup, iterations); + ref_total += ref_ms; + cand_total += cand_ms; + std::printf( + "seed=%lld bitwise=OK q8_ms=%.6f ref_matvec_ms=%.6f cand_matvec_ms=%.6f speedup=%.3fx\n", + seed_idx, + q8_ms, + ref_ms, + cand_ms, + ref_ms / cand_ms); + + CUDA_CHECK(cudaFree(dA)); + CUDA_CHECK(cudaFree(dqdata)); + CUDA_CHECK(cudaFree(dscale)); + CUDA_CHECK(cudaFree(dscale_step)); + CUDA_CHECK(cudaFree(dzero)); + CUDA_CHECK(cudaFree(dzero_step)); + CUDA_CHECK(cudaFree(dq8)); + CUDA_CHECK(cudaFree(dref)); + CUDA_CHECK(cudaFree(dcand)); + } + + std::printf( + "avg ref_matvec_ms=%.6f cand_matvec_ms=%.6f speedup=%.3fx\n", + ref_total / seeds, + cand_total / seeds, + ref_total / cand_total); + return 0; +} From 0b48798c77e6f3a5c7a5fa7e4475f7fded081bb9 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 29 Jul 2026 22:32:36 -0700 Subject: [PATCH 2/3] [executorch][cuda] Optimize short-query INT6 matvec kernels Pull Request resolved: https://github.com/pytorch/executorch/pull/21474 Add compile-time M=1/2/3/4 short-query paths and share their safe implementation through a ROWS/USE_SUM/GS16 device template. These kernels were generated and tuned with KernelAgent. Dispatch uses the regular gs16 kernels for M=1/2 and the precomputed-activation-sum kernels for M=3/4. The M=3 non-sum entry remains explicit as a performance/correctness comparison because templating that entry increased register pressure. A100 op-level results (N=6656, gs=16, three seeds): - M1, generic reference -> gs16: 0.039416 -> 0.038211 ms (1.032x; no regression) - M2, generic reference -> gs16: 0.055348 -> 0.032540 ms (1.701x) - M3, optimized non-sum -> activation-sum: 0.040003 -> 0.039949 ms (1.001x; neutral incremental change) - M4, generic reference -> activation-sum: 0.100204 -> 0.045694 ms (2.193x) The M3 baseline is intentionally the optimized non-sum kernel rather than the generic reference; this isolates the activation-sum optimization. All candidates are bitwise-correct across the three seeds, and ptxas reports zero spills. ghstack-source-id: 408539235 @exported-using-ghexport Differential Revision: [D114032330](https://our.internmc.facebook.com/intern/diff/D114032330/) --- backends/cuda/runtime/shims/int6_plain_mm.cuh | 461 ++++++++++++++--- .../cuda/runtime/shims/tests/CMakeLists.txt | 30 ++ .../shims/tests/benchmark_int6_plain_mm.cu | 486 ++++++++++++++++++ 3 files changed, 912 insertions(+), 65 deletions(-) create mode 100644 backends/cuda/runtime/shims/tests/benchmark_int6_plain_mm.cu diff --git a/backends/cuda/runtime/shims/int6_plain_mm.cuh b/backends/cuda/runtime/shims/int6_plain_mm.cuh index 6a2dc6e9020..125f707d1f9 100644 --- a/backends/cuda/runtime/shims/int6_plain_mm.cuh +++ b/backends/cuda/runtime/shims/int6_plain_mm.cuh @@ -111,6 +111,7 @@ __device__ __forceinline__ uint32_t spread2_i6(uint32_t b) { struct alignas(16) Q8Block_i6 { int8_t qs_even[Q8_BLOCK_SIZE_I6 / 2]; int8_t qs_odd[Q8_BLOCK_SIZE_I6 / 2]; + int16_t sum8[4]; float d; // scale }; @@ -145,6 +146,15 @@ __global__ void quantize_activations_q8_i6_kernel( else dst->qs_odd[lane / 2] = static_cast(q); + int32_t sum8 = q; +#pragma unroll + for (int offset = 4; offset > 0; offset >>= 1) { + sum8 += __shfl_xor_sync(0xffffffff, sum8, offset); + } + if ((lane & 7) == 0) { + dst->sum8[lane >> 3] = static_cast(sum8); + } + if (lane == 0) dst->d = d; } @@ -166,6 +176,156 @@ __global__ void quantize_activations_q8_i6_kernel( // per word. Register-only (no smem => no occupancy cliff). // --------------------------------------------------------------------------- +__device__ __forceinline__ uint32_t uint4_at_i6(uint4 v, int32_t i) { + return i == 0 ? v.x : (i == 1 ? v.y : (i == 2 ? v.z : v.w)); +} + + +__device__ __forceinline__ void accum_i6( + int32_t vfull_even, + int32_t vfull_odd, + uint32_t a_even, + uint32_t a_odd, + float scale, + float& sum) { + int32_t dp = __dp4a(vfull_even, static_cast(a_even), 0); + dp = __dp4a(vfull_odd, static_cast(a_odd), dp); + int32_t a_sum = __dp4a(0x01010101, static_cast(a_even), 0); + a_sum = __dp4a(0x01010101, static_cast(a_odd), a_sum); + sum += scale * (static_cast(dp) - 32.0f * static_cast(a_sum)); +} + +__device__ __forceinline__ void accum_i6_sum( + int32_t vfull_even, + int32_t vfull_odd, + uint32_t a_even, + uint32_t a_odd, + float scale, + int32_t a_sum, + float& sum) { + int32_t dp = __dp4a(vfull_even, static_cast(a_even), 0); + dp = __dp4a(vfull_odd, static_cast(a_odd), dp); + sum += scale * (static_cast(dp) - 32.0f * static_cast(a_sum)); +} + +template +__device__ __forceinline__ void int6_w6a8_matvec_body( + const uint8_t* __restrict__ ql, + const uint8_t* __restrict__ qh, + const int8_t* __restrict__ w_scale, + const __half* __restrict__ w_scale_step, + const Q8Block_i6* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t n, + int32_t N, + int32_t K, + int32_t gs_shift, + int32_t n_groups, + int32_t n_super) { + const int32_t K_half = K / 2; + const int32_t K_quarter = K / 4; + const int32_t lane_id = threadIdx.x; + const int32_t n_q8_blocks = K / Q8_BLOCK_SIZE_I6; + + const uint8_t* qlrow = ql + static_cast(n) * K_half; + const uint8_t* qhrow = qh + static_cast(n) * K_quarter; + const int8_t* scale_row = w_scale + static_cast(n) * n_groups; + const __half* scale_step_row = + w_scale_step + static_cast(n) * n_super; + + const uint4* qlrow16 = reinterpret_cast(qlrow); + const uint2* qhrow8 = reinterpret_cast(qhrow); + const int32_t K_half_16 = K_half / 16; + const int32_t sb_shift = SUPER_BLOCK_SHIFT_I6 - gs_shift; + const int32_t wpg_shift = gs_shift - 3; + float sums[ROWS] = {}; + + for (int32_t i = lane_id; i < K_half_16; i += MV6_WARP_SIZE) { + const uint4 packed16 = __ldg(&qlrow16[i]); + const uint2 qh_chunk = __ldg(&qhrow8[i]); + const int32_t k_base = i * 32; + const uint32_t words[4] = {packed16.x, packed16.y, packed16.z, packed16.w}; + const uint32_t hi_even_word = qh_chunk.x; + const uint32_t hi_odd_word = qh_chunk.y; + + const int32_t g_base = GS16 ? (k_base >> 4) : (k_base >> gs_shift); + const float scale_step = __half2float( + __ldg(&scale_step_row[GS16 ? (g_base >> 4) : (g_base >> sb_shift)])); + const float ws0 = + static_cast(__ldg(&scale_row[g_base])) * scale_step; + const float ws1 = + static_cast(__ldg(&scale_row[g_base + 1])) * scale_step; + + uint4 activations_even[ROWS]; + uint4 activations_odd[ROWS]; + float activation_scales[ROWS]; + int16_t activation_sums[ROWS][4]; +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + const Q8Block_i6* qb = q8 + static_cast(row) * n_q8_blocks + i; + activations_even[row] = *reinterpret_cast(qb->qs_even); + activations_odd[row] = *reinterpret_cast(qb->qs_odd); + activation_scales[row] = qb->d; + if constexpr (USE_SUM) { + activation_sums[row][0] = qb->sum8[0]; + activation_sums[row][1] = qb->sum8[1]; + activation_sums[row][2] = qb->sum8[2]; + activation_sums[row][3] = qb->sum8[3]; + } + } + +#pragma unroll + for (int32_t w = 0; w < 4; w++) { + const uint32_t packed = words[w]; + const int32_t vi_lo = static_cast(packed & 0x0F0F0F0F); + const int32_t vi_hi = static_cast((packed >> 4) & 0x0F0F0F0F); + const uint32_t hi_even_byte = (hi_even_word >> (w * 8)) & 0xFF; + const uint32_t hi_odd_byte = (hi_odd_word >> (w * 8)) & 0xFF; + const int32_t vfull_even = + vi_lo | static_cast(spread2_i6(hi_even_byte) << 4); + const int32_t vfull_odd = + vi_hi | static_cast(spread2_i6(hi_odd_byte) << 4); + const float ws = GS16 ? ((w >= 2) ? ws1 : ws0) + : ((w >> wpg_shift) ? ws1 : ws0); + +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + if constexpr (USE_SUM) { + accum_i6_sum( + vfull_even, + vfull_odd, + uint4_at_i6(activations_even[row], w), + uint4_at_i6(activations_odd[row], w), + ws * activation_scales[row], + static_cast(activation_sums[row][w]), + sums[row]); + } else { + accum_i6( + vfull_even, + vfull_odd, + uint4_at_i6(activations_even[row], w), + uint4_at_i6(activations_odd[row], w), + ws * activation_scales[row], + sums[row]); + } + } + } + } + + for (int offset = MV6_WARP_SIZE / 2; offset > 0; offset >>= 1) { +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + sums[row] += __shfl_xor_sync(0xffffffff, sums[row], offset); + } + } + if (lane_id == 0) { +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + out[static_cast(row) * N + n] = __float2bfloat16(sums[row]); + } + } +} + __global__ void __launch_bounds__(MV6_THREADS) int6_w6a8_matvec_kernel( const uint8_t* __restrict__ ql, // [N, K/2] const uint8_t* __restrict__ qh, // [N, K/4] @@ -175,11 +335,69 @@ __global__ void __launch_bounds__(MV6_THREADS) int6_w6a8_matvec_kernel( __nv_bfloat16* __restrict__ out, int32_t N, int32_t K, + int32_t M, int32_t gs_shift, int32_t n_groups, int32_t n_super) { const int32_t n = blockIdx.x * MV6_NWARPS + threadIdx.y; - const int32_t m = blockIdx.y; + if (n >= N) { + return; + } + if (M == 1) { + int6_w6a8_matvec_body<1, false, false>( + ql, qh, w_scale, w_scale_step, q8, out, n, N, K, gs_shift, n_groups, n_super); + return; + } + if (M == 2) { + int6_w6a8_matvec_body<2, false, false>( + ql, qh, w_scale, w_scale_step, q8, out, n, N, K, gs_shift, n_groups, n_super); + return; + } + if (M == 3) { + int6_w6a8_matvec_body<3, false, false>( + ql, qh, w_scale, w_scale_step, q8, out, n, N, K, gs_shift, n_groups, n_super); + return; + } + int6_w6a8_matvec_body<4, false, false>( + ql, qh, w_scale, w_scale_step, q8, out, n, N, K, gs_shift, n_groups, n_super); +} + +#define DEFINE_INT6_GS16_KERNEL(ROWS) \ + __global__ void __launch_bounds__(MV6_THREADS) \ + int6_w6a8_matvec_m##ROWS##_gs16_kernel( \ + const uint8_t* __restrict__ ql, \ + const uint8_t* __restrict__ qh, \ + const int8_t* __restrict__ w_scale, \ + const __half* __restrict__ w_scale_step, \ + const Q8Block_i6* __restrict__ q8, \ + __nv_bfloat16* __restrict__ out, \ + int32_t N, \ + int32_t K, \ + int32_t n_groups, \ + int32_t n_super) { \ + const int32_t n = blockIdx.x * MV6_NWARPS + threadIdx.y; \ + if (n >= N) { \ + return; \ + } \ + int6_w6a8_matvec_body( \ + ql, qh, w_scale, w_scale_step, q8, out, n, N, K, 4, n_groups, n_super); \ + } + +DEFINE_INT6_GS16_KERNEL(1) +DEFINE_INT6_GS16_KERNEL(2) + +__global__ void __launch_bounds__(MV6_THREADS) int6_w6a8_matvec_m3_gs16_kernel( + const uint8_t* __restrict__ ql, + const uint8_t* __restrict__ qh, + const int8_t* __restrict__ w_scale, + const __half* __restrict__ w_scale_step, + const Q8Block_i6* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t N, + int32_t K, + int32_t n_groups, + int32_t n_super) { + const int32_t n = blockIdx.x * MV6_NWARPS + threadIdx.y; if (n >= N) return; @@ -193,104 +411,139 @@ __global__ void __launch_bounds__(MV6_THREADS) int6_w6a8_matvec_kernel( const int8_t* scale_row = w_scale + static_cast(n) * n_groups; const __half* scale_step_row = w_scale_step + static_cast(n) * n_super; - const Q8Block_i6* q8_row = q8 + static_cast(m) * n_q8_blocks; - // Vectorized loads: one uint4 of ql (32 weights) + one uint2 of qh (the - // 8 high-bit bytes for the same 32-weight chunk) per iteration. const uint4* qlrow16 = reinterpret_cast(qlrow); const uint2* qhrow8 = reinterpret_cast(qhrow); const int32_t K_half_16 = K_half / 16; - - float sum = 0.0f; - - // Per-uint4 __ldg of the per-256 fp16 step: each lane's uint4 is one 256 - // super-block per 8 lanes, so within a warp iteration the 32 lanes touch only - // 4 distinct scale_step addresses (8 lanes share each). The [N, K/256] step - // tensor is tiny and L1-resident, so those 8-way-shared __ldg's hit the - // read-only cache (hardware broadcast). No warp-shuffle needed. - const int32_t sb_shift = SUPER_BLOCK_SHIFT_I6 - gs_shift; // group -> super-block + float s0 = 0.0f; + float s1 = 0.0f; + float s2 = 0.0f; for (int32_t i = lane_id; i < K_half_16; i += MV6_WARP_SIZE) { uint4 packed16 = __ldg(&qlrow16[i]); uint2 qh_chunk = __ldg(&qhrow8[i]); int32_t k_base = i * 32; uint32_t words[4] = {packed16.x, packed16.y, packed16.z, packed16.w}; - // qh_chunk.x bytes = hi_even_packed[0..3], qh_chunk.y = - // hi_odd_packed[0..3]. uint32_t hi_even_word = qh_chunk.x; uint32_t hi_odd_word = qh_chunk.y; - // Load the per-256 fp16 step for this uint4's super-block once (constant - // across its 4 dp4a words: gs=16 => a uint4 spans 2 groups but only ONE - // super-block). 8 lanes share this address => read-only-cache broadcast. - int32_t g_base = k_base >> gs_shift; // first group of this uint4 + int32_t g_base = k_base >> 4; float scale_step = - __half2float(__ldg(&scale_step_row[g_base >> sb_shift])); - - // One uint4 (32 weights) maps to exactly one Q8 activation block (32 - // activations), i.e. q8_block_idx == i. Load the whole block with two - // vectorized uint4 loads. - const Q8Block_i6* qb = &q8_row[i]; - uint4 ae = *reinterpret_cast(qb->qs_even); - uint4 ao = *reinterpret_cast(qb->qs_odd); - float a_scale = qb->d; - const uint32_t a_even[4] = {ae.x, ae.y, ae.z, ae.w}; - const uint32_t a_odd[4] = {ao.x, ao.y, ao.z, ao.w}; - - // Hoist the per-group int8 scale-code loads out of the w-loop. A uint4 (32 - // weights) spans 32/gs groups: gs=16 -> exactly 2 (words {0,1} in g_base, - // words {2,3} in g_base+1). Load those 2 codes once here instead of 4 times - // inside the w-loop (the OLD/naive kernel re-__ldg'd the code every word). - // This mirrors int4's per-uint4 code hoist (the real op-level win there); - // gs=16 => 2 code loads/uint4 vs 4. Two int8 scalars (not a float[4] array) - // keep register pressure at the OLD kernel's level. gs is asserted a - // multiple of 8, and K a multiple of 256, so a uint4 spans at most 2 groups - // for gs in {16, 32}; gs=8 (4 groups) is out of the supported Q6_K range. + __half2float(__ldg(&scale_step_row[g_base >> 4])); const float ws0 = static_cast(__ldg(&scale_row[g_base])) * scale_step; const float ws1 = static_cast(__ldg(&scale_row[g_base + 1])) * scale_step; - const int32_t wpg_shift = gs_shift - 3; // log2(words per group) + + const Q8Block_i6* qb0 = q8 + i; + const Q8Block_i6* qb1 = q8 + n_q8_blocks + i; + const Q8Block_i6* qb2 = q8 + static_cast(2) * n_q8_blocks + i; + uint4 ae0 = *reinterpret_cast(qb0->qs_even); + uint4 ao0 = *reinterpret_cast(qb0->qs_odd); + uint4 ae1 = *reinterpret_cast(qb1->qs_even); + uint4 ao1 = *reinterpret_cast(qb1->qs_odd); + uint4 ae2 = *reinterpret_cast(qb2->qs_even); + uint4 ao2 = *reinterpret_cast(qb2->qs_odd); + float as0 = qb0->d; + float as1 = qb1->d; + float as2 = qb2->d; #pragma unroll for (int32_t w = 0; w < 4; w++) { uint32_t packed = words[w]; - // Effective per-group scale: coalesced int8 code * the per-256 step (the - // step factors out of the dp4a sum, so the dot products stay bit-identical - // to the bf16-metadata kernel modulo the step dtype). At gs=16 words - // {0,1}->ws0, {2,3}->ws1 (w >> wpg_shift selects the group within the - // uint4). - float ws = (w >> wpg_shift) ? ws1 : ws0; - int32_t vi_lo = static_cast(packed & 0x0F0F0F0F); int32_t vi_hi = static_cast((packed >> 4) & 0x0F0F0F0F); - uint32_t hi_even_byte = (hi_even_word >> (w * 8)) & 0xFF; uint32_t hi_odd_byte = (hi_odd_word >> (w * 8)) & 0xFF; - - // Reconstruct full 6-bit weight bytes (u in [0, 63]). int32_t vfull_even = vi_lo | static_cast(spread2_i6(hi_even_byte) << 4); int32_t vfull_odd = vi_hi | static_cast(spread2_i6(hi_odd_byte) << 4); + float ws = (w >= 2) ? ws1 : ws0; + accum_i6( + vfull_even, + vfull_odd, + uint4_at_i6(ae0, w), + uint4_at_i6(ao0, w), + ws * as0, + s0); + accum_i6( + vfull_even, + vfull_odd, + uint4_at_i6(ae1, w), + uint4_at_i6(ao1, w), + ws * as1, + s1); + accum_i6( + vfull_even, + vfull_odd, + uint4_at_i6(ae2, w), + uint4_at_i6(ao2, w), + ws * as2, + s2); + } + } - int32_t dp = __dp4a(vfull_even, static_cast(a_even[w]), 0); - dp = __dp4a(vfull_odd, static_cast(a_odd[w]), dp); - - int32_t a_sum = __dp4a(0x01010101, static_cast(a_even[w]), 0); - a_sum = __dp4a(0x01010101, static_cast(a_odd[w]), a_sum); + for (int offset = MV6_WARP_SIZE / 2; offset > 0; offset >>= 1) { + s0 += __shfl_xor_sync(0xffffffff, s0, offset); + s1 += __shfl_xor_sync(0xffffffff, s1, offset); + s2 += __shfl_xor_sync(0xffffffff, s2, offset); + } + if (lane_id == 0) { + out[n] = __float2bfloat16(s0); + out[static_cast(N) + n] = __float2bfloat16(s1); + out[static_cast(2) * N + n] = __float2bfloat16(s2); + } +} - // q = u - 32, so the -32 offset replaces the per-group zero point. - sum += ws * a_scale * - (static_cast(dp) - 32.0f * static_cast(a_sum)); - } +DEFINE_INT6_GS16_KERNEL(4) + +#undef DEFINE_INT6_GS16_KERNEL + +#define DEFINE_INT6_SUM_GS16_KERNEL(ROWS) \ + __global__ void __launch_bounds__(MV6_THREADS) \ + int6_w6a8_matvec_m##ROWS##_sum_gs16_kernel( \ + const uint8_t* __restrict__ ql, \ + const uint8_t* __restrict__ qh, \ + const int8_t* __restrict__ w_scale, \ + const __half* __restrict__ w_scale_step, \ + const Q8Block_i6* __restrict__ q8, \ + __nv_bfloat16* __restrict__ out, \ + int32_t N, \ + int32_t K, \ + int32_t n_groups, \ + int32_t n_super) { \ + const int32_t n = blockIdx.x * MV6_NWARPS + threadIdx.y; \ + if (n >= N) { \ + return; \ + } \ + int6_w6a8_matvec_body( \ + ql, qh, w_scale, w_scale_step, q8, out, n, N, K, 4, n_groups, n_super); \ } - for (int offset = MV6_WARP_SIZE / 2; offset > 0; offset >>= 1) - sum += __shfl_xor_sync(0xffffffff, sum, offset); +DEFINE_INT6_SUM_GS16_KERNEL(3) +DEFINE_INT6_SUM_GS16_KERNEL(4) + +#undef DEFINE_INT6_SUM_GS16_KERNEL - if (lane_id == 0) - out[static_cast(m) * N + n] = __float2bfloat16(sum); +__global__ void __launch_bounds__(MV6_THREADS) int6_w6a8_matvec_m4_kernel( + const uint8_t* __restrict__ ql, + const uint8_t* __restrict__ qh, + const int8_t* __restrict__ w_scale, + const __half* __restrict__ w_scale_step, + const Q8Block_i6* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t N, + int32_t K, + int32_t gs_shift, + int32_t n_groups, + int32_t n_super) { + const int32_t n = blockIdx.x * MV6_NWARPS + threadIdx.y; + if (n >= N) { + return; + } + int6_w6a8_matvec_body<4, false, false>( + ql, qh, w_scale, w_scale_step, q8, out, n, N, K, gs_shift, n_groups, n_super); } // --------------------------------------------------------------------------- @@ -348,6 +601,7 @@ inline void _int6_plain_mm_cuda( int32_t N = ql.size(0); ET_CHECK(A.dtype() == c10::ScalarType::BFloat16); + ET_CHECK_MSG(M >= 1 && M <= 4, "int6 short-query M=%d must be in [1, 4]", M); ET_CHECK( ql.dtype() == c10::ScalarType::Byte || ql.dtype() == c10::ScalarType::Char); @@ -408,11 +662,87 @@ inline void _int6_plain_mm_cuda( reinterpret_cast(A.data_ptr()), q8_buf, K); // dp4a matvec - dim3 grid((N + MV6_NWARPS - 1) / MV6_NWARPS, M); + dim3 grid((N + MV6_NWARPS - 1) / MV6_NWARPS); dim3 block(MV6_WARP_SIZE, MV6_NWARPS); int32_t n_groups = static_cast(scale.size(1)); int32_t n_super = static_cast(steps.size(1)); + if (M == 4 && gs == 16) { + int6_w6a8_matvec_m4_sum_gs16_kernel<<>>( + reinterpret_cast(ql.data_ptr()), + reinterpret_cast(qh.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(steps.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + n_groups, + n_super); + return; + } + + if (M == 3 && gs == 16 && N >= 1024) { + int6_w6a8_matvec_m3_sum_gs16_kernel<<>>( + reinterpret_cast(ql.data_ptr()), + reinterpret_cast(qh.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(steps.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + n_groups, + n_super); + return; + } + + if (M == 2 && gs == 16 && N >= 1024) { + int6_w6a8_matvec_m2_gs16_kernel<<>>( + reinterpret_cast(ql.data_ptr()), + reinterpret_cast(qh.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(steps.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + n_groups, + n_super); + return; + } + + if (M == 1 && gs == 16 && N >= 1024 && K >= 8192) { + int6_w6a8_matvec_m1_gs16_kernel<<>>( + reinterpret_cast(ql.data_ptr()), + reinterpret_cast(qh.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(steps.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + n_groups, + n_super); + return; + } + + if (M == 4) { + int6_w6a8_matvec_m4_kernel<<>>( + reinterpret_cast(ql.data_ptr()), + reinterpret_cast(qh.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(steps.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + gs_shift, + n_groups, + n_super); + return; + } + int6_w6a8_matvec_kernel<<>>( reinterpret_cast(ql.data_ptr()), reinterpret_cast(qh.data_ptr()), @@ -422,6 +752,7 @@ inline void _int6_plain_mm_cuda( reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), N, K, + M, gs_shift, n_groups, n_super); diff --git a/backends/cuda/runtime/shims/tests/CMakeLists.txt b/backends/cuda/runtime/shims/tests/CMakeLists.txt index 950fd8f1d91..7464805d88b 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -104,6 +104,36 @@ add_test(NAME benchmark_int4_plain_mm_m4 --warmup=1 --iters=1 --seeds=3 ) +add_executable(benchmark_int6_plain_mm benchmark_int6_plain_mm.cu) +target_include_directories( + benchmark_int6_plain_mm PRIVATE ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT} + ${CUDAToolkit_INCLUDE_DIRS} +) +target_compile_definitions(benchmark_int6_plain_mm PRIVATE CUDA_AVAILABLE=1) +set_property(TARGET benchmark_int6_plain_mm PROPERTY CUDA_ARCHITECTURES 80) +target_compile_options( + benchmark_int6_plain_mm PRIVATE $<$:-Xptxas=-v> +) +target_link_libraries( + benchmark_int6_plain_mm PRIVATE aoti_cuda_shims executorch_core CUDA::cudart +) +add_test(NAME benchmark_int6_plain_mm_m1 + COMMAND benchmark_int6_plain_mm --M=1 --N=1024 --K=8192 --gs=16 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int6_plain_mm_m2 + COMMAND benchmark_int6_plain_mm --M=2 --N=1024 --K=256 --gs=16 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int6_plain_mm_m3 + COMMAND benchmark_int6_plain_mm --M=3 --N=1024 --K=256 --gs=16 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int6_plain_mm_m4 + COMMAND benchmark_int6_plain_mm --M=4 --N=64 --K=256 --gs=16 + --warmup=1 --iters=1 --seeds=3 +) + foreach(test_name ${CUDA_KERNEL_TESTS}) add_executable(${test_name} ${test_name}.cpp) diff --git a/backends/cuda/runtime/shims/tests/benchmark_int6_plain_mm.cu b/backends/cuda/runtime/shims/tests/benchmark_int6_plain_mm.cu new file mode 100644 index 00000000000..d25476d6190 --- /dev/null +++ b/backends/cuda/runtime/shims/tests/benchmark_int6_plain_mm.cu @@ -0,0 +1,486 @@ +/* Copyright (c) Meta Platforms, Inc. and affiliates. */ + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuda_shims = executorch::backends::cuda; + +namespace { + +#define CUDA_CHECK(expr) \ + do { \ + cudaError_t err__ = (expr); \ + if (err__ != cudaSuccess) { \ + std::fprintf( \ + stderr, \ + "CUDA error %s:%d: %s\n", \ + __FILE__, \ + __LINE__, \ + cudaGetErrorString(err__)); \ + std::exit(1); \ + } \ + } while (0) + +__global__ void __launch_bounds__(cuda_shims::MV6_THREADS) + int6_w6a8_matvec_reference_kernel( + const uint8_t* __restrict__ ql, + const uint8_t* __restrict__ qh, + const int8_t* __restrict__ w_scale, + const __half* __restrict__ w_scale_step, + const cuda_shims::Q8Block_i6* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t N, + int32_t K, + int32_t M, + int32_t gs_shift, + int32_t n_groups, + int32_t n_super) { + const int32_t n = blockIdx.x * cuda_shims::MV6_NWARPS + threadIdx.y; + const int32_t m = blockIdx.y; + if (n >= N || m >= M) { + return; + } + + const int32_t K_half = K / 2; + const int32_t K_quarter = K / 4; + const int32_t lane_id = threadIdx.x; + const int32_t n_q8_blocks = K / cuda_shims::Q8_BLOCK_SIZE_I6; + + const uint8_t* qlrow = ql + static_cast(n) * K_half; + const uint8_t* qhrow = qh + static_cast(n) * K_quarter; + const int8_t* scale_row = w_scale + static_cast(n) * n_groups; + const __half* scale_step_row = + w_scale_step + static_cast(n) * n_super; + const cuda_shims::Q8Block_i6* q8_row = + q8 + static_cast(m) * n_q8_blocks; + + const uint4* qlrow16 = reinterpret_cast(qlrow); + const uint2* qhrow8 = reinterpret_cast(qhrow); + const int32_t K_half_16 = K_half / 16; + const int32_t sb_shift = cuda_shims::SUPER_BLOCK_SHIFT_I6 - gs_shift; + const int32_t wpg_shift = gs_shift - 3; + + float sum = 0.0f; + for (int32_t i = lane_id; i < K_half_16; i += cuda_shims::MV6_WARP_SIZE) { + uint4 packed16 = __ldg(&qlrow16[i]); + uint2 qh_chunk = __ldg(&qhrow8[i]); + int32_t k_base = i * 32; + uint32_t words[4] = {packed16.x, packed16.y, packed16.z, packed16.w}; + uint32_t hi_even_word = qh_chunk.x; + uint32_t hi_odd_word = qh_chunk.y; + int32_t g_base = k_base >> gs_shift; + float scale_step = + __half2float(__ldg(&scale_step_row[g_base >> sb_shift])); + float ws0 = static_cast(__ldg(&scale_row[g_base])) * scale_step; + float ws1 = static_cast(__ldg(&scale_row[g_base + 1])) * scale_step; + const cuda_shims::Q8Block_i6* qb = &q8_row[i]; + uint4 ae = *reinterpret_cast(qb->qs_even); + uint4 ao = *reinterpret_cast(qb->qs_odd); + float a_scale = qb->d; + uint32_t a_even[4] = {ae.x, ae.y, ae.z, ae.w}; + uint32_t a_odd[4] = {ao.x, ao.y, ao.z, ao.w}; + +#pragma unroll + for (int32_t w = 0; w < 4; ++w) { + uint32_t packed = words[w]; + int32_t vi_lo = static_cast(packed & 0x0F0F0F0F); + int32_t vi_hi = static_cast((packed >> 4) & 0x0F0F0F0F); + uint32_t hi_even_byte = (hi_even_word >> (w * 8)) & 0xFF; + uint32_t hi_odd_byte = (hi_odd_word >> (w * 8)) & 0xFF; + int32_t vfull_even = + vi_lo | static_cast(cuda_shims::spread2_i6(hi_even_byte) << 4); + int32_t vfull_odd = + vi_hi | static_cast(cuda_shims::spread2_i6(hi_odd_byte) << 4); + int32_t dp = __dp4a(vfull_even, static_cast(a_even[w]), 0); + dp = __dp4a(vfull_odd, static_cast(a_odd[w]), dp); + int32_t a_sum = + __dp4a(0x01010101, static_cast(a_even[w]), 0); + a_sum = __dp4a(0x01010101, static_cast(a_odd[w]), a_sum); + float ws = (w >> wpg_shift) ? ws1 : ws0; + sum += ws * a_scale * + (static_cast(dp) - 32.0f * static_cast(a_sum)); + } + } + + for (int offset = cuda_shims::MV6_WARP_SIZE / 2; offset > 0; offset >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, offset); + } + if (lane_id == 0) { + out[static_cast(m) * N + n] = __float2bfloat16(sum); + } +} + +int32_t log2_pow2_host(int32_t v) { + int32_t r = 0; + while (v > 1) { + v >>= 1; + ++r; + } + return r; +} + +uint16_t float_to_bf16(float x) { + uint32_t bits; + std::memcpy(&bits, &x, sizeof(bits)); + return static_cast(bits >> 16); +} + +void fill_case( + int64_t M, + int64_t N, + int64_t K, + int64_t gs, + uint32_t seed, + std::vector& A, + std::vector& ql, + std::vector& qh, + std::vector& scale, + std::vector& steps) { + std::mt19937 rng(seed); + std::uniform_real_distribution adist(-2.0f, 2.0f); + std::uniform_int_distribution qdist(0, 63); + std::uniform_int_distribution sdist(-8, 8); + std::uniform_real_distribution stepdist(0.003f, 0.035f); + + A.resize(M * K); + ql.assign(N * (K / 2), 0); + qh.assign(N * (K / 4), 0); + scale.resize(N * (K / gs)); + steps.resize(N * (K / 256)); + + for (auto& x : A) { + x = float_to_bf16(adist(rng)); + } + for (auto& x : scale) { + int v = sdist(rng); + x = static_cast(v == 0 ? 1 : v); + } + for (auto& x : steps) { + x = __half_as_ushort(__float2half(stepdist(rng))); + } + + std::vector u(K); + for (int64_t n = 0; n < N; ++n) { + uint8_t* qlrow = ql.data() + n * (K / 2); + uint8_t* qhrow = qh.data() + n * (K / 4); + for (int64_t k = 0; k < K; ++k) { + u[k] = static_cast(qdist(rng)); + } + for (int64_t k = 0; k < K; k += 2) { + qlrow[k / 2] = static_cast((u[k] & 0xF) | ((u[k + 1] & 0xF) << 4)); + } + for (int64_t k = 0; k < K; k += 32) { + uint8_t* chunk = qhrow + k / 4; + for (int w = 0; w < 4; ++w) { + uint8_t he = 0; + uint8_t ho = 0; + for (int j = 0; j < 4; ++j) { + he |= static_cast(((u[k + w * 8 + j * 2] >> 4) & 0x3) << (j * 2)); + ho |= static_cast(((u[k + w * 8 + j * 2 + 1] >> 4) & 0x3) << (j * 2)); + } + chunk[w] = he; + chunk[4 + w] = ho; + } + } + } +} + +template +float time_ms(Fn fn, int warmup, int iterations) { + for (int i = 0; i < warmup; ++i) { + fn(); + } + CUDA_CHECK(cudaDeviceSynchronize()); + cudaEvent_t start, stop; + CUDA_CHECK(cudaEventCreate(&start)); + CUDA_CHECK(cudaEventCreate(&stop)); + CUDA_CHECK(cudaEventRecord(start)); + for (int i = 0; i < iterations; ++i) { + fn(); + } + CUDA_CHECK(cudaEventRecord(stop)); + CUDA_CHECK(cudaEventSynchronize(stop)); + float elapsed = 0.0f; + CUDA_CHECK(cudaEventElapsedTime(&elapsed, start, stop)); + CUDA_CHECK(cudaEventDestroy(start)); + CUDA_CHECK(cudaEventDestroy(stop)); + return elapsed / iterations; +} + +void* device_alloc_copy(const void* src, size_t bytes) { + void* dst = nullptr; + CUDA_CHECK(cudaMalloc(&dst, bytes)); + CUDA_CHECK(cudaMemcpy(dst, src, bytes, cudaMemcpyHostToDevice)); + return dst; +} + +int64_t arg_value(int argc, char** argv, const char* name, int64_t fallback) { + std::string key = std::string("--") + name + "="; + for (int i = 1; i < argc; ++i) { + std::string arg(argv[i]); + if (arg.rfind(key, 0) == 0) { + return std::stoll(arg.substr(key.size())); + } + } + return fallback; +} + +} // namespace + +int main(int argc, char** argv) { + const int64_t M = arg_value(argc, argv, "M", 4); + const int64_t K = arg_value(argc, argv, "K", 6656); + const int64_t N = arg_value(argc, argv, "N", 6656); + const int64_t gs = arg_value(argc, argv, "gs", 16); + const int64_t warmup = arg_value(argc, argv, "warmup", 30); + const int64_t iterations = arg_value(argc, argv, "iters", 200); + const int64_t seeds = arg_value(argc, argv, "seeds", 3); + + if (M < 1 || M > 4 || K % 256 != 0 || K % 32 != 0 || gs < 8 || (gs & (gs - 1))) { + std::fprintf(stderr, "unsupported shape M=%lld N=%lld K=%lld gs=%lld\n", M, N, K, gs); + return 2; + } + + CUDA_CHECK(cudaSetDevice(0)); + int32_t gs_shift = log2_pow2_host(static_cast(gs)); + int32_t n_groups = static_cast(K / gs); + int32_t n_super = static_cast(K / 256); + int32_t n_q8_blocks = static_cast(K / cuda_shims::Q8_BLOCK_SIZE_I6); + dim3 q8_grid((n_q8_blocks + cuda_shims::MV6_NWARPS - 1) / cuda_shims::MV6_NWARPS, M); + dim3 block(cuda_shims::MV6_WARP_SIZE, cuda_shims::MV6_NWARPS); + dim3 ref_grid((N + cuda_shims::MV6_NWARPS - 1) / cuda_shims::MV6_NWARPS, M); + dim3 cand_grid((N + cuda_shims::MV6_NWARPS - 1) / cuda_shims::MV6_NWARPS); + + std::printf( + "shape M=%lld N=%lld K=%lld gs=%lld q8_blocks=%d warmup=%lld iters=%lld seeds=%lld\n", + M, + N, + K, + gs, + n_q8_blocks, + warmup, + iterations, + seeds); + + double ref_total = 0.0; + double cand_total = 0.0; + for (int64_t seed_idx = 0; seed_idx < seeds; ++seed_idx) { + std::vector hA; + std::vector hql; + std::vector hqh; + std::vector hscale; + std::vector hsteps; + fill_case(M, N, K, gs, 1009 + seed_idx * 17, hA, hql, hqh, hscale, hsteps); + + auto* dA = static_cast<__nv_bfloat16*>(device_alloc_copy(hA.data(), hA.size() * sizeof(uint16_t))); + auto* dql = static_cast(device_alloc_copy(hql.data(), hql.size() * sizeof(uint8_t))); + auto* dqh = static_cast(device_alloc_copy(hqh.data(), hqh.size() * sizeof(uint8_t))); + auto* dscale = static_cast(device_alloc_copy(hscale.data(), hscale.size() * sizeof(int8_t))); + auto* dsteps = static_cast<__half*>(device_alloc_copy(hsteps.data(), hsteps.size() * sizeof(uint16_t))); + cuda_shims::Q8Block_i6* dq8 = nullptr; + __nv_bfloat16* dref = nullptr; + __nv_bfloat16* dcand = nullptr; + CUDA_CHECK(cudaMalloc(&dq8, M * n_q8_blocks * sizeof(cuda_shims::Q8Block_i6))); + CUDA_CHECK(cudaMalloc(&dref, M * N * sizeof(__nv_bfloat16))); + CUDA_CHECK(cudaMalloc(&dcand, M * N * sizeof(__nv_bfloat16))); + + cuda_shims::quantize_activations_q8_i6_kernel<<>>(dA, dq8, static_cast(K)); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + + if (M == 3 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m3_gs16_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dref, + static_cast(N), + static_cast(K), + n_groups, + n_super); + } else { + int6_w6a8_matvec_reference_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dref, + static_cast(N), + static_cast(K), + static_cast(M), + gs_shift, + n_groups, + n_super); + } + if (M == 4 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m4_sum_gs16_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dcand, + static_cast(N), + static_cast(K), + n_groups, + n_super); + } else if (M == 3 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m3_sum_gs16_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dcand, + static_cast(N), + static_cast(K), + n_groups, + n_super); + } else if (M == 2 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m2_gs16_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dcand, + static_cast(N), + static_cast(K), + n_groups, + n_super); + } else if (M == 1 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m1_gs16_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dcand, + static_cast(N), + static_cast(K), + n_groups, + n_super); + } else if (M == 4) { + cuda_shims::int6_w6a8_matvec_m4_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dcand, + static_cast(N), + static_cast(K), + gs_shift, + n_groups, + n_super); + } else { + cuda_shims::int6_w6a8_matvec_kernel<<>>( + dql, + dqh, + dscale, + dsteps, + dq8, + dcand, + static_cast(N), + static_cast(K), + static_cast(M), + gs_shift, + n_groups, + n_super); + } + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + + std::vector href(M * N); + std::vector hcand(M * N); + CUDA_CHECK(cudaMemcpy(href.data(), dref, href.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaMemcpy(hcand.data(), dcand, hcand.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + size_t mismatches = 0; + for (size_t i = 0; i < href.size(); ++i) { + if (href[i] != hcand[i]) { + if (mismatches < 8) { + std::printf("mismatch seed=%lld idx=%zu ref=0x%04x cand=0x%04x\n", seed_idx, i, href[i], hcand[i]); + } + ++mismatches; + } + } + if (mismatches != 0) { + std::printf("bitwise=FAIL mismatches=%zu/%zu\n", mismatches, href.size()); + return 1; + } + + float q8_ms = time_ms([&]() { + cuda_shims::quantize_activations_q8_i6_kernel<<>>(dA, dq8, static_cast(K)); + }, warmup, iterations); + float ref_ms = time_ms([&]() { + if (M == 3 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m3_gs16_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dref, static_cast(N), static_cast(K), n_groups, n_super); + } else { + int6_w6a8_matvec_reference_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dref, static_cast(N), static_cast(K), static_cast(M), gs_shift, n_groups, n_super); + } + }, warmup, iterations); + float cand_ms = time_ms([&]() { + if (M == 4 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m4_sum_gs16_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dcand, static_cast(N), static_cast(K), n_groups, n_super); + } else if (M == 3 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m3_sum_gs16_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dcand, static_cast(N), static_cast(K), n_groups, n_super); + } else if (M == 2 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m2_gs16_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dcand, static_cast(N), static_cast(K), n_groups, n_super); + } else if (M == 1 && gs == 16) { + cuda_shims::int6_w6a8_matvec_m1_gs16_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dcand, static_cast(N), static_cast(K), n_groups, n_super); + } else if (M == 4) { + cuda_shims::int6_w6a8_matvec_m4_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dcand, static_cast(N), static_cast(K), gs_shift, n_groups, n_super); + } else { + cuda_shims::int6_w6a8_matvec_kernel<<>>( + dql, dqh, dscale, dsteps, dq8, dcand, static_cast(N), static_cast(K), static_cast(M), gs_shift, n_groups, n_super); + } + }, warmup, iterations); + ref_total += ref_ms; + cand_total += cand_ms; + std::printf( + "seed=%lld bitwise=OK q8_ms=%.6f ref_matvec_ms=%.6f cand_matvec_ms=%.6f speedup=%.3fx\n", + seed_idx, + q8_ms, + ref_ms, + cand_ms, + ref_ms / cand_ms); + + CUDA_CHECK(cudaFree(dA)); + CUDA_CHECK(cudaFree(dql)); + CUDA_CHECK(cudaFree(dqh)); + CUDA_CHECK(cudaFree(dscale)); + CUDA_CHECK(cudaFree(dsteps)); + CUDA_CHECK(cudaFree(dq8)); + CUDA_CHECK(cudaFree(dref)); + CUDA_CHECK(cudaFree(dcand)); + } + + std::printf( + "avg ref_matvec_ms=%.6f cand_matvec_ms=%.6f speedup=%.3fx\n", + ref_total / seeds, + cand_total / seeds, + ref_total / cand_total); + return 0; +} From 34b05fd0a18f873bde364a3a554e8f6186a7f1ac Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 29 Jul 2026 22:32:37 -0700 Subject: [PATCH 3/3] [executorch][cuda] Optimize short-query INT5 matvec kernels Pull Request resolved: https://github.com/pytorch/executorch/pull/21475 Add specialized M=1 and M=4 short-query paths for the planar INT5 W5A8 matvec operator. These are the latency-sensitive single-row and four-row shapes. M=2/M=3 retain the generic fallback path; this change does not claim those shapes as optimized. The M=1 implementation shares safe gs32 decode/accumulate logic through a compile-time ROWS device template. M=4 remains explicit because KernelAgent experiments with the generalized template increased register pressure and regressed latency. The final kernels were generated and tuned with KernelAgent. A100 kernel microbenchmarks versus the previous generic path: - M1: 1.09971 -> 0.70441 ms (1.56x; no regression) - M4: 1.96445 -> 1.12976 ms (1.74x) - M2/M3: no specialized dispatch; generic fallback remains unchanged Correctness tests cover M=1,2,3,4 with three randomized seeds; ptxas reports zero spills. ghstack-source-id: 408539258 @exported-using-ghexport Differential Revision: [D114032323](https://our.internmc.facebook.com/intern/diff/D114032323/) --- backends/cuda/runtime/shims/int5_plain_mm.cuh | 430 +++++++++++++++++- .../cuda/runtime/shims/tests/CMakeLists.txt | 30 ++ .../shims/tests/benchmark_int5_plain_mm.cu | 367 +++++++++++++++ 3 files changed, 804 insertions(+), 23 deletions(-) create mode 100644 backends/cuda/runtime/shims/tests/benchmark_int5_plain_mm.cu diff --git a/backends/cuda/runtime/shims/int5_plain_mm.cuh b/backends/cuda/runtime/shims/int5_plain_mm.cuh index d85f69eeb1a..8d3003b2378 100644 --- a/backends/cuda/runtime/shims/int5_plain_mm.cuh +++ b/backends/cuda/runtime/shims/int5_plain_mm.cuh @@ -167,11 +167,11 @@ __global__ void __launch_bounds__(MV5_THREADS) int5_w5a8_matvec_kernel( __nv_bfloat16* __restrict__ out, int32_t N, int32_t K, + int32_t M, int32_t gs_shift, int32_t n_groups, int32_t n_super) { const int32_t n = blockIdx.x * MV5_NWARPS + threadIdx.y; - const int32_t m = blockIdx.y; if (n >= N) return; @@ -191,7 +191,6 @@ __global__ void __launch_bounds__(MV5_THREADS) int5_w5a8_matvec_kernel( // below), so the dp4a dot products stay bit-identical to the scale-only // kernel. zero = zero_code * zero_point_step[super-block]. const __half* zero_point_step_row = w_zero_point_step + static_cast(n) * n_super; - const Q8Block_i5* q8_row = q8 + static_cast(m) * n_q8_blocks; // Vectorized loads: one uint4 of ql (32 weights) + one uint (the 4 high-bit // bytes for the same 32-weight chunk) per iteration. @@ -199,7 +198,7 @@ __global__ void __launch_bounds__(MV5_THREADS) int5_w5a8_matvec_kernel( const uint32_t* qhrow4 = reinterpret_cast(qhrow); const int32_t K_half_16 = K_half / 16; - float sum = 0.0f; + float sum[4] = {0.0f, 0.0f, 0.0f, 0.0f}; // z_pack: within a warp iteration the 32 lanes cover uint4 indices it..it+31 // = up to 4 consecutive super-blocks, split into 8-lane subgroups (lanes @@ -259,13 +258,6 @@ __global__ void __launch_bounds__(MV5_THREADS) int5_w5a8_matvec_kernel( // One uint4 (32 weights) maps to exactly one Q8 activation block (32 // activations), i.e. q8_block_idx == i. Load the whole block with two // vectorized uint4 loads (+ one scale/zero load). - const Q8Block_i5* qb = &q8_row[i]; - uint4 ae = *reinterpret_cast(qb->qs_even); - uint4 ao = *reinterpret_cast(qb->qs_odd); - float a_scale = qb->d; - const uint32_t a_even[4] = {ae.x, ae.y, ae.z, ae.w}; - const uint32_t a_odd[4] = {ao.x, ao.y, ao.z, ao.w}; - #pragma unroll for (int32_t w = 0; w < 4; w++) { uint32_t packed = words[w]; @@ -286,24 +278,381 @@ __global__ void __launch_bounds__(MV5_THREADS) int5_w5a8_matvec_kernel( int32_t vfull_odd = vi_hi | static_cast(spread1_i5(hi_odd_nib) << 4); - int32_t dp = __dp4a(vfull_even, static_cast(a_even[w]), 0); - dp = __dp4a(vfull_odd, static_cast(a_odd[w]), dp); + for (int32_t m = 0; m < M; ++m) { + const Q8Block_i5* qb = + &q8[static_cast(m) * n_q8_blocks + i]; + const uint4 ae = *reinterpret_cast(qb->qs_even); + const uint4 ao = *reinterpret_cast(qb->qs_odd); + const uint32_t a_even[4] = {ae.x, ae.y, ae.z, ae.w}; + const uint32_t a_odd[4] = {ao.x, ao.y, ao.z, ao.w}; + int32_t dp = + __dp4a(vfull_even, static_cast(a_even[w]), 0); + dp = __dp4a(vfull_odd, static_cast(a_odd[w]), dp); + int32_t a_sum = + __dp4a(0x01010101, static_cast(a_even[w]), 0); + a_sum = __dp4a( + 0x01010101, static_cast(a_odd[w]), a_sum); + sum[m] += ws * qb->d * + (static_cast(dp) - wz * static_cast(a_sum)); + } + } + } + + for (int32_t m = 0; m < M; ++m) { + for (int offset = MV5_WARP_SIZE / 2; offset > 0; offset >>= 1) { + sum[m] += __shfl_xor_sync(0xffffffff, sum[m], offset); + } + if (lane_id == 0) { + out[static_cast(m) * N + n] = __float2bfloat16(sum[m]); + } + } +} + + + +__device__ __forceinline__ void accum_i5_gs32_word( + int32_t vfull_even, + int32_t vfull_odd, + uint32_t a_even, + uint32_t a_odd, + float ws, + float wz, + float a_scale, + float& sum) { + int32_t dp = __dp4a(vfull_even, static_cast(a_even), 0); + dp = __dp4a(vfull_odd, static_cast(a_odd), dp); + int32_t a_sum = __dp4a(0x01010101, static_cast(a_even), 0); + a_sum = __dp4a(0x01010101, static_cast(a_odd), a_sum); + sum += ws * a_scale * + (static_cast(dp) - wz * static_cast(a_sum)); +} + +__device__ __forceinline__ void accum_i5_gs32_row( + int32_t v_even0, + int32_t v_odd0, + int32_t v_even1, + int32_t v_odd1, + int32_t v_even2, + int32_t v_odd2, + int32_t v_even3, + int32_t v_odd3, + uint4 ae, + uint4 ao, + float ws, + float wz, + float a_scale, + float& sum) { + accum_i5_gs32_word(v_even0, v_odd0, ae.x, ao.x, ws, wz, a_scale, sum); + accum_i5_gs32_word(v_even1, v_odd1, ae.y, ao.y, ws, wz, a_scale, sum); + accum_i5_gs32_word(v_even2, v_odd2, ae.z, ao.z, ws, wz, a_scale, sum); + accum_i5_gs32_word(v_even3, v_odd3, ae.w, ao.w, ws, wz, a_scale, sum); +} + +template +__device__ __forceinline__ void int5_w5a8_matvec_gs32_body( + const uint8_t* __restrict__ ql, + const uint8_t* __restrict__ qh, + const uint8_t* __restrict__ w_scale, + const __half* __restrict__ w_scale_step, + const uint8_t* __restrict__ w_zero, + const __half* __restrict__ w_zero_point_step, + const Q8Block_i5* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t n, + int32_t N, + int32_t K, + int32_t n_groups, + int32_t n_super) { + const int32_t K_half = K / 2; + const int32_t K_eighth = K / 8; + const int32_t lane_id = threadIdx.x; + const int32_t n_q8_blocks = K / Q8_BLOCK_SIZE_I5; - int32_t a_sum = __dp4a(0x01010101, static_cast(a_even[w]), 0); - a_sum = __dp4a(0x01010101, static_cast(a_odd[w]), a_sum); + const uint8_t* qlrow = ql + static_cast(n) * K_half; + const uint8_t* qhrow = qh + static_cast(n) * K_eighth; + const uint8_t* scale_row = w_scale + static_cast(n) * n_groups; + const __half* scale_step_row = + w_scale_step + static_cast(n) * n_super; + const uint8_t* zero_row = w_zero + static_cast(n) * n_groups; + const __half* zero_point_step_row = + w_zero_point_step + static_cast(n) * n_super; - // Asymmetric: w = scale * (u - zero), so the per-group zero point is - // subtracted from the dp4a sum (weighted by the activation sum). - sum += ws * a_scale * - (static_cast(dp) - wz * static_cast(a_sum)); + const uint4* qlrow16 = reinterpret_cast(qlrow); + const uint32_t* qhrow4 = reinterpret_cast(qhrow); + const int32_t K_half_16 = K_half / 16; + + float sums[ROWS] = {}; + const int32_t leader = lane_id & ~7; + const int32_t n_iters = + ((K_half_16 + MV5_WARP_SIZE - 1) / MV5_WARP_SIZE) * MV5_WARP_SIZE; + + for (int32_t it = 0; it < n_iters; it += MV5_WARP_SIZE) { + const int32_t i = it + lane_id; + const bool active = i < K_half_16; + const int32_t i_safe = active ? i : 0; + + const uint4 packed16 = __ldg(&qlrow16[i_safe]); + const uint32_t qh_word = __ldg(&qhrow4[i_safe]); + + uint32_t steps_packed = 0; + if (lane_id == leader) { + const int32_t sb = i_safe >> 3; + const unsigned short s_bits = __half_as_ushort(__ldg(&scale_step_row[sb])); + const unsigned short z_bits = + __half_as_ushort(__ldg(&zero_point_step_row[sb])); + steps_packed = + static_cast(s_bits) | (static_cast(z_bits) << 16); + } + steps_packed = __shfl_sync(0xffffffff, steps_packed, leader); + if (!active) { + continue; + } + + const float scale_step = __half2float( + __ushort_as_half(static_cast(steps_packed & 0xFFFF))); + const float zero_point_step = __half2float( + __ushort_as_half(static_cast(steps_packed >> 16))); + const float ws = static_cast(__ldg(&scale_row[i])) * scale_step; + const float wz = static_cast(__ldg(&zero_row[i])) * zero_point_step; + + const uint32_t hi_byte0 = qh_word & 0xFF; + const uint32_t hi_byte1 = (qh_word >> 8) & 0xFF; + const uint32_t hi_byte2 = (qh_word >> 16) & 0xFF; + const uint32_t hi_byte3 = (qh_word >> 24) & 0xFF; + const int32_t v_even0 = static_cast(packed16.x & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte0 & 0xF) << 4); + const int32_t v_odd0 = static_cast((packed16.x >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte0 >> 4) << 4); + const int32_t v_even1 = static_cast(packed16.y & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte1 & 0xF) << 4); + const int32_t v_odd1 = static_cast((packed16.y >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte1 >> 4) << 4); + const int32_t v_even2 = static_cast(packed16.z & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte2 & 0xF) << 4); + const int32_t v_odd2 = static_cast((packed16.z >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte2 >> 4) << 4); + const int32_t v_even3 = static_cast(packed16.w & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte3 & 0xF) << 4); + const int32_t v_odd3 = static_cast((packed16.w >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte3 >> 4) << 4); + +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + const Q8Block_i5* qb = + &q8[static_cast(row) * n_q8_blocks + i]; + const uint4 ae = *reinterpret_cast(qb->qs_even); + const uint4 ao = *reinterpret_cast(qb->qs_odd); + accum_i5_gs32_row( + v_even0, + v_odd0, + v_even1, + v_odd1, + v_even2, + v_odd2, + v_even3, + v_odd3, + ae, + ao, + ws, + wz, + qb->d, + sums[row]); } } - for (int offset = MV5_WARP_SIZE / 2; offset > 0; offset >>= 1) - sum += __shfl_xor_sync(0xffffffff, sum, offset); + for (int offset = MV5_WARP_SIZE / 2; offset > 0; offset >>= 1) { +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + sums[row] += __shfl_xor_sync(0xffffffff, sums[row], offset); + } + } + if (lane_id == 0) { +#pragma unroll + for (int32_t row = 0; row < ROWS; ++row) { + out[static_cast(row) * N + n] = __float2bfloat16(sums[row]); + } + } +} - if (lane_id == 0) - out[static_cast(m) * N + n] = __float2bfloat16(sum); +__global__ void __launch_bounds__(MV5_THREADS) int5_w5a8_matvec_m1_gs32_kernel( + const uint8_t* __restrict__ ql, // [N, K/2] + const uint8_t* __restrict__ qh, // [N, K/8] + const uint8_t* __restrict__ w_scale, // [N, K/32] uint8 codes + const __half* __restrict__ w_scale_step, // [N, K/256] fp16 + const uint8_t* __restrict__ w_zero, // [N, K/32] uint8 codes + const __half* __restrict__ w_zero_point_step, // [N, K/256] fp16 + const Q8Block_i5* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t N, + int32_t K, + int32_t n_groups, + int32_t n_super) { + const int32_t n = blockIdx.x * MV5_NWARPS + threadIdx.y; + if (n >= N) { + return; + } + int5_w5a8_matvec_gs32_body<1>( + ql, + qh, + w_scale, + w_scale_step, + w_zero, + w_zero_point_step, + q8, + out, + n, + N, + K, + n_groups, + n_super); +} + +__global__ void __launch_bounds__(MV5_THREADS) int5_w5a8_matvec_m4_gs32_kernel( + const uint8_t* __restrict__ ql, // [N, K/2] + const uint8_t* __restrict__ qh, // [N, K/8] + const uint8_t* __restrict__ w_scale, // [N, K/32] uint8 codes + const __half* __restrict__ w_scale_step, // [N, K/256] fp16 + const uint8_t* __restrict__ w_zero, // [N, K/32] uint8 codes + const __half* __restrict__ w_zero_point_step, // [N, K/256] fp16 + const Q8Block_i5* __restrict__ q8, + __nv_bfloat16* __restrict__ out, + int32_t N, + int32_t K, + int32_t n_groups, + int32_t n_super) { + const int32_t n = blockIdx.x * MV5_NWARPS + threadIdx.y; + if (n >= N) { + return; + } + + const int32_t K_half = K / 2; + const int32_t K_eighth = K / 8; + const int32_t lane_id = threadIdx.x; + const int32_t n_q8_blocks = K / Q8_BLOCK_SIZE_I5; + + const uint8_t* qlrow = ql + static_cast(n) * K_half; + const uint8_t* qhrow = qh + static_cast(n) * K_eighth; + const uint8_t* scale_row = w_scale + static_cast(n) * n_groups; + const __half* scale_step_row = + w_scale_step + static_cast(n) * n_super; + const uint8_t* zero_row = w_zero + static_cast(n) * n_groups; + const __half* zero_point_step_row = + w_zero_point_step + static_cast(n) * n_super; + + const uint4* qlrow16 = reinterpret_cast(qlrow); + const uint32_t* qhrow4 = reinterpret_cast(qhrow); + const int32_t K_half_16 = K_half / 16; + + float sum0 = 0.0f; + float sum1 = 0.0f; + float sum2 = 0.0f; + float sum3 = 0.0f; + const int32_t leader = lane_id & ~7; + const int32_t n_iters = + ((K_half_16 + MV5_WARP_SIZE - 1) / MV5_WARP_SIZE) * MV5_WARP_SIZE; + + for (int32_t it = 0; it < n_iters; it += MV5_WARP_SIZE) { + const int32_t i = it + lane_id; + const bool active = i < K_half_16; + const int32_t i_safe = active ? i : 0; + + const uint4 packed16 = __ldg(&qlrow16[i_safe]); + const uint32_t qh_word = __ldg(&qhrow4[i_safe]); + + uint32_t steps_packed = 0; + if (lane_id == leader) { + const int32_t sb = i_safe >> 3; // gs=32: 8 groups per 256-weight super-block. + const unsigned short s_bits = __half_as_ushort(__ldg(&scale_step_row[sb])); + const unsigned short z_bits = + __half_as_ushort(__ldg(&zero_point_step_row[sb])); + steps_packed = + static_cast(s_bits) | (static_cast(z_bits) << 16); + } + steps_packed = __shfl_sync(0xffffffff, steps_packed, leader); + if (!active) { + continue; + } + + const float scale_step = __half2float( + __ushort_as_half(static_cast(steps_packed & 0xFFFF))); + const float zero_point_step = __half2float( + __ushort_as_half(static_cast(steps_packed >> 16))); + const float ws = static_cast(__ldg(&scale_row[i])) * scale_step; + const float wz = static_cast(__ldg(&zero_row[i])) * zero_point_step; + + const uint32_t hi_byte0 = qh_word & 0xFF; + const uint32_t hi_byte1 = (qh_word >> 8) & 0xFF; + const uint32_t hi_byte2 = (qh_word >> 16) & 0xFF; + const uint32_t hi_byte3 = (qh_word >> 24) & 0xFF; + const int32_t v_even0 = static_cast(packed16.x & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte0 & 0xF) << 4); + const int32_t v_odd0 = static_cast((packed16.x >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte0 >> 4) << 4); + const int32_t v_even1 = static_cast(packed16.y & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte1 & 0xF) << 4); + const int32_t v_odd1 = static_cast((packed16.y >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte1 >> 4) << 4); + const int32_t v_even2 = static_cast(packed16.z & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte2 & 0xF) << 4); + const int32_t v_odd2 = static_cast((packed16.z >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte2 >> 4) << 4); + const int32_t v_even3 = static_cast(packed16.w & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte3 & 0xF) << 4); + const int32_t v_odd3 = static_cast((packed16.w >> 4) & 0x0F0F0F0F) | + static_cast(spread1_i5(hi_byte3 >> 4) << 4); + +#define ACCUM_INT5_M4_ROW(row_, sum_) \ + do { \ + const Q8Block_i5* qb = \ + &q8[static_cast(row_) * n_q8_blocks + i]; \ + const uint4 ae = *reinterpret_cast(qb->qs_even); \ + const uint4 ao = *reinterpret_cast(qb->qs_odd); \ + const float a_scale = qb->d; \ + int32_t dp = __dp4a(v_even0, static_cast(ae.x), 0); \ + dp = __dp4a(v_odd0, static_cast(ao.x), dp); \ + int32_t a_sum = __dp4a(0x01010101, static_cast(ae.x), 0); \ + a_sum = __dp4a(0x01010101, static_cast(ao.x), a_sum); \ + (sum_) += ws * a_scale * \ + (static_cast(dp) - wz * static_cast(a_sum)); \ + dp = __dp4a(v_even1, static_cast(ae.y), 0); \ + dp = __dp4a(v_odd1, static_cast(ao.y), dp); \ + a_sum = __dp4a(0x01010101, static_cast(ae.y), 0); \ + a_sum = __dp4a(0x01010101, static_cast(ao.y), a_sum); \ + (sum_) += ws * a_scale * \ + (static_cast(dp) - wz * static_cast(a_sum)); \ + dp = __dp4a(v_even2, static_cast(ae.z), 0); \ + dp = __dp4a(v_odd2, static_cast(ao.z), dp); \ + a_sum = __dp4a(0x01010101, static_cast(ae.z), 0); \ + a_sum = __dp4a(0x01010101, static_cast(ao.z), a_sum); \ + (sum_) += ws * a_scale * \ + (static_cast(dp) - wz * static_cast(a_sum)); \ + dp = __dp4a(v_even3, static_cast(ae.w), 0); \ + dp = __dp4a(v_odd3, static_cast(ao.w), dp); \ + a_sum = __dp4a(0x01010101, static_cast(ae.w), 0); \ + a_sum = __dp4a(0x01010101, static_cast(ao.w), a_sum); \ + (sum_) += ws * a_scale * \ + (static_cast(dp) - wz * static_cast(a_sum)); \ + } while (0) + + ACCUM_INT5_M4_ROW(0, sum0); + ACCUM_INT5_M4_ROW(1, sum1); + ACCUM_INT5_M4_ROW(2, sum2); + ACCUM_INT5_M4_ROW(3, sum3); +#undef ACCUM_INT5_M4_ROW + } + + for (int offset = MV5_WARP_SIZE / 2; offset > 0; offset >>= 1) { + sum0 += __shfl_xor_sync(0xffffffff, sum0, offset); + sum1 += __shfl_xor_sync(0xffffffff, sum1, offset); + sum2 += __shfl_xor_sync(0xffffffff, sum2, offset); + sum3 += __shfl_xor_sync(0xffffffff, sum3, offset); + } + if (lane_id == 0) { + out[n] = __float2bfloat16(sum0); + out[static_cast(N) + n] = __float2bfloat16(sum1); + out[static_cast(2) * N + n] = __float2bfloat16(sum2); + out[static_cast(3) * N + n] = __float2bfloat16(sum3); + } } // --------------------------------------------------------------------------- @@ -363,6 +712,7 @@ inline void _int5_plain_mm_cuda( int32_t N = ql.size(0); ET_CHECK(A.dtype() == c10::ScalarType::BFloat16); + ET_CHECK_MSG(M >= 1 && M <= 4, "int5 short-query M=%d must be in [1, 4]", M); ET_CHECK( ql.dtype() == c10::ScalarType::Byte || ql.dtype() == c10::ScalarType::Char); @@ -431,11 +781,44 @@ inline void _int5_plain_mm_cuda( reinterpret_cast(A.data_ptr()), q8_buf, K); // dp4a matvec - dim3 grid((N + MV5_NWARPS - 1) / MV5_NWARPS, M); + dim3 grid((N + MV5_NWARPS - 1) / MV5_NWARPS); dim3 block(MV5_WARP_SIZE, MV5_NWARPS); int32_t n_groups = static_cast(scale.size(1)); int32_t n_super = static_cast(scale_step.size(1)); + if (M == 1 && gs == 32) { + int5_w5a8_matvec_m1_gs32_kernel<<>>( + reinterpret_cast(ql.data_ptr()), + reinterpret_cast(qh.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(scale_step.data_ptr()), + reinterpret_cast(zero.data_ptr()), + reinterpret_cast(zero_point_step.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + n_groups, + n_super); + return; + } + if (M == 4 && gs == 32) { + int5_w5a8_matvec_m4_gs32_kernel<<>>( + reinterpret_cast(ql.data_ptr()), + reinterpret_cast(qh.data_ptr()), + reinterpret_cast(scale.data_ptr()), + reinterpret_cast(scale_step.data_ptr()), + reinterpret_cast(zero.data_ptr()), + reinterpret_cast(zero_point_step.data_ptr()), + q8_buf, + reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), + N, + K, + n_groups, + n_super); + return; + } + int5_w5a8_matvec_kernel<<>>( reinterpret_cast(ql.data_ptr()), reinterpret_cast(qh.data_ptr()), @@ -447,6 +830,7 @@ inline void _int5_plain_mm_cuda( reinterpret_cast<__nv_bfloat16*>(output->data_ptr()), N, K, + M, gs_shift, n_groups, n_super); diff --git a/backends/cuda/runtime/shims/tests/CMakeLists.txt b/backends/cuda/runtime/shims/tests/CMakeLists.txt index 7464805d88b..a5c01f42939 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -134,6 +134,36 @@ add_test(NAME benchmark_int6_plain_mm_m4 --warmup=1 --iters=1 --seeds=3 ) +add_executable(benchmark_int5_plain_mm benchmark_int5_plain_mm.cu) +target_include_directories( + benchmark_int5_plain_mm PRIVATE ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT} + ${CUDAToolkit_INCLUDE_DIRS} +) +target_compile_definitions(benchmark_int5_plain_mm PRIVATE CUDA_AVAILABLE=1) +set_property(TARGET benchmark_int5_plain_mm PROPERTY CUDA_ARCHITECTURES 80) +target_compile_options( + benchmark_int5_plain_mm PRIVATE $<$:-Xptxas=-v> +) +target_link_libraries( + benchmark_int5_plain_mm PRIVATE aoti_cuda_shims executorch_core CUDA::cudart +) +add_test(NAME benchmark_int5_plain_mm_m1 + COMMAND benchmark_int5_plain_mm --M=1 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int5_plain_mm_m2 + COMMAND benchmark_int5_plain_mm --M=2 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int5_plain_mm_m3 + COMMAND benchmark_int5_plain_mm --M=3 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) +add_test(NAME benchmark_int5_plain_mm_m4 + COMMAND benchmark_int5_plain_mm --M=4 --N=64 --K=256 --gs=32 + --warmup=1 --iters=1 --seeds=3 +) + foreach(test_name ${CUDA_KERNEL_TESTS}) add_executable(${test_name} ${test_name}.cpp) diff --git a/backends/cuda/runtime/shims/tests/benchmark_int5_plain_mm.cu b/backends/cuda/runtime/shims/tests/benchmark_int5_plain_mm.cu new file mode 100644 index 00000000000..cec8e28e762 --- /dev/null +++ b/backends/cuda/runtime/shims/tests/benchmark_int5_plain_mm.cu @@ -0,0 +1,367 @@ +/* Copyright (c) Meta Platforms, Inc. and affiliates. */ + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuda_shims = executorch::backends::cuda; + +namespace { + +#define CUDA_CHECK(expr) \ + do { \ + cudaError_t err__ = (expr); \ + if (err__ != cudaSuccess) { \ + std::fprintf( \ + stderr, \ + "CUDA error %s:%d: %s\n", \ + __FILE__, \ + __LINE__, \ + cudaGetErrorString(err__)); \ + std::exit(1); \ + } \ + } while (0) + +int32_t log2_pow2_host(int32_t v) { + int32_t r = 0; + while (v > 1) { + v >>= 1; + ++r; + } + return r; +} + +uint16_t float_to_bf16(float x) { + uint32_t bits; + std::memcpy(&bits, &x, sizeof(bits)); + return static_cast(bits >> 16); +} + +void fill_case( + int64_t M, + int64_t N, + int64_t K, + int64_t gs, + uint32_t seed, + std::vector& A, + std::vector& ql, + std::vector& qh, + std::vector& scale, + std::vector& scale_step, + std::vector& zero, + std::vector& zero_step) { + std::mt19937 rng(seed); + std::uniform_real_distribution adist(-2.0f, 2.0f); + std::uniform_int_distribution qdist(0, 31); + std::uniform_int_distribution codedist(1, 15); + std::uniform_real_distribution sdist(0.003f, 0.035f); + std::uniform_real_distribution zdist(0.25f, 1.25f); + + A.resize(M * K); + ql.assign(N * (K / 2), 0); + qh.assign(N * (K / 8), 0); + scale.resize(N * (K / gs)); + scale_step.resize(N * (K / 256)); + zero.resize(N * (K / gs)); + zero_step.resize(N * (K / 256)); + + for (auto& x : A) { + x = float_to_bf16(adist(rng)); + } + for (auto& x : scale) { + x = static_cast(codedist(rng)); + } + for (auto& x : zero) { + x = static_cast(codedist(rng)); + } + for (auto& x : scale_step) { + x = __half_as_ushort(__float2half(sdist(rng))); + } + for (auto& x : zero_step) { + x = __half_as_ushort(__float2half(zdist(rng))); + } + + std::vector u(K); + for (int64_t n = 0; n < N; ++n) { + uint8_t* qlrow = ql.data() + n * (K / 2); + uint8_t* qhrow = qh.data() + n * (K / 8); + for (int64_t k = 0; k < K; ++k) { + u[k] = static_cast(qdist(rng)); + } + for (int64_t k = 0; k < K; k += 2) { + qlrow[k / 2] = static_cast((u[k] & 0xF) | ((u[k + 1] & 0xF) << 4)); + } + for (int64_t k = 0; k < K; k += 32) { + uint8_t* chunk = qhrow + k / 8; + for (int w = 0; w < 4; ++w) { + uint8_t hi_even = 0; + uint8_t hi_odd = 0; + for (int j = 0; j < 4; ++j) { + hi_even |= static_cast(((u[k + w * 8 + j * 2] >> 4) & 0x1) << j); + hi_odd |= static_cast(((u[k + w * 8 + j * 2 + 1] >> 4) & 0x1) << j); + } + chunk[w] = static_cast(hi_even | (hi_odd << 4)); + } + } + } +} + +template +float time_ms(Fn fn, int warmup, int iterations) { + for (int i = 0; i < warmup; ++i) { + fn(); + } + CUDA_CHECK(cudaDeviceSynchronize()); + cudaEvent_t start, stop; + CUDA_CHECK(cudaEventCreate(&start)); + CUDA_CHECK(cudaEventCreate(&stop)); + CUDA_CHECK(cudaEventRecord(start)); + for (int i = 0; i < iterations; ++i) { + fn(); + } + CUDA_CHECK(cudaEventRecord(stop)); + CUDA_CHECK(cudaEventSynchronize(stop)); + float elapsed = 0.0f; + CUDA_CHECK(cudaEventElapsedTime(&elapsed, start, stop)); + CUDA_CHECK(cudaEventDestroy(start)); + CUDA_CHECK(cudaEventDestroy(stop)); + return elapsed / iterations; +} + +void* device_alloc_copy(const void* src, size_t bytes) { + void* dst = nullptr; + CUDA_CHECK(cudaMalloc(&dst, bytes)); + CUDA_CHECK(cudaMemcpy(dst, src, bytes, cudaMemcpyHostToDevice)); + return dst; +} + +int64_t arg_value(int argc, char** argv, const char* name, int64_t fallback) { + std::string key = std::string("--") + name + "="; + for (int i = 1; i < argc; ++i) { + std::string arg(argv[i]); + if (arg.rfind(key, 0) == 0) { + return std::stoll(arg.substr(key.size())); + } + } + return fallback; +} + +} // namespace + +int main(int argc, char** argv) { + const int64_t M = arg_value(argc, argv, "M", 1); + const int64_t K = arg_value(argc, argv, "K", 6656); + const int64_t N = arg_value(argc, argv, "N", 202048); + const int64_t gs = arg_value(argc, argv, "gs", 32); + const int64_t warmup = arg_value(argc, argv, "warmup", 30); + const int64_t iterations = arg_value(argc, argv, "iters", 200); + const int64_t seeds = arg_value(argc, argv, "seeds", 3); + + if (M < 1 || M > 4 || K % 256 != 0 || K % 32 != 0 || gs < 32 || + (gs & (gs - 1))) { + std::fprintf(stderr, "unsupported shape M=%lld N=%lld K=%lld gs=%lld\n", M, N, K, gs); + return 2; + } + + CUDA_CHECK(cudaSetDevice(0)); + const int32_t gs_shift = log2_pow2_host(static_cast(gs)); + const int32_t n_groups = static_cast(K / gs); + const int32_t n_super = static_cast(K / 256); + const int32_t n_q8_blocks = static_cast(K / cuda_shims::Q8_BLOCK_SIZE_I5); + dim3 q8_grid((n_q8_blocks + cuda_shims::MV5_NWARPS - 1) / cuda_shims::MV5_NWARPS, M); + dim3 block(cuda_shims::MV5_WARP_SIZE, cuda_shims::MV5_NWARPS); + dim3 grid((N + cuda_shims::MV5_NWARPS - 1) / cuda_shims::MV5_NWARPS); + + std::printf( + "shape M=%lld N=%lld K=%lld gs=%lld q8_blocks=%d warmup=%lld iters=%lld seeds=%lld\n", + M, + N, + K, + gs, + n_q8_blocks, + warmup, + iterations, + seeds); + + double ref_total = 0.0; + double cand_total = 0.0; + for (int64_t seed_idx = 0; seed_idx < seeds; ++seed_idx) { + std::vector hA; + std::vector hql; + std::vector hqh; + std::vector hscale; + std::vector hscale_step; + std::vector hzero; + std::vector hzero_step; + fill_case( + M, + N, + K, + gs, + 3001 + seed_idx * 17, + hA, + hql, + hqh, + hscale, + hscale_step, + hzero, + hzero_step); + + auto* dA = static_cast<__nv_bfloat16*>(device_alloc_copy(hA.data(), hA.size() * sizeof(uint16_t))); + auto* dql = static_cast(device_alloc_copy(hql.data(), hql.size() * sizeof(uint8_t))); + auto* dqh = static_cast(device_alloc_copy(hqh.data(), hqh.size() * sizeof(uint8_t))); + auto* dscale = static_cast(device_alloc_copy(hscale.data(), hscale.size() * sizeof(uint8_t))); + auto* dscale_step = static_cast<__half*>(device_alloc_copy(hscale_step.data(), hscale_step.size() * sizeof(uint16_t))); + auto* dzero = static_cast(device_alloc_copy(hzero.data(), hzero.size() * sizeof(uint8_t))); + auto* dzero_step = static_cast<__half*>(device_alloc_copy(hzero_step.data(), hzero_step.size() * sizeof(uint16_t))); + cuda_shims::Q8Block_i5* dq8 = nullptr; + __nv_bfloat16* dref = nullptr; + __nv_bfloat16* dcand = nullptr; + CUDA_CHECK(cudaMalloc(&dq8, M * n_q8_blocks * sizeof(cuda_shims::Q8Block_i5))); + CUDA_CHECK(cudaMalloc(&dref, M * N * sizeof(__nv_bfloat16))); + CUDA_CHECK(cudaMalloc(&dcand, M * N * sizeof(__nv_bfloat16))); + + cuda_shims::quantize_activations_q8_i5_kernel<<>>(dA, dq8, static_cast(K)); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + + cuda_shims::int5_w5a8_matvec_kernel<<>>( + dql, + dqh, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dref, + static_cast(N), + static_cast(K), + static_cast(M), + gs_shift, + n_groups, + n_super); + if (M == 1 && gs == 32) { + cuda_shims::int5_w5a8_matvec_m1_gs32_kernel<<>>( + dql, + dqh, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dcand, + static_cast(N), + static_cast(K), + n_groups, + n_super); + } else if (M == 4 && gs == 32) { + cuda_shims::int5_w5a8_matvec_m4_gs32_kernel<<>>( + dql, + dqh, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dcand, + static_cast(N), + static_cast(K), + n_groups, + n_super); + } else { + cuda_shims::int5_w5a8_matvec_kernel<<>>( + dql, + dqh, + dscale, + dscale_step, + dzero, + dzero_step, + dq8, + dcand, + static_cast(N), + static_cast(K), + static_cast(M), + gs_shift, + n_groups, + n_super); + } + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + + std::vector href(M * N); + std::vector hcand(M * N); + CUDA_CHECK(cudaMemcpy(href.data(), dref, href.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaMemcpy(hcand.data(), dcand, hcand.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost)); + size_t mismatches = 0; + for (size_t i = 0; i < href.size(); ++i) { + if (href[i] != hcand[i]) { + if (mismatches < 8) { + std::printf("mismatch seed=%lld idx=%zu ref=0x%04x cand=0x%04x\n", seed_idx, i, href[i], hcand[i]); + } + ++mismatches; + } + } + if (mismatches != 0) { + std::printf("bitwise=FAIL mismatches=%zu/%zu\n", mismatches, href.size()); + return 1; + } + + float q8_ms = time_ms([&]() { + cuda_shims::quantize_activations_q8_i5_kernel<<>>(dA, dq8, static_cast(K)); + }, warmup, iterations); + float ref_ms = time_ms([&]() { + cuda_shims::int5_w5a8_matvec_kernel<<>>( + dql, dqh, dscale, dscale_step, dzero, dzero_step, dq8, dref, static_cast(N), static_cast(K), static_cast(M), gs_shift, n_groups, n_super); + }, warmup, iterations); + float cand_ms = time_ms([&]() { + if (M == 1 && gs == 32) { + cuda_shims::int5_w5a8_matvec_m1_gs32_kernel<<>>( + dql, dqh, dscale, dscale_step, dzero, dzero_step, dq8, dcand, static_cast(N), static_cast(K), n_groups, n_super); + } else if (M == 4 && gs == 32) { + cuda_shims::int5_w5a8_matvec_m4_gs32_kernel<<>>( + dql, dqh, dscale, dscale_step, dzero, dzero_step, dq8, dcand, static_cast(N), static_cast(K), n_groups, n_super); + } else { + cuda_shims::int5_w5a8_matvec_kernel<<>>( + dql, dqh, dscale, dscale_step, dzero, dzero_step, dq8, dcand, static_cast(N), static_cast(K), static_cast(M), gs_shift, n_groups, n_super); + } + }, warmup, iterations); + ref_total += ref_ms; + cand_total += cand_ms; + std::printf( + "seed=%lld bitwise=OK q8_ms=%.6f ref_matvec_ms=%.6f cand_matvec_ms=%.6f speedup=%.3fx\n", + seed_idx, + q8_ms, + ref_ms, + cand_ms, + ref_ms / cand_ms); + + CUDA_CHECK(cudaFree(dA)); + CUDA_CHECK(cudaFree(dql)); + CUDA_CHECK(cudaFree(dqh)); + CUDA_CHECK(cudaFree(dscale)); + CUDA_CHECK(cudaFree(dscale_step)); + CUDA_CHECK(cudaFree(dzero)); + CUDA_CHECK(cudaFree(dzero_step)); + CUDA_CHECK(cudaFree(dq8)); + CUDA_CHECK(cudaFree(dref)); + CUDA_CHECK(cudaFree(dcand)); + } + + std::printf( + "avg ref_matvec_ms=%.6f cand_matvec_ms=%.6f speedup=%.3fx\n", + ref_total / seeds, + cand_total / seeds, + ref_total / cand_total); + return 0; +}