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 044a86dfe28..678c0eacc28 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -130,6 +130,34 @@ add_test( COMMAND benchmark_int6_plain_mm --M=4 --N=64 --K=256 --gs=16 --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; +}