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;
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 {
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,
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;
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);
// 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,
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,
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;
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
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
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];
+ }
+ } else {
+ dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+ GGML_UNUSED(y_scale_used);
+ }
}
}
}
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) {
}
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);
}
}
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,
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];
offset_y = 0;
offset_dst = 0;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = 0;
+ }
if (jt*J >= col_diff) {
return;
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;
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;
}
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];
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;
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;
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;
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];
offset_y = 0;
offset_dst = 0;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = 0;
+ }
if (jt*J >= col_diff) {
return;
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;
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);
}
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;
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,
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,
#include "quantize.cuh"
#include <cstdint>
+#if defined(BLACKWELL_MMA_AVAILABLE)
+// this maps to 256-bit loads in PTX on supported devices,
+// and otherwise falls back to 2 128-bit loads
+struct __builtin_align__(32) float8 {
+ float x; float y; float z; float w;
+ float p; float q; float r; float s;
+};
+#endif
+
+#if CUDART_VERSION >= 12080
+static __device__ __forceinline__ float nvfp4_native_scale_error(
+ const float vals[QK_NVFP4_SUB], const float inv_col_scale, const float inv_scale, const float scale) {
+ const float scale_dequant = 2.0f * scale;
+ float err = 0.0f;
+
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB; k += 4) {
+ const float v0 = vals[k + 0] * inv_col_scale;
+ const float v1 = vals[k + 1] * inv_col_scale;
+ const float v2 = vals[k + 2] * inv_col_scale;
+ const float v3 = vals[k + 3] * inv_col_scale;
+
+ const __nv_fp4x4_e2m1 q(make_float4(v0 * inv_scale, v1 * inv_scale, v2 * inv_scale, v3 * inv_scale));
+ const __nv_fp4x4_storage_t q_storage = q.__x;
+ const __nv_fp4x2_storage_t q_lo = static_cast<__nv_fp4x2_storage_t>(q_storage);
+ const __nv_fp4x2_storage_t q_hi = static_cast<__nv_fp4x2_storage_t>(q_storage >> 8U);
+
+ const __half2_raw hraw2_lo = __nv_cvt_fp4x2_to_halfraw2(q_lo, __NV_E2M1);
+ const __half2_raw hraw2_hi = __nv_cvt_fp4x2_to_halfraw2(q_hi, __NV_E2M1);
+ const __half2 h2_lo = static_cast<__half2>(hraw2_lo);
+ const __half2 h2_hi = static_cast<__half2>(hraw2_hi);
+ const float2 dq_lo = __half22float2(h2_lo);
+ const float2 dq_hi = __half22float2(h2_hi);
+
+ const float err0 = fabsf(v0) - fabsf(dq_lo.x) * scale_dequant;
+ const float err1 = fabsf(v1) - fabsf(dq_lo.y) * scale_dequant;
+ const float err2 = fabsf(v2) - fabsf(dq_hi.x) * scale_dequant;
+ const float err3 = fabsf(v3) - fabsf(dq_hi.y) * scale_dequant;
+
+ err = fmaf(err0, err0, err);
+ err = fmaf(err1, err1, err);
+ err = fmaf(err2, err2, err);
+ err = fmaf(err3, err3, err);
+ }
+
+ return err;
+}
+#endif // CUDART_VERSION >= 12080
+
__launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
static __global__ void quantize_q8_1(
const float * x_ptr, void * vy_ptr,
return static_cast<uint8_t>(biased);
}
-
// scatter: grid over tokens, quantize once, write to all the token's compact rows
-template <bool scatter>
+template <bool scatter, bool use_aligned_float8>
static __global__ void quantize_mmq_nvfp4(
- const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+ const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, float * __restrict__ scale,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) {
#if defined(BLACKWELL_MMA_AVAILABLE)
- const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB;
- if (i0_base >= ne0) {
- return;
- }
-
- const int64_t k_block = i0_base / QK_FP4_MMQ;
const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ;
- if (k_block >= blocks_per_col) {
- return;
- }
- const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
int64_t base_idx;
if constexpr (scatter) {
base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
} else {
- const int64_t i2 = blockIdx.z % ne2;
- const int64_t i3 = blockIdx.z / ne2;
+ const int64_t i2 = blockIdx.y % ne2;
+ const int64_t i3 = blockIdx.y / ne2;
const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
base_idx = i3 * s03 + i2 * s02 + i01 * s01;
}
+ const float * __restrict__ x_row = x + base_idx;
+
+ float amax = 0.0f;
+ if constexpr (use_aligned_float8) {
+ for (int64_t i0 = 8 * threadIdx.x; i0 < ne00; i0 += 8 * blockDim.x) {
+ const float * x_base = x_row + i0;
+ const float8 v = reinterpret_cast<const float8 *>(x_base)[0];
+ amax = fmaxf(amax, fabsf(v.x));
+ amax = fmaxf(amax, fabsf(v.y));
+ amax = fmaxf(amax, fabsf(v.z));
+ amax = fmaxf(amax, fabsf(v.w));
+ amax = fmaxf(amax, fabsf(v.p));
+ amax = fmaxf(amax, fabsf(v.q));
+ amax = fmaxf(amax, fabsf(v.r));
+ amax = fmaxf(amax, fabsf(v.s));
+ }
+ } else {
+ for (int64_t i0 = threadIdx.x; i0 < ne00; i0 += blockDim.x) {
+ amax = fmaxf(amax, fabsf(x_row[i0]));
+ }
+ }
+
+ amax = warp_reduce_max<WARP_SIZE>(amax);
- float vals_raw[QK_NVFP4_SUB];
- float amax_raw = 0.0f;
+ __shared__ float warp_amax[CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE];
+ const int lane = threadIdx.x % WARP_SIZE;
+ const int warp = threadIdx.x / WARP_SIZE;
+
+ if (lane == 0) {
+ warp_amax[warp] = amax;
+ }
+ __syncthreads();
+
+ if (warp == 0) {
+ amax = threadIdx.x < int(CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE) ? warp_amax[lane] : 0.0f;
+ amax = warp_reduce_max<WARP_SIZE>(amax);
+ if (lane == 0) {
+ warp_amax[0] = amax / (6.0f * 448.0f);
+ if constexpr (scatter) {
#pragma unroll
- for (int k = 0; k < QK_NVFP4_SUB; k++) {
- const int64_t i00 = i0_base + k;
- if (i00 < ne00) {
- const float v = x[base_idx + i00];
- vals_raw[k] = v;
- amax_raw = fmaxf(amax_raw, fabsf(v));
- } else {
- vals_raw[k] = 0.0f;
+ for (int slot = 0; slot < n_expert_used; ++slot) {
+ const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+ scale[i] = warp_amax[0];
+ }
+ } else {
+ scale[blockIdx.y * ne1 + blockIdx.x] = warp_amax[0];
+ }
}
}
+ __syncthreads();
- static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2};
- const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_raw / 6.0f);
+ block_fp4_mmq * y = (block_fp4_mmq *) vy;
+ const int64_t n_subblocks = (ne0 + QK_NVFP4_SUB - 1) / QK_NVFP4_SUB;
+
+ for (int64_t isb = threadIdx.x; isb < n_subblocks; isb += blockDim.x) {
+ const int64_t i0_base = isb * QK_NVFP4_SUB;
+ const int64_t k_block = i0_base / QK_FP4_MMQ;
+ const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
+
+ const float row_scale = warp_amax[0];
+ const float inv_col_scale = row_scale > 0.0f ? 1.0f / row_scale : 0.0f;
+
+ float vals[QK_NVFP4_SUB];
+ if constexpr (use_aligned_float8) {
+ const float * x_base = x_row + i0_base;
+ const float8 v0 = i0_base + 7 < ne00 ? reinterpret_cast<const float8 *>(x_base)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+ const float8 v1 = i0_base + 15 < ne00 ? reinterpret_cast<const float8 *>(x_base + 8)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+ vals[0] = v0.x; vals[1] = v0.y; vals[2] = v0.z; vals[3] = v0.w;
+ vals[4] = v0.p; vals[5] = v0.q; vals[6] = v0.r; vals[7] = v0.s;
+ vals[8] = v1.x; vals[9] = v1.y; vals[10] = v1.z; vals[11] = v1.w;
+ vals[12] = v1.p; vals[13] = v1.q; vals[14] = v1.r; vals[15] = v1.s;
+ } else {
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+ const int64_t i00 = i0_base + k;
+ vals[k] = i00 < ne00 ? x_row[i00] : 0.0f;
+ }
+ }
- float best_err = FLT_MAX;
- uint8_t fp8_code = 0;
- float subblock_scale = 0.0f;
+ uint32_t q0 = 0;
+ uint32_t q1 = 0;
-#pragma unroll // Check +/- 2 to find best code to reduce NVFP4 activation loss. Negligible overhead on Blackwell.
- for (int i = 0; i < 5; i++) {
- const int test_code = first_fp8_code + test_offsets[i];
- if (test_code < 0 || test_code > 0x7e) {
- continue;
- }
- const uint8_t code = (uint8_t) test_code;
- const float test_scale = ggml_cuda_ue4m3_to_fp32(code);
- const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
- float cur_err = 0.0f;
+ float amax_sub = 0.0f;
#pragma unroll
for (int k = 0; k < QK_NVFP4_SUB; ++k) {
- const float v = vals_raw[k];
- const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
- const float err_diff = fabsf(v) - fabsf(kvalues_mxfp4[q & 0x7]) * test_scale;
- cur_err = fmaf(err_diff, err_diff, cur_err);
+ amax_sub = fmaxf(amax_sub, fabsf(vals[k] * inv_col_scale));
}
- if (cur_err < best_err) {
- best_err = cur_err;
- fp8_code = test_code;
- subblock_scale = test_scale;
+ static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2 };
+ const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_sub / 6.0f);
+
+ uint8_t fp8_code = (uint8_t) first_fp8_code;
+ float subblock_scale = ggml_cuda_ue4m3_to_fp32(fp8_code);
+ float inv_scale_err = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+ float best_err = nvfp4_native_scale_error(vals, inv_col_scale, inv_scale_err, subblock_scale);
+#else
+ float best_err = 0.0f;
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+ const float v = vals[k] * inv_col_scale;
+ const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, inv_scale_err);
+ const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * subblock_scale;
+ best_err = fmaf(err_diff, err_diff, best_err);
}
- }
+#endif // CUDART_VERSION >= 12080
- const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
- uint32_t q0 = 0;
- uint32_t q1 = 0;
-#pragma unroll // this is faster than the previous __nv_fp4x4_e2m1
- for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
- q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 0], inv_scale) << (8 * k);
- q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 8], inv_scale) << (8 * k + 4);
- q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 4], inv_scale) << (8 * k);
- q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
- }
+#pragma unroll
+ for (int i = 1; i < 5; ++i) {
+ const int test_code = first_fp8_code + test_offsets[i];
+ if (test_code < 0 || test_code > 0x7e) {
+ continue;
+ }
- block_fp4_mmq * y = (block_fp4_mmq *) vy;
- if constexpr (scatter) {
+ const float test_scale = ggml_cuda_ue4m3_to_fp32((uint8_t) test_code);
+ const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+ const float cur_err = nvfp4_native_scale_error(vals, inv_col_scale, test_inv_scale, test_scale);
+#else
+ float cur_err = 0.0f;
#pragma unroll
- for (int slot = 0; slot < n_expert_used; ++slot) {
- const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
- block_fp4_mmq * yb = y + (k_block * ne1 + i);
+ for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+ const float v = vals[k] * inv_col_scale;
+ const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
+ const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * test_scale;
+ cur_err = fmaf(err_diff, err_diff, cur_err);
+ }
+#endif // CUDART_VERSION >= 12080
+
+ if (cur_err < best_err) {
+ best_err = cur_err;
+ fp8_code = (uint8_t) test_code;
+ subblock_scale = test_scale;
+ }
+ }
+#if CUDART_VERSION >= 12080
+ const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+ const float s = inv_col_scale * inv_scale;
+
+ __nv_fp4x4_e2m1 q0_lo(make_float4(vals[0] * s, vals[8] * s, vals[1] * s, vals[9] * s));
+ __nv_fp4x4_e2m1 q0_hi(make_float4(vals[2] * s, vals[10] * s, vals[3] * s, vals[11] * s));
+ __nv_fp4x4_e2m1 q1_lo(make_float4(vals[4] * s, vals[12] * s, vals[5] * s, vals[13] * s));
+ __nv_fp4x4_e2m1 q1_hi(make_float4(vals[6] * s, vals[14] * s, vals[7] * s, vals[15] * s));
+
+ const char2 q0_lo_c = *reinterpret_cast<char2 *>(&q0_lo);
+ const char2 q0_hi_c = *reinterpret_cast<char2 *>(&q0_hi);
+ const char2 q1_lo_c = *reinterpret_cast<char2 *>(&q1_lo);
+ const char2 q1_hi_c = *reinterpret_cast<char2 *>(&q1_hi);
+
+ q0 = uint32_t(uint8_t(q0_lo_c.x)) | (uint32_t(uint8_t(q0_lo_c.y)) << 8) |
+ (uint32_t(uint8_t(q0_hi_c.x)) << 16) | (uint32_t(uint8_t(q0_hi_c.y)) << 24);
+ q1 = uint32_t(uint8_t(q1_lo_c.x)) | (uint32_t(uint8_t(q1_lo_c.y)) << 8) |
+ (uint32_t(uint8_t(q1_hi_c.x)) << 16) | (uint32_t(uint8_t(q1_hi_c.y)) << 24);
+#else
+ const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
+ q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 0] * inv_col_scale, inv_scale)) << (8 * k);
+ q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 8] * inv_col_scale, inv_scale)) << (8 * k + 4);
+ q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 4] * inv_col_scale, inv_scale)) << (8 * k);
+ q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 12] * inv_col_scale, inv_scale)) << (8 * k + 4);
+ }
+#endif // CUDART_VERSION >= 12080
+
+ if constexpr (scatter) {
+#pragma unroll
+ for (int slot = 0; slot < n_expert_used; ++slot) {
+ const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+ block_fp4_mmq * yb = y + (k_block * ne1 + i);
+ uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+ yqs[2 * sub + 0] = q0;
+ yqs[2 * sub + 1] = q1;
+ reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+ }
+ } else {
+ block_fp4_mmq * yb = y + (blockIdx.y * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
yqs[2 * sub + 0] = q0;
yqs[2 * sub + 1] = q1;
reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
}
- } else {
- block_fp4_mmq * yb = y + (blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
- uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
- yqs[2 * sub + 0] = q0;
- yqs[2 * sub + 1] = q1;
- reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
}
- GGML_UNUSED(n_expert_used);
#else
- GGML_UNUSED(n_expert_used);
+ GGML_UNUSED_VARS(x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, n_expert_used);
NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
#endif // defined(BLACKWELL_MMA_AVAILABLE)
// scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row)
void quantize_scatter_mmq_fp4_cuda(
- const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+ const float * x, const int32_t * ids_src1_inv, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
const int64_t ne00, const int64_t stride_token, const int64_t ne0,
const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
GGML_ASSERT(ne0 > 0);
if (type_src0 == GGML_TYPE_NVFP4) {
+ GGML_ASSERT(scale);
GGML_ASSERT(ne00 % QK_NVFP4 == 0);
- constexpr int nvfp4_block_size = 128;
- const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
- const dim3 block_size(nvfp4_block_size, 1, 1);
- const dim3 num_blocks(n_tokens, block_num_y, 1);
- quantize_mmq_nvfp4<true><<<num_blocks, block_size, 0, stream>>>(
- x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+ const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+ const dim3 num_blocks(n_tokens, 1, 1);
+ if (use_aligned_float8) {
+ quantize_mmq_nvfp4<true, true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+ } else {
+ quantize_mmq_nvfp4<true, false><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+ }
} else {
GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4);
constexpr int nwarps = 8;
}
void quantize_mmq_fp4_cuda(
- const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+ const float * x, const int32_t * ids, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4);
GGML_ASSERT(ne0 > 0);
if (type_src0 == GGML_TYPE_NVFP4) {
+ GGML_ASSERT(scale);
GGML_ASSERT(ne00 % QK_NVFP4 == 0);
- constexpr int nvfp4_block_size = 128;
- const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
- const dim3 block_size(nvfp4_block_size, 1, 1);
- const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
- quantize_mmq_nvfp4<false><<<num_blocks, block_size, 0, stream>>>(
- x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+ const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+ const dim3 num_blocks(ne1, ne2 * ne3, 1);
+ if (use_aligned_float8) {
+ quantize_mmq_nvfp4<false, true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+ } else {
+ quantize_mmq_nvfp4<false, false><<<num_blocks, block_size, 0, stream>>>(
+ x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+ }
} else {
GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
void quantize_mmq_fp4_cuda(const float * x,
const int32_t * ids,
void * vy,
+ float * scale,
ggml_type type_src0,
+ bool use_aligned_float8,
int64_t ne00,
int64_t s01,
int64_t s02,
void quantize_scatter_mmq_fp4_cuda(const float * x,
const int32_t * ids_src1_inv,
void * vy,
+ float * scale,
ggml_type type_src0,
+ bool use_aligned_float8,
int64_t ne00,
int64_t stride_token,
int64_t ne0,