MoonshotAI / MoonshotAI/MoonEP
Docs vs code: README weight contract omits the H/H' multiple-of-128 rule; test generator docstring contradicts its own seeding
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.1k
- Forks
- 134
- PR merge metrics
- No merged PRs in 30d
Description
Two places where a stated contract disagrees with the code. Both are small; grouping them since they're the same class of problem rather than filing separately.
Line numbers are master @ 0f385f0.
1. The README weight-tensor contract omits the H/H' multiple-of-128 requirement
README.md:45 states the framework-facing contract:
MoonEP's contract with a training or inference framework is one contiguous symmetric-memory weight tensor per expert projection, plus a planner-produced
cu_seqlens. The VM group GEMM consumes a single[E+B, H, H']weight tensor
and :51 restates it, calling contiguity out explicitly:
every layer holds one contiguous VMM range
[E+B, H, H'], identically laid out on every rank. Contiguity is a hard requirement
Neither mentions any constraint on H or H' beyond the shape, and :43's notation block just defines them as "hidden size" and "expert FFN intermediate size". But three kernels require both to be multiples of 128:
# moonep/combine.py:604 (ACC_THREADS = 128, combine.py:63)
assert ctx['H'] % CombineKernel.ACC_THREADS == 0, \
f"H must be a multiple of ACC_THREADS={CombineKernel.ACC_THREADS} " \
"(also covers the 16-B bulk-copy alignment requirement)"
# moonep/prefetch.py:353
assert H % PrefetchKernel.M_BLOCK == 0 and Hp % PrefetchKernel.N_BLOCK == 0, \
f"H and H' must be multiples of ({PrefetchKernel.M_BLOCK}, {PrefetchKernel.N_BLOCK}), got ({H}, {Hp})"
plus the equivalent in grad_reduce.py:497.
To be fair to the code, this is documented — just not where a framework integrator reading the contract would look. moonep/prefetch.py:11-12:
The initial tile shape is fixed at 128 x 128 bf16 elements. H and H' are
therefore required to be multiples of 128 for this first implementation.
So the ask is narrow: add the constraint to the README section that presents [E+B, H, H'] as the integration contract, since that is the document an integrator sizing their model against MoonEP will read. A concrete H like 2880 (a multiple of 8 but not of 128) satisfies everything the contract states and everything dispatch checks, and only fails later inside launch_prefetch/launch_combine.
Whether it is also worth validating H/H' early — at Buffer construction rather than at first kernel launch — is a maintainer call; I have not traced exactly which validations run where, so I am not claiming there is no early check, only that the README does not state the rule.
2. generate_topk_routing docstring contradicts the generator it uses
tests/generate_topk_routing.py:17-18 states:
seedseeds rank-shared state (the expert-logit distribution and the round-robin expert permutation);rankseeds the per-token draws.
The two generators are set up accordingly (:24-25):
g_shared = torch.Generator(device=dev).manual_seed(seed)
g_local = torch.Generator(device=dev).manual_seed(rank)
but the round-robin permutation — named in the docstring as rank-shared — is drawn from the rank-local generator (:32):
perm = torch.randperm(epn, device=dev, generator=g_local)
g_shared is consumed exactly once more in the file, at :36, inside the biased branch. So on the bias_ratio == 0.0 path nothing reads g_shared at all and the seed argument has no effect on the generated routing — which covers every routing="balanced" case in the kernel tests, since kernel_test_utils.py passes bias=0.0 for them.
The fix is one of two one-line changes depending on intent: :32 should use g_shared if the permutation is meant to be rank-shared as documented, or the docstring should stop claiming it is.
Worth noting the practical impact on test validity is small: perm is a bijection on [0, epn), so per-rank expert-load multisets are unchanged either way, and when S/R is not divisible by epn a rank-local permutation actually spreads the residual +1 counts across different global experts, which smooths load rather than degrading it. So this is a correctness-of-contract issue and a dead argument, not a broken baseline.
Verification
- All quotes and line numbers read from
master@0f385f0this session, clean working tree; confirmedg_sharedhas exactly two uses in the file and that:32is not one of them. - Not executed. No GPU here, so I have not run the test generator or any kernel test — the
bias=0.0routing for balanced cases is read fromkernel_test_utils.py, not observed.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Read README.md around lines 43-51 and compare its weight-tensor contract with the kernel assertions cited in moonep/combine.py, moonep/prefetch.py, and moonep/grad_reduce.py. Then inspect tests/generate_topk_routing.py and its uses in kernel_test_utils.py to resolve whether the permutation or docstring should change. Done means the README and generator documentation agree with the implemented contract; GPU tests are unavailable here.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- documentation, testing
- Issue type
- Documentation
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100