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

[yurekami]

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.