Skip to content

perf(ggml): reuse ROCmFP4 weights across q=4 columns #534

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
davide221 wants to merge 1 commit into
base: main
Choose a base branch
from
+113 −6
Open

perf(ggml): reuse ROCmFP4 weights across q=4 columns #534

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
119 changes: 113 additions & 6 deletions server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -501,8 +501,40 @@ static bool is_gfx1151(const int cc) {
return cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151;
}

// Dense four-column projections use the same weight row for every column.
// Decode each packed ROCmFP4 fragment once while preserving the original
// per-column DP4A and floating-point accumulation order.
static __device__ __forceinline__ void vec_dot_rocmfp4_fast_q8_1_4cols(
const void * __restrict__ vbq, const block_q8_1 * __restrict__ y,
const uint32_t stride_col_y, const int kby, const int kbx, const int iqs,
float (&result)[4]) {
const block_rocmfp4_fast * bq4 = (const block_rocmfp4_fast *) vbq + kbx;

int2 weights[VDR_ROCMFP4_FAST_Q8_1_MMVQ];
#pragma unroll
for (int l = 0; l < VDR_ROCMFP4_FAST_Q8_1_MMVQ; ++l) {
const int packed = rocmfp4_get_qs_i32(bq4->qs, iqs + l);
weights[l] = rocmfp4_get_int_from_codebook_16(packed, kvalues_rocmfp4);
}

const float weight_scale = rocmfp4_ue4m3_to_fp32_half_finite(bq4->e);
int sums[4] = {};
#pragma unroll
for (int j = 0; j < 4; ++j) {
const block_q8_1 * bq8 = &y[j*stride_col_y + kby];
const int * q8 = (const int *) bq8->qs + iqs;
#pragma unroll
for (int l = 0; l < VDR_ROCMFP4_FAST_Q8_1_MMVQ; ++l) {
sums[j] = ggml_cuda_dp4a(weights[l].x, q8[l + 0], sums[j]);
sums[j] = ggml_cuda_dp4a(weights[l].y, q8[l + 4], sums[j]);
}
result[j] = __low2float(bq8->ds) * weight_scale * sums[j];
}
}

template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false,
int fixed_ncols_x = 0, bool unroll_k_loop_2 = false>
int fixed_ncols_x = 0, bool unroll_k_loop_2 = false,
bool reuse_rocmfp4_weights = false>
__launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mul_mat_vec_q(
const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst,
Expand Down Expand Up @@ -533,6 +565,10 @@ static __global__ void mul_mat_vec_q(
"fixed-K decode specialization is limited to profiled ROCmFP2/FP3 types");
static_assert(!unroll_k_loop_2 || type == GGML_TYPE_Q4_0_ROCMFP4_FAST,
"partial K-loop unrolling is limited to the profiled FP4FAST path");
static_assert(!reuse_rocmfp4_weights ||
(type == GGML_TYPE_Q4_0_ROCMFP4_FAST &&
ncols_dst == 4 && fixed_ncols_x == 0 && !unroll_k_loop_2),
"weight reuse is limited to the four-column FP4FAST path");

const uint32_t channel_dst = blockIdx.y;

Expand Down Expand Up @@ -613,7 +649,38 @@ static __global__ void mul_mat_vec_q(

const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y;

if constexpr (fixed_ncols_x > 0) {
if constexpr (reuse_rocmfp4_weights) {
const int kbx_offset = sample_x*stride_sample_x +
channel_x*stride_channel_x + row0*stride_row_x;
for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) {
const int kby = kbx * (qk/QK8_1);
const int kqs = vdr * (tid % (qi/vdr));

#pragma unroll
for (int i = 0; i < rows_per_cuda_block; ++i) {
float dots[4];
vec_dot_rocmfp4_fast_q8_1_4cols(
vx, y, stride_col_y, kby,
kbx_offset + i*stride_row_x + kbx, kqs, dots);
#pragma unroll
for (int j = 0; j < 4; ++j) {
tmp[j][i] += dots[j];
}
if constexpr (has_fusion) {
if (use_gate) {
float gate_dots[4];
vec_dot_rocmfp4_fast_q8_1_4cols(
vgate, y, stride_col_y, kby,
kbx_offset + i*stride_row_x + kbx, kqs, gate_dots);
#pragma unroll
for (int j = 0; j < 4; ++j) {
tmp_gate[j][i] += gate_dots[j];
}
}
}
}
}
} else if constexpr (fixed_ncols_x > 0) {
static_assert(fixed_ncols_x % qk == 0, "fixed K must contain whole quant blocks");
constexpr int fixed_blocks = fixed_ncols_x/qk;
constexpr int fixed_iters = (fixed_blocks + blocks_per_iter - 1) / blocks_per_iter;
Expand Down Expand Up @@ -1204,7 +1271,8 @@ static std::pair<dim3, dim3> calc_launch_params(
}

template<ggml_type type, int c_ncols_dst, bool small_k = false,
int fixed_ncols_x = 0, bool unroll_k_loop_2 = false>
int fixed_ncols_x = 0, bool unroll_k_loop_2 = false,
bool reuse_rocmfp4_weights = false>
static void mul_mat_vec_q_switch_fusion(
const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y,
Expand All @@ -1216,14 +1284,18 @@ static void mul_mat_vec_q_switch_fusion(

const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
if (has_fusion) {
mul_mat_vec_q<type, c_ncols_dst, true, small_k, fixed_ncols_x, unroll_k_loop_2><<<block_nums, block_dims, nbytes_shared, stream>>>
mul_mat_vec_q<type, c_ncols_dst, true, small_k, fixed_ncols_x,
unroll_k_loop_2, reuse_rocmfp4_weights>
<<<block_nums, block_dims, nbytes_shared, stream>>>
(vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, ids_tokenwise_samples);
return;
}

mul_mat_vec_q<type, c_ncols_dst, false, small_k, fixed_ncols_x, unroll_k_loop_2><<<block_nums, block_dims, nbytes_shared, stream>>>
mul_mat_vec_q<type, c_ncols_dst, false, small_k, fixed_ncols_x,
unroll_k_loop_2, reuse_rocmfp4_weights>
<<<block_nums, block_dims, nbytes_shared, stream>>>
(vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, ids_tokenwise_samples);
Expand Down Expand Up @@ -1280,6 +1352,31 @@ static void mul_mat_vec_rocmfp4_unroll2_launch(
ids_stride, stream);
}

static void mul_mat_vec_rocmfp4_4col_reuse_launch(
const void * vx, const void * vy,
const ggml_cuda_mm_fusion_args_device fusion, float * dst,
const uint32_t ncols_x, const uint32_t nrows_x, const uint3 nchannels_y,
const uint32_t nchannels_dst, const uint32_t stride_row_x,
const uint32_t stride_col_y, const uint32_t stride_col_dst,
const uint3 channel_ratio, const uint32_t stride_channel_x,
const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
const uint32_t nsamples_dst, const uint3 sample_ratio,
const uint32_t stride_sample_x, const uint32_t stride_sample_y,
const uint32_t stride_sample_dst, const uint32_t ids_stride,
const int warp_size, const mmvq_parameter_table_id table_id,
cudaStream_t stream) {
constexpr int ncols_dst = 4;
const auto dims = calc_launch_params<GGML_TYPE_Q4_0_ROCMFP4_FAST>(
ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id);
mul_mat_vec_q_switch_fusion<
GGML_TYPE_Q4_0_ROCMFP4_FAST, ncols_dst, false, 0, false, true>(
vx, vy, nullptr, fusion, dst, ncols_x, nchannels_y, stride_row_x,
stride_col_y, stride_col_dst, channel_ratio, stride_channel_x,
stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x,
stride_sample_y, stride_sample_dst, dims.first, dims.second, 0,
ids_stride, stream);
}

template <ggml_type type>
static void mul_mat_vec_q_moe_launch(
const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
Expand Down Expand Up @@ -1450,10 +1547,20 @@ static void mul_mat_vec_q_switch_ncols_dst(
return;
}

// Decode-only gfx1151 specializations. Each one preserves the original
// Decode-only gfx1151 specializations. Each preserves the original
// per-lane K traversal and accumulation order, so exact validated shapes
// can use the faster kernel without a serving-time tuning flag.
if constexpr (type == GGML_TYPE_Q4_0_ROCMFP4_FAST) {
if (is_gfx1151(cc) && !has_ids && ncols_dst == 4) {
mul_mat_vec_rocmfp4_4col_reuse_launch(
vx, vy, fusion, dst, ncols_x, nrows_x, nchannels_y_fd,
nchannels_dst, stride_row_x, stride_col_y, stride_col_dst,
channel_ratio_fd, stride_channel_x, stride_channel_y,
stride_channel_dst, nsamples_dst, sample_ratio_fd,
stride_sample_x, stride_sample_y, stride_sample_dst,
ids_stride, warp_size, table_id, stream);
return;
}
if (is_gfx1151(cc) && ncols_dst == 1) {
mul_mat_vec_rocmfp4_unroll2_launch(
vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd,
Expand Down