Buffer() rejects all odd EP sizes including R=1: plan multimem publish has no scalar tail

Summary

Buffer() raises AssertionError for every odd world size whose epn is not a multiple of 4 — including single-rank (R=1). Seven of the cases already declared in tests/test_planning.py cannot construct a Buffer at R=1, and seven at R=3.

The assert is not protecting an invariant of the algorithm. It exists because one publish loop vectorizes without a scalar tail.

Root cause

moonep/api.py:302-305:

broadcast_elems = 3 * E * R
assert broadcast_elems % 4 == 0, (
    f"broadcast_elems ({broadcast_elems}) must be divisible by 4"
)

Since E = R * epn, broadcast_elems = 3 * R**2 * epn, which is divisible by 4 iff R is even, or epn % 4 == 0. So R=1, epn=1 gives 3 and construction aborts before any CUDA work.

The requirement comes from the plan broadcast in moonep/planning.py:955-960:

nb = 3 * E * R; nvec = nb // 4
for i in cutlass.range(pid * num_threads + tid, nvec, num_sms * num_threads):
    a0 = meta[PB + i * 4 + 0]; a1 = meta[PB + i * 4 + 1]
    a2 = meta[PB + i * 4 + 2]; a3 = meta[PB + i * 4 + 3]
    addr = (mc.iterator + (PLAN_OFF + i * 4)).toint()
    multimem_st_v4(addr.ir_value(), a0, a1, a2, a3)

nvec = nb // 4 truncates, so the loop publishes 4 * (nb // 4) int32s and the trailing nb % 4 are never stored. The next statement (line 961) is unrelated — there is no tail loop. The host-side assert compensates by forbidding every shape where a tail would exist.

Given the region layout at planning.py:546-549:

ALLOC_SUB = 0
TPE_SUB   = E * R
EOFF_SUB  = 2 * E * R
CU_SUB    = 3 * E * R

the broadcast covers exactly ALLOC_SUB | TPE_SUB | EOFF_SUB, so the untransferred tail would be the last nb % 4 entries of EOFF_SUB (expert offsets).

The same file already implements the correct pattern. copy_v4_remote at planning.py:284-302 does scalar head, vector body, and scalar tail, and its own comment says so:

# src[n] -> dst[dst_off:]: scalar head pads to 16B, int4 body does one
# 128bit store (transactions /4), scalar tail.

It is called with a non-multiple-of-4 length (n=E) at planning.py:603 and :608.

Why odd R is in scope

Nothing else in the library restricts R to even values:

• planning.py:962-972 has an explicit single-rank branch — if cutlass.const_expr(R > 1): ... else: self.run_c1(...).
• inter_rank_sync.py:87-88 asserts only R > 0 and R <= 1024.
• api.py:261 constrains only E % R == 0.
• planning.py:114-116 log2_r(R) = max(R.bit_length(), 1) is a fixed-trip-count search valid for arbitrary R, not powers of two.
• tests/kernel_test_utils.py defaults min_R: int = 1.

Affected cases already in the repo

Computed from tests/test_planning.py PLANNING_CASES, after applying min_R/max_R gating:

World size	Cases that run	Cases that abort in Buffer()
R=1	14	7 — tiny_s1_k1, no_padding, non_power_mild_bias, small_balanced_with_prefetch, near_degenerate_bias, heavy_bias, experts_gt_block_size
R=3	16	7 — the above minus experts_gt_block_size, plus tiny_biased_with_prefetch
R=2,4,8	—	0

Reproduce the arithmetic without a GPU:

# E = R * epn ; broadcast_elems = 3 * E * R
for R, epn in [(1, 1), (1, 2), (1, 3), (3, 3), (2, 3), (8, 1)]:
    print(R, epn, (3 * (R * epn) * R) % 4 == 0)
# 1 1 False / 1 2 False / 1 3 False / 3 3 False / 2 3 True / 8 1 True

Suggested fix

Mirror copy_v4_remote: keep the multimem_st_v4 body over nvec groups and add a scalar multimem store loop for the remaining nb - nvec * 4 elements, then drop the broadcast_elems % 4 assert in api.py.

A question rather than a claim

For R=1 the publish looks like a self-copy — the source is meta[PB + ...] with PB = PLAN_OFF (planning.py:553) and the readback is plo = rank * ms + PLAN_OFF (planning.py:986), which is the same address at rank 0 — so dropping the tail there appears harmless and the assert purely spurious.

For odd R > 1 the trailing EOFF_SUB entries would not be published to peer chunks by this loop. I could not determine by reading whether another path republishes them. If it does, the fix is only to relax the assert; if not, the tail store is required. Could you confirm which?

Testing

• Arithmetic and affected-case table verified statically against tests/test_planning.py.
• All cited line numbers verified against master @ 0f385f0.
• Not executed. I do not have access to a multi-GPU NVLink machine, so I could not build moonep._C or run any kernel test.

Happy to send a PR with the tail loop if you confirm the intended behavior for odd R > 1.

Start with the broadcast assertion in moonep/api.py:302-305 and the vectorized publish loop in moonep/planning.py:955-960. Compare that loop with copy_v4_remote at planning.py:284-302, then inspect the affected cases in tests/test_planning.py and the R=1 branch at planning.py:962-972. Done means odd-R cases no longer abort and the full plan region, including any required expert-offset tail, is published correctly; GPU execution is needed to verify the kernel behavior.