[RDY] feat: add basic MoE support - #55
Open
JulienBalianSonos wants to merge 42 commits into
Open
Conversation
JulienBalianSonos
changed the title
JulienBalianSonos
changed the title
JulienBalianSonos
changed the title
JulienBalianSonos
force-pushed
the
feat/moe-ffn-export
branch
from
June 23, 2026 07:07
860f46f to
f9bba50
Compare
JulienBalianSonos
force-pushed
the
feat/moe-ffn-export
branch
5 times, most recently
from
July 10, 2026 09:00
ecfd12a to
933a04a
Compare
…was emitting [1,T,D] and breaking tract IO check)
…pert biases, interleaved gate-up, and clamped-swiglu activation
…ze_gates and reject the unsupported norm_topk_prob=False
…n top of tract_moe_ffn; add Mixtral + Qwen2 shared-expert end-to-end tests
…resh supported-arch docstring
…stale harness regen script (norm_topk_prob + seeded init + built-tract); guard shared-expert non-SiLU activation
…ffn; 0.23.0-0.23.3 already shipped without it), keep explicit-tract escape
…m_topk_prob=False) now exports instead of raising
…ng the Qwen fused-expert layout (input_linear/output_linear, router.layer)
…below lint complexity/statement thresholds
…nemo[asr] 2.7.2 needs setuptools>=79 / py3.13 numba, absent on the torch index)
…a resolves (pre-installed numpy 2.5 forced backtrack to unbuildable numba 0.53.1)
JulienBalianSonos
force-pushed
the
feat/moe-ffn-export
branch
from
July 21, 2026 16:25
61b621d to
aff303d
Compare
… multimodal guard
The GPT-OSS handler only remapped the experts implementation; it never
overrode build_forward_inputs, so it inherited the default handler's
single 4D causal mask.
That silently loses the sliding window. GPT-OSS alternates
sliding_attention and full_attention layers, and
masking_utils._preprocess_mask_arguments returns any 4D mask as-is, so
both create_causal_mask and create_sliding_window_causal_mask hit that
early exit. The model's mask mapping ends up holding the same unwindowed
tensor under both keys and every layer attends over the whole context.
The failure is invisible below the window, where the two masks agree,
and grows as the sequence extends past it, so short smoke tests do not
catch it. On gpt-oss-20b (window 128, 12 of 24 layers sliding) output
stays correct for roughly the first 128 positions, then decays into
repeated fragments; a prompt longer than the window is degenerate from
the first generated token.
Build both masks from token positions and pass them as the
{full_attention, sliding_attention} mapping the model accepts directly.
Masks use float arithmetic rather than boolean ops, for the same tracing
reason noted in gemma3_vl.
Models with no sliding layers keep the base handler's single mask.
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
1 participant
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.
Implementation of #16.
Compatible with the tract PR sonos/tract#2084 .
Waiting for more efficient 'tract' implementation of a kind of 'routed_matmul' that handle the combine (scatter-add weighted results) internally.
This PR is complete.
As of now we have in tract a good enough implementation, but still a bit suboptimal.
Hence PR is 'blocked' since no urgent "needs" identified.