Skip to content

[executorch][cuda] Optimize short-query INT4 matvec kernels #21456

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Gasoonjia wants to merge 1 commit into
base: gh/gasoonjia/175/base
Choose a base branch
from
+619 −94
Open

[executorch][cuda] Optimize short-query INT4 matvec kernels #21456

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
323 changes: 229 additions & 94 deletions backends/cuda/runtime/shims/int4_plain_mm.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <int32_t ROWS>
__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<int64_t>(n) * K_half;
const uint8_t* scale_row = w_scale_t + static_cast<int64_t>(n) * n_groups;
const __half* scale_step_row =
w_scale_step + static_cast<int64_t>(n) * n_super;
const uint8_t* zero_row = w_zero_t + static_cast<int64_t>(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<int64_t>(n) * n_super;
const Q8Block* q8_row = q8 + static_cast<int64_t>(m) * n_q8_blocks;

const uint4* qrow16 = reinterpret_cast<const uint4*>(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<uint32_t>(s_bits) |
(static_cast<uint32_t>(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<unsigned short>(steps_packed & 0xFFFF)));
float zero_point_step = __half2float(
const float zero_point_step = __half2float(
__ushort_as_half(static_cast<unsigned short>(steps_packed >> 16)));
// Effective per-group scale/zero (one coalesced code byte each per group).
float ws = static_cast<float>(__ldg(&scale_row[g])) * scale_step;
float wz = static_cast<float>(__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<const uint4*>(qb->qs_even);
uint4 ao = *reinterpret_cast<const uint4*>(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<float>(__ldg(&scale_row[g])) * scale_step;
const float wz = static_cast<float>(__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<int64_t>(row) * n_q8_blocks + i_safe;
activations_even[row] =
*reinterpret_cast<const uint4*>(qb->qs_even);
activations_odd[row] =
*reinterpret_cast<const uint4*>(qb->qs_odd);
activation_scales[row] = qb->d;
}

int32_t dp = __dp4a(vi_lo, static_cast<int32_t>(a_even[w]), 0);
dp = __dp4a(vi_hi, static_cast<int32_t>(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<int32_t>(a_even[w]), 0);
a_sum8 = __dp4a(0x01010101, static_cast<int32_t>(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<int32_t>(a_even), 0);
dp = __dp4a(vi_hi, static_cast<int32_t>(a_odd), dp);
int32_t a_sum8 =
__dp4a(0x01010101, static_cast<int32_t>(a_even), 0);
a_sum8 =
__dp4a(0x01010101, static_cast<int32_t>(a_odd), a_sum8);
sums[row] += ws * activation_scales[row] *
(static_cast<float>(dp) - wz * static_cast<float>(a_sum8));
}
}
}

sum += ws * a_scale *
(static_cast<float>(dp) - wz * static_cast<float>(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<int64_t>(row) * N + n] =
__float2bfloat16(sums[row]);
}
}
}

if (lane_id == 0)
out[static_cast<int64_t>(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<int64_t>(m) * (K / Q8_BLOCK_SIZE),
out + static_cast<int64_t>(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<ROWS>( \
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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -383,6 +460,60 @@ void _int4_plain_mm_cuda(

int32_t n_groups = static_cast<int32_t>(scale.size(1));
int32_t n_super = static_cast<int32_t>(scale_step.size(1));
if (M == 4) {
dim3 m4_grid((N + MV_NWARPS - 1) / MV_NWARPS);
int4_w4a8_matvec_m4_coalesced_kernel<<<m4_grid, block, 0, stream>>>(
reinterpret_cast<const uint8_t*>(qdata.data_ptr()),
reinterpret_cast<const uint8_t*>(scale.data_ptr()),
reinterpret_cast<const __half*>(scale_step.data_ptr()),
reinterpret_cast<const uint8_t*>(zero.data_ptr()),
reinterpret_cast<const __half*>(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<<<m3_grid, block, 0, stream>>>(
reinterpret_cast<const uint8_t*>(qdata.data_ptr()),
reinterpret_cast<const uint8_t*>(scale.data_ptr()),
reinterpret_cast<const __half*>(scale_step.data_ptr()),
reinterpret_cast<const uint8_t*>(zero.data_ptr()),
reinterpret_cast<const __half*>(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<<<m2_grid, block, 0, stream>>>(
reinterpret_cast<const uint8_t*>(qdata.data_ptr()),
reinterpret_cast<const uint8_t*>(scale.data_ptr()),
reinterpret_cast<const __half*>(scale_step.data_ptr()),
reinterpret_cast<const uint8_t*>(zero.data_ptr()),
reinterpret_cast<const __half*>(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<<<grid, block, 0, stream>>>(
reinterpret_cast<const uint8_t*>(qdata.data_ptr()),
reinterpret_cast<const uint8_t*>(scale.data_ptr()),
Expand All @@ -391,7 +522,11 @@ void _int4_plain_mm_cuda(
reinterpret_cast<const __half*>(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
Loading
Loading