From 41b30bfccb64b93733cb066c5512f9182d315f33 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 29 Jul 2026 03:27:34 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/runtime/shims/int6_plain_mm.cuh | 461 ++++++++++++++--- .../cuda/runtime/shims/tests/CMakeLists.txt | 28 + .../shims/tests/benchmark_int6_plain_mm.cu | 486 ++++++++++++++++++ 3 files changed, 910 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 cd9abb54cad..044a86dfe28 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -102,6 +102,34 @@ add_test( COMMAND benchmark_int4_plain_mm --M=4 --N=64 --K=256 --gs=32 --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; +}