[Fix] Publish the plan tail for non-vector-aligned world sizes

[0z5a]

Incremental review: Files changed against master.

Refs #16.

What

PlanningKernel broadcasts ALLOC | TPE | EOFF (3 * E * R int32) from rank 0 into every peer chunk with a vectorized multimem body, which stored only 4 * (3 * E * R // 4) elements. A world size whose broadcast length is not a multiple of 4 therefore lost the trailing entries, and api.py rejected exactly those shapes with assert broadcast_elems % 4 == 0 -- R = 1 included.

This stores the trailing 3 * E * R % 4 elements with a scalar multimem.st.b32 and drops the assert.

Why the tail store is required (the open question in #16)

Only rank 0 publishes, and it publishes into peer chunks. The two sides of the plan disagree about where it is read from:

sub-region	written by	read by
ALLOC | TPE | EOFF (first 3 * E * R int32)	rank 0, into its own chunk (planning.py:649/721/815), then replicated by this publish	every rank, from its own chunk (plo = rank * ms + PLAN_OFF: tpe_cumsum_view, alloc_cumsum_view, expert_off_view)
CU | ZFR | ETC | STATS	rank 0 only	every rank, straight out of rank 0's chunk (cu_src_begin = PB + CU_SUB + rank * pd_group_count + ..., remote_stats)

No other path republishes the peer copies, so the elements this loop drops stay stale in every peer chunk -- and they are the last expert offsets, i.e. the top experts' destination offsets. The scalar tail is therefore required; relaxing the assert alone would be wrong for odd R > 1.

At R = 1 the multicast view maps rank 0's own chunk, which already holds the plan, so the publish is a pure self-copy and is now skipped under cutlass.const_expr(R > 1).

Tests

tests/test_planning.py gains an index model of the publish:

• every element of [0, nb) is written exactly once and the vector body never overlaps the scalar region, for nb in [1, 256) and workers in {1, 8, 32, 128, 256, 512};
• the tail length nb % 4 stays reachable at R = 1 and R = 3 for the existing PLANNING_CASES;
• negative control: dropping the scalar tail loses exactly the last nb % 4 elements.

The distributed planning/dispatch/combine/e2e tests that already exist for those shapes now construct a Buffer instead of aborting in api.py.

Verification performed

• host lane pytest tests/test_planning.py -q: 3 passed, 18 skipped;
• Buffer construction at R = 1 for the shapes the old assert rejected (nb = 3, 9, 6);
• planning kernel lowers for sm_90 at nb = 3, 36, 54;
• trace-level store count: nb % 4 == 0 emits no scalar store, nb = 54 emits one, R = 1 emits none.

Limitation: the machine used for this work has no SM90 device (the kernels refuse sm_89 at trace time) and no multicast support, so the multimem tail was compiled but not executed. The distributed tests above are what exercise it end to end on supported hardware.

Speed comparison

The tail adds work only where 3 * E * R % 4 != 0; every other shape compiles to the pre-patch instruction stream (byte-identical SASS), so the table below is the whole cost of the fix.

group	samples	max abs delta	median abs delta
tail = 0 (instruction-identical, null control)	10	51.6 ns	31.5 ns
tail = 1	6	29.6 ns	21.1 ns
tail = 2	4	8.8 ns	7.7 ns
tail = 3	6	31.7 ns	27.7 ns

Method: 9 paired APPA blocks per shape, 20000 launches per arm, grid 1 and 32 blocks x 512 threads, shapes from R in {1,2,3,4,8} x epn in {1,2,3,4,5,7}. One launch is ~1.9 us, so the tail is below the harness resolution: under 1.7% and with no consistent sign. The tail-free group is the null control -- its true delta is exactly zero, and it shows the largest spread, so none of the residual is attributable to this patch.

Two further points from the same run:

• SASS: a tail-free nb produces an instruction-identical kernel (nb is a cutlass.const_expr, so the guarded loop is compiled away); a tail-carrying nb adds exactly one load and one store, the rest being loop control.
• Availability: at R = 1 and R = 3, epn in {1,2,3} the old code raised AssertionError: broadcast_elems (3/6/9/27/54/81) must be divisible by 4 before any CUDA work; the patched tree reaches the allocator and stops only at the hardware capability check.

Caveat: this is an instruction-shape proxy measured on an sm_89 host -- st.global.v4.s32 / st.global.s32 in place of multimem.st, which ptxas rejects below sm_90. It measures instruction and store cost, not multicast fan-out semantics, and it is not a substitute for the end-to-end numbers the distributed lane will produce on supported hardware.

Test and benchmark sources are retained locally; the validation results below refer to those local files.

[0z5a]

Draft: the fix is verified down to the compiled kernel, but the multimem tail has not been executed on hardware yet.

Verified so far:

• the removed host gate reproduces the issue exactly -- R = 1 rejects 7 of the PLANNING_CASES, R = 3 rejects 7;
• the publish body lowers for sm_90 at nb = 3 / 36 / 54;
• a trace-level count of the publish stores shows the scalar store is emitted exactly when nb % 4 != 0, and the whole publish is gone at R = 1.

Not yet run: the planning/dispatch/combine tests at odd R, and the negative control (assert relaxed, tail store absent). The machine used for this work has no SM90 device and no multicast support, so the native run is pending; I will mark this ready once it lands.