From c5a07ba23a602afd34c7ceb5060e0fe9c5f8e97d Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 29 Jul 2026 03:27:30 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/runtime/shims/int4_plain_mm.cuh | 323 +++++++++++----- .../cuda/runtime/shims/tests/CMakeLists.txt | 28 ++ .../shims/tests/benchmark_int4_plain_mm.cu | 362 ++++++++++++++++++ 3 files changed, 619 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..cd9abb54cad 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -74,6 +74,34 @@ 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; +}