Add persistent dense-expert MoE path via spyre::moe_ffn - #2
Draft
ani300 wants to merge 14 commits into
Draft
Conversation
Add a fourth, opt-in Gemma4 MoE FFN formulation (_MOE_PERSISTENT_ONDEVICE) alongside the existing chunked / loop-on-topk paths. The router runs once, then the dense all-expert value path lowers through torch-spyre's spyre::moe_ffn semantic op as a single counted device loop: the activation and output accumulator stay resident in LX while the gate/up/down weights and the routing scalar advance one expert bank per trip. Routing is materialized as a full 64-wide broadcast stick per token/expert scalar (relu(unsqueeze.expand(-1,-1,64)) @ I64, then lane-0 view), which is what the persistent planner requires (routing stride (E*64,64,1)); a bare unsqueeze produced a degenerate size-1 stick and wrong output. Weights are stored with contiguous [K,E,N] backing and exposed as logical expert-major views. The three device modes are made mutually exclusive with an enabled_modes <= 1 guard. Requires the torch-spyre persistent-expert compiler stack (PR#3818); the existing PR293 paths are unchanged when the flag is disabled. Validated prefill-only A/B on gemma-4-26B-A4B-it (T=512, K=4): both persistent and chunked emit the same token; persistent warm prefill 2.84s vs chunked 13.32s (~4.7x faster). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
Prepare the sparse (loop-on-topk) and chunked MoE paths for a batch of torch-spyre compiler fixes, per moe-implementation-notes-aug2026.md. * Loop-on-topk region kept in [T,K]: consume topk's index in its native [T,K] shape with no flatten to N=T*K. Gathers and batched matmuls run in [T,K,...] form; the [T,K]->[N] flatten for the eager index_add scatter-combine moves to the host glue in _moe_ffn_loop. * B5-B8 index address-prep: the fp16 topk index is replicated over a dummy 64-lane stick (B5), restickified onto the stick (B6), widened fp16 -> fp32 (B7) and sliced to 32 elements (B8) to produce the idx_addr [T,K,32] fp32 the backend's idx2Addr consumes. idx_addr feeds all three indirect consumers (per_expert_scale, gate_up_dev, down_dev); idx2Addr itself is inserted by the backend when it lowers the index_selects. * index_mask in the chunked router: replace the torch.where(probs >= kth, probs, 0) expert threshold with torch.ops.spyre.index_mask(probs, kth), the device op form. Shared by the chunked and persistent routers. * Lift the K==4 bring-up pin: torch-spyre #3782 raised the topk ceiling from 4 to 128, so use the checkpoint's real top_k_experts (8 for gemma-4-26B-A4B) and assert only K <= 128. These paths cannot be traced end-to-end until the compiler batch lands (idx2Addr accepting fp32 input, the indirect-access weight matmul, and the index_mask op); the adapter side is written to be ready for them. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
Consolidate the four MoE FFN formulations (persistent/chunked/split/loop) to two, dispatched by seq_len: persistent for prefill (seq_len>1), loop-on- topk for decode (seq_len==1). Delete split + chunked. Rework decode combine to an on-device K-axis reduction; wire index_mask ahead of the op landing. Whole-layer compile deferred to a follow-up (KV scatter already lowers on device when compiled). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
Six-task plan implementing the two-method restructure spec: delete split + chunked, dispatch by seq_len, move the decode combine on-device (K-axis sum), swap the router to spyre::index_mask last. Staged for a green prefill checkpoint before the tree depends on unlanded ops. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
Remove the host/device-split MoE path (_moe_ffn_split and its grouped-GEMM helpers, device-gather/expert-region compiles, and the host-resident expert stacks). Superseded by the persistent (prefill) and loop-on-topk (decode) methods. Working-tree checkpoint; not pushed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
Remove the all-device chunked mask-reduce MoE path (_moe_ffn_chunked, _moe_expert_chunk, the standalone router compile, and the per-chunk offset-0 expert-weight materialization) plus its _MOE_GEMM_4B/_MOE_EC globals. The _MOE_CHUNKED_ONDEVICE flag is left in place for Task 3 to remove with the other mode selectors. Working-tree checkpoint; not pushed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
…set materialization) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
…dual weight materialization) Commits the prior phase-dispatch plan's Task 3+4 Step-1 edits as the starting base for the single-shared-layout plan: - forward dispatches persistent (prefill, seq_len>1) vs loop (decode) by seq_len - collapsed the three _MOE_*_ONDEVICE mode-flag globals + mutual-excl assert NOTE: this commit still carries the dual weight-set materialization in prepare_for_spyre (both persistent + loop device tensors), which OOMs the card (0x340f). That block is REPLACED by the single-shared-layout plan's Task 3; committing it first keeps each task's review diff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
… paths Replace the dual expert-weight materialization (which OOM'd the card at ~85 GiB) with ONE shared device layout per weight, read by both the prefill persistent hint-body and the decode loop-on-topk gather. De-fuse gate_up into separate gate/up. Lay each expert weight out E-outermost with the free dim on the stick via torch-spyre dma_moe_expert_weight_to_spyre -- simultaneously the gather-source layout (expert dim outermost) and the matmul weight-operand layout (sticked on free dim). Restore the hint-body persistent matmuls, dropping the absent spyre::moe_ffn op. One weight set is ~42.5 GiB/30 layers, fits the card. Validated at load: the 26B MoE model loads fully with ONE weight set and NO 0x340f FlexAllocator OutOfMemory (75.8s); the expert-weight device layout is confirmed [E=128, H=2816, M//64=11, 64] (E-outermost + 64 stick), proving the shared layout applied. This retires the dual-materialization OOM the refactor targeted. Full end-to-end prefill trace is deferred behind two pre-existing, out-of-scope blockers, neither introduced by this change: (A) torch.ops.spyre.index_mask, called by the router since before this work, is not registered in the current torch-spyre checkout; (B) dma_moe_expert_weight_to_spyre is now the first device touch during load and torch-spyre's context is created lazily, so a device warm-up is needed until spyre_empty_with_layout self-triggers it. Both are upstream (missing op + context-init ordering), tracked separately. Working-tree checkpoint; not pushed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
The module docstring still described the retired four-mode flag scheme (_MOE_PERSISTENT_ONDEVICE lowering through spyre::moe_ffn, _MOE_CHUNKED_ONDEVICE, _MOE_LOOP_ON_TOPK, default _moe_ffn_split). The live architecture is two paths selected per-forward by seq_len: _moe_ffn_persistent (prefill, hinted plain-matmul dense body, no custom op) and _moe_ffn_loop (decode, on-device gather), both reading ONE shared expert-weight layout produced by dma_moe_expert_weight_to_spyre. Rewrite the docstring to match; no code change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
torch.ops.spyre.index_mask is not yet registered in torch-spyre, so the persistent-prefill router (_moe_route_padded) aborted at trace time with '_OpNamespace spyre has no attribute index_mask'. Fall back to the equivalent torch.where(probs >= kth, probs, 0) threshold form, which traces and lowers on-device. Switch back to index_mask once the op lands. Validated prefill-only on-card (gemma-4-26B-A4B, 512-token prefill): load 77.9s (single weight set, no 0x340f OOM), cold gen 163.8s -> ' the', warm gen 6.8s -> ' the'. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
The loop-on-topk decode combine summed each token's K expert-row outputs via a host index_add (row_out.cpu().index_add over token_of_row[t,k]==t), forcing a device->host->device round-trip every decode step. Move the combine on device by reducing over the K axis inside the compiled region: row_out.sum(dim=1) [T,K,H] -> [T,H], the same reduction the host index_add emulated (all K rows of token t scatter to slot t == a per-token sum over K). This mirrors the persistent path's proven (down_out*route).sum device reduction, and drops the now-dead token_of_row/token_ids plumbing. Not a swap to the new on-device spyre_index_add (torch-spyre #3753): that op is a gather+add+overwrite read-modify-write with a no-duplicate-index precondition, and the K rows per token are all the same index t, which it would silently get wrong. The K-axis reduction has no such hazard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
…ckify
Express the decode-path topk-index address prep (moe-implementation-notes
B5-B8, Path B) in plain traceable ops so the layout pass inserts the
restickify itself, instead of calling torch.ops.spyre.restickify directly
(pass-inserted-only op, no fake impl -> untraceable from user code).
B5 identity idx[T,K] -> [T,K,64] expand a dummy 64-lane axis
B6 restickify [T,K,64] .contiguous() -> row-major relayout
moves stick dim onto the size-64 axis
B7 df16tofp32 [T,K,64] fp32 widen for address arithmetic
B8 slice -> [T,K,32] fp32 one fp32 stick (32 elems)
post-B8 -> int32 [T,K] integer gather index
Decode leg compiles through layout propagation (B5-B8 restickify + the
decode-shape gathers trace clean); aborts upstream at the batched
[T,K,1,*] matmul layout: propagate_layouts _matmul_layouts ->
find_matmul_generated_var "expected exactly 1 generated variable, got
{d0,d1}" (same family as E-tiling #3888 -- both T and K batch axes survive
into the matmul generated-var set). Awaiting upstream compiler fix.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
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.
Summary
Adds a fourth, opt-in Gemma4 MoE FFN formulation —
_MOE_PERSISTENT_ONDEVICE— on top of the class-structure rework in the base branch (gemma4-moe-design). The router runs once, then the dense all-expert value path lowers through torch-spyre'sspyre::moe_ffnsemantic op as a single counted device loop: the activation and output accumulator stay resident in LX while the gate/up/down weights and routing scalar advance one expert bank per trip.This is a second draft PR stacked on the first (
gemma4-moe-design); the diff here is only the persistent-path delta.Key details
[token, expert](relu(unsqueeze.expand(-1,-1,64)) @ I64, then a lane-0 view). This is what the persistent planner requires (routing stride(E*64, 64, 1)); a bareunsqueeze(-1)gives a degenerate size-1 stick and produced wrong output.[K,E,N]backing, exposed as logical expert-major views tomoe_ffn.enabled_modes <= 1assert at load.Validation
Prefill-only A/B on
gemma-4-26B-A4B-it(T=512, K=4), same torch-spyre head (PR#38182262ca13):' the'✓' the'✓Both modes emit the same token; persistent is ~4.7× faster warm and ~1.8× faster to compile.
Note
Draft: depends on torch-spyre PR#3818 (persistent expert stack) being present.
🤖 Generated with Claude Code