__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;
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) {
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");
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;
}
}
}
+// 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;
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;
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)
// 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,
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
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) {
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);
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;
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;
- const float4 * x4 = (const float4 *) x;
+ 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));
q.y = roundf(xi.y*d_inv);
q.z = roundf(xi.z*d_inv);
q.w = roundf(xi.w*d_inv);
+ const float d = 1.0f / 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;
+ // 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;
}
- y[ib].d2s6[2 + iqs/16] = sum;
-
- if (iqs % 64 != 0) {
- return;
+ // 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;
+ }
}
-
- 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;
}
+ GGML_UNUSED(n_expert_used);
}
void quantize_row_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");
}
}
+// 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,
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);
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);
}
}
overview / pkg / ggml / sources / llama.cpp / commitdiff