Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
9 changes: 9 additions & 0 deletions ggml/src/ggml-cuda/common.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -362,6 +362,15 @@ static bool blackwell_mma_available(const int cc) {
ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN;
}

// Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) {
GGML_ASSERT(tensor != nullptr);
return (reinterpret_cast<uintptr_t>(tensor->data) % alignment) == 0 &&
tensor->nb[1] % alignment == 0 &&
tensor->nb[2] % alignment == 0 &&
tensor->nb[3] % alignment == 0;
}

static constexpr __device__ int ggml_cuda_get_physical_warp_size() {
#if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__))
return 64;
Expand Down
33 changes: 23 additions & 10 deletions ggml/src/ggml-cuda/mmq.cu
Original file line number Diff line number Diff line change
Expand Up @@ -130,14 +130,20 @@ void ggml_cuda_mul_mat_q(
const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block +
ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
src1_scale.alloc(ne13*ne12*ne11);
}

{
const int64_t s11 = src1->nb[1] / ts_src1;
const int64_t s12 = src1->nb[2] / ts_src1;
const int64_t s13 = src1->nb[3] / ts_src1;
if (use_native_fp4) {
static constexpr size_t align_float8 = 32;
const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1));
quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded,
ne11, ne12, ne13, stream);

} else {
Expand All @@ -155,6 +161,7 @@ void ggml_cuda_mul_mat_q(

const mmq_args args = {
src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
ne00, ne01, ne1, s01, ne11, s1,
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
Expand Down Expand Up @@ -192,6 +199,10 @@ void ggml_cuda_mul_mat_q(
const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block +
ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
src1_scale.alloc(ne12*n_expert_used);
}

const int64_t ne11_flat = ne12*n_expert_used;
const int64_t ne12_flat = 1;
Expand All @@ -202,18 +213,19 @@ void ggml_cuda_mul_mat_q(
const int64_t s12 = src1->nb[2] / ts_src1;
const int64_t s13 = src1->nb[3] / ts_src1;

if (dedup_bcast) {
// quantize each token once, scatter its block to all n_expert_used slots
if (use_native_fp4) {
quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
if (use_native_fp4) {
static constexpr size_t align_float8 = 32;
const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
if (dedup_bcast) {
quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10,
/*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
} else {
quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
/*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13,
ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
}
} else if (use_native_fp4) {
quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
} else if (dedup_bcast) {
quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
/*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
} else {
quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
Expand All @@ -229,6 +241,7 @@ void ggml_cuda_mul_mat_q(
// Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
const mmq_args args = {
src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
src1_scale.ptr,
ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
ne02, ne02, s02, s12, s2,
ne03, ne13, s03, s13, s3,
Expand Down
107 changes: 89 additions & 18 deletions ggml/src/ggml-cuda/mmq.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
float * __restrict__ dst, const int stride, const int i_max, const int j_max);
float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max);

enum mmq_q8_1_ds_layout {
MMQ_Q8_1_DS_LAYOUT_D4,
Expand Down Expand Up @@ -413,11 +413,13 @@ static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config

template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
const int stride, const int i_max, const int j_max) {
const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);

const bool y_scale_used = y_scale != nullptr;
Comment thread
ORippler marked this conversation as resolved.

#pragma unroll
for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
Expand All @@ -434,15 +436,25 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
continue;
}

dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
if constexpr (type == GGML_TYPE_NVFP4) {
if (y_scale_used) {
dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
} else {
dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
}
} else {
dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
GGML_UNUSED(y_scale_used);
}
}
}
}

template<ggml_type type, int J, bool fallback>
static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
const int stride, const int i_max, const int j_max) {
const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {

#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
#else
Expand All @@ -457,6 +469,8 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(

const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I);

const bool y_scale_used = y_scale != nullptr;

#pragma unroll
for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
#pragma unroll
Expand All @@ -475,7 +489,16 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
continue;
}

dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
if constexpr (type == GGML_TYPE_NVFP4) {
if (y_scale_used) {
dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
} else {
dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
}
Comment on lines +493 to +497

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is something like (y_scale_used ? y_scale[j] : 1.0) going to be useful to avoid the branch?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generally yes, in this instance no. On BW, I have seen the compiler fail to do LDCU on incoming pointers if they are only ever conditionally accessed in nested loops. You can see in the attached screen that for each instance of j <= j_max, there is a predicated load of LDC from the same address. This is the load of y_scale's base pointer, which can be hoisted outside the loop. Called the OC commit before rebasing "compiler massaging" for a reason 😄 Will follow-up on the performance bug I filed with yet-another repo (being this PR)

image

} else {
dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
GGML_UNUSED(y_scale_used);
}
}
}
}
Expand Down Expand Up @@ -819,6 +842,7 @@ template <ggml_type type, int J, bool fallback, bool fixup>
static __device__ __forceinline__ void mul_mat_q_process_tile(
const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
const float * __restrict__ y_scale,
const int stride_row_x, const int ncols_y, const int stride_col_dst,
const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {

Expand Down Expand Up @@ -884,9 +908,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
}

if (fixup) {
write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), I, I, J);
write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
} else {
write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j);
write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
}
}

Expand All @@ -898,6 +922,7 @@ __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_g
static __global__ void mul_mat_q(
const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
const float * __restrict__ y_scale,
const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
Expand Down Expand Up @@ -943,8 +968,14 @@ static __global__ void mul_mat_q(
int col_low = 0;
int col_high = ncols_dst;
int col_diff = ncols_dst;
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
int offset_y_scale;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
} else {
GGML_UNUSED(offset_y_scale);
}

if (ids_dst) {
col_low = expert_bounds[zt + 0];
Expand All @@ -953,6 +984,9 @@ static __global__ void mul_mat_q(

offset_y = 0;
offset_dst = 0;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale = 0;
}

if (jt*J >= col_diff) {
return;
Expand All @@ -974,6 +1008,11 @@ static __global__ void mul_mat_q(

offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
offset_dst += it*I;
const float * y_scale_tile = nullptr;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale += col_low + jt*J;
y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
}

const int tile_x_max_i = nrows_x - it*I - 1;
const int tile_y_max_j = col_diff - jt*J - 1;
Expand All @@ -982,7 +1021,8 @@ static __global__ void mul_mat_q(

constexpr bool fixup = false;
mul_mat_q_process_tile<type, J, fallback, fixup>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
return;
}
Expand Down Expand Up @@ -1016,8 +1056,14 @@ static __global__ void mul_mat_q(
int col_low = 0;
int col_high = ncols_dst;
int col_diff = ncols_dst;
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
int offset_y_scale;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
} else {
GGML_UNUSED(offset_y_scale);
}

if (ids_dst) {
col_low = expert_bounds[zt + 0];
Expand All @@ -1026,6 +1072,9 @@ static __global__ void mul_mat_q(

offset_y = 0;
offset_dst = 0;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale = 0;
}

if (jt*J >= col_diff) {
kbc += blocks_per_ne00.z;
Expand Down Expand Up @@ -1053,6 +1102,11 @@ static __global__ void mul_mat_q(

offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
offset_dst += it*I;
const float * y_scale_tile = nullptr;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale += col_low + jt * J;
y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
}

const int tile_x_max_i = nrows_x - it*I - 1;
const int tile_y_max_j = col_diff - jt*J - 1;
Expand All @@ -1061,7 +1115,8 @@ static __global__ void mul_mat_q(

constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
mul_mat_q_process_tile<type, J, fallback, fixup>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);

kbc += blocks_per_ne00.z;
Expand Down Expand Up @@ -1090,8 +1145,14 @@ static __global__ void mul_mat_q(
int col_low = 0;
int col_high = ncols_dst;
int col_diff = ncols_dst;
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
int offset_y_scale;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
} else {
GGML_UNUSED(offset_y_scale);
}

if (ids_dst) {
col_low = expert_bounds[zt + 0];
Expand All @@ -1100,6 +1161,9 @@ static __global__ void mul_mat_q(

offset_y = 0;
offset_dst = 0;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale = 0;
}

if (jt*J >= col_diff) {
return;
Expand All @@ -1122,6 +1186,11 @@ static __global__ void mul_mat_q(

offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
offset_dst += it*I;
const float * y_scale_tile = nullptr;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale += col_low + jt * J;
y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
}

const int tile_x_max_i = nrows_x - it*I - 1;
const int tile_y_max_j = col_diff - jt*J - 1;
Expand All @@ -1130,7 +1199,8 @@ static __global__ void mul_mat_q(

constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
mul_mat_q_process_tile<type, J, fallback, fixup>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
}

Expand Down Expand Up @@ -1274,6 +1344,7 @@ static __global__ void mul_mat_q_stream_k_fixup(

struct mmq_args {
const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst;
const float * y_scale;
int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
Expand Down Expand Up @@ -1323,7 +1394,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a

if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) {
mul_mat_q<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr,
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
Expand Down Expand Up @@ -1352,7 +1423,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);

mul_mat_q<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr,
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
Expand Down
Loading
Loading