Skip to content

Extend MoEPacked for Bolt ReLU-square routing and BF16 IQ2 precision #33963

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Rui Ren (rui-ren) wants to merge 9 commits into from
+1,119 −89
Open

Extend MoEPacked for Bolt ReLU-square routing and BF16 IQ2 precision #33963

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
243376e
Support ReLU-square experts and preselected FP32 MoEPacked routes
rui-ren Nov 7, 2026
7a24246
Preserve projected sigmoid routing in packed ReLU-square MoE
rui-ren Nov 7, 2026
3b37af1
Preserve contiguous BF16 IQ2 ReLU-square MoE reductions
rui-ren Nov 7, 2026
02fa4df
Make MoE cancellation regression portable across CUDA architectures
rui-ren Nov 7, 2026
40cd165
Refresh generated packed MoE operator documentation
rui-ren Nov 7, 2026
e5adafd
Format combined packed MoE rebase resolution
rui-ren Nov 9, 2026
edb5e69
Preserve 32-route MoEPacked modes after upstream validation rebase
rui-ren Nov 9, 2026
ea0e854
Harden MoEPacked routing contracts and independent IQ2 regression tests
rui-ren Nov 9, 2026
79067e3
Verify malformed MoE outputs against per-element NaN expectations
rui-ren Nov 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 29 additions & 11 deletions docs/ContribOperators.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4976,7 +4976,7 @@ This version of the operator has been available since version 1 of the 'com.micr

### <a name="com.microsoft.MoEPacked"></a><a name="com.microsoft.moepacked">**com.microsoft.MoEPacked**</a>

Mixture-of-Experts layer consuming 3D block-packed gate, up, and down expert weight blobs. CUDA execution requires weights in device memory; host-backed external initializers are copied to the device without changing their packed layout.
Mixture-of-Experts layer consuming 3D block-packed gate and down expert weight blobs, plus up expert weights for SwiGLU. CUDA execution requires weights in device memory; host-backed external initializers are copied to the device without changing their packed layout.
Supported format IDs and their stable byte layouts are defined in the [packed-weight format contract](https://github.com/microsoft/onnxruntime/blob/main/docs/PackedWeightFormats.md).

#### Version
Expand All @@ -4987,32 +4987,46 @@ This version of the operator has been available since version 1 of the 'com.micr

<dl>
<dt><tt>activation_type</tt> : string</dt>
<dd>Expert activation function ('swiglu').</dd>
<dd>Expert activation function: 'swiglu' or 'relu2' (squared ReLU; 'relu_squared' is an alias). relu2 requires IQ2_XXS (type 16) gate and down weights and omitted up weights.</dd>
<dt><tt>down_type</tt> : int (required)</dt>
<dd>Stable packed-format ID of down_experts_weights. Only contract-listed IDs are supported.</dd>
<dt><tt>gate_type</tt> : int (required)</dt>
<dd>Stable packed-format ID of gate_experts_weights. Only contract-listed IDs are supported.</dd>
<dt><tt>k</tt> : int (required)</dt>
<dd>Number of selected experts per token.</dd>
<dd>Number of selected experts per token: 1..16 for internal softmax routing, 1..32 for projected sigmoid or preselected routing; at most num_experts.</dd>
<dt><tt>normalize_routing_weights</tt> : int</dt>
<dd>Whether to normalize top-k routing probabilities.</dd>
<dt><tt>up_type</tt> : int (required)</dt>
<dd>Stable packed-format ID of up_experts_weights. Only contract-listed IDs are supported.</dd>
<dd>Whether to renormalize selected softmax probabilities. Must be 1 for projected sigmoid routing and 0 for preselected routing.</dd>
<dt><tt>router_scale</tt> : float</dt>
<dd>Finite post-normalization scale for projected sigmoid routing; must be 1 for softmax or preselected routing.</dd>
<dt><tt>routing_function</tt> : string</dt>
<dd>Internal routing: softmax logits or projected FP32 sigmoid with selection-only bias.</dd>
<dt><tt>up_type</tt> : int</dt>
<dd>Stable packed-format ID of up_experts_weights; required for swiglu. Only contract-listed IDs are supported.</dd>
</dl>

#### Inputs
#### Inputs (5 - 10)

<dl>
<dt><tt>input</tt> : T</dt>
<dd>Input hidden states [..., hidden_size].</dd>
<dt><tt>router_probs</tt> : T</dt>
<dd>Router logits [num_rows, num_experts].</dd>
<dt><tt>router_probs</tt> (optional) : T</dt>
<dd>Router logits [num_rows, num_experts]; omit for projected sigmoid or preselected routing.</dd>
<dt><tt>gate_experts_weights</tt> : T2</dt>
<dd>3D packed gate weight blob [num_experts, inter_size, gate_row_bytes].</dd>
<dt><tt>up_experts_weights</tt> : T2</dt>
<dd>3D packed up weight blob [num_experts, inter_size, up_row_bytes].</dd>
<dt><tt>up_experts_weights</tt> (optional) : T2</dt>
<dd>3D packed up weight blob [num_experts, inter_size, up_row_bytes]; required for swiglu, omit for relu2.</dd>
<dt><tt>down_experts_weights</tt> : T2</dt>
<dd>3D packed down weight blob [num_experts, hidden_size, down_row_bytes].</dd>
<dt><tt>expert_indices</tt> (optional) : Tind</dt>
<dd>Preselected distinct expert IDs [num_rows, k], k <= 32. Requires expert_weights, no router_probs and normalize_routing_weights=0. Invalid or duplicate IDs and nonfinite weights poison only their token row.</dd>
<dt><tt>expert_weights</tt> (optional) : Tweight</dt>
<dd>Preselected FP32 weights [num_rows, k], preserved without normalization and mixed in supplied order.</dd>
<dt><tt>router_input</tt> (optional) : T</dt>
<dd>Projected sigmoid routing input [num_rows, router_width]; omit other routing inputs. router_width must be positive, fit int32 and be divisible by 256. Requires relu2, routing_function='sigmoid' and normalize_routing_weights=1.</dd>
<dt><tt>router_weight</tt> (optional) : T</dt>
<dd>Sigmoid router weight [num_experts, router_width].</dd>
<dt><tt>router_bias</tt> (optional) : Tweight</dt>
<dd>FP32 correction bias [num_experts], used only for expert selection.</dd>
</dl>

#### Outputs
Expand All @@ -5029,6 +5043,10 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Activation type.</dd>
<dt><tt>T2</tt> : tensor(uint8)</dt>
<dd>Packed byte storage.</dd>
<dt><tt>Tind</tt> : tensor(int32)</dt>
<dd>Preselected expert index type.</dd>
<dt><tt>Tweight</tt> : tensor(float)</dt>
<dd>Routing weight and correction bias type.</dd>
</dl>


Expand Down
2 changes: 1 addition & 1 deletion docs/OperatorKernels.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1159,7 +1159,7 @@ The **OpSet Version** column uses the following notation:
|MatMulPacked|*in* A:**T1**<br> *in* B:**T2**<br> *in* global_scale:**Tscale**<br> *out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float16)<br/> **T2** = tensor(uint8)<br/> **Tscale** = tensor(float)|
|MoE|*in* input:**T**<br> *in* router_probs:**T**<br> *in* fc1_experts_weights:**T**<br> *in* fc1_experts_bias:**T**<br> *in* fc2_experts_weights:**T**<br> *in* fc2_experts_bias:**T**<br> *in* fc3_experts_weights:**T**<br> *in* fc3_experts_bias:**T**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|MoEGroupRouter|*in* logits:**T**<br> *in* expert_bias:**TF**<br> *out* router_probs:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **TF** = tensor(float)|
|MoEPacked|*in* input:**T**<br> *in* router_probs:**T**<br> *in* gate_experts_weights:**T2**<br> *in* up_experts_weights:**T2**<br> *in* down_experts_weights:**T2**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)<br/> **T2** = tensor(uint8)|
|MoEPacked|*in* input:**T**<br> *in* router_probs:**T**<br> *in* gate_experts_weights:**T2**<br> *in* up_experts_weights:**T2**<br> *in* down_experts_weights:**T2**<br> *in* expert_indices:**Tind**<br> *in* expert_weights:**Tweight**<br> *in* router_input:**T**<br> *in* router_weight:**T**<br> *in* router_bias:**Tweight**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)<br/> **T2** = tensor(uint8)<br/> **Tind** = tensor(int32)<br/> **Tweight** = tensor(float)|
|MultiHeadAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* bias:**T**<br> *in* key_padding_mask:**M**<br> *in* attention_bias:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* past_sequence_length:**M**<br> *in* cache_indirection:**M**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**<br> *out* qk:**QK**|1+|**QK** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)|
|NGramHashMapping|*in* input_ids:**M**<br> *in* multipliers:**M**<br> *in* vocab_sizes:**M**<br> *in* past_ids:**M**<br> *in* head_offsets:**M**<br> *in* eos_token_id:**M**<br> *in* segment_ids:**tensor(int32)**<br> *out* hash_ids:**M**<br> *out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)|
|NGramRepeatBlock|*in* input_ids:**Tid**<br> *in* scores:**T**<br> *out* scores_out:**T**|1+|**T** = tensor(float)<br/> **Tid** = tensor(int64)|
Expand Down
Loading
Loading