Plan-reuse dispatch launches 96 threads, capping R at 96 while planning declares R <= 128

The dispatch kernel launches a different thread count depending on whether it builds the dedup map, and the smaller of the two silently caps the supported rank count at 96 — while planning.py declares and enforces R <= 128.

Line numbers are master @ 0f385f0.

The two thread counts

moonep/dispatch.py:113-115:

        self.num_threads = (
            96 + 32 * DEDUP_BUILDER_WARPS if build_dedup_map else 96
        )

with DEDUP_BUILDER_WARPS = 4 (moonep/constants.py:17), so:

• fresh plan (build_dedup_map=True) → 96 + 128 = 224 threads
• reused plan (build_dedup_map=False) → 96 threads

That value is forwarded to cross_rank_barrier, which requires one thread per rank — moonep/_common.py:304:

    assert num_threads >= num_ranks, (
        "cross_rank_barrier requires blockDim.x >= num_ranks: "
        f"num_threads={num_threads}, num_ranks={num_ranks}"
    )

The assert is load-bearing rather than defensive: only threads with tid < num_ranks signal peers, so a thread count below R would silently fail to signal some ranks.

The conflict

moonep/planning.py:1282 declares the supported ceiling:

    assert R <= 128, (

Nothing anywhere caps R at 96. So on the plan-reuse path any R in [97, 128] fails the cross_rank_barrier assert at cute.compile trace time, while the same configuration works on the fresh-plan path in the same run.

The reuse path is reachable from the public API — api.py:742-752 sets planning_args = None whenever a plan is passed, and api.py:618 derives build_dedup_map = planning_args is not None. benchmarks/bench_comm.py:244 exercises exactly this for the backward re-dispatch.

I checked the other cross_rank_barrier callers and none is below 96: combine 192/224 (combine.py:91-92), grad_reduce 160 (grad_reduce.py:53), planning 512, inter_rank_sync max(32, ceil(R/32)*32) with R <= 1024 (inter_rank_sync.py:86-89). dispatch_epilogue.py:63's num_threads = 64 is not a counterexample — it never calls cross_rank_barrier.

Severity

Low, deliberately. The failure is a loud compile-time AssertionError, not corruption, and no shipping NVLink domain exceeds 72 ranks, so nothing reachable today trips it. It is a latent inconsistency between two paths of the same kernel rather than a live break.

Filing it because the two declared ceilings disagree, and if R > 96 ever becomes real the failure will appear only on the plan-reuse path, which is the harder one to notice.

Suggested fix

Either raise the reuse-path launch to at least the planning ceiling, or lower the declared ceiling and document 96 as the real limit. If the 96 is a deliberate occupancy choice for the reuse path, a comment saying so at :113 would prevent the next reader concluding it's a typo.

One correction to something that might come up: constants.py:7 defines RANK_BITS = 7 (2^7 = 128), which looks like it documents the 128 ceiling — but grep -rn RANK_BITS returns only that definition, so it is dead and the ceiling rests solely on the planning.py:1282 assert.

Verification

• All quotes and line numbers read from master @ 0f385f0 this session, clean working tree; swept every cross_rank_barrier caller for a lower ceiling.
• Not executed. No GPU/NVLink here, so I have not run cute.compile at any R, let alone R > 96. The trace-time failure is inferred from the assert operating on two Constexpr ints inside a @cute.jit, not observed.

Start with moonep/dispatch.py:113-115 and moonep/planning.py:1282, then trace plan reuse through api.py:618 and api.py:742-752 and inspect the barrier assertion at moonep/_common.py:304. Determine which rank ceiling is intended and make the launch and validation behavior consistent. Done means fresh and reused plans agree for the supported rank range, with verification covering the affected path.