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

[yurekami]

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.