CUDA: dedup MoE gate/up activation quantization - #25441
Conversation
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.
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.
github-actions
Bot
added
ggml
|
@gaugarg-nv can you please help review this PR? Thanks. |
|
Hi @praneshgo, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
This is addressed now. |
|
How is this different from #18538? |
Shouldn't this technique be generally applicable to Q8_1 as well? Also, do you have some perf measurements you could share to justify this change? |
These changes are complementary but not duplicates. Both attempt to remove redundant activation quantization for MoE MMQ, but along orthogonal axes. The goal of #18538 seems to be to not re-quantize the same tensor when a later op reuses it. Let T = tokens, E = n_expert_used. For models with separate gate/up (two mul_mat_id on the same activation)
For models with merged gate_up (single mul_mat_id)
|
Yes, this should be applicable to Q8_1 as well. The corresponding changes are in progress, and I can add them as a commit to this PR soon, if we want to track both of them here. Performance numbers and the associated perf gain with the changes in this PR are updated in the description. |
JohannesGaessler
left a comment
JohannesGaessler
left a comment
There was a problem hiding this comment.
If I understand this PR correctly it adds too much complexity for what it does. On master the quantization kernel is doing duplicated work proportional to n_expert_used. The correct way to fix that is to change grid size and indexing logic of the quantization kernel to instead quantize the activations once and to write back the result n_expert_used times. No extra kernels should be added.
- 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.
Fair point. I have refactored the code based on your suggestions, by adding indexing logic such that activations are quantized once and the result is written back |
JohannesGaessler
left a comment
JohannesGaessler
left a comment
There was a problem hiding this comment.
If at all possible, modify the existing kernels both for the FP4 and the q8_0 path to use the deduplicated pattern unconditionally. It is not necessary to preserve the pattern on master for MMQ or to make the feature configurable.
| __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, | ||
| int32_t * __restrict__ ids_src1_inv, |
There was a problem hiding this comment.
A kernel will only ever need one out of the current or the inverse mapping. Therefore, rather than a pointer pass a boolean to control whether the inverse mapping should be written to ids_src1_inv.
There was a problem hiding this comment.
Both fair points. Made the dedup unconditional for both quantization types.
Also, instead of passing the pointer, I am now passing a boolean, which when set to true writes the inverse mapping to ids_src1 (and the forward mapping when false), so it only ever builds the one that is needed.
praneshgo
changed the title
Given we intend to change NVFP4 to do dynamic per-channel amax-quant, does it make sense to withhold this PR until those changes have been made and reevaluate their impact on the new quantization logic? |
I think they are orthogonal on correctness. The dedup removes the redundant per-expert re-quant within one MUL_MAT_ID, so it composes with per-channel amax and stays bit-exact. The perf side is worth re-checking once amax is in, and I can recheck it then. However, I do not see any harm in this change landing now. It is self-contained and gives a gain under the current kernel. |
JohannesGaessler
left a comment
JohannesGaessler
left a comment
There was a problem hiding this comment.
I'll check this PR for performance regressions but other than that I would consider it essentially good to merge.
| return; | ||
| // write the block once (normal) or to each of the token's compact rows (scatter) | ||
| const int nwrite = scatter ? n_expert_used : 1; | ||
| for (int slot = 0; slot < nwrite; ++slot) { |
There was a problem hiding this comment.
| for (int slot = 0; slot < nwrite; ++slot) { | |
| #pragma unroll | |
| for (int slot = 0; slot < nwrite; ++slot) { |
Please add explicit #pragma unrolls.
There was a problem hiding this comment.
Added explicit pragma unrolls in the latest commit.
| if (dedup_bcast) { | ||
| CUDA_CHECK(cudaMemsetAsync(ids_src1.get(), 0xFF, ne_get_rows*sizeof(int32_t), stream)); // -1 for unused slots | ||
| } |
There was a problem hiding this comment.
Why did you add this call to cudaMemsetAsync? I don't think it's needed, ids_src1 should be completely initialized even with the inverse mapping. I am not getting incorrect results from removing this call, I do however see an end-to-end performance difference of up to 5% on my RTX 5090.
There was a problem hiding this comment.
Good catch, it is in fact redundant. I had initially added it to initialize -1 for unwritten slots, but mm_ids_helper fills every (token, slot) entry, so all the slots are initialized anyway. Removed it in the latest commit.
…gonegandla/moe_quant_optimize
Performance
|
* 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>
* 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>
* 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>
* 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>
5 participants
Overview
For MoE gate/up projections the src1 activation is broadcast across the routed experts (ne11 == 1), so the MMQ path re-quantized each token's activation once per routed expert. This deduplicates it: each unique token's activation is quantized once and reused for all of its experts.
Additional information
Output is bit-identical to the previous per-expert path (same source data, deterministic per-block quantization). Down projection is unaffected (its rows differ per expert), and dense (non-MoE) models never trigger this path. This removes redundant activation-quant work in MoE prefill.
Verification:
Perf result comparison on running with GPU: RTX 5090 ·
llama-bench -p 8192 -n 128 -r 10 -ngl 99, BS=1, CUDA graphs onQ8_1 is perf-neutral, since its activation quant is a cheap, memory-bound amax + int8 round, so there is little redundant work to remove (unlike NVFP4's per-sub-block scale search). We apply the dedup there anyway for parity across quantization types.
Requirements