Fix: preserve A2A3 early dispatch on queue overflow - #1812
Open
sjduan wants to merge 1 commit into
Open
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughThe change expands the early-dispatch queue and adds explicit overflow handling. Failed initial enqueueing restores candidates to ChangesEarly-dispatch overflow handling
Estimated code review effort: 2 (Simple) | ~10 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
zhangqi-chen
pushed a commit
to hw-native-sys/pypto-lib
that referenced
this pull request
Aug 13, 2026
- Add prefill_cp_layer.py, a single-layer CP prefill entry for SWA, HCA, and CSA that composes the inline CP attention cores with four fixed-size baseline MoE waves, keeping the attention, layer-stage, and MoE communication windows separately owned. - Add prefill_cp_fwd.py, a synthetic multi-layer CP prefill forward whose schedule is SWA, SWA, a pl.range pair loop over CSA/HCA, and an import-time final-CSA branch, selected by --num-layers (6, 8, 10, or the production 43; default 8) with no Python layer unroll. - Reuse one communication bank per domain across layers, ordered by monotonically increasing epochs, with the HCA and CSA compact signal banks kept distinct from the tail, stage, and MoE banks. - Give the SWA, HCA, and CSA cores a completion_token Out param plus tail_comm_epoch and compact_comm_epoch_base scalars, and capture the leaf commit/transport TaskIds so a terminal rank_complete task fans them in through pl.system.task_dummy. - Split the dual tail-exchange wave's epoch argument into payload_epoch and comm_epoch so the cross-layer tail counters advance independently of invocation-local payload rows. - Publish the last hidden through an SPMD final clear, then an inlined hc_head and rms_norm, then a host-launched grouped device lm_head. - Keep raw KV, compressed KV, index caches, and compressor states rank-local, propagated as InOut tensors. CP == EP holds for CP2, CP4, and CP8, with TP2/TP4 final-head groups. The forward calls the attention cores and moe directly rather than the single-layer @pl.jit children, and the baseline MoE early-resolve hints are untouched, so CP does not fork the MoE scheduling policy. EP8 needs the A2A3 early-dispatch queue fix in hw-native-sys/simpler#1812: a routed expert producer can publish 128 early-dispatch candidates against a former per-shape queue capacity of 64. This builds on the CP attention cores from #889, uses synthetic weights for the multi-layer forward, and leaves real-weight serving, CP != EP, full 43-layer golden compare, and performance overlap out of scope.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
STAGINGback toNONEwhen itsinitial queue publication fails, preserving the ordinary readiness path;
Root cause
One wide producer may publish substantially more early-dispatch candidates
than there are physical cores. DeepSeek V4 EP8 routed experts expose the
boundary clearly: balanced 64-token routing creates 64 local gate/up MM
candidates and passes, while 96 or 128 tokens create a burst of 128 candidates
and can stall with the former 64-entry queue.
The scheduler also ignored
push_tagged()failure after changing a candidatefrom
NONEtoSTAGING. An unqueued candidate could therefore remain in aspeculative state without an early-dispatch drain owning it. The rollback uses
a compare-exchange so a concurrent producer release continues to own any state
transition it has already made.
The larger queue absorbs the known wide-producer burst. The fallback keeps
overflow semantically safe rather than relying on capacity alone.
Validation
test_wiring: all 37 tests pass, including the newEarlyDispatchQueueOverflowFallsBackToNormalDispatchcase.allow_early_resolvehints: PASS.finite/nonzero hidden states, and valid logits from both TP4 groups.
A diagnostic build retaining queue capacity 64 but adding the rollback made
the 96-token workload pass, while 128 tokens still stalled. This confirms that
the state fallback and sufficient publication-burst capacity are both needed.