upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-23 11:10:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

CUDA: dedup MoE gate/up activation quantization (#25441)

Browse Source 
* CUDA: dedup MoE gate/up activation quantization (fp4)

For MoE gate/up projections the src1 activation is broadcast across the
routed experts (ne11 == 1), so ids_src1 maps every one of a token's
n_expert_used slots to the same physical row. The MMQ path therefore
re-quantized each token's activation n_expert_used times.

For fp4 (NVFP4/MXFP4) src0, quantize each unique token row once instead of
once per expert. For NVFP4 a single quantize+scatter kernel
(quantize_scatter_mmq_nvfp4) quantizes each token once and writes the
resulting block_fp4_mmq straight to all n_expert_used slots, using an
inverse token->compact-row map (build_tok2c). MXFP4, and
GGML_CUDA_MOE_QUANT_GATHER=1, use a two-kernel variant: quantize unique
rows then gather into the expert-sorted layout (gather_mmq_fp4_blocks).
Both are bit-identical to the previous gather-then-quantize path (identical
source data, deterministic per-block quantization), verified by
test-backend-ops MUL_MAT_ID (type_a=nvfp4, broadcast b=1; 790/790 for the
default, gather, and per-expert paths) and by coherent end-to-end
generation. Set GGML_CUDA_NO_MOE_QUANT_DEDUP=1 to force the original
per-expert path.

Same-binary A/B on RTX 5090 (sm_120), Qwen3.6-35B-A3B-NVFP4 prefill @8192
(nsys, graphs-off; the unchanged mul_mat_q GEMM confirms stable clocks):
activation-quant GPU-busy drops 61% (78.2 -> 30.4 ms) with the fused
quantize+scatter, vs 33% (78.2 -> 52.8 ms) for the two-kernel gather. The
fused path avoids materializing and re-reading the 8x compact buffer,
writing the expert copies directly from registers.

* CUDA: bounds-check token ids in build_tok2c_kernel

Guard against malformed ids_src1: skip out-of-range token ids (t < 0 or
t >= n_tokens) and drop entries beyond n_expert_used per token instead of
writing past the token's tok2c region. No behavior change for valid MoE
routing data; test-backend-ops MUL_MAT_ID 790/790.

* Refactor the code based on review comments

- Removed previously added kernels that were not necessary anymore\
- Added an inverse mapping from (token, slot) to compact row. Each token is quantized once and scattered to its compact rows.

* Adding q8_1 support for dedup and addressing review comments

* Add pragma unrolls

* Remove redundant cudaMemsetAsync call

* Removing follow up redundancies

---------

Co-authored-by: praneshgo <227579474+praneshgo@users.noreply.github.com>
This commit is contained in:
Pranesh Gonegandla
2026-07-16 07:02:25 +00:00
committed by GitHub
parent a320cbfcb7
commit 5839ba3524
6 changed files with 255 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ggml/src/ggml-cuda/mmf.cu
View File

@@ -85,7 +85,7 @@ void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * sr
GGML_ASSERT(sis1 > 0);
ggml_cuda_launch_mm_ids_helper(ids_d, ids_src_compact_dev.get(), ids_dst_compact_dev.get(), expert_bounds_dev.get(),
static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, ctx.stream());
static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, /*write_inverse =*/ false, ctx.stream());
CUDA_CHECK(cudaGetLastError());
ids_info.ids_src_compact = ids_src_compact_dev.get();

31
ggml/src/ggml-cuda/mmid.cu
View File

@@ -27,7 +27,7 @@ template <int n_expert_used_template>
__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
@@ -98,8 +98,13 @@ static __global__ void mm_ids_helper(
const mm_ids_helper_store store_it = store[itc];
const int it = store_it.it();
const int iex_used = store_it.iex_used();
ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
ids_dst[nex_prev + itc] = it*n_expert_used + iex_used;
// ids_src1 holds the forward map, or the inverse map (token slot -> compact row) for quant dedup
if (write_inverse) {
ids_src1[it*n_expert_used + iex_used] = nex_prev + itc;
} else {
ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
}
}
if (threadIdx.x != 0) {
@@ -118,7 +123,7 @@ static __global__ void mm_ids_helper(
template <int n_expert_used_template>
static void launch_mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
GGML_ASSERT(n_tokens < (1 << 22) && "too few bits in mm_ids_helper_store");
GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
@@ -132,33 +137,33 @@ static void launch_mm_ids_helper(
const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
GGML_ASSERT(nbytes_shared <= smpbo);
mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
(ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1);
(ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
}
void ggml_cuda_launch_mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
switch (n_expert_used) {
case 2:
launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 4:
launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 6:
launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 8:
launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 16:
launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 32:
launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
default:
launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
}
}

2
ggml/src/ggml-cuda/mmid.cuh
View File

@@ -2,4 +2,4 @@
void ggml_cuda_launch_mm_ids_helper(
const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, cudaStream_t stream);
int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);

17
ggml/src/ggml-cuda/mmq.cu
View File

@@ -175,13 +175,17 @@ void ggml_cuda_mul_mat_q(
ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows);
ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
// gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
// scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).
const bool dedup_bcast = ne11 == 1 && n_expert_used > 1;
{
GGML_ASSERT(ids->nb[0] == ggml_element_size(ids));
const int si1 = ids->nb[1] / ggml_element_size(ids);
const int sis1 = nb12 / nb11;
ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
ne02, ne12, n_expert_used, ne11, si1, sis1, stream);
ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
CUDA_CHECK(cudaGetLastError());
}
@@ -198,7 +202,16 @@ 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 (use_native_fp4) {
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,
/*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);
}
} 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 {

294
ggml/src/ggml-cuda/quantize.cu
View File

@@ -75,10 +75,12 @@ __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
}
// scatter: grid over tokens, quantize once, write to all the token's compact rows
template <bool scatter>
static __global__ void quantize_mmq_nvfp4(
const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
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 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;
@@ -86,25 +88,25 @@ static __global__ void quantize_mmq_nvfp4(
return;
}
const int64_t i1 = blockIdx.x;
const int64_t i2 = blockIdx.z % ne2;
const int64_t i3 = blockIdx.z / ne2;
const int64_t i01 = ids ? ids[i1] : i1;
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 int64_t ib = blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x;
block_fp4_mmq * y = (block_fp4_mmq *) vy;
block_fp4_mmq * yb = y + ib;
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 i01 = ids ? ids[blockIdx.x] : blockIdx.x;
base_idx = i3 * s03 + i2 * s02 + i01 * s01;
}
float vals_raw[QK_NVFP4_SUB];
float amax_raw = 0.0f;
const int64_t base_idx = i3 * s03 + i2 * s02 + i01 * s01;
#pragma unroll
for (int k = 0; k < QK_NVFP4_SUB; k++) {
const int64_t i00 = i0_base + k;
@@ -160,11 +162,27 @@ static __global__ void quantize_mmq_nvfp4(
q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
}
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;
block_fp4_mmq * y = (block_fp4_mmq *) vy;
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.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);
NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
#endif // defined(BLACKWELL_MMA_AVAILABLE)
@@ -172,6 +190,8 @@ static __global__ void quantize_mmq_nvfp4(
// quantize values in the format mxfp4 is stored which is interleaved nibbles
// i.e. a block a0-a31 is represented as a0a16,a1a17 ...a15a31
// scatter: grid over tokens, quantize once, write to all the token's compact rows
template <bool scatter>
static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
const int32_t * __restrict__ ids,
void * __restrict__ vy,
@@ -181,7 +201,8 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
const int64_t s03,
const int64_t ne0,
const int ne1,
const int ne2) {
const int ne2,
const int n_expert_used) {
constexpr int vals_per_scale = 32;
constexpr int vals_per_warp = 2 * vals_per_scale; // Each warp processes 2 blocks of 32 = 64 values
@@ -196,30 +217,27 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
return;
}
const int64_t i1 = blockIdx.x;
const int64_t i2 = blockIdx.z % ne2;
const int64_t i3 = blockIdx.z / ne2;
ggml_cuda_pdl_sync();
const int64_t i01 = ids ? ids[i1] : i1;
const int64_t i02 = i2;
const int64_t i03 = i3;
block_fp4_mmq * y = (block_fp4_mmq *) vy;
const int64_t block_fp4_mmq_size = QK_FP4_MMQ;
const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
const int64_t ib = ib0 + (warp_start_offset / block_fp4_mmq_size) * ne1 + blockIdx.x;
const int64_t k_block = warp_start_offset / block_fp4_mmq_size;
const int64_t quad_idx_in_block = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp;
const int group_id = lane_id_32 / 4;
const int lane_in_group = lane_id_32 % 4;
const int base = group_id * 2;
char2 * yqs2 = (char2 *) y[ib].qs;
const int64_t base_pos = i03 * s03 + i02 * s02 + i01 * s01;
ggml_cuda_pdl_sync();
int64_t base_pos;
if constexpr (scatter) {
base_pos = (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 i01 = ids ? ids[blockIdx.x] : blockIdx.x;
base_pos = i3 * s03 + i2 * s02 + i01 * s01;
}
uint8_t scales[2];
char2 packed[2];
#pragma unroll
for (int b = 0; b < 2; ++b) {
@@ -244,11 +262,8 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
const float val2 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 1, WARP_SIZE);
const float val3 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 17, WARP_SIZE);
if (lane_in_group == 0) {
__nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = *(char2 *) &fp4_packed;
}
__nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
packed[b] = *(char2 *) &fp4_packed;
#else
// Fallback: manual FP4 conversion using LUT
const uint8_t q_val = ggml_cuda_float_to_fp4_e2m1(xi, inv_s);
@@ -258,26 +273,49 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
const uint8_t q_hi_0 = __shfl_sync(0xFFFFFFFF, q_val, base + 16, WARP_SIZE);
const uint8_t q_hi_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 17, WARP_SIZE);
if (lane_in_group == 0) {
char2 q;
q.x = (q_hi_0 << 4) | q_lo_0;
q.y = (q_hi_1 << 4) | q_lo_1;
yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = q;
}
char2 q;
q.x = (q_hi_0 << 4) | q_lo_0;
q.y = (q_hi_1 << 4) | q_lo_1;
packed[b] = q;
#endif // CUDART_VERSION >= 12080
}
if (lane_id_32 == 0) {
// Store 2 scales packed into 1 uint32
y[ib].d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
block_fp4_mmq * y = (block_fp4_mmq *) vy;
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);
char2 * yqs2 = (char2 *) yb->qs;
if (lane_in_group == 0) {
yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
}
if (lane_id_32 == 0) {
yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
}
}
} else {
const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
block_fp4_mmq * yb = y + (ib0 + k_block * ne1 + blockIdx.x);
char2 * yqs2 = (char2 *) yb->qs;
if (lane_in_group == 0) {
yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
}
if (lane_id_32 == 0) {
yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
}
}
GGML_UNUSED(n_expert_used);
}
template <mmq_q8_1_ds_layout ds_layout>
// scatter: grid over tokens, quantize once, write to all the token's compact rows
template <mmq_q8_1_ds_layout ds_layout, bool scatter>
static __global__ void quantize_mmq_q8_1(
const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int ne1, const int ne2) {
const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) {
constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;
constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;
@@ -288,26 +326,27 @@ static __global__ void quantize_mmq_q8_1(
return;
}
const int64_t i1 = blockIdx.x;
const int64_t i2 = blockIdx.z % ne2;
const int64_t i3 = blockIdx.z / ne2;
const int64_t i00 = i0;
ggml_cuda_pdl_sync();
const int64_t i01 = ids ? ids[i1] : i1;
const int64_t i02 = i2;
const int64_t i03 = i3;
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 i01 = ids ? ids[blockIdx.x] : blockIdx.x;
base_idx = i3*s03 + i2*s02 + i01*s01;
}
const float4 * x4 = (const float4 *) x;
block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
const int64_t ib = ib0 + (i0 / QK8_1_MMQ)*ne1 + blockIdx.x; // block index in channel
const int64_t iqs = i0 % QK8_1_MMQ; // quant index in block
const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel
const int64_t iqs = i0 % QK8_1_MMQ; // quant index in block
// Load 4 floats per thread and calculate max. abs. value between them:
const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
const float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float amax = fabsf(xi.x);
amax = fmaxf(amax, fabsf(xi.y));
amax = fmaxf(amax, fabsf(xi.z));
@@ -336,40 +375,41 @@ static __global__ void quantize_mmq_q8_1(
q.y = roundf(xi.y*d_inv);
q.z = roundf(xi.z*d_inv);
q.w = roundf(xi.w*d_inv);
// Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
char4 * yqs4 = (char4 *) y[ib].qs;
yqs4[iqs/4] = q;
if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
if (iqs % 16 != 0 || iqs >= 96) {
return;
}
y[ib].d2s6[2 + iqs/16] = sum;
if (iqs % 64 != 0) {
return;
}
const float d = 1.0f / d_inv;
y[ib].d2s6[iqs/64] = d;
return;
}
if (iqs % 32 != 0) {
return;
}
const float d = 1.0f / d_inv;
if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
y[ib].ds4[iqs/32] = make_half2(d, sum);
} else {
y[ib].d4[iqs/32] = d;
// write the block once (normal) or to each of the token's compact rows (scatter)
const int nwrite = scatter ? n_expert_used : 1;
#pragma unroll
for (int slot = 0; slot < nwrite; ++slot) {
int64_t ib;
if constexpr (scatter) {
const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
ib = k_block*ne1 + i;
} else {
const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
ib = ib0 + k_block*ne1 + blockIdx.x;
}
// Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
char4 * yqs4 = (char4 *) y[ib].qs;
yqs4[iqs/4] = q;
if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
if (iqs % 16 == 0 && iqs < 96) {
y[ib].d2s6[2 + iqs/16] = sum;
if (iqs % 64 == 0) {
y[ib].d2s6[iqs/64] = d;
}
}
} else if (iqs % 32 == 0) {
if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
y[ib].ds4[iqs/32] = make_half2(d, sum);
} else {
y[ib].d4[iqs/32] = d;
}
}
}
GGML_UNUSED(n_expert_used);
}
void quantize_row_q8_1_cuda(
@@ -402,16 +442,16 @@ void quantize_mmq_q8_1_cuda(
const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
switch (mmq_get_q8_1_ds_layout(type_src0)) {
case MMQ_Q8_1_DS_LAYOUT_D4:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, false>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
break;
case MMQ_Q8_1_DS_LAYOUT_DS4:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, false>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
break;
case MMQ_Q8_1_DS_LAYOUT_D2S6:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, false>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
break;
default:
GGML_ABORT("fatal error");
@@ -419,6 +459,62 @@ void quantize_mmq_q8_1_cuda(
}
}
// scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row)
void quantize_scatter_mmq_q8_1_cuda(
const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
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(ne00 % 4 == 0);
GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
const dim3 num_blocks(n_tokens, block_num_y, 1);
const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
switch (mmq_get_q8_1_ds_layout(type_src0)) {
case MMQ_Q8_1_DS_LAYOUT_D4:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, true><<<num_blocks, block_size, 0, stream>>>(
x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
break;
case MMQ_Q8_1_DS_LAYOUT_DS4:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, true><<<num_blocks, block_size, 0, stream>>>(
x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
break;
case MMQ_Q8_1_DS_LAYOUT_D2S6:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, true><<<num_blocks, block_size, 0, stream>>>(
x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
break;
default:
GGML_ABORT("fatal error");
break;
}
}
// 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 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(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);
} else {
GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4);
constexpr int nwarps = 8;
constexpr int vals_per_block = nwarps * 2 * QK_MXFP4;
const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block;
const dim3 block_size(WARP_SIZE, nwarps, 1);
const dim3 num_blocks(n_tokens, block_num_y, 1);
quantize_mmq_mxfp4<true><<<num_blocks, block_size, 0, stream>>>(
x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
}
}
void quantize_mmq_fp4_cuda(
const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
@@ -432,8 +528,8 @@ void quantize_mmq_fp4_cuda(
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<<<num_blocks, block_size, 0, stream>>>(
x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
quantize_mmq_nvfp4<false><<<num_blocks, block_size, 0, stream>>>(
x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
} else {
GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
@@ -445,6 +541,6 @@ void quantize_mmq_fp4_cuda(
const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
const dim3 block_size(WARP_SIZE, nwarps, 1);
quantize_mmq_mxfp4<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
quantize_mmq_mxfp4<false><<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
}
}

25
ggml/src/ggml-cuda/quantize.cuh
View File

@@ -39,3 +39,28 @@ void quantize_mmq_fp4_cuda(const float * x,
int64_t ne2,
int64_t ne3,
cudaStream_t stream);
// quantize each token once and scatter the block to its compact rows (via the inverse map)
void quantize_scatter_mmq_fp4_cuda(const float * x,
const int32_t * ids_src1_inv,
void * vy,
ggml_type type_src0,
int64_t ne00,
int64_t stride_token,
int64_t ne0,
int64_t n_tokens,
int64_t nrows_dst,
int n_expert_used,
cudaStream_t stream);
void quantize_scatter_mmq_q8_1_cuda(const float * x,
const int32_t * ids_src1_inv,
void * vy,
ggml_type type_src0,
int64_t ne00,
int64_t stride_token,
int64_t ne0,
int64_t n_tokens,
int64_t nrows_dst,
int n_expert_used,
cudaStream_t stream);