Skip to content

Extract a shared routed-expert FFN helper (pure refactor, byte-identical) #9

Description

@lineape

What to build

Extract the four-step routed-expert FFN (expert_gate_up → siluf → e8_rot_rows → matmul_qt) into ONE shared helper, and replace the ~6 inlined copies in moe() with calls. Pure refactor: byte-identical output (token-exact oracle green before AND after), one file, no new knobs.

Two things this forces, worth doing now:

  1. Resolve an existing inconsistency. Two of the inlined copies apply the fmt=6 down-rotation (if(d.fmt==6) e8_rot_rows); four don't (CPU-resident, Vulkan dev0/dev2 device-lost fallbacks, CUDA-take fallback). Determine whether the non-rotating fallback paths can ever see a fmt=6 expert, and settle the canonical behavior. This is either harmless (those paths never reach fmt=6) or a latent bug — decide and document it.
  2. Shape the helper for reuse. It must operate on (gate/up/down tensors + input rows + count), not on moe()'s locals, so the distributed worker can call the same helper later instead of carrying a parallel copy. This is the seam that removes the maintainer's "re-implemented core kernel" objection at the root.

This change is a self-contained, independently-submittable refactor: single concern, byte-identical, no new knobs — the exact shape the maintainer has praised — and should go upstream as its own PR, separate from the distributed-workers PR.

Acceptance criteria

  • One shared FFN helper; all moe() call sites route through it (no remaining inlined four-step sequence).
  • Byte-identical: token-exact oracle (~30-32/32 TF + 20/20 greedy) identical before and after.
  • The fmt=6 down-rotation inconsistency is resolved and documented (which paths can see fmt=6; what's canonical).
  • make check green, 0 warnings.
  • Helper signature is worker-reusable (takes gate/up/down tensors + input rows + count, not moe() locals).
  • Submitted upstream as its own byte-identical refactor PR (single concern, dev target), independent of the distributed-workers PR.

Environment: codegraph-indexed — run codegraph explore first; refs in /home/eli/dev/_reference-docs-and-repos/colibri.

Metadata

Metadata

Assignees

No one assigned

    Labels

    effort/highNeeds frontier — new modules/core-kernel/complex state machinesready-for-agentFully specified, AFK-ready

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions