MoonshotAI / MoonshotAI/MoonEP

Docs vs code: README weight contract omits the H/H' multiple-of-128 rule; test generator docstring contradicts its own seeding

Open Beginner friendly
#19 1 comment 0 reactions 0 assignees View on GitHub

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:

seed seeds rank-shared state (the expert-logit distribution and the round-robin expert permutation); rank seeds 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 @ 0f385f0 this session, clean working tree; confirmed g_shared has exactly two uses in the file and that :32 is not one of them.
  • Not executed. No GPU here, so I have not run the test generator or any kernel test — the bias=0.0 routing for balanced cases is read from kernel_test_utils.py, not observed.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.