refactor: extract shared routed-expert FFN helper (byte-identical) - #1035
Conversation
Extract the four-step routed-expert FFN sequence (expert_gate_up -> silu(gate)*up -> [fmt=6 down-input rotation] -> down matmul) into one shared expert_ffn helper and replace the six inlined copies in moe() with calls to it. Single concern, no new knobs. The helper takes the expert's own tensors plus buffers/row-count (no moe() locals), so a distributed worker can call it instead of carrying a parallel copy. Input arrives already rotated by the caller via E8_XE (per-call cache, JustVugg#452); the fmt=6 down-input rotation lives inside the helper and is applied whenever d->fmt==6. Byte-identical on the oracle and on any single-format model (what a normal conversion produces). Four of the six copies previously omitted the down-rotation. Two of those -- the Vulkan device-lost fallbacks -- are provably safe: they reload only registry-resident experts, and vk_registry_fill admits only fmt 2/4/5, so d->fmt can never be 6 there. The other two (the Vulkan CPU-share path and the CUDA early-issued-take fallback) can in principle reach a fmt=6 expert in a mixed-format container (fmt is derived per-tensor by qt_resolve_fmt, with no uniformity enforced across experts at load time); routing them through the helper makes their down-rotation consistent with the canonical path, closing a latent omission rather than changing any single-format model's output.
|
Quick context for reviewers: this is a straight extraction — the four-step FFN body is moved verbatim into expert_ffn, six call sites now call it, and the fmt=6 down-input rotation is applied unconditionally inside the helper (a pure function of the down tensor's format). The only behavioral nuance is that two of the six copies previously skipped that rotation; both of those (Vulkan device-lost fallback, CUDA early-take fallback) either provably can't see fmt=6 or only could in a mixed-format container, so output is unchanged on any single-format model and the oracle is byte-identical. No new knobs, single concern. I'm also shipping a follow-up PR (#1036) that consumes this helper from a distributed worker — mentioned only so the "so a downstream worker can call it" framing doesn't read as speculative generality. |
2 participants
Summary
moe()inlines the four-step routed-expert FFN (expert_gate_up → silu(gate)*up → [fmt=6 down-input rotation] → down matmul) six separate times, and those six copies were written inconsistently: four omit the fmt=6 down-input rotation, two apply it. This extracts the sequence into oneexpert_ffnhelper and routes all six call sites through it.The helper takes the expert's own tensors plus buffers and a row count — no
moe()locals — so a downstream distributed worker can call the same helper instead of carrying a parallel copy. Input arrives already rotated by the caller viaE8_XE(the per-call cache is preserved, #452); the fmt=6 down-input rotation lives inside the helper and is applied wheneverd->fmt==6.Validation
make→ 0 warnings.make check→ green (C suiteALL PASS (0 failures), PythonRan 473 tests OK).gcc -fsyntax-only -DCOLI_CUDA -DCOLI_VULKAN -DCOLI_METAL -Wall -Wextra c/colibri.c→ 0 warnings (the#ifdef-gated call sites still compile).Compatibility
No new knobs, no behaviour change on any single-format model (the oracle
SNAP=./glm_tiny TF=1 ./colibri 64 16 16is byte-identical before and after). The helper applies the fmt=6 down-input rotation unconditionally — it is a pure function of the down tensor's format. Of the four copies that previously skipped it, the two Vulkan device-lost fallbacks provably cannot see fmt=6 (the Vulkan tier loads only fmt 2/4/5), and the other two (Vulkan CPU-share, CUDA early-issued-take fallback) can reach fmt=6 only in a mixed-format container, where this now matches the canonical path. Single-machine path unchanged.