[RDY] feat: add basic MoE support - #55

Open

JulienBalianSonos wants to merge 42 commits into

mainfrom

feat/moe-ffn-export

Commits on Jul 21, 2026

feat: add basic MoE support
JulienBalianSonos
committed
e808199
tests: gate MoE tests to tract > 0.22.0 to avoid CI failures with official release
JulienBalianSonos
committed
f7ede9f
fix: add export script for moe test asset and update extractor
JulienBalianSonos
committed
b3861bf
fix: lint issues in export_moe_test_asset.py
JulienBalianSonos
committed
6264b9b
fix: lint export_moe_test_asset.py under updated ruff (post-rebase)
JulienBalianSonos
committed
75cc5e1
fix: skip rank-alignment for tract_moe_ffn so 2D MoE input stays 2D (was emitting [1,T,D] and breaking tract IO check)
JulienBalianSonos
committed
d614e5b
feat: export gpt-oss MoE (GptOssMLP) via tract_moe_ffn with router/expert biases, interleaved gate-up, and clamped-swiglu activation
JulienBalianSonos
committed
9783c02
test: end-to-end Qwen3-MoE export test; map norm_topk_prob to normalize_gates and reject the unsupported norm_topk_prob=False
JulienBalianSonos
committed
8eb0468
feat: shared-expert MoE (Qwen2/Qwen3.5) decomposed as NNEF subgraph on top of tract_moe_ffn; add Mixtral + Qwen2 shared-expert end-to-end tests
JulienBalianSonos
committed
b07a6ae
feat: register OLMoE MoE (reuses Qwen adapter) + end-to-end test; refresh supported-arch docstring
JulienBalianSonos
committed
455ac5c
fix: skip MoE tests on op-less official tract (>0.23.0 gate); repair stale harness regen script (norm_topk_prob + seeded init + built-tract); guard shared-expert non-SiLU activation
JulienBalianSonos
committed
b8d62cc
test: gate MoE tests to tract >=0.23.4 (first release with tract_moe_ffn; 0.23.0-0.23.3 already shipped without it), keep explicit-tract escape
JulienBalianSonos
committed
ebdf7f9
style: wrap long comment in MoE test gate
JulienBalianSonos
committed
bc6f691
feat: gate-mode enum replaces normalize_gates; OLMoE softmax_all (norm_topk_prob=False) now exports instead of raising
JulienBalianSonos
committed
5e90000
feat: Granite MoE adapter (GraniteMoeMoE / GraniteMoeSharedMoE) reusing the Qwen fused-expert layout (input_linear/output_linear, router.layer)
JulienBalianSonos
committed
2802902
fix(moe): extract shared-expert emission to cut _convert_moe_to_nnef below lint complexity/statement thresholds
JulienBalianSonos
committed
6930381
fix(ci,nemo): install torch from CPU index then nemo deps from PyPI (nemo[asr] 2.7.2 needs setuptools>=79 / py3.13 numba, absent on the torch index)
JulienBalianSonos
committed
a9d2d7e
fix(ci,nemo): pin numpy<2.5 for nemo install so a py3.13-capable numba resolves (pre-installed numpy 2.5 forced backtrack to unbuildable numba 0.53.1)
JulienBalianSonos
committed
39b4d07
fix(llm): keep tract attention scores in f32
JulienBalianSonos
committed
ff3925e
fix(moe): export split experts as Q40
JulienBalianSonos
committed
b72c227
chore: clean up moe export CI
JulienBalianSonos
committed
a76fd0a
fix(moe): emit linear Q40 expert layout
JulienBalianSonos
committed
3d73c5c
test(moe): decouple expert layout from q40
JulienBalianSonos
committed
7276712
fix(moe): materialize opaque expert views before export
JulienBalianSonos
committed
b5afd13
feat: allow reified sdpa f16 inputs with f32 accum
JulienBalianSonos
committed
96fd430
feat(llm): forward reified sdpa input upcast policy
JulienBalianSonos
committed
7a50b50
fix(moe): support current Granite expert layout
JulienBalianSonos
committed
e7fc9aa
fix(opaque): trace float refs with fake device tensors
JulienBalianSonos
committed
b432305
Reduce Q40 export memory during NNEF writes
JulienBalianSonos
committed
d89fca3
Fix implicit cast shape for broadcast operands
JulienBalianSonos
committed
b19549e
fix(div): restore graph-output naming and keep integer rounding-division in float
JulienBalianSonos
committed
3205aac
fix(opaque): let NamedTensor compose with fake-traced float weights
JulienBalianSonos
committed
b55f29f
chore: bump torch-to-nnef to 0.24.3 in example uv.lock files
JulienBalianSonos
committed
fc6b620
fix(opaque): construct trace FakeTensorMode compatibly on torch < 2.0
JulienBalianSonos
committed
3543376
fix(opaque): unwrap NamedTensor operands so quantized weights compose with fake trace
JulienBalianSonos
committed
bd8c15a
fix(llm): avoid grouped_mm during gpt-oss export trace
JulienBalianSonos
committed
a4b64fb
fix(moe): ignore broken optional transformer imports
JulienBalianSonos
committed
aff303d

Commits on Jul 24, 2026

feat: moe custome extractor
JulienBalianSonos
committed
0cf5294

Commits on Jul 25, 2026

Commits on Jul 28, 2026

fix(llm): give GPT-OSS sliding layers their attention window
JulienBalianSonos
committed
e610817