[CPU] QMoE: float zero-points for fractional/asymmetric quantization - #32199
Draft
Thiago Pereira Rocha (thpereir) wants to merge 3 commits into
Draft
[CPU] QMoE: float zero-points for fractional/asymmetric quantization#32199Thiago Pereira Rocha (thpereir) wants to merge 3 commits into
Thiago Pereira Rocha (thpereir) wants to merge 3 commits into
Conversation
Gemma-style MoE experts use GeGLU (gelu_pytorch_tanh(gate) * up), which the
fused CPU QMoE op could not express: its only gated path was SwiGLU
(silu-gated), and plain gelu is non-gated (single-width fc1). Add a
first-class geglu activation that reuses the interleaved doubled-fc1 SwiGLU
machinery but applies the gelu-tanh gate. CPU QMoE only; CUDA and float MoE
paths are unchanged.
- moe_base_cpu.h: ActivationType::GeGLU + "geglu" parser branch
- moe_utils.{h,cc}: ApplyGeGLUActivation (tanh-gelu gate, matches HF
gelu_pytorch_tanh)
- moe_quantization_cpu.cc: treat GeGLU as a gated activation (doubled fc1,
swiglu_fusion=1) and dispatch to ApplyGeGLUActivation
Mirror the existing SwiGLU QMoE CPU parity suite for the new GeGLU activation: geglu() reference (gelu-tanh gate matching ApplyGeGLUActivation), GegluMlp / GegluMoeConfig / GegluMoEBlock, and TestGegluQMoECPU covering row-wise and block-wise, symmetric and asymmetric, at 2/4/8-bit. The shared ONNX graph builder now emits activation_type="geglu" with a plain gelu gate (alpha=1, beta=0, no clamp), distinct from SwiGLU's scaled/clamped gate.
Enables the CPU QMoE kernel to run the fused Quark uint2 gemma-4-26B-A4B-it
export, whose zero-points are float (a constant 1.5 -> codes {0,1,2,3} map to
{-1.5,-0.5,0.5,1.5}*scale), unrepresentable on the integer code grid.
Float zero-points (this change):
- contrib_defs.cc: new TZ type constraint so fc*_zero_points accept float/
float16/bfloat16 in addition to packed uint8
- moe_quantization_cpu.cc: detect float zp by dtype, read it as a parallel
unpacked one-per-group pointer (scales layout), gate off the integer-only
fast paths (LUT GEMM, prepack, direct Q4), and dequantize as (code-zp)*scale
- moe_helper.h: per-tensor zp packing factor so unpacked float zp passes shape
validation (pack factor 1 instead of 8/bits)
GeGLU CPU QMoE support (from the geglu-qmoe-cpu baseline; the worktree base
commit predates it): the ORT_ENFORCE now accepts GeGLU as well as SwiGLU with
swiglu_fusion=1, and the gated-activation path dispatches to
ApplyGeGLUActivation for the geglu activation type.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| ) | ||
| # Swap SwiGLU experts for GeGLU experts and flip the activation flags. The graph is | ||
| # (re)built lazily in recreate_onnx_model(), which reads self.use_geglu. | ||
| self.use_swiglu = False |
| # (re)built lazily in recreate_onnx_model(), which reads self.use_geglu. | ||
| self.use_swiglu = False | ||
| self.use_geglu = True | ||
| self.experts = nn.ModuleList([GegluMlp(config) for _ in range(self.num_experts)]) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds fractional float zero-points to CPU QMoE, alongside stacked GeGLU changes from #32153.
Changes:
- Extends QMoE zero-point schema types and shape validation.
- Adds float zero-point dequantization routing.
- Adds CPU GeGLU implementation and parity tests.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
contrib_defs.cc |
Extends QMoE zero-point schema types. |
moe_base_cpu.h |
Adds GeGLU activation parsing. |
moe_helper.h |
Validates unpacked float zero-point shapes. |
moe_quantization_cpu.cc |
Implements float zero-point and GeGLU dispatch. |
moe_utils.h |
Declares GeGLU helper. |
moe_utils.cc |
Implements GeGLU activation. |
test_qmoe_cpu.py |
Adds GeGLU parity coverage. |
Suppressed comments (1)
onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc:1225
- No test in this PR supplies a non-
uint8zero-point:create_cpu_moe_onnx_graphstill unconditionally casts both zero-point tensors tonumpy.uint8and emitsTensorProto.UINT8(test_qmoe_cpu.py:594-617). Thus the new branch and the motivating fractional value 1.5 are never exercised, which leaves the unsupported 4-bit, 8-bit, and row-wise paths undetected. Add fractional float-zero-point parity tests covering the intended bit widths/layouts and fast-path settings.
const bool fc1_zp_is_float = fc1_zero_points != nullptr &&
fc1_zero_points->template IsDataType<T>();
const bool fc2_zp_is_float = fc2_zero_points != nullptr &&
fc2_zero_points->template IsDataType<T>();
const uint8_t* fc1_zp_data = (fc1_zero_points && !fc1_zp_is_float) ? fc1_zero_points->template Data<uint8_t>() : nullptr;
const uint8_t* fc2_zp_data = (fc2_zero_points && !fc2_zp_is_float) ? fc2_zero_points->template Data<uint8_t>() : nullptr;
const T* fc1_zp_fp_data = fc1_zp_is_float ? fc1_zero_points->template Data<T>() : nullptr;
const T* fc2_zp_fp_data = fc2_zp_is_float ? fc2_zero_points->template Data<T>() : nullptr;
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+46
to
+47
| } else if (activation_type_str == "geglu") { | ||
| activation_type_ = ActivationType::GeGLU; |
Comment on lines
+313
to
+318
| const TScale* zero_points_fp) { | ||
| // zero_points_fp (optional): float zero-points laid out identically to `scales` | ||
| // (unpacked, one value per group, same index). When present it takes precedence over the | ||
| // packed integer `zero_points` and dequant is w = (code - zero_point) * scale. This supports | ||
| // fractional/asymmetric schemes (e.g. Quark uint2 with a constant zp of 1.5) that cannot be | ||
| // represented on the integer code grid. Only wired for the 2/4-bit block-wise branch. |
Comment on lines
+1673
to
+1675
| .TypeConstraint("TZ", {"tensor(uint8)", "tensor(float8e4m3fn)", "tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, | ||
| "Constrain zero-point types. Integer zero-points use packed uint8 (same layout as weights). " | ||
| "Float zero-points are unpacked (one per group, scales layout) and support fractional/asymmetric schemes.") |
Comment on lines
+1604
to
+1608
| "3D tensor with shape (num_experts, fusion_size * inter_size, hidden_size / block_size / pack_size) when block_size is provided. " | ||
| "Integer zero-points use the same packed uint8 layout as the weights (T1). Float zero-points " | ||
| "(TZ) are unpacked, one value per group, matching the scales layout, and dequantize as " | ||
| "w = (code - zero_point) * scale.", | ||
| "TZ", |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds float zero-point support to the CPU QMoE kernel, enabling fractional/asymmetric quantization schemes that cannot be represented on the integer code grid. Motivating case: the Quark
uint2gemma-4-26B-A4B-itexport, whose zero-points are a constant1.5(codes{0,1,2,3}→{-1.5,-0.5,0.5,1.5} * scale).What changes
contrib_defs.cc: newTZtype constraint sofc*_zero_pointsacceptfloat/float16/bfloat16in addition to packeduint8; schema docs describe the integer (T1) vs float (TZ) layouts.moe_quantization_cpu.cc: detect float zp by element type, read it as a parallel unpacked one-per-group pointer (same layout asscales), gate off the integer-only fast paths (LUT GEMM, prepack, direct Q4), and dequantize asw = (code - zero_point) * scale.moe_helper.h: per-tensor zp packing factor so unpacked float zp passes shape validation (pack factor 1 instead of8/bits).Integer zero-points are unchanged — float zp is detected purely by dtype and only alters the block-wise dequant branch.
Test plan