CUDA: dedup MoE gate/up activation quantization (#25441)
* 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:
co-authored by
praneshgo
parent
a320cbfcb7
commit
5839ba3524
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user