mega-MoE prefill: cycle-chunked dispatch pull (ring < num_min)

[ipiszy]

Enable prefill (large seqlen) for the packed-FP4 mega-MoE kernel by chunking the dispatch pull into num_cycles rounds, so the per-rank token ring can be ratiomax_seq_len instead of num_ranksmax_seq_len (the num_min that forced a seqlen cliff ~98K/8 ranks, ~49K/16 ranks).

Design: the cycle loop lives in the pull only. The GEMM roles run the ORIGINAL single-cycle for_each_block, gated by the ring full/empty counts (A-load waits for full_count per block; pull waits for empty_count per cycle). The pull cycle loop (cycle_end_pool bound) + a cross-CTA comm::grid_sync at each cycle boundary keeps the global ring counts consistent across CTAs. For multi-rank, ring counts are per-rank so no cross-rank cycle barrier is needed (only the num_cycles all-reduce ensures all ranks run the same cycle count).

A per-role GEMM cycle loop + shared mbarrier does NOT work: the 5 GEMM roles are pipelined (A-load ahead of epilogue) and reach the cycle boundary at different cycles, so a single counted-arrival mbarrier mixes arrivals and desyncs. Abandoned after printf isolation.

Root cause of the prior cycle-1+ deadlock: stored_rank_count was only reloaded when the expert changed within a for-loop iteration. CTAs that pulled 0 tokens in a cycle (broke at the cycle bound before loading) carried stale/zero counts into the next cycle -> num_active_ranks=0 -> the round-robin rank-selection while(true) spun forever. Fixed by tracking counts_loaded_expert and reloading whenever current_expert_idx != counts_loaded_expert.

Verification: baseline num_cycles=1 mxfp4 0.00075 / nvfp4 0.00058 (unchanged); 1rank all pass; prefill 1-rank ratio=0.5 num_cycles=4 diff=0.00079, EP2 ratio=1.3 num_cycles=2 diff=0.00079, EP8 ratio=1.3 num_cycles=3 diff=0.00064; bench EP8 ratio 1.3 227us vs baseline 193us (~18% slower). Full design + debugging notes in doc/chunking.md.

Remaining (follow-up): mirror to sm100_fp8_fp4_mega_moe.cuh; right-size the SF ring; combine_token_buffer chunking for very large seqlen (out of scope, kept as-is).

Generated with Claude Code

---

Note

High Risk
Changes distributed mega-MoE dispatch synchronization, ring sizing, and CUDA kernel control flow; regressions would show as deadlocks or wrong MoE outputs on multi-GPU prefill, though baseline num_cycles==1 paths are intended to be preserved.

Overview
Enables large prefill for packed-FP4 mega-MoE by allowing a per-rank token ring smaller than the old num_ranks × max_seq_len minimum, trading symmetric-buffer memory for optional multi-round dispatch.

Host API: get_symm_buffer_for_mega_moe gains chunk_ratio (None = prior behavior, float = ceil(ratio × tokens) ring, 'auto' = ring sized to max_tokens × topk for no-wrap baseline-like latency with less memory). Chunked rings may be below num_min; sizing only clamps to the upper ring limit.

Kernel / runtime: When ring < num_min, the dispatch pull runs num_cycles = ceil(total_pool_blocks / ring_blocks) with pool-block bounds, cross-CTA grid_sync each cycle, cross-rank max num_cycles via a new workspace slot, and reload of per-rank expert counts via counts_loaded_expert (fixes hang when a CTA skipped loading counts). GEMM roles stay on the original single-cycle for_each_block, synchronized by ring full/empty counts—not a per-role cycle barrier. Wave heuristics force one expert per wave when chunking; buffer assert no longer requires ring >= num_min.

Tests & benchmarks: Prefill scopes sweep chunk_ratio, capture kernel num_cycles printf, and assert correctness + cycle behavior; bench_packed_fp4 adds a multi-GPU prefill sweep (latency, sym-buffer size, peak memory). Design notes in doc/chunking.md.

^{Reviewed by Cursor Bugbot for commit e0b66df. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Debug printf left in kernel

Medium Severity

An unconditional printf on SM0 warp0 runs on every mega-MoE kernel launch to print rank, num_cycles, and ring capacity. That is leftover instrumentation, not gated by a debug flag, and adds avoidable overhead and device log noise in production runs.

 

^{Reviewed by Cursor Bugbot for commit ae9031d. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 93ff23f. Configure here.}

[cursor]

Printf logs pre-allreduce cycles

Medium Severity

The kernel printf for num_cycles runs before the cross-rank max all-reduce updates num_cycles, so logged and test-captured values reflect per-rank local cycle counts, not the synchronized count the pull loop actually executes under EP.

 

^{Reviewed by Cursor Bugbot for commit 93ff23f. Configure here.}