Introducing OptimizedMoeTransform by ochougul · Pull Request #1190 · quic/efficient-transformers

ochougul · 2026-07-17T11:26:55Z

This PR standardizes QEff MoE model wrappers around a single explicit
QEffMoEBlockMixin contract.

Key changes:

Enforces every QEffMoEBlockMixin subclass to explicitly declare
supported_moe_flavours.
Collapses DeepSeek V3 MoE into a single QEffDeepseekV3MoE block and
removes the separate prefill-only class.
Adds explicit MoE flavour declarations for DeepSeek V3 and Grok1.
Introduces a single Gemma4 text MoE block backed by QEffMoEBlockMixin.
Updates external MoE mapper wiring to source support metadata from the
QEff block classes.
Removes Gemma4’s prefill-only expert mapping in favor of the shared MoE
flavour path.

Impact For New MoE Models

When adding a new MoE model, the model should expose one QEff MoE block that
inherits QEffMoEBlockMixin and explicitly declares:

supported_moe_flavours = (...)

The block should implement the common variation points:

route(...)
build_moe_weights(...) / get_moe_weights(...)
moe_profile
optional shared expert handling

This keeps flavour selection, export config, and future optimizations
centralized instead of spreading model-specific prefill/decode paths across
multiple classes.

Future Scope

The longer-term direction is to support all MoE flavours across all MoE
models where the model architecture allows it, especially:

decode_bmm
simple_loop
expert_parallel

This PR makes that easier by making supported flavour declarations explicit
and discoverable.

Tests

Added/updated coverage for:

explicit supported_moe_flavours enforcement
DeepSeek V3 single-block/external mapper behavior
Grok1 explicit flavour metadata
Gemma4 MoE block parity for decode_bmm and simple_loop
DeepSeek legacy prefill fallback behavior

Validation run:

pytest tests/transformers/models/test_moe_prefill_blocked.py -k 'moe_block
or deepseek or grok or gemma4'
pytest tests/unit_test/transforms/test_transform_accuracy.py -k
'OptimizedMoE or deepseek or external_mapper'
pytest tests/unit_test/models/test_model_quickcheck.py -k 'moe_prefill'

ochougul · 2026-07-21T07:38:03Z

-
-        hidden_states = hidden_states + self.shared_experts(residuals)
-        return hidden_states
+QEffPrefillChunkedGlm4MoeMoE = QEffGlm4MoeMoE


not needed.

ochougul · 2026-07-21T07:38:37Z

+        if hasattr(self.experts, "gate_up_proj"):
+            self.moe_weights = build_canonical_expert_weights(
+                gate_up=self.experts.gate_up_proj,
+                down=self.experts.down_proj,
+                fused=True,
+                fused_split_dim=1,
+                transpose_gate_up=True,
+                transpose_down=True,
+            )
+            delete_module_attrs(self.experts, "gate_up_proj", "down_proj")
+        else:


this if-else is not needed, as we always know what weights model is going to have, let's fix it instead of handling un-necessary cases.

ochougul · 2026-07-21T07:39:49Z

+    order = torch.argsort(out, dim=1)
+    last_positions = order[:, -state_len:]


ochougul · 2026-07-21T07:39:57Z

+        mask_causal = torch.ones(chunk_size, chunk_size, dtype=torch.bool)
        for i in range(chunk_size):
-            for j in range(i, chunk_size):
-                mask_causal[i, j] = True
+            for j in range(i + 1):
+                mask_causal[i, j] = False


ochougul · 2026-07-21T07:40:11Z

        g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
-        # mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)
-        mask = mask_causal
+        mask = mask_causal.to(device=query.device)


ochougul · 2026-07-28T08:00:45Z

TODO: allow cores_per_expert to be fraction
TODO: what happens when we do prefill -> decode compile back to back on same object?

quic-rishinr · 2026-08-06T09:19:30Z

        if kv_cache_prefix:
            kwargs["kv_cache_prefix"] = kv_cache_prefix
-
+        num_devices = int(compiler_options.pop("num_devices", 1))


nit: any specific reason why we are popping the num_devices? similar to how we take the num_cores we can use get right?

quic-rishinr · 2026-08-06T09:21:40Z

            qaic_config=qaic_config,
+            prefill_only=prefill_only,
+            enable_chunking=enable_chunking,
+            num_cores=kwargs.get("num_cores", compiler_options.get("aic_num_cores", constants.DEFAULT_AIC_NUM_CORES)),


variables are not added to transfrom method, in the transform method is still fetching the num_cores from compiler args.

quic-rishinr · 2026-08-06T09:23:56Z

            self.hash_params["qaic_config"] = qaic_config
        self.hash_params["num_replicate_kv_heads"] = effective_num_replicate_kv_heads

+        num_cores = compiler_options.get("num_cores", compiler_options.get("aic_num_cores"))


num_cores, prefill_seq, prefill_only length is passed as arguments. Need not extrat it again.

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

ochougul requested review from anujgupt-github, quic-rishinr and vbaddi July 21, 2026 07:14

ochougul force-pushed the moe_rewrite branch from f9932d5 to c9c907e Compare July 21, 2026 07:31

ochougul commented Jul 21, 2026

View reviewed changes

abhishek-singh591 mentioned this pull request Jul 28, 2026

Gpt_oss prefill: avoid packed expert-axis slicing in chunked MoE to restore MXFP6 constant compression #1165

Closed

ochougul marked this pull request as draft August 4, 2026 16:40

quic-rishinr requested changes Aug 6, 2026

View reviewed changes

ochougul marked this pull request as ready for review August 7, 2026 09:01

ochougul added 16 commits August 9, 2026 17:56

Refactor MoE optimization transforms

57355b9

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Consolidate MoE adapter cleanup

51f9496

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Centralize MoE flavour execution

d2176f8

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Unify MoE transform execution paths

bf34778

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Centralize MoE transforms

d216523

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Refactor optimized MoE cleanup

23abe5b

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Apply optimized MoE during model transform

727b5b0

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Make optimized MoE transform explicit

400608c

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

added SimpleMoeTransform for keeping legacy default behaviour

85d2f1a

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Refactor MoE flavour config handling

43b611c

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Clean up unused export kwargs

940fea8

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Keep legacy MoE option rejection in transform

e21774f

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Inline export argument binding

7e79fcf

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Close module-level MoE flavour coverage gaps

54bfc6c

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Refactor QEff MoE block contracts

8bc4043

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Refactor MoE weight transform contract

726201f

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

ochougul added 6 commits August 9, 2026 17:59

removed qaic_config un-necessary checks

d4aab01

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Pack expert-parallel MoE weights during transform

8a361bd

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Materialize expert-parallel MoE weights

6a7f9b1

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Remove unused MoE packed chunk module size

a85a978

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

Refactor MoE expert-parallel layout

67daf9b

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

fixed RAM BLOAT issue

718b8e2

Signed-off-by: ochougul <ochougul@qti.qualcomm.com>

abhishek-singh591 force-pushed the moe_rewrite branch from 07007f5 to 718b8e2 Compare August 9, 2026 14:55

Navigation Menu

Uh oh!

Introducing OptimizedMoeTransform - #1190

Introducing OptimizedMoeTransform#1190
ochougul wants to merge 22 commits into
quic:mainfrom
ochougul:moe_rewrite

ochougul commented Jul 17, 2026

Uh oh!

ochougul Jul 21, 2026

Uh oh!

ochougul Jul 21, 2026

Uh oh!

ochougul Jul 21, 2026

Uh oh!

ochougul Jul 21, 2026

Uh oh!

ochougul Jul 21, 2026

Uh oh!

ochougul commented Jul 28, 2026

Uh oh!

quic-rishinr Aug 6, 2026

Uh oh!

quic-rishinr Aug 6, 2026

Uh oh!

quic-rishinr Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		order = torch.argsort(out, dim=1)
		last_positions = order[:, -state_len:]

Uh oh!

Conversation

ochougul commented Jul 17, 2026

Impact For New MoE Models

Future Scope

Tests

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ochougul commented Jul 28, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants