-
Notifications
You must be signed in to change notification settings - Fork 21.9k
CUDA: Improve NVFP4 W4A4 activation quantization #25730
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.
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
Changes from all commits
46caff6
8a6ca53
b1ea605
598f69b
1394d6e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, | ||
|
|
@@ -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; | ||
|
|
||
| #pragma unroll | ||
| for (int j0 = 0; j0 < J; j0 += nwarps) { | ||
| const int j = j0 + threadIdx.y; | ||
|
|
@@ -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 | ||
|
|
@@ -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 | ||
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. is something like
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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)
|
||
| } else { | ||
| dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; | ||
| GGML_UNUSED(y_scale_used); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
@@ -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) { | ||
|
|
||
|
|
@@ -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); | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -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, | ||
|
|
@@ -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]; | ||
|
|
@@ -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; | ||
|
|
@@ -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; | ||
|
|
@@ -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; | ||
| } | ||
|
|
@@ -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]; | ||
|
|
@@ -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; | ||
|
|
@@ -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; | ||
|
|
@@ -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; | ||
|
|
@@ -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]; | ||
|
|
@@ -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; | ||
|
|
@@ -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; | ||
|
|
@@ -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); | ||
| } | ||
|
|
||
|
|
@@ -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; | ||
|
|
@@ -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, | ||
|
|
@@ -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, | ||
|
|
||

Uh oh!
There was an error while loading. Please reload this page.