Skip to content

Commit d411562

Browse files
Copilotjustinchuby
andauthored
Set QMoE raw weight prepack mode
Emit the new weights_prepacked=0 QMoE attribute for raw int expert weights now that ORT can prepack them internally. Add focused assertions for both component-built and rewrite-fused QMoE nodes. Signed-off-by: GitHub <noreply@github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
1 parent bbad5b1 commit d411562

4 files changed

Lines changed: 10 additions & 0 deletions

File tree

‎src/mobius/components/_moe.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,8 @@ def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value):
366366
expert_weight_bits=quantization.bits,
367367
block_size=quantization.group_size,
368368
swiglu_fusion=2,
369+
quant_type="int",
370+
weights_prepacked=0,
369371
_domain="com.microsoft",
370372
)
371373
if output_scale != 1.0: # noqa: RUF069

‎src/mobius/components/_moe_test.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,8 @@ def test_int4_olive_moe_emits_expert_major_qmoe(self):
160160
assert qmoe.attributes["block_size"].value == 32
161161
assert qmoe.attributes["swiglu_fusion"].value == 2
162162
assert qmoe.attributes["expert_weight_bits"].value == 4
163+
assert qmoe.attributes["quant_type"].value == "int"
164+
assert qmoe.attributes["weights_prepacked"].value == 0
163165
assert layer.experts is None
164166
assert layer.fc1_experts_weights.shape == ir.Shape([64, 64, 32])
165167
assert layer.fc2_experts_weights.shape == ir.Shape([64, 64, 16])
@@ -209,6 +211,8 @@ def test_int4_moe_emits_expert_major_qmoe(self):
209211
assert qmoe.attributes["k"].value == 6
210212
assert qmoe.attributes["block_size"].value == 32
211213
assert qmoe.attributes["swiglu_fusion"].value == 2
214+
assert qmoe.attributes["quant_type"].value == "int"
215+
assert qmoe.attributes["weights_prepacked"].value == 0
212216
assert layer.fc1_experts_weights.shape == ir.Shape([64, 64, 32])
213217
assert layer.fc1_scales.shape == ir.Shape([64, 64, 2])
214218
assert layer.fc1_experts_zero_points.shape == ir.Shape([64, 64, 1])

‎src/mobius/rewrite_rules/_qmoe_fusion.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,8 @@ def _fuse_layer(graph: ir.Graph, layer: _DenseMoELayer, index: int) -> None:
404404
"expert_weight_bits": bits,
405405
"block_size": block_size,
406406
"swiglu_fusion": 2,
407+
"quant_type": "int",
408+
"weights_prepacked": 0,
407409
},
408410
domain=_MS_DOMAIN,
409411
num_outputs=1,

‎src/mobius/rewrite_rules/_qmoe_fusion_test.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -478,6 +478,8 @@ def test_attributes_match_qmoe_abi() -> None:
478478
assert attrs["swiglu_fusion"].value == 2
479479
assert attrs["normalize_routing_weights"].value == 1
480480
assert attrs["activation_type"].value == "swiglu"
481+
assert attrs["quant_type"].value == "int"
482+
assert attrs["weights_prepacked"].value == 0
481483

482484

483485
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)